@tduncan / Nvidia-Optimization-Stack-VLLM / commits / a4eb3f25d6

[Spec Decode] Enable Gemma4 DSpark adaptive verification with FlashInfer (#57263)

Signed-off-by: zixi-qi <zixi@inferact.ai> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>

qizixi committed Sep 26, 2026 at 02:01 UTC a4eb3f25d6f9b3cad7ecf5390423d853935fcaeb
24 files changed +894 -135
docs/design/cuda_graphs.md
+2
@@ -172,6 +172,8 @@ class AttentionCGSupport(enum.Enum):
172
173 Suppose we have hybrid attention backends (e.g., in mamba mixer models). In that case, we seek the minimum capability of all backends to determine the final capability of the model, and we might resolve the incompatible CUDA Graphs mode by downgrading the mode to the best fit one. For example, downgrading `FULL` mode to `FULL_AND_PIECEWISE` mode if the minimum capability is `UNIFORM_BATCH`, or `PIECEWISE` mode if the minimum capability is `NEVER` for -O3 compilation mode. For the complete fallback policy, please see the code for [this][vllm.v1.worker.gpu_model_runner.GPUModelRunner._check_and_update_cudagraph_mode].
174
175 +Variable-length decode batches, where each request carries a different number of query tokens read from the device `query_start_loc` (as in adaptive verification), are declared separately and not ordered against the enum: `AttentionMetadataBuilder.get_varlen_cudagraph_max_query_len()` returns the largest per-request query length a FULL decode graph can replay. It returns `None` for builders reporting `ALWAYS`, which replay any batch, and for builders that cannot replay variable-length batches; FlashInfer returns `1 + num_speculative_tokens` when its TRTLLM-GEN varlen decode path is active. Batches with a prefill never replay these graphs.
176 +
177 The following table lists backends that support full CUDA Graphs at the time of writing.
178
179 | Attention Backend | cudagraph_support | Comments |
tests/distributed/test_dcp_direct_a2a_lse_reduce.py
+19 -3
@@ -455,13 +455,29 @@ def test_mla_dcp_manager_selects_fallback_backends(monkeypatch):
455 )
456
457
458 -def test_dcp_workspace_covers_parallel_drafting():
458 +@pytest.mark.parametrize(
459 + ("parallel_drafting", "use_v2_model_runner", "expected"),
460 + [
461 + # Model runner V1's parallel drafter widens requests to 1 + 2 * 3.
462 + (True, False, 28),
463 + (True, True, 16),
464 + # A diffusion canvas of 3 without a speculative config.
465 + (None, True, 16),
466 + ],
467 +)
468 +def test_dcp_workspace_covers_decode_width(
469 + parallel_drafting, use_v2_model_runner, expected
470 +):
471 config = _manager_config()
472 config.scheduler_config.max_num_batched_tokens = 128
473 config.num_speculative_tokens = 3
462 - config.speculative_config = MagicMock(parallel_drafting=True)
474 + config.use_v2_model_runner = use_v2_model_runner
475 + if parallel_drafting is not None:
476 + config.speculative_config = MagicMock(
477 + num_speculative_tokens=3, parallel_drafting=parallel_drafting
478 + )
479
464 - assert dcp.get_dcp_workspace_max_num_tokens(config) == 28
480 + assert dcp.get_dcp_workspace_max_num_tokens(config) == expected
481
482
483 def test_mla_dcp_manager_selects_pcp_combine(monkeypatch):
tests/models/language/generation/test_gemma.py
+67 -2
@@ -7,7 +7,8 @@ import numpy as np
7 import pytest
8 import torch
9
10 -from vllm.config import VllmConfig
10 +from vllm.config import CompilationConfig, VllmConfig
11 +from vllm.config.compilation import CompilationMode
12 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
13 from vllm.model_executor.models import gemma
14 from vllm.model_executor.models.gemma3n import (
@@ -18,10 +19,72 @@ from vllm.model_executor.models.gemma4 import (
19 Gemma4ForCausalLM,
20 _gemma4_layer_weights_mapper,
21 )
22 +from vllm.model_executor.models.gemma4_dspark import (
23 + Gemma4DSparkForCausalLM,
24 + Gemma4DSparkModel,
25 +)
26
27 MODELS = ["google/gemma-2b", "google/gemma-2-2b", "google/gemma-3-4b-it"]
28
29
30 +@pytest.mark.usefixtures("dist_init")
31 +@pytest.mark.parametrize(
32 + "enabled,has_weights,with_markov",
33 + [
34 + pytest.param(True, True, False, id="hidden-only"),
35 + pytest.param(True, True, True, id="with-markov"),
36 + pytest.param(True, False, True, id="missing-weights"),
37 + pytest.param(False, True, True, id="disabled-head"),
38 + ],
39 +)
40 +def test_gemma4_dspark_loads_confidence_head(
41 + monkeypatch, enabled, has_weights, with_markov
42 +) -> None:
43 + """Use checkpoint confidence parameters, or disable an unavailable head."""
44 + config = SimpleNamespace(
45 + vocab_size=64,
46 + hidden_size=8,
47 + target_layer_ids=[0, 1],
48 + num_hidden_layers=0,
49 + rms_norm_eps=1e-6,
50 + markov_rank=4,
51 + enable_confidence_head=enabled,
52 + confidence_head_with_markov=with_markov,
53 + )
54 + vllm_config = SimpleNamespace(
55 + compilation_config=CompilationConfig(mode=CompilationMode.NONE),
56 + model_config=SimpleNamespace(dtype=torch.bfloat16),
57 + speculative_config=SimpleNamespace(
58 + draft_model_config=SimpleNamespace(hf_config=config),
59 + ),
60 + )
61 + monkeypatch.setattr(Gemma4DSparkModel, "_build_fused_kv_buffers", lambda _: None)
62 + model = Gemma4DSparkForCausalLM(vllm_config=cast(VllmConfig, vllm_config))
63 + width = config.hidden_size + (config.markov_rank if with_markov else 0)
64 + weight = torch.arange(width, dtype=torch.bfloat16).reshape(1, -1) / 16
65 + bias = torch.tensor([-0.5], dtype=torch.bfloat16)
66 + loaded = model.load_weights(
67 + [("confidence_head.proj.weight", weight), ("confidence_head.proj.bias", bias)]
68 + if has_weights
69 + else []
70 + )
71 + if not (enabled and has_weights):
72 + assert model.model.confidence_head is None
73 + assert not loaded
74 + return
75 +
76 + assert loaded == {
77 + "model.confidence_head.proj.weight",
78 + "model.confidence_head.proj.bias",
79 + }
80 + assert model.model.confidence_head.proj.weight.dtype == torch.float32
81 + hidden = torch.full((3, config.hidden_size), 0.25, dtype=torch.bfloat16)
82 + markov = torch.full((3, config.markov_rank), -0.5, dtype=torch.bfloat16)
83 + inputs = torch.cat([hidden, markov], dim=-1) if with_markov else hidden
84 + expected = (inputs.float() @ weight.float().T + bias.float()).sigmoid().squeeze(-1)
85 + torch.testing.assert_close(model.compute_confidence(hidden, markov), expected)
86 +
87 +
88 @pytest.mark.cpu_test
89 @pytest.mark.usefixtures("dist_init")
90 def test_checkpoint_lm_head_can_override_tied_config(monkeypatch) -> None:
@@ -164,7 +227,9 @@ def test_dummy_loader(vllm_runner, monkeypatch, model: str) -> None:
227 ) as llm:
228 if model == "google/gemma-3-4b-it":
229 normalizers = llm.llm.collective_rpc(
167 - lambda self: self.model_runner.model.language_model.model.normalizer.cpu().item() # noqa: E501
230 + lambda self: (
231 + self.model_runner.model.language_model.model.normalizer.cpu().item()
232 + ) # noqa: E501
233 )
234 config = llm.llm.llm_engine.model_config.hf_config.text_config
235 else:
tests/test_config.py
+12
@@ -92,6 +92,16 @@ def test_nested_rope_validation_patch_preserves_flat_rope_parameters(monkeypatch
92 assert len(calls) == 2
93
94
95 +def test_dspark_adaptive_verification_separates_graph_cache():
96 + config = object.__new__(SpeculativeConfig)
97 + config.method = "dspark"
98 + config.draft_model_config = None
99 + config.enable_adaptive_verification = False
100 + fixed_hash = config.compute_hash()
101 + config.enable_adaptive_verification = True
102 + assert config.compute_hash() != fixed_hash
103 +
104 +
105 def _write_json(path: Path, value: object) -> None:
106 path.parent.mkdir(parents=True, exist_ok=True)
107 path.write_text(json.dumps(value), encoding="utf-8")
@@ -878,6 +888,8 @@ def test_resolve_cudagraph_mode_adjusts_spec_decode_sizes_only_for_v1(
888 ("FULL_AND_PIECEWISE", False, "ALWAYS", "FULL_DECODE_ONLY"),
889 ("FULL_DECODE_ONLY", False, "ALWAYS", "FULL_DECODE_ONLY"),
890 ("FULL_DECODE_ONLY", False, "NEVER", "NONE"),
891 + ("FULL_AND_PIECEWISE", True, "UNIFORM_BATCH", "FULL_AND_PIECEWISE"),
892 + ("FULL", True, "UNIFORM_BATCH", "FULL_DECODE_ONLY"),
893 ],
894 )
895 def test_resolve_cudagraph_mode_uses_loaded_piecewise_provider(
tests/transformers_utils/test_config.py
+45
@@ -5,6 +5,7 @@ only get the `eos_token_id` from the tokenizer as defined by
5 `BaseRenderer.get_eos_token_id`.
6 """
7
8 +import json
9 import math
10 from types import SimpleNamespace
11 from typing import cast
@@ -31,6 +32,50 @@ from vllm.transformers_utils.configs.glm5_next import (
32 from vllm.transformers_utils.configs.mistral import adapt_config_dict
33
34
35 +@pytest.mark.parametrize("layout", ["mixed", "flat"])
36 +def test_gemma4_dspark_rope_config_preserves_parameters(tmp_path, layout):
37 + """Remove redundant shared entries while preserving per-layer and flat RoPE."""
38 + from transformers import Gemma4TextConfig
39 +
40 + per_layer = {
41 + "full_attention": {
42 + "rope_type": "proportional",
43 + "partial_rotary_factor": 0.25,
44 + "rope_theta": 1000000.0,
45 + },
46 + "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0},
47 + }
48 + rope_parameters: dict[str, object] = dict(per_layer)
49 + if layout == "mixed":
50 + rope_parameters.update(rope_type="default", rope_theta=None)
51 + else:
52 + rope_parameters = {"rope_type": "default", "rope_theta": 12345.0}
53 + (tmp_path / "config.json").write_text(
54 + json.dumps(
55 + {
56 + "model_type": "gemma4_text",
57 + "architectures": ["Gemma4DSparkModel"],
58 + "num_hidden_layers": 1,
59 + "layer_types": ["full_attention"],
60 + "rope_parameters": rope_parameters,
61 + }
62 + )
63 + )
64 + _, config = config_module.HFConfigParser().parse(
65 + tmp_path, trust_remote_code=False, max_position_embeddings=8192
66 + )
67 + assert isinstance(config, Gemma4TextConfig)
68 + assert config.name_or_path == str(tmp_path)
69 + assert config.max_position_embeddings == 8192
70 + if layout == "flat":
71 + assert config.rope_parameters["rope_theta"] == 12345.0
72 + else:
73 + for layer_type, expected in per_layer.items():
74 + for key, value in expected.items():
75 + assert config.rope_parameters[layer_type][key] == value
76 + assert set(config.rope_parameters) == set(per_layer)
77 +
78 +
79 def test_patch_legacy_rope_type_preserves_nope_layers():
80 """NoPE layers stay disabled while later RoPE layers are normalized."""
81 rope_parameters = {
tests/v1/attention/test_attention_backends.py
+376
@@ -2,12 +2,14 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for v1 attention backends without GPUModelRunner dependency."""
4
5 +from dataclasses import replace
6 from functools import partial
7 from types import SimpleNamespace
8
9 import pytest
10 import torch
11 from torch.nn.attention.flex_attention import create_block_mask, flex_attention
12 +from transformers import LlamaConfig
13
14 from tests.v1.attention.utils import (
15 BatchSpec,
@@ -28,8 +30,10 @@ from vllm.utils.torch_utils import (
30 )
31 from vllm.v1.attention.backend import (
32 AttentionCGSupport,
33 + AttentionMetadataBuilder,
34 AttentionType,
35 CommonAttentionMetadata,
36 + max_decode_query_len,
37 )
38 from vllm.v1.attention.backends.registry import AttentionBackendEnum
39 from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheLayout
@@ -1066,6 +1070,378 @@ def test_flashinfer_trtllm_gen_padded_decode_uses_varlen_offsets(
1070 assert q_cu_seq_lens.tolist() == qo_indptr_values
1071
1072
1073 +@pytest.mark.skipif(
1074 + AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,
1075 + reason="FlashInfer is not available.",
1076 +)
1077 +@pytest.mark.parametrize(
1078 + "adaptive,decode_kernel,dcp_size,expected_bound",
1079 + [
1080 + (None, "TRTLLM_GEN", 1, None),
1081 + (False, "TRTLLM_GEN", 1, None),
1082 + (True, "TRTLLM_GEN", 1, 8),
1083 + (True, "XQA", 1, None),
1084 + (True, "TRTLLM_GEN", 2, None),
1085 + ],
1086 +)
1087 +def test_flashinfer_varlen_cudagraph_capability(
1088 + monkeypatch, adaptive, decode_kernel, dcp_size, expected_bound
1089 +):
1090 + """Only the trtllm-gen path that adaptive verification enables replays
1091 + varlen decode graphs, up to the decode width; the uniform level is fixed."""
1092 + from vllm.model_executor.layers.attention.chunked_local_attention import (
1093 + create_chunked_local_attention_backend,
1094 + )
1095 + from vllm.v1.attention.backends import flashinfer as fi
1096 +
1097 + config = SimpleNamespace(
1098 + attention_config=SimpleNamespace(use_non_causal=False),
1099 + speculative_config=SimpleNamespace(
1100 + enable_adaptive_verification=adaptive,
1101 + num_speculative_tokens=7,
1102 + parallel_drafting=True,
1103 + )
1104 + if adaptive is not None
1105 + else None,
1106 + parallel_config=SimpleNamespace(decode_context_parallel_size=dcp_size),
1107 + model_config=SimpleNamespace(get_num_attention_heads=lambda _: 32),
1108 + use_v2_model_runner=True,
1109 + )
1110 + monkeypatch.setattr(fi, "can_use_trtllm_attention", lambda *_, **__: True)
1111 + monkeypatch.setattr(
1112 + fi.FlashInferMetadataBuilder,
1113 + "_get_flashinfer_trtllm_api_decode_kernel",
1114 + staticmethod(lambda: fi.FlashInferDecodeKernel[decode_kernel]),
1115 + )
1116 + spec = FullAttentionSpec(
1117 + block_size=16,
1118 + num_kv_heads=4,
1119 + head_size=64,
1120 + dtype=torch.bfloat16,
1121 + )
1122 + builder_cls = fi.FlashInferMetadataBuilder
1123 + assert builder_cls.get_cudagraph_support(config, spec) == (
1124 + AttentionCGSupport.UNIFORM_BATCH
1125 + if dcp_size == 1
1126 + else AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
1127 + )
1128 + bound = builder_cls.get_varlen_cudagraph_max_query_len(config, spec)
1129 + assert bound == expected_bound
1130 + if bound is not None:
1131 + # A wrapper forcing NEVER inherits the override but not the bound.
1132 + chunked_local = create_chunked_local_attention_backend(fi.FlashInferBackend, 16)
1133 + wrapped_cls = chunked_local.get_builder_cls()
1134 + assert wrapped_cls.get_varlen_cudagraph_max_query_len(config, spec) is None
1135 +
1136 +
1137 +@pytest.mark.parametrize(
1138 + "use_v2_model_runner,parallel_drafting,expected",
1139 + [
1140 + (True, None, 1),
1141 + (True, False, 8),
1142 + (True, True, 8),
1143 + (False, False, 8),
1144 + (False, True, 15),
1145 + ],
1146 +)
1147 +def test_spec_as_decode_width_follows_model_runner(
1148 + use_v2_model_runner, parallel_drafting, expected
1149 +):
1150 + """Only model runner V1's parallel drafter widens decode requests."""
1151 +
1152 + class Builder(AttentionMetadataBuilder):
1153 + def __init__(self, vllm_config):
1154 + self.vllm_config = vllm_config
1155 +
1156 + def build(self, common_prefix_len, common_attn_metadata, fast_build=False):
1157 + raise NotImplementedError
1158 +
1159 + config = SimpleNamespace(
1160 + speculative_config=None
1161 + if parallel_drafting is None
1162 + else SimpleNamespace(
1163 + num_speculative_tokens=7, parallel_drafting=parallel_drafting
1164 + ),
1165 + parallel_config=SimpleNamespace(decode_context_parallel_size=1),
1166 + use_v2_model_runner=use_v2_model_runner,
1167 + )
1168 + assert max_decode_query_len(config) == expected
1169 + builder = Builder(config)
1170 + builder._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
1171 + assert builder.reorder_batch_threshold == expected
1172 +
1173 +
1174 +@pytest.mark.skipif(
1175 + AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST
1176 + or not current_platform.is_device_capability_family(100),
1177 + reason="TRTLLM-GEN varlen decode graphs require FlashInfer on SM100-family GPUs",
1178 +)
1179 +@pytest.mark.parametrize(
1180 + "head_size,sliding_window,num_tokens,num_reqs,num_heads,num_kv_heads",
1181 + [
1182 + pytest.param(64, None, 8, 8, 32, 4, id="single-token-capture"),
1183 + pytest.param(
1184 + 256,
1185 + 1024,
1186 + 40,
1187 + 32,
1188 + 16,
1189 + 8,
1190 + id="gemma4-sliding-gqa",
1191 + ),
1192 + pytest.param(
1193 + 512,
1194 + None,
1195 + 40,
1196 + 32,
1197 + 16,
1198 + 1,
1199 + id="gemma4-global-mqa",
1200 + ),
1201 + ],
1202 +)
1203 +def test_flashinfer_varlen_decode_graph_replays_changed_layout(
1204 + tmp_path,
1205 + monkeypatch,
1206 + head_size,
1207 + sliding_window,
1208 + num_tokens,
1209 + num_reqs,
1210 + num_heads,
1211 + num_kv_heads,
1212 +):
1213 + """Replay changes device offsets, active slots, context, and KV writes.
1214 +
1215 + The graph is captured with evenly spread dummy queries and an explicit
1216 + bound of eight. CPU offsets then remain stale while device offsets change.
1217 + Compare real output rows with independent FP32 causal/windowed attention,
1218 + and verify cache writes including preservation of the null padding block.
1219 + """
1220 + from vllm.v1.attention.backends import flashinfer as fi
1221 + from vllm.v1.attention.backends.utils import PerLayerParameters
1222 +
1223 + set_random_seed(0)
1224 + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
1225 + block_size = 16
1226 + dtype = torch.bfloat16
1227 + LlamaConfig(
1228 + architectures=["LlamaForCausalLM"],
1229 + hidden_size=num_heads * head_size,
1230 + head_dim=head_size,
1231 + num_attention_heads=num_heads,
1232 + num_key_value_heads=num_kv_heads,
1233 + num_hidden_layers=1,
1234 + max_position_embeddings=4096,
1235 + ).save_pretrained(tmp_path)
1236 + num_blocks = 1 + num_reqs * cdiv(2112, block_size)
1237 + config = create_vllm_config(
1238 + model_name=str(tmp_path),
1239 + max_model_len=4096,
1240 + dtype=dtype,
1241 + num_gpu_blocks=num_blocks,
1242 + block_size=block_size,
1243 + max_num_seqs=32,
1244 + max_num_batched_tokens=256,
1245 + )
1246 + config.cache_config.kv_cache_layout = "BLHNC"
1247 + config.speculative_config = SimpleNamespace(
1248 + num_speculative_tokens=7,
1249 + enable_adaptive_verification=True,
1250 + parallel_drafting=False,
1251 + )
1252 + spec = create_standard_kv_cache_spec(config)
1253 + scale = head_size**-0.5
1254 + monkeypatch.setattr(
1255 + fi,
1256 + "get_per_layer_parameters",
1257 + lambda *_: {
1258 + "layer": PerLayerParameters(
1259 + window_left=-1 if sliding_window is None else sliding_window - 1,
1260 + logits_soft_cap=0.0,
1261 + sm_scale=scale,
1262 + has_sinks=False,
1263 + ),
1264 + },
1265 + )
1266 + device = torch.device("cuda")
1267 + layout = KVCacheLayout.BLHNC
1268 + query = torch.empty(num_tokens, num_heads, head_size, dtype=dtype, device=device)
1269 + key = torch.empty(num_tokens, num_kv_heads, head_size, dtype=dtype, device=device)
1270 + value = torch.empty_like(key)
1271 + output = torch.empty_like(query)
1272 + dummy_lens = [
1273 + num_tokens // num_reqs + (i < num_tokens % num_reqs) for i in range(num_reqs)
1274 + ]
1275 + common = create_common_attn_metadata(
1276 + BatchSpec(seq_lens=[2112] * num_reqs, query_lens=dummy_lens),
1277 + block_size,
1278 + device,
1279 + )
1280 + common.max_query_len = 8
1281 + cache = None
1282 +
1283 + def prepare(query_lens):
1284 + nonlocal cache
1285 + query_lens = query_lens + [0] * (num_reqs - len(query_lens))
1286 + query.normal_()
1287 + key.normal_()
1288 + value.normal_()
1289 + # Alternating short prefills, nonempty decode contexts, and contexts
1290 + # beyond the sliding window; inactive requests have zero sequence length.
1291 + contexts = [
1292 + ([0, 17, 2079, 39][i % 4] if q else 0) for i, q in enumerate(query_lens)
1293 + ]
1294 + ks = [
1295 + torch.randn(c, num_kv_heads, head_size, dtype=dtype, device=device)
1296 + for c in contexts
1297 + ]
1298 + vs = [torch.randn_like(k) for k in ks]
1299 + actual = create_common_attn_metadata(
1300 + BatchSpec(
1301 + seq_lens=[c + q for c, q in zip(contexts, query_lens)],
1302 + query_lens=query_lens,
1303 + ),
1304 + block_size,
1305 + device,
1306 + )
1307 + fresh_cache = create_and_prepopulate_kv_cache(
1308 + ks,
1309 + vs,
1310 + block_size,
1311 + num_kv_heads,
1312 + head_size,
1313 + dtype,
1314 + device,
1315 + num_blocks,
1316 + actual,
1317 + layout,
1318 + randomize_blocks=True,
1319 + )
1320 + if cache is None:
1321 + cache = fresh_cache
1322 + else:
1323 + cache.copy_(fresh_cache)
1324 + common.query_start_loc.copy_(actual.query_start_loc)
1325 + common.seq_lens.copy_(actual.seq_lens)
1326 + common.block_table_tensor.zero_()
1327 + common.block_table_tensor[:, : actual.block_table_tensor.shape[1]].copy_(
1328 + actual.block_table_tensor
1329 + )
1330 + common.slot_mapping.fill_(-1)
1331 + common.slot_mapping[: sum(query_lens)].copy_(actual.slot_mapping)
1332 + return contexts, ks, vs
1333 +
1334 + prepare(dummy_lens)
1335 + with set_current_vllm_config(config):
1336 + builder = fi.FlashInferMetadataBuilder(spec, ["layer"], config, device)
1337 + metadata = builder.build(0, common)
1338 + assert metadata.num_prefills == 0
1339 + assert metadata.decode.q_len_per_req == 8
1340 + impl = fi.FlashInferImpl(
1341 + num_heads=num_heads,
1342 + num_kv_heads=num_kv_heads,
1343 + head_size=head_size,
1344 + scale=scale,
1345 + alibi_slopes=None,
1346 + sliding_window=sliding_window,
1347 + kv_cache_dtype="auto",
1348 + attn_type=AttentionType.DECODER,
1349 + )
1350 + layer = MockAttentionLayer(device)
1351 +
1352 + def forward(attn_metadata=metadata):
1353 + impl.do_kv_cache_update(layer, key, value, cache, common.slot_mapping)
1354 + impl.forward(layer, query, key, value, cache, attn_metadata, output=output)
1355 +
1356 + for _ in range(3):
1357 + forward()
1358 + torch.accelerator.synchronize()
1359 + graph = torch.cuda.CUDAGraph()
1360 + with torch.cuda.graph(graph):
1361 + forward()
1362 +
1363 + layouts = [[8], [1] * num_reqs]
1364 + if num_tokens >= 16:
1365 + layouts.append([8, 4, 3, 1])
1366 + if num_tokens >= 40:
1367 + layouts.append([8] * 5)
1368 +
1369 + def check(query_lens, run):
1370 + contexts, ks, vs = prepare(query_lens)
1371 + assert cache is not None
1372 + output.fill_(float("nan"))
1373 + run()
1374 + start = 0
1375 + for i, q_len in enumerate(query_lens):
1376 + if q_len == 0:
1377 + continue
1378 + end = start + q_len
1379 + full_k = (
1380 + torch.cat((ks[i], key[start:end]))
1381 + .float()
1382 + .repeat_interleave(num_heads // num_kv_heads, dim=1)
1383 + )
1384 + full_v = (
1385 + torch.cat((vs[i], value[start:end]))
1386 + .float()
1387 + .repeat_interleave(num_heads // num_kv_heads, dim=1)
1388 + )
1389 + scores = (
1390 + torch.einsum("qhd,khd->hqk", query[start:end].float(), full_k) * scale
1391 + )
1392 + q_pos = torch.arange(q_len, device=device) + contexts[i]
1393 + k_pos = torch.arange(full_k.shape[0], device=device)
1394 + allowed = k_pos[None, :] <= q_pos[:, None]
1395 + if sliding_window is not None:
1396 + allowed &= k_pos[None, :] > q_pos[:, None] - sliding_window
1397 + probs = scores.masked_fill(~allowed[None], float("-inf")).softmax(-1)
1398 + reference = torch.einsum("hqk,khd->qhd", probs, full_v)
1399 + torch.testing.assert_close(
1400 + output[start:end].float(),
1401 + reference,
1402 + atol=2e-2,
1403 + rtol=2e-2,
1404 + )
1405 + slots = common.slot_mapping[start:end]
1406 + written = cache[slots // block_size, :, slots % block_size, :]
1407 + torch.testing.assert_close(written[..., :head_size], key[start:end])
1408 + torch.testing.assert_close(written[..., head_size:], value[start:end])
1409 + start = end
1410 + assert torch.count_nonzero(cache[0]) == 0
1411 +
1412 + for query_lens in layouts:
1413 + check(query_lens, graph.replay)
1414 +
1415 + if num_tokens == 40:
1416 + # A longer prefill forces a separate launch in eager/piecewise mode.
1417 + # Its length must not become the decode kernel's query bound, and
1418 + # compacted CPU placeholders still need not match the GPU allocation.
1419 + common.num_reqs = 5
1420 + common.max_query_len = 24
1421 + common.query_start_loc_cpu[:6] = torch.tensor([0, 4, 8, 12, 16, 40])
1422 + common.query_start_loc_cpu[6:] = 40
1423 +
1424 + def mixed_forward():
1425 + with set_current_vllm_config(config):
1426 + metadata = builder.build(
1427 + 0,
1428 + replace(
1429 + common,
1430 + query_start_loc=common.query_start_loc[:6],
1431 + query_start_loc_cpu=common.query_start_loc_cpu[:6],
1432 + seq_lens=common.seq_lens[:5],
1433 + seq_lens_cpu_upper_bound=common.seq_lens_cpu_upper_bound[:5],
1434 + block_table_tensor=common.block_table_tensor[:5],
1435 + ),
1436 + )
1437 + assert metadata.num_decodes == 4
1438 + assert metadata.num_prefills == 1
1439 + assert metadata.decode.q_len_per_req == 8
1440 + forward(metadata)
1441 +
1442 + check([8, 4, 3, 1, 24], mixed_forward)
1443 +
1444 +
1445 @pytest.mark.skipif(
1446 AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,
1447 reason="FlashInfer is not available.",
tests/v1/cudagraph/test_cudagraph_manager.py
+24
@@ -294,6 +294,7 @@ def _make_spec_decode_manager(
294 dynamic_spec_schedule: list[tuple[int, int, int]] | None = None,
295 max_num_seqs: int = 8,
296 use_kda_recoverssm: bool = False,
297 + varlen_decode: bool = False,
298 ) -> gpu_cudagraph_utils.CudaGraphManager:
299 monkeypatch.setattr(
300 gpu_cudagraph_utils,
@@ -316,6 +317,7 @@ def _make_spec_decode_manager(
317 device=torch.device("cpu"),
318 cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE,
319 decode_query_len=decode_query_len,
320 + varlen_decode=varlen_decode,
321 )
322 manager._graphs_captured = True
323 return manager
@@ -373,6 +375,28 @@ def test_mixed_batch_never_selects_a_uniform_decode_graph(monkeypatch):
375 assert desc.num_tokens == 16
376
377
378 +def test_varlen_decode_graph_takes_bounded_decode_batches_only(monkeypatch):
379 + """A varlen decode graph replays any decode batch within its query-length
380 + bound. The runner passes no bound for a batch with a prefill, which must not
381 + replay one however short the prefill is."""
382 + manager = _make_spec_decode_manager(monkeypatch, varlen_decode=True)
383 +
384 + def dispatch(max_query_len):
385 + return manager.dispatch(
386 + num_reqs=3,
387 + num_tokens=6,
388 + uniform_token_count=None,
389 + num_active_loras=0,
390 + max_query_len=max_query_len,
391 + )
392 +
393 + desc = dispatch(_DECODE_QUERY_LEN)
394 + assert desc.cg_mode == CUDAGraphMode.FULL
395 + assert desc.max_query_len == _DECODE_QUERY_LEN
396 + assert dispatch(None).cg_mode == CUDAGraphMode.PIECEWISE
397 + assert dispatch(_DECODE_QUERY_LEN + 1).cg_mode == CUDAGraphMode.PIECEWISE
398 +
399 +
400 def test_mixed_batch_at_decode_only_token_count_still_gets_a_graph(monkeypatch):
401 """A mixed batch must not fall to eager where only decode graphs are staged.
402
tests/v1/spec_decode/test_adaptive_verification.py
+86 -34
@@ -1,15 +1,15 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from contextlib import nullcontext
5 from types import SimpleNamespace
6
7 import numpy as np
8 import pytest
9
10 from vllm.config.compilation import CUDAGraphMode
10 -from vllm.v1.attention.backend import AttentionCGSupport
11 +from vllm.v1.attention.backend import AttentionCGSupport, AttentionMetadataBuilder
12 from vllm.v1.worker.gpu.async_utils import StepTimingSample
12 -from vllm.v1.worker.gpu.attn_utils import AttentionCGSupportInfo
13 from vllm.v1.worker.gpu.spec_decode import adaptive_verification as adaptive_module
14 from vllm.v1.worker.gpu.spec_decode.adaptive_verification import (
15 AdaptiveVerificationManager,
@@ -58,37 +58,90 @@ def test_resolve_adaptive_cudagraph_mode(mode, piecewise_capture_available, expe
58 )
59
60
61 -def test_manager_scopes_varlen_check_without_weakening_runner_cg_mode(monkeypatch):
62 - class Backend:
63 - @classmethod
64 - def supports_device_cpu_query_lens_mismatch(cls):
65 - return True
61 +@pytest.mark.parametrize(
62 + "target_support,target_bound,device_offsets,additional,error",
63 + [
64 + pytest.param(AttentionCGSupport.ALWAYS, None, True, None, None, id="always"),
65 + pytest.param(
66 + AttentionCGSupport.UNIFORM_BATCH, 8, True, None, None, id="bounded"
67 + ),
68 + pytest.param(
69 + AttentionCGSupport.UNIFORM_BATCH,
70 + None,
71 + True,
72 + None,
73 + "TargetBackend allows none",
74 + id="unsupported",
75 + ),
76 + pytest.param(
77 + AttentionCGSupport.UNIFORM_BATCH,
78 + 7,
79 + True,
80 + None,
81 + "up to 8, but TargetBackend allows at most 7",
82 + id="too-narrow",
83 + ),
84 + pytest.param(
85 + AttentionCGSupport.UNIFORM_BATCH,
86 + 8,
87 + False,
88 + None,
89 + "trims verification requests",
90 + id="host-query-lens",
91 + ),
92 + pytest.param(
93 + AttentionCGSupport.ALWAYS,
94 + None,
95 + True,
96 + (AttentionCGSupport.UNIFORM_BATCH, "EncoderBackend"),
97 + "EncoderBackend allows none",
98 + id="additional-group",
99 + ),
100 + ],
101 +)
102 +def test_manager_checks_target_varlen_cudagraph_bound(
103 + monkeypatch, target_support, target_bound, device_offsets, additional, error
104 +):
105 + """Target builders must replay varlen decode graphs of the verification
106 + width; draft-only groups are not checked."""
107 +
108 + def group(backend_name, layer_name, support, bound=None, mismatch=True):
109 + class Builder(AttentionMetadataBuilder):
110 + _cudagraph_support = support
111 +
112 + def __init__(self):
113 + pass
114
67 - class Builder:
68 - def __init__(self, support):
69 - self.support = support
115 + def build(self, common_prefix_len, common_attn_metadata, fast_build=False):
116 + raise NotImplementedError
117
71 - def get_cudagraph_support(self, *_args):
72 - return self.support
118 + if bound is not None:
119
74 - def group(layer_name, support):
75 - builder = Builder(support)
120 + @classmethod
121 + def get_varlen_cudagraph_max_query_len(cls, *_args):
122 + return bound
123 +
124 + backend = type(
125 + backend_name,
126 + (),
127 + {"supports_device_cpu_query_lens_mismatch": staticmethod(lambda: mismatch)},
128 + )
129 + builder = Builder()
130 return SimpleNamespace(
131 layer_names=[layer_name],
78 - backend=Backend,
132 + backend=backend,
133 kv_cache_spec=None,
134 get_metadata_builder=lambda _index: builder,
135 )
136
137 groups = [
138 [
85 - group("target", AttentionCGSupport.ALWAYS),
86 - group("draft", AttentionCGSupport.UNIFORM_BATCH),
139 + group(
140 + "TargetBackend", "target", target_support, target_bound, device_offsets
141 + ),
142 + group("DraftBackend", "draft", AttentionCGSupport.UNIFORM_BATCH),
143 ]
144 ]
89 - runner_support = AttentionCGSupportInfo(
90 - AttentionCGSupport.UNIFORM_BATCH, "DraftBackend"
91 - )
145 created = object()
146 monkeypatch.setattr(
147 adaptive_module,
@@ -96,20 +149,19 @@ def test_manager_scopes_varlen_check_without_weakening_runner_cg_mode(monkeypatc
149 lambda *_args, **_kwargs: created,
150 )
151
99 - manager = maybe_create_adaptive_verification_manager(
100 - enable_adaptive_verification=True,
101 - attn_groups=groups,
102 - attn_cg_support=runner_support,
103 - req_states=object(),
104 - query_start_loc=object(),
105 - num_bonus_tokens=1,
106 - max_total_logits=1,
107 - vllm_config=None,
108 - target_layer_names={"target"},
109 - )
110 -
111 - assert manager is created
112 - assert runner_support.min_cg_support == AttentionCGSupport.UNIFORM_BATCH
152 + with pytest.raises(ValueError, match=error) if error else nullcontext():
153 + manager = maybe_create_adaptive_verification_manager(
154 + enable_adaptive_verification=True,
155 + attn_groups=groups,
156 + req_states=SimpleNamespace(num_speculative_steps=7),
157 + query_start_loc=object(),
158 + num_bonus_tokens=1,
159 + max_total_logits=1,
160 + vllm_config=None,
161 + target_layer_names={"target"},
162 + additional_attn_cg_support=additional,
163 + )
164 + assert manager is created
165
166
167 def test_budget_stops_where_marginal_drafts_stop_paying_for_themselves():
tests/v1/worker/test_attn_utils.py
+42 -1
@@ -8,6 +8,7 @@ never addressed by the logical view.
8 """
9
10 from types import SimpleNamespace
11 +from typing import Any
12
13 import numpy as np
14 import pytest
@@ -99,12 +100,16 @@ def test_get_kv_cache_spec_resolves_hisparse_block_size(
100
101
102 class _FakeMetadataBuilder:
102 - def __init__(self, support: AttentionCGSupport):
103 + def __init__(self, support: AttentionCGSupport, varlen_bound: int | None = None):
104 self.support = support
105 + self.varlen_bound = varlen_bound
106
107 def get_cudagraph_support(self, *_args):
108 return self.support
109
110 + def get_varlen_cudagraph_max_query_len(self, *_args):
111 + return self.varlen_bound
112 +
113
114 class _TargetBackend:
115 @classmethod
@@ -183,6 +188,42 @@ def test_attention_checks_preserve_global_and_target_scoped_support():
188 )
189
190
191 +def test_varlen_cudagraph_unsupported_backend_checks_scoped_bounds():
192 + """ALWAYS passes without a bound, other builders need one at least as wide as
193 + the requested length, and NEVER fails whatever bound a builder reports."""
194 + config: Any = SimpleNamespace()
195 + spec = FullAttentionSpec(
196 + block_size=16,
197 + num_kv_heads=1,
198 + head_size=128,
199 + dtype=torch.bfloat16,
200 + )
201 +
202 + def group(
203 + backend: Any,
204 + layer_name: str,
205 + support: AttentionCGSupport,
206 + bound: int | None = None,
207 + ):
208 + builder: Any = _FakeMetadataBuilder(support, bound)
209 + attn_group = AttentionGroup(backend, [layer_name], spec, 0)
210 + attn_group.metadata_builders = [builder]
211 + return attn_group
212 +
213 + target = group(_TargetBackend, "target", AttentionCGSupport.ALWAYS)
214 + draft = group(_DraftBackend, "draft", AttentionCGSupport.UNIFORM_BATCH, 8)
215 + never = group(_DraftBackend, "never", AttentionCGSupport.NEVER, 8)
216 +
217 + unsupported = attn_utils.get_varlen_cudagraph_unsupported_backend
218 + assert unsupported([[target, draft]], config, 8) is None
219 + assert unsupported([[target, draft]], config, 9) == ("_DraftBackend", 8)
220 + assert (
221 + unsupported([[target, draft]], config, 9, checked_layer_names={"target"})
222 + is None
223 + )
224 + assert unsupported([[target, never]], config, 1) == ("_DraftBackend", None)
225 +
226 +
227 def test_get_kv_sharing_fast_prefill_eligible_layers(monkeypatch: pytest.MonkeyPatch):
228 """Fast prefill applies to the contiguous suffix of KV-sharing layers.
229
vllm/config/speculative.py
+2
@@ -631,6 +631,8 @@ class SpeculativeConfig:
631 "dspark",
632 )
633 factors.append(uses_aux_hidden_states)
634 + if self.method == "dspark":
635 + factors.append(self.enable_adaptive_verification)
636
637 if self.draft_model_config is not None:
638 factors.append(self.draft_model_config.compute_hash())
vllm/model_executor/models/gemma4_dspark.py
+22 -3
@@ -24,7 +24,11 @@ from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config
24
25 from .gemma4_mtp import Gemma4MTPAttention, Gemma4MTPDecoderLayer
26 from .qwen3_dflash import DFlashQwen3Model, _dflash_layer_causal
27 -from .qwen3_dspark import DSparkMarkovHead, Qwen3DSparkForCausalLM
27 +from .qwen3_dspark import (
28 + DSparkConfidenceHead,
29 + DSparkMarkovHead,
30 + Qwen3DSparkForCausalLM,
31 +)
32 from .utils import extract_layer_index, maybe_prefix
33
34
@@ -187,6 +191,18 @@ class Gemma4DSparkModel(DFlashQwen3Model):
191 config.markov_rank,
192 prefix=maybe_prefix(prefix, "markov_head"),
193 )
194 + self.confidence_head: DSparkConfidenceHead | None = None
195 + if getattr(config, "enable_confidence_head", False):
196 + with_markov = getattr(config, "confidence_head_with_markov", False)
197 + input_dim = config.hidden_size
198 + if with_markov:
199 + input_dim += config.markov_rank
200 + self.confidence_head = DSparkConfidenceHead(
201 + input_dim,
202 + prefix=maybe_prefix(prefix, "confidence_head"),
203 + bias=True,
204 + with_markov=with_markov,
205 + )
206
207 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
208 return self.embed_tokens(input_ids) * self.normalizer
@@ -295,9 +311,8 @@ class Gemma4DSparkForCausalLM(Qwen3DSparkForCausalLM):
311 params = dict(self.named_parameters())
312 params.update(dict(self.named_buffers()))
313 loaded: set[str] = set()
314 + includes_confidence_head = False
315 for name, w in weights:
299 - if "confidence_head" in name:
300 - continue
316 if "lm_head" not in name:
317 name = "model." + name
318 for pn, wn, shard in stacked:
@@ -310,5 +325,9 @@ class Gemma4DSparkForCausalLM(Qwen3DSparkForCausalLM):
325 p = params[name]
326 getattr(p, "weight_loader", default_weight_loader)(p, w)
327 loaded.add(name)
328 + if "confidence_head" in name:
329 + includes_confidence_head = True
330 + if not includes_confidence_head:
331 + self.model.confidence_head = None
332 self.model._build_fused_kv_buffers()
333 return loaded
vllm/models/minimax_m3/amd/indexer_aiter.py
+6 -5
@@ -52,7 +52,11 @@ from vllm.models.minimax_m3.common.sparse_attention import (
52 _minimax_m3_aiter_sparse_pa_requested,
53 )
54 from vllm.platforms import current_platform
55 -from vllm.v1.attention.backend import AttentionBackend, CommonAttentionMetadata
55 +from vllm.v1.attention.backend import (
56 + AttentionBackend,
57 + CommonAttentionMetadata,
58 + max_decode_query_len,
59 +)
60 from vllm.v1.attention.backends.utils import split_decodes_and_prefills
61 from vllm.v1.kv_cache_interface import AttentionSpec
62
@@ -97,10 +101,7 @@ def aiter_indexer_max_decode_query_len(vllm_config: VllmConfig) -> int:
101 since that is what the builder splits the batch on and therefore what the
102 decode kernel will actually be handed.
103 """
100 - spec = vllm_config.speculative_config
101 - if spec is None or spec.num_speculative_tokens is None:
102 - return 1
103 - return 1 + (2 if spec.parallel_drafting else 1) * spec.num_speculative_tokens
104 + return max_decode_query_len(vllm_config)
105
106
107 @cache
vllm/transformers_utils/config.py
+14 -1
@@ -387,7 +387,20 @@ class HFConfigParser(ConfigParserBase):
387 if extra_layer_types := _PATCH_HF_ALLOWED_LAYER_TYPES.get(model_type):
388 _patch_hf_transformers_allowed_layer_types(extra_layer_types)
389
390 - if model_type in _SPECULATIVE_DECODING_CONFIGS:
390 + rope_parameters = config_dict.get("rope_parameters") or {}
391 + if model_type == "gemma4_text" and "full_attention" in rope_parameters:
392 + from transformers import Gemma4TextConfig
393 +
394 + # Published DSpark configs mix redundant scalar entries with per-layer
395 + # RoPE dicts. Remove them before Transformers' constructor validates.
396 + config_dict["rope_parameters"] = {
397 + k: v
398 + for k, v in rope_parameters.items()
399 + if k not in ("rope_type", "rope_theta")
400 + }
401 + kwargs.setdefault("name_or_path", str(model))
402 + config = Gemma4TextConfig.from_dict(config_dict, **kwargs)
403 + elif model_type in _SPECULATIVE_DECODING_CONFIGS:
404 config_class = _CONFIG_REGISTRY[model_type]
405 config = config_class.from_pretrained(
406 model,
vllm/v1/attention/backend.py
+45 -6
@@ -573,6 +573,26 @@ class AttentionCGSupport(Enum):
573 """NO cudagraph support"""
574
575
576 +def max_decode_query_len(vllm_config: "VllmConfig") -> int:
577 + """Widest request a spec-as-decode builder treats as a decode.
578 +
579 + On model runner V2 this is a verification request, 1 +
580 + num_speculative_tokens: no V2 speculator builds a wider query. The reorder
581 + threshold and the code that mirrors the builders' decode/prefill split all
582 + read this, so they cannot drift apart.
583 + """
584 + speculative_config = getattr(vllm_config, "speculative_config", None)
585 + if speculative_config is None or speculative_config.num_speculative_tokens is None:
586 + return 1
587 + num_speculative_tokens = speculative_config.num_speculative_tokens
588 + max_query_len = 1 + num_speculative_tokens
589 + if speculative_config.parallel_drafting and not vllm_config.use_v2_model_runner:
590 + # Model runner V1 only; remove with it. Its parallel drafter appends up
591 + # to num_speculative_tokens mask slots to each verified request.
592 + max_query_len += num_speculative_tokens
593 + return max_query_len
594 +
595 +
596 class AttentionMetadataBuilder(ABC, Generic[M]):
597 # Does this backend/builder support CUDA Graphs for attention (default: no).
598 # Do not access directly. Call get_cudagraph_support() instead.
@@ -616,6 +636,30 @@ class AttentionMetadataBuilder(ABC, Generic[M]):
636 """Get the cudagraph support level of this builder class."""
637 return cls._cudagraph_support
638
639 + @classmethod
640 + def get_varlen_cudagraph_max_query_len(
641 + cls: type["AttentionMetadataBuilder"],
642 + vllm_config: "VllmConfig",
643 + kv_cache_spec: "KVCacheSpec",
644 + ) -> int | None:
645 + """Get the largest per-request query length L of the variable-length
646 + decode batches a FULL cudagraph of this builder class can replay.
647 +
648 + The graph must replay any decode batch in which every real request has
649 + between 1 and L query tokens and every padding request has 0. Lengths
650 + come from the device query_start_loc; host metadata carries only the
651 + token count and an upper bound on the per-request length. Batches with
652 + a prefill never replay these graphs. Whether host metadata may
653 + understate device query lengths at all is a separate backend question;
654 + see supports_device_cpu_query_lens_mismatch().
655 +
656 + Returns:
657 + L, or None for builders reporting ALWAYS, which replay any batch,
658 + and for builders that cannot replay variable-length batches.
659 +
660 + """
661 + return None
662 +
663 def _init_reorder_batch_threshold(
664 self,
665 reorder_batch_threshold: int | None = 1,
@@ -632,14 +676,9 @@ class AttentionMetadataBuilder(ABC, Generic[M]):
676 speculative_config is not None
677 and speculative_config.num_speculative_tokens is not None
678 ):
635 - max_num_queries_for_spec = (
636 - 1
637 - + (2 if speculative_config.parallel_drafting else 1)
638 - * speculative_config.num_speculative_tokens
639 - )
679 self.reorder_batch_threshold = max(
680 self.reorder_batch_threshold,
642 - max_num_queries_for_spec,
681 + max_decode_query_len(self.vllm_config),
682 )
683
684 if (
vllm/v1/attention/backends/flashinfer.py
+56 -18
@@ -67,6 +67,7 @@ from vllm.v1.attention.backend import (
67 AttentionType,
68 CommonAttentionMetadata,
69 MultipleOf,
70 + max_decode_query_len,
71 )
72 from vllm.v1.attention.backends.utils import (
73 get_dcp_local_seq_lens,
@@ -856,15 +857,8 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
857 self.use_xqa = (
858 self.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.XQA
859 )
859 - # Adaptive verification trims drafts on device, so decode query lengths
860 - # must come from the device qo_indptr; only trtllm-gen supports that
861 - # (the selector already rejects the other configurations).
862 - self.use_trtllm_gen_varlen_decode = (
863 - speculative_config is not None
864 - and speculative_config.enable_adaptive_verification
865 - and self.flashinfer_trtllm_api_decode_kernel
866 - == FlashInferDecodeKernel.TRTLLM_GEN
867 - and not self.use_dcp
860 + self.use_trtllm_gen_varlen_decode = self._uses_trtllm_gen_varlen_decode(
861 + vllm_config, self.flashinfer_trtllm_api_decode_kernel, self.use_dcp
862 )
863 self._init_reorder_batch_threshold(
864 1,
@@ -1011,6 +1005,31 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1005 else:
1006 return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
1007
1008 + @override # type: ignore[misc]
1009 + @classmethod
1010 + def get_varlen_cudagraph_max_query_len(
1011 + cls: type["FlashInferMetadataBuilder"],
1012 + vllm_config: VllmConfig,
1013 + kv_cache_spec: KVCacheSpec,
1014 + ) -> int | None:
1015 + # The uniform gate already covers DCP, head geometry, the XQA head-dim
1016 + # limit and causality, and yields None for subclasses that force NEVER.
1017 + if (
1018 + cls.get_cudagraph_support(vllm_config, kv_cache_spec)
1019 + != AttentionCGSupport.UNIFORM_BATCH
1020 + ):
1021 + return None
1022 + # The kernel selector asserts off the trtllm-capable architectures, so
1023 + # call it only past the gate, as get_cudagraph_support does.
1024 + if not cls._uses_trtllm_gen_varlen_decode(
1025 + vllm_config,
1026 + cls._get_flashinfer_trtllm_api_decode_kernel(),
1027 + use_dcp=vllm_config.parallel_config.decode_context_parallel_size > 1,
1028 + ):
1029 + return None
1030 + # Longer requests are split off as prefills.
1031 + return max_decode_query_len(vllm_config)
1032 +
1033 def _get_workspace_buffer(self):
1034 if self._workspace_buffer is None:
1035 buffer_size = envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE
@@ -1039,6 +1058,24 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1058 def set_workspace_buffer(self, workspace_buffer: torch.Tensor):
1059 self._workspace_buffer = workspace_buffer
1060
1061 + @staticmethod
1062 + def _uses_trtllm_gen_varlen_decode(
1063 + vllm_config: VllmConfig,
1064 + decode_kernel: FlashInferDecodeKernel | None,
1065 + use_dcp: bool,
1066 + ) -> bool:
1067 + """Whether decode reads per-request query lengths from the device
1068 + qo_indptr. Adaptive verification trims drafts on device, and only
1069 + trtllm-gen supports that (the selector already rejects the other
1070 + configurations)."""
1071 + speculative_config = vllm_config.speculative_config
1072 + return (
1073 + speculative_config is not None
1074 + and speculative_config.enable_adaptive_verification
1075 + and decode_kernel == FlashInferDecodeKernel.TRTLLM_GEN
1076 + and not use_dcp
1077 + )
1078 +
1079 @staticmethod
1080 def _get_flashinfer_trtllm_api_decode_kernel() -> FlashInferDecodeKernel:
1081 if current_platform.is_device_capability(
@@ -1302,7 +1339,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1339 split_decodes_and_prefills(
1340 common_attn_metadata,
1341 decode_threshold=self.reorder_batch_threshold,
1305 - require_uniform=not self.use_xqa,
1342 + require_uniform=not (
1343 + self.use_xqa or self.use_trtllm_gen_varlen_decode
1344 + ),
1345 )
1346 )
1347 else:
@@ -1704,15 +1743,14 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1743 bool(causal),
1744 )
1745 elif self.use_trtllm_gen_varlen_decode:
1707 - # CPU lens are only an upper bound, device qo_indptr is
1708 - # the truth.
1709 - decode_q_lens = (
1710 - qo_indptr_cpu[1 : num_decodes + 1] - qo_indptr_cpu[:num_decodes]
1746 + # CPU placeholders may underestimate query lengths;
1747 + # use a safe upper bound.
1748 + q_len_per_req = min(
1749 + common_attn_metadata.max_query_len, self.reorder_batch_threshold
1750 )
1712 - max_q_upper = int(decode_q_lens.max())
1713 - if max_q_upper > 1:
1714 - q_len_per_req = max_q_upper
1715 - q_cu_seq_lens = qo_indptr[: num_decodes + 1]
1751 + # Always pass GPU offsets so a graph captured with uniform
1752 + # query lengths can replay batches with variable query lengths.
1753 + q_cu_seq_lens = qo_indptr[: num_decodes + 1]
1754 attn_metadata.decode = FlashInferTrtllmAPIDecode(
1755 kernel=self.flashinfer_trtllm_api_decode_kernel,
1756 block_tables=block_table_tensor[:num_decodes],
vllm/v1/attention/backends/mla/index_group.py
+2 -10
@@ -12,6 +12,7 @@ from vllm import _custom_ops as ops
12 from vllm.config import VllmConfig
13 from vllm.forward_context import in_piecewise_cudagraph
14 from vllm.utils.torch_utils import current_stream
15 +from vllm.v1.attention.backend import max_decode_query_len
16 from vllm.v1.attention.backends.mla.sparse_utils import (
17 triton_convert_req_index_to_global_index,
18 )
@@ -496,17 +497,8 @@ class SparseMLAIndexGroupBuilder:
497
498
499 def get_sparse_mla_index_group_max_rows(vllm_config: VllmConfig) -> int:
499 - max_query_len = 1
500 - speculative_config = vllm_config.speculative_config
501 - if (
502 - speculative_config is not None
503 - and speculative_config.num_speculative_tokens is not None
504 - ):
505 - max_query_len += speculative_config.num_speculative_tokens * (
506 - 2 if speculative_config.parallel_drafting else 1
507 - )
500 scheduler_config = vllm_config.scheduler_config
501 return min(
502 scheduler_config.max_num_batched_tokens,
511 - scheduler_config.max_num_seqs * max_query_len,
503 + scheduler_config.max_num_seqs * max_decode_query_len(vllm_config),
504 )
vllm/v1/attention/backends/mla/sparse_swa.py
+4 -7
@@ -26,6 +26,7 @@ from vllm.v1.attention.backend import (
26 AttentionMetadataBuilder,
27 CommonAttentionMetadata,
28 MultipleOf,
29 + max_decode_query_len,
30 )
31 from vllm.v1.attention.backends.mla.compressor_utils import (
32 get_dspark_swa_index_width,
@@ -456,14 +457,10 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
457 self.num_speculative_tokens = (
458 spec_config.num_speculative_tokens if spec_config else 0
459 )
459 - # Decode can have query_len up to
460 - # 1 + (2 if parallel drafting else 1) * num_speculative_tokens.
460 # sparse_swa has no MQA-vs-dense-MHA routing, so multi-token queries take
462 - # the prefill path and the decode/prefill split stays at that width.
463 - spec_mult = (
464 - 2 if (spec_config is not None and spec_config.parallel_drafting) else 1
465 - )
466 - self.decode_threshold = 1 + spec_mult * self.num_speculative_tokens
461 + # the prefill path and the decode/prefill split stays at the widest
462 + # decode.
463 + self.decode_threshold = max_decode_query_len(self.vllm_config)
464 self.reorder_batch_threshold = None
465
466 hf_config = self.vllm_config.model_config.hf_config
vllm/v1/attention/ops/dcp.py
+4 -10
@@ -27,6 +27,7 @@ from vllm.model_executor.warmup.jit_warmup_triton_helper import (
27 triton_scalar_specialization_rep,
28 )
29 from vllm.triton_utils import tl, triton
30 +from vllm.v1.attention.backend import max_decode_query_len
31 from vllm.v1.attention.ops.cp_common import (
32 DirectCPWorkspace,
33 direct_cp_enabled,
@@ -1011,16 +1012,9 @@ def dcp_a2a_lse_reduce(
1012
1013 def get_dcp_workspace_max_num_tokens(vllm_config: VllmConfig) -> int:
1014 scheduler_config = vllm_config.scheduler_config
1014 - speculative_config = vllm_config.speculative_config
1015 - speculative_tokens = vllm_config.num_speculative_tokens
1016 - tokens_per_seq = (
1017 - 1
1018 - + (
1019 - 2
1020 - if speculative_config is not None and speculative_config.parallel_drafting
1021 - else 1
1022 - )
1023 - * speculative_tokens
1015 + # num_speculative_tokens also counts a diffusion canvas.
1016 + tokens_per_seq = max(
1017 + 1 + vllm_config.num_speculative_tokens, max_decode_query_len(vllm_config)
1018 )
1019 return min(
1020 scheduler_config.max_num_batched_tokens,
vllm/v1/hisparse/runtime.py
+3 -14
@@ -21,6 +21,7 @@ from vllm.logger import init_logger
21 from vllm.platforms import current_platform
22 from vllm.utils.math_utils import round_up
23 from vllm.utils.torch_utils import current_stream
24 +from vllm.v1.attention.backend import max_decode_query_len
25 from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
26 from vllm.v1.simple_kv_offload.cuda_mem_ops import pin_tensor
27
@@ -46,20 +47,8 @@ def is_hisparse_decode_batch(attn_metadata: Any | None) -> bool:
47 )
48
49
49 -def _get_max_decode_query_len(vllm_config: VllmConfig) -> int:
50 - speculative_config = getattr(vllm_config, "speculative_config", None)
51 - if (
52 - speculative_config is not None
53 - and speculative_config.num_speculative_tokens is not None
54 - ):
55 - return 1 + speculative_config.num_speculative_tokens * (
56 - 2 if speculative_config.parallel_drafting else 1
57 - )
58 - return 1
59 -
60 -
50 def _get_max_swap_rows(vllm_config: VllmConfig) -> int:
62 - max_query_len = _get_max_decode_query_len(vllm_config)
51 + max_query_len = max_decode_query_len(vllm_config)
52 scheduler_config = vllm_config.scheduler_config
53 return min(
54 scheduler_config.max_num_batched_tokens,
@@ -120,7 +109,7 @@ class ResolvedHiSparseConfig:
109 if config is None:
110 return None
111
123 - max_query_len = _get_max_decode_query_len(vllm_config)
112 + max_query_len = max_decode_query_len(vllm_config)
113 min_device_buffer_size = max_query_len * model_top_k
114 # Retain the worst-case union of every speculative step's top-k and
115 # leave one additional top-k of LRU slack.
vllm/v1/worker/gpu/attn_utils.py
+34
@@ -328,6 +328,40 @@ def get_query_lens_mismatch_unsupported_backend(
328 return None
329
330
331 +def get_varlen_cudagraph_unsupported_backend(
332 + attn_groups: list[list[AttentionGroup]],
333 + vllm_config: VllmConfig,
334 + max_query_len: int,
335 + checked_layer_names: set[str] | None = None,
336 +) -> tuple[str, int | None] | None:
337 + """Name the first backend whose FULL cudagraphs cannot replay decode batches
338 + of 1 to max_query_len tokens per request, with its bound, if any.
339 +
340 + ALWAYS builders replay any batch; the others need a bound of at least
341 + max_query_len. See AttentionMetadataBuilder.get_varlen_cudagraph_max_query_len().
342 + """
343 + for groups in attn_groups:
344 + for group in groups:
345 + if checked_layer_names is not None and checked_layer_names.isdisjoint(
346 + group.layer_names
347 + ):
348 + continue
349 + builder = group.get_metadata_builder(0)
350 + spec = group.kv_cache_spec
351 + support = builder.get_cudagraph_support(vllm_config, spec)
352 + if support == AttentionCGSupport.ALWAYS:
353 + continue
354 + bound = (
355 + # NEVER means never, whatever an inherited override says.
356 + None
357 + if support == AttentionCGSupport.NEVER
358 + else builder.get_varlen_cudagraph_max_query_len(vllm_config, spec)
359 + )
360 + if bound is None or bound < max_query_len:
361 + return group.backend.__name__, bound
362 + return None
363 +
364 +
365 def init_kv_cache(
366 forward_context: dict[str, Any],
367 kv_cache_config: KVCacheConfig,
vllm/v1/worker/gpu/cudagraph_utils.py
+2 -1
@@ -69,7 +69,8 @@ class BatchExecutionDescriptor:
69 num_reqs: int | None # None means no request padding is needed (PIECEWISE graphs)
70 uniform_token_count: int | None = None
71 # Upper bound on per-request query length. Varlen decode graphs leave
72 - # uniform_token_count unset, so this is what keeps a prefill batch out of one.
72 + # uniform_token_count unset, so this is what keeps a prefill batch out of one:
73 + # the runner passes None for any batch with a prefill.
74 max_query_len: int | None = None
75 num_active_loras: int = 0
76 # Number of microbatches the batch is split into (DBO). 1 means no splitting.
vllm/v1/worker/gpu/dp_utils.py
+2 -1
@@ -243,7 +243,8 @@ def dispatch_cg_and_sync_dp(
243 dp_size: Data-parallel world size. 1 skips all cross-rank work.
244 dp_rank: This rank's index in the DP group.
245 max_query_len: Upper bound on per-request query length, for selecting
246 - varlen decode graphs. None means the graph must not constrain it.
246 + varlen decode graphs. None, as for a batch with a prefill, keeps
247 + the batch out of graphs that constrain it.
248 need_eager: Force `CUDAGraphMode.NONE` instead of dispatching.
249 num_active_loras: Active LoRA count for this rank. Does not need
250 cross-rank agreement; it never changes a bucket's token count.
vllm/v1/worker/gpu/model_runner.py
+5 -2
@@ -640,7 +640,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
640 self.speculator, "enable_adaptive_verification", False
641 ),
642 attn_groups=self.attn_groups,
643 - attn_cg_support=attn_cg_support,
643 req_states=self.req_states,
644 query_start_loc=self.input_buffers.query_start_loc,
645 num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step,
@@ -1646,12 +1645,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1645 # Get batch descriptor and sync across DP ranks.
1646 num_reqs = len(scheduler_output.num_scheduled_tokens)
1647 num_toks = scheduler_output.total_num_scheduled_tokens
1649 - max_query_len = max(scheduler_output.num_scheduled_tokens.values())
1648 + max_query_len: int | None = max(scheduler_output.num_scheduled_tokens.values())
1649 batch_req_state, uniform_tok_count = self.gather_batch_req_state(
1650 scheduler_output, dummy_run
1651 )
1652 if batch_req_state is not None:
1653 num_toks = batch_req_state.num_tokens
1654 + if batch_req_state.has_prefill:
1655 + # Varlen decode graphs replay decodes only, and their bound
1656 + # alone would admit a short prefill.
1657 + max_query_len = None
1658 if self.pcp_manager is not None:
1659 num_toks = self.pcp_manager.get_num_tokens_for_dispatch(
1660 batch_req_state.num_scheduled_tokens,
vllm/v1/worker/gpu/spec_decode/adaptive_verification.py
+20 -17
@@ -19,8 +19,8 @@ from vllm.v1.attention.backend import AttentionCGSupport
19 from vllm.v1.utils import CpuGpuBuffer
20 from vllm.v1.worker.gpu.async_utils import StepTimingSample, stream
21 from vllm.v1.worker.gpu.attn_utils import (
22 - get_attn_cg_support,
22 get_query_lens_mismatch_unsupported_backend,
23 + get_varlen_cudagraph_unsupported_backend,
24 )
25
26 logger = init_logger(__name__)
@@ -28,7 +28,6 @@ _PROFILE_REPLAYS = 5
28
29 if TYPE_CHECKING:
30 from vllm.config import VllmConfig
31 - from vllm.v1.worker.gpu.attn_utils import AttentionCGSupportInfo
31 from vllm.v1.worker.gpu.input_batch import InputBatch
32 from vllm.v1.worker.gpu.states import RequestState
33 from vllm.v1.worker.utils import AttentionGroup
@@ -452,7 +451,6 @@ def maybe_create_adaptive_verification_manager(
451 *,
452 enable_adaptive_verification: bool,
453 attn_groups: list[list["AttentionGroup"]],
455 - attn_cg_support: "AttentionCGSupportInfo",
454 req_states: "RequestState",
455 query_start_loc: torch.Tensor,
456 num_bonus_tokens: int,
@@ -477,25 +475,30 @@ def maybe_create_adaptive_verification_manager(
475 "use a backend that does."
476 )
477
480 - target_attn_cg_support = attn_cg_support
481 - if target_layer_names is not None:
482 - target_attn_cg_support = get_attn_cg_support(
478 + # The runner's decode_query_len, the width varlen decode graphs capture.
479 + max_query_len = req_states.num_speculative_steps + num_bonus_tokens
480 + unsupported: tuple[str | None, int | None] | None = (
481 + get_varlen_cudagraph_unsupported_backend(
482 attn_groups,
483 vllm_config,
484 + max_query_len,
485 checked_layer_names=target_layer_names,
486 )
487 - if additional_attn_cg_support is not None:
488 - target_attn_cg_support = target_attn_cg_support.narrow(
489 - *additional_attn_cg_support
490 - )
491 - if target_attn_cg_support.min_cg_support != AttentionCGSupport.ALWAYS:
487 + )
488 + if unsupported is None and additional_attn_cg_support is not None:
489 + # Groups built outside init_attn_backend report only their support
490 + # level, and without a bound only ALWAYS replays varlen batches.
491 + additional_support, additional_backend = additional_attn_cg_support
492 + if additional_support != AttentionCGSupport.ALWAYS:
493 + unsupported = additional_backend, None
494 + if unsupported is not None:
495 + backend, bound = unsupported
496 + allowed = "none" if bound is None else f"at most {bound}"
497 raise ValueError(
493 - "Adaptive verification captures varlen decode cudagraphs, so every"
494 - " target attention builder must report AttentionCGSupport.ALWAYS, but "
495 - f"{target_attn_cg_support.min_cg_attn_backend} reports "
496 - f"{target_attn_cg_support.min_cg_support}. Pass "
497 - "enable_adaptive_verification=false in the speculative config, or "
498 - "use a backend that does."
498 + "Adaptive verification replays decode cudagraphs whose per-request "
499 + f"query lengths vary up to {max_query_len}, but {backend} allows "
500 + f"{allowed}. Pass enable_adaptive_verification=false in the "
501 + "speculative config, or use a backend that does."
502 )
503
504 return AdaptiveVerificationManager(