@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 7f7a32cfec

[Spec Decode] DSpark confidence-scheduled verification (#47808)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Signed-off-by: Benjamin Chislett <chislett.ben@gmail.com> Signed-off-by: Lucas Wilkinson <wilkinson.lucas@gmail.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Benjamin Chislett <chislett.ben@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: OpenAI Codex <noreply@openai.com>

Lucas Wilkinson committed Aug 12, 2026 at 15:27 UTC 7f7a32cfec0f1bc5b73c37200b86631523a1ea8f

40 files changed +1643 -111

.buildkite/test_areas/lm_eval.yaml

+20

@@ -101,6 +101,26 @@ steps:

101 commands:

102 - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-pcp.txt

103

104 +- label: LM Eval Spec Decode (4xB200)

105 + key: lm-eval-spec-decode-4xb200

106 + timeout_in_minutes: 120

107 + device: b200-k8s

108 + num_devices: 4

109 + optional: true

110 + source_file_dependencies:

111 + - tests/evals/gsm8k/configs/DeepSeek-V4-Flash-DSpark-confidence-TP4.yaml

112 + - tests/evals/gsm8k/configs/models-spec-decode.txt

113 + - vllm/models/deepseek_v4/

114 + - vllm/model_executor/models/qwen3_dspark.py

115 + - vllm/v1/spec_decode/metrics.py

116 + - vllm/v1/worker/gpu/spec_decode/adaptive_verification.py

117 + - vllm/v1/worker/gpu/spec_decode/dflash/

118 + - vllm/v1/worker/gpu/spec_decode/dspark/

119 + - vllm/v1/worker/gpu/spec_decode/rejection_sampler.py

120 + - vllm/v1/worker/gpu/spec_decode/speculator.py

121 + commands:

122 + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-spec-decode.txt

123 +

124 - label: LM Eval Large Models EP (2xB200)

125 key: lm-eval-large-models-ep-2xb200

126 timeout_in_minutes: 60

docs/features/speculative_decoding/README.md

@@ -18,6 +18,7 @@ vLLM supports a variety of methods of speculative decoding. Model-based methods

18 - [Hidden State Extraction](extract_hidden_states.md)

19 - [Custom Proposer Backend (Experimental)](#custom-proposer-backend-experimental)

20 - [Dynamic Speculative Decoding](dynamic_speculative_decoding.md)

21 +- [Adaptive Verification](adaptive_verification.md)

23 ## Method Selection at a Glance

@@ -35,6 +36,7 @@ depend on your model family, traffic pattern, hardware, and sampling settings.

41 For reproducible measurements in your environment, use

42 [`examples/features/speculative_decoding/spec_decode_offline.py`](../../../examples/features/speculative_decoding/spec_decode_offline.py)

docs/features/speculative_decoding/adaptive_verification.md new

+47

@@ -0,0 +1,47 @@

1 +# Adaptive Verification

2 +

3 +Speculative decoding buys fewer decode steps with more compute. At batch size 1 that is a good trade: the GPU is memory-bound with spare compute, so the extra draft tokens are close to free. At batch size 256 it is a much more delicate one. Draft tokens now compete with real tokens for the same compute, and every rejected token is compute wasted; with enough of them, throughput drops.

4 +

5 +That matters because per-position acceptance decays fast. While the GPU is memory-bound that slot is effectively free and worth the gamble; once it saturates the gamble has a real throughput cost. The crossover moves with load and with workload-dependent acceptance rates, so no static `num_speculative_tokens` is right across concurrencies.

6 +

7 +Adaptive verification decides per step how much of the draft to verify instead. Every (request, position) draft slot is scored by its *survival probability*, the running product of that request's per-position confidences, and the highest-scoring slots are admitted until a global budget is spent. Slots compete across requests: position 5 of a confident request can outrank position 1 of a doubtful one, so one request keeps its full block while another could be trimmed after a token or two.

8 +

9 +The budget itself comes from a cost model profiled at startup. vLLM measures what a step costs at each shape, then picks the token count that maximizes expected accepted tokens per second.

10 +

11 +The practical effect is that one configuration holds up across the whole load range, which removes most of the need to tune `num_speculative_tokens` per deployment.

12 +

13 +## Support

14 +

15 +Adaptive verification needs per-position acceptance estimates, so today it is only supported for DSpark with a **confidence head**.

16 +

17 +## Usage

18 +

19 +It is off by default. Enable it in the speculative config:

20 +

21 +```bash

22 +vllm serve deepseek-ai/DeepSeek-V4-Flash-DSpark \

23 + --tokenizer-mode deepseek_v4 --trust-remote-code \

24 + --speculative-config '{

25 + "method": "dspark",

26 + "model": "deepseek-ai/DeepSeek-V4-Flash-DSpark",

27 + "num_speculative_tokens": 7,

28 + "draft_sample_method": "probabilistic",

29 + "enable_adaptive_verification": true

30 + }'

31 +```

32 +

33 +Set `enable_adaptive_verification: false` to verify the full block for every request.

34 +

35 +## Requirements and limitations

36 +

37 +- The attention backend must tolerate device-decided query lengths, since the CPU lengths only bound them from above. Backends that plan off the CPU lengths are excluded by the attention selector, and rejected at startup for models that hard-wire their backend.

38 +- Full cudagraphs are required: step costs are profiled from captured graphs, so `--enforce-eager` is rejected at startup.

39 +- Not supported with LoRA (the per-token LoRA mapping is built from CPU-side boundaries), pipeline parallelism (cost curves and confidences exist only on the last rank), or output logprobs (to be fixed).

40 +

41 +## Tuning the cost profile

42 +

43 +Step costs are profiled against a synthetic KV context, 8192 tokens by default. Deployments serving much longer contexts may want to raise it so the profiled step reads a more realistic amount of cache (this matters a bit less for sparse attention models like DeepSeek-v4 since the cheap indexer is the main cost that scales with context length).

44 +

45 +```bash

46 +export VLLM_ADAPTIVE_VERIFICATION_PROFILE_CONTEXT_LEN=131072

47 +```

tests/evals/gsm8k/configs/DeepSeek-V4-Flash-DSpark-confidence-TP4.yaml new

+25

@@ -0,0 +1,25 @@

1 +model_name: "deepseek-ai/DeepSeek-V4-Flash-DSpark"

2 +accuracy_threshold: 0.92

3 +# lower acceptance length since trimming drafts at higher concurrency reduces

4 +# acceptance length, this is very loose lower bound for smoke testing

5 +min_acceptance_length: 1.2

6 +num_questions: 1319

7 +num_fewshot: 5

8 +startup_max_wait_seconds: 1800

9 +server_args: >-

10 + --tokenizer-mode deepseek_v4

11 + --trust-remote-code

12 + --max-model-len 8192

13 + --tensor-parallel-size 4

14 + --enable-expert-parallel

15 + --block-size 256

16 + --gpu-memory-utilization 0.5

17 + --kv-cache-dtype fp8

18 + --attention_config.use_fp4_indexer_cache=True

19 + --max-num-batched-tokens 16384

20 + --max-num-seqs 128

21 + --speculative-config '{"method":"dspark",

22 + "model":"deepseek-ai/DeepSeek-V4-Flash-DSpark",

23 + "attention_backend":"FLASH_ATTN","num_speculative_tokens":7,

24 + "draft_sample_method":"probabilistic",

25 + "enable_adaptive_verification":true}'

tests/evals/gsm8k/configs/models-spec-decode.txt new

@@ -0,0 +1 @@

1 +DeepSeek-V4-Flash-DSpark-confidence-TP4.yaml

tests/evals/gsm8k/test_gsm8k_correctness.py

+35

@@ -12,6 +12,7 @@ pytest -s -v tests/evals/gsm8k/test_gsm8k_correctness.py \

12 import shlex

14 import pytest

15 +import requests

16 import yaml

18 from tests.utils import RemoteOpenAIServer

@@ -65,6 +66,28 @@ def run_gsm8k_eval(eval_config: dict, server_url: str) -> dict:

66 return results

69 +def get_acceptance_length(server_url: str) -> float:

70 + """Mean tokens emitted per verification step, from the server's counters.

71 +

72 + 1.0 means every draft was rejected (speculation bought nothing); the

73 + theoretical maximum is 1 + num_speculative_tokens.

74 + """

75 + response = requests.get(f"{server_url.rstrip('/').removesuffix('/v1')}/metrics")

76 + response.raise_for_status()

77 + counters: dict[str, float] = {}

78 + for line in response.text.splitlines():

79 + if line.startswith("vllm:spec_decode_num_"):

80 + name, _, value = line.partition(" ")

81 + counters[name.split("{")[0]] = float(value)

82 +

83 + num_drafts = counters.get("vllm:spec_decode_num_drafts_total", 0.0)

84 + num_accepted = counters.get("vllm:spec_decode_num_accepted_tokens_total", 0.0)

85 + assert num_drafts > 0, (

86 + "no drafts recorded; speculative decoding did not run for this config"

87 + )

88 + return 1.0 + num_accepted / num_drafts

89 +

90 +

91 def test_gsm8k_correctness(config_filename):

92 """Test GSM8K correctness for a given model configuration."""

93 eval_config = yaml.safe_load(config_filename.read_text(encoding="utf-8"))

@@ -168,4 +191,16 @@ def test_gsm8k_correctness(config_filename):

191 f"{expected_metric:.4f} - {tol:.4f} = {expected_metric - tol:.4f}"

192 )

193

194 + # Speculative configs additionally assert that drafts are actually

195 + # landing: accuracy alone passes even when every draft is rejected.

196 + min_acceptance_length = eval_config.get("min_acceptance_length")

197 + if min_acceptance_length is not None:

198 + acceptance_length = get_acceptance_length(server_url)

199 + print(f" Mean acceptance length: {acceptance_length:.3f}")

200 + print(f" Minimum acceptance length: {min_acceptance_length:.3f}")

201 + assert acceptance_length >= min_acceptance_length, (

202 + f"Acceptance length too low: {acceptance_length:.3f} < "

203 + f"{min_acceptance_length:.3f}"

204 + )

205 +

206 print(f"✅ GSM8K test passed for {eval_config['model_name']}")

tests/v1/spec_decode/test_adaptive_verification.py new

+170

@@ -0,0 +1,170 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +from types import SimpleNamespace

5 +

6 +import numpy as np

7 +

8 +from vllm.v1.worker.gpu.async_utils import StepTimingSample

9 +from vllm.v1.worker.gpu.spec_decode.adaptive_verification import (

10 + AdaptiveVerificationManager,

11 +)

12 +

13 +

14 +def make_manager(

15 + confidences: np.ndarray, verify_cost_ms: np.ndarray

16 +) -> AdaptiveVerificationManager:

17 + num_reqs, num_steps = confidences.shape

18 + manager = AdaptiveVerificationManager.__new__(AdaptiveVerificationManager)

19 + manager.num_speculative_steps = num_steps

20 + manager._stale_confidences = [SimpleNamespace(np=confidences)]

21 + manager._stale_idx = 0

22 + manager.req_states = SimpleNamespace(

23 + req_id_to_index={"low": 0, "high": 1},

24 + num_computed_tokens_np=np.ones(num_reqs, dtype=np.int32),

25 + prefill_len=SimpleNamespace(np=np.ones(num_reqs, dtype=np.int32)),

26 + )

27 + manager.cost_tables = (np.zeros(num_reqs + 1), verify_cost_ms)

28 + manager._max_total_logits = 1 << 30

29 + manager.num_bonus_tokens = 1

30 + return manager

31 +

32 +

33 +def test_budget_stops_where_marginal_drafts_stop_paying_for_themselves():

34 + # Verification is cheap up to two extra tokens, then jumps 100x; only the

35 + # highest-confidence draft is worth the cheap slot.

36 + manager = make_manager(

37 + np.array([[0.1, 0.1], [0.9, 0.9]], dtype=np.float32),

38 + np.array([1.0, 1.0, 1.0, 1.0, 100.0, 100.0, 100.0]),

39 + )

40 +

41 + manager.get_num_tokens(

42 + {"low": 3, "high": 3},

43 + {"low": [1, 2], "high": [3, 4]},

44 + )

45 + valid_drafts, num_non_draft_tokens, draft_budget = manager._batch_budget

46 +

47 + assert draft_budget == 1

48 + assert valid_drafts == {"low": 2, "high": 2}

49 + assert num_non_draft_tokens == {"low": 1, "high": 1}

50 +

51 +

52 +def test_profiled_batches_seed_cost_curves_via_consumer():

53 + manager = AdaptiveVerificationManager.__new__(AdaptiveVerificationManager)

54 + manager.req_states = SimpleNamespace(max_num_batched_tokens=4096, max_num_reqs=64)

55 + manager.num_speculative_steps = 7

56 + manager.num_bonus_tokens = 1

57 + curves: dict[str, list[tuple[int, float]]] = {}

58 + manager.set_cost_curves = lambda draft, verify: curves.update(

59 + draft=draft, verify=verify

60 + )

61 +

62 + timings = [

63 + StepTimingSample(

64 + forward_ms=float(batch["num_tokens"]),

65 + drafter_ms=1.0,

66 + num_target_tokens=batch["num_tokens"],

67 + num_reqs=batch["num_tokens"] // 8,

68 + # Only the captured sizes replay a graph; the tail sizes run eager.

69 + full_cudagraph=batch["num_tokens"] <= 1024,

70 + )

71 + for batch in manager.batches_to_profile([8, 1024])

72 + ]

73 + manager.set_initial_cost_curves(timings)

74 +

75 + # Tail beyond the last capture size: 1.5x then doubling to the max.

76 + assert curves["verify"] == [

77 + (8, 8.0),

78 + (1024, 1024.0),

79 + (1536, 1536.0),

80 + (2048, 2048.0),

81 + (4096, 4096.0),

82 + ]

83 + # Eager batches must not contribute to the draft curve: keyed by request

84 + # count they would land inside the captured range and, once made monotonic,

85 + # smear that eager cost across every larger request count.

86 + assert curves["draft"] == [(1, 1.0), (128, 1.0)]

87 +

88 +

89 +def test_compact_batch_preserves_totals_and_bounds():

90 + # The CPU placeholder layout must keep the batch total equal to the GPU

91 + # total and every verification row within decode_query_len, or downstream

92 + # CPU metadata desyncs from the reallocated GPU boundaries.

93 + manager = make_manager(

94 + np.array([[0.9, 0.9], [0.9, 0.9], [1.0, 1.0]], dtype=np.float32),

95 + np.array([1.0] * 44 + [100.0] * 3),

96 + )

97 + manager.req_states.req_id_to_index["prefill"] = 2

98 + manager.req_states.num_computed_tokens_np = np.zeros(3, dtype=np.int32)

99 + manager.req_states.prefill_len.np = np.array([0, 0, 60], dtype=np.int32)

100 + num_tokens = manager.get_num_tokens(

101 + {"low": 3, "high": 3, "prefill": 40},

102 + {"low": [1, 2], "high": [3, 4]},

103 + )

104 + scheduled = np.array([3, 3, 40], dtype=np.int32)

105 + drafts = np.array([2, 2, 0], dtype=np.int32)

106 + cu_num_logits_np = np.array([0, 3, 6, 7], dtype=np.int32)

107 + compacted, _ = manager.compact_batch(drafts, scheduled, cu_num_logits_np)

108 +

109 + assert int(compacted.sum()) == num_tokens

110 + num_steps = manager.num_speculative_steps

111 + assert (compacted[:2] <= 1 + num_steps).all()

112 + assert compacted[2] == 40

113 +

114 +

115 +def test_budget_caps_at_one_rejection_sampler_chunk():

116 + # The chunked verification path cannot address the compacted logits

117 + # layout, so the budget must keep total logits within a single chunk.

118 + manager = make_manager(

119 + np.array([[0.9, 0.9], [0.9, 0.9]], dtype=np.float32),

120 + np.ones(7),

121 + )

122 + manager._max_total_logits = 3 # 2 bonus logits + at most 1 draft

123 + manager.get_num_tokens(

124 + {"low": 3, "high": 3},

125 + {"low": [1, 2], "high": [3, 4]},

126 + )

127 + _, _, draft_budget = manager._batch_budget

128 + assert draft_budget <= 1

129 +

130 +

131 +def test_zero_budget_rebuilds_cpu_cu_num_logits():

132 + # When one bonus row per request already overflows a verification chunk, the

133 + # budget clamps to zero but the batch still needs chunking. Every capacity is

134 + # zeroed on device, so the CPU can name that layout exactly -- and must, since

135 + # _iter_request_chunks slices the compacted logits with these offsets.

136 + #

137 + # The third request is a chunked prefill (no drafts, still mid-prompt). The

138 + # runner gives *every* request num_bonus_tokens logits rows regardless

139 + # (num_logits = num_draft_tokens_per_req + num_bonus_tokens), so the rebuilt

140 + # offsets stay uniform rather than skipping non-verification rows.

141 + manager = make_manager(

142 + np.array([[0.9, 0.9], [0.9, 0.9], [1.0, 1.0]], dtype=np.float32),

143 + np.ones(64),

144 + )

145 + manager.req_states.req_id_to_index["prefill"] = 2

146 + manager.req_states.num_computed_tokens_np = np.zeros(3, dtype=np.int32)

147 + manager.req_states.prefill_len.np = np.array([0, 0, 60], dtype=np.int32)

148 + manager._max_total_logits = 2 # < 3 requests * 1 bonus token

149 +

150 + manager.get_num_tokens(

151 + {"low": 3, "high": 3, "prefill": 40},

152 + {"low": [1, 2], "high": [3, 4]},

153 + )

154 + _, _, draft_budget = manager._batch_budget

155 + assert draft_budget == 0

156 +

157 + scheduled = np.array([3, 3, 40], dtype=np.int32)

158 + drafts = np.array([2, 2, 0], dtype=np.int32)

159 + scheduled_cu_num_logits = np.array([0, 3, 6, 7], dtype=np.int32)

160 + compacted, cu_num_logits_np = manager.compact_batch(

161 + drafts, scheduled, scheduled_cu_num_logits

162 + )

163 +

164 + # One bonus row per request, matching cumsum(capacities + num_bonus_tokens)

165 + # with every capacity zeroed -- the prefill row included.

166 + expected = np.arange(4, dtype=np.int32) * manager.num_bonus_tokens

167 + assert np.array_equal(cu_num_logits_np, expected)

168 + assert cu_num_logits_np.dtype == scheduled_cu_num_logits.dtype

169 + # The prefill keeps its scheduled tokens; only drafts are dropped.

170 + assert np.array_equal(compacted, np.array([1, 1, 40], dtype=np.int32))

tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py

@@ -42,6 +42,7 @@ def mock_model_runner_with_req_states():

42 runner.prompt_logprobs_worker = None

43 runner.is_last_pp_rank = False

44 runner.pooling_runner = None

45 + runner.adaptive_verification = None

47 # Mock staged writes — they use Triton kernels that require GPU

48 runner.req_states.apply_staged_writes = Mock()

tests/v1/worker/test_gpu_batch_ordering.py

+34 -3

@@ -55,6 +55,7 @@ def _make_runner(

55 num_computed = np.array([s[0] for s in req_states.values()], dtype=np.int32)

56 runner: Any = GPUModelRunner.__new__(GPUModelRunner)

57 runner.decode_query_len = decode_query_len

58 + runner.adaptive_verification = None

59 runner.req_states = SimpleNamespace(

60 req_id_to_index={req_id: i for i, req_id in enumerate(req_states)},

61 # The runner keeps this as min(num_computed_tokens, prefill_len).

@@ -73,6 +74,7 @@ def _uniform_token_count(

74 scheduler_output = SimpleNamespace(

75 num_scheduled_tokens={req_id: query_len for req_id in req_states},

76 total_num_scheduled_tokens=query_len * len(req_states),

77 + scheduled_spec_decode_tokens={},

78 )

79 runner = _make_runner(req_states, decode_query_len=query_len)

80 _, uniform_tok_count = runner.gather_batch_req_state(scheduler_output, dummy_run)

@@ -85,6 +87,33 @@ def test_spec_decode_batch_is_uniform_decode():

87 assert _uniform_token_count(decodes, 8) == 8

90 +def test_adaptive_verification_sizes_only_batches_with_drafts():

91 + decodes = {"d0": (16, 16), "d1": (16, 16)}

92 + runner = _make_runner(decodes, decode_query_len=8)

93 + manager = SimpleNamespace(

94 + get_num_tokens=lambda _num_tokens_per_req, _draft_tokens: 12

95 + )

96 + runner.adaptive_verification = manager

97 + scheduler_output = SimpleNamespace(

98 + num_scheduled_tokens={req_id: 8 for req_id in decodes},

99 + total_num_scheduled_tokens=16,

100 + scheduled_spec_decode_tokens={},

101 + )

102 +

103 + state, uniform_tok_count = runner.gather_batch_req_state(scheduler_output, False)

104 + assert state is not None

105 + assert state.num_tokens == 16

106 + assert uniform_tok_count == 8

107 +

108 + scheduler_output.scheduled_spec_decode_tokens = {

109 + req_id: [1, 2] for req_id in decodes

110 + }

111 + state, uniform_tok_count = runner.gather_batch_req_state(scheduler_output, False)

112 + assert state is not None

113 + assert state.num_tokens == 12

114 + assert uniform_tok_count is None

115 +

116 +

117 def test_prompt_chunk_of_decode_query_len_is_not_uniform_decode():

118 # Two prompt chunks whose query length coincides with the K+1 spec-decode

119 # query length must reject the batch (issue #49918).

@@ -116,14 +145,14 @@ def test_dummy_batches_stay_uniform_decode():

145 def test_sort_batch_req_ids_no_spec():

146 # decode_query_len == 1: plain ascending order (decodes first).

147 num_tokens_per_req = {"p1": 100, "d1": 1, "p2": 7, "d2": 1}

119 - assert sort_batch_req_ids(num_tokens_per_req, 1) == ["d1", "d2", "p2", "p1"]

148 + assert sort_batch_req_ids(num_tokens_per_req, {}, 1) == ["d1", "d2", "p2", "p1"]

149

150

151 def test_sort_batch_req_ids_spec_decode():

152 # decode_query_len == 2 (MTP k=1): uniform decodes lead, then the 1-token

153 # chunked-prefill tail, then longer prefills.

154 num_tokens_per_req = {"tail": 1, "d1": 2, "p1": 100, "d2": 2}

126 - assert sort_batch_req_ids(num_tokens_per_req, 2) == ["d1", "d2", "tail", "p1"]

155 + assert sort_batch_req_ids(num_tokens_per_req, {}, 2) == ["d1", "d2", "tail", "p1"]

156

157

158 def test_spec_decodes_lead_short_prefill_tail():

@@ -131,7 +160,7 @@ def test_spec_decodes_lead_short_prefill_tail():

160 # uniform 2-token decodes as decodes even when a 1-token prefill tail is

161 # in the batch (indexer-style: require_uniform, threshold=1+k).

162 num_tokens_per_req = {"tail": 1, **{f"d{i}": 2 for i in range(8)}}

134 - req_ids = sort_batch_req_ids(num_tokens_per_req, 2)

163 + req_ids = sort_batch_req_ids(num_tokens_per_req, {}, 2)

164 query_lens = [num_tokens_per_req[r] for r in req_ids]

165 assert query_lens == [2] * 8 + [1]

166

@@ -155,6 +184,7 @@ def test_uniform_decode_uses_state_index_not_batch_position():

184 """

185 runner: Any = GPUModelRunner.__new__(GPUModelRunner)

186 runner.decode_query_len = 8

187 + runner.adaptive_verification = None

188 # State arrays in state-index order: a prefilling request, then two decodes.

189 runner.req_states = SimpleNamespace(

190 req_id_to_index={"prefilling": 0, "decode_a": 1, "decode_b": 2},

@@ -164,6 +194,7 @@ def test_uniform_decode_uses_state_index_not_batch_position():

194 scheduler_output = SimpleNamespace(

195 num_scheduled_tokens={"decode_a": 8, "decode_b": 8},

196 total_num_scheduled_tokens=16,

197 + scheduled_spec_decode_tokens={},

198 )

199

200 state, uniform_tok_count = runner.gather_batch_req_state(scheduler_output, False)

vllm/config/speculative.py

+11

@@ -238,6 +238,10 @@ class SpeculativeConfig:

238 synthetic_acceptance_rates. Only valid when rejection_sample_method is 'synthetic'.

239 Mutually exclusive with synthetic_acceptance_rates."""

240

241 + enable_adaptive_verification: bool = False

242 + """Whether to adaptively size the draft-verification budget from per-request

243 + confidence. Currently only supported for method="dspark"."""

244 +

245 @staticmethod

246 def _acceptance_length_to_rates(length: float, n: int) -> list[float]:

247 """Mean acceptance length to unconditional per-position rates, using

@@ -1014,6 +1018,9 @@ class SpeculativeConfig:

1018 self.draft_model_config.hf_config.architectures = [

1019 "DSparkDraftModel"

1020 ]

1021 + self.draft_model_config.quantization = (

1022 + self.target_model_config.quantization

1023 + )

1024 self.update_arch_()

1025 elif (

1026 self.method == "dspark"

@@ -1159,6 +1166,10 @@ class SpeculativeConfig:

1166 self.target_parallel_config, self.draft_tensor_parallel_size

1167 )

1168 )

1169 +

1170 + if self.method != "dspark" and self.enable_adaptive_verification:

1171 + raise ValueError("Adaptive verification only supported with DSpark")

1172 +

1173 return self

1174

1175 def _validate_suffix_decoding(self):

vllm/config/vllm.py

+28

@@ -2403,6 +2403,34 @@ class VllmConfig:

2403 ):

2404 unsupported.append("EAGLE3 with pipeline parallelism")

2405

2406 + if (

2407 + speculative_config.enable_adaptive_verification

2408 + and self.lora_config is not None

2409 + ):

2410 + # The per-token LoRA mapping is built from CPU placeholder boundaries,

2411 + # while the trimmed batch's true boundaries are decided on the GPU.

2412 + unsupported.append("adaptive verification with LoRA")

2413 +

2414 + if (

2415 + speculative_config.enable_adaptive_verification

2416 + and self.compilation_config.cudagraph_mode == CUDAGraphMode.NONE

2417 + ):

2418 + # The draft budget divides by step costs profiled from captured

2419 + # cudagraphs; eager execution captures none.

2420 + unsupported.append(

2421 + "adaptive verification with enforce_eager/cudagraph_mode=none"

2422 + )

2423 +

2424 + if (

2425 + speculative_config.enable_adaptive_verification

2426 + and self.parallel_config.pipeline_parallel_size > 1

2427 + ):

2428 + # Cost curves and confidences currently only exist on the last PP rank;

2429 + # earlier ranks would diverge on the trimmed batch shape.

2430 + # TODO: we should be able to support adaptive verification with PP by

2431 + # broadcasting the cost curves and confidences to all ranks.

2432 + unsupported.append("adaptive verification with pipeline parallelism")

2433 +

2434 if self.parallel_config.enable_dbo:

2435 unsupported.append("dual batch overlap")

2436

vllm/envs.py

@@ -57,6 +57,7 @@ if TYPE_CHECKING:

57 VLLM_XLA_CACHE_PATH: str = os.path.join(VLLM_CACHE_ROOT, "xla_cache")

58 VLLM_XLA_CHECK_RECOMPILATION: bool = False

59 VLLM_SPARSE_INDEXER_MAX_LOGITS_MB: int = 512

60 + VLLM_ADAPTIVE_VERIFICATION_PROFILE_CONTEXT_LEN: int = 8192

61 VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE: Literal["auto", "nccl", "shm"] = "auto"

62 VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM: bool = False

63 VLLM_USE_RAY_WRAPPED_PP_COMM: bool = True

@@ -1067,6 +1068,14 @@ environment_variables: dict[str, Callable[[], Any]] = {

1068 "VLLM_SPARSE_INDEXER_MAX_LOGITS_MB": lambda: int(

1069 os.getenv("VLLM_SPARSE_INDEXER_MAX_LOGITS_MB", "512")

1070 ),

1071 + # KV context length each adaptive-verification profiling request pretends to

1072 + # carry, so the profiled step reads a realistic amount of cache.

1073 + # Raise it for long-context deployments, where step cost is dominated by

1074 + # attention over a much larger KV cache than the default assumes.

1075 + # Clamped to max_model_len - query_len. Default: 8192 tokens

1076 + "VLLM_ADAPTIVE_VERIFICATION_PROFILE_CONTEXT_LEN": lambda: int(

1077 + os.getenv("VLLM_ADAPTIVE_VERIFICATION_PROFILE_CONTEXT_LEN", "8192")

1078 + ),

1079 # If set, the OpenAI API server will stay alive even after the underlying

1080 # AsyncLLMEngine errors and stops serving requests

1081 "VLLM_KEEP_ALIVE_ON_ENGINE_DEATH": lambda: bool(

vllm/model_executor/layers/sparse_attn_indexer.py

@@ -609,6 +609,7 @@ def sparse_attn_indexer(

609 decode_metadata.schedule_metadata,

610 max_model_len=max_model_len,

611 clean_logits=False,

612 + indices=decode_metadata.indices,

613 )

614 num_rows = logits.shape[0]

615 topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]

vllm/model_executor/models/qwen3_dspark.py

+58 -4

@@ -10,8 +10,9 @@ The parallel backbone is a standard Qwen3 decoder stack reused from the

10 DFlash Qwen3 draft (see qwen3_dflash.py). DSpark adds:

11 * ``markov_head``: low-rank V x r / r x V transition bias added to the base

12 logits, sampled left-to-right by the speculator (the sequential stage).

13 + * ``confidence_head``: per-position acceptance-probability estimate.

14 -DSparkMarkovHead is shared with the DSV4-style DSpark model.

15 +DSparkMarkovHead and DSparkConfidenceHead are shared with the DSV4-style DSpark model.

16 """

18 from collections.abc import Iterable

@@ -21,6 +22,7 @@ import torch.nn as nn

23 from vllm.config import VllmConfig

24 from vllm.logger import init_logger

25 +from vllm.model_executor.layers.linear import ReplicatedLinear

26 from vllm.model_executor.layers.logits_processor import LogitsProcessor

27 from vllm.model_executor.layers.quantization import QuantizationConfig

28 from vllm.model_executor.layers.vocab_parallel_embedding import (

@@ -102,8 +104,36 @@ class DSparkMarkovHead(nn.Module):

104 return logits.scatter_(1, index, corrected.squeeze(-1))

105

106

107 +class DSparkConfidenceHead(nn.Module):

108 + """DSpark acceptance-confidence head."""

109 +

110 + def __init__(

111 + self,

112 + input_dim: int,

113 + prefix: str,

114 + bias: bool = False,

115 + with_markov: bool = True,

116 + ) -> None:

117 + super().__init__()

118 + self.with_markov = with_markov

119 + self.proj = ReplicatedLinear(

120 + input_dim,

121 + 1,

122 + bias=bias,

123 + return_bias=False,

124 + params_dtype=torch.float32,

125 + prefix=maybe_prefix(prefix, "proj"),

126 + )

127 +

128 + def forward(self, hidden: torch.Tensor, markov_embed: torch.Tensor) -> torch.Tensor:

129 + x = (

130 + torch.cat([hidden, markov_embed], dim=-1) if self.with_markov else hidden

131 + ).float()

132 + return self.proj(x).squeeze(-1)

133 +

134 +

135 class Qwen3DSparkModel(DFlashQwen3Model):

106 - """DFlash Qwen3 backbone + DSpark Markov head."""

136 + """DFlash Qwen3 backbone + DSpark Markov / confidence heads."""

137

138 def __init__(

139 self,

@@ -126,6 +156,18 @@ class Qwen3DSparkModel(DFlashQwen3Model):

156 prefix=maybe_prefix(prefix, "markov_head"),

157 quant_config=self.quant_config,

158 )

159 + self.confidence_head: DSparkConfidenceHead | None = None

160 + if getattr(config, "enable_confidence_head", False):

161 + with_markov = getattr(config, "confidence_head_with_markov", False)

162 + input_dim = config.hidden_size

163 + if with_markov:

164 + input_dim += config.markov_rank

165 + self.confidence_head = DSparkConfidenceHead(

166 + input_dim,

167 + prefix=maybe_prefix(prefix, "confidence_head"),

168 + bias=True,

169 + with_markov=with_markov,

170 + )

171

172

173 class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM):

@@ -197,11 +239,19 @@ class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM):

239 self.logits_processor.scale,

240 )

241

242 + def compute_confidence(

243 + self, head_hidden: torch.Tensor, markov_embed: torch.Tensor

244 + ) -> torch.Tensor:

245 + """Per-position acceptance probability for each drafted token."""

246 + assert self.model.confidence_head is not None

247 + return torch.sigmoid(self.model.confidence_head(head_hidden, markov_embed))

248 +

249 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

250 model_weights = {}

251 includes_embed_tokens = False

252 includes_lm_head = False

253 includes_draft_id_mapping = False

254 + includes_confidence_head = False

255 for name, loaded_weight in weights:

256 # t2d is training-only; the draft remaps via d2t at sampling time.

257 if "t2d" in name:

@@ -215,22 +265,26 @@ class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM):

265 includes_embed_tokens = True

266 if "lm_head" in name:

267 includes_lm_head = True

268 + if "confidence_head" in name:

269 + includes_confidence_head = True

270 model_weights[name] = loaded_weight

271 # Sets has_own_embed_tokens / has_own_lm_head so load_dspark_model

272 # knows whether to keep these or alias the target's.

273 process_eagle_weight(self, name)

274

275 # mask_embedding is an unused placeholder param; DSpark masks via the vocab row.

224 - # confidence_head is not wired into inference yet; skip its weights.

276 # embed_tokens / lm_head are optional; when omitted they are shared from

277 # the target by load_dspark_model, so skip the unloaded params here.

227 - skip_substrs = ["mask_embedding", "confidence_head"]

278 + skip_substrs = ["mask_embedding"]

279 if not includes_embed_tokens:

280 skip_substrs.append("embed_tokens")

281 if not includes_lm_head:

282 skip_substrs.append("lm_head")

283 if not includes_draft_id_mapping:

284 skip_substrs.append("draft_id_to_target_id")

285 + if self.model.confidence_head is None or not includes_confidence_head:

286 + self.model.confidence_head = None

287 + skip_substrs.append("confidence_head")

288 loader = AutoWeightsLoader(self, skip_substrs=skip_substrs)

289 loader.load_weights(model_weights.items())

290 self.model._build_fused_kv_buffers()

vllm/models/deepseek_v4/amd/rocm.py

+4 -4

@@ -14,9 +14,9 @@ from vllm.forward_context import get_forward_context

14 from vllm.models.deepseek_v4.attention import DeepseekV4Attention

15 from vllm.models.deepseek_v4.common.ops import dequantize_and_gather_k_cache

16 from vllm.models.deepseek_v4.sparse_mla import (

17 - DeepseekV4FlashMLABackend,

17 DeepseekV4FlashMLAMetadata,

19 - DeepseekV4FlashMLAMetadataBuilder,

18 + DeepseekV4SparseMLABackend,

19 + DeepseekV4SparseMLAMetadataBuilder,

20 )

21 from vllm.platforms import current_platform

22 from vllm.triton_utils import tl, triton

@@ -314,7 +314,7 @@ class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata):

314 decode_swa_ragged_indptr: torch.Tensor | None = None

315

316

317 -class DeepseekV4ROCMAiterMLASparseMetadataBuilder(DeepseekV4FlashMLAMetadataBuilder):

317 +class DeepseekV4ROCMAiterMLASparseMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder):

318 def __init__(self, *args, **kwargs):

319 super().__init__(*args, **kwargs)

320 self.c128a_decode_topk_ragged_indices_buffer: torch.Tensor | None = None

@@ -435,7 +435,7 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuild

435 )

436

437

438 -class DeepseekV4ROCMAiterMLASparseBackend(DeepseekV4FlashMLABackend):

438 +class DeepseekV4ROCMAiterMLASparseBackend(DeepseekV4SparseMLABackend):

439 @staticmethod

440 def get_name() -> str:

441 return "ROCM_FLASHMLA_SPARSE_DSV4"

vllm/models/deepseek_v4/attention.py

@@ -134,6 +134,8 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):

134

135 # Provided by the platform subclass.

136 backend_cls: ClassVar[type[AttentionBackend]]

137 + # Backend for the SWA cache layer; None uses the default SWA backend.

138 + swa_backend_cls: ClassVar[type[AttentionBackend] | None] = None

139 # KV-cache per-token block format (both layouts are paged). True (default)

140 # = fp8_ds_mla (UE8M0 block-scaled fp8 packed as uint8); False = plain

141 # bf16 / per-tensor fp8 KV row. Backends can override the instance hook when

@@ -322,6 +324,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):

324 dtype=self.kv_cache_torch_dtype,

325 prefix=f"{prefix}.swa_cache",

326 cache_config=cache_config,

327 + backend_cls=self.swa_backend_cls,

328 )

329

330 # Register with compilation context for metadata lookup.

vllm/models/deepseek_v4/nvidia/dspark.py

+24 -5

@@ -39,6 +39,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

39 )

40 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

41 from vllm.model_executor.models.qwen3_dspark import (

42 + DSparkConfidenceHead,

43 DSparkMarkovHead,

44 )

45 from vllm.model_executor.models.utils import maybe_prefix

@@ -113,7 +114,7 @@ class DSparkDeepseekV4Model(nn.Module):

114 ]

115 )

116

116 - # Heads: final norm + hc_head, and the Markov head

117 + # Heads: final norm + hc_head, and the Markov + confidence heads

118 # Loaded from the "final" MTP layer weights (mtp.*) in the target checkpoint

119 self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

120 hc_dim = self.hc_mult * config.hidden_size

@@ -136,6 +137,12 @@ class DSparkDeepseekV4Model(nn.Module):

137 config.dspark_markov_rank,

138 prefix=maybe_prefix(prefix, "markov_head"),

139 )

140 + self.confidence_head: DSparkConfidenceHead | None = None

141 + if getattr(config, "enable_confidence_head", True):

142 + self.confidence_head = DSparkConfidenceHead(

143 + config.hidden_size + config.dspark_markov_rank,

144 + prefix=maybe_prefix(prefix, "confidence_head"),

145 + )

146

147 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

148 return self.embed_tokens(input_ids)

@@ -369,6 +376,13 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):

376 def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor:

377 return self.model.markov_head.bias(markov_embed, self.logits_processor)

378

379 + def compute_confidence(

380 + self, head_hidden: torch.Tensor, markov_embed: torch.Tensor

381 + ) -> torch.Tensor:

382 + """Per-position acceptance probability for each drafted token."""

383 + assert self.model.confidence_head is not None

384 + return torch.sigmoid(self.model.confidence_head(head_hidden, markov_embed))

385 +

386 # --- Weight loading ----------------------------------------------------

387

388 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

@@ -407,6 +421,7 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):

421

422 params_dict = dict(self.named_parameters())

423 loaded_params: set[str] = set()

424 + loaded_confidence_head = False

425

426 tp_size = get_tensor_model_parallel_world_size()

427 tp_rank = get_tensor_model_parallel_rank()

@@ -419,6 +434,8 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):

434 if mapped is None:

435 continue

436 name = mapped

437 + if "confidence_head." in name:

438 + loaded_confidence_head = True

439

440 # ``.scale`` -> per-method scale suffix.

441 if name.endswith(".scale"):

@@ -461,8 +478,8 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):

478 continue

479

480 # Stacked rules only apply to decoder-layer weights. Head-stack params

464 - # (main_proj/norm/hc_head/markov_head) load directly — otherwise e.g.

465 - # "markov_w1" would collide with the "w1" shard rule.

481 + # (main_proj/norm/hc_head/markov_head/confidence_head) load directly —

482 + # otherwise e.g. "markov_w1" would collide with the "w1" shard rule.

483 is_layer_param = name.startswith("model.layers.")

484 for param_name, weight_name, stacked_shard_id in stacked_params_mapping:

485 if not is_layer_param or weight_name not in name:

@@ -488,6 +505,8 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):

505 loaded_params.add(name)

506

507 self._finalize_moe()

508 + if self.model.confidence_head is not None and not loaded_confidence_head:

509 + self.model.confidence_head = None

510 logger.info_once("DSpark draft model loaded: %d params", len(loaded_params))

511 return loaded_params

512

@@ -505,8 +524,7 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):

524 return None

525 stage = int(m.group(1))

526 rest = m.group(2)

508 - # The confidence head is not wired into inference yet; drop its weights.

509 - if rest.startswith("confidence_head."):

527 + if rest.startswith("confidence_head.") and self.model.confidence_head is None:

528 return None

529 # Head-stack params live at model level (mtp.last), context combiner at

530 # model level (mtp.0); everything else is a per-layer decoder block.

@@ -516,6 +534,7 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):

534 "hc_head_base",

535 "hc_head_scale",

536 "markov_head.",

537 + "confidence_head.",

538 )

539 if rest.startswith(("main_proj.", "main_norm.")) or rest.startswith(

540 head_prefixes

vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py

+5 -8

@@ -18,8 +18,8 @@ from vllm.models.deepseek_v4.nvidia.ops.o_proj import (

18 deep_gemm_fp8_o_proj,

19 )

20 from vllm.models.deepseek_v4.sparse_mla import (

21 - DeepseekV4FlashMLABackend,

21 DeepseekV4FlashMLAMetadata,

22 + DeepseekV4SparseMLABackend,

23 )

24 from vllm.platforms import current_platform

25 from vllm.platforms.interface import DeviceCapability

@@ -74,11 +74,10 @@ def _pad_to_supported_q_heads(num_heads: int) -> int:

74 )

77 -class DeepseekV4FlashInferMLASparseBackend(DeepseekV4FlashMLABackend):

77 +class DeepseekV4FlashInferMLASparseBackend(DeepseekV4SparseMLABackend):

78 """FlashInfer backend using the DSv4 sparse metadata/cache layout.

80 - Inheriting from the FlashMLA V4 backend reuses its

81 - ``DeepseekV4FlashMLAMetadata`` builder.

80 + Inherits the base and backend reuses its``DeepseekV4SparseMLAMetadataBuilder``

81 """

83 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]

@@ -159,7 +158,7 @@ class DeepseekV4FlashInferMLASparseBackend(DeepseekV4FlashMLABackend):

158 ) -> tuple[int, ...]:

159 device_capability = current_platform.get_device_capability()

160 if device_capability is not None and device_capability.major == 12:

162 - return DeepseekV4FlashMLABackend.get_kv_cache_shape(

161 + return DeepseekV4SparseMLABackend.get_kv_cache_shape(

162 num_blocks,

163 block_size,

164 num_kv_heads,

@@ -489,8 +488,6 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):

488 # uniform-q batches, and this avoids flattening mixed batches into one call.

489 if num_decode_tokens > 0:

490 decode_cu = query_start_loc[: num_decodes + 1]

492 - decode_cu_cpu = query_start_loc_cpu[: num_decodes + 1]

493 - decode_lens_cpu = decode_cu_cpu[1:] - decode_cu_cpu[:-1]

491 flashinfer_trtllm_batch_decode_sparse_mla_dsv4(

492 query=query[:num_decode_tokens],

493 swa_kv_cache=swa_k_cache,

@@ -504,7 +501,7 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):

501 bmm2_scale=bmm2_scale,

502 sinks=self.attn_sink,

503 cum_seq_lens_q=decode_cu,

507 - max_q_len=int(decode_lens_cpu.max().item()),

504 + max_q_len=swa_metadata.max_decode_query_len,

505 )

506

507 if num_prefill_tokens > 0:

vllm/models/deepseek_v4/nvidia/flashmla.py

+19 -1

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -from typing import TYPE_CHECKING, cast

4 +from typing import TYPE_CHECKING, ClassVar, cast

6 import torch

@@ -21,6 +21,11 @@ from vllm.models.deepseek_v4.sparse_mla import (

21 DeepseekV4FlashMLAMetadata,

22 )

23 from vllm.utils.math_utils import round_up

24 +from vllm.v1.attention.backend import AttentionCGSupport

25 +from vllm.v1.attention.backends.mla.sparse_swa import (

26 + DeepseekSparseSWABackend,

27 + DeepseekSparseSWAMetadataBuilder,

28 +)

29 from vllm.v1.attention.ops.flashmla import (

30 flash_mla_sparse_fwd,

31 flash_mla_with_kvcache,

@@ -31,10 +36,23 @@ if TYPE_CHECKING:

36 from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata

39 +class DeepseekSparseSWAFlashMLAMetadataBuilder(DeepseekSparseSWAMetadataBuilder):

40 + """SWA metadata for the FlashMLA decode path, which allows varlen decode."""

41 +

42 + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS

43 +

44 +

45 +class DeepseekSparseSWAFlashMLABackend(DeepseekSparseSWABackend):

46 + @staticmethod

47 + def get_builder_cls() -> type[DeepseekSparseSWAFlashMLAMetadataBuilder]:

48 + return DeepseekSparseSWAFlashMLAMetadataBuilder

49 +

50 +

51 class DeepseekV4FlashMLAAttention(DeepseekV4Attention):

52 """FlashMLA sparse MLA attention layer for DeepSeek V4 (CUDA)."""

54 backend_cls = DeepseekV4FlashMLABackend

55 + swa_backend_cls = DeepseekSparseSWAFlashMLABackend

57 def __init__(self, *args, **kwargs) -> None:

58 super().__init__(*args, **kwargs)

vllm/models/deepseek_v4/sparse_mla.py

+20 -10

@@ -32,8 +32,8 @@ from vllm.v1.kv_cache_interface import AttentionSpec

32 _C128A_TOPK_ALIGNMENT = 128

35 -class DeepseekV4FlashMLABackend(AttentionBackend):

36 - """DeepSeek-V4 sparse-MLA backend.

35 +class DeepseekV4SparseMLABackend(AttentionBackend):

36 + """DeepSeek-V4 sparse-MLA backend base.

38 Subclasses ``AttentionBackend`` directly (not the V3.2

39 ``FlashMLASparseBackend``): DeepSeek-V4 runs its own attention layer

@@ -54,12 +54,8 @@ class DeepseekV4FlashMLABackend(AttentionBackend):

54 return [256]

56 @staticmethod

57 - def get_name() -> str:

58 - return "FLASHMLA_SPARSE_DSV4"

59 -

60 - @staticmethod

61 - def get_builder_cls() -> type["DeepseekV4FlashMLAMetadataBuilder"]:

62 - return DeepseekV4FlashMLAMetadataBuilder

57 + def get_builder_cls() -> type["DeepseekV4SparseMLAMetadataBuilder"]:

58 + return DeepseekV4SparseMLAMetadataBuilder

60 @staticmethod

61 def get_impl_cls() -> type[Any]:

@@ -67,7 +63,7 @@ class DeepseekV4FlashMLABackend(AttentionBackend):

63 # not the generic ``Attention``/``MLAAttention`` layer, so the backend's

64 # impl class is never instantiated.

65 raise NotImplementedError(

70 - "DeepseekV4FlashMLABackend has no separate impl class; DeepSeek-V4 "

66 + "DeepseekV4SparseMLABackend has no separate impl class; DeepSeek-V4 "

67 "attention runs through DeepseekV4Attention."

68 )

@@ -131,7 +127,7 @@ class DeepseekV4FlashMLAMetadata(AttentionMetadata):

127 c128a_prefill_topk_indices: torch.Tensor | None = None

128

129

134 -class DeepseekV4FlashMLAMetadataBuilder(

130 +class DeepseekV4SparseMLAMetadataBuilder(

131 AttentionMetadataBuilder[DeepseekV4FlashMLAMetadata]

132 ):

133 _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH

@@ -294,6 +290,20 @@ class DeepseekV4FlashMLAMetadataBuilder(

290 return result

291

292

293 +class DeepseekV4FlashMLAMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder):

294 + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS

295 +

296 +

297 +class DeepseekV4FlashMLABackend(DeepseekV4SparseMLABackend):

298 + @staticmethod

299 + def get_name() -> str:

300 + return "FLASHMLA_SPARSE_DSV4"

301 +

302 + @staticmethod

303 + def get_builder_cls() -> type[DeepseekV4FlashMLAMetadataBuilder]:

304 + return DeepseekV4FlashMLAMetadataBuilder

305 +

306 +

307 def build_c128a_topk_metadata(

308 positions: torch.Tensor,

309 compress_ratio: int,

vllm/models/deepseek_v4/xpu/xpu_sparse.py

+2 -2

@@ -19,8 +19,8 @@ from vllm.models.deepseek_v4.common.ops import (

19 dequantize_and_gather_k_cache,

20 )

21 from vllm.models.deepseek_v4.sparse_mla import (

22 - DeepseekV4FlashMLABackend,

22 DeepseekV4FlashMLAMetadata,

23 + DeepseekV4SparseMLABackend,

24 )

25 from vllm.models.deepseek_v4.xpu.xpu_sparse_decode_fp8 import (

26 xpu_sparse_decode_fp8,

@@ -32,7 +32,7 @@ if TYPE_CHECKING:

32 from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata

35 -class DeepseekV4XPUSparseBackend(DeepseekV4FlashMLABackend):

35 +class DeepseekV4XPUSparseBackend(DeepseekV4SparseMLABackend):

36 @staticmethod

37 def get_name() -> str:

38 return "XPU_V4_MLA_SPARSE"

vllm/sampling_params.py

+14

@@ -908,6 +908,20 @@ class SamplingParams(

908 if speculative_config is None:

909 return

910

911 + # Adaptive verification compacts logits after the forward pass, while

912 + # compute_topk_scores uses the scheduled layout in cu_num_logits_np.

913 + # TODO(lucas): lift this restriction. The true boundaries exist on device as

914 + # cu_num_logits; cu_num_generated_tokens is only read host-side after

915 + # the logprobs D2H, so it could ride along on that copy.

916 + if (

917 + speculative_config.enable_adaptive_verification

918 + and self.num_logprobs is not None

919 + ):

920 + raise ValueError(

921 + "Output logprobs are not supported with DSpark confidence-based "

922 + "verification."

923 + )

924 +

925 # Some sampling parameters are not yet compatible with spec decoding.

926 if self.min_p > _SAMPLING_EPS or self.logit_bias:

927 raise VLLMValidationError(

vllm/utils/deep_gemm.py

+13 -2

@@ -553,7 +553,10 @@ def fp8_fp4_mqa_logits(

553

554

555 def get_paged_mqa_logits_metadata(

556 - context_lens: torch.Tensor, block_size: int, num_sms: int

556 + context_lens: torch.Tensor,

557 + block_size: int,

558 + num_sms: int,

559 + indices: torch.Tensor | None = None,

560 ) -> torch.Tensor:

561 """Build scheduling metadata for paged MQA logits.

562

@@ -562,6 +565,7 @@ def get_paged_mqa_logits_metadata(

565 per batch element.

566 block_size: KV-cache block size in tokens (e.g., 64).

567 num_sms: Number of SMs available. 132 for Hopper

568 + indices: Optional request index for each varlen row.

569

570 Returns:

571 Backend-specific tensor consumed by `fp8_fp4_paged_mqa_logits` to

@@ -570,7 +574,10 @@ def get_paged_mqa_logits_metadata(

574 _lazy_init()

575 if _get_paged_mqa_logits_metadata_impl is None:

576 return _missing()

573 - return _get_paged_mqa_logits_metadata_impl(context_lens, block_size, num_sms)

577 + kwargs = {} if indices is None else {"indices": indices}

578 + return _get_paged_mqa_logits_metadata_impl(

579 + context_lens, block_size, num_sms, **kwargs

580 + )

581

582

583 def fp8_fp4_paged_mqa_logits(

@@ -582,6 +589,7 @@ def fp8_fp4_paged_mqa_logits(

589 schedule_metadata: torch.Tensor,

590 max_model_len: int,

591 clean_logits: bool,

592 + indices: torch.Tensor | None = None,

593 ) -> torch.Tensor:

594 """Compute MQA logits using a paged KV-cache.

595

@@ -606,6 +614,7 @@ def fp8_fp4_paged_mqa_logits(

614 used to distribute work across SMs.

615 max_model_len: Maximum sequence length used to size the logits output.

616 clean_logits: Whether to clean the unfilled logits into `-inf`.

617 + indices: Optional request index for each varlen row.

618

619 Returns:

620 Logits tensor of shape [B * next_n, max_model_len], dtype

@@ -614,6 +623,7 @@ def fp8_fp4_paged_mqa_logits(

623 _lazy_init()

624 if _fp8_fp4_paged_mqa_logits_impl is None:

625 return _missing()

626 + kwargs = {} if indices is None else {"indices": indices}

627 return _fp8_fp4_paged_mqa_logits_impl(

628 q,

629 kv_cache,

@@ -623,6 +633,7 @@ def fp8_fp4_paged_mqa_logits(

633 schedule_metadata,

634 max_model_len,

635 clean_logits=clean_logits,

636 + **kwargs,

637 )

638

639

vllm/v1/attention/backend.py

+35 -10

@@ -16,7 +16,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

16 kFp8StaticTensorSym,

17 kNvfp4Dynamic,

18 )

19 -from vllm.utils.torch_utils import np_to_pinned_tensor

20 if TYPE_CHECKING:

21 from vllm.config import VllmConfig

@@ -281,6 +280,21 @@ class AttentionBackend(ABC):

280 def supports_kv_connector(cls) -> bool:

281 return True

282

283 + @classmethod

284 + def supports_device_cpu_query_lens_mismatch(cls) -> bool:

285 + """Whether this backend can run a batch whose device query_start_loc disagrees

286 + with the CPU one; backends that plan off the CPU query lengths must opt out.

287 +

288 + Currently only verification requests are affected: adaptive verification trims

289 + their drafts on device. On the CPU the draft budget is evenly distributed across

290 + requests, so the total draft budget, the decode/prefill split point and the CPU

291 + prefill query lengths all stay correct.

292 +

293 + SSM backends opt out: their recurrent-state planning is built from the CPU

294 + per-request boundaries, which the trimmed batch no longer matches.

295 + """

296 + return not cls.is_ssm()

297 +

298 @classmethod

299 def supports_pcp(cls) -> bool:

300 try:

@@ -335,6 +349,7 @@ class AttentionBackend(ABC):

349 use_batch_invariant: bool = False,

350 use_kv_connector: bool = False,

351 use_pcp: bool = False,

352 + use_adaptive_verification: bool = False,

353 ) -> list[str]:

354 invalid_reasons = []

355 if not cls.supports_head_size(head_size):

@@ -377,6 +392,14 @@ class AttentionBackend(ABC):

392 invalid_reasons.append("KV connector not supported")

393 if use_pcp and not cls.supports_pcp():

394 invalid_reasons.append("PCP not supported")

395 + if (

396 + use_adaptive_verification

397 + and not cls.supports_device_cpu_query_lens_mismatch()

398 + ):

399 + invalid_reasons.append(

400 + "device-cpu query lens mismatch not supported, "

401 + "this is needed for adaptive verification"

402 + )

403 combination_reason = cls.supports_combination(

404 head_size,

405 dtype,

@@ -549,17 +572,19 @@ class CommonAttentionMetadata:

572 assert self._token_to_req_indices_cache.shape[0] >= num_tokens

573 return self._token_to_req_indices_cache[:num_tokens]

574

552 - starts = np.asarray(self.query_start_loc_cpu, dtype=np.int32)

553 - query_lens = np.diff(starts)

554 - token_to_req_indices = np.repeat(

555 - np.arange(query_lens.shape[0], dtype=np.int32), query_lens

556 - )

557 - num_mapped_tokens = token_to_req_indices.shape[0]

575 + # Built from the device query_start_loc: adaptive verification decides the

576 + # per-request draft split on device, so the CPU copy carries the right total

577 + # but not the right per-request boundaries. Padding requests have a query

578 + # length of zero and drop out of the repeat.

579 + num_mapped_tokens = int(self.query_start_loc_cpu[-1])

580 + query_lens = self.query_start_loc[1:] - self.query_start_loc[:-1]

581 assert buffer.shape[0] >= max(num_mapped_tokens, num_tokens)

559 - # copy from CPU to GPU

560 - buffer[:num_mapped_tokens].copy_(

561 - np_to_pinned_tensor(token_to_req_indices), non_blocking=True

582 + token_to_req_indices = torch.repeat_interleave(

583 + torch.arange(query_lens.shape[0], dtype=torch.int32, device=buffer.device),

584 + query_lens,

585 + output_size=num_mapped_tokens,

586 )

587 + buffer[:num_mapped_tokens].copy_(token_to_req_indices)

588 if num_mapped_tokens < num_tokens:

589 buffer[num_mapped_tokens:num_tokens].zero_()

590 self._token_to_req_indices_cache = buffer[: max(num_mapped_tokens, num_tokens)]

vllm/v1/attention/backends/flashinfer.py

@@ -428,6 +428,12 @@ class FlashInferBackend(AttentionBackend):

428 def supports_non_causal(cls) -> bool:

429 return True

430

431 + @classmethod

432 + def supports_device_cpu_query_lens_mismatch(cls) -> bool:

433 + # The wrappers are planned from qo_indptr_cpu, so the CPU query offsets

434 + # have to be the ones the kernel runs on.

435 + return False

436 +

437 @classmethod

438 def supports_sliding_window(cls) -> bool:

439 return True

vllm/v1/attention/backends/mla/indexer.py

+71 -24

@@ -128,6 +128,12 @@ class DeepseekV32IndexerBackend(AttentionBackend):

128 def supports_pcp(cls) -> bool:

129 return True

130

131 + @classmethod

132 + def supports_device_cpu_query_lens_mismatch(cls) -> bool:

133 + # Only the varlen paged MQA logits kernel takes per-request query

134 + # lengths from device tensors; otherwise the indexer needs uniform ones.

135 + return _supports_varlen_paged_mqa_logits()

136 +

137 @staticmethod

138 def get_name() -> str:

139 return "DEEPSEEK_V32_INDEXER"

@@ -418,6 +424,7 @@ class DeepSeekV32IndexerDecodeMetadata:

424 requires_padding: bool

425 schedule_metadata: torch.Tensor

426 global_seq_lens: torch.Tensor | None = None

427 + indices: torch.Tensor | None = None

428

429

430 @dataclass

@@ -452,6 +459,14 @@ def get_max_prefill_buffer_size(vllm_config: VllmConfig):

459 return max_model_len * 40

460

461

462 +def _supports_varlen_paged_mqa_logits() -> bool:

463 + return (

464 + current_platform.is_cuda()

465 + and current_platform.is_device_capability_family(100)

466 + and has_deep_gemm()

467 + )

468 +

469 +

470 class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

471 # The indexer opts out of the shared reorder-threshold vote (see __init__),

472 # so this is None; its own split uses self.decode_threshold.

@@ -464,6 +479,8 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

479 vllm_config: VllmConfig,

480 kv_cache_spec: KVCacheSpec,

481 ) -> AttentionCGSupport:

482 + if _supports_varlen_paged_mqa_logits():

483 + return AttentionCGSupport.ALWAYS

484 return AttentionCGSupport.UNIFORM_BATCH

485

486 def __init__(self, *args, block_table_width: int, **kwargs) -> None:

@@ -515,10 +532,12 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

532 self.use_flattening = not current_platform.is_device_capability_family(

533 100

534 ) and next_n not in (1, 2)

535 + self.supports_varlen = _supports_varlen_paged_mqa_logits()

536 logger.info_once(

519 - "DSA indexer decode path: use_flattening=%s "

537 + "DSA indexer decode path: use_flattening=%s supports_varlen=%s "

538 "(next_n=%d, use_fp4_indexer_cache=%s)",

539 self.use_flattening,

540 + self.supports_varlen,

541 next_n,

542 self.use_fp4_indexer_cache,

543 )

@@ -547,6 +566,11 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

566 dtype=torch.int32,

567 device=self.device,

568 )

569 + self.decode_indices_buffer = torch.zeros(

570 + (scheduler_config.max_num_batched_tokens,),

571 + dtype=torch.int32,

572 + device=self.device,

573 + )

574 self.arange_buffer = torch.arange(

575 max(

576 scheduler_config.max_num_seqs * next_n,

@@ -626,25 +650,16 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

650 next_n: int,

651 max_decode_len: int,

652 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, bool]:

629 - """Expand seq_lens/block_table/decode_lens for the decode kernels.

630 -

631 - Flatten path (not use_native, max_decode_len > 1):

632 - Each multi-token decode request is expanded into individual

633 - single-token entries so the kernel always sees next_n=1.

634 -

635 - Native path (use_native or max_decode_len == 1):

636 - Plain decode or spec-decode with 2D per-token context lengths.

637 -

638 - Returns (seq_lens, block_table, decode_lens, batch_size, requires_padding).

639 - seq_lens is 1D (batch_size,) for flatten/plain, 2D (B, max_decode_len)

640 - for native MTP.

641 - """

653 + """Prepare native or per-token flattened decode tensors."""

654 min_decode_len = int(decode_lens_cpu.min().item())

643 - if not use_native and max_decode_len > 1:

655 + if not use_native:

656 assert self.decode_seq_lens_buffer.dim() == 1

645 - if min_decode_len == max_decode_len:

646 - # Uniform decode lengths.

647 - num_decode_tokens = num_decodes * max_decode_len

657 + if (

658 + not self.supports_varlen

659 + and min_decode_len == max_decode_len

660 + and num_decodes * max_decode_len == num_decode_tokens

661 + ):

662 + # Uniform decode lengths with no cudagraph token padding.

663 _prepare_uniform_decode_kernel[(num_decode_tokens,)](

664 seq_lens,

665 self.decode_seq_lens_buffer,

@@ -761,6 +776,28 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

776 self.global_decode_seq_lens_buffer[actual_expanded:num_decode_tokens] = 0

777 return self.global_decode_seq_lens_buffer[:num_decode_tokens]

778

779 + def _build_varlen_decode_indices(

780 + self,

781 + decode_lens: torch.Tensor,

782 + decode_lens_cpu: torch.Tensor,

783 + num_decode_tokens: int,

784 + ) -> torch.Tensor:

785 + """Build request ids for flattened SM100 varlen rows."""

786 + indices = self.decode_indices_buffer[:num_decode_tokens]

787 + actual_expanded = int(decode_lens_cpu.sum().item())

788 + num_decodes = decode_lens.shape[0]

789 + indices[:actual_expanded] = torch.repeat_interleave(

790 + self.arange_buffer[:num_decodes],

791 + decode_lens,

792 + output_size=actual_expanded,

793 + )

794 + if actual_expanded < num_decode_tokens:

795 + pad = num_decode_tokens - actual_expanded

796 + indices[actual_expanded:num_decode_tokens] = (

797 + num_decodes + self.arange_buffer[:pad]

798 + )

799 + return indices

800 +

801 def build(

802 self,

803 common_prefix_len: int,

@@ -775,12 +812,11 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

812 slot_mapping = common_attn_metadata.slot_mapping

813 block_table = common_attn_metadata.block_table_tensor

814 dcp_local_seq_lens = common_attn_metadata.dcp_local_seq_lens

778 -

815 num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (

816 split_decodes_and_prefills(

817 common_attn_metadata,

818 decode_threshold=self.decode_threshold,

783 - require_uniform=not self.use_flattening,

819 + require_uniform=not (self.use_flattening or self.supports_varlen),

820 treat_short_extends_as_decodes=not self.use_pcp,

821 )

822 )

@@ -888,7 +924,10 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

924

925 max_decode_len = int(decode_lens_cpu.max().item())

926 next_n = 1 + self.num_speculative_tokens

891 - use_native = not self.use_flattening and max_decode_len <= next_n

927 + use_native = (

928 + not (self.use_flattening or self.supports_varlen)

929 + and max_decode_len <= next_n

930 + )

931

932 global_seq_lens_for_decode = self._prepare_global_decode_seq_lens(

933 global_seq_lens=global_seq_lens_for_decode,

@@ -900,6 +939,14 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

939 max_decode_len=max_decode_len,

940 )

941

942 + decode_indices = None

943 + if self.supports_varlen:

944 + decode_indices = self._build_varlen_decode_indices(

945 + decode_lens=decode_lens,

946 + decode_lens_cpu=decode_lens_cpu,

947 + num_decode_tokens=num_decode_tokens,

948 + )

949 +

950 seq_lens, block_table, decode_lens, batch_size, requires_padding = (

951 self._prepare_decode_tensors(

952 seq_lens=seq_lens,

@@ -915,9 +962,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

962 )

963 )

964

918 - seq_lens_is_buffer_view = (use_native and next_n > 1) or (

919 - not use_native and max_decode_len > 1

920 - )

965 + seq_lens_is_buffer_view = not use_native or next_n > 1

966

967 # DCP: localize the now-expanded per-token global bounds to this

968 # rank's owned KV. Done here (after expansion) so each token's global

@@ -952,6 +997,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

997 seq_lens,

998 self.kv_cache_spec.storage_block_size,

999 self.num_sms,

1000 + indices=decode_indices,

1001 )

1002

1003 decode_metadata = DeepSeekV32IndexerDecodeMetadata(

@@ -960,6 +1006,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

1006 decode_lens=decode_lens,

1007 requires_padding=requires_padding,

1008 schedule_metadata=self.scheduler_metadata_buffer,

1009 + indices=decode_indices,

1010 global_seq_lens=global_seq_lens_for_decode,

1011 )

1012

vllm/v1/attention/backends/mla/sparse_swa.py

+27 -1

@@ -61,8 +61,10 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):

61 dtype: torch.dtype,

62 prefix: str,

63 cache_config: CacheConfig,

64 + backend_cls: "type[AttentionBackend] | None" = None,

65 ):

66 super().__init__()

67 + self.backend_cls = backend_cls or DeepseekSparseSWABackend

68 self.kv_cache = torch.tensor([])

69 self.head_dim = head_dim

70 self.window_size = window_size

@@ -104,7 +106,7 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):

106 def forward(self): ...

107

108 def get_attn_backend(self) -> type[AttentionBackend]:

107 - return DeepseekSparseSWABackend

109 + return self.backend_cls

110

111

112 class DeepseekSparseSWABackend(AttentionBackend):

@@ -183,6 +185,7 @@ class DeepseekSparseSWAMetadata:

185 num_prefills: int = 0

186 num_decode_tokens: int = 0

187 num_prefill_tokens: int = 0

188 + max_decode_query_len: int = 1

189

190 # Pre-computed prefill metadata shared across all DeepseekV4 attention layers.

191 prefill_seq_lens: torch.Tensor | None = None

@@ -644,6 +647,13 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):

647 num_prefills=num_prefills,

648 num_decode_tokens=num_decode_tokens,

649 num_prefill_tokens=num_prefill_tokens,

650 + # Upper bound on decode-split rows for the kernel's max_q_len

651 + # hint. common max_query_len bounds every row (scheduled max under

652 + # adaptive verification), clamped to what the split can admit so a

653 + # mixed batch's prefill max does not inflate decode scheduling.

654 + max_decode_query_len=min(

655 + common_attn_metadata.max_query_len, self.decode_threshold

656 + ),

657 tile_sched_swaonly=tile_sched[_LAYER_TYPE_SWAONLY],

658 tile_sched_c4a=tile_sched[_LAYER_TYPE_C4A],

659 tile_sched_c128a=tile_sched[_LAYER_TYPE_C128A],

@@ -787,6 +797,14 @@ def _compute_swa_indices_and_lens_kernel(

797 is_valid = tl.load(is_valid_token_ptr + token_idx)

798 if not is_valid:

799 tl.store(swa_lens_ptr + pid, 0)

800 + # Clear the row so a padded token cannot gather through stale indices.

801 + for i in range(0, window_size, TRITON_BLOCK_SIZE):

802 + offset = i + tl.arange(0, TRITON_BLOCK_SIZE)

803 + tl.store(

804 + swa_indices_ptr + pid * swa_indices_stride + offset,

805 + -1,

806 + mask=offset < window_size,

807 + )

808 return

809

810 req_idx = tl.load(token_to_req_indices_ptr + token_idx)

@@ -853,6 +871,14 @@ def _compute_dspark_noncausal_swa_indices_kernel(

871 is_valid = tl.load(is_valid_token_ptr + token_idx)

872 if not is_valid:

873 tl.store(swa_lens_ptr + pid, 0)

874 + # Clear the row so a padded token cannot gather through stale indices.

875 + for i in range(0, index_width, TRITON_BLOCK_SIZE):

876 + offset = i + tl.arange(0, TRITON_BLOCK_SIZE)

877 + tl.store(

878 + swa_indices_ptr + pid * swa_indices_stride + offset,

879 + -1,

880 + mask=offset < index_width,

881 + )

882 return

883

884 req_idx = tl.load(token_to_req_indices_ptr + token_idx)

vllm/v1/attention/selector.py

+15

@@ -37,6 +37,7 @@ class AttentionSelectorConfig(NamedTuple):

37 use_batch_invariant: bool = False

38 use_kv_connector: bool = False

39 use_pcp: bool = False

40 + use_adaptive_verification: bool = False

42 def __repr__(self):

43 return (

@@ -54,6 +55,7 @@ class AttentionSelectorConfig(NamedTuple):

55 f"use_non_causal={self.use_non_causal}, "

56 f"use_batch_invariant={self.use_batch_invariant}, "

57 f"use_kv_connector={self.use_kv_connector}, "

58 + f"use_adaptive_verification={self.use_adaptive_verification}, "

59 f"use_pcp={self.use_pcp})"

60 )

@@ -136,6 +138,18 @@ def get_attn_backend(

138 kv_transfer_config is not None and kv_transfer_config.is_kv_transfer_instance

139 )

140

141 + speculative_config = vllm_config.speculative_config

142 + use_adaptive_verification = (

143 + speculative_config is not None

144 + and speculative_config.enable_adaptive_verification

145 + )

146 + if use_adaptive_verification:

147 + from vllm.compilation.backends import model_tag

148 +

149 + # The drafter always runs full-length blocks; only the verifier sees

150 + # the trimmed, device-decided query lengths.

151 + use_adaptive_verification = model_tag != "dspark_head"

152 +

153 attn_type = attn_type or AttentionType.DECODER

154 attn_selector_config = AttentionSelectorConfig(

155 head_size=head_size,

@@ -153,6 +167,7 @@ def get_attn_backend(

167 use_batch_invariant=envs.VLLM_BATCH_INVARIANT,

168 use_kv_connector=use_kv_connector,

169 use_pcp=vllm_config.parallel_config.prefill_context_parallel_size > 1,

170 + use_adaptive_verification=use_adaptive_verification,

171 )

172

173 # A per-KV-group override (keyed by KVCacheSpecKind) takes precedence over

vllm/v1/worker/gpu/async_utils.py

+94

@@ -1,6 +1,9 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import contextlib

4 +from collections.abc import Iterator

5 +from dataclasses import dataclass, field

6 +from typing import TYPE_CHECKING

8 import numpy as np

9 import torch

@@ -17,6 +20,97 @@ from vllm.v1.outputs import (

20 from vllm.v1.worker.gpu.sample.output import SamplerOutput

21 from vllm.v1.worker.utils import raise_if_nan_logits

23 +if TYPE_CHECKING:

24 + from vllm.v1.worker.gpu.input_batch import InputBatch

25 +

26 +

27 +@dataclass(frozen=True)

28 +class StepTimingSample:

29 + forward_ms: float

30 + drafter_ms: float

31 + num_target_tokens: int

32 + num_reqs: int

33 + full_cudagraph: bool

34 +

35 +

36 +def _timing_event() -> torch.cuda.Event:

37 + return torch.cuda.Event(enable_timing=True)

38 +

39 +

40 +@dataclass

41 +class StepTimingEvents:

42 + forward_start: torch.cuda.Event = field(default_factory=_timing_event)

43 + forward_end: torch.cuda.Event = field(default_factory=_timing_event)

44 + drafter_start: torch.cuda.Event = field(default_factory=_timing_event)

45 + drafter_end: torch.cuda.Event = field(default_factory=_timing_event)

46 +

47 +

48 +class StepTimingCollector:

49 + """Times the steps run inside ``collect``; record calls no-op outside it.

50 +

51 + Every step gets its own events, so the steps queue back-to-back and the

52 + block resolves them all behind one sync on the way out.

53 + """

54 +

55 + def __init__(self):

56 + self._collecting = False

57 + self._step: StepTimingEvents | None = None

58 + self._batch = (False, 0, 0)

59 + self._timed: list[tuple[StepTimingEvents, tuple[bool, int, int]]] = []

60 +

61 + @contextlib.contextmanager

62 + def collect(self) -> Iterator[list[StepTimingSample]]:

63 + """Time every step run in this block.

64 +

65 + The yielded list holds one sample per timed step once the block exits;

66 + it stays empty inside the block, where the timings are still on device.

67 + """

68 + samples: list[StepTimingSample] = []

69 + self._collecting = True

70 + try:

71 + yield samples

72 + finally:

73 + self._collecting = False

74 + timed, self._timed, self._step = self._timed, [], None

75 + if not timed:

76 + return

77 + # Same stream, issue order: once the last step is done, so are the rest.

78 + timed[-1][0].drafter_end.synchronize()

79 + samples.extend(

80 + StepTimingSample(

81 + events.forward_start.elapsed_time(events.forward_end),

82 + events.drafter_start.elapsed_time(events.drafter_end),

83 + num_target_tokens,

84 + num_reqs,

85 + full_cudagraph,

86 + )

87 + for events, (full_cudagraph, num_target_tokens, num_reqs) in timed

88 + )

89 +

90 + def record_batch(self, input_batch: "InputBatch", full_cudagraph: bool) -> None:

91 + """Costs from different execution modes must not share a cost curve."""

92 + self._batch = (full_cudagraph, input_batch.num_tokens, input_batch.num_reqs)

93 +

94 + def forward_start(self) -> None:

95 + if self._collecting:

96 + self._step = StepTimingEvents()

97 + self._step.forward_start.record()

98 +

99 + def forward_end(self) -> None:

100 + if self._step is not None:

101 + self._step.forward_end.record()

102 +

103 + def drafter_start(self) -> None:

104 + if self._step is not None:

105 + self._step.drafter_start.record()

106 +

107 + def drafter_end(self) -> None:

108 + """Ends the step: only steps that reach here have a draft cost."""

109 + if self._step is not None:

110 + self._step.drafter_end.record()

111 + self._timed.append((self._step, self._batch))

112 + self._step = None

113 +

114

115 class AsyncOutput(AsyncModelRunnerOutput):

116 def __init__(

vllm/v1/worker/gpu/attn_utils.py

+16

@@ -180,6 +180,22 @@ def init_attn_backend(

180 return attn_groups, attn_cg_support_info, kernel_block_sizes

181

182

183 +def get_query_lens_mismatch_unsupported_backend(

184 + attn_groups: list[list[AttentionGroup]],

185 +) -> str | None:

186 + """Name the first backend needing the CPU query lengths to be exact, if any.

187 +

188 + The attention selector already excludes these when adaptive verification is

189 + enabled, but models that hard-wire their backend never consult it. See

190 + AttentionBackend.supports_device_cpu_query_lens_mismatch().

191 + """

192 + for groups in attn_groups:

193 + for group in groups:

194 + if not group.backend.supports_device_cpu_query_lens_mismatch():

195 + return group.backend.__name__

196 + return None

197 +

198 +

199 def _allocate_kv_cache(

200 kv_cache_config: KVCacheConfig, shared_layers: dict[str, str], device: torch.device

201 ):

vllm/v1/worker/gpu/cudagraph_utils.py

+43 -2

@@ -54,6 +54,9 @@ class BatchExecutionDescriptor:

54 num_tokens: int

55 num_reqs: int | None # None means no request padding is needed (PIECEWISE graphs)

56 uniform_token_count: int | None = None

57 + # Upper bound on per-request query length. Varlen decode graphs leave

58 + # uniform_token_count unset, so this is what keeps a prefill batch out of one.

59 + max_query_len: int | None = None

60 num_active_loras: int = 0

@@ -75,14 +78,21 @@ def _is_compatible(

78 num_tokens: int,

79 uniform_token_count: int | None,

80 num_active_loras: int,

81 + max_query_len: int | None,

82 ) -> bool:

83 # desc.uniform_token_count=None (PIECEWISE) can handle any uniform_token_count

84 # desc.num_reqs=None means no request padding needed (PIECEWISE)

85 + # desc.max_query_len=None means the graph does not constrain query length; a

86 + # caller that does not track max_query_len must not match one that does

87 return (

88 (

89 desc.uniform_token_count is None

90 or desc.uniform_token_count == uniform_token_count

91 )

92 + and (

93 + desc.max_query_len is None

94 + or (max_query_len is not None and desc.max_query_len >= max_query_len)

95 + )

96 and (desc.num_reqs is None or desc.num_reqs >= num_reqs)

97 and desc.num_tokens >= num_tokens

98 and desc.num_active_loras == num_active_loras

@@ -97,6 +107,7 @@ class CudaGraphManager:

107 cudagraph_mode: CUDAGraphMode,

108 decode_query_len: int,

109 lora_capture_cases: list[int] | None = None,

110 + varlen_decode: bool = False,

111 ):

112 self.vllm_config = vllm_config

113 self.device = device

@@ -105,6 +116,7 @@ class CudaGraphManager:

116 assert self.compilation_config is not None

117 self.cudagraph_mode = cudagraph_mode

118 self.decode_query_len = decode_query_len

119 + self.varlen_decode = varlen_decode

120

121 self.dp_size = vllm_config.parallel_config.data_parallel_size

122 self.tp_size = vllm_config.parallel_config.tensor_parallel_size

@@ -204,12 +216,27 @@ class CudaGraphManager:

216 else:

217 decode_query_lens = [self.decode_query_len]

218

219 + capture_varlen_decode = (

220 + separate_decode_routine and bool(decode_mode) and self.varlen_decode

221 + )

222 for num_tokens, num_active_loras in product(

223 capture_sizes, self.lora_capture_cases

224 ):

225 + # Varlen decode graphs take any mix of 1..decode_query_len tokens per

226 + # request, worst case 1 token per request (or max_num_reqs)

227 + if capture_varlen_decode and num_tokens <= max_decode_tokens:

228 + desc = BatchExecutionDescriptor(

229 + cg_mode=decode_mode,

230 + num_tokens=num_tokens,

231 + num_reqs=min(num_tokens, self.max_num_reqs),

232 + max_query_len=self.decode_query_len,

233 + num_active_loras=num_active_loras,

234 + )

235 + descs_by_mode[decode_mode].append(desc)

236 + descs_by_token_lora[(num_tokens, num_active_loras)].append(desc)

237 # Capture uniform decode specfifc graphs if required

238 # (i.e. separate decode routine)

212 - if separate_decode_routine and decode_mode:

239 + elif separate_decode_routine and decode_mode and not self.varlen_decode:

240 for decode_query_len in decode_query_lens:

241 rounded_num_tokens = round_up(num_tokens, decode_query_len)

242 rounded_num_reqs = rounded_num_tokens // decode_query_len

@@ -346,12 +373,19 @@ class CudaGraphManager:

373 compilation_counter.num_cudagraph_captured += 1

374 self._graphs_captured = True

375

376 + def captured_token_counts(self) -> list[int]:

377 + """Sorted token counts with a captured graph, ignoring LoRA variants."""

378 + return sorted(

379 + {desc.num_tokens for desc in self.graphs if desc.num_active_loras == 0}

380 + )

381 +

382 def dispatch(

383 self,

384 num_reqs: int,

385 num_tokens: int,

386 uniform_token_count: int | None,

387 num_active_loras: int,

388 + max_query_len: int | None = None,

389 ) -> BatchExecutionDescriptor:

390 """Find matching cudagraph descriptor from priority-ordered candidates."""

391

@@ -365,6 +399,7 @@ class CudaGraphManager:

399 num_tokens,

400 uniform_token_count,

401 effective_loras,

402 + max_query_len,

403 ):

404 return desc

405 return BatchExecutionDescriptor(

@@ -413,6 +448,7 @@ class ModelCudaGraphManager(CudaGraphManager):

448 cudagraph_mode: CUDAGraphMode,

449 decode_query_len: int,

450 lora_capture_cases: list[int] | None = None,

451 + varlen_decode: bool = False,

452 ):

453 super().__init__(

454 vllm_config,

@@ -420,6 +456,7 @@ class ModelCudaGraphManager(CudaGraphManager):

456 cudagraph_mode,

457 decode_query_len,

458 lora_capture_cases=lora_capture_cases,

459 + varlen_decode=varlen_decode,

460 )

461 self.hidden_states: torch.Tensor | None = None

462 self.aux_hidden_states: list[torch.Tensor] = []

@@ -483,6 +520,7 @@ class ModelCudaGraphManager(CudaGraphManager):

520 attn_groups,

521 kv_cache_config,

522 full_cudagraph=desc.cg_mode == CUDAGraphMode.FULL,

523 + max_query_len=desc.max_query_len,

524 )

525

526 # Capture with dummy rows marked as padding.

@@ -574,8 +612,11 @@ def prepare_inputs_to_capture(

612 attn_groups: list[list[AttentionGroup]],

613 kv_cache_config: KVCacheConfig,

614 full_cudagraph: bool,

615 + max_query_len: int | None = None,

616 ) -> AttentionState:

578 - input_batch = InputBatch.make_dummy(num_reqs, num_tokens, input_buffers)

617 + input_batch = InputBatch.make_dummy(

618 + num_reqs, num_tokens, input_buffers, max_query_len=max_query_len

619 + )

620 input_block_tables = block_tables.get_dummy_block_tables(num_reqs)

621 slot_mappings = block_tables.get_dummy_slot_mappings(num_tokens)

622 slot_mappings_by_layer = build_slot_mappings_by_layer(

vllm/v1/worker/gpu/dp_utils.py

+14 -1

@@ -21,6 +21,7 @@ def sync_cudagraph_and_dp_padding(

21 uniform_token_count: int | None,

22 dp_size: int,

23 dp_rank: int,

24 + max_query_len: int | None = None,

25 num_active_loras: int = 0,

26 ) -> tuple[BatchExecutionDescriptor, torch.Tensor | None]:

27 """

@@ -30,15 +31,17 @@ def sync_cudagraph_and_dp_padding(

31 """

32 assert dp_size > 1, "DP size must be greater than 1"

33 group = get_dp_group().cpu_group

33 - tensor = torch.zeros(3, dp_size, dtype=torch.int32, device="cpu")

34 + tensor = torch.zeros(4, dp_size, dtype=torch.int32, device="cpu")

35 tensor[0][dp_rank] = num_tokens

36 tensor[1][dp_rank] = desired_batch_desc.cg_mode.value

37 tensor[2][dp_rank] = uniform_token_count or 0 # (0 means None)

38 + tensor[3][dp_rank] = max_query_len or -1 # (-1 means None)

39 dist.all_reduce(tensor, group=group)

41 num_tokens_across_dp = tensor[0]

42 cg_mode_across_dp = tensor[1]

43 uniform_token_counts_across_dp = tensor[2]

44 + max_query_lens_across_dp = tensor[3]

46 if torch.all(num_tokens_across_dp == 0).item():

47 synced_desc = BatchExecutionDescriptor(

@@ -69,6 +72,12 @@ def sync_cudagraph_and_dp_padding(

72 ):

73 synced_uniform_token_count = None

75 + # Varlen decode graphs are selected by the query-length bound, so ranks must agree

76 + # on it or they pad to different token counts below.

77 + synced_max_query_len: int | None = None

78 + if bool(torch.all(max_query_lens_across_dp != -1).item()):

79 + synced_max_query_len = int(max_query_lens_across_dp.max().item())

80 +

81 # Dispatch for the final synced values, use num_reqs instead of synced_num_reqs

82 # so we don't perform request padding for PIECEWISE graphs.

83 # num_active_loras is per-rank and doesn't need cross-rank agreement.

@@ -77,6 +86,7 @@ def sync_cudagraph_and_dp_padding(

86 synced_num_tokens,

87 synced_uniform_token_count,

88 num_active_loras=num_active_loras,

89 + max_query_len=synced_max_query_len,

90 )

92 # Update num_tokens_across_dp to reflect padded size.

@@ -92,6 +102,7 @@ def dispatch_cg_and_sync_dp(

102 uniform_token_count: int | None,

103 dp_size: int,

104 dp_rank: int,

105 + max_query_len: int | None = None,

106 need_eager: bool = False,

107 num_active_loras: int = 0,

108 ) -> tuple[BatchExecutionDescriptor, torch.Tensor | None]:

@@ -112,6 +123,7 @@ def dispatch_cg_and_sync_dp(

123 num_tokens,

124 uniform_token_count,

125 num_active_loras=num_active_loras,

126 + max_query_len=max_query_len,

127 )

128

129 if dp_size == 1:

@@ -125,5 +137,6 @@ def dispatch_cg_and_sync_dp(

137 uniform_token_count,

138 dp_size,

139 dp_rank,

140 + max_query_len=max_query_len,

141 num_active_loras=num_active_loras,

142 )

vllm/v1/worker/gpu/input_batch.py

+68 -5

@@ -1,12 +1,17 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from dataclasses import dataclass

4 +from typing import TYPE_CHECKING

6 import numpy as np

7 import torch

9 from vllm.triton_utils import tl, triton

10 from vllm.utils import random_uuid

11 +from vllm.utils.math_utils import cdiv

12 +

13 +if TYPE_CHECKING:

14 + from vllm.v1.worker.gpu.block_table import BlockTables

17 class InputBuffers:

@@ -49,9 +54,10 @@ class InputBatch:

54 expanded_local_pos: torch.Tensor

56 # [num_reqs]

52 - # batch_idx -> num_scheduled_tokens

57 + # batch_idx -> num_scheduled_tokens, (upper bound when using adaptive verification)

58 num_scheduled_tokens: np.ndarray

54 - # sum(num_scheduled_tokens)

59 + # number of tokens in the batch,

60 + # may be < sum(num_scheduled_tokens) when using adaptive verification

61 num_tokens: int

62 num_tokens_after_padding: int

63 # Sum of draft tokens scheduled across requests.

@@ -101,9 +107,18 @@ class InputBatch:

107 # [num_reqs] per-request prompt length, only populated for R-SWA.

108 prompt_lens: torch.Tensor | None

109

110 + # Longest query the batch may contain. Set when a cudagraph descriptor promises

111 + # a query length this batch's own split does not reach, so attention metadata

112 + # stays valid for every replay the graph serves.

113 + max_query_len: int | None = None

114 +

115 @classmethod

116 def make_dummy(

106 - cls, num_reqs: int, num_tokens: int, input_buffers: InputBuffers

117 + cls,

118 + num_reqs: int,

119 + num_tokens: int,

120 + input_buffers: InputBuffers,

121 + max_query_len: int | None = None,

122 ) -> "InputBatch":

123 assert 0 < num_reqs <= num_tokens

124 device = input_buffers.device

@@ -115,9 +130,12 @@ class InputBatch:

130 expanded_local_pos = torch.zeros(num_reqs, dtype=torch.int32, device=device)

131

132 # Distribute the remainder evenly so that no dummy request exceeds

118 - # ceil(num_tokens / num_reqs) <= max_model_len tokens.

133 + # ceil(num_tokens / num_reqs) <= max_model_len tokens. Varlen graphs

134 + # accept any split with non-empty slots, so this shape works for them

135 + # too; attention metadata is built from the promised max_query_len.

136 base_tokens = num_tokens // num_reqs

137 num_extra = num_tokens % num_reqs

138 + assert max_query_len is None or base_tokens + (num_extra > 0) <= max_query_len

139 num_scheduled_tokens = np.full(num_reqs, base_tokens, dtype=np.int32)

140 if num_extra > 0:

141 num_scheduled_tokens[-num_extra:] += 1

@@ -150,7 +168,8 @@ class InputBatch:

168 logits_indices = query_start_loc[1:] - 1

169 cu_num_logits = torch.arange(num_reqs + 1, device=device, dtype=torch.int32)

170 cu_num_logits_np = np.arange(num_reqs + 1, dtype=np.int32)

153 - # Dummy: seq_len == query_len (fresh-prefill shape).

171 + # Copy so set_dummy_context can add context in place without touching

172 + # num_scheduled_tokens.

173 seq_lens_cpu_upper_bound = torch.from_numpy(num_scheduled_tokens.copy())

174 return cls(

175 req_ids=req_ids,

@@ -184,9 +203,53 @@ class InputBatch:

203 cu_num_logits_np=cu_num_logits_np,

204 has_structured_output_reqs=False,

205 prompt_lens=None,

206 + max_query_len=max_query_len,

207 )

208

209

210 +def set_dummy_context(

211 + input_batch: InputBatch,

212 + block_tables: "BlockTables",

213 + context_len: int,

214 + num_kv_blocks: int,

215 + max_model_len: int,

216 +) -> None:

217 + """Give each dummy request context_len of context, used when profiling step cost."""

218 + if not block_tables.input_block_tables:

219 + # Attention-free models have no KV context to fabricate.

220 + return

221 + num_reqs = input_batch.num_reqs

222 + query_len = input_batch.max_query_len or int(input_batch.num_scheduled_tokens.max())

223 + context_len = max(min(context_len, max_model_len - query_len), 0)

224 + if not context_len:

225 + return

226 +

227 + # Decode-like shape: each request continues after context_len

228 + # already-computed tokens.

229 + input_batch.seq_lens += context_len

230 + input_batch.seq_lens_cpu_upper_bound += context_len

231 + input_batch.num_computed_tokens_np.fill(context_len)

232 + input_batch.num_computed_prefill_tokens_np.fill(context_len)

233 + local_pos = np.arange(input_batch.num_tokens, dtype=np.int64) - np.repeat(

234 + input_batch.query_start_loc_np[:-1], input_batch.num_scheduled_tokens

235 + )

236 + input_batch.positions.copy_(torch.from_numpy(local_pos + context_len))

237 +

238 + seq_len = context_len + query_len

239 + for block_table, block_size, bpk in zip(

240 + block_tables.input_block_tables,

241 + block_tables.kernel_block_sizes,

242 + block_tables.blocks_per_kv_block,

243 + ):

244 + num_blocks = min(cdiv(seq_len, block_size), block_table.shape[1])

245 + # Spans are disjoint until the pool runs out, then they wrap and share

246 + # blocks: profiling only needs the reads to be realistic, not distinct.

247 + block_ids = torch.arange(

248 + num_reqs * num_blocks, dtype=block_table.dtype, device=block_table.device

249 + ) % (num_kv_blocks * bpk)

250 + block_table[:num_reqs, :num_blocks] = block_ids.view(num_reqs, num_blocks)

251 +

252 +

253 @triton.jit

254 def _prepare_prefill_inputs_kernel(

255 input_ids_ptr,

vllm/v1/worker/gpu/model_runner.py

+134 -16

@@ -72,7 +72,11 @@ from vllm.v1.outputs import (

72 from vllm.v1.worker.block_table import get_block_table_width

73 from vllm.v1.worker.cp_utils import check_attention_cp_compatibility

74 from vllm.v1.worker.gpu import pcp_manager as pcp

75 -from vllm.v1.worker.gpu.async_utils import AsyncOutput, AsyncPoolingOutput

75 +from vllm.v1.worker.gpu.async_utils import (

76 + AsyncOutput,

77 + AsyncPoolingOutput,

78 + StepTimingCollector,

79 +)

80 from vllm.v1.worker.gpu.attn_utils import (

81 build_slot_mappings_by_layer,

82 get_kv_cache_spec,

@@ -101,6 +105,7 @@ from vllm.v1.worker.gpu.input_batch import (

105 post_update_num_computed_tokens,

106 prepare_pos_seq_lens,

107 prepare_prefill_inputs,

108 + set_dummy_context,

109 )

110 from vllm.v1.worker.gpu.kv_connector import (

111 NO_OP_KV_CONNECTOR,

@@ -123,10 +128,17 @@ from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker

128 from vllm.v1.worker.gpu.sample.sampler import Sampler

129 from vllm.v1.worker.gpu.shutdown import free_before_shutdown

130 from vllm.v1.worker.gpu.spec_decode import init_speculator

131 +from vllm.v1.worker.gpu.spec_decode.adaptive_verification import (

132 + AdaptiveVerificationManager,

133 + maybe_create_adaptive_verification_manager,

134 +)

135 from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (

136 set_eagle3_aux_hidden_state_layers,

137 )

129 -from vllm.v1.worker.gpu.spec_decode.rejection_sampler import RejectionSampler

138 +from vllm.v1.worker.gpu.spec_decode.rejection_sampler import (

139 + RejectionSampler,

140 + get_max_chunk_logits,

141 +)

142 from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator

143 from vllm.v1.worker.gpu.spec_decode.utils import DraftTokensHandler

144 from vllm.v1.worker.gpu.states import RequestState

@@ -252,6 +264,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):

264 and self.speculative_config.use_multi_module_mtp()

265 else 1

266 )

267 +

268 + self.step_timing = StepTimingCollector()

269 +

270 # General request states.

271 self.req_states = RequestState(

272 max_num_reqs=self.max_num_reqs,

@@ -262,6 +277,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

277 device=self.device,

278 num_prefill_lookahead=num_prefill_lookahead,

279 )

280 + self.adaptive_verification: AdaptiveVerificationManager | None = None

281 self.input_buffers = InputBuffers(

282 max_num_reqs=self.max_num_reqs,

283 max_num_tokens=self.max_num_tokens,

@@ -401,6 +417,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

417 max_num_logits=self.max_num_reqs * self.decode_query_len,

418 vocab_size=self.vocab_size,

419 device=self.device,

420 + mask_stride=self.decode_query_len,

421 )

422

423 if self.is_pooling_model and self.is_last_pp_rank:

@@ -510,6 +527,20 @@ class GPUModelRunner(LoRAModelRunnerMixin):

527 attn_cg_support = attn_cg_support.narrow(

528 *self.model_state.get_additional_cg_support()

529 )

530 + # The speculator clears the flag at load time when the checkpoint has

531 + # no confidence head, so it holds the effective value.

532 + self.adaptive_verification = maybe_create_adaptive_verification_manager(

533 + enable_adaptive_verification=getattr(

534 + self.speculator, "enable_adaptive_verification", False

535 + ),

536 + attn_groups=self.attn_groups,

537 + attn_cg_support=attn_cg_support,

538 + req_states=self.req_states,

539 + query_start_loc=self.input_buffers.query_start_loc,

540 + num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step,

541 + max_total_logits=get_max_chunk_logits(self.vocab_size),

542 + )

543 +

544 self.block_tables = BlockTables(

545 block_sizes=block_sizes,

546 max_num_reqs=self.max_num_reqs,

@@ -532,6 +563,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):

563 initialize_mamba_ssu_backend(

564 self.vllm_config.mamba_config, self.kv_cache_config

565 )

566 + if self.adaptive_verification is not None:

567 + self.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE

568 cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes(

569 attn_cg_support.min_cg_support,

570 attn_cg_support.min_cg_attn_backend,

@@ -547,6 +580,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

580 cudagraph_mode,

581 decode_query_len=self.decode_query_len,

582 lora_capture_cases=self.lora_capture_cases,

583 + varlen_decode=self.adaptive_verification is not None,

584 )

585 check_attention_cp_compatibility(self.vllm_config)

586 if isinstance(self.speculator, DraftModelSpeculator):

@@ -594,6 +628,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

628 *args,

629 skip_attn: bool = False,

630 uniform_decode: bool = False,

631 + context_len: int = 0,

632 skip_eplb: bool = False,

633 is_profile: bool = False,

634 **kwargs,

@@ -655,6 +690,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

690 dummy_run=True,

691 skip_attn_for_dummy_run=skip_attn,

692 is_profile=is_profile,

693 + context_len=context_len,

694 )

695 self.kv_connector.set_disabled(False)

696

@@ -670,9 +706,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):

706 aux_hidden_states = self.execute_model_state.aux_hidden_states

707 self.execute_model_state = None

708

709 + self.step_timing.forward_end()

710 +

711 # dummy run the eagle speculator's propose to ensure DP/EP sync.

712 if self.speculator is not None:

713 assert self.sampler is not None

714 + self.step_timing.drafter_start()

715 mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None

716 if self.speculator.supports_mm_inputs:

717 mm_inputs = (

@@ -713,6 +752,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

752 mm_inputs=mm_inputs,

753 is_profile=is_profile,

754 )

755 + self.step_timing.drafter_end()

756

757 assert hidden_states is not None # Last PP rank always has hidden_states

758 sample_hidden_states = hidden_states[input_batch.logits_indices]

@@ -830,6 +870,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):

870 )

871 if self.speculator is not None:

872 self.speculator.capture()

873 + if self.adaptive_verification is not None:

874 + with self.step_timing.collect() as timings:

875 + for batch in self.adaptive_verification.batches_to_profile(

876 + self.cudagraph_manager.captured_token_counts()

877 + ):

878 + self._dummy_run(**batch)

879 + self.adaptive_verification.set_initial_cost_curves(timings)

880

881 end_time = time.perf_counter()

882 end_free_gpu_memory = torch.accelerator.get_memory_info()[0]

@@ -906,6 +953,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):

953 max_tokens=sampling_params.max_tokens if sampling_params else 1, # type: ignore[arg-type]

954 )

955 req_index = self.req_states.req_id_to_index[req_id]

956 + if self.adaptive_verification is not None:

957 + self.adaptive_verification.add_request(req_index)

958

959 if self.pooling_runner is not None:

960 assert new_req_data.pooling_params is not None

@@ -994,8 +1043,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1043 num_reqs, num_toks, max_query_len, has_prefill=False

1044 )

1045

1046 + draft_tokens = scheduler_output.scheduled_spec_decode_tokens

1047 # batch_idx -> req_id

998 - req_ids = sort_batch_req_ids(num_tokens_per_req, self.decode_query_len)

1048 + req_ids = sort_batch_req_ids(

1049 + num_tokens_per_req, draft_tokens, self.decode_query_len

1050 + )

1051 +

1052 numtoks_iter = map(num_tokens_per_req.__getitem__, req_ids)

1053 num_scheduled_tokens = np.fromiter(numtoks_iter, dtype=np.int32, count=num_reqs)

1054

@@ -1007,9 +1060,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1060 ]

1061 is_prefilling_np = num_computed_prefill_tokens_np < prefill_len_np

1062

1063 + if self.adaptive_verification is not None and draft_tokens:

1064 + num_toks = self.adaptive_verification.get_num_tokens(

1065 + num_tokens_per_req, draft_tokens

1066 + )

1067 +

1068 batch_state = BatchReqState(

1069 req_ids=req_ids,

1070 num_scheduled_tokens=num_scheduled_tokens,

1071 + num_tokens=num_toks,

1072 idx_mapping_np=idx_mapping_np,

1073 prefill_len_np=prefill_len_np,

1074 num_computed_prefill_tokens_np=num_computed_prefill_tokens_np,

@@ -1026,7 +1085,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1085 batch_req_state: "BatchReqState",

1086 batch_desc: BatchExecutionDescriptor,

1087 ) -> InputBatch:

1029 - num_tokens = scheduler_output.total_num_scheduled_tokens

1088 + num_tokens = batch_req_state.num_tokens

1089 num_tokens_after_padding = batch_desc.num_tokens

1090 assert num_tokens > 0

1091 if envs.VLLM_MOE_SKIP_PADDING:

@@ -1072,9 +1131,20 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1131 np.cumsum(num_logits, out=cu_num_logits_np[1:])

1132 cu_num_logits = async_copy_to_gpu(cu_num_logits_np, device=self.device)

1133

1075 - max_expand_len = self.decode_query_len

1076 - expanded_idx_mapping, expanded_local_pos = expand_idx_mapping(

1077 - idx_mapping, total_num_logits, cu_num_logits, max_expand_len

1134 + adaptive_verification = (

1135 + self.adaptive_verification if num_draft_tokens_per_req is not None else None

1136 + )

1137 + num_scheduled_tokens_upper_bound = num_scheduled_tokens_np

1138 + if adaptive_verification is not None:

1139 + # num_scheduled_tokens represents the draft budget evenly distributed across

1140 + # all verification requests, `reallocate_drafts` will unevenly assign the

1141 + # draft budget to requests on the GPU side only.

1142 + num_scheduled_tokens_np, cu_num_logits_np = (

1143 + adaptive_verification.compact_batch(

1144 + num_draft_tokens_per_req,

1145 + num_scheduled_tokens_np,

1146 + cu_num_logits_np,

1147 + )

1148 )

1149

1150 # Get query_start_loc.

@@ -1086,9 +1156,19 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1156 # Pad for full CUDA graph mode.

1157 # Some attention backends like FA3 require query_start_loc to be non-decreasing.

1158 query_start_loc_np[num_reqs + 1 :] = num_tokens

1089 - async_copy_to_gpu(query_start_loc_np, out=self.input_buffers.query_start_loc)

1159 + query_start_loc = self.input_buffers.query_start_loc

1160 + async_copy_to_gpu(query_start_loc_np, out=query_start_loc)

1161 + if adaptive_verification is not None:

1162 + cu_num_logits, query_start_loc, total_num_draft_tokens = (

1163 + adaptive_verification.reallocate_drafts(req_ids, idx_mapping)

1164 + )

1165 + total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens

1166 + if draft_tokens:

1167 + expanded_idx_mapping, expanded_local_pos = expand_idx_mapping(

1168 + idx_mapping, total_num_logits, cu_num_logits, self.decode_query_len

1169 + )

1170 query_start_loc_np = query_start_loc_np[: num_reqs_padded + 1]

1091 - query_start_loc = self.input_buffers.query_start_loc[: num_reqs_padded + 1]

1171 + query_start_loc = query_start_loc[: num_reqs_padded + 1]

1172

1173 # Get prefill tokens if any.

1174 if batch_req_state.has_prefill:

@@ -1145,7 +1225,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1225 seq_lens_cpu_upper_bound_np = np.zeros(num_reqs_padded, dtype=np.int32)

1226 np.add(

1227 num_computed_tokens_np,

1148 - num_scheduled_tokens_np,

1228 + num_scheduled_tokens_upper_bound,

1229 out=seq_lens_cpu_upper_bound_np[:num_reqs],

1230 )

1231 seq_lens_cpu_upper_bound = torch.from_numpy(seq_lens_cpu_upper_bound_np)

@@ -1168,7 +1248,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1248 idx_mapping_np=idx_mapping_np,

1249 expanded_idx_mapping=expanded_idx_mapping,

1250 expanded_local_pos=expanded_local_pos,

1171 - num_scheduled_tokens=num_scheduled_tokens_np,

1251 + num_scheduled_tokens=num_scheduled_tokens_upper_bound,

1252 num_tokens=num_tokens,

1253 num_tokens_after_padding=num_tokens_after_padding,

1254 num_draft_tokens=total_num_draft_tokens,

@@ -1192,6 +1272,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1272 cu_num_logits_np=cu_num_logits_np,

1273 has_structured_output_reqs=scheduler_output.has_structured_output_requests,

1274 prompt_lens=prompt_lens,

1275 + max_query_len=(

1276 + int(num_scheduled_tokens_upper_bound.max())

1277 + if adaptive_verification is not None

1278 + else None

1279 + ),

1280 )

1281 return pcp.maybe_partition_pcp_batch(self.pcp_manager, input_batch)

1282

@@ -1298,6 +1383,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1383 dummy_run: bool = False,

1384 skip_attn_for_dummy_run: bool = False,

1385 is_profile: bool = False,

1386 + context_len: int = 0,

1387 ) -> ModelRunnerOutput | IntermediateTensors | None:

1388 if not dummy_run:

1389 # Update the request states.

@@ -1315,9 +1401,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1401 # Get batch descriptor and sync across DP ranks.

1402 num_reqs = len(scheduler_output.num_scheduled_tokens)

1403 num_toks = scheduler_output.total_num_scheduled_tokens

1404 + max_query_len = max(scheduler_output.num_scheduled_tokens.values())

1405 batch_req_state, uniform_tok_count = self.gather_batch_req_state(

1406 scheduler_output, dummy_run

1407 )

1408 + if batch_req_state is not None:

1409 + num_toks = batch_req_state.num_tokens

1410

1411 num_active_loras = 0

1412 if self.lora_config:

@@ -1340,6 +1429,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1429 uniform_tok_count,

1430 self.dp_size,

1431 self.dp_rank,

1432 + max_query_len=max_query_len,

1433 need_eager=is_profile or skip_compiled,

1434 num_active_loras=num_active_loras,

1435 )

@@ -1378,13 +1468,23 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1468 self._set_active_loras(*lora_inputs)

1469 else:

1470 # No actual tokens to run. A dummy run for DP or memory profiling.

1471 + dummy_num_reqs = batch_desc.num_reqs or num_reqs

1472 input_batch = InputBatch.make_dummy(

1382 - batch_desc.num_reqs or num_reqs,

1473 + dummy_num_reqs,

1474 batch_desc.num_tokens,

1475 self.input_buffers,

1476 + max_query_len=batch_desc.max_query_len,

1477 )

1478 if not skip_attn_for_dummy_run:

1479 block_tables, slot_mappings = self.prepare_dummy_attn(input_batch)

1480 + if context_len:

1481 + set_dummy_context(

1482 + input_batch,

1483 + self.block_tables,

1484 + context_len,

1485 + self.kv_cache_config.num_blocks,

1486 + self.max_model_len,

1487 + )

1488 else:

1489 assert batch_desc.cg_mode != CUDAGraphMode.FULL, (

1490 "Attention metadata must be prepared for dummy runs when using "

@@ -1488,6 +1588,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1588 # Update the EPLB meta.

1589 self.eplb.prepare_forward(self.model_config, input_batch.num_tokens)

1590

1591 + self.step_timing.record_batch(

1592 + input_batch, batch_desc.cg_mode == CUDAGraphMode.FULL

1593 + )

1594 + self.step_timing.forward_start()

1595 +

1596 # Run model.

1597 if batch_desc.cg_mode == CUDAGraphMode.FULL:

1598 # Use explicit cudagraph replay for FULL mode.

@@ -1698,6 +1803,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1803 mm_inputs=mm_inputs,

1804 )

1805 self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens

1806 + if self.adaptive_verification is not None:

1807 + self.adaptive_verification.record_confidences(

1808 + self.speculator.draft_token_confidence_probs, input_batch

1809 + )

1810

1811 if self.num_speculative_steps > 0:

1812 # Spec-decode and diffusion LLMs both use draft tokens but the latter does

@@ -1839,6 +1948,9 @@ class BatchReqState(NamedTuple):

1948

1949 req_ids: list[str]

1950 num_scheduled_tokens: np.ndarray # [num_reqs]

1951 + # May be less than scheduler_output.total_num_scheduled_tokens:

1952 + # adaptive verification trims the draft budget before running.

1953 + num_tokens: int

1954 idx_mapping_np: np.ndarray # [num_reqs]

1955 prefill_len_np: np.ndarray # [num_reqs]

1956 num_computed_prefill_tokens_np: np.ndarray # [num_reqs]

@@ -1847,9 +1959,15 @@ class BatchReqState(NamedTuple):

1959

1960

1961 def sort_batch_req_ids(

1850 - num_tokens_per_req: dict[str, int], decode_query_len: int

1962 + num_tokens_per_req: dict[str, int],

1963 + draft_tokens: dict[str, list[int]],

1964 + decode_query_len: int,

1965 ) -> list[str]:

1852 - # Order decode -> short_extend -> prefill; split_decodes_and_prefills

1853 - # relies on uniform decodes (query_len == decode_query_len) leading.

1854 - key = lambda r: ((num := num_tokens_per_req[r]) != decode_query_len, num)

1966 + # Order verification/decode -> short_extend -> prefill;

1967 + # split_decodes_and_prefills relies on decode-like requests leading.

1968 + key = lambda r: (

1969 + not draft_tokens.get(r),

1970 + (num := num_tokens_per_req[r]) != decode_query_len,

1971 + num,

1972 + )

1973 return sorted(num_tokens_per_req, key=key)

vllm/v1/worker/gpu/model_states/default.py

+3 -1

@@ -144,7 +144,9 @@ class DefaultModelState(ModelState):

144 input_batch.query_start_loc_np[: num_reqs + 1]

145 )

146 query_start_loc_gpu = input_batch.query_start_loc[: num_reqs + 1]

147 - max_query_len = input_batch.num_scheduled_tokens.max().item()

147 + max_query_len = input_batch.max_query_len

148 + if max_query_len is None:

149 + max_query_len = input_batch.num_scheduled_tokens.max().item()

150 seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound

151 if for_capture:

152 # Capture with worst-case max_seq_len so the graph is valid at any replay.

vllm/v1/worker/gpu/spec_decode/adaptive_verification.py new

+477

@@ -0,0 +1,477 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Adaptive verification for DSpark speculative decoding."""

4 +

5 +from collections import defaultdict

6 +from collections.abc import Iterable, Iterator

7 +from typing import TYPE_CHECKING

8 +

9 +import numpy as np

10 +import torch

11 +

12 +import vllm.envs as envs

13 +from vllm.distributed.parallel_state import get_tp_group

14 +from vllm.logger import init_logger

15 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

16 +from vllm.v1.attention.backend import AttentionCGSupport

17 +from vllm.v1.utils import CpuGpuBuffer

18 +from vllm.v1.worker.gpu.async_utils import StepTimingSample, stream

19 +from vllm.v1.worker.gpu.attn_utils import (

20 + get_query_lens_mismatch_unsupported_backend,

21 +)

22 +from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu

23 +

24 +logger = init_logger(__name__)

25 +_PROFILE_REPLAYS = 5

26 +

27 +if TYPE_CHECKING:

28 + from vllm.v1.worker.gpu.attn_utils import AttentionCGSupportInfo

29 + from vllm.v1.worker.gpu.input_batch import InputBatch

30 + from vllm.v1.worker.gpu.states import RequestState

31 + from vllm.v1.worker.utils import AttentionGroup

32 +

33 +

34 +def _assign_draft_token_budget(

35 + confidence_probs: torch.Tensor,

36 + idx_mapping: torch.Tensor,

37 + capacities: torch.Tensor,

38 + draft_budget: int,

39 + num_steps: int,

40 +) -> None:

41 + """Admit the globally best `draft_budget` draft slots, in place.

42 +

43 + Every (request, step) slot is scored by its survival probability, the running

44 + product of that request's per-position confidences, and the highest scores win.

45 + Survival only decreases along a request, so a global top-k always admits

46 + continuously along steps with a request.

47 +

48 + `capacities` enters holding each request's scheduled draft count (which bounds its

49 + eligible slots) and leaves holding the admitted count. The caller only calls this

50 + when `draft_budget < capacities.sum()`, so every winner is a real draft slot.

51 + """

52 + survival = confidence_probs[idx_mapping].cumprod(dim=1)

53 + steps = torch.arange(num_steps, device=survival.device)

54 + # Out-of-range slots score -inf so they never outrank a real draft.

55 + out_of_range = steps[None, :] >= capacities[:, None]

56 + survival = survival.masked_fill(out_of_range, -float("inf"))

57 + flat = survival.flatten()

58 + winners = flat.topk(draft_budget).indices

59 + admitted = torch.zeros_like(flat, dtype=torch.bool).index_fill_(0, winners, True)

60 + torch.sum(admitted.view_as(survival), dim=1, dtype=capacities.dtype, out=capacities)

61 +

62 +

63 +_assign_draft_token_budget_compiled = torch.compile(

64 + _assign_draft_token_budget, dynamic=True

65 +)

66 +

67 +

68 +def build_cost_tables_from_curves(

69 + draft_curve: list[tuple[int, float]],

70 + verify_curve: list[tuple[int, float]],

71 + max_num_reqs: int,

72 + max_batch_tokens: int,

73 + cudagraph_limit: int = 0,

74 +) -> tuple[np.ndarray, np.ndarray]:

75 + """Build cost tables: graph-padded below the capture limit, smooth above.

76 +

77 + Args:

78 + cudagraph_limit: Largest cudagraph-captured size. At or below it,

79 + execution pads up to the next captured size, so cost is a step

80 + function. Above it there is no padding, so cost is continuous.

81 + """

82 +

83 + def build_table(limit: int, curve: list[tuple[int, float]]) -> np.ndarray:

84 + xs, ys = np.asarray(curve, dtype=np.float64).T

85 + ys = np.maximum.accumulate(ys)

86 + values = np.arange(limit + 1)

87 + # Execution pads to the next captured size, so cost is a step

88 + # function of the padded size: smooth interpolation would invent

89 + # marginal per-token costs that don't exist within a pad bucket.

90 + idx = np.searchsorted(xs, values, side="left")

91 + result = ys[np.minimum(idx, len(xs) - 1)]

92 + # Past the capture limit nothing pads, so cost really is continuous in

93 + # size; snapping to the next profiled point overestimates badly when

94 + # the profiled points are far apart. Interpolate only between points

95 + # that are themselves past the limit: crossing the limit loses

96 + # cudagraphs entirely, which is a genuine discontinuity, so the first

97 + # point above it must not be blended with the last one below.

98 + if cudagraph_limit:

99 + smooth = values > cudagraph_limit

100 + above = xs > cudagraph_limit

101 + if smooth.any() and above.any():

102 + result[smooth] = np.interp(values[smooth], xs[above], ys[above])

103 + if len(xs) > 1:

104 + after = values > xs[-1]

105 + slope = (ys[-1] - ys[-2]) / (xs[-1] - xs[-2])

106 + result[after] = ys[-1] + slope * (values[after] - xs[-1])

107 + return result

108 +

109 + draft_table = np.maximum(build_table(max_num_reqs, draft_curve), 0.0)

110 + verify_table = np.maximum(build_table(max_batch_tokens, verify_curve), 1e-6)

111 + return draft_table, verify_table

112 +

113 +

114 +class AdaptiveVerificationManager:

115 + def __init__(

116 + self,

117 + req_states: "RequestState",

118 + query_start_loc: torch.Tensor,

119 + num_bonus_tokens: int,

120 + max_total_logits: int,

121 + ):

122 + self.req_states = req_states

123 + self.num_speculative_steps = req_states.num_speculative_steps

124 + device = req_states.device

125 + self._copy_stream = torch.cuda.Stream(device)

126 +

127 + self.num_bonus_tokens = num_bonus_tokens

128 + # Rejection sampling verifies logits in one contiguous chunk; the

129 + # chunked path indexes by scheduled (untrimmed) offsets and cannot

130 + # address the compacted layout, so the budget must fit one chunk.

131 + self._max_total_logits = max_total_logits

132 + self.query_start_loc = query_start_loc

133 + self.cost_tables: tuple[np.ndarray, np.ndarray] | None = None

134 + # Largest cudagraph-captured token count; above it nothing pads.

135 + self._cudagraph_limit = 0

136 + self._batch_budget: tuple[dict[str, int], dict[str, int], int] | None = None

137 + max_num_reqs = req_states.max_num_reqs

138 + # Current per-slot confidences

139 + self._confidence_probs = torch.empty(

140 + (max_num_reqs, self.num_speculative_steps),

141 + dtype=torch.float32,

142 + device=device,

143 + )

144 + self._batch_draft_capacity = torch.empty(

145 + max_num_reqs, dtype=torch.int32, device=device

146 + )

147 + self._num_non_draft_tokens = torch.empty_like(query_start_loc[:-1])

148 + self._cu_num_logits = torch.empty_like(query_start_loc)

149 +

150 + # Two D2H slots preserve stale inputs for budget selection.

151 + self._stale_confidences = [

152 + CpuGpuBuffer(

153 + max_num_reqs,

154 + self.num_speculative_steps,

155 + dtype=torch.float32,

156 + device=device,

157 + )

158 + for _ in range(2)

159 + ]

160 + self._copy_events = [torch.cuda.Event(blocking=True) for _ in range(2)]

161 + self._pending_resets: list[int] = []

162 + self._stale_idx = 0

163 + for slot in self._stale_confidences:

164 + slot.np.fill(1.0)

165 +

166 + def add_request(self, req_idx: int) -> None:

167 + self._stale_confidences[self._stale_idx].np[req_idx].fill(1.0)

168 + self._pending_resets.append(req_idx)

169 + self._confidence_probs[req_idx].fill_(1.0)

170 +

171 + def batches_to_profile(self, capture_sizes: list[int]) -> Iterator[dict[str, int]]:

172 + """Dummy-run kwargs whose step timings seed the cost tables.

173 +

174 + Run these inside StepTimingCollector.collect(), then hand the block's

175 + timings to set_initial_cost_curves."""

176 + max_num_tokens = self.req_states.max_num_batched_tokens

177 + size = self._cudagraph_limit = capture_sizes[-1] if capture_sizes else 0

178 + # Also profile beyond the capture limit: real steps run there

179 + # (piecewise/eager) and linear extrapolation badly underestimates

180 + # them. These runs double as JIT warmup for the piecewise shapes.

181 + tail_sizes: set[int] = set()

182 + if size:

183 + tail_sizes.add(min(size + size // 2, max_num_tokens))

184 + while size < max_num_tokens:

185 + size = min(size * 2, max_num_tokens)

186 + tail_sizes.add(size)

187 + tail_sizes -= set(capture_sizes)

188 + for num_tokens in capture_sizes + sorted(tail_sizes):

189 + for _ in range(_PROFILE_REPLAYS):

190 + yield {

191 + "num_tokens": num_tokens,

192 + "context_len": envs.VLLM_ADAPTIVE_VERIFICATION_PROFILE_CONTEXT_LEN,

193 + }

194 +

195 + def set_initial_cost_curves(self, samples: list[StepTimingSample]) -> None:

196 + def median_curve(

197 + points: Iterable[tuple[int, float]],

198 + ) -> list[tuple[int, float]]:

199 + grouped: defaultdict[int, list[float]] = defaultdict(list)

200 + for key, value in points:

201 + grouped[key].append(value)

202 + return [(k, float(np.median(v))) for k, v in sorted(grouped.items())]

203 +

204 + # Draft curve: eager-target steps inflate drafter timings (the CPU is

205 + # still launching kernels, opening gaps between the drafter's events),

206 + # and request counts — unlike token counts — collide across execution

207 + # modes, so only graph-replay samples may price the draft curve.

208 + draft_curve = median_curve(

209 + (s.num_reqs, s.drafter_ms) for s in samples if s.full_cudagraph

210 + )

211 + verify_curve = median_curve(

212 + (s.num_target_tokens, s.forward_ms) for s in samples

213 + )

214 + self.set_cost_curves(draft_curve, verify_curve)

215 +

216 + def set_cost_curves(

217 + self,

218 + draft_curve: list[tuple[int, float]],

219 + verify_curve: list[tuple[int, float]],

220 + ) -> None:

221 + draft_curve, verify_curve = get_tp_group().broadcast_object(

222 + (draft_curve, verify_curve), src=0

223 + )

224 + if not draft_curve or not verify_curve:

225 + raise RuntimeError(

226 + "Adaptive verification could not profile step costs. Pass "

227 + "`enable_adaptive_verification=false` in the speculative config to "

228 + "verify a fixed number of drafts instead."

229 + )

230 + self.cost_tables = build_cost_tables_from_curves(

231 + draft_curve,

232 + verify_curve,

233 + self.req_states.max_num_reqs,

234 + self.req_states.max_num_batched_tokens,

235 + self._cudagraph_limit,

236 + )

237 + logger.debug("DSpark cost tables: %s", self.cost_tables)

238 +

239 + def record_confidences(

240 + self,

241 + confidence_probs: torch.Tensor,

242 + input_batch: "InputBatch",

243 + ) -> None:

244 + """Publish this step's raw confidences for the ranking kernel and start

245 + copying them to the CPU, where a later step's budget reads them."""

246 + num_reqs = input_batch.num_reqs

247 + ready_idx = self._stale_idx ^ 1

248 + with gpu_sync_allowed():

249 + self._copy_events[ready_idx].synchronize()

250 + if self._pending_resets:

251 + self._stale_confidences[ready_idx].np[self._pending_resets] = 1.0

252 + self._pending_resets.clear()

253 + # Last step's copy has landed: budgets read it, this step overwrites the

254 + # slot they were reading before.

255 + self._stale_idx, write_idx = ready_idx, self._stale_idx

256 +

257 + self._confidence_probs[input_batch.idx_mapping] = confidence_probs[:num_reqs]

258 + write_slot = self._stale_confidences[write_idx]

259 + write_slot.gpu.copy_(self._confidence_probs)

260 +

261 + current_stream = torch.cuda.current_stream(self.req_states.device)

262 + self._copy_stream.wait_stream(current_stream)

263 + with stream(self._copy_stream, current_stream):

264 + write_slot.copy_to_cpu()

265 + self._copy_events[write_idx].record()

266 +

267 + def get_num_tokens(

268 + self,

269 + num_tokens_per_req: dict[str, int],

270 + draft_tokens: dict[str, list[int]],

271 + ) -> int:

272 + """Token count once the draft budget is trimmed to fit.

273 +

274 + Stashes the chosen budget in ``_batch_budget`` for the compaction and

275 + reallocation that follow in the same step.

276 + """

277 + assert self.cost_tables is not None

278 + req_ids = list(num_tokens_per_req)

279 + num_reqs = len(req_ids)

280 + scheduled_tokens = np.fromiter(

281 + num_tokens_per_req.values(), dtype=np.int32, count=num_reqs

282 + )

283 + scheduled_drafts = np.fromiter(

284 + (len(draft_tokens.get(req_id, ())) for req_id in req_ids),

285 + dtype=np.int32,

286 + count=num_reqs,

287 + )

288 + num_non_draft_tokens = scheduled_tokens - scheduled_drafts

289 + slots = np.fromiter(

290 + (self.req_states.req_id_to_index[req_id] for req_id in req_ids),

291 + dtype=np.int32,

292 + count=len(req_ids),

293 + )

294 + stale_confidences = self._stale_confidences[self._stale_idx].np[slots]

295 + survival_probability = np.cumprod(stale_confidences.astype(np.float64), axis=1)

296 + steps = np.arange(self.num_speculative_steps)

297 + valid = steps[None, :] < scheduled_drafts[:, None]

298 + scores = np.sort(survival_probability[valid])[::-1]

299 + num_non_draft_tokens_total = int(num_non_draft_tokens.sum())

300 + max_draft_budget = min(

301 + int(scheduled_drafts.sum()),

302 + max(0, self._max_total_logits - num_reqs * self.num_bonus_tokens),

303 + )

304 + scores = scores[:max_draft_budget]

305 + draft_cost_ms, verify_cost_ms = self.cost_tables

306 + num_sampling_requests = np.count_nonzero(

307 + self.req_states.num_computed_tokens_np[slots] + num_non_draft_tokens

308 + >= self.req_states.prefill_len.np[slots]

309 + )

310 + num_tokens_to_estimated_accepted_tokens = np.concatenate(

311 + ([num_sampling_requests], num_sampling_requests + np.cumsum(scores))

312 + )

313 + costs = (

314 + draft_cost_ms[len(req_ids)]

315 + + verify_cost_ms[

316 + num_non_draft_tokens_total : num_non_draft_tokens_total

317 + + max_draft_budget

318 + + 1

319 + ]

320 + )

321 + num_drafts_per_req = {

322 + req_id: int(num_drafts)

323 + for req_id, num_drafts in zip(req_ids, scheduled_drafts, strict=True)

324 + }

325 + num_non_draft_tokens_per_req = {

326 + req_id: int(num_tokens)

327 + for req_id, num_tokens in zip(req_ids, num_non_draft_tokens, strict=True)

328 + }

329 + draft_budget = int(np.argmax(num_tokens_to_estimated_accepted_tokens / costs))

330 + self._batch_budget = (

331 + num_drafts_per_req,

332 + num_non_draft_tokens_per_req,

333 + draft_budget,

334 + )

335 + return sum(num_non_draft_tokens_per_req.values()) + draft_budget

336 +

337 + def compact_batch(

338 + self,

339 + num_draft_tokens_per_req: np.ndarray,

340 + num_scheduled_tokens: np.ndarray,

341 + cu_num_logits_np: np.ndarray,

342 + ) -> tuple[np.ndarray, np.ndarray]:

343 + """Compact the CPU batch to the chosen draft budget.

344 +

345 + Returns the compacted per-request token counts and the CPU cu_num_logits_np.

346 + If the draft budget is 0, we can know cu_num_logits_np exactly, otherwise

347 + its unchanged/an-upper-bound.

348 + """

349 + batch_budget = self._batch_budget

350 + assert batch_budget is not None

351 + _, _, draft_budget = batch_budget

352 + num_drafts = int(num_draft_tokens_per_req.sum())

353 + if draft_budget == num_drafts:

354 + return num_scheduled_tokens, cu_num_logits_np

355 +

356 + num_non_draft_tokens = num_scheduled_tokens - num_draft_tokens_per_req

357 + if draft_budget == 0:

358 + # The draft budget is 0, so we can know cu_num_logits_np exactly. This helps

359 + # when we would exceed the sampler logit chunk size.

360 + num_reqs = num_scheduled_tokens.shape[0]

361 + cu_num_logits_np = (

362 + np.arange(num_reqs + 1, dtype=cu_num_logits_np.dtype)

363 + * self.num_bonus_tokens

364 + )

365 + return num_non_draft_tokens, cu_num_logits_np

366 +

367 + is_verification_request = num_draft_tokens_per_req > 0

368 + num_verification_reqs = int(is_verification_request.sum())

369 + # sort_batch_req_ids keeps verification requests at the front.

370 + assert np.all(is_verification_request[:num_verification_reqs])

371 + # for the CPU side buffer we distribute draft tokens evenly

372 + draft_lens_cpu = np.zeros_like(num_non_draft_tokens)

373 + draft_lens_cpu[:num_verification_reqs] = draft_budget // num_verification_reqs

374 + draft_lens_cpu[: draft_budget % num_verification_reqs] += 1

375 + return num_non_draft_tokens + draft_lens_cpu, cu_num_logits_np

376 +

377 + def reallocate_drafts(

378 + self, req_ids: list[str], idx_mapping: torch.Tensor

379 + ) -> tuple[torch.Tensor, torch.Tensor, int]:

380 + batch_budget, self._batch_budget = self._batch_budget, None

381 + assert batch_budget is not None

382 + num_drafts_per_req, num_non_draft_tokens_per_req, draft_budget = batch_budget

383 + num_reqs = idx_mapping.shape[0]

384 + scheduled_drafts = np.fromiter(

385 + (num_drafts_per_req[req_id] for req_id in req_ids),

386 + dtype=np.int32,

387 + count=num_reqs,

388 + )

389 + num_non_draft_tokens = np.fromiter(

390 + (num_non_draft_tokens_per_req[req_id] for req_id in req_ids),

391 + dtype=np.int32,

392 + count=num_reqs,

393 + )

394 + num_tokens = int(num_non_draft_tokens.sum()) + draft_budget

395 +

396 + # Rank draft slots by survival probability and admit the best prefix.

397 + # capacities enters holding each request's valid draft count (the kernel

398 + # uses it to bound eligible slots) and leaves holding the admitted count.

399 + capacities = self._batch_draft_capacity[:num_reqs]

400 + if draft_budget == 0:

401 + capacities.zero_()

402 + else:

403 + async_copy_to_gpu(scheduled_drafts, out=capacities)

404 + if draft_budget < int(scheduled_drafts.sum()):

405 + _assign_draft_token_budget_compiled(

406 + self._confidence_probs,

407 + idx_mapping,

408 + capacities,

409 + draft_budget,

410 + self.num_speculative_steps,

411 + )

412 +

413 + num_non_draft_tokens_gpu = self._num_non_draft_tokens[:num_reqs]

414 + async_copy_to_gpu(

415 + num_non_draft_tokens,

416 + out=num_non_draft_tokens_gpu,

417 + )

418 + self._cu_num_logits[:1].zero_()

419 + torch.cumsum(

420 + capacities + self.num_bonus_tokens,

421 + dim=0,

422 + out=self._cu_num_logits[1 : num_reqs + 1],

423 + )

424 + self.query_start_loc[:1].zero_()

425 + torch.cumsum(

426 + capacities + num_non_draft_tokens_gpu,

427 + dim=0,

428 + out=self.query_start_loc[1 : num_reqs + 1],

429 + )

430 + self.query_start_loc[num_reqs + 1 :].fill_(num_tokens)

431 + return (

432 + self._cu_num_logits[: num_reqs + 1],

433 + self.query_start_loc,

434 + draft_budget,

435 + )

436 +

437 +

438 +def maybe_create_adaptive_verification_manager(

439 + *,

440 + enable_adaptive_verification: bool,

441 + attn_groups: list[list["AttentionGroup"]],

442 + attn_cg_support: "AttentionCGSupportInfo",

443 + req_states: "RequestState",

444 + query_start_loc: torch.Tensor,

445 + num_bonus_tokens: int,

446 + max_total_logits: int,

447 +) -> AdaptiveVerificationManager | None:

448 + if not enable_adaptive_verification:

449 + return None

450 +

451 + # The selector rejects unsupported backends, but models that

452 + # hard-wire theirs (e.g. DeepSeek-V4) never go through it.

453 + backend = get_query_lens_mismatch_unsupported_backend(attn_groups)

454 + if backend is not None:

455 + raise ValueError(

456 + "Adaptive verification trims verification requests on device, which"

457 + f" the {backend} attention backend does not support. Pass "

458 + "enable_adaptive_verification=false in the speculative config, or "

459 + "use a backend that does."

460 + )

461 +

462 + if attn_cg_support.min_cg_support != AttentionCGSupport.ALWAYS:

463 + raise ValueError(

464 + "Adaptive verification captures varlen decode cudagraphs, so every"

465 + " attention builder must report AttentionCGSupport.ALWAYS, but "

466 + f"{attn_cg_support.min_cg_attn_backend} reports "

467 + f"{attn_cg_support.min_cg_support}. Pass "

468 + "enable_adaptive_verification=false in the speculative config, or "

469 + "use a backend that does."

470 + )

471 +

472 + return AdaptiveVerificationManager(

473 + req_states,

474 + query_start_loc,

475 + num_bonus_tokens,

476 + max_total_logits=max_total_logits,

477 + )

vllm/v1/worker/gpu/spec_decode/dflash/speculator.py

+6 -4

@@ -135,11 +135,10 @@ class DFlashSpeculator(DraftModelSpeculator):

135

136 def capture(self) -> None:

137 logger.info("Capturing model for %s speculator...", self._speculator_name)

138 - # Reset sampling indices to zero to prevent stale values from prior

139 - # dummy runs from being baked into the captured graph.

138 + # Padded sample rows must not scatter into a live request during capture.

139 self.sample_indices.zero_()

140 self.sample_pos.zero_()

142 - self.sample_idx_mapping.zero_()

141 + self.sample_idx_mapping.fill_(-1)

142 assert self.query_cudagraph_manager is not None

143 self.query_cudagraph_manager.capture(

144 self._generate_draft,

@@ -589,7 +588,10 @@ def _prepare_dflash_inputs_kernel(

588 # seq_lens is the absolute sequence length the draft attention

589 # reads up to (context + query), not just the count of accepted

590 # tokens this step.

592 - tl.store(out_seq_lens_ptr + req_idx, last_valid_pos + 1 + num_query_per_req)

591 + tl.store(

592 + out_seq_lens_ptr + req_idx,

593 + tl.minimum(last_valid_pos + 1 + num_query_per_req, max_model_len),

594 + )

595 # Copy sampling state.

596 tl.store(

597 out_temperature_ptr + req_state_idx,

vllm/v1/worker/gpu/spec_decode/dspark/speculator.py

+41

@@ -29,10 +29,13 @@ import torch

30 from vllm.config import VllmConfig

31 from vllm.config.compilation import CUDAGraphMode

32 +from vllm.logger import init_logger

33 from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample

34 from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator

35 from vllm.v1.worker.gpu.spec_decode.dspark.utils import load_dspark_model

37 +logger = init_logger(__name__)

38 +

40 class DSparkSpeculator(DFlashSpeculator):

41 _speculator_name = "DSpark"

@@ -76,6 +79,13 @@ class DSparkSpeculator(DFlashSpeculator):

79 self.draft_model_config.hf_config, "dspark_draft_topk", None

80 )

82 + self.draft_token_confidence_probs = torch.empty_like(

83 + self.draft_tokens, dtype=torch.float32

84 + )

85 + self.enable_adaptive_verification = (

86 + self.speculative_config.enable_adaptive_verification

87 + )

88 +

89 def load_draft_model(

90 self,

91 target_model: torch.nn.Module,

@@ -98,6 +108,13 @@ class DSparkSpeculator(DFlashSpeculator):

108 dtype=self.draft_logits.dtype,

109 device=self.device,

110 )

111 + if self.enable_adaptive_verification and model.model.confidence_head is None:

112 + raise ValueError(

113 + "Adaptive verification needs a DSpark checkpoint with a confidence "

114 + "head, and this one has none. Pass "

115 + "enable_adaptive_verification=false in the speculative config to verify"

116 + " a fixed number of drafts instead."

117 + )

118 return model

119

120 def _sample_logits(

@@ -149,6 +166,7 @@ class DSparkSpeculator(DFlashSpeculator):

166

167 idx_map = self.sample_idx_mapping[:num_sample].view(num_reqs, n_spec)

168 sample_pos = self.sample_pos[:num_sample].view(num_reqs, n_spec)

169 + confidence_markov_embeds = []

170

171 # Anchor (bonus) token per request = the input id at query offset 0,

172 # read via the precomputed persistent index (fixed buffer for capture).

@@ -157,6 +175,8 @@ class DSparkSpeculator(DFlashSpeculator):

175 for i in range(n_spec):

176 # Sequential stage: Markov bias from the previously sampled token.

177 markov_embed = self.model.markov_embed(prev)

178 + if self.enable_adaptive_verification:

179 + confidence_markov_embeds.append(markov_embed)

180 bias = self.model.markov_bias(markov_embed)

181 logits_i = base_logits[:, i] + bias

182 draft_sampled_i = self._sample_logits(

@@ -165,6 +185,15 @@ class DSparkSpeculator(DFlashSpeculator):

185 self.draft_tokens[:num_reqs, i] = draft_sampled_i

186 prev = draft_sampled_i

187

188 + if self.enable_adaptive_verification:

189 + confidence = self.model.compute_confidence(

190 + sample_hidden,

191 + torch.stack(confidence_markov_embeds, dim=1).flatten(0, 1),

192 + )

193 + self.draft_token_confidence_probs[:num_reqs] = confidence.view(

194 + num_reqs, n_spec

195 + )

196 +

197 def _sample_sequential_topk(self, num_reqs: int, head_hidden: torch.Tensor) -> None:

198 """Apply the sequential Markov head only to top-k base-logit candidates.

199

@@ -186,10 +215,13 @@ class DSparkSpeculator(DFlashSpeculator):

215 base_logits.fill_(float("-inf"))

216 idx_map = self.sample_idx_mapping[:num_sample].view(num_reqs, n_spec)

217 sample_pos = self.sample_pos[:num_sample].view(num_reqs, n_spec)

218 + confidence_markov_embeds = []

219 prev = self.input_buffers.input_ids[self._anchor_idx[:num_reqs]]

220

221 for i in range(n_spec):

222 markov_embed = self.model.markov_embed(prev)

223 + if self.enable_adaptive_verification:

224 + confidence_markov_embeds.append(markov_embed)

225 logits_i = self.model.apply_markov_bias_gathered(

226 markov_embed,

227 base_logits[:, i],

@@ -202,6 +234,15 @@ class DSparkSpeculator(DFlashSpeculator):

234 self.draft_tokens[:num_reqs, i] = draft_sampled_i

235 prev = draft_sampled_i

236

237 + if self.enable_adaptive_verification:

238 + confidence = self.model.compute_confidence(

239 + sample_hidden,

240 + torch.stack(confidence_markov_embeds, dim=1).flatten(0, 1),

241 + )

242 + self.draft_token_confidence_probs[:num_reqs] = confidence.view(

243 + num_reqs, n_spec

244 + )

245 +

246 def _generate_draft(

247 self,

248 num_reqs: int,

vllm/v1/worker/gpu/spec_decode/rejection_sampler.py

+20 -4

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from collections.abc import Iterator

3 +from collections.abc import Iterable, Iterator

5 import numpy as np

6 import torch

@@ -32,6 +32,11 @@ MAX_CHUNK_BYTES = 2**30 # 1GB

32 _FP32_BYTES = 4

35 +def get_max_chunk_logits(vocab_size: int) -> int:

36 + """Largest number of logits rows one verification chunk may hold."""

37 + return max(1, MAX_CHUNK_BYTES // (vocab_size * _FP32_BYTES))

38 +

39 +

40 def _iter_request_chunks(

41 cu_num_logits: np.ndarray, max_chunk_logits: int

42 ) -> Iterator[tuple[int, int]]:

@@ -177,11 +182,22 @@ class RejectionSampler:

182 ) -> tuple[torch.Tensor, torch.Tensor, LogprobsTensors | None]:

183 cu_num_logits_np = input_batch.cu_num_logits_np

184 use_processed_logits = self.sampler.logprobs_mode in PROCESSED_LOGPROBS_MODES

185 + num_reqs = input_batch.num_reqs

186 +

187 + if logits.shape[0] <= max_chunk_logits:

188 + # One chunk covers the batch. Adaptive verification compacts the logits

189 + # without updating cu_num_logits_np (it keeps the pre-compacted layout),

190 + # so the stale sums must not pick chunk boundaries; its budget cap

191 + # guarantees the compacted batch always lands here.

192 + request_chunks: Iterable[tuple[int, int]] = ((0, num_reqs),)

193 + else:

194 + request_chunks = _iter_request_chunks(cu_num_logits_np, max_chunk_logits)

195 +

196 sampled_chunks: list[torch.Tensor] = []

197 num_sampled_chunks: list[torch.Tensor] = []

198 logprobs_chunks: list[LogprobsTensors] = []

199

184 - for start, end in _iter_request_chunks(cu_num_logits_np, max_chunk_logits):

200 + for start, end in request_chunks:

201 lo = int(cu_num_logits_np[start])

202 hi = int(cu_num_logits_np[end])

203 chunk_cu_num_logits_np = cu_num_logits_np[start : end + 1] - lo

@@ -246,14 +262,14 @@ class RejectionSampler:

262 max_num_logprobs = self.sampler.sampling_states.max_num_logprobs(

263 input_batch.idx_mapping_np

264 )

249 - max_chunk_logits = max(1, MAX_CHUNK_BYTES // (logits.shape[1] * _FP32_BYTES))

265 + chunk_logit_limit = get_max_chunk_logits(logits.shape[1])

266 sampled, num_sampled, logprobs_tensors = self._verify_in_chunks(

267 logits,

268 input_batch,

269 draft_logits,

270 draft_sampled,

271 pos,

256 - max_chunk_logits,

272 + chunk_logit_limit,

273 max_num_logprobs,

274 )

275

vllm/v1/worker/gpu/structured_outputs.py

+27 -4

@@ -10,7 +10,13 @@ from vllm.v1.worker.gpu.input_batch import InputBatch

12 class StructuredOutputsWorker:

13 - def __init__(self, max_num_logits: int, vocab_size: int, device: torch.device):

13 + def __init__(

14 + self,

15 + max_num_logits: int,

16 + vocab_size: int,

17 + device: torch.device,

18 + mask_stride: int,

19 + ):

20 self.logits_indices = torch.zeros(

21 max_num_logits, dtype=torch.int32, device=device

22 )

@@ -19,6 +25,7 @@ class StructuredOutputsWorker:

25 )

26 self.device = device

27 self.copy_stream = torch.cuda.Stream()

28 + self.mask_stride = mask_stride

30 def apply_grammar_bitmask(

31 self,

@@ -45,7 +52,13 @@ class StructuredOutputsWorker:

52 req_idx = req_id_to_idx[grammar_req_id]

53 logits_start_idx = cu_num_logits[req_idx]

54 logits_end_idx = cu_num_logits[req_idx + 1]

48 - mapping.extend(range(logits_start_idx, logits_end_idx))

55 + # Key by (request, position) rather than absolute logit index:

56 + # adaptive verification finalizes per-request logit offsets on

57 + # device, so the kernel resolves them from the GPU cu_num_logits.

58 + mapping.extend(

59 + req_idx * self.mask_stride + position

60 + for position in range(logits_end_idx - logits_start_idx)

61 + )

63 # Asynchronously copy the mapping to GPU.

64 with torch.cuda.stream(self.copy_stream):

@@ -69,9 +82,11 @@ class StructuredOutputsWorker:

82 logits,

83 logits.stride(0),

84 logits_indices,

85 + input_batch.cu_num_logits,

86 bitmask,

87 bitmask.stride(0),

88 vocab_size,

89 + MASK_STRIDE=self.mask_stride,

90 BLOCK_SIZE=BLOCK_SIZE,

91 )

@@ -87,13 +102,21 @@ def _apply_grammar_bitmask_kernel(

102 logits_ptr,

103 logits_stride,

104 logits_indices_ptr,

105 + cu_num_logits_ptr,

106 bitmask_ptr,

107 bitmask_stride,

108 vocab_size,

109 + MASK_STRIDE: tl.constexpr,

110 BLOCK_SIZE: tl.constexpr,

111 ):

112 bitmask_idx = tl.program_id(0)

96 - logits_idx = tl.load(logits_indices_ptr + bitmask_idx)

113 + mapping_idx = tl.load(logits_indices_ptr + bitmask_idx)

114 + req_idx = mapping_idx // MASK_STRIDE

115 + position_idx = mapping_idx % MASK_STRIDE

116 + logits_idx = tl.load(cu_num_logits_ptr + req_idx)

117 + num_req_logits = tl.load(cu_num_logits_ptr + req_idx + 1) - logits_idx

118 + logits_idx += position_idx

119 + position_is_active = position_idx < num_req_logits

120

121 # Load the bitmask.

122 block_id = tl.program_id(1)

@@ -111,5 +134,5 @@ def _apply_grammar_bitmask_kernel(

134 tl.store(

135 logits_ptr + logits_idx * logits_stride + block_offset,

136 -float("inf"),

114 - mask=bitmask & (block_offset < vocab_size),

137 + mask=position_is_active & bitmask & (block_offset < vocab_size),

138 )