@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 97388c44f9

[Bugfix] Make DSV4 sparse MLA work end-to-end for plain decode, MTP, and DSpark (#51538)

Gabriel Wu committed Aug 15, 2026 at 22:01 UTC 97388c44f9c608f83318e0c9540a536cd7de3e0d
20 files changed +797 -120
csrc/libtorch_stable/cooperative_topk.cuh
+4 -1
@@ -484,7 +484,10 @@ __device__ void large_topk(const float* __restrict__ row_input,
484 template <uint32_t TopK, uint32_t CS>
485 __device__ void cooperative_topk_body(CooperativeTopKParams<TopK> params) {
486 const auto rank = blockIdx.y, row = blockIdx.x, tx = threadIdx.x;
487 - const auto sl = params.lengths[row];
487 + // Clamp at 0: `sl` is compared signed here but cast to uint32 below, so a
488 + // negative length would otherwise emit indices 0..TopK-1 as valid instead
489 + // of the -1 padding.
490 + const int32_t sl = params.lengths[row] > 0 ? params.lengths[row] : 0;
491 int32_t* out = params.output + row * TopK;
492 const float* in = params.input + row * params.stride;
493
csrc/libtorch_stable/persistent_topk.cuh
+20 -1
@@ -906,7 +906,26 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
906 uint32_t row_idx = group_id + iter * num_groups;
907 if (row_idx >= params.num_rows) break;
908
909 - const uint32_t seq_len = params.lengths[row_idx];
909 + // Clamp the row length before any decision is made on it.
910 + //
911 + // `lengths` is int32 and is consumed here as uint32, so a negative value
912 + // (e.g. a padded decode slot whose per-token context length underflowed)
913 + // would reinterpret as ~4e9 and sail past every threshold below. Any
914 + // value beyond the row width would also read into the next row.
915 + //
916 + // Clamping to max_seq_len additionally keeps this per-row decision
917 + // consistent with the `cta_in_group != 0` early exit above, which is
918 + // taken from the host-side scalar: when max_seq_len <= RADIX_THRESHOLD
919 + // the non-leader CTAs return immediately, so a leader that reached the
920 + // cooperative radix path would wait on the inter-CTA barrier for peers
921 + // that no longer exist and spin until the kernel is killed.
922 + const int32_t raw_len = params.lengths[row_idx];
923 + const uint32_t row_bound =
924 + params.stride < params.max_seq_len ? params.stride : params.max_seq_len;
925 + const uint32_t non_negative_len =
926 + raw_len > 0 ? static_cast<uint32_t>(raw_len) : 0u;
927 + const uint32_t seq_len =
928 + non_negative_len < row_bound ? non_negative_len : row_bound;
929 int32_t* row_output = params.output + row_idx * params.top_k;
930 const float* row_input = params.input + row_idx * params.stride;
931
tests/kernels/attention/test_flashmla_sparse.py
+118
@@ -225,6 +225,7 @@ def test_flashinfer_sparse_indices_cache(monkeypatch):
225 token_to_req_indices=torch.tensor([0, 1, 1], dtype=torch.int32),
226 decode_swa_indices=torch.tensor([[5, 6, -1, -1]], dtype=torch.int32),
227 decode_swa_lens=torch.tensor([2], dtype=torch.int32),
228 + decode_swa_width=4,
229 is_valid_token=torch.tensor([True], dtype=torch.bool),
230 num_decodes=1,
231 num_prefills=1,
@@ -331,3 +332,120 @@ def test_flashinfer_sparse_indices_cache(monkeypatch):
332 assert builder_calls == 4
333 assert sparse_indices_third is not sparse_indices_fourth
334 assert sparse_lens_third is not sparse_lens_fourth
335 +
336 +
337 +def test_flashinfer_sparse_index_preserves_logical_window(monkeypatch):
338 + from vllm.models.deepseek_v4.nvidia import flashinfer_sparse as flashinfer_mod
339 + from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
340 +
341 + captured_shapes_and_windows: list[tuple[int, int]] = []
342 +
343 + def fake_build(*args, **kwargs):
344 + # window_size is the 12th positional arg of
345 + # build_flashinfer_mixed_sparse_indices.
346 + captured_shapes_and_windows.append((args[0].shape[-1], args[11]))
347 + num_tokens = args[0].shape[0] + args[3].shape[0]
348 + return (
349 + torch.zeros((num_tokens, 1), dtype=torch.int32),
350 + torch.zeros((num_tokens,), dtype=torch.int32),
351 + )
352 +
353 + monkeypatch.setattr(
354 + flashinfer_mod, "build_flashinfer_mixed_sparse_indices", fake_build
355 + )
356 +
357 + attn = object.__new__(flashinfer_mod.DeepseekV4FlashInferMLAAttention)
358 + attn.compress_ratio = 1
359 + attn.window_size = 4
360 + attn.topk_indices_buffer = torch.zeros((4, 0), dtype=torch.int32)
361 +
362 + wide_width = 8
363 + wide_indices = torch.full((1, wide_width), -1, dtype=torch.int32)
364 + wide_indices[0, :2] = torch.tensor([5, 6], dtype=torch.int32)
365 + wide_metadata = DeepseekSparseSWAMetadata(
366 + block_table=torch.tensor([[0, 1]], dtype=torch.int32),
367 + slot_mapping=torch.tensor([0], dtype=torch.int64),
368 + block_size=64,
369 + seq_lens=torch.tensor([8], dtype=torch.int32),
370 + query_start_loc=torch.tensor([0, 1], dtype=torch.int32),
371 + query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.int32),
372 + token_to_req_indices=torch.tensor([0], dtype=torch.int32),
373 + decode_swa_indices=wide_indices,
374 + decode_swa_lens=torch.tensor([2], dtype=torch.int32),
375 + decode_swa_width=wide_width,
376 + is_valid_token=torch.tensor([True], dtype=torch.bool),
377 + num_decodes=1,
378 + num_prefills=0,
379 + num_decode_tokens=1,
380 + num_prefill_tokens=0,
381 + )
382 + attn._build_sparse_index_metadata(
383 + kv_cache=None,
384 + swa_k_cache=torch.empty((1, 64, 512), dtype=torch.bfloat16),
385 + swa_metadata=wide_metadata,
386 + attn_metadata=None,
387 + swa_only=True,
388 + )
389 + assert captured_shapes_and_windows == [(wide_width, attn.window_size)]
390 +
391 + empty_width = 8
392 + empty_metadata = DeepseekSparseSWAMetadata(
393 + block_table=torch.tensor([[0, 1]], dtype=torch.int32),
394 + slot_mapping=torch.tensor([0, 1], dtype=torch.int64),
395 + block_size=64,
396 + seq_lens=torch.tensor([8], dtype=torch.int32),
397 + query_start_loc=torch.tensor([0, 2], dtype=torch.int32),
398 + query_start_loc_cpu=torch.tensor([0, 2], dtype=torch.int32),
399 + token_to_req_indices=torch.tensor([0, 0], dtype=torch.int32),
400 + decode_swa_indices=torch.empty((0, 1, empty_width), dtype=torch.int32),
401 + decode_swa_lens=torch.empty((0,), dtype=torch.int32),
402 + decode_swa_width=empty_width,
403 + is_valid_token=torch.tensor([True, True], dtype=torch.bool),
404 + num_decodes=0,
405 + num_prefills=1,
406 + num_decode_tokens=0,
407 + num_prefill_tokens=2,
408 + )
409 + attn._build_sparse_index_metadata(
410 + kv_cache=None,
411 + swa_k_cache=torch.empty((1, 64, 512), dtype=torch.bfloat16),
412 + swa_metadata=empty_metadata,
413 + attn_metadata=None,
414 + swa_only=True,
415 + )
416 + assert captured_shapes_and_windows == [
417 + (wide_width, attn.window_size),
418 + (empty_width, attn.window_size),
419 + ]
420 +
421 +
422 +def test_flashinfer_mixed_sparse_indices_separates_window_and_padded_width():
423 + from vllm.models.deepseek_v4.common.ops.cache_utils import (
424 + build_flashinfer_mixed_sparse_indices,
425 + )
426 +
427 + device = torch.device("cuda")
428 + padded_width = 8
429 + logical_window = 4
430 + sparse_indices, sparse_lens = build_flashinfer_mixed_sparse_indices(
431 + decode_swa_indices=torch.empty(
432 + (0, padded_width), dtype=torch.int32, device=device
433 + ),
434 + decode_compressed_indices=None,
435 + decode_compressed_topk_lens=None,
436 + prefill_topk_indices=torch.empty((1, 0), dtype=torch.int32, device=device),
437 + query_start_loc=torch.tensor([0, 1], dtype=torch.int32, device=device),
438 + seq_lens=torch.tensor([logical_window], dtype=torch.int32, device=device),
439 + token_to_req_indices=torch.tensor([0], dtype=torch.int32, device=device),
440 + swa_block_table=torch.tensor([[0]], dtype=torch.int32, device=device),
441 + swa_block_size=64,
442 + compressed_block_table=None,
443 + compressed_block_size=64,
444 + window_size=logical_window,
445 + compress_ratio=1,
446 + topk=0,
447 + )
448 +
449 + assert sparse_indices.shape == (1, padded_width)
450 + assert sparse_indices[0].cpu().tolist() == [0, 1, 2, 3, -1, -1, -1, -1]
451 + assert sparse_lens.cpu().tolist() == [padded_width]
tests/kernels/moe/test_ocp_mx_moe.py
+95
@@ -1095,6 +1095,101 @@ def test_convert_standard_mxfp4_weights_for_flashinfer_cutlass(
1095 torch.testing.assert_close(interleaved_scales[1], w2_scale)
1096
1097
1098 +def test_gpt_oss_quant_config_supplies_clamped_swiglu_params():
1099 + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend
1100 + from vllm.model_executor.layers.quantization.mxfp4 import GptOssMxfp4MoEMethod
1101 +
1102 + fake_method = types.SimpleNamespace(
1103 + mxfp4_backend=Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
1104 + w13_precision_config=None,
1105 + w2_precision_config=None,
1106 + )
1107 + fake_layer = types.SimpleNamespace(
1108 + w13_weight_scale=torch.zeros(2, 4, 2, dtype=torch.uint8),
1109 + w2_weight_scale=torch.zeros(2, 4, 2, dtype=torch.uint8),
1110 + )
1111 +
1112 + quant_config = GptOssMxfp4MoEMethod.get_fused_moe_quant_config(
1113 + fake_method, fake_layer
1114 + )
1115 +
1116 + assert quant_config is not None
1117 + assert quant_config.gemm1_alpha == 1.702
1118 + assert quant_config.gemm1_beta == 1.0
1119 + assert quant_config.gemm1_clamp_limit == 7.0
1120 +
1121 +
1122 +@pytest.mark.skipif(
1123 + not current_platform.is_cuda(),
1124 + reason="CUDA is required for FlashInferExperts parameter storage",
1125 +)
1126 +@pytest.mark.parametrize("supplied", [False, True])
1127 +def test_flashinfer_experts_swiglu_params_follow_quant_config(supplied: bool):
1128 + from vllm.model_executor.layers.fused_moe.activation import MoEActivation
1129 + from vllm.model_executor.layers.fused_moe.config import (
1130 + FusedMoEConfig,
1131 + FusedMoEParallelConfig,
1132 + RoutingMethodType,
1133 + mxfp4_w4a16_moe_quant_config,
1134 + )
1135 + from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import (
1136 + FlashInferExperts,
1137 + )
1138 +
1139 + num_experts, intermediate_size, hidden_size = 2, 128, 256
1140 + w1_scale = torch.zeros(
1141 + (num_experts, 2 * intermediate_size, hidden_size // 32),
1142 + dtype=torch.uint8,
1143 + device="cuda",
1144 + )
1145 + w2_scale = torch.zeros(
1146 + (num_experts, hidden_size, intermediate_size // 32),
1147 + dtype=torch.uint8,
1148 + device="cuda",
1149 + )
1150 + extra = (
1151 + {"gemm1_alpha": 0.5, "gemm1_beta": 0.25, "gemm1_clamp_limit": 3.0}
1152 + if supplied
1153 + else {}
1154 + )
1155 + quant_config = mxfp4_w4a16_moe_quant_config(
1156 + w1_scale=w1_scale, w2_scale=w2_scale, **extra
1157 + )
1158 + moe_config = FusedMoEConfig(
1159 + num_experts=num_experts,
1160 + experts_per_token=1,
1161 + hidden_dim=hidden_size,
1162 + intermediate_size=intermediate_size,
1163 + num_local_experts=num_experts,
1164 + num_logical_experts=num_experts,
1165 + activation=MoEActivation.SILU,
1166 + device="cuda",
1167 + routing_method=RoutingMethodType.TopK,
1168 + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
1169 + in_dtype=torch.bfloat16,
1170 + )
1171 + experts = FlashInferExperts(moe_config=moe_config, quant_config=quant_config)
1172 +
1173 + if supplied:
1174 + for name, value in (
1175 + ("gemm1_alpha", 0.5),
1176 + ("gemm1_beta", 0.25),
1177 + ("gemm1_clamp_limit", 3.0),
1178 + ):
1179 + parameter = getattr(experts, name)
1180 + assert isinstance(parameter, torch.Tensor)
1181 + assert parameter.dtype == torch.float32
1182 + assert parameter.shape == (num_experts,)
1183 + torch.testing.assert_close(
1184 + parameter.cpu(),
1185 + torch.full((num_experts,), value, dtype=torch.float32),
1186 + )
1187 + else:
1188 + assert experts.gemm1_alpha is None
1189 + assert experts.gemm1_beta is None
1190 + assert experts.gemm1_clamp_limit is None
1191 +
1192 +
1193 @pytest.mark.parametrize("topk", [1, 4])
1194 @pytest.mark.parametrize("num_experts", [32])
1195 @pytest.mark.parametrize("num_tokens", [1, 128])
tests/kernels/test_compressor_kv_cache.py
+13
@@ -31,10 +31,23 @@ from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import (
31 )
32 from vllm.models.deepseek_v4.compressor import _get_c128_boundary
33 from vllm.platforms import current_platform
34 +from vllm.v1.attention.backends.mla.compressor_utils import (
35 + get_dspark_swa_index_width,
36 +)
37
38 from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4
39
40
41 +@pytest.mark.parametrize(
42 + ("window_size", "num_speculative_tokens", "expected"),
43 + [(128, 5, 192), (512, 5, 576), (1024, 0, 1024)],
44 +)
45 +def test_get_dspark_swa_index_width(
46 + window_size: int, num_speculative_tokens: int, expected: int
47 +):
48 + assert get_dspark_swa_index_width(window_size, num_speculative_tokens) == expected
49 +
50 +
51 def test_compute_global_topk_reuses_output_buffers():
52 device = "cuda"
53 topk_indices = torch.tensor(
tests/v1/attention/test_flashinfer_sparse_mla_sm120_api.py
+33
@@ -7,6 +7,9 @@ from types import SimpleNamespace
7 import torch
8
9 from vllm.config import set_current_vllm_config
10 +from vllm.models.deepseek_v4.nvidia.flashinfer_sparse import (
11 + _required_sm120_sparse_topk,
12 +)
13 from vllm.platforms.interface import DeviceCapability
14 from vllm.utils import flashinfer as fi_utils
15 from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
@@ -52,3 +55,33 @@ def test_v32_glm_sm120_backend_accepts_glm_block_size(
55 )
56
57 assert invalid_reasons == []
58 +
59 +
60 +def test_sm120_dsv4_capability_checks_exact_dispatch_shape(monkeypatch) -> None:
61 + fake_module = SimpleNamespace(
62 + _DECODE_DSV4_DISPATCH=frozenset({(32, 128), (32, 192)})
63 + )
64 + monkeypatch.setattr(fi_utils, "has_flashinfer_sparse_mla_sm120", lambda: True)
65 + monkeypatch.setattr(fi_utils, "_get_submodule", lambda _name: fake_module)
66 + fi_utils.has_flashinfer_sparse_mla_sm120_config.cache_clear()
67 +
68 + assert fi_utils.has_flashinfer_sparse_mla_sm120_config(32, 128)
69 + assert fi_utils.has_flashinfer_sparse_mla_sm120_config(32, 192)
70 + assert not fi_utils.has_flashinfer_sparse_mla_sm120_config(32, 256)
71 + assert not fi_utils.has_flashinfer_sparse_mla_sm120_config(16, 192)
72 +
73 + fi_utils.has_flashinfer_sparse_mla_sm120_config.cache_clear()
74 +
75 +
76 +def test_sm120_dsv4_required_topk_tracks_dspark_width() -> None:
77 + causal = SimpleNamespace(
78 + attention_config=SimpleNamespace(use_non_causal=False),
79 + speculative_config=SimpleNamespace(num_speculative_tokens=5),
80 + )
81 + dspark = SimpleNamespace(
82 + attention_config=SimpleNamespace(use_non_causal=True),
83 + speculative_config=SimpleNamespace(num_speculative_tokens=5),
84 + )
85 +
86 + assert _required_sm120_sparse_topk(causal, 128) == 128
87 + assert _required_sm120_sparse_topk(dspark, 128) == 192
tests/v1/spec_decode/test_dflash_prepare_inputs.py new
+152
@@ -0,0 +1,152 @@
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 +import pytest
8 +import torch
9 +
10 +from vllm.v1.attention.backends.utils import PAD_SLOT_ID
11 +from vllm.v1.worker.gpu.spec_decode.dflash.speculator import (
12 + prepare_dflash_inputs,
13 +)
14 +
15 +pytestmark = pytest.mark.skipif(
16 + not torch.cuda.is_available(), reason="requires a CUDA device"
17 +)
18 +
19 +
20 +def _run_prepare(*, target_positions: list[int], block_table_values: list[int]):
21 + device = torch.device("cuda")
22 + max_num_reqs = 4
23 + max_num_tokens = 16
24 + num_speculative_steps = 3
25 +
26 + input_buffers = SimpleNamespace(
27 + input_ids=torch.full((max_num_tokens,), -1, dtype=torch.int32, device=device),
28 + positions=torch.full((max_num_tokens,), -1, dtype=torch.int64, device=device),
29 + query_start_loc=torch.full(
30 + (max_num_reqs + 1,), -1, dtype=torch.int32, device=device
31 + ),
32 + seq_lens=torch.full((max_num_reqs,), -1, dtype=torch.int32, device=device),
33 + )
34 + input_batch = SimpleNamespace(
35 + num_reqs=1,
36 + num_scheduled_tokens=np.array([4], dtype=np.int32),
37 + positions=torch.tensor(target_positions, dtype=torch.int64, device=device),
38 + query_start_loc=torch.tensor([0, 4], dtype=torch.int32, device=device),
39 + idx_mapping=torch.tensor([2], dtype=torch.int32, device=device),
40 + )
41 + query_slot_mapping = torch.full(
42 + (max_num_tokens,), -2, dtype=torch.int64, device=device
43 + )
44 + context_positions = torch.full(
45 + (max_num_tokens,), -1, dtype=torch.int64, device=device
46 + )
47 + context_slot_mapping = torch.full(
48 + (max_num_tokens,), -2, dtype=torch.int64, device=device
49 + )
50 + sample_indices = torch.full(
51 + (max_num_reqs * num_speculative_steps,),
52 + -1,
53 + dtype=torch.int64,
54 + device=device,
55 + )
56 + sample_pos = torch.full_like(sample_indices, -1)
57 + sample_idx_mapping = torch.full(
58 + sample_indices.shape, -1, dtype=torch.int32, device=device
59 + )
60 + temperature = torch.zeros(max_num_reqs, dtype=torch.float32, device=device)
61 + seeds = torch.zeros(max_num_reqs, dtype=torch.int64, device=device)
62 + input_temperature = torch.tensor(
63 + [0.0, 0.0, 1.0, 0.0], dtype=torch.float32, device=device
64 + )
65 + input_seeds = torch.tensor([0, 0, 17, 0], dtype=torch.int64, device=device)
66 + last_sampled = torch.tensor([0, 0, 99, 0], dtype=torch.int64, device=device)
67 + next_prefill_tokens = torch.zeros_like(last_sampled)
68 + block_table = torch.tensor([block_table_values], dtype=torch.int32, device=device)
69 +
70 + prepare_dflash_inputs(
71 + input_buffers,
72 + query_slot_mapping,
73 + context_positions,
74 + context_slot_mapping,
75 + sample_indices,
76 + sample_pos,
77 + sample_idx_mapping,
78 + temperature,
79 + seeds,
80 + input_batch,
81 + torch.tensor([1], dtype=torch.int32, device=device),
82 + torch.tensor([2], dtype=torch.int32, device=device),
83 + last_sampled,
84 + next_prefill_tokens,
85 + input_temperature,
86 + input_seeds,
87 + block_table,
88 + 4,
89 + 123,
90 + num_speculative_steps,
91 + num_speculative_steps,
92 + max_num_reqs,
93 + max_num_tokens,
94 + 128,
95 + sample_from_anchor=True,
96 + )
97 + torch.accelerator.synchronize()
98 + return SimpleNamespace(
99 + input_buffers=input_buffers,
100 + query_slot_mapping=query_slot_mapping.cpu(),
101 + context_positions=context_positions.cpu(),
102 + context_slot_mapping=context_slot_mapping.cpu(),
103 + sample_indices=sample_indices.cpu(),
104 + sample_pos=sample_pos.cpu(),
105 + sample_idx_mapping=sample_idx_mapping.cpu(),
106 + temperature=temperature.cpu(),
107 + seeds=seeds.cpu(),
108 + )
109 +
110 +
111 +def test_prepare_dflash_inputs_excludes_rejected_context_suffix():
112 + # Positions 10/11 use physical block 7. Rejected positions 12/13 would use
113 + # block 8, but must be PAD context rather than contaminating draft KV.
114 + out = _run_prepare(
115 + target_positions=[10, 11, 12, 13],
116 + block_table_values=[0, 0, 7, 8, 9, 10, 11, 12],
117 + )
118 +
119 + assert out.context_positions[:4].tolist() == [10, 11, 0, 0]
120 + assert out.context_slot_mapping[:4].tolist() == [30, 31, PAD_SLOT_ID, PAD_SLOT_ID]
121 +
122 + # The replacement query starts immediately after the two valid rows and
123 + # advances from the last accepted position (11).
124 + assert out.input_buffers.input_ids[:3].cpu().tolist() == [99, 123, 123]
125 + assert out.input_buffers.positions[:3].cpu().tolist() == [12, 13, 14]
126 + assert out.query_slot_mapping[:3].tolist() == [32, 33, 34]
127 + assert out.sample_indices[:3].tolist() == [0, 1, 2]
128 + assert out.sample_pos[:3].tolist() == [13, 14, 15]
129 + assert out.sample_idx_mapping[:3].tolist() == [2, 2, 2]
130 + assert out.temperature[2].item() == 1.0
131 + assert out.seeds[2].item() == 17
132 +
133 +
134 +def test_prepare_dflash_inputs_never_writes_the_null_block():
135 + # The valid context uses logical block 0 and the replacement query uses
136 + # logical block 1. Both map to the null block and must remain unwritable.
137 + out = _run_prepare(
138 + target_positions=[2, 3, 4, 5],
139 + block_table_values=[0, 0, 7, 8, 9, 10, 11, 12],
140 + )
141 +
142 + assert out.context_slot_mapping[:4].tolist() == [
143 + PAD_SLOT_ID,
144 + PAD_SLOT_ID,
145 + PAD_SLOT_ID,
146 + PAD_SLOT_ID,
147 + ]
148 + assert out.query_slot_mapping[:3].tolist() == [
149 + PAD_SLOT_ID,
150 + PAD_SLOT_ID,
151 + PAD_SLOT_ID,
152 + ]
tests/v1/worker/test_workspace.py new
+102
@@ -0,0 +1,102 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from typing import cast
5 +
6 +import pytest
7 +import torch
8 +
9 +import vllm.v1.worker.workspace as workspace
10 +from vllm.config import VllmConfig
11 +from vllm.v1.worker.gpu_worker import _num_workspace_lanes
12 +
13 +
14 +class _SpecConfig:
15 + def __init__(self, dspark: bool) -> None:
16 + self._dspark = dspark
17 +
18 + def use_dspark(self) -> bool:
19 + return self._dspark
20 +
21 +
22 +class _VllmConfig:
23 + def __init__(self, spec_config: _SpecConfig | None) -> None:
24 + self.speculative_config = spec_config
25 +
26 +
27 +@pytest.mark.parametrize(
28 + ("use_v2_model_runner", "spec_config", "expected"),
29 + [
30 + (True, _SpecConfig(True), 2),
31 + (False, _SpecConfig(True), 1),
32 + (True, _SpecConfig(False), 1),
33 + (True, None, 1),
34 + ],
35 +)
36 +def test_workspace_lane_count_is_dspark_only(
37 + use_v2_model_runner: bool,
38 + spec_config: _SpecConfig | None,
39 + expected: int,
40 +) -> None:
41 + config = cast(VllmConfig, _VllmConfig(spec_config))
42 + assert _num_workspace_lanes(config, use_v2_model_runner) == expected
43 +
44 +
45 +def test_workspace_lanes_do_not_alias_and_restore_context(monkeypatch) -> None:
46 + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0)
47 + manager = workspace.WorkspaceManager(
48 + torch.device("cpu"), num_ubatches=2, num_lanes=2
49 + )
50 +
51 + assert manager._current_workspaces == [None, None, None, None]
52 +
53 + (target,) = manager.get_simultaneous(((512,), torch.uint8))
54 + with workspace.use_workspace_lane(1):
55 + (draft,) = manager.get_simultaneous(((256,), torch.uint8))
56 + (draft_reused,) = manager.get_simultaneous(((8,), torch.uint8))
57 + (target_reused,) = manager.get_simultaneous(((8,), torch.uint8))
58 +
59 + assert manager._current_workspaces[0].numel() == 512 # type: ignore[union-attr]
60 + assert manager._current_workspaces[1].numel() == 256 # type: ignore[union-attr]
61 + assert manager._current_workspaces[2:] == [None, None]
62 + assert target.data_ptr() != draft.data_ptr()
63 + assert draft.data_ptr() == draft_reused.data_ptr()
64 + assert target.data_ptr() == target_reused.data_ptr()
65 +
66 +
67 +def test_workspace_lanes_compose_with_ubatches(monkeypatch) -> None:
68 + active_ubatch = [0]
69 + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: active_ubatch[0])
70 + manager = workspace.WorkspaceManager(
71 + torch.device("cpu"), num_ubatches=2, num_lanes=2
72 + )
73 +
74 + pointers = set()
75 + for ubatch_id in range(2):
76 + active_ubatch[0] = ubatch_id
77 + for lane in range(2):
78 + with workspace.use_workspace_lane(lane):
79 + (buffer,) = manager.get_simultaneous(((16,), torch.uint8))
80 + pointers.add(buffer.data_ptr())
81 +
82 + assert len(pointers) == 4
83 +
84 +
85 +def test_workspace_lane_validation(monkeypatch) -> None:
86 + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0)
87 + manager = workspace.WorkspaceManager(torch.device("cpu"), num_lanes=1)
88 +
89 + with (
90 + pytest.raises(ValueError, match="non-negative"),
91 + workspace.use_workspace_lane(-1),
92 + ):
93 + pass
94 +
95 + with (
96 + workspace.use_workspace_lane(1),
97 + pytest.raises(RuntimeError, match="is not configured"),
98 + ):
99 + manager.get_simultaneous(((1,), torch.uint8))
100 +
101 + with pytest.raises(ValueError, match="at least one"):
102 + workspace.WorkspaceManager(torch.device("cpu"), num_lanes=0)
vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py
+9 -18
@@ -111,21 +111,15 @@ class FlashInferExperts(mk.FusedMoEExpertsModular):
111 self.gemm1_alpha = _per_expert(quant_config.gemm1_alpha)
112 self.gemm1_beta = _per_expert(quant_config.gemm1_beta)
113
114 - if quant_config.weight_quant_dtype == "mxfp4":
115 - # This value is used specifically for gpt-oss,
116 - # Need to revisit this for other models
117 - if self.gemm1_alpha is None:
118 - self.gemm1_alpha = _per_expert(1.702)
119 - if self.gemm1_beta is None:
120 - self.gemm1_beta = _per_expert(1.0)
121 - if self.gemm1_clamp_limit is None:
122 - self.gemm1_clamp_limit = _per_expert(7.0)
123 - if quant_config.quant_dtype == "mxfp8":
124 - self.fake_input_scale = torch.ones(
125 - self.num_experts,
126 - device=self.device,
127 - dtype=torch.float32,
128 - )
114 + if (
115 + quant_config.weight_quant_dtype == "mxfp4"
116 + and quant_config.quant_dtype == "mxfp8"
117 + ):
118 + self.fake_input_scale = torch.ones(
119 + self.num_experts,
120 + device=self.device,
121 + dtype=torch.float32,
122 + )
123
124 @property
125 def expects_unquantized_inputs(self) -> bool:
@@ -324,9 +318,6 @@ class FlashInferExperts(mk.FusedMoEExpertsModular):
318 elif self.weight_quant_dtype == "mxfp4":
319 assert self.w1_scale is not None and self.w2_scale is not None
320 assert w1.is_contiguous() and w2.is_contiguous()
327 - assert self.gemm1_alpha is not None
328 - assert self.gemm1_beta is not None
329 - assert self.gemm1_clamp_limit is not None
321 assert topk_ids.is_contiguous()
322
323 fc1_expert_biases = self.w1_bias
vllm/models/deepseek_v4/amd/rocm.py
+4 -4
@@ -414,7 +414,9 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuild
414 and base.decode_swa_lens is not None
415 ):
416 ragged_indices, ragged_indptr = build_ragged_indices_from_dense(
417 - base.decode_swa_indices.reshape(base.num_decode_tokens, -1),
417 + base.decode_swa_indices.reshape(
418 + base.num_decode_tokens, base.decode_swa_width
419 + ),
420 base.decode_swa_lens,
421 )
422 ragged_indices, ragged_indptr = _copy_ragged_to_graph_buffers(
@@ -423,9 +425,7 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuild
425 self.decode_swa_ragged_indices_buffer,
426 self.decode_swa_ragged_indptr_buffer,
427 base.num_decode_tokens,
426 - # Actual dense width for this build: window_size (causal) or
427 - # noncausal_index_width (DSpark non-causal draft).
428 - base.decode_swa_indices.shape[-1],
428 + base.decode_swa_width,
429 )
430
431 return DeepseekV4ROCMAiterSparseSWAMetadata(
vllm/models/deepseek_v4/common/ops/cache_utils.py
+20 -17
@@ -839,16 +839,17 @@ def build_flashinfer_mixed_sparse_indices(
839 ) -> tuple[torch.Tensor, torch.Tensor]:
840 """Build the FlashInfer DSV4 sparse-index matrix for decode-first batches.
841
842 - Produces ``sparse_indices`` of shape ``[num_tokens, window_size +
843 - padded_topk]`` (the first ``window_size`` columns are SWA slot ids, the rest
844 - are compressed/top-k slot ids) and ``sparse_topk_lens`` (active length per
845 - token). Decode tokens read precomputed SWA/compressed indices; prefill tokens
846 - derive their SWA window from the position and translate local compressed
847 - indices to global slots via the block tables.
842 + Produces ``sparse_indices`` of shape ``[num_tokens, swa_index_width +
843 + padded_topk]`` (the first ``swa_index_width`` columns are SWA slot ids, the
844 + rest are compressed/top-k slot ids) and ``sparse_topk_lens`` (active length
845 + per token). Decode tokens read precomputed SWA/compressed indices; prefill
846 + tokens derive their SWA window from the position and translate local
847 + compressed indices to global slots via the block tables.
848 """
849 assert decode_swa_indices.dtype == torch.int32
850 assert decode_swa_indices.dim() == 2
851 - assert decode_swa_indices.shape[-1] == window_size
851 + swa_index_width = decode_swa_indices.shape[-1]
852 + assert swa_index_width >= window_size
853 if decode_compressed_topk_lens is not None:
854 assert decode_compressed_topk_lens.dtype == torch.int32
855 assert prefill_topk_indices.dtype == torch.int32
@@ -897,7 +898,7 @@ def build_flashinfer_mixed_sparse_indices(
898 padded_topk = max(topk, decode_compressed_topk)
899 padded_topk = (padded_topk + 3) // 4 * 4
900 sparse_indices = torch.empty(
900 - (num_tokens, window_size + padded_topk),
901 + (num_tokens, swa_index_width + padded_topk),
902 dtype=torch.int32,
903 device=decode_swa_indices.device,
904 )
@@ -907,7 +908,7 @@ def build_flashinfer_mixed_sparse_indices(
908 if num_tokens == 0:
909 return sparse_indices, sparse_topk_lens
910
910 - window_block_size = triton.next_power_of_2(max(window_size, 1))
911 + window_block_size = triton.next_power_of_2(max(swa_index_width, 1))
912 topk_block_size = triton.next_power_of_2(max(padded_topk, 1))
913 max_block_size = max(window_block_size, topk_block_size)
914 num_warps = 4 if max_block_size >= 256 else 1
@@ -944,6 +945,7 @@ def build_flashinfer_mixed_sparse_indices(
945 compressed_span,
946 NUM_DECODE_TOKENS=num_decode_tokens,
947 WINDOW_SIZE=window_size,
948 + SWA_INDEX_WIDTH=swa_index_width,
949 COMPRESS_RATIO=compress_ratio,
950 TOP_K=topk,
951 PADDED_TOP_K=padded_topk,
@@ -1011,6 +1013,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel(
1013 compressed_block_span,
1014 NUM_DECODE_TOKENS,
1015 WINDOW_SIZE: tl.constexpr,
1016 + SWA_INDEX_WIDTH: tl.constexpr,
1017 COMPRESS_RATIO: tl.constexpr,
1018 TOP_K: tl.constexpr,
1019 PADDED_TOP_K: tl.constexpr,
@@ -1024,9 +1027,9 @@ def _build_flashinfer_mixed_sparse_indices_kernel(
1027 token_idx = tl.program_id(0)
1028
1029 if token_idx < NUM_DECODE_TOKENS:
1027 - for i in range(0, WINDOW_SIZE, WINDOW_BLOCK_SIZE):
1030 + for i in range(0, SWA_INDEX_WIDTH, WINDOW_BLOCK_SIZE):
1031 offset = i + tl.arange(0, WINDOW_BLOCK_SIZE)
1029 - mask = offset < WINDOW_SIZE
1032 + mask = offset < SWA_INDEX_WIDTH
1033 values = tl.load(
1034 decode_swa_indices_ptr + token_idx * decode_swa_stride + offset,
1035 mask=mask,
@@ -1072,7 +1075,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel(
1075 tl.store(
1076 sparse_indices_ptr
1077 + token_idx * sparse_indices_stride
1075 - + WINDOW_SIZE
1078 + + SWA_INDEX_WIDTH
1079 + offset,
1080 values,
1081 mask=mask,
@@ -1086,7 +1089,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel(
1089 else:
1090 compressed_len = tl.full((), DECODE_COMPRESSED_TOPK, dtype=tl.int32)
1091
1089 - tl.store(sparse_topk_lens_ptr + token_idx, WINDOW_SIZE + compressed_len)
1092 + tl.store(sparse_topk_lens_ptr + token_idx, SWA_INDEX_WIDTH + compressed_len)
1093 return
1094
1095 prefill_idx = token_idx - NUM_DECODE_TOKENS
@@ -1102,9 +1105,9 @@ def _build_flashinfer_mixed_sparse_indices_kernel(
1105 swa_start_pos = pos - swa_len + 1
1106 topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
1107
1105 - for i in range(0, WINDOW_SIZE, WINDOW_BLOCK_SIZE):
1108 + for i in range(0, SWA_INDEX_WIDTH, WINDOW_BLOCK_SIZE):
1109 offset = i + tl.arange(0, WINDOW_BLOCK_SIZE)
1107 - mask = offset < WINDOW_SIZE
1110 + mask = offset < SWA_INDEX_WIDTH
1111 pos_offset = swa_start_pos + offset
1112 block_indices = pos_offset // swa_block_size
1113 block_numbers = tl.load(
@@ -1148,10 +1151,10 @@ def _build_flashinfer_mixed_sparse_indices_kernel(
1151 tl.store(
1152 sparse_indices_ptr
1153 + token_idx * sparse_indices_stride
1151 - + WINDOW_SIZE
1154 + + SWA_INDEX_WIDTH
1155 + offset,
1156 slot_ids,
1157 mask=mask,
1158 )
1159
1157 - tl.store(sparse_topk_lens_ptr + token_idx, WINDOW_SIZE + topk_len)
1160 + tl.store(sparse_topk_lens_ptr + token_idx, SWA_INDEX_WIDTH + topk_len)
vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py
+28 -7
@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, ClassVar, cast
6
7 import torch
8
9 +from vllm.config import VllmConfig
10 from vllm.config.cache import CacheDType
11 from vllm.forward_context import get_forward_context
12 from vllm.models.deepseek_v4.attention import DeepseekV4Attention
@@ -25,6 +26,9 @@ from vllm.platforms import current_platform
26 from vllm.platforms.interface import DeviceCapability
27 from vllm.utils.flashinfer import flashinfer_trtllm_batch_decode_sparse_mla_dsv4
28 from vllm.v1.attention.backend import MultipleOf
29 +from vllm.v1.attention.backends.mla.compressor_utils import (
30 + get_dspark_swa_index_width,
31 +)
32
33 if TYPE_CHECKING:
34 from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
@@ -74,6 +78,19 @@ def _pad_to_supported_q_heads(num_heads: int) -> int:
78 )
79
80
81 +def _required_sm120_sparse_topk(vllm_config: VllmConfig, window_size: int) -> int:
82 + """Return the SM120 DSV4 SWA specialization needed by this model."""
83 + if not vllm_config.attention_config.use_non_causal:
84 + return window_size
85 + speculative_config = vllm_config.speculative_config
86 + if speculative_config is None:
87 + return window_size
88 + return get_dspark_swa_index_width(
89 + window_size,
90 + speculative_config.num_speculative_tokens,
91 + )
92 +
93 +
94 class DeepseekV4FlashInferMLASparseBackend(DeepseekV4SparseMLABackend):
95 """FlashInfer backend using the DSv4 sparse metadata/cache layout.
96
@@ -309,7 +326,7 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
326 assert swa_metadata.block_table is not None
327
328 decode_swa_indices = swa_metadata.decode_swa_indices.reshape(
312 - num_decode_tokens, self.window_size
329 + num_decode_tokens, swa_metadata.decode_swa_width
330 )
331 decode_compressed_topk_lens = None
332 decode_compressed_indices_are_local = False
@@ -568,14 +585,18 @@ class DeepseekV4FlashInferSM120Attention(DeepseekV4Attention):
585 tma_aligned_scales=self._tma_aligned_scales,
586 )
587
571 - def __init__(self, *args, **kwargs) -> None:
572 - super().__init__(*args, **kwargs)
573 - from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120
588 + def __init__(self, vllm_config: VllmConfig, *args, **kwargs) -> None:
589 + super().__init__(vllm_config, *args, **kwargs)
590 + from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120_config
591
575 - if not has_flashinfer_sparse_mla_sm120():
592 + required_topk = _required_sm120_sparse_topk(vllm_config, self.window_size)
593 + if not has_flashinfer_sparse_mla_sm120_config(self.padded_heads, required_topk):
594 raise RuntimeError(
577 - "FLASHINFER_MLA_SPARSE_DSV4 on SM120 requires FlashInfer's "
578 - "sparse MLA decode API."
595 + "FLASHINFER_MLA_SPARSE_DSV4 on SM120 requires a FlashInfer "
596 + "DSV4 sparse MLA decode specialization for "
597 + f"(num_q_heads={self.padded_heads}, top_k={required_topk}). "
598 + "Install a FlashInfer build containing "
599 + "flashinfer-ai/flashinfer#4380."
600 )
601 self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe()
602 # Per-tensor FP8 cache path scales.
vllm/utils/flashinfer.py
+16
@@ -237,6 +237,22 @@ def has_flashinfer_sparse_mla_sm120() -> bool:
237 )
238
239
240 +@functools.cache
241 +def has_flashinfer_sparse_mla_sm120_config(num_q_heads: int, top_k: int) -> bool:
242 + """Return whether FlashInfer ships an SM120 DSV4 decode specialization.
243 +
244 + The public sparse MLA API predates some DSV4 shapes, so checking only that
245 + the callable exists can select a package that later aborts or rejects a
246 + valid vLLM configuration. Inspect FlashInfer's dispatch table until it
247 + exposes a public capability query.
248 + """
249 + if not has_flashinfer_sparse_mla_sm120():
250 + return False
251 + mod = _get_submodule("flashinfer.mla._sparse_mla_sm120")
252 + dispatch = getattr(mod, "_DECODE_DSV4_DISPATCH", None) if mod else None
253 + return dispatch is not None and (int(num_q_heads), int(top_k)) in dispatch
254 +
255 +
256 @functools.cache
257 def has_flashinfer_cutedsl() -> bool:
258 """Return ``True`` if FlashInfer cutedsl module is available."""
vllm/v1/attention/backends/mla/compressor_utils.py
+12
@@ -3,6 +3,18 @@
3 import torch
4
5 from vllm.triton_utils import tl, triton
6 +from vllm.utils.math_utils import cdiv
7 +
8 +_DSPARK_SWA_INDEX_ALIGNMENT = 64
9 +
10 +
11 +def get_dspark_swa_index_width(
12 + window_size: int,
13 + num_speculative_tokens: int,
14 +) -> int:
15 + """Return the padded width of non-causal DSpark SWA indices."""
16 + width = max(int(window_size), 0) + max(int(num_speculative_tokens), 0)
17 + return cdiv(width, _DSPARK_SWA_INDEX_ALIGNMENT) * _DSPARK_SWA_INDEX_ALIGNMENT
18
19
20 @triton.jit
vllm/v1/attention/backends/mla/indexer.py
+9 -3
@@ -57,9 +57,12 @@ def _prepare_uniform_decode_kernel(
57 req_id = idx // max_decode_len
58 local_idx = idx % max_decode_len
59
60 - # Compute number of KVs attended to by this token.
60 + # Compute number of KVs attended to by this token. Padding requests have
61 + # seq_len == 0, which would otherwise make the first token of each padded
62 + # request negative (e.g. next_n=2 gives 0-2+0+1 = -1). Downstream kernels
63 + # read these as uint32, turning -1 into ~4e9.
64 seq_len = tl.load(seq_lens_ptr + req_id)
62 - per_token_seq_len = seq_len - max_decode_len + local_idx + 1
65 + per_token_seq_len = tl.maximum(seq_len - max_decode_len + local_idx + 1, 0)
66 tl.store(decode_seq_lens_ptr + idx, per_token_seq_len)
67
68 # Copy block table row.
@@ -748,12 +751,15 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
751 seq_lens_buffer = self.decode_seq_lens_buffer[
752 : num_decodes * max_decode_len
753 ].view(num_decodes, max_decode_len)
754 + # Clamp at 0: padding requests have seq_len == 0, which would
755 + # otherwise make token 0 negative (next_n=2 gives 0-2+1+0 = -1).
756 + # Downstream kernels read these as uint32, turning -1 into ~4e9.
757 seq_lens_buffer[:] = (
758 seq_lens.unsqueeze(1)
759 - max_decode_len
760 + 1
761 + self.offsets_buffer[:max_decode_len]
756 - )
762 + ).clamp_(min=0)
763 seq_lens = seq_lens_buffer
764 return seq_lens, block_table, decode_lens, num_decodes, requires_padding
765
vllm/v1/attention/backends/mla/sparse_swa.py
+16 -4
@@ -22,6 +22,9 @@ from vllm.v1.attention.backend import (
22 CommonAttentionMetadata,
23 MultipleOf,
24 )
25 +from vllm.v1.attention.backends.mla.compressor_utils import (
26 + get_dspark_swa_index_width,
27 +)
28 from vllm.v1.attention.backends.utils import split_decodes_and_prefills
29 from vllm.v1.attention.ops.flashmla import FlashMLASchedMeta, get_mla_metadata
30 from vllm.v1.kv_cache_interface import (
@@ -174,8 +177,10 @@ class DeepseekSparseSWAMetadata:
177
178 is_valid_token: torch.Tensor | None = None # [num_tokens]
179 token_to_req_indices: torch.Tensor | None = None # [num_tokens]
177 - decode_swa_indices: torch.Tensor | None = None # [num_decode_tokens, window_size]
180 + decode_swa_indices: torch.Tensor | None = None # [num_decode_tokens, width]
181 decode_swa_lens: torch.Tensor | None = None # [num_decode_tokens]
182 + # window_size (causal) or noncausal_index_width (DSpark non-causal).
183 + decode_swa_width: int = 0
184 # Paged-coordinate prefill SWA indices/lens (FP8 paged-direct prefill).
185 prefill_swa_indices: torch.Tensor | None = (
186 None # [num_prefill_tokens, 1, window_size]
@@ -487,11 +492,14 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
492 # DSpark draft: the block is non-causal (every query attends to the
493 # trailing window of context PLUS all query tokens, including future ones),
494 # so its per-token index list is wider than `window_size`. The kernel pads
490 - # the q-head count to B_TOPK (64/128), which requires the index width to be
491 - # a multiple of 128.
495 + # the q-head count to B_TOPK. Pad to a kernel-supported width; the logical
496 + # SWA window remains unchanged when the padded matrix is built.
497 self.is_dspark = spec_config is not None and spec_config.use_dspark()
498 self.noncausal_index_width = (
494 - cdiv(self.window_size + self.num_speculative_tokens, 128) * 128
499 + get_dspark_swa_index_width(
500 + self.window_size,
501 + self.num_speculative_tokens,
502 + )
503 if self.is_dspark
504 else 0
505 )
@@ -536,6 +544,9 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
544 is_valid_token.copy_(slot_mapping >= 0)
545
546 non_causal = not common_attn_metadata.causal
547 + decode_swa_width = (
548 + self.noncausal_index_width if non_causal else self.window_size
549 + )
550 decode_swa_indices = self.decode_swa_indices
551 if num_decode_tokens > 0:
552 self.decode_swa_lens[num_decode_tokens:] = 0
@@ -634,6 +645,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
645 token_to_req_indices=token_to_req_indices,
646 decode_swa_indices=decode_swa_indices[:num_decode_tokens],
647 decode_swa_lens=self.decode_swa_lens[:num_decode_tokens],
648 + decode_swa_width=decode_swa_width,
649 prefill_swa_indices=(
650 self.prefill_swa_indices[:num_prefill_tokens]
651 if num_prefill_tokens > 0
vllm/v1/worker/gpu/model_runner.py
+48 -40
@@ -149,6 +149,7 @@ from vllm.v1.worker.utils import (
149 copy_kv_cache_blocks_inplace,
150 get_uniform_decode_token_count,
151 )
152 +from vllm.v1.worker.workspace import use_workspace_lane
153
154 logger = init_logger(__name__)
155
@@ -164,6 +165,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
165 self.parallel_config = vllm_config.parallel_config
166 self.scheduler_config = vllm_config.scheduler_config
167 self.speculative_config = vllm_config.speculative_config
168 + self._draft_workspace_lane = int(
169 + self.speculative_config is not None and self.speculative_config.use_dspark()
170 + )
171 self.observability_config = vllm_config.observability_config
172
173 self.device = device
@@ -364,10 +368,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
368 assert self.speculative_config is not None
369 set_eagle3_aux_hidden_state_layers(self.model, self.speculative_config)
370 if isinstance(self.speculator, DraftModelSpeculator):
367 - self.speculator.load_model(self.model)
368 - eplb_models_added = self.eplb.maybe_register_speculator(
369 - self.speculator, self.speculative_config, load_dummy_weights
370 - )
371 + with use_workspace_lane(self._draft_workspace_lane):
372 + self.speculator.load_model(self.model)
373 + eplb_models_added = self.eplb.maybe_register_speculator(
374 + self.speculator, self.speculative_config, load_dummy_weights
375 + )
376 time_after_load = time.perf_counter()
377
378 self.model_memory_usage = m.consumed_memory
@@ -732,27 +737,28 @@ class GPUModelRunner(LoRAModelRunnerMixin):
737 if hasattr(self.model, "get_mtp_target_hidden_states"):
738 pre_hc_hidden_states = self.model.get_mtp_target_hidden_states()
739 spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr]
735 - self.speculator.propose(
736 - input_batch=input_batch,
737 - attn_metadata=attn_metadata,
738 - slot_mappings=slot_mappings_by_layer,
739 - last_hidden_states=spec_hidden_states,
740 - aux_hidden_states=aux_hidden_states,
741 - num_sampled=torch.ones(
742 - input_batch.num_reqs, dtype=torch.int32, device=self.device
743 - ),
744 - num_rejected=torch.zeros(
745 - input_batch.num_reqs, dtype=torch.int32, device=self.device
746 - ),
747 - last_sampled=self.req_states.last_sampled_tokens,
748 - next_prefill_tokens=self.req_states.next_prefill_tokens,
749 - temperature=self.sampler.sampling_states.temperature.gpu,
750 - seeds=self.sampler.sampling_states.seeds.gpu,
751 - dummy_run=True,
752 - skip_attn_for_dummy_run=skip_attn,
753 - mm_inputs=mm_inputs,
754 - is_profile=is_profile,
755 - )
740 + with use_workspace_lane(self._draft_workspace_lane):
741 + self.speculator.propose(
742 + input_batch=input_batch,
743 + attn_metadata=attn_metadata,
744 + slot_mappings=slot_mappings_by_layer,
745 + last_hidden_states=spec_hidden_states,
746 + aux_hidden_states=aux_hidden_states,
747 + num_sampled=torch.ones(
748 + input_batch.num_reqs, dtype=torch.int32, device=self.device
749 + ),
750 + num_rejected=torch.zeros(
751 + input_batch.num_reqs, dtype=torch.int32, device=self.device
752 + ),
753 + last_sampled=self.req_states.last_sampled_tokens,
754 + next_prefill_tokens=self.req_states.next_prefill_tokens,
755 + temperature=self.sampler.sampling_states.temperature.gpu,
756 + seeds=self.sampler.sampling_states.seeds.gpu,
757 + dummy_run=True,
758 + skip_attn_for_dummy_run=skip_attn,
759 + mm_inputs=mm_inputs,
760 + is_profile=is_profile,
761 + )
762 self.step_timing.drafter_end()
763
764 assert hidden_states is not None # Last PP rank always has hidden_states
@@ -880,7 +886,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
886 lora_capture_hook=create_lora_capture_hook(self.lora_config, self),
887 )
888 if self.speculator is not None:
883 - self.speculator.capture()
889 + with use_workspace_lane(self._draft_workspace_lane):
890 + self.speculator.capture()
891 if self.adaptive_verification is not None:
892 with self.step_timing.collect() as timings:
893 for batch in self.adaptive_verification.batches_to_profile(
@@ -1799,20 +1806,21 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1806 if hasattr(self.model, "get_mtp_target_hidden_states"):
1807 pre_hc_hidden_states = self.model.get_mtp_target_hidden_states()
1808 spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr]
1802 - draft_tokens = self.speculator.propose(
1803 - input_batch,
1804 - attn_metadata,
1805 - slot_mappings_by_layer,
1806 - spec_hidden_states,
1807 - aux_hidden_states,
1808 - num_sampled,
1809 - num_rejected,
1810 - self.req_states.last_sampled_tokens,
1811 - self.req_states.next_prefill_tokens,
1812 - self.sampler.sampling_states.temperature.gpu,
1813 - self.sampler.sampling_states.seeds.gpu,
1814 - mm_inputs=mm_inputs,
1815 - )
1809 + with use_workspace_lane(self._draft_workspace_lane):
1810 + draft_tokens = self.speculator.propose(
1811 + input_batch,
1812 + attn_metadata,
1813 + slot_mappings_by_layer,
1814 + spec_hidden_states,
1815 + aux_hidden_states,
1816 + num_sampled,
1817 + num_rejected,
1818 + self.req_states.last_sampled_tokens,
1819 + self.req_states.next_prefill_tokens,
1820 + self.sampler.sampling_states.temperature.gpu,
1821 + self.sampler.sampling_states.seeds.gpu,
1822 + mm_inputs=mm_inputs,
1823 + )
1824 self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
1825 if self.adaptive_verification is not None:
1826 self.adaptive_verification.record_confidences(
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
+33 -9
@@ -84,8 +84,11 @@ class DFlashSpeculator(DraftModelSpeculator):
84 self.sample_pos = torch.zeros(
85 max_num_sampled_tokens, dtype=torch.int64, device=device
86 )
87 - self.sample_idx_mapping = torch.zeros(
88 - max_num_sampled_tokens, dtype=torch.int32, device=device
87 + # -1 marks an inert sampling row. CUDA graph capture can execute the
88 + # full buffer before a real batch has populated it, so zero would make
89 + # every padding row scatter into request slot 0.
90 + self.sample_idx_mapping = torch.full(
91 + (max_num_sampled_tokens,), -1, dtype=torch.int32, device=device
92 )
93 # [0, 1, ..., N-1, 0, 1, ..., N-1, ...] -> the per-token column index into
94 # draft_logits[req, step, :].
@@ -255,7 +258,6 @@ class DFlashSpeculator(DraftModelSpeculator):
258 num_tokens_across_dp,
259 cudagraph_runtime_mode,
260 )
258 -
261 num_sample = num_reqs * self.num_speculative_steps
262 sample_hidden_states = last_hidden_states[self.sample_indices[:num_sample]]
263 # sample_pos is the predicted token's position Q; verification keys
@@ -520,6 +522,7 @@ def _prepare_dflash_inputs_kernel(
522
523 num_rejected = tl.load(num_rejected_ptr + req_idx)
524 valid_ctx_end = ctx_end - num_rejected
525 + num_valid_ctx = valid_ctx_end - ctx_start
526
527 num_sampled = tl.load(num_sampled_ptr + req_idx)
528 if num_sampled > 0:
@@ -533,20 +536,34 @@ def _prepare_dflash_inputs_kernel(
536
537 j = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
538 is_ctx = j < num_ctx
536 - is_query = (j >= num_ctx) & (j < num_ctx + num_query_per_req)
537 - query_off = j - num_ctx
539 + is_valid_ctx = j < num_valid_ctx
540 + is_query = (j >= num_valid_ctx) & (j < num_valid_ctx + num_query_per_req)
541 + query_off = j - num_valid_ctx
542
543 # --- Context positions / slots ---
544 ctx_pos_idx = ctx_start + tl.where(is_ctx, j, 0)
541 - ctx_pos = tl.load(target_positions_ptr + ctx_pos_idx, mask=is_ctx, other=0)
545 + ctx_pos = tl.load(target_positions_ptr + ctx_pos_idx, mask=is_valid_ctx, other=0)
546 ctx_block_num = ctx_pos // block_size
547 ctx_block_num = tl.minimum(ctx_block_num, block_table_stride - 1)
548 ctx_block_id = tl.load(
549 block_table_ptr + req_idx * block_table_stride + ctx_block_num,
546 - mask=is_ctx,
550 + mask=is_valid_ctx,
551 other=0,
552 ).to(tl.int64)
549 - ctx_slot = ctx_block_id * block_size + (ctx_pos % block_size)
553 + # Block 0 is the null block. Old sliding-window context positions can map
554 + # to it after eviction; rejected suffix rows are invalid context as well.
555 + # Neither kind of row may write draft KV into physical block 0.
556 + ctx_resident = is_valid_ctx & (ctx_block_id != 0)
557 + ctx_slot = tl.where(
558 + ctx_resident,
559 + ctx_block_id * block_size + (ctx_pos % block_size),
560 + PAD_SLOT_ID,
561 + )
562 + # Stored over the full [0, num_ctx) span while the loads above are masked to
563 + # [0, num_valid_ctx): the rejected suffix rows in between get position 0 and
564 + # PAD_SLOT_ID. That is intentional — those rows write no KV and their
565 + # positions are never consumed, but the span must stay fully initialized so
566 + # a replayed graph cannot observe a stale value from an earlier batch.
567 tl.store(out_context_positions_ptr + ctx_start + j, ctx_pos, mask=is_ctx)
568 tl.store(out_context_slot_mapping_ptr + ctx_start + j, ctx_slot, mask=is_ctx)
569
@@ -563,7 +580,14 @@ def _prepare_dflash_inputs_kernel(
580 mask=is_query,
581 other=0,
582 ).to(tl.int64)
566 - q_slot = q_block_id * block_size + (query_pos % block_size)
583 + # A null block is never a writable cache slot. This can occur when a
584 + # sliding-window block table contains evicted/global padding entries.
585 + q_resident = is_query & (q_block_id != 0)
586 + q_slot = tl.where(
587 + q_resident,
588 + q_block_id * block_size + (query_pos % block_size),
589 + PAD_SLOT_ID,
590 + )
591
592 tl.store(out_input_ids_ptr + query_idx, input_id, mask=is_query)
593 clamped_query_pos = tl.minimum(query_pos, max_model_len - 1)
vllm/v1/worker/gpu_worker.py
+16 -2
@@ -91,6 +91,16 @@ from .utils import request_memory
91
92 logger = init_logger(__name__)
93
94 +
95 +def _num_workspace_lanes(vllm_config: VllmConfig, use_v2_model_runner: bool) -> int:
96 + spec_config = vllm_config.speculative_config
97 + return (
98 + 2
99 + if use_v2_model_runner and spec_config is not None and spec_config.use_dspark()
100 + else 1
101 + )
102 +
103 +
104 if TYPE_CHECKING:
105 from vllm.device_allocator.sleep_mode_backend import SleepModeBackend
106 from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
@@ -402,9 +412,13 @@ class Worker(WorkerBase):
412 else:
413 raise RuntimeError(f"Unsupported device type: {self.device_config.device}")
414
405 - # Initialize workspace manager
415 + # DSpark target and draft CUDA graphs retain workspace views concurrently.
416 num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1
407 - init_workspace_manager(self.device, num_ubatches)
417 + init_workspace_manager(
418 + self.device,
419 + num_ubatches,
420 + _num_workspace_lanes(self.vllm_config, self.use_v2_model_runner),
421 + )
422
423 # Construct the model runner
424 if self.use_v2_model_runner:
vllm/v1/worker/workspace.py
+49 -14
@@ -3,6 +3,9 @@
3
4 import inspect
5 import os
6 +from collections.abc import Iterator
7 +from contextlib import contextmanager
8 +from contextvars import ContextVar
9 from itertools import accumulate
10 from math import prod
11
@@ -26,22 +29,43 @@ _GiB = 1024**3
29
30 # Global workspace manager instance
31 _manager: "WorkspaceManager | None" = None
32 +_workspace_lane: ContextVar[int] = ContextVar("vllm_workspace_lane", default=0)
33 +
34 +
35 +@contextmanager
36 +def use_workspace_lane(lane: int) -> Iterator[None]:
37 + """Select an independent workspace owner for this execution context."""
38 + if lane < 0:
39 + raise ValueError(f"Workspace lane must be non-negative, got {lane}.")
40 + token = _workspace_lane.set(lane)
41 + try:
42 + yield
43 + finally:
44 + _workspace_lane.reset(token)
45
46
47 class WorkspaceManager:
48 """Manager for workspace allocation.
49
34 - Manages one workspace buffer per active ubatch slot.
50 + Manages one workspace buffer per active ``(ubatch, lane)`` slot.
51 Can be locked to prevent further growth during execution.
52 """
53
38 - def __init__(self, device: torch.device, num_ubatches: int | None = None):
54 + def __init__(
55 + self,
56 + device: torch.device,
57 + num_ubatches: int | None = None,
58 + num_lanes: int = 1,
59 + ):
60 self._device = device
61 # Cache num ubatches at init based on configuration (default to 1)
62 self._num_ubatches = num_ubatches if num_ubatches is not None else 1
42 - self._current_workspaces: list[torch.Tensor | None] = [
43 - None
44 - ] * self._num_ubatches
63 + if num_lanes < 1:
64 + raise ValueError(f"num_lanes must be at least one, got {num_lanes}.")
65 + self._num_lanes = num_lanes
66 + self._current_workspaces: list[torch.Tensor | None] = [None] * (
67 + self._num_ubatches * self._num_lanes
68 + )
69 self._locked: bool = False
70
71 @staticmethod
@@ -126,7 +150,14 @@ class WorkspaceManager:
150 The current workspace tensor.
151 """
152 ubatch_id = dbo_current_ubatch_id()
129 - current_workspace = self._current_workspaces[ubatch_id]
153 + lane = _workspace_lane.get()
154 + if lane >= self._num_lanes:
155 + raise RuntimeError(
156 + f"Workspace lane {lane} is not configured; manager has "
157 + f"{self._num_lanes} lane(s)."
158 + )
159 + workspace_id = ubatch_id * self._num_lanes + lane
160 + current_workspace = self._current_workspaces[workspace_id]
161 current_size = self._workspace_size_bytes(current_workspace)
162
163 if current_size < required_bytes:
@@ -161,11 +192,11 @@ class WorkspaceManager:
192 "Workspace growth is not allowed after locking."
193 )
194
164 - # Only resize the requesting ubatch's workspace. Other
165 - # ubatches resize lazily on their next get_simultaneous call.
195 + # Only resize the requesting ubatch/lane workspace. Other slots
196 + # resize lazily on their next get_simultaneous call.
197 # Resizing all ubatches here would orphan the other ubatch's
198 # old tensor when it still holds views into it (DBO leak).
168 - self._current_workspaces[ubatch_id] = None
199 + self._current_workspaces[workspace_id] = None
200 del current_workspace
201 # Release the freed segment back to CUDA so the caching
202 # allocator can reuse the GPU memory for the larger
@@ -173,19 +204,20 @@ class WorkspaceManager:
204 # dead segment in reserved memory which can cause higher peak
205 # memory usage.
206 torch.accelerator.empty_cache()
176 - self._current_workspaces[ubatch_id] = torch.empty(
207 + self._current_workspaces[workspace_id] = torch.empty(
208 (required_bytes,), dtype=torch.uint8, device=self._device
209 )
179 - current_workspace = self._current_workspaces[ubatch_id]
210 + current_workspace = self._current_workspaces[workspace_id]
211
212 if envs.VLLM_DEBUG_WORKSPACE:
213 logger.info(
214 "[WORKSPACE DEBUG] Resized workspace from '%s': %.2f MB -> "
184 - "%.2f MB (ubatch %d)",
215 + "%.2f MB (ubatch %d, lane %d)",
216 get_caller_info(),
217 current_size / _MB,
218 required_bytes / _MB,
219 ubatch_id,
220 + lane,
221 )
222
223 return current_workspace
@@ -214,7 +246,9 @@ def current_workspace_manager() -> "WorkspaceManager":
246
247
248 def init_workspace_manager(
217 - device: torch.device, num_ubatches: int | None = None
249 + device: torch.device,
250 + num_ubatches: int | None = None,
251 + num_lanes: int = 1,
252 ) -> None:
253 """Initialize the workspace manager with a device.
254
@@ -224,6 +258,7 @@ def init_workspace_manager(
258 Args:
259 device: The device to allocate workspace on.
260 num_ubatches: Number of workspace ubatch slots. Defaults to 1.
261 + num_lanes: Number of independent execution lanes per ubatch. Defaults to 1.
262 """
263 global _manager
264 if _manager is not None:
@@ -233,7 +268,7 @@ def init_workspace_manager(
268 _manager._device,
269 device,
270 )
236 - _manager = WorkspaceManager(device, num_ubatches)
271 + _manager = WorkspaceManager(device, num_ubatches, num_lanes)
272
273
274 def lock_workspace() -> None: