[RFC]: Context-length-aware speculative token scheduling — extending num_speculative_tokens_per_batch_size with a context-length axis

[seongyun1104]

Summary

num_speculative_tokens_per_batch_size schedules the speculative depth K by batch size. We propose extending each entry with an optional, backward-compatible context-length range, so the runtime picks K from a (batch, ctx) table instead of a batch-only table:

// today (unchanged, still valid):
"num_speculative_tokens_per_batch_size": [[1,64,3],[65,128,1],[129,512,0]]

// proposed (5-tuple entries; 3-tuple entries mean "all context lengths"):
"num_speculative_tokens_per_batch_size": [
  [1,   64,  0,     32768, 3],
  [65,  256, 0,     768,   0],   // short-context high batch: spec off
  [65,  256, 769,   32768, 3],   // long-context high batch: spec back on
  [257, 512, 0,     32768, 0]
]

(Both tables are illustrative of the API shape, not prescriptive; K values are deployment-specific — see §Evidence for the measured Gemma-4-31B / H100 numbers this shape is drawn from.)

The scheduler change is one lookup dimension; per-K CUDA graphs and buffers are unaffected (they are keyed by K, not by table shape), so the extension is resource-neutral.

Motivation

The batch-only rule "speculation stops paying at high batch" is a short-context artifact, not a general property. We measured this on Gemma-4-31B (FP8, hybrid sliding+global attention) with its MTP drafter on a single H100 NVL 96GB, vLLM 0.23/0.24, greedy, fixed output length, 120s steady-state windows with Prometheus counter cross-checks.

1. Short-context crossover (the motivation for today's batch table) — reproduced:

client concurrency	K=0 (tok/s)	best K (tok/s)	gain
30	1,400	2,500 (K=3)	1.79×
60	2,200	3,000 (K=3)	1.36×
128	3,413	3,413	1.00× (converged)

2. Same high batch, growing context (prefix-cache-hit regime, ~98% APC hit, prefill amortized) — the gain returns and grows:

decode-time ctx (tok)	K=0 (tok/s)	K=3 (tok/s)	gain	K=0 TPOT	K=3 TPOT
~460	3,107	3,640	1.17×	81.0 ms	62.9 ms
~970	2,320	2,897	1.25×	109.3 ms	77.8 ms
~1,990	2,110	2,915	1.38×	119.9 ms	78.0 ms
~4,096	1,768	2,397	1.36× [1]	137.8 ms	96.9 ms [2]

(client concurrency 256 fixed; at concurrency 192 the ctx≈2k cell gives 1.45×.)

[1] The 4k throughput ratio is a 3rd-run warm value; APC cache accumulation nudged it 1.28→1.32→1.36 run-to-run, so 1.36× is a conservative floor. The TPOT ratio (~1.4×, stable across 3 runs) is the more robust estimate.
[2] 4k is a single point; we characterize it as onset of decline, not a decline curve — mapping ctx > 2k needs more points.

The mechanism is visible in the TPOT column: as ctx doubles (970 → 1,990), K=3 TPOT stays flat (77.8 → 78.0 ms) while K=0 TPOT keeps climbing (109 → 120 ms). Long-context decode is memory-bandwidth-bound on the per-step KV read; verifying K drafted tokens amortizes that read across K+1 tokens. The saving grows with ctx — but not without bound. The gain peaks around ctx ≈ 1.5–2k (1.38× throughput / 1.54× TPOT) and begins to recede by 4k (1.36× / ~1.42×), as K=3's own per-step cost starts rising past 2k (its flat 78 ms breaks to 97 ms), narrowing the gap to K=0. (We did not decompose that rise; candidates include the drafter's own long-context decode cost — SWA window growth, draft-layer FLOPs — and target-side KV growth.) The practical reading: speculation's sweet spot is the mid-to-long band (a few hundred to ~2k tokens); in the very-long regime the amortization gain converges to a ceiling rather than growing indefinitely. This is a stronger claim than monotone growth — it is bounded, mechanistic, and does not invite the "then why not always speculate at long ctx" objection.

Consequence: optimal K is a function of (batch, ctx) and is not separable. A batch-only table forces one K per batch tier:

• If the tier says K=0 (tuned for short-context saturation), long-context traffic at that batch loses a measured 1.2–1.45×.
• If the tier says K>0 (tuned for long-context), short-context traffic at that batch pays verify overhead for no gain (and TTFT tail inflation).

This matters most for exactly the workloads the ecosystem is optimizing for — agentic / RAG / multi-turn traffic with long, prefix-shared contexts at high concurrency. Related discussion where a member raised the "MTP may only help for small batch size" rule and we posted this datapoint: #47277.

A secondary observation reinforces that even the batch axis is currently scheduled coarsely: on the same stack, short-context K=3 stays optimal further up the batch range than a typical hand-tuned table assumes — at 110 scheduled requests K=3 (3,300 tok/s) still beats K=1 (3,185) with a healthy TTFT p99 (796 ms), only converging to K=0 at the ~128 crossover. A static per-batch table with a conservative middle tier (e.g. dropping to K=1 at 65) leaves throughput on the table. This is orthogonal to the ctx axis but points the same way: the optimal-K surface is finer than a coarse batch-only table captures, along both axes.

Proposed Change

1. Schema (vllm/config/speculative.py, v1/spec_decode/dynamic/utils.py): accept 5-tuple entries (bs_lo, bs_hi, ctx_lo, ctx_hi, K) alongside today's 3-tuples (interpreted as ctx_lo=0, ctx_hi=max_model_len). Validation extends the existing rules: inclusive ranges, bs coverage from 1, non-overlapping, and per-bs-range full ctx coverage (rectangular grid).

2. Runtime lookup (v1/core/sched/scheduler.py): today the scheduler does dynamic_sd_lookup[len(num_scheduled_tokens)]. We extend the dense lookup to two dimensions: dense[B][ctx_bucket], where the batch's context representative is the p50 of decode-time sequence lengths of the scheduled requests — information the scheduler already holds, so the new signal costs nothing. SchedulerOutput.num_spec_tokens_to_schedule is unchanged (still a scalar per step).

3. Resource neutrality: CUDA graphs and runtime buffers are keyed by the K values appearing in the schedule, not by the number of table cells (this is already how the per-K capture works). A 2D table with palette {0,1,3} captures exactly the same graphs as a 1D table with the same palette. The context axis is resource-free.

4. Docs: two clarifications we found necessary in practice:

• The table index is the per-step scheduled request count, not client concurrency; it fluctuates through admission ramps, so tier boundaries near a workload's steady running level leak (~2% draft volume in our measurement). Guidance: place boundaries outside the running-distribution tail.
• Because the amortization gain saturates around ~2k (see Motivation), ctx buckets need not be fine-grained in the long tail: a single bucket covering "≳2k" is sufficient in our data, keeping the table small.
• K=0 tiers do not currently make speculation free (drafter obligations persist; see the overhead reports in [Performance]: Qwen3.5 native MTP can be slower than no-MTP CUDA graph baseline despite good acceptance #47277 and our isolation summary there). In our measurement the penalty concentrates in TTFT (prefill/scheduling: p50 0.28s → 2.0s) rather than TPOT (+6.7%), consistent with the input-preparation hotspots reported in [Performance]: Qwen3.5 native MTP can be slower than no-MTP CUDA graph baseline despite good acceptance #47277. That is an orthogonal issue, but users sizing K=0 tiers for high-batch production should know it.

Evidence & reproducibility

Full tables, the measurement protocol (120s windows / 30s warmup / generated-token deltas / spec-counter cross-checks), harness, and an engine-agnostic reference controller (declarative (B, ctx) table + acceptance-rate correction + overload override) are public: https://github.com/seongyun1104/depthchart. MTP × DSD runtime tier switching on 0.24.0 was verified with spec-decode counters (c30 → K=3 with drafts/step ≈ 3.0; c400 → K=0 with zero drafts); the DSD docs currently note testing with Eagle/E3 only, so this doubles as an MTP datapoint (capture-path bug on 0.25 filed as #48494, backend-selection issue as #48495).

Alternatives considered

• Acceptance-driven adaptation only (SGLang --speculative-adaptive style): reacts to drafter quality but not to the regime. Two structural limits: at K=0 the acceptance signal vanishes and requires periodic probing to escape; and acceptance does not encode the ctx-dependent verify economics at all (our AR was flat 86/66/50 per position across c=30→128 while the gain moved from 1.79× to 1.00×). A (B, ctx) table's inputs never vanish; acceptance works better as a correction layer on top.
• Per-sequence K: strictly more expressive, but requires variable-K verify batching and straggler control (cf. DSDE's per-sequence SL with a cap). The per-batch 2D table is the minimal change that captures the measured effect; per-sequence can layer on later.
• Entropy/complexity signals (HeteroSpec-style): orthogonal — they estimate acceptance, not the verify-side KV-read economics that the ctx axis captures.

Prior art & scope honesty

Batch-size-conditioned K is established (this feature; SGLang ships batch-tiered candidate sets by default). "Context-aware speculative decoding" (CASD, 2024) names a different technique — drafting by retrieval from the context. We have not found prior art that uses context length as a first-class scheduling axis for K, nor measurements of the (batch, ctx) interaction; we would welcome pointers if they exist.

Known limitations of our data: one model/drafter pair (Gemma-4-31B FP8 + QAT-matched MTP head — drafter lineage alone moved acceptance 51.6% → 67.7% in our A/B, so table values are deployment-specific and need calibration; the shape of the surface is what this RFC relies on); the ctx sweep is from a prefix-cache-hit regime (miss-heavy traffic unmeasured); single GPU, TP=1.

Implementation

We are happy to contribute the patch: the change is contained (schema validation + dense-builder in dynamic/utils.py, one lookup site in scheduler.py, docs), backward compatible, and covered by the validation rules above. A second-engine replication of the mechanism (SGLang _route extension) is in progress on our side and can inform the design review.

Feedback Period

2 weeks, or as maintainers prefer.

CC

@ekagra-ranjan @benchislett @luccafong @MatthewBonanni @JumpingRain