Context-Aware Precision Scheduling for Small Language Models

The gist: A calibrated predictor for mixed-precision latency and an allocator that uses it to schedule tensor precisions as context grows.

Context-Aware Precision Scheduling for Small Language Models

Small language models used as on-device agents must balance accuracy and latency as their context grows, which is a problem addressed by calibrating a timing model from measurements at different precisions and context lengths and combining it with tensor sensitivities to choose configurations under a latency or accuracy budget Page 1.

Predicting Latency of Mixed-Precision Configurations

The paper introduces a calibrated predictor for mixed-precision latency that combines attention, feed-forward, and output costs using tensors’ expected work Page 3. This predictor estimates latency from calibrated execution costs rather than weight bytes alone, arguing that fewer bits do not always mean faster execution Page 3. The timing model is calibrated by measuring reference latency as T0(L) = a + c L/1024, where 'a' is the fitted intercept in milliseconds and 'c' is the added time per 1024 context tokens Page 3. Furthermore, timing changes are estimated by distributing group changes by weight count: ∆T(u, p) = W(u) / P v∈g ∆Tg(p), u ∈ g Page 3.

Choosing Precision at a Given Context

The allocation process involves combining timing changes with measured accuracy costs to select a supported precision pu for each unit u Page 3. Under a latency budget D, the allocator solves min X u ∆A(u, pu) s.t. T0(L) +X u ∆T(u, pu) ≤ D, where T0(L) is the reference latency and ∆T is the timing change Page 3. Under an accuracy budget instead, the allocator chooses the fastest configuration allowed by the additive accuracy estimate Page 3. The additive accuracy change is approximated as ∆A(u, p) = A(u, p) − A u, p0(u), where A u, p is the mean KL between reference and quantized next-token distributions Page 3.

Executing a Context-Dependent Schedule

The runtime executes tensor changes while retaining the KV cache, and the schedule is determined by repeating the allocation at later context lengths to predict which tensors to change and when Page 4. The prototype keeps model weights at multiple precisions in memory so that switching does not require reloading them Page 2. The switching time is included in the latency of the first token after the change, with a median switching time being 1.81 ms Page 9.

Evaluation and Results

The experiments evaluate executed schedules on a laptop GPU, including a longer-context schedule with three weight formats and a tooluse benchmark measuring full response time and correctness Page 2. The results show that the cost schedule has 16.8% lower joint KL and 3.3% lower latency than the byte schedule Page 9. Furthermore, when evaluating under an accuracy budget, Higher weight precision can improve future predictions but leaves stored keys and values unchanged Page 5. The additive estimate ranks configurations well, with Spearman correlation of 0.86 against measured joint KL Page 13.

Tool-Use Evaluation

The paper evaluates whether tensor scheduling helps an agent respond to tool-use requests using Qwen2.5-3B-Instruct on six previously unused BFCL multi-turn tasks Page 6. The cost schedule answers the same 13 turns correctly as Q8, with a median response time of 6.74 s versus 7.52 s: 10.4% lower Page 7. This demonstrates that the method provides improved latency–fidelity tradeoffs and faster tool-use responses at matching correctness Page 8.

Limitations of the Work

The paper discusses limitations such as extra weight storage, calibration for the device and software, and limits of the additive accuracy estimate, noting that Recovery is partial, and these results do not validate additive weight-plus-cache error as a reliable absolute budget constraint Page 5. The paper also notes that using a different device or software stack requires new timing calibration Page 7.

Broader Impacts

The method improves the efficiency of language-model serving, supporting lower-cost, lower-energy, and private on-device inference Page 8. The paper identifies no direct negative societal impacts specific to this precision-allocation method Page 9. The researchers hope serving systems can use the predictor and scheduling approach to choose precision as a session grows Page 8.

Conclusion

Our calibrated predictor estimates the latency of mixed-precision configurations and captures execution costs that weight bytes alone can miss Page 8. Combining context-dependent timing with tensor sensitivities gives a precision schedule that can be prepared in advance and executed while retaining the KV cache Page 8. The experiments show improved latency–fidelity tradeoffs and faster tool-use responses at matching correctness, alongside the limits of recovering cache error through weight precision Page 8. The paper reports no formal theorems requiring proof; the allocator is a standard multiple-choice knapsack solved by an exact dynamic program, whose exactness for the discretized additive model is stated in Section 3.2 Page 18. The shared core-code supplement documents the core implementation and CPU examples, but it does not reproduce every GPU or agent experiment Page 18. The paper involves no research with human subjects Page 19. The model releases, llama.cpp, WikiText, and BFCL are cited, and the supplement records source provenance and does not redistribute model weights or text corpora Page 12. The research conforms with the NeurIPS Code of Ethics; it uses only publicly available models and text corpora and involves no human subjects or sensitive data Page 19. The paper releases no models or datasets with high misuse potential; it studies the serving efficiency of existing, publicly released models Page 8. The code supplement documents setup, inputs, a CPU example, sources, checks, and omitted runtime components Page 13. The paper identifies no direct negative societal impacts specific to this precision-allocation method Page 9. The researchers hope serving systems can use the predictor and scheduling approach to choose precision as a session grows Page 8.

The gist

A calibrated predictor for mixed-precision latency and an allocator that uses it to schedule tensor precisions as context grows.

How it works

1. The predictor estimates latency from calibrated execution costs rather than weight bytes alone Page 3.

2. The timing model is calibrated by measuring reference latency as T0(L) = a + c L/1024, where 'a' is the fitted intercept in milliseconds and 'c' is the added time per 1024 context tokens Page 3.

3. Timing changes are estimated by distributing group changes by weight count: ∆T(u, p) = W(u) / P v∈g ∆Tg(p), u ∈ g Page 3.

4. Under a latency budget D, the allocator solves min X u ∆A(u, pu) s.t.

1. Bold header: Calibrated Latency Prediction for Mixed Precision

This provides a calibrated predictor for mixed-precision latency by combining a timing model from measurements at different precisions and context lengths with tensor sensitivities, allowing the system to choose configurations under a latency budget rather than relying solely on weight bytes alone.

2. Bold header: Context-Aware Tensor Scheduling

The allocator can be improved to select tensor changes based on predicted future needs, as the paper shows that fewer bits do not always mean faster execution, by precomputing which tensors should change and when under a latency budget.

3. Bold header: Dynamic Precision Allocation Under Accuracy Budget

The system can implement an allocator that solves for precision by minimizing time change subject to an accuracy constraint, as defined in Equation 7: min min pu X u ∆T(u, pu) s.t. X u ∆A(u, pu) ≤ ϵ.

4. Bold header: KV Cache Preservation during Tensor Switching

The runtime can be enhanced to execute the schedule while retaining the session’s KV cache, ensuring that switching between weight formats does not require a costly rebuild of the cache, which is crucial for maintaining performance over long sessions.

5. Bold header: Improved Tool-Use Response Time

The system can achieve faster tool-use responses by utilizing schedules that balance latency and quality, as demonstrated by the Cost schedule yielding 10.4% faster response times at an 85 ms decode budget compared to fixed configurations.

6. Bold header: Cache Error Budget Management

The system can dynamically adjust weight precision based on cache performance using the estimate: measured cache error determines the remaining weight-error budget, allowing precision to move in either direction, though the paper notes this recovery is limited.