[Model] Support DeepSeek-V4.1-Flash (#56214)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> Signed-off-by: andyluo7 <andy.luo@amd.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: andyluo7 <andy.luo@amd.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Yongye Zhu committed
Sep 11, 2026 at 02:11 UTC
e77daef89e18e08321ae7b8b24827eedd5fe8673
49 files changed
+3996
-138
.buildkite/test_areas/kernels.yaml
+3
@@ -102,6 +102,8 @@ steps:
102
- "!vllm/distributed/kv_transfer/"
103
- tests/kernels/
104
commands:
105
+ # Isolate mHC from FlashInfer's import-time ptxas override.
106
+ - pytest -v -s kernels/test_mhc_kernels.py
107
- pytest -v -s kernels/
108
--ignore=kernels/attention
109
--ignore=kernels/core
@@ -117,6 +119,7 @@ steps:
119
--ignore=kernels/test_kda.py
120
--ignore=kernels/test_bf16x3_router_gemm_cutedsl.py
121
--ignore=kernels/test_ll_bf16_gemm.py
122
+ --ignore=kernels/test_mhc_kernels.py
123
--ignore=kernels/test_shuffle_rows.py
124
# BROKEN on main, pending kernel fixes (B200):
125
# test_shuffle_rows.py (1: test_shuffle_rows_edge_cases)
tests/fusion/test_quant_activation_contract.py
+8
@@ -9,8 +9,13 @@ from vllm.model_executor.kernels.linear import (
9
_POSSIBLE_FP8_BLOCK_KERNELS,
10
_POSSIBLE_FP8_KERNELS,
11
_POSSIBLE_INT8_KERNELS,
12
+ _POSSIBLE_MXFP8_KERNELS,
13
_POSSIBLE_NVFP4_KERNELS,
14
)
15
+from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
16
+ FlashInferCutedslMxfp8LinearKernel,
17
+ FlashInferCutlassMxfp8LinearKernel,
18
+)
19
from vllm.model_executor.kernels.linear.nvfp4.base import (
20
NvFp4LinearKernel,
21
NvFp4LinearLayerConfig,
@@ -50,6 +55,8 @@ SUPPORTING = {
55
FlashInferFP8ScaledMMLinearKernel,
56
FlashInferCutlassNvFp4LinearKernel,
57
PerTensorTorchFP8ScaledMMLinearKernel,
58
+ FlashInferCutedslMxfp8LinearKernel,
59
+ FlashInferCutlassMxfp8LinearKernel,
60
}
61
62
@@ -60,6 +67,7 @@ def _all_kernel_classes() -> list[type]:
67
_POSSIBLE_FP8_BLOCK_KERNELS,
68
_POSSIBLE_INT8_KERNELS,
69
_POSSIBLE_NVFP4_KERNELS,
70
+ _POSSIBLE_MXFP8_KERNELS,
71
):
72
for kernels in registry.values():
73
for cls in kernels:
tests/kernels/attention/test_rocm_triton_attn_dsv4.py
+25
@@ -1268,6 +1268,31 @@ def test_get_cached_wo_a_bf16_plain_caches() -> None:
1268
torch.testing.assert_close(out2, expected, atol=0, rtol=0)
1269
1270
1271
+@pytest.mark.parametrize("scale_attr", ["weight_scale", "weight_scale_inv"])
1272
+@torch.inference_mode()
1273
+def test_get_cached_wo_a_bf16_dequantized_ignores_retained_scale(
1274
+ scale_attr: str,
1275
+) -> None:
1276
+ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _get_cached_wo_a_bf16
1277
+
1278
+ device = torch.device("cuda")
1279
+ torch.manual_seed(6)
1280
+ n_local_groups, o_lora_rank, hidden_dim = 2, 4, 8
1281
+ weight = torch.randn(
1282
+ n_local_groups * o_lora_rank, hidden_dim, dtype=torch.bfloat16, device=device
1283
+ )
1284
+ retained_scale = torch.full(
1285
+ (n_local_groups, 1), 0.25, dtype=torch.float32, device=device
1286
+ )
1287
+ wo_a = _FakeWoA(weight)
1288
+ setattr(wo_a, scale_attr, retained_scale)
1289
+
1290
+ out = _get_cached_wo_a_bf16(wo_a, n_local_groups, o_lora_rank, hidden_dim)
1291
+
1292
+ expected = weight.view(n_local_groups, o_lora_rank, hidden_dim)
1293
+ torch.testing.assert_close(out, expected, atol=0, rtol=0)
1294
+
1295
+
1296
@torch.inference_mode()
1297
def test_get_cached_wo_a_bf16_fp8_blockscale_caches() -> None:
1298
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _get_cached_wo_a_bf16
tests/kernels/core/test_fused_q_kv_rmsnorm.py
+189
@@ -112,3 +112,192 @@ def test_fused_q_kv_rmsnorm_launches_past_grid_y_cap(num_tokens: int):
112
rtol=1e-2,
113
atol=1e-2,
114
)
115
+
116
+
117
+@pytest.mark.skipif(
118
+ not current_platform.has_device_capability(100), reason="MXFP8 needs Blackwell"
119
+)
120
+@pytest.mark.parametrize("num_tokens", [0, 1, 17, 128, 129, 1024])
121
+@pytest.mark.parametrize("q_size", [384, 1280, 1344])
122
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
123
+def test_fused_q_kv_rmsnorm_quant_matches_separate(num_tokens, q_size, dtype):
124
+ """Preserve native MXFP8 rounding and all swizzled scale padding bytes."""
125
+ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
126
+ mxfp8_e4m3_quantize,
127
+ )
128
+ from vllm.models.deepseek_v4_1.common.ops.query_quant import (
129
+ fused_q_kv_rmsnorm_quant,
130
+ )
131
+
132
+ torch.manual_seed(0)
133
+ kv_size = 512
134
+ x = torch.randn(num_tokens, q_size + kv_size, device="cuda", dtype=dtype)
135
+ x[:1] = 0
136
+ x[1:2] *= 1e-30
137
+ qr, kv = x.split([q_size, kv_size], -1)
138
+ qw = torch.randn(q_size, device="cuda", dtype=x.dtype)
139
+ kvw = torch.randn(kv_size, device="cuda", dtype=x.dtype)
140
+ eps = 1e-20
141
+ result = fused_q_kv_rmsnorm_quant(qr, kv, qw, kvw, eps)
142
+ assert result[0].orig_dtype == dtype
143
+ assert result[0].orig_shape == qr.shape
144
+ if not num_tokens:
145
+ assert result[0].data.shape == qr.shape
146
+ assert result[0].scale.numel() == 0
147
+ assert result[1].shape == kv.shape
148
+ return
149
+ qr_ref, kv_ref = fused_q_kv_rmsnorm(qr, kv, qw, kvw, eps)
150
+ q_ref, scale_ref = mxfp8_e4m3_quantize(qr_ref, is_sf_swizzled_layout=True)
151
+ for replay in (False, True):
152
+ if replay:
153
+ graph = torch.cuda.CUDAGraph()
154
+ with torch.cuda.graph(graph):
155
+ result = fused_q_kv_rmsnorm_quant(qr, kv, qw, kvw, eps)
156
+ graph.replay()
157
+ torch.testing.assert_close(
158
+ result[0].data.view(torch.uint8), q_ref.view(torch.uint8), rtol=0, atol=0
159
+ )
160
+ torch.testing.assert_close(result[0].scale, scale_ref, rtol=0, atol=0)
161
+ torch.testing.assert_close(result[1], kv_ref, rtol=0, atol=0)
162
+
163
+
164
+@pytest.mark.skipif(
165
+ not current_platform.has_device_capability(100), reason="MXFP8 needs Blackwell"
166
+)
167
+@pytest.mark.parametrize("backend", ["cute-dsl", "cutlass"])
168
+@pytest.mark.parametrize("num_tokens", [1, 17, 256])
169
+@pytest.mark.parametrize("projection", ["attention", "indexer"])
170
+@pytest.mark.parametrize("bias", [False, True])
171
+def test_shared_query_quant_preserves_projection(
172
+ num_tokens, backend, projection, bias, monkeypatch
173
+):
174
+ """The native-checkpoint projection keeps identical GEMM outputs."""
175
+ from dataclasses import replace
176
+ from types import SimpleNamespace
177
+ from typing import Any, cast
178
+
179
+ from vllm.config import (
180
+ CompilationConfig,
181
+ CompilationMode,
182
+ VllmConfig,
183
+ set_current_vllm_config,
184
+ )
185
+ from vllm.model_executor.kernels.linear.mxfp8 import Mxfp8LinearLayerConfig
186
+ from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
187
+ FlashInferCutedslMxfp8LinearKernel,
188
+ FlashInferCutlassMxfp8LinearKernel,
189
+ )
190
+ from vllm.model_executor.layers.fusion.quant_activation import (
191
+ expose_input_quant_key,
192
+ )
193
+ from vllm.model_executor.layers.linear import ColumnParallelLinear, ReplicatedLinear
194
+ from vllm.model_executor.layers.quantization.modelopt import (
195
+ ModelOptLinearMethod,
196
+ )
197
+ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
198
+ mxfp8_e4m3_quantize,
199
+ )
200
+ from vllm.model_executor.layers.quantization.utils.quant_utils import kNvfp4Dynamic
201
+ from vllm.models.deepseek_v4_1.attention import (
202
+ DeepseekV4Attention,
203
+ DeepseekV4Indexer,
204
+ )
205
+ from vllm.models.deepseek_v4_1.common.ops.query_quant import (
206
+ can_fuse_query_quant,
207
+ fused_q_kv_rmsnorm_quant,
208
+ )
209
+ from vllm.models.deepseek_v4_1.quant_config import DeepseekV4FP8Config
210
+
211
+ cls: (
212
+ type[FlashInferCutedslMxfp8LinearKernel]
213
+ | type[FlashInferCutlassMxfp8LinearKernel]
214
+ )
215
+ if backend == "cute-dsl":
216
+ cls = FlashInferCutedslMxfp8LinearKernel
217
+ else:
218
+ cls = FlashInferCutlassMxfp8LinearKernel
219
+ supported, reason = cls.is_supported()
220
+ if not supported:
221
+ pytest.skip(reason)
222
+ torch.manual_seed(0)
223
+ with set_current_vllm_config(
224
+ VllmConfig(compilation_config=CompilationConfig(mode=CompilationMode.NONE))
225
+ ):
226
+ kernel = cls(Mxfp8LinearLayerConfig())
227
+ # Match the checkpoint's TP4 attention and replicated indexer widths.
228
+ if projection == "attention":
229
+ linear = ColumnParallelLinear.__new__(ColumnParallelLinear)
230
+ output_size = 8192
231
+ else:
232
+ linear = ReplicatedLinear.__new__(ReplicatedLinear)
233
+ output_size = 4096
234
+ torch.nn.Module.__init__(linear)
235
+ linear.bias = (
236
+ torch.nn.Parameter(
237
+ torch.randn(output_size, device="cuda", dtype=torch.bfloat16)
238
+ )
239
+ if bias
240
+ else None
241
+ )
242
+ linear.skip_bias_add = False
243
+ linear.return_bias = projection == "indexer"
244
+ linear.gather_output = False
245
+ w, ws = mxfp8_e4m3_quantize(
246
+ torch.randn(output_size, 1280, device="cuda", dtype=torch.bfloat16)
247
+ )
248
+ linear.weight = torch.nn.Parameter(w, requires_grad=False)
249
+ linear.weight_scale = torch.nn.Parameter(ws, requires_grad=False)
250
+ quant_config = DeepseekV4FP8Config.from_config(
251
+ {
252
+ "quant_method": "fp8",
253
+ "activation_scheme": "dynamic",
254
+ "weight_block_size": [32, 32],
255
+ "scale_fmt": "ue8m0",
256
+ "expert_dtype": "fp4",
257
+ }
258
+ )
259
+ method = quant_config.get_quant_method(linear, "model.layers.2.attn.wq_b")
260
+ assert type(method) is ModelOptLinearMethod
261
+ method.out_dtype = torch.bfloat16
262
+ method.kernel = kernel
263
+ linear.quant_method = method
264
+ kernel.process_weights_after_loading(linear)
265
+ expose_input_quant_key(linear, kernel)
266
+ assert can_fuse_query_quant([linear, linear])
267
+ x = torch.randn(num_tokens, 1792, device="cuda", dtype=torch.bfloat16)
268
+ qr, kv = x.split([1280, 512], -1)
269
+ qw = torch.randn(1280, device="cuda", dtype=x.dtype)
270
+ kvw = torch.randn(512, device="cuda", dtype=x.dtype)
271
+ qr_ref, _ = fused_q_kv_rmsnorm(qr, kv, qw, kvw, 1e-20)
272
+ q, _ = fused_q_kv_rmsnorm_quant(qr, kv, qw, kvw, 1e-20)
273
+ owner = cast(Any, SimpleNamespace(wq_b=linear))
274
+ project = (
275
+ DeepseekV4Attention._wq_b_proj
276
+ if projection == "attention"
277
+ else DeepseekV4Indexer._wq_b_proj
278
+ )
279
+ expected = project(owner, qr_ref)
280
+
281
+ def reject_requantization(*args, **kwargs):
282
+ raise AssertionError("pre-quantized input must skip quantization")
283
+
284
+ monkeypatch.setattr(
285
+ "vllm.model_executor.kernels.linear.mxfp8.flashinfer.mxfp8_e4m3_quantize",
286
+ reject_requantization,
287
+ )
288
+ actual = project(owner, q)
289
+ torch.testing.assert_close(actual, expected, rtol=0, atol=0)
290
+ graph = torch.cuda.CUDAGraph()
291
+ with torch.cuda.graph(graph):
292
+ actual = project(owner, q)
293
+ graph.replay()
294
+ torch.testing.assert_close(actual, expected, rtol=0, atol=0)
295
+ assert not can_fuse_query_quant(
296
+ [linear, torch.nn.Linear(1280, 256, bias=False)]
297
+ )
298
+ reshaped = project(
299
+ owner, replace(q, orig_shape=torch.Size((1, num_tokens, 1280)))
300
+ )
301
+ torch.testing.assert_close(reshaped, expected.unsqueeze(0), rtol=0, atol=0)
302
+ with pytest.raises(AssertionError, match="consumer kernel"):
303
+ linear(replace(q, quant_key=kNvfp4Dynamic))
tests/kernels/quantization/test_rocm_mxfp8_linear.py
new
+144
@@ -0,0 +1,144 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Native CDNA4 MXFP8 linear (``RocmDotScaledMxfp8LinearKernel``).
4
+
5
+DeepSeek V4.1 routes every attention linear through this kernel. It used to
6
+carry two BF16-dequantize workarounds on ROCm (one for all ``.attn.`` linears,
7
+one for the row-parallel ``wo_b``, whose comment blamed NaNs at small token
8
+counts). These tests pin the two properties those workarounds stood in for --
9
+finite output and MXFP8-level accuracy -- across the V4.1 attention shapes at
10
+TP=4 and the small token counts decode actually runs.
11
+"""
12
+
13
+import pytest
14
+import torch
15
+
16
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
17
+ _mxfp8_e4m3_quantize_torch,
18
+ _mxfp8_e4m3_quantize_triton,
19
+ dequant_mxfp8_to_bf16,
20
+ mxfp8_e4m3_quantize,
21
+)
22
+from vllm.platforms import current_platform
23
+
24
+pytestmark = pytest.mark.skipif(
25
+ not current_platform.is_rocm() or not current_platform.supports_mx(),
26
+ reason="native MX linear requires CDNA4 (gfx95x)",
27
+)
28
+
29
+# (name, N, K) for the DeepSeek V4.1 attention linears at tensor_parallel=4:
30
+# fused wq_a|wkv is replicated, wq_b/wo_a are column-parallel, wo_b is
31
+# row-parallel (so K, not N, is sharded).
32
+V41_ATTN_SHAPES = [
33
+ ("fused_wqa_wkv", 1280 + 512, 5120),
34
+ ("wq_b", 64 * 512 // 4, 1280),
35
+ ("wo_a", 8 * 1024 // 4, 4096),
36
+ ("wo_b", 5120, 8192 // 4),
37
+]
38
+
39
+
40
+def _make_layer(n: int, k: int, device: str) -> torch.nn.Module:
41
+ weight_bf16 = torch.randn(n, k, device=device, dtype=torch.bfloat16) * 0.05
42
+ weight, weight_scale = _mxfp8_e4m3_quantize_torch(weight_bf16)
43
+ layer = torch.nn.Module()
44
+ layer.weight = torch.nn.Parameter(weight, requires_grad=False)
45
+ layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
46
+ return layer
47
+
48
+
49
+@pytest.mark.parametrize("name,n,k", V41_ATTN_SHAPES, ids=lambda v: str(v))
50
+@pytest.mark.parametrize("num_tokens", [1, 2, 7, 16, 31, 32, 128, 1024])
51
+@torch.inference_mode()
52
+def test_rocm_mxfp8_linear_matches_dequant_reference(name, n, k, num_tokens):
53
+ """Output is finite and matches an exact MXFP8 GEMM of the same inputs.
54
+
55
+ The reference quantizes the activation with the same quantizer and then
56
+ contracts in FP32, so the tolerance covers only kernel/accumulation error,
57
+ not the MXFP8 quantization error both sides share.
58
+ """
59
+ from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import (
60
+ Mxfp8LinearLayerConfig,
61
+ )
62
+ from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (
63
+ RocmDotScaledMxfp8LinearKernel,
64
+ )
65
+
66
+ torch.manual_seed(num_tokens)
67
+ device = "cuda"
68
+ layer = _make_layer(n, k, device)
69
+ kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(bmm_batch_size=None))
70
+ kernel.process_weights_after_loading(layer)
71
+
72
+ x = torch.randn(num_tokens, k, device=device, dtype=torch.bfloat16) * 0.5
73
+ out = kernel.apply_weights(layer, x)
74
+
75
+ assert out.shape == (num_tokens, n)
76
+ assert torch.isfinite(out).all(), f"{name}: non-finite output"
77
+
78
+ x_q, x_scale = mxfp8_e4m3_quantize(x)
79
+ expected = (
80
+ dequant_mxfp8_to_bf16(x_q, x_scale).float()
81
+ @ dequant_mxfp8_to_bf16(layer.weight, layer.weight_scale).float().T
82
+ )
83
+ rel = (out.float() - expected).norm() / expected.norm()
84
+ assert rel < 5e-3, f"{name}: relative error {rel:.4f}"
85
+
86
+
87
+@torch.inference_mode()
88
+def test_rocm_mxfp8_linear_handles_extreme_magnitudes():
89
+ """Rows that are all-zero or near the E4M3 range must not produce NaNs.
90
+
91
+ A degenerate block makes the dynamic activation quantizer pick an extreme
92
+ E8M0 scale; a zero row makes it pick the tiny-clamped one. Both feed
93
+ ``tl.dot_scaled`` directly.
94
+ """
95
+ from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import (
96
+ Mxfp8LinearLayerConfig,
97
+ )
98
+ from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (
99
+ RocmDotScaledMxfp8LinearKernel,
100
+ )
101
+
102
+ torch.manual_seed(0)
103
+ device = "cuda"
104
+ n, k = 5120, 2048 # wo_b at TP=4
105
+ layer = _make_layer(n, k, device)
106
+ kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(bmm_batch_size=None))
107
+ kernel.process_weights_after_loading(layer)
108
+
109
+ x = torch.randn(8, k, device=device, dtype=torch.bfloat16)
110
+ x[0] = 0.0
111
+ x[1] = 1e-30
112
+ x[2] = 400.0
113
+ x[3, ::2] = 0.0
114
+
115
+ out = kernel.apply_weights(layer, x)
116
+ assert torch.isfinite(out).all()
117
+ assert (out[0] == 0).all(), "an all-zero activation row must give a zero row"
118
+
119
+
120
+@pytest.mark.parametrize("num_rows", [1, 64, 65])
121
+@torch.inference_mode()
122
+def test_rocm_mxfp8_quantizer_matches_torch_on_degenerate_blocks(num_rows):
123
+ """The fused ROCm quantizer must agree with the torch one, zeros included.
124
+
125
+ An all-zero 32-element block clamps the E8M0 scale to 0, i.e. 2**-127 --
126
+ subnormal in fp32 and flushed to zero on CDNA. Dividing by it turned the
127
+ block into 0/0 == NaN, which then poisoned every output column of the GEMM
128
+ that consumed it.
129
+ """
130
+ torch.manual_seed(num_rows)
131
+ device = "cuda"
132
+ x = torch.randn(num_rows, 128, device=device, dtype=torch.bfloat16)
133
+ x[0] = 0.0 # all-zero row: every block clamps to scale byte 0
134
+ x[-1, :32] = 0.0 # single zero block in an otherwise normal row
135
+ x[-1, 32:64] = 1e-38 # denormal-ish block
136
+
137
+ q_triton, s_triton = _mxfp8_e4m3_quantize_triton(x)
138
+ q_torch, s_torch = _mxfp8_e4m3_quantize_torch(x)
139
+
140
+ assert not torch.isnan(q_triton.float()).any()
141
+ torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0)
142
+ torch.testing.assert_close(
143
+ q_triton.view(torch.uint8), q_torch.view(torch.uint8), rtol=0, atol=0
144
+ )
tests/kernels/test_compressor_kv_cache.py
+432
@@ -42,6 +42,438 @@ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
42
from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4
43
44
45
+@pytest.mark.skipif(
46
+ not current_platform.is_cuda_alike(), reason="graph capture coverage"
47
+)
48
+@pytest.mark.parametrize("compress_ratio", [1, 2])
49
+@pytest.mark.parametrize("use_graph", [False, True])
50
+@pytest.mark.parametrize(
51
+ "lengths",
52
+ [(1, 5, 13), (1, 1, 1), (1, 2, 0), (7, 8, 8), (2, 3, 3)],
53
+ ids=["prefill", "decode", "empty", "wrap_threshold", "packed_boundary"],
54
+)
55
+def test_v41_fused_save_compress_and_insert(
56
+ compress_ratio: int, use_graph: bool, lengths: tuple[int, ...]
57
+):
58
+ """Ring states: a group boundary whose previous token is outside the chunk
59
+ pools the ring row that chunk left behind, even when this chunk overwrites
60
+ that row; only the last ``capacity`` tokens of a chunk are stored.
61
+
62
+ Exercise historical group starts, adjacent requests, a chunk longer than
63
+ the ring, single-token decode pairs, empty requests, wrap thresholds at
64
+ odd and even starts, padded graph rows, and padded physical pages against
65
+ a Torch reference with BF16 rounding before RoPE. Replays change both
66
+ inputs and history so stale scratch/cache values cannot satisfy the test.
67
+ """
68
+ from vllm.models.deepseek_v4_1.common.ops.fused_compress_quant_cache import (
69
+ fused_save_compress_norm,
70
+ rope_quant_insert,
71
+ )
72
+
73
+ torch.manual_seed(42)
74
+ device = "cuda"
75
+ capacity = 8
76
+ starts = [7, 8, 13]
77
+ actual = sum(lengths)
78
+ num_tokens = actual + 5
79
+ positions = torch.zeros(num_tokens, dtype=torch.int64, device=device)
80
+ req_ids = torch.zeros(actual + 2, dtype=torch.int32, device=device)
81
+ state_slots = torch.full((actual + 2,), -1, dtype=torch.int64, device=device)
82
+ cache_slots = torch.full_like(state_slots, -1)
83
+ query_start_loc = torch.tensor(
84
+ [0, *torch.tensor(lengths).cumsum(0).tolist()], dtype=torch.int32
85
+ ).to(device)
86
+ ring_blocks = torch.randperm(6, device=device)[:3]
87
+ raw = torch.randn(num_tokens, 512 * compress_ratio, device=device)
88
+ norm = torch.randn(512, dtype=torch.bfloat16, device=device)
89
+ latent = torch.empty(num_tokens, 512, dtype=torch.bfloat16, device=device)
90
+ angles = torch.randn(32, 32, device=device)
91
+ cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1)
92
+ # capacity x 1024 floats rounded up to 576-byte alignment.
93
+ state_backing = torch.empty(6, 8208, device=device)
94
+ state = state_backing.as_strided((6, capacity, 1024), (8208, 1024, 1))
95
+ cache_block = 256 // compress_ratio
96
+ cache_stride = math.ceil(cache_block * 584 / 576) * 576
97
+ cache_backing = torch.empty(3, cache_stride, dtype=torch.uint8, device=device)
98
+ cache = cache_backing.as_strided((3, cache_block, 584), (cache_stride, 584, 1))
99
+ cursor = 0
100
+ for req, (start, length) in enumerate(zip(starts, lengths)):
101
+ pos = torch.arange(start, start + length, device=device)
102
+ rows = slice(cursor, cursor + length)
103
+ positions[rows] = pos
104
+ req_ids[rows] = req
105
+ state_slots[rows] = ring_blocks[req] * capacity + pos % capacity
106
+ cache_slots[rows] = (2 - req) * cache_block + pos // compress_ratio
107
+ cursor += length
108
+
109
+ has_ring = compress_ratio == 2
110
+
111
+ def run():
112
+ fused_save_compress_norm(
113
+ raw,
114
+ positions,
115
+ state if has_ring else None,
116
+ state_slots,
117
+ query_start_loc if has_ring else None,
118
+ req_ids if has_ring else None,
119
+ norm,
120
+ 1e-20,
121
+ compress_ratio,
122
+ latent,
123
+ )
124
+ rope_quant_insert(
125
+ latent, positions, cos_sin, cache, cache_slots, compress_ratio
126
+ )
127
+
128
+ state_backing.normal_()
129
+ run()
130
+ graph = None
131
+ if use_graph:
132
+ stream = torch.cuda.Stream()
133
+ stream.wait_stream(torch.cuda.current_stream())
134
+ with torch.cuda.stream(stream):
135
+ run()
136
+ torch.cuda.current_stream().wait_stream(stream)
137
+ graph = torch.cuda.CUDAGraph()
138
+ with torch.cuda.graph(graph):
139
+ run()
140
+
141
+ for _ in range(3):
142
+ raw.normal_()
143
+ state_backing.normal_()
144
+ history = state.clone()
145
+ expected_state = state_backing.clone()
146
+ expected_rows = expected_state.as_strided((6, capacity, 1024), (8208, 1024, 1))
147
+ cache_backing.fill_(165)
148
+ expected_cache = cache_backing.clone()
149
+ latent.fill_(float("nan"))
150
+ # Ratio 1 has no ring; ratio 2 stores each chunk's last `capacity` rows.
151
+ for t in range(actual if has_ring else 0):
152
+ req = req_ids[t].item()
153
+ if query_start_loc[req + 1].item() - t > capacity:
154
+ continue
155
+ slot = state_slots[t].item()
156
+ expected_rows[slot // capacity, slot % capacity, :512] = raw[t, :512]
157
+ expected_rows[slot // capacity, slot % capacity, 512:] = raw[t, 512:]
158
+
159
+ if graph is None:
160
+ run()
161
+ else:
162
+ graph.replay()
163
+
164
+ torch.testing.assert_close(state_backing, expected_state, rtol=0, atol=0)
165
+ for t in range(num_tokens):
166
+ pos = positions[t].item()
167
+ if t >= actual or (pos + 1) % compress_ratio:
168
+ assert torch.isnan(latent[t]).all()
169
+ continue
170
+ req = req_ids[t].item()
171
+ group_rows = []
172
+ for k in range(compress_ratio - 1, -1, -1):
173
+ if t - k >= query_start_loc[req].item():
174
+ kv_score = raw[t - k]
175
+ if compress_ratio == 1:
176
+ kv_score = torch.cat((kv_score, torch.zeros_like(kv_score)))
177
+ group_rows.append(kv_score)
178
+ else:
179
+ ring = ring_blocks[req].item()
180
+ group_rows.append(history[ring, (pos - k) % capacity])
181
+ group = torch.stack(group_rows)
182
+ pooled = (group[:, :512] * group[:, 512:].softmax(0)).sum(0)
183
+ normed = pooled * torch.rsqrt(pooled.square().mean() + 1e-20) * norm
184
+ torch.testing.assert_close(latent[t].float(), normed, rtol=0.004, atol=1e-6)
185
+
186
+ slot = cache_slots[t].item()
187
+ page, row = divmod(slot, cache_block)
188
+ values = expected_cache[page, row * 576 : (row + 1) * 576]
189
+ quantized, scales = _ue8m0_reference(latent[t, :448], 64, 448.0)
190
+ values[:448] = quantized.view(torch.uint8)
191
+ c, s = cos_sin[pos // compress_ratio * compress_ratio].chunk(2)
192
+ rope_input = latent[t, 448:].float()
193
+ rotated = (
194
+ torch.stack(
195
+ (
196
+ rope_input[0::2] * c - rope_input[1::2] * s,
197
+ rope_input[1::2] * c + rope_input[0::2] * s,
198
+ ),
199
+ dim=-1,
200
+ )
201
+ .flatten()
202
+ .to(torch.bfloat16)
203
+ )
204
+ actual_rope = cache_backing[page, row * 576 + 448 : (row + 1) * 576].view(
205
+ torch.bfloat16
206
+ )
207
+ torch.testing.assert_close(actual_rope, rotated, rtol=0.008, atol=1e-6)
208
+ values[448:] = actual_rope.view(torch.uint8)
209
+ scale_offset = cache_block * 576 + row * 8
210
+ expected_cache[page, scale_offset : scale_offset + 7] = (
211
+ scales.log2() + 127
212
+ ).to(torch.uint8)
213
+ expected_cache[page, scale_offset + 7] = 0
214
+ # Includes unwritten rows, page padding, quantized values and scales.
215
+ torch.testing.assert_close(cache_backing, expected_cache, rtol=0, atol=0)
216
+
217
+
218
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
219
+@pytest.mark.parametrize("compress_ratio", [1, 2])
220
+@pytest.mark.parametrize("store_fp8", [False, True])
221
+def test_v41_rope_insert_plain_row(compress_ratio: int, store_fp8: bool):
222
+ """The FlashInfer plain-row compressed cache gets [448 NoPE | 64 RoPE] rows.
223
+
224
+ bf16 caches store the latent verbatim with the RoPE tail rotated; per-tensor
225
+ fp8 caches scale the bf16-rounded row by 1/fp8_scale and clamp to e4m3.
226
+ Non-boundary positions, negative slots and untouched rows must keep their
227
+ prior contents.
228
+ """
229
+ from vllm.models.deepseek_v4_1.common.ops.fused_compress_quant_cache import (
230
+ rope_quant_insert,
231
+ )
232
+
233
+ torch.manual_seed(3)
234
+ device = "cuda"
235
+ num_tokens = 9
236
+ cache_block = 4
237
+ positions = torch.arange(5, 5 + num_tokens, dtype=torch.int64, device=device)
238
+ cache_slots = torch.randperm(3 * cache_block, device=device)[:num_tokens]
239
+ cache_slots[2] = -1
240
+ latent = torch.randn(num_tokens, 512, device=device).to(torch.bfloat16)
241
+ angles = torch.randn(32, 32, device=device)
242
+ cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1)
243
+ dtype = torch.float8_e4m3fn if store_fp8 else torch.bfloat16
244
+ # Pad the page stride so the kernel must honour strides, not shape.
245
+ cache_backing = torch.full((3, (cache_block + 1) * 512), 3.0, device=device)
246
+ cache_backing = cache_backing.to(dtype)
247
+ cache = cache_backing.as_strided(
248
+ (3, cache_block, 512), ((cache_block + 1) * 512, 512, 1)
249
+ )
250
+ fp8_scale = torch.tensor([0.5], dtype=torch.float32, device=device)
251
+ expected = cache_backing.clone().float()
252
+
253
+ rope_quant_insert(
254
+ latent,
255
+ positions,
256
+ cos_sin,
257
+ cache,
258
+ cache_slots,
259
+ compress_ratio,
260
+ fp8_scale=fp8_scale if store_fp8 else None,
261
+ )
262
+
263
+ for t in range(num_tokens):
264
+ pos = positions[t].item()
265
+ slot = cache_slots[t].item()
266
+ if slot < 0 or (pos + 1) % compress_ratio:
267
+ continue
268
+ c, s = cos_sin[pos // compress_ratio * compress_ratio].chunk(2)
269
+ rope_input = latent[t, 448:].float()
270
+ rotated = torch.stack(
271
+ (
272
+ rope_input[0::2] * c - rope_input[1::2] * s,
273
+ rope_input[1::2] * c + rope_input[0::2] * s,
274
+ ),
275
+ dim=-1,
276
+ ).flatten()
277
+ row = torch.cat((latent[t, :448], rotated.to(torch.bfloat16))).float()
278
+ if store_fp8:
279
+ row = (row * (1.0 / fp8_scale)).clamp(-448.0, 448.0)
280
+ row = row.to(torch.float8_e4m3fn).float()
281
+ page, idx = divmod(slot, cache_block)
282
+ expected[page, idx * 512 : (idx + 1) * 512] = row
283
+
284
+ actual = cache_backing.float()
285
+ if store_fp8:
286
+ # One e4m3 ulp of slack for FMA-vs-separate rounding in the RoPE tail.
287
+ torch.testing.assert_close(actual, expected, rtol=0.13, atol=1e-2)
288
+ else:
289
+ nope = torch.arange(512, device=device) < 448
290
+ nope = nope.repeat(cache_block + 1)
291
+ torch.testing.assert_close(actual[:, nope], expected[:, nope], rtol=0, atol=0)
292
+ torch.testing.assert_close(
293
+ actual[:, ~nope], expected[:, ~nope], rtol=0.008, atol=1e-6
294
+ )
295
+
296
+
297
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Triton kernel")
298
+def test_v41_compressor_metadata_maps_tokens_to_their_ring():
299
+ """The ring group's generic slot mapping is disabled (all PAD), so the
300
+ builder must map every real token to ``ring_block * capacity + pos %
301
+ capacity`` and keep padding tokens at PAD."""
302
+ from unittest.mock import MagicMock
303
+
304
+ from vllm.models.deepseek_v4_1.compressor import CompressorMetadataBuilder
305
+ from vllm.v1.attention.backend import CommonAttentionMetadata
306
+ from vllm.v1.kv_cache_interface import CircularBufferSpec
307
+
308
+ capacity = 8
309
+ vllm_config = MagicMock()
310
+ vllm_config.scheduler_config.max_num_batched_tokens = 16
311
+ spec = CircularBufferSpec(
312
+ block_size=capacity,
313
+ num_kv_heads=1,
314
+ head_size=1024,
315
+ head_size_v=0,
316
+ dtype=torch.float32,
317
+ )
318
+ device = torch.device("cuda")
319
+ builder = CompressorMetadataBuilder(spec, ["state"], vllm_config, device)
320
+
321
+ # Two requests: 3 tokens at positions 13..15 on ring block 5, then 2
322
+ # tokens at positions 7..8 on ring block 2; three padding tokens.
323
+ query_start_loc = torch.tensor([0, 3, 5], dtype=torch.int32, device=device)
324
+ positions = torch.tensor([13, 14, 15, 7, 8, 0, 0, 0], device=device)
325
+ block_table = torch.tensor([[5], [2]], dtype=torch.int32, device=device)
326
+ common = CommonAttentionMetadata(
327
+ query_start_loc=query_start_loc,
328
+ query_start_loc_cpu=query_start_loc.cpu(),
329
+ seq_lens=torch.tensor([16, 9], dtype=torch.int32, device=device),
330
+ num_reqs=2,
331
+ num_actual_tokens=5,
332
+ max_query_len=3,
333
+ max_seq_len=16,
334
+ block_table_tensor=block_table,
335
+ slot_mapping=torch.full((8,), -1, dtype=torch.int64, device=device),
336
+ positions=positions,
337
+ )
338
+ metadata = builder.build(0, common)
339
+
340
+ expected = [5 * 8 + 5, 5 * 8 + 6, 5 * 8 + 7, 2 * 8 + 7, 2 * 8 + 0, -1, -1, -1]
341
+ assert metadata.slot_mapping.tolist() == expected
342
+ assert metadata.query_start_loc is query_start_loc
343
+ assert metadata.token_to_req_indices.tolist() == [0, 0, 0, 1, 1]
344
+
345
+
346
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA stream coverage")
347
+@pytest.mark.parametrize(
348
+ "use_aux,use_graph", [(False, False), (True, False), (True, True)]
349
+)
350
+@torch.inference_mode()
351
+def test_v41_attention_joins_cache_writes_before_consumption(use_aux, use_graph):
352
+ """Both reused-event joins must publish this forward's states and cache rows."""
353
+ from vllm.forward_context import ForwardContext, override_forward_context
354
+ from vllm.models.deepseek_v4_1.attention import (
355
+ DeepseekV4Attention,
356
+ DeepseekV4Indexer,
357
+ )
358
+ from vllm.models.deepseek_v4_1.compressor import DeepseekCompressor
359
+
360
+ torch.manual_seed(43)
361
+ raw = torch.randn(19, 1024, device="cuda")
362
+ q = torch.zeros(19, 512, dtype=torch.bfloat16, device="cuda")
363
+ positions = torch.arange(7, 26, device="cuda")
364
+ state_slots = positions.clone()
365
+ state_slots[-2:] = -1
366
+ cache_slots = torch.where(state_slots >= 0, positions // 2, -1)
367
+ state = torch.empty(4, 8, 1024, device="cuda")
368
+ main = torch.empty(1, 128, 584, dtype=torch.uint8, device="cuda")
369
+ index = torch.empty(1, 128, 132, dtype=torch.uint8, device="cuda")
370
+ caches = (state, main, index)
371
+ observed = tuple(torch.empty_like(cache) for cache in caches)
372
+ rotary = SimpleNamespace(cos_sin_cache=torch.randn(32, 64, device="cuda"))
373
+ metadata = {
374
+ "state": SimpleNamespace(
375
+ slot_mapping=state_slots,
376
+ query_start_loc=torch.tensor([0, 19], dtype=torch.int32, device="cuda"),
377
+ token_to_req_indices=torch.zeros(19, dtype=torch.int32, device="cuda"),
378
+ ),
379
+ "main": SimpleNamespace(slot_mapping=cache_slots),
380
+ "index": SimpleNamespace(slot_mapping=cache_slots),
381
+ }
382
+ context = ForwardContext({}, metadata, {})
383
+ compressor = DeepseekCompressor.__new__(DeepseekCompressor)
384
+ torch.nn.Module.__init__(compressor)
385
+ compressor.head_dim, compressor.rope_head_dim, compressor.compress_ratio = (
386
+ 512,
387
+ 64,
388
+ 2,
389
+ )
390
+ compressor.rms_norm_eps = 1e-20
391
+ compressor.norm = SimpleNamespace(
392
+ weight=torch.ones(512, dtype=torch.bfloat16, device="cuda")
393
+ )
394
+ compressor.state_cache = SimpleNamespace(prefix="state", kv_cache=state)
395
+ compressor.k_cache_prefix = "main"
396
+ compressor._static_forward_context = {"main": SimpleNamespace(kv_cache=main)}
397
+ indexer_weight = torch.randn(128, 512, dtype=torch.bfloat16, device="cuda")
398
+ indexer = SimpleNamespace(
399
+ owns_k=True,
400
+ wk=lambda latent: (torch.nn.functional.linear(latent, indexer_weight), None),
401
+ k_norm=SimpleNamespace(
402
+ weight=torch.ones(128, dtype=torch.bfloat16, device="cuda"),
403
+ variance_epsilon=1e-20,
404
+ ),
405
+ k_cache=SimpleNamespace(prefix="index", kv_cache=index),
406
+ compress_ratio=2,
407
+ use_fp4_kv=False,
408
+ )
409
+
410
+ def prepare_indexer(qr, latent, weights, positions, rotary, qr_scale):
411
+ DeepseekV4Indexer._produce_k(indexer, latent, positions, rotary)
412
+ return None, None, None
413
+
414
+ def observe(*args):
415
+ for output, cache in zip(observed, caches):
416
+ output.copy_(cache)
417
+
418
+ auxiliary = [torch.cuda.Stream()] if use_aux else None
419
+ attention = SimpleNamespace(
420
+ compressor=compressor,
421
+ indexer=prepare_indexer,
422
+ aux_stream_list=auxiliary,
423
+ ln_events=[torch.cuda.Event(), torch.cuda.Event()],
424
+ rotary_emb=rotary,
425
+ indexer_rotary_emb=rotary,
426
+ n_local_heads=1,
427
+ head_dim=512,
428
+ _wq_b_proj=lambda qr, scale: qr.clone(),
429
+ _fused_qnorm_rope_kv_insert=lambda q, kv, pos, meta: q,
430
+ _sparse_indexer_and_attn=observe,
431
+ )
432
+
433
+ def run():
434
+ with override_forward_context(context):
435
+ DeepseekV4Attention._prepare_and_attn(
436
+ attention, q, q, q, None, raw, q, positions, q
437
+ )
438
+
439
+ def reset():
440
+ state.zero_()
441
+ main.fill_(165)
442
+ index.fill_(165)
443
+
444
+ reset()
445
+ run()
446
+ graph = None
447
+ if use_graph:
448
+ stream = torch.cuda.Stream()
449
+ stream.wait_stream(torch.cuda.current_stream())
450
+ with torch.cuda.stream(stream):
451
+ run()
452
+ torch.cuda.current_stream().wait_stream(stream)
453
+ graph = torch.cuda.CUDAGraph()
454
+ with torch.cuda.graph(graph):
455
+ run()
456
+
457
+ previous = None
458
+ for _ in range(3):
459
+ raw.normal_()
460
+ reset()
461
+ attention.aux_stream_list = None
462
+ run()
463
+ expected = tuple(output.clone() for output in observed)
464
+ reset()
465
+ attention.aux_stream_list = auxiliary
466
+ if graph is None:
467
+ run()
468
+ else:
469
+ graph.replay()
470
+ for actual, reference in zip(observed, expected):
471
+ torch.testing.assert_close(actual, reference, rtol=0, atol=0)
472
+ if previous is not None:
473
+ assert all(not torch.equal(a, b) for a, b in zip(expected, previous))
474
+ previous = expected
475
+
476
+
477
def _on_gfx950() -> bool:
478
if not current_platform.is_rocm():
479
return False
tests/kernels/test_engram.py
new
+778
@@ -0,0 +1,778 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import bisect
4
+from types import SimpleNamespace
5
+
6
+import pytest
7
+import torch
8
+
9
+from vllm.models.deepseek_v4_1.common import engram as engram_ops
10
+from vllm.models.deepseek_v4_1.common.engram import (
11
+ Engram,
12
+ NgramHashState,
13
+ ParallelEngramEmbedding,
14
+)
15
+from vllm.platforms import current_platform
16
+
17
+
18
+def _reference_engram_post_wkv(
19
+ hidden_states,
20
+ kv,
21
+ q_weight,
22
+ k_weight,
23
+ token_mask,
24
+ eps,
25
+ clamp_value,
26
+ kv_start,
27
+):
28
+ num_tokens, hc_mult, dim = hidden_states.shape
29
+ source = torch.arange(num_tokens, device=kv.device) + kv_start
30
+ valid = source < kv.shape[0]
31
+ local_kv = kv.new_zeros(num_tokens, kv.shape[1])
32
+ local_kv[valid] = kv[source[valid]]
33
+
34
+ key = local_kv[:, : hc_mult * dim].float().view(num_tokens, hc_mult, dim)
35
+ value = local_kv[:, hc_mult * dim :].float()
36
+ h = hidden_states.float()
37
+ weight = q_weight.float() * k_weight.float()
38
+ rstd = torch.rsqrt(h.square().mean(-1) + eps) * torch.rsqrt(
39
+ key.square().mean(-1) + eps
40
+ )
41
+ dot = (h * weight * key).sum(-1) * rstd * dim**-0.5
42
+ gate = torch.sigmoid(torch.copysign(dot.abs().clamp_min(clamp_value).sqrt(), dot))
43
+ if token_mask is not None:
44
+ local_mask = torch.zeros(num_tokens, dtype=torch.bool, device=kv.device)
45
+ local_mask[valid] = token_mask[source[valid]]
46
+ gate = gate.masked_fill(~local_mask.unsqueeze(-1), 0)
47
+ return (h + gate.unsqueeze(-1) * value.unsqueeze(-2)).to(hidden_states.dtype)
48
+
49
+
50
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
51
+@pytest.mark.parametrize(
52
+ "num_tokens,num_kv_tokens,hc_mult,dim,tp_size,tp_rank,use_mask",
53
+ [
54
+ (1, 1, 4, 5120, 1, 0, False),
55
+ (3, 5, 4, 96, 2, 0, True),
56
+ (3, 5, 2, 257, 2, 1, False),
57
+ (3, 5, 4, 96, 2, 1, True),
58
+ (1, 1, 4, 96, 4, 3, True),
59
+ (0, 0, 4, 96, 4, 0, False),
60
+ ],
61
+)
62
+def test_fused_engram_post_wkv_matches_reference(
63
+ num_tokens,
64
+ num_kv_tokens,
65
+ hc_mult,
66
+ dim,
67
+ tp_size,
68
+ tp_rank,
69
+ use_mask,
70
+ monkeypatch,
71
+):
72
+ """Preserve gate math, WKV layout, masking, and SP padding semantics."""
73
+ torch.manual_seed(0)
74
+ device = "cuda"
75
+ hidden_states = torch.randn(
76
+ num_tokens, hc_mult, dim, dtype=torch.bfloat16, device=device
77
+ )
78
+ kv = torch.randn(
79
+ num_kv_tokens,
80
+ (hc_mult + 1) * dim,
81
+ dtype=torch.bfloat16,
82
+ device=device,
83
+ )
84
+ q_weight = torch.randn(hc_mult, dim, dtype=torch.bfloat16, device=device)
85
+ k_weight = torch.randn(hc_mult, dim, dtype=torch.bfloat16, device=device)
86
+ token_mask = None
87
+ if use_mask:
88
+ token_mask = torch.arange(num_kv_tokens, device=device) % 2 == 0
89
+ eps = 1e-20
90
+ clamp_value = 1e-6
91
+ use_sequence_parallel = tp_size > 1
92
+ shard_size = (num_kv_tokens + tp_size - 1) // tp_size
93
+ kv_start = tp_rank * shard_size if use_sequence_parallel else 0
94
+
95
+ expected = _reference_engram_post_wkv(
96
+ hidden_states,
97
+ kv,
98
+ q_weight,
99
+ k_weight,
100
+ token_mask,
101
+ eps,
102
+ clamp_value,
103
+ kv_start,
104
+ )
105
+ monkeypatch.setattr(
106
+ engram_ops, "get_tensor_model_parallel_world_size", lambda: tp_size
107
+ )
108
+ monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: tp_rank)
109
+
110
+ module = Engram.__new__(Engram)
111
+ torch.nn.Module.__init__(module)
112
+ module.dim = dim
113
+ module.hc_mult = hc_mult
114
+ module.eps = eps
115
+ module.clamp_value = clamp_value
116
+ module.use_sequence_parallel = use_sequence_parallel
117
+ module.embed_tokens = torch.nn.Identity()
118
+ # `forward` reads rows staged by `prepare_embeddings`, so inject kv there.
119
+ module.embed_tokens.tp_size = 1
120
+ module.staged_rows = kv.unsqueeze(1)
121
+ if use_sequence_parallel:
122
+ padded = torch.nn.functional.pad(kv, (0, 0, 0, (-num_kv_tokens) % tp_size))
123
+ module.staged_rows = padded[
124
+ tp_rank * shard_size : (tp_rank + 1) * shard_size
125
+ ].unsqueeze(1)
126
+ module.wkv = torch.nn.Identity()
127
+
128
+ def check_local_projection(_, args):
129
+ assert args[0].shape[0] == num_tokens
130
+
131
+ module.wkv.register_forward_pre_hook(check_local_projection)
132
+ module.q_weight = torch.nn.Parameter(q_weight, requires_grad=False)
133
+ module.k_weight = torch.nn.Parameter(k_weight, requires_grad=False)
134
+
135
+ actual = module(hidden_states, kv.unsqueeze(1), token_mask)
136
+
137
+ assert actual.shape == hidden_states.shape
138
+ assert actual.dtype == hidden_states.dtype
139
+ torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2)
140
+
141
+
142
+def _hash_state(use_slot_cache: bool) -> NgramHashState:
143
+ state = NgramHashState.__new__(NgramHashState)
144
+ torch.nn.Module.__init__(state)
145
+ state.block_size = 64
146
+ state.use_slot_cache = use_slot_cache
147
+ state._cache = None
148
+ state._kv_cache_ref = None
149
+ state.swa_cache_module = torch.nn.Module()
150
+ state.swa_cache_module.kv_cache = torch.empty(0)
151
+ return state
152
+
153
+
154
+@pytest.mark.parametrize("num_blocks", [4, 100])
155
+def test_engram_cache_follows_kv_cache_binding(num_blocks):
156
+ """Discard profiling history when the real KV cache replaces its storage."""
157
+ state = _hash_state(use_slot_cache=True)
158
+ assert not state.ensure_cache()
159
+
160
+ state.swa_cache_module.kv_cache = torch.empty(4, 64, 1)
161
+ assert state.ensure_cache()
162
+ profiling_cache = state._cache
163
+ profiling_cache.fill_(7)
164
+ assert state.ensure_cache()
165
+ assert state._cache is profiling_cache
166
+ assert torch.all(state._cache == 7)
167
+
168
+ state.swa_cache_module.kv_cache = torch.empty(num_blocks, 64, 1)
169
+ assert state.ensure_cache()
170
+ assert state._cache is not profiling_cache
171
+ assert state._cache.shape == (num_blocks * 64,)
172
+ assert torch.count_nonzero(state._cache) == 0
173
+
174
+ state.swa_cache_module.kv_cache = torch.empty(0)
175
+ assert not state.ensure_cache()
176
+ assert state._cache is None
177
+
178
+
179
+def test_engram_without_slot_cache_only_gates_on_kv_binding():
180
+ """The V2 runner supplies every lookback, so no slot cache is allocated."""
181
+ state = _hash_state(use_slot_cache=False)
182
+ assert not state.ensure_cache()
183
+ state.swa_cache_module.kv_cache = torch.empty(4, 64, 1)
184
+ assert state.ensure_cache()
185
+ assert state._cache is None
186
+
187
+
188
+def _reference_hashes(
189
+ ids,
190
+ positions,
191
+ slots,
192
+ starts,
193
+ tables,
194
+ dead,
195
+ cache,
196
+ token_map,
197
+ multipliers,
198
+ primes,
199
+ offsets,
200
+):
201
+ """Scalar n-gram oracle with persistent, physically addressed token history."""
202
+ for token, slot, is_dead in zip(ids, slots, dead):
203
+ if slot >= 0:
204
+ cache[slot] = -1 if is_dead else token_map[token]
205
+ result = []
206
+ for i, position in enumerate(positions):
207
+ req = min(bisect.bisect_right(starts[1:], i), len(tables) - 1)
208
+ history, blocked = [], False
209
+ for shift in range(4):
210
+ lookback = position - shift
211
+ p = max(0, min(lookback, len(tables[req]) * 16 - 1))
212
+ source = cache[tables[req][p // 16] * 16 + p % 16]
213
+ blocked |= lookback < 0 or source == -1
214
+ history.append(0 if blocked else source)
215
+ layers = []
216
+ for layer, row in enumerate(multipliers):
217
+ rolling, hashes = history[0] * row[0], []
218
+ for j, head_primes in enumerate(primes[layer], start=1):
219
+ rolling ^= history[j] * row[j]
220
+ for head, prime in enumerate(head_primes):
221
+ col = (j - 1) * len(head_primes) + head
222
+ hashes.append(rolling % prime + offsets[layer][col])
223
+ layers.append(hashes)
224
+ result.append(layers)
225
+ return torch.tensor(result, dtype=torch.int32)
226
+
227
+
228
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
229
+@pytest.mark.parametrize("query_len", [1, 17, 257])
230
+@pytest.mark.parametrize("capture", [False, True])
231
+@pytest.mark.parametrize("strided_slots", [False, True])
232
+@pytest.mark.parametrize("large_hashes", [False, True])
233
+def test_engram_hash_cache_replay(query_len, capture, strided_slots, large_hashes):
234
+ """Preserve history, dead tokens and padding as inputs change across replay."""
235
+ num_tokens = 2 * query_len + 3
236
+ blocks = (3 * query_len + 15) // 16
237
+ tables = [list(range(blocks, 2 * blocks)), list(range(blocks))]
238
+ starts = [0, query_len, 2 * query_len]
239
+
240
+ def tensor(data, dtype=torch.int64):
241
+ return torch.tensor(data, dtype=dtype, device="cuda")
242
+
243
+ input_ids = tensor([0] * num_tokens)
244
+ positions = torch.zeros_like(input_ids)
245
+ slots = tensor([-1] * (num_tokens * (2 if strided_slots else 1)))
246
+ if strided_slots:
247
+ slots = slots[::2]
248
+ dead_mask = tensor([False] * num_tokens, torch.bool)
249
+ cache = tensor([-1] * (2 * blocks * 16), torch.int32)
250
+ token_map = [i * 3091 if large_hashes else i % 7 for i in range(32)]
251
+ multipliers = (
252
+ engram_ops.compute_hash_multipliers((1, 14), 4, 99092).tolist()
253
+ if large_hashes
254
+ else [[3, 5, 7, 9], [11, 13, 15, 17]]
255
+ )
256
+ primes = (
257
+ [
258
+ [
259
+ [16000057 + 2 * (layer * 24 + j * 8 + h) for h in range(8)]
260
+ for j in range(3)
261
+ ]
262
+ for layer in range(2)
263
+ ]
264
+ if large_hashes
265
+ else [[[19], [23], [29]]] * 2
266
+ )
267
+ offsets = []
268
+ for layer in primes:
269
+ row, offset = [], 0
270
+ for heads in layer:
271
+ for prime in heads:
272
+ row.append(offset)
273
+ offset += prime
274
+ offsets.append(row)
275
+ state = NgramHashState.__new__(NgramHashState)
276
+ torch.nn.Module.__init__(state)
277
+ state.token_map = tensor(token_map, torch.int32)
278
+ state.multipliers = tensor(multipliers)
279
+ state.primes = tensor(primes)
280
+ state.offsets = tensor(offsets)
281
+ state._cache = cache
282
+ state.use_slot_cache = True
283
+ state.pad_id = 0
284
+ state.block_size = 16
285
+ lookback = tensor([[-1] * 3] * 2, torch.int32)
286
+ args = (
287
+ input_ids,
288
+ positions,
289
+ tensor(starts, torch.int32),
290
+ dead_mask,
291
+ lookback,
292
+ torch.zeros_like(lookback, dtype=torch.bool),
293
+ slots,
294
+ tensor(tables, torch.int32),
295
+ )
296
+ op = state
297
+ graph = None
298
+ if capture:
299
+ stream = torch.cuda.Stream()
300
+ stream.wait_stream(torch.cuda.current_stream())
301
+ with torch.cuda.stream(stream):
302
+ for _ in range(3):
303
+ op(*args)
304
+ torch.cuda.current_stream().wait_stream(stream)
305
+ graph = torch.cuda.CUDAGraph()
306
+ with torch.cuda.graph(graph, stream=stream):
307
+ output = op(*args)
308
+ torch.cuda.current_stream().wait_stream(stream)
309
+ cache.fill_(-1)
310
+ expected_cache = [-1] * cache.numel()
311
+
312
+ for step in range(3):
313
+ ids = [(i + 3 * step) % 32 for i in range(num_tokens)]
314
+ pos = list(range(step * query_len, (step + 1) * query_len)) * 2
315
+ physical_slots = [
316
+ tables[i // query_len][p // 16] * 16 + p % 16 for i, p in enumerate(pos)
317
+ ]
318
+ pos += [0] * 3
319
+ physical_slots += [-1] * 3
320
+ # A fully padded replay must leave every cache slot untouched.
321
+ if step == 2:
322
+ physical_slots = [-1] * num_tokens
323
+ dead = [(i + step) % 5 == 0 for i in range(num_tokens)]
324
+ input_ids.copy_(tensor(ids))
325
+ positions.copy_(tensor(pos))
326
+ slots.copy_(tensor(physical_slots))
327
+ dead_mask.copy_(tensor(dead, torch.bool))
328
+ if graph is None:
329
+ output = op(*args)
330
+ else:
331
+ graph.replay()
332
+ expected = _reference_hashes(
333
+ ids,
334
+ pos,
335
+ physical_slots,
336
+ starts,
337
+ tables,
338
+ dead,
339
+ expected_cache,
340
+ token_map,
341
+ multipliers,
342
+ primes,
343
+ offsets,
344
+ )
345
+ if step != 2:
346
+ real = 2 * query_len
347
+ torch.testing.assert_close(
348
+ output[:real].cpu(), expected[:real], rtol=0, atol=0
349
+ )
350
+ assert cache.cpu().tolist() == expected_cache
351
+
352
+ empty_output = state(
353
+ input_ids[:0], positions[:0], args[2], dead_mask[:0], *args[4:]
354
+ )
355
+ assert empty_output.shape == (0, 2, len(offsets[0]))
356
+ assert cache.cpu().tolist() == expected_cache
357
+
358
+
359
+def _hash_ids(requests, tables_meta, block_table, cache=None, capture=False):
360
+ """Run the op on a batch of (token_ids, start, window) requests; `window`
361
+ is the runner's [depth] lookback for that request (None = all unknown).
362
+ Token IDs 14 and 22 stand in for image tokens that break n-grams.
363
+ Returns one [len(token_ids), ...] hash tensor per request."""
364
+ token_map, multipliers, primes, offsets, block_size = tables_meta
365
+ depth = multipliers.shape[1] - 1
366
+
367
+ def tensor(data, dtype=torch.int64):
368
+ return torch.tensor(data, dtype=dtype, device="cuda")
369
+
370
+ input_ids, positions, slots, windows, starts = [], [], [], [], [0]
371
+ for req, (token_ids, start, window) in enumerate(requests):
372
+ pos = list(range(start, start + len(token_ids)))
373
+ input_ids += token_ids
374
+ positions += pos
375
+ slots += [
376
+ int(block_table[req, p // block_size]) * block_size + p % block_size
377
+ for p in pos
378
+ ]
379
+ windows.append(window or [-1] * depth)
380
+ starts.append(starts[-1] + len(token_ids))
381
+ window = tensor(windows, torch.int32)
382
+ state = _hash_state(use_slot_cache=cache is not None)
383
+ state.token_map = token_map
384
+ state.multipliers = multipliers
385
+ state.primes = primes
386
+ state.offsets = offsets
387
+ state.pad_id = 0
388
+ state.block_size = block_size
389
+ state._cache = cache
390
+ input_tensor = tensor(input_ids, torch.int32)
391
+ args = (
392
+ input_tensor,
393
+ tensor(positions),
394
+ tensor(starts, torch.int32),
395
+ (input_tensor == 14) | (input_tensor == 22),
396
+ window,
397
+ (window == 14) | (window == 22),
398
+ tensor(slots) if cache is not None else None,
399
+ block_table if cache is not None else None,
400
+ )
401
+ out = state(*args)
402
+ if capture:
403
+ graph = torch.cuda.CUDAGraph()
404
+ with torch.cuda.graph(graph):
405
+ out = state(*args)
406
+ graph.replay()
407
+ return [out[a:b] for a, b in zip(starts, starts[1:])]
408
+
409
+
410
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
411
+@pytest.mark.parametrize("runner", ["v2_window_only", "v1_prompt_window_plus_cache"])
412
+@pytest.mark.parametrize("capture", [False, True])
413
+def test_lookback_window_reproduces_single_instance(runner, capture):
414
+ """Decoding on a fresh instance whose prompt KV came from elsewhere (P/D,
415
+ offload) must hash like the instance that processed the whole sequence,
416
+ with two requests batched and chunks of 1-2 tokens so in-batch, window
417
+ and (V1) cache lookbacks mix within one step.
418
+
419
+ V2 supplies every lookback from its device token history and needs no
420
+ slot cache. V1 supplies prompt positions only and reads generated
421
+ positions from the slot cache it fills itself. Without any window the
422
+ first decode token hashes stale slots (negative control)."""
423
+ block_size = 4
424
+ block_table = torch.tensor(
425
+ [[3, 1, 5, 0], [2, 4, 6, 7]], dtype=torch.int32, device="cuda"
426
+ )
427
+ tables_meta = (
428
+ torch.arange(50, dtype=torch.int32, device="cuda"),
429
+ torch.tensor([[3, 5, 7, 9]], device="cuda"),
430
+ torch.tensor([[[97, 89], [83, 79], [73, 71]]], device="cuda"),
431
+ torch.tensor([[0, 97, 186, 269, 348, 421]], device="cuda"),
432
+ block_size,
433
+ )
434
+ depth = 3
435
+ prompts = [[11, 12, 13, 14, 15, 16], [31, 32, 33, 34, 35]]
436
+ decodes = [[21, 22, 23, 24], [41, 42, 43, 44]]
437
+ chunk_sizes = [1, 2, 1]
438
+ v2 = runner.startswith("v2")
439
+
440
+ def fresh_cache():
441
+ return torch.zeros(8 * block_size, dtype=torch.int32, device="cuda")
442
+
443
+ reference = _hash_ids(
444
+ [(p + d, 0, None) for p, d in zip(prompts, decodes)],
445
+ tables_meta,
446
+ block_table,
447
+ cache=fresh_cache(),
448
+ )
449
+
450
+ histories = [list(p) for p in prompts]
451
+ cache = None if v2 else fresh_cache()
452
+ consumed = 0
453
+ for chunk in chunk_sizes:
454
+ batch = []
455
+ for history, prompt, decode in zip(histories, prompts, decodes):
456
+ start = len(history)
457
+ tokens = decode[consumed : consumed + chunk]
458
+ window = [
459
+ history[start - 1 - j] if v2 or start - 1 - j < len(prompt) else -1
460
+ for j in range(depth)
461
+ ]
462
+ batch.append((tokens, start, window))
463
+ got = _hash_ids(batch, tables_meta, block_table, cache=cache, capture=capture)
464
+ for req, (tokens, start, _) in enumerate(batch):
465
+ assert torch.equal(got[req], reference[req][start : start + chunk]), (
466
+ f"{runner} request {req} chunk at {start}"
467
+ )
468
+ histories[req] += tokens
469
+ consumed += chunk
470
+
471
+ unseeded = _hash_ids(
472
+ [(decodes[0][:1], len(prompts[0]), None)],
473
+ tables_meta,
474
+ block_table,
475
+ cache=fresh_cache(),
476
+ )
477
+ assert not torch.equal(unseeded[0][0], reference[0][len(prompts[0])])
478
+
479
+
480
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
481
+def test_v2_model_state_gathers_lookback_window():
482
+ """The window is gathered on device from the runner's token history."""
483
+ from vllm.models.deepseek_v4_1.nvidia.model_state import DeepseekV41ModelState
484
+
485
+ depth, max_num_reqs, max_model_len = 3, 4, 16
486
+ state = DeepseekV41ModelState.__new__(DeepseekV41ModelState)
487
+ state.rope_state = None
488
+ state.supports_mm_inputs = False
489
+ state.prompt_embeds_state = None
490
+ state.lookback_token_ids = torch.full(
491
+ (max_num_reqs, depth), -1, dtype=torch.int32, device="cuda"
492
+ )
493
+
494
+ all_token_ids = torch.zeros(max_num_reqs, max_model_len, dtype=torch.int32)
495
+ all_token_ids[2, :8] = torch.tensor([11, 12, 13, 14, 15, 16, 21, 22])
496
+ all_token_ids[0, :5] = torch.arange(1, 6)
497
+ req_states = SimpleNamespace(
498
+ num_computed_tokens=SimpleNamespace(
499
+ gpu=torch.tensor([0, 9, 8, 6], dtype=torch.int32, device="cuda")
500
+ ),
501
+ all_token_ids=SimpleNamespace(gpu=all_token_ids.cuda()),
502
+ )
503
+ # Batch rows -> request state rows: a request two decode steps in, a fresh
504
+ # request, and one at the end of a 5-token prompt.
505
+ input_batch = SimpleNamespace(
506
+ idx_mapping=torch.tensor([2, 0, 3], dtype=torch.int64, device="cuda")
507
+ )
508
+ window = state.prepare_inputs(input_batch, req_states)["lookback_token_ids"]
509
+ assert window.cpu().tolist() == [
510
+ [22, 21, 16],
511
+ [-1, -1, -1],
512
+ [0, 0, 0],
513
+ [-1, -1, -1],
514
+ ]
515
+ # Graph capture runs on the dummy inputs and replays read the same buffer.
516
+ dummy = state.prepare_dummy_inputs(num_reqs=3, num_tokens=3)["lookback_token_ids"]
517
+ assert dummy is window and torch.all(dummy == -1)
518
+
519
+
520
+def _reference_lookup(weight, scale_inv, ids, start, end, block=32):
521
+ """The torch expression the fused kernel replaces."""
522
+ mask = (ids < start) | (ids >= end)
523
+ local = (ids - start).masked_fill(mask, 0).long()
524
+ values = torch.nn.functional.embedding(local, weight)
525
+ scales = torch.nn.functional.embedding(local, scale_inv)
526
+ scales = (scales.to(torch.int32) << 23).view(torch.float32)
527
+ values = values.float().unflatten(-1, (-1, block))
528
+ values = (values * scales.unsqueeze(-1)).flatten(-2).to(torch.bfloat16)
529
+ return values.masked_fill(mask.unsqueeze(-1), 0)
530
+
531
+
532
+def _make_embedding(cpu_offload, rows=4096, dim=256, block=32):
533
+ layer = ParallelEngramEmbedding.__new__(ParallelEngramEmbedding)
534
+ torch.nn.Module.__init__(layer)
535
+ layer.dim, layer.block_size, layer.tp_size = dim, block, 1
536
+ layer.n_hash_cols = layer.part_n_hash_cols = 24
537
+ layer.head_start = 0
538
+ layer.cpu_offload = cpu_offload
539
+ layer.part_num_embeddings = rows
540
+ # A window strictly inside the table, so unowned rows are exercised too.
541
+ layer.vocab_start_idx, layer.vocab_end_idx = rows // 4, rows // 4 + rows // 2
542
+ layer._views = layer._view_src = None
543
+ layer._num_sms = torch.cuda.get_device_properties(
544
+ torch.accelerator.current_device_index()
545
+ ).multi_processor_count
546
+ kwargs = (
547
+ {"device": "cpu", "pin_memory": True} if cpu_offload else {"device": "cuda"}
548
+ )
549
+ owned = layer.vocab_end_idx - layer.vocab_start_idx
550
+ layer.weight = torch.empty(owned, dim, dtype=torch.float8_e4m3fn, **kwargs)
551
+ layer.weight_scale_inv = torch.empty(
552
+ owned, dim // block, dtype=torch.uint8, **kwargs
553
+ )
554
+ torch.manual_seed(0)
555
+ layer.weight.copy_((torch.randn(owned, dim) * 4).to(torch.float8_e4m3fn))
556
+ layer.weight_scale_inv.copy_(
557
+ torch.randint(120, 134, (owned, dim // block), dtype=torch.uint8)
558
+ )
559
+ return layer
560
+
561
+
562
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
563
+@pytest.mark.parametrize("cpu_offload", [False, True])
564
+@pytest.mark.parametrize("tp_size", [1, 2, 4, 8])
565
+def test_engram_head_shards_reconstruct_checkpoint(cpu_offload, tp_size, monkeypatch):
566
+ """Keep complete buckets and reconstruct head order, including TP padding."""
567
+ head_sizes = (17, 19, 23, 29, 31, 37)
568
+ num_rows, dim = sum(head_sizes), 64
569
+ torch.manual_seed(0)
570
+ weight = torch.randn(num_rows + 7, dim).to(torch.float8_e4m3fn)
571
+ scales = torch.randint(120, 134, (num_rows + 7, dim // 32), dtype=torch.uint8)
572
+ hashes = torch.empty(7, 2, len(head_sizes), dtype=torch.int32, device="cuda")
573
+ start = 0
574
+ for head, size in enumerate(head_sizes):
575
+ hashes[:, :, head].random_(start, start + size)
576
+ hashes[0, :, head] = start
577
+ hashes[-1, :, head] = start + size - 1
578
+ start += size
579
+ ids = hashes[:, 1]
580
+ expected = _reference_lookup(weight.cuda(), scales.cuda(), ids, 0, num_rows)
581
+ monkeypatch.setattr(
582
+ engram_ops, "get_tensor_model_parallel_world_size", lambda: tp_size
583
+ )
584
+ shards = []
585
+ layers = []
586
+ for rank in range(tp_size):
587
+ monkeypatch.setattr(
588
+ engram_ops, "get_tensor_model_parallel_rank", lambda rank=rank: rank
589
+ )
590
+ with torch.device("cuda"):
591
+ layer = ParallelEngramEmbedding(
592
+ num_rows + 7, dim, head_sizes, cpu_offload=cpu_offload
593
+ )
594
+ layer.weight.weight_loader(layer.weight, weight)
595
+ layer.weight_scale_inv.weight_loader(
596
+ layer.weight_scale_inv, scales.view(torch.float8_e8m0fnu)
597
+ )
598
+ out = torch.empty(
599
+ len(ids),
600
+ layer.part_n_hash_cols,
601
+ dim,
602
+ device="cuda",
603
+ dtype=torch.bfloat16,
604
+ )
605
+ layer.lookup(ids, out)
606
+ shards.append(out)
607
+ layers.append(layer)
608
+ gathered = torch.cat(shards, dim=1)
609
+ torch.testing.assert_close(gathered[:, : len(head_sizes)], expected, rtol=0, atol=0)
610
+ assert torch.count_nonzero(gathered[:, len(head_sizes) :]) == 0
611
+ assert sum(layer.part_num_embeddings for layer in layers) == num_rows
612
+
613
+ def gather(local, dim):
614
+ torch.testing.assert_close(local, shards[0], rtol=0, atol=0)
615
+ return torch.cat(shards, dim=dim)
616
+
617
+ monkeypatch.setattr(engram_ops, "tensor_model_parallel_all_gather", gather)
618
+ torch.testing.assert_close(layers[0](ids), expected, rtol=0, atol=0)
619
+ module = Engram.__new__(Engram)
620
+ torch.nn.Module.__init__(module)
621
+ module.embed_tokens = layers[0]
622
+ module.use_sequence_parallel = False
623
+ module.staged_rows = torch.empty_like(shards[0])
624
+ module.prepare_embeddings(ids)
625
+ torch.testing.assert_close(module.embed(ids), expected, rtol=0, atol=0)
626
+ # Slicing before head reordering must preserve padded heads and empty owners.
627
+ module.use_sequence_parallel = True
628
+ chunk = (len(ids) + tp_size - 1) // tp_size
629
+ padded = torch.nn.functional.pad(expected, (0, 0, 0, 0, 0, (-len(ids)) % tp_size))
630
+ for rank in range(tp_size):
631
+ monkeypatch.setattr(
632
+ engram_ops, "get_tensor_model_parallel_rank", lambda rank=rank: rank
633
+ )
634
+ torch.testing.assert_close(
635
+ module.embed(ids), padded[rank * chunk : (rank + 1) * chunk], rtol=0, atol=0
636
+ )
637
+
638
+
639
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
640
+@pytest.mark.parametrize("cpu_offload", [False, True])
641
+@pytest.mark.parametrize("background", [False, True])
642
+@pytest.mark.parametrize("num_tokens", [1, 7, 256])
643
+def test_engram_lookup_matches_torch(cpu_offload, background, num_tokens):
644
+ """The fused gather must be bit-exact with the torch dequant path it
645
+ replaces, from HBM and from pinned host memory alike, and must contribute
646
+ zeros for rows another TP rank owns."""
647
+ layer = _make_embedding(cpu_offload)
648
+ cols, rows = 24, layer.part_num_embeddings
649
+ ids = torch.randint(0, rows, (num_tokens, cols), dtype=torch.int32, device="cuda")
650
+ expected = _reference_lookup(
651
+ layer.weight.cuda(),
652
+ layer.weight_scale_inv.cuda(),
653
+ ids,
654
+ layer.vocab_start_idx,
655
+ layer.vocab_end_idx,
656
+ )
657
+ out = torch.empty(num_tokens, cols, layer.dim, dtype=torch.bfloat16, device="cuda")
658
+ layer.lookup(ids, out, background=background)
659
+ assert torch.equal(out, expected)
660
+
661
+
662
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
663
+@pytest.mark.parametrize("cpu_offload", [False, True])
664
+@pytest.mark.parametrize("capture", ["eager", "full", "breakable"])
665
+def test_engram_prepared_rows_survive_graph_breaks(cpu_offload, capture):
666
+ """Consume early lookup results after a break, with fresh IDs each replay."""
667
+ _run_engram_prepared_rows(cpu_offload, capture)
668
+
669
+
670
+def _run_engram_prepared_rows(
671
+ cpu_offload, capture, tp_size=1, rank=0, use_sequence_parallel=False
672
+):
673
+ from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
674
+
675
+ layer = _make_embedding(cpu_offload)
676
+ cols, num_tokens = (23, 65) if use_sequence_parallel else (24, 64)
677
+ layer.n_hash_cols = cols
678
+ layer.tp_size = tp_size
679
+ layer.part_n_hash_cols = (cols + tp_size - 1) // tp_size
680
+ layer.head_start = rank * layer.part_n_hash_cols
681
+ engram = Engram.__new__(Engram)
682
+ torch.nn.Module.__init__(engram)
683
+ engram.embed_tokens = layer
684
+ engram.use_sequence_parallel = use_sequence_parallel
685
+ engram.staged_rows = torch.empty(
686
+ num_tokens,
687
+ layer.part_n_hash_cols,
688
+ layer.dim,
689
+ dtype=torch.bfloat16,
690
+ device="cuda",
691
+ )
692
+ # Match the non-contiguous per-layer slice of the model hash tensor.
693
+ hashes = torch.randint(
694
+ 0,
695
+ layer.part_num_embeddings,
696
+ (num_tokens, 2, cols),
697
+ dtype=torch.int32,
698
+ device="cuda",
699
+ )
700
+ src = hashes[:, 1]
701
+ local_tokens = (
702
+ (num_tokens + tp_size - 1) // tp_size if use_sequence_parallel else num_tokens
703
+ )
704
+ out = torch.empty(
705
+ local_tokens, cols, layer.dim, dtype=torch.bfloat16, device="cuda"
706
+ )
707
+ embed = engram.embed
708
+ if capture == "compiled":
709
+ embed = torch.compile(embed, backend="eager", fullgraph=True, dynamic=True)
710
+
711
+ def step(cap=None):
712
+ engram.prepare_embeddings(src)
713
+ if cap is not None:
714
+ cap.add_eager(lambda: None)
715
+ out.copy_(embed(src))
716
+
717
+ warmup = torch.cuda.Stream()
718
+ warmup.wait_stream(torch.cuda.current_stream())
719
+ with torch.cuda.stream(warmup):
720
+ step()
721
+ torch.cuda.current_stream().wait_stream(warmup)
722
+
723
+ graph = None
724
+ if capture == "full":
725
+ graph = torch.cuda.CUDAGraph()
726
+ with torch.cuda.graph(graph):
727
+ step()
728
+ elif capture == "breakable":
729
+ graph = BreakableCUDAGraphCapture()
730
+ with torch.cuda.stream(warmup), graph:
731
+ step(graph)
732
+ torch.cuda.current_stream().wait_stream(warmup)
733
+ assert graph.num_graphs == 2
734
+
735
+ for _ in range(3):
736
+ hashes.random_(0, layer.part_num_embeddings)
737
+ if graph is None:
738
+ step()
739
+ else:
740
+ graph.replay()
741
+ expected = _reference_lookup(
742
+ layer.weight.cuda(),
743
+ layer.weight_scale_inv.cuda(),
744
+ src,
745
+ layer.vocab_start_idx,
746
+ layer.vocab_end_idx,
747
+ )
748
+ if use_sequence_parallel:
749
+ expected = torch.nn.functional.pad(
750
+ expected, (0, 0, 0, 0, 0, (-num_tokens) % tp_size)
751
+ )[rank * local_tokens : (rank + 1) * local_tokens]
752
+ torch.testing.assert_close(out, expected, rtol=0, atol=0)
753
+
754
+
755
+def _engram_tp_worker(rank, tp_size, port):
756
+ from tests.utils import init_test_distributed_environment
757
+ from vllm.distributed import cleanup_dist_env_and_memory
758
+
759
+ torch.accelerator.set_device_index(rank)
760
+ init_test_distributed_environment(tp_size, 1, rank, str(port), local_rank=rank)
761
+ try:
762
+ for cpu_offload in (False, True):
763
+ for sp in (False, True):
764
+ for capture in ("eager", "compiled", "full", "breakable"):
765
+ _run_engram_prepared_rows(cpu_offload, capture, tp_size, rank, sp)
766
+ finally:
767
+ cleanup_dist_env_and_memory()
768
+
769
+
770
+@pytest.mark.distributed(num_gpus=2)
771
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
772
+def test_engram_head_collectives_survive_graph_breaks():
773
+ """All-gather preserves head order and local SP tokens across graph replay."""
774
+ from vllm.utils.network_utils import get_open_port
775
+
776
+ if torch.accelerator.device_count() < 2:
777
+ pytest.skip("Requires two GPUs")
778
+ torch.multiprocessing.spawn(_engram_tp_worker, args=(2, get_open_port()), nprocs=2)
tests/kernels/test_fused_indexer_q_rope_quant.py
+359
-9
@@ -14,6 +14,7 @@ Expects bit-exact equality on both q_fp8 and weights_out.
14
"""
15
16
import contextlib
17
+import importlib
18
from unittest import mock
19
20
import pytest
@@ -29,9 +30,10 @@ from vllm.utils.import_utils import has_cutedsl
30
31
HEAD_DIM = 128
32
ROPE_DIM = 64
32
-N_HEAD = 64
33
MAX_POS = 4096
34
35
+CUTEDSL_MODULE = "vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl"
36
+
37
38
def quantize_to_mxfp4(
39
x: torch.Tensor,
@@ -90,6 +92,7 @@ def _reference(
92
weights: torch.Tensor,
93
softmax_scale: float,
94
head_scale: float,
95
+ n_head: int,
96
use_fp4: bool = False,
97
) -> tuple[torch.Tensor, torch.Tensor]:
98
q_rot = q.clone()
@@ -105,7 +108,7 @@ def _reference(
108
)
109
110
if use_fp4:
108
- q_packed, ue8m0 = quantize_to_mxfp4(q_rot.view(-1, N_HEAD, HEAD_DIM))
111
+ q_packed, ue8m0 = quantize_to_mxfp4(q_rot.view(-1, n_head, HEAD_DIM))
112
# Pack 4 ue8m0 bytes into 1 int32
113
q_scale = ue8m0.view(torch.int32).squeeze(-1)
114
# FP4 path: q_scale stays separate (cannot be folded into a per-token scalar)
@@ -118,8 +121,8 @@ def _reference(
121
HEAD_DIM,
122
use_ue8m0=True,
123
)
121
- q_fp8 = q_fp8.view(-1, N_HEAD, HEAD_DIM)
122
- q_scale = q_scale.view(-1, N_HEAD)
124
+ q_fp8 = q_fp8.view(-1, n_head, HEAD_DIM)
125
+ q_scale = q_scale.view(-1, n_head)
126
127
weights_out = weights.to(torch.float32) * q_scale * softmax_scale * head_scale
128
return q_fp8, weights_out
@@ -144,9 +147,10 @@ def _reference(
147
],
148
)
149
@pytest.mark.parametrize("use_cutedsl", [False, True])
150
+@pytest.mark.parametrize("n_head", [32, 64])
151
@torch.inference_mode()
152
def test_fused_indexer_q_rope_quant_matches_unfused(
149
- num_tokens, cache_dtype, use_fp4, use_cutedsl
153
+ num_tokens, cache_dtype, use_fp4, use_cutedsl, n_head
154
):
155
if use_cutedsl and not has_cutedsl():
156
pytest.skip("cutedsl (cutlass) not installed")
@@ -154,17 +158,29 @@ def test_fused_indexer_q_rope_quant_matches_unfused(
158
device = "cuda"
159
torch.manual_seed(0)
160
157
- q = torch.randn(num_tokens, N_HEAD, HEAD_DIM, dtype=torch.bfloat16, device=device)
161
+ q = torch.randn(num_tokens, n_head, HEAD_DIM, dtype=torch.bfloat16, device=device)
162
positions = torch.randint(
163
0, MAX_POS, (num_tokens,), dtype=torch.int64, device=device
164
)
165
cos_sin_cache = torch.randn(MAX_POS, ROPE_DIM, dtype=cache_dtype, device=device)
162
- weights = torch.randn(num_tokens, N_HEAD, dtype=torch.bfloat16, device=device)
166
+ weights = torch.randn(num_tokens, n_head, dtype=torch.bfloat16, device=device)
167
softmax_scale = HEAD_DIM**-0.5
164
- head_scale = N_HEAD**-0.5
168
+ # head_scale must be an exact power of two for bit-exactness to be a fair
169
+ # expectation: the triton kernel multiplies by softmax_scale and head_scale
170
+ # separately (two roundings) while the cutedsl kernels fold them into one
171
+ # fp32 scalar first (one rounding). Those agree only when one factor is a
172
+ # pow2. n_head**-0.5 is a pow2 at 64 heads but not at 32.
173
+ head_scale = 0.125
174
175
q_quant_ref, weights_ref = _reference(
167
- positions, q, cos_sin_cache, weights, softmax_scale, head_scale, use_fp4
176
+ positions,
177
+ q,
178
+ cos_sin_cache,
179
+ weights,
180
+ softmax_scale,
181
+ head_scale,
182
+ n_head,
183
+ use_fp4,
184
)
185
# use_cutedsl=False: force the triton path even when cutedsl is installed
186
# by patching the dispatcher's has_cutedsl() binding to return False.
@@ -210,3 +226,337 @@ def test_fused_indexer_q_rope_quant_matches_unfused(
226
f"weights mismatch: max abs diff "
227
f"{(weights_ref - weights_fused).abs().max().item()}"
228
)
229
+
230
+
231
+@pytest.mark.skipif(not has_cutedsl(), reason="cutedsl (cutlass) not installed")
232
+@pytest.mark.parametrize(
233
+ "use_fp4",
234
+ [
235
+ False,
236
+ pytest.param(
237
+ True,
238
+ marks=pytest.mark.skipif(
239
+ not (
240
+ current_platform.is_cuda()
241
+ and current_platform.is_device_capability_family(100)
242
+ ),
243
+ reason="MXFP4 indexer cache requires an SM100-family GPU",
244
+ ),
245
+ ),
246
+ ],
247
+)
248
+@pytest.mark.parametrize("n_head", [32, 64])
249
+@torch.inference_mode()
250
+def test_cutedsl_indexer_q_writes_stay_within_num_tokens(use_fp4, n_head):
251
+ """The cutedsl kernels must not store past ``num_tokens``.
252
+
253
+ The launcher rounds ``num_tokens * threads_per_token`` up to whole 128-thread
254
+ blocks, so the trailing block is only partially in range unless
255
+ ``threads_per_token`` is itself a multiple of the block size. With 32 heads
256
+ and coarsen=4 (picked for >= 512 tokens) a token needs 64 threads, so an odd
257
+ token count leaves half of the last block addressing token ``num_tokens``.
258
+ Every output gets one extra sentinel row that the kernel must leave alone.
259
+ """
260
+ mod = importlib.import_module(CUTEDSL_MODULE)
261
+ device = "cuda"
262
+ torch.manual_seed(0)
263
+ num_tokens = 935
264
+
265
+ sentinel_u8, sentinel_f32 = 0x55, -12345.0
266
+ # positions/q are padded too, so an out-of-bounds *load* stays addressable
267
+ # and the assertions below are about stores only.
268
+ positions = torch.zeros(num_tokens + 1, dtype=torch.int64, device=device)
269
+ q = torch.randn(
270
+ num_tokens + 1, n_head, HEAD_DIM, dtype=torch.bfloat16, device=device
271
+ )
272
+ cos_sin_cache = torch.randn(MAX_POS, ROPE_DIM, dtype=torch.float32, device=device)
273
+ weights = torch.randn(num_tokens + 1, n_head, dtype=torch.bfloat16, device=device)
274
+ weights_out = torch.full(
275
+ (num_tokens + 1, n_head), sentinel_f32, dtype=torch.float32, device=device
276
+ )
277
+ outputs = {"weights_out": weights_out}
278
+
279
+ if use_fp4:
280
+ outputs["q_packed"] = torch.full(
281
+ (num_tokens + 1, n_head, HEAD_DIM // 2),
282
+ sentinel_u8,
283
+ dtype=torch.uint8,
284
+ device=device,
285
+ )
286
+ outputs["q_scale"] = torch.full(
287
+ (num_tokens + 1, n_head, HEAD_DIM // 32),
288
+ sentinel_u8,
289
+ dtype=torch.uint8,
290
+ device=device,
291
+ )
292
+ mod.fused_indexer_q_rope_quant_mxfp4_cutedsl(
293
+ positions[:num_tokens],
294
+ q[:num_tokens],
295
+ cos_sin_cache,
296
+ weights[:num_tokens],
297
+ 1.0,
298
+ 1.0,
299
+ outputs["q_packed"][:num_tokens],
300
+ outputs["q_scale"][:num_tokens],
301
+ weights_out[:num_tokens],
302
+ )
303
+ else:
304
+ outputs["q_fp8"] = torch.full(
305
+ (num_tokens + 1, n_head, HEAD_DIM),
306
+ sentinel_u8,
307
+ dtype=torch.uint8,
308
+ device=device,
309
+ )
310
+ mod.fused_indexer_q_rope_quant_fp8_cutedsl(
311
+ positions[:num_tokens],
312
+ q[:num_tokens],
313
+ cos_sin_cache,
314
+ weights[:num_tokens],
315
+ 1.0,
316
+ 1.0,
317
+ outputs["q_fp8"][:num_tokens].view(torch.float8_e4m3fn),
318
+ weights_out[:num_tokens],
319
+ )
320
+ torch.accelerator.synchronize()
321
+
322
+ clobbered = {}
323
+ for name, tensor in outputs.items():
324
+ sentinel = sentinel_f32 if tensor.dtype == torch.float32 else sentinel_u8
325
+ n_changed = (tensor[num_tokens] != sentinel).sum().item()
326
+ if n_changed:
327
+ clobbered[name] = n_changed
328
+ assert not clobbered, (
329
+ f"kernel wrote past token {num_tokens - 1} into the sentinel row: {clobbered}"
330
+ )
331
+
332
+
333
+def _indexer_k_reference(
334
+ k_pre, positions, cos_sin_cache, rms_norm_weight, compress_ratio, use_fp4
335
+):
336
+ k = k_pre.float()
337
+ k = k * torch.rsqrt(k.square().mean(dim=-1, keepdim=True) + 1e-20)
338
+ k = (k * rms_norm_weight.float()).to(torch.bfloat16).float()
339
+ group_positions = positions // compress_ratio * compress_ratio
340
+ cos, sin = cos_sin_cache[group_positions].float().chunk(2, dim=-1)
341
+ even, odd = k[:, 64::2], k[:, 65::2]
342
+ rotated = torch.stack((even * cos - odd * sin, odd * cos + even * sin), dim=-1)
343
+ k = torch.cat((k[:, :64], rotated.flatten(1)), dim=-1).to(torch.bfloat16)
344
+ if use_fp4:
345
+ return quantize_to_mxfp4(k)
346
+ exponent = k.float().abs().amax(dim=-1, keepdim=True).clamp(min=1e-4) / 448
347
+ scale = exponent.log2().ceil().exp2()
348
+ values = (k.float() / scale).clamp(-448, 448).to(torch.float8_e4m3fn)
349
+ return values.view(torch.uint8), scale.view(torch.uint8)
350
+
351
+
352
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA kernel")
353
+@pytest.mark.parametrize("num_tokens", [1, 17, 257])
354
+@pytest.mark.parametrize("compress_ratio", [1, 2])
355
+@pytest.mark.parametrize("use_fp4", [False, True])
356
+@pytest.mark.parametrize("cache_dtype", [torch.float32, torch.bfloat16])
357
+@torch.inference_mode()
358
+def test_indexer_k_inserts_only_valid_groups_into_padded_pages(
359
+ num_tokens, compress_ratio, use_fp4, cache_dtype
360
+):
361
+ """Insert group ends and preserve skipped slots, graph rows, and page padding."""
362
+ from vllm.models.deepseek_v4_1.common.ops.indexer_k_store import (
363
+ indexer_k_norm_rope_store,
364
+ )
365
+
366
+ if use_fp4 and not current_platform.is_device_capability_family(100):
367
+ pytest.skip("MXFP4 indexer cache requires an SM100-family GPU")
368
+ torch.manual_seed(17)
369
+ device = "cuda"
370
+ k_pre = torch.randn(num_tokens + 3, 128, device=device, dtype=torch.bfloat16)
371
+ norm_weight = torch.randn(128, device=device, dtype=torch.bfloat16)
372
+ positions = torch.arange(7, num_tokens + 10, device=device)
373
+ slots = torch.randperm(2 * num_tokens, device=device)[:num_tokens].clone()
374
+ slots[5::7] = -1
375
+ valid = (slots >= 0) & ((positions[:num_tokens] + 1) % compress_ratio == 0)
376
+ k_pre[3::7] = 0
377
+ k_pre[8::11] *= 1e-10
378
+ k_pre[:num_tokens][~valid] = float("nan")
379
+ k_pre[num_tokens:] = float("nan")
380
+ positions[num_tokens:] = MAX_POS + 10
381
+ angles = torch.randn(MAX_POS, 32, device=device)
382
+ cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1).to(cache_dtype)
383
+ block_size = 16
384
+ value_bytes = 64 if use_fp4 else 128
385
+ page_bytes = block_size * (value_bytes + 4)
386
+ num_blocks = (2 * num_tokens + block_size - 1) // block_size + 1
387
+ storage = torch.full(
388
+ (num_blocks, page_bytes + 128), 0xA5, device=device, dtype=torch.uint8
389
+ )
390
+ cache = storage.as_strided(
391
+ (num_blocks, block_size, value_bytes + 4),
392
+ (storage.stride(0), value_bytes + 4, 1),
393
+ )
394
+ expected = storage.clone()
395
+ values, scales = _indexer_k_reference(
396
+ k_pre[:num_tokens][valid],
397
+ positions[:num_tokens][valid],
398
+ cos_sin,
399
+ norm_weight,
400
+ compress_ratio,
401
+ use_fp4,
402
+ )
403
+ block_ids = slots[valid] // block_size
404
+ block_offsets = slots[valid] % block_size
405
+ value_offsets = block_offsets[:, None] * value_bytes + torch.arange(
406
+ value_bytes, device=device
407
+ )
408
+ scale_offsets = block_size * value_bytes + block_offsets[:, None] * 4
409
+ scale_offsets = scale_offsets + torch.arange(4, device=device)
410
+ expected[block_ids[:, None], value_offsets] = values
411
+ expected[block_ids[:, None], scale_offsets] = scales
412
+
413
+ indexer_k_norm_rope_store(
414
+ k_pre,
415
+ positions,
416
+ cos_sin,
417
+ norm_weight,
418
+ 1e-20,
419
+ cache,
420
+ slots,
421
+ compress_ratio,
422
+ use_fp4,
423
+ )
424
+ # Fused RoPE can change zero signs without changing their numerical values.
425
+ actual_values = storage[block_ids[:, None], value_offsets]
426
+ if use_fp4:
427
+ for shift in (0, 4):
428
+ zero_mask = (actual_values >> shift) & 7 == 0
429
+ actual_values &= ~(zero_mask.to(torch.uint8) << (shift + 3))
430
+ zero_mask = (values >> shift) & 7 == 0
431
+ values &= ~(zero_mask.to(torch.uint8) << (shift + 3))
432
+ else:
433
+ actual_values[actual_values == 128] = 0
434
+ values[values == 128] = 0
435
+ storage[block_ids[:, None], value_offsets] = actual_values
436
+ expected[block_ids[:, None], value_offsets] = values
437
+ torch.testing.assert_close(storage, expected, rtol=0, atol=0)
438
+
439
+
440
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA kernel")
441
+@pytest.mark.parametrize("use_fp4", [False, True])
442
+@torch.inference_mode()
443
+def test_indexer_k_cuda_graph_replay_reads_current_projection(use_fp4):
444
+ """Replay must consume updated keys and slot mapping through the captured API."""
445
+ from vllm.models.deepseek_v4_1.common.ops.indexer_k_store import (
446
+ indexer_k_norm_rope_store,
447
+ )
448
+
449
+ if use_fp4 and not current_platform.is_device_capability_family(100):
450
+ pytest.skip("MXFP4 indexer cache requires an SM100-family GPU")
451
+ torch.manual_seed(18)
452
+ k_pre = torch.randn(19, 128, device="cuda", dtype=torch.bfloat16)
453
+ norm_weight = torch.ones(128, device="cuda", dtype=torch.bfloat16)
454
+ positions = torch.arange(19, device="cuda")
455
+ cos_sin = torch.zeros(32, 64, device="cuda")
456
+ cos_sin[:, :32] = 1
457
+ slots = torch.arange(17, device="cuda")
458
+ cache = torch.full(
459
+ (2, 16, (64 if use_fp4 else 128) + 4),
460
+ 0xA5,
461
+ device="cuda",
462
+ dtype=torch.uint8,
463
+ )
464
+
465
+ def run():
466
+ indexer_k_norm_rope_store(
467
+ k_pre,
468
+ positions,
469
+ cos_sin,
470
+ norm_weight,
471
+ 1e-20,
472
+ cache,
473
+ slots,
474
+ 2,
475
+ use_fp4,
476
+ )
477
+
478
+ run()
479
+ graph = torch.cuda.CUDAGraph()
480
+ with torch.cuda.graph(graph):
481
+ run()
482
+ k_pre.neg_()
483
+ slots.add_(4)
484
+ cache.fill_(0xA5)
485
+ graph.replay()
486
+ replayed = cache.clone()
487
+ cache.fill_(0xA5)
488
+ run()
489
+ torch.testing.assert_close(cache, replayed, rtol=0, atol=0)
490
+
491
+
492
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm cache layout")
493
+@pytest.mark.parametrize("block_size", [16, 64, 128])
494
+@pytest.mark.parametrize("compress_ratio", [1, 2])
495
+@torch.inference_mode()
496
+def test_indexer_k_store_roundtrips_through_rocm_gather(block_size, compress_ratio):
497
+ """ROCm reads the indexer K cache back with the 16x16-tiled value layout.
498
+
499
+ ``cp_gather_indexer_k_quant_cache_triton`` (prefill) and aiter's
500
+ ``Preshuffle=True`` paged kernel (decode) both pick that layout whenever
501
+ ``block_size > 1``, so a row-major store would feed the indexer permuted
502
+ key bytes and randomize its top-k. Assert the write/read pair is exact.
503
+ """
504
+ from vllm.models.deepseek_v4_1.common.ops.indexer_k_store import (
505
+ indexer_k_norm_rope_store,
506
+ )
507
+ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
508
+ cp_gather_indexer_k_quant_cache_triton,
509
+ )
510
+
511
+ torch.manual_seed(0)
512
+ device = "cuda"
513
+ num_tokens = 3 * block_size + 5
514
+ num_blocks = num_tokens // block_size + 2
515
+ k_pre = torch.randn(num_tokens, HEAD_DIM, device=device, dtype=torch.bfloat16)
516
+ norm_weight = torch.randn(HEAD_DIM, device=device, dtype=torch.bfloat16)
517
+ positions = torch.arange(num_tokens, device=device)
518
+ angles = torch.randn(MAX_POS, ROPE_DIM // 2, device=device)
519
+ cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1).float()
520
+ cache = torch.zeros(
521
+ num_blocks, block_size, HEAD_DIM + 4, device=device, dtype=torch.uint8
522
+ )
523
+ slots = torch.arange(num_tokens, device=device)
524
+
525
+ indexer_k_norm_rope_store(
526
+ k_pre,
527
+ positions,
528
+ cos_sin,
529
+ norm_weight,
530
+ 1e-20,
531
+ cache,
532
+ slots,
533
+ compress_ratio,
534
+ False,
535
+ )
536
+
537
+ k_fp8 = torch.zeros(num_tokens, HEAD_DIM, device=device, dtype=torch.float8_e4m3fn)
538
+ k_scale = torch.zeros(num_tokens, 4, device=device, dtype=torch.uint8)
539
+ cp_gather_indexer_k_quant_cache_triton(
540
+ cache,
541
+ k_fp8,
542
+ k_scale,
543
+ torch.arange(num_blocks, device=device, dtype=torch.int32).view(1, -1),
544
+ torch.tensor([0, num_tokens], device=device, dtype=torch.int32),
545
+ torch.zeros(num_tokens, device=device, dtype=torch.int32),
546
+ )
547
+
548
+ expected_values, expected_scale = _indexer_k_reference(
549
+ k_pre, positions, cos_sin, norm_weight, compress_ratio, False
550
+ )
551
+ emitted = (positions + 1) % compress_ratio == 0
552
+ gathered = k_fp8.view(torch.uint8)
553
+ gathered = torch.where(gathered == 128, gathered.new_zeros(()), gathered)
554
+ expected_values = torch.where(
555
+ expected_values == 128, expected_values.new_zeros(()), expected_values
556
+ )
557
+ torch.testing.assert_close(
558
+ gathered[emitted], expected_values[emitted], rtol=0, atol=0
559
+ )
560
+ torch.testing.assert_close(
561
+ k_scale[emitted], expected_scale[emitted], rtol=0, atol=0
562
+ )
tests/kernels/test_fused_inv_rope_fp8_quant.py
+88
-17
@@ -159,13 +159,12 @@ def _ref_ue8m0_quant_block(x_f32: torch.Tensor) -> tuple[torch.Tensor, torch.Ten
159
absmax -> 2^ceil(log2(absmax / fp8_max)) -> clamp(x / scale) -> fp8
160
161
Args:
162
- x_f32: [..., quant_group_size] float32 — one or more 128-element blocks.
162
+ x_f32: [..., quant_group_size] float32 quantization blocks.
163
164
Returns:
165
x_fp8: same shape, float8_e4m3fn
166
scales: [...] float32, one scale per block
167
"""
168
- assert x_f32.shape[-1] == QUANT_GROUP_SIZE
168
absmax = x_f32.abs().amax(dim=-1, keepdim=True).clamp(min=EPS)
169
scale_raw = absmax * (1.0 / FP8_MAX)
170
scale = torch.exp2(torch.ceil(torch.log2(scale_raw)))
@@ -253,6 +252,56 @@ def reference_inv_rope_fp8_quant(
252
# =========================================================================
253
254
255
+@pytest.mark.parametrize("num_tokens", [1, 7, 128, 4096])
256
+@pytest.mark.parametrize("quant_group_size", [32, 128])
257
+@torch.inference_mode()
258
+def test_packed_quantization_preserves_rotary_dimensions(num_tokens, quant_group_size):
259
+ """RoPE may span multiple scale blocks; pack every scale and padding row."""
260
+ from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import (
261
+ fused_inv_rope_fp8_quant as quantize,
262
+ )
263
+
264
+ torch.manual_seed(42)
265
+ groups, heads_per_group = 2, 2
266
+ x = torch.randn(
267
+ num_tokens,
268
+ groups * heads_per_group,
269
+ HEAD_DIM,
270
+ device="cuda",
271
+ dtype=torch.bfloat16,
272
+ )
273
+ cache = make_cos_sin_cache(4096)
274
+ positions = torch.randint(4096, (num_tokens,), device="cuda")
275
+ q, packed_scales = quantize(
276
+ x,
277
+ positions,
278
+ cache,
279
+ groups,
280
+ heads_per_group,
281
+ quant_group_size=quant_group_size,
282
+ tma_aligned_scales=True,
283
+ )
284
+ scale_bytes = torch.stack(
285
+ [(packed_scales >> (8 * i)) & 0xFF for i in range(4)], dim=-1
286
+ ).flatten(-2)
287
+ scales = torch.exp2(scale_bytes.float() - 127)
288
+ rotated = reference_inv_rope(x.float(), positions, cache)
289
+ ref_q, ref_scales = _ref_ue8m0_quant_block(
290
+ rotated.reshape(num_tokens, groups, -1, quant_group_size)
291
+ )
292
+ torch.testing.assert_close(scales, ref_scales, rtol=0, atol=0)
293
+ actual = q.float().unflatten(-1, (-1, quant_group_size)) * scales.unsqueeze(-1)
294
+ expected = ref_q.float() * ref_scales.unsqueeze(-1)
295
+ assert (actual - expected).norm() / expected.norm() < 1e-3
296
+
297
+ aligned_tokens = (num_tokens + 3) // 4 * 4
298
+ full_scales = packed_scales.transpose(0, 1).as_strided(
299
+ (groups, aligned_tokens, packed_scales.shape[-1]),
300
+ (packed_scales.stride(1), 1, packed_scales.stride(2)),
301
+ )
302
+ assert torch.count_nonzero(full_scales[:, num_tokens:]) == 0
303
+
304
+
305
@pytest.mark.parametrize("num_tokens", [1, 7, 32, 128])
306
@pytest.mark.parametrize(
307
"num_heads,n_groups",
@@ -260,8 +309,9 @@ def reference_inv_rope_fp8_quant(
309
ids=["H64_G8", "H32_G4", "H128_G8"],
310
)
311
@pytest.mark.parametrize("seed", [0, 42])
312
+@pytest.mark.parametrize("quantize", [True, False], ids=["fp8", "bf16"])
313
@torch.inference_mode()
264
-def test_correctness(num_tokens, num_heads, n_groups, seed):
314
+def test_correctness(num_tokens, num_heads, n_groups, seed, quantize):
315
"""Compare fused kernel against reference for FP8 values and scales."""
316
torch.manual_seed(seed)
317
@@ -280,6 +330,26 @@ def test_correctness(num_tokens, num_heads, n_groups, seed):
330
max_pos, ROPE_DIM, dtype=torch.float32, device=device
331
)
332
333
+ if not quantize:
334
+ ref_output = reference_inv_rope(o, positions, cos_sin_cache).view(
335
+ num_tokens, n_groups, -1
336
+ )
337
+ fused_output, fused_scale = fused_inv_rope_fp8_quant(
338
+ o.clone(),
339
+ positions,
340
+ cos_sin_cache,
341
+ n_groups,
342
+ heads_per_group,
343
+ quantize=False,
344
+ )
345
+ d = heads_per_group * HEAD_DIM
346
+ assert fused_output.shape == (num_tokens, n_groups, d)
347
+ assert fused_scale.numel() == 0
348
+ assert fused_output.dtype == torch.bfloat16
349
+ assert fused_output.stride() == (d, num_tokens * d, 1)
350
+ torch.testing.assert_close(fused_output, ref_output, rtol=1e-2, atol=1e-2)
351
+ return
352
+
353
# Reference
354
ref_fp8, ref_scale = reference_inv_rope_fp8_quant(
355
o.clone(),
@@ -726,6 +796,7 @@ def test_einsum_end_to_end(num_tokens, num_heads, n_groups):
796
This catches stride/layout bugs that only manifest when the einsum
797
kernel actually consumes the quantized activations.
798
"""
799
+ from vllm.models.deepseek_v4.nvidia.ops.o_proj import deep_gemm_fp8_o_proj
800
from vllm.utils.deep_gemm import (
801
fp8_einsum,
802
is_deep_gemm_supported,
@@ -793,23 +864,23 @@ def test_einsum_end_to_end(num_tokens, num_heads, n_groups):
864
)
865
866
# -- FUSED path --
796
- fused_fp8, fused_scale = fused_inv_rope_fp8_quant(
867
+ wo_a = torch.nn.Module()
868
+ wo_a.register_buffer("weight", w_fp8)
869
+ wo_a.register_buffer("weight_scale", w_scale_t)
870
+ z_fused = deep_gemm_fp8_o_proj(
871
o.clone(),
872
positions,
873
cos_sin_cache,
800
- n_groups,
801
- heads_per_group,
802
- )
803
- z_fused = torch.empty(
804
- num_tokens, n_groups, o_lora_rank, device=device, dtype=torch.bfloat16
805
- )
806
- fp8_einsum(
807
- "bhr,hdr->bhd",
808
- (fused_fp8, fused_scale),
809
- (w_fp8, w_scale_t),
810
- z_fused,
811
- recipe=recipe,
812
- )
874
+ wo_a,
875
+ torch.nn.Identity(),
876
+ n_groups=n_groups,
877
+ heads_per_group=heads_per_group,
878
+ nope_dim=NOPE_DIM,
879
+ rope_dim=ROPE_DIM,
880
+ o_lora_rank=o_lora_rank,
881
+ einsum_recipe=recipe,
882
+ tma_aligned_scales=True,
883
+ ).view_as(z_ref)
884
885
# -- Checks --
886
# Einsum output: Triton and CUDA both rotate in fp32 now, so diffs
tests/kernels/test_mhc_kernels.py
+285
-1
@@ -12,7 +12,10 @@ import vllm.model_executor.layers.mhc as mhc_layers
12
from vllm.model_executor.kernels.mhc.tilelang import (
13
_tilelang_hc_prenorm_gemm,
14
_torch_hc_prenorm_gemm,
15
+ mhc_pre_delayed_tilelang,
16
)
17
+from vllm.model_executor.kernels.mhc.torch import mhc_pre_delayed_torch
18
+from vllm.model_executor.kernels.mhc.triton import hc_collapse_triton
19
from vllm.model_executor.layers.mhc import (
20
HAS_AITER_MHC,
21
HAS_AITER_MHC_FUSED,
@@ -26,12 +29,293 @@ from vllm.models.deepseek_v4.nvidia.model import (
29
DeepseekV4DecoderLayer,
30
DeepseekV4Model,
31
)
32
+from vllm.models.deepseek_v4_1.nvidia.model import (
33
+ DeepseekV4DecoderLayer as DeepseekV41DecoderLayer,
34
+)
35
from vllm.platforms import current_platform
36
from vllm.utils.torch_utils import set_random_seed
37
38
DEVICE = current_platform.device_type
39
40
41
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
42
+@pytest.mark.parametrize(
43
+ "num_tokens,hc_mult,hidden_size",
44
+ [
45
+ (0, 4, 5120),
46
+ (1, 4, 5120),
47
+ (33, 4, 5137),
48
+ (128, 4, 5120),
49
+ (1024, 4, 5120),
50
+ (2, 1, 513),
51
+ (7, 8, 5137),
52
+ (65537, 4, 65),
53
+ ],
54
+)
55
+@pytest.mark.parametrize("strided", [False, True])
56
+def test_hc_collapse_preserves_weighted_residual_sum(
57
+ num_tokens, hc_mult, hidden_size, strided
58
+):
59
+ """Handle tile tails and strided streams without changing FP32 mixing."""
60
+ set_random_seed(0)
61
+ x = torch.randn(
62
+ num_tokens,
63
+ hc_mult,
64
+ hidden_size * (2 if strided else 1),
65
+ dtype=torch.bfloat16,
66
+ device=DEVICE,
67
+ )
68
+ pre = torch.rand(num_tokens, hc_mult, device=DEVICE)
69
+ if strided:
70
+ x = x[..., ::2]
71
+ pre = pre.t().contiguous().t()
72
+ if num_tokens:
73
+ pre[0].zero_()
74
+ pre[0, -1] = 1
75
+ if num_tokens > 1 and hc_mult == 4:
76
+ pre[1].fill_(1)
77
+ x[1, 1] = -x[1, 0]
78
+ x[1, 3] = -x[1, 2]
79
+ expected = (pre.unsqueeze(-1) * x.float()).sum(dim=1).to(x.dtype)
80
+ actual = hc_collapse_triton(x, pre)
81
+ assert actual.is_contiguous()
82
+ # BF16 rounding only: ptxas <13.1 contracts the FP32 mix into FFMA despite
83
+ # enable_fp_fusion=False, so compare within a ULP and pin the exact cases below.
84
+ torch.testing.assert_close(actual, expected, atol=1.6e-2, rtol=1e-2)
85
+ if num_tokens:
86
+ torch.testing.assert_close(actual[0], x[0, -1], atol=0, rtol=0)
87
+ if num_tokens > 1 and hc_mult == 4:
88
+ torch.testing.assert_close(
89
+ actual[1], torch.zeros_like(actual[1]), atol=0, rtol=0
90
+ )
91
+
92
+
93
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
94
+def test_hc_collapse_custom_op_supports_compile():
95
+ x = torch.randn(2, 4, 5120, dtype=torch.bfloat16, device=DEVICE)
96
+ pre = torch.rand(2, 4, device=DEVICE)
97
+ torch.library.opcheck(torch.ops.vllm.hc_collapse_triton.default, (x, pre))
98
+
99
+
100
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
101
+@pytest.mark.parametrize("num_tokens", [1, 7])
102
+def test_v41_dspark_head_collapses_with_last_ffn_mix(num_tokens, monkeypatch):
103
+ """Match reference DSparkBlock.forward_head's hc_pre before its RMSNorm."""
104
+ from vllm.models.deepseek_v4_1.nvidia import dspark
105
+
106
+ set_random_seed(0)
107
+ hidden_size, hc_mult = 5120, 4
108
+ streams = torch.randn(
109
+ num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=DEVICE
110
+ )
111
+ mixes = [torch.rand(num_tokens, hc_mult, device=DEVICE) for _ in range(2)]
112
+ carried_mixes = []
113
+
114
+ def make_layer(mix):
115
+ def forward(hidden, positions, ids, pre, post, res, residual):
116
+ carried_mixes.append(pre)
117
+ return hidden, None, None, None, mix
118
+
119
+ return forward
120
+
121
+ monkeypatch.setattr(dspark, "mhc_post_tilelang", lambda *args: streams)
122
+ draft = SimpleNamespace(
123
+ use_sequence_parallel=False,
124
+ hc_mult=hc_mult,
125
+ layers=[make_layer(mix) for mix in mixes],
126
+ )
127
+ actual = dspark.DSparkDeepseekV4Model.forward(
128
+ draft,
129
+ input_ids=torch.zeros(num_tokens, dtype=torch.long, device=DEVICE),
130
+ positions=torch.arange(num_tokens, device=DEVICE),
131
+ inputs_embeds=torch.zeros(
132
+ num_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE
133
+ ),
134
+ )
135
+ expected = (streams.float() * mixes[-1].unsqueeze(-1)).sum(dim=1)
136
+ assert carried_mixes[0] is None
137
+ assert carried_mixes[1] is mixes[0]
138
+ torch.testing.assert_close(
139
+ actual, expected.to(streams.dtype), atol=1.6e-2, rtol=1e-2
140
+ )
141
+
142
+
143
+@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required")
144
+@pytest.mark.parametrize(
145
+ "num_tokens,sinkhorn_repeat,norm_eps",
146
+ [
147
+ (0, 20, 1e-6),
148
+ (1, 1, 1e-6),
149
+ (1, 20, 1e-6),
150
+ (33, 20, 1e-20),
151
+ (128, 20, 1e-6),
152
+ (1024, 20, 1e-6),
153
+ ],
154
+)
155
+@pytest.mark.parametrize("entry", ["broadcast", "identity", "carried"])
156
+@pytest.mark.parametrize("use_deep_gemm", [False, True])
157
+@pytest.mark.parametrize("fused_norm", [False, True])
158
+def test_deepseek_v41_mhc_pre_delayed(
159
+ num_tokens,
160
+ sinkhorn_repeat,
161
+ norm_eps,
162
+ entry,
163
+ use_deep_gemm,
164
+ fused_norm,
165
+ monkeypatch,
166
+):
167
+ """Collapse with the carried pre-mix and save the new one for the next sublayer."""
168
+ from vllm.utils.deep_gemm import is_deep_gemm_supported
169
+
170
+ if use_deep_gemm and not is_deep_gemm_supported():
171
+ pytest.skip("DeepGEMM not available")
172
+ monkeypatch.setattr(
173
+ "vllm.utils.deep_gemm.is_deep_gemm_supported", lambda: use_deep_gemm
174
+ )
175
+ set_random_seed(0)
176
+ hc_mult, hidden_size = 4, 5120
177
+ residual = torch.randn(
178
+ num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=DEVICE
179
+ )
180
+ if num_tokens > 1:
181
+ residual[0].zero_()
182
+ fn = torch.randn(24, hc_mult, hidden_size, device=DEVICE) * 0.02
183
+ x = None
184
+ pre_mix = None
185
+ if entry == "broadcast":
186
+ x = residual[:, 0].contiguous()
187
+ residual = x.unsqueeze(1).expand(-1, hc_mult, -1).contiguous()
188
+ fn = fn.sum(1)
189
+ else:
190
+ fn = fn.flatten(1)
191
+ if entry == "carried":
192
+ pre_mix = torch.rand(num_tokens, hc_mult, device=DEVICE)
193
+ if num_tokens > 1:
194
+ # A constant collapse just above 1 rounds to BF16 1 before RMSNorm.
195
+ residual[1].fill_(1)
196
+ pre_mix[1].zero_()
197
+ pre_mix[1, 0] = 1.003
198
+ scale = torch.tensor([0.5, 0.25, 1.0], device=DEVICE)
199
+ base = torch.randn(24, device=DEVICE)
200
+ args = (residual, fn, scale, base, 1e-20, 1e-6, 1e-6, 2.0, sinkhorn_repeat)
201
+ expected = mhc_pre_delayed_torch(*args, pre_mix=pre_mix, x=x)
202
+ norm_kwargs = {}
203
+ if fused_norm:
204
+ weight = torch.empty(hidden_size, dtype=torch.bfloat16, device=DEVICE)
205
+ weight.uniform_(0.5, 1.5)
206
+ norm_kwargs = dict(norm_weight=weight, norm_eps=norm_eps)
207
+ # Compare to the original collapse + CUDA RMSNorm, including BF16 rounding.
208
+ unfused = mhc_pre_delayed_tilelang(*args, pre_mix=pre_mix, x=x)
209
+ normalized = torch.empty_like(unfused[2])
210
+ if num_tokens:
211
+ torch.ops._C.rms_norm(normalized, unfused[2], weight, norm_eps)
212
+ expected = (*unfused[:2], normalized, unfused[3])
213
+ actual = mhc_pre_delayed_tilelang(*args, pre_mix=pre_mix, x=x, **norm_kwargs)
214
+ if fused_norm and num_tokens:
215
+ graph = torch.cuda.CUDAGraph()
216
+ with torch.cuda.graph(graph):
217
+ captured = mhc_pre_delayed_tilelang(
218
+ *args, pre_mix=pre_mix, x=x, **norm_kwargs
219
+ )
220
+ graph.replay()
221
+ for eager, replayed in zip(actual, captured, strict=True):
222
+ torch.testing.assert_close(eager, replayed, atol=0, rtol=0)
223
+ if fused_norm and entry == "carried" and num_tokens > 1:
224
+ torch.testing.assert_close(actual[2][1], expected[2][1], atol=0, rtol=0)
225
+ # DeepGEMM uses TF32 for the projection; the TileLang fallback uses FP32.
226
+ atol, rtol = (1e-3, 1e-3) if use_deep_gemm else (1e-5, 1e-4)
227
+ for i in (0, 1, 3):
228
+ torch.testing.assert_close(actual[i], expected[i], atol=atol, rtol=rtol)
229
+ if pre_mix is None and not fused_norm:
230
+ torch.testing.assert_close(actual[2], expected[2], atol=0, rtol=0)
231
+ else:
232
+ torch.testing.assert_close(actual[2], expected[2], atol=1.6e-2, rtol=1e-2)
233
+
234
+
235
+@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required")
236
+@pytest.mark.parametrize("entry", ["broadcast", "pipeline", "residual", "engram"])
237
+def test_deepseek_v41_decoder_mixes_match_torch(
238
+ entry, monkeypatch, default_vllm_config
239
+):
240
+ """Preserve carried pre-mixes and Engram ordering at every decoder entry."""
241
+ set_random_seed(0)
242
+ decoder = DeepseekV41DecoderLayer.__new__(DeepseekV41DecoderLayer)
243
+ nn.Module.__init__(decoder)
244
+ decoder.hc_mult = 4
245
+ decoder.hc_sinkhorn_iters = 20
246
+ decoder.hc_eps = 1e-6
247
+ decoder.rms_norm_eps = 1e-20
248
+ decoder.hc_post_alpha = 2.0
249
+ decoder.use_sequence_parallel = False
250
+ decoder.engram = None
251
+ from vllm.model_executor.layers.layernorm import RMSNorm
252
+
253
+ decoder.attn_norm = decoder.ffn_norm = RMSNorm(5120, 1e-6).to(
254
+ device=DEVICE, dtype=torch.bfloat16
255
+ )
256
+ decoder.attn = lambda positions, x, _: x * 0.5
257
+ decoder.ffn = lambda x, input_ids: x * 0.25
258
+ with torch.device(DEVICE):
259
+ decoder.hc_attn_fn = torch.randn(24, 20480) * 0.02
260
+ decoder.hc_ffn_fn = torch.randn(24, 20480) * 0.02
261
+ decoder.hc_attn_fn_broadcast = decoder.hc_attn_fn.view(24, 4, 5120).sum(1)
262
+ decoder.hc_attn_scale = decoder.hc_ffn_scale = torch.ones(3)
263
+ decoder.hc_attn_base = torch.randn(24)
264
+ decoder.hc_ffn_base = torch.randn(24)
265
+ x = torch.randn(3, 5120, dtype=torch.bfloat16)
266
+ positions = torch.arange(3)
267
+ kwargs = {}
268
+ if entry != "broadcast":
269
+ kwargs["pre_mix"] = torch.rand(3, 4)
270
+ if entry == "pipeline":
271
+ x = torch.randn(3, 4, 5120, dtype=torch.bfloat16)
272
+ if entry in ("residual", "engram"):
273
+ kwargs.update(
274
+ residual=torch.randn(3, 4, 5120, dtype=torch.bfloat16),
275
+ post_mix=torch.rand(3, 4, 1),
276
+ res_mix=torch.rand(3, 4, 4),
277
+ )
278
+ if entry == "engram":
279
+
280
+ class FakeEngram(nn.Module):
281
+ layer_hash_index = 0
282
+
283
+ def forward(self, residual, hashes, mask):
284
+ return residual + 0.125
285
+
286
+ decoder.engram = FakeEngram()
287
+ kwargs["engram_hashes"] = torch.zeros(3, 1, 1, dtype=torch.int32)
288
+
289
+ actual = decoder(x, positions, None, **kwargs)
290
+
291
+ def reference(*args, norm_weight, norm_eps, **kwargs):
292
+ post, res, collapsed, pre = mhc_pre_delayed_torch(*args, **kwargs)
293
+ return post, res, decoder.attn_norm(collapsed), pre
294
+
295
+ monkeypatch.setattr(
296
+ "vllm.models.deepseek_v4_1.nvidia.model.mhc_pre_delayed_tilelang",
297
+ reference,
298
+ )
299
+ expected = decoder(x, positions, None, **kwargs)
300
+ for result, ref in zip(actual, expected, strict=True):
301
+ torch.testing.assert_close(result, ref, atol=2e-2, rtol=1e-2)
302
+
303
+
304
+@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required")
305
+@pytest.mark.parametrize("carried", [False, True])
306
+def test_mhc_pre_delayed_custom_op_supports_compile(carried):
307
+ set_random_seed(0)
308
+ x = torch.randn(2, 4, 5120, dtype=torch.bfloat16, device=DEVICE)
309
+ fn = torch.randn(24, 20480, device=DEVICE) * 0.02
310
+ pre_mix = torch.rand(2, 4, device=DEVICE) if carried else None
311
+ scale = torch.ones(3, device=DEVICE)
312
+ base = torch.zeros(24, device=DEVICE)
313
+ torch.library.opcheck(
314
+ torch.ops.vllm.mhc_pre_delayed_tilelang.default,
315
+ (x, fn, scale, base, 1e-20, 1e-6, 1e-6, 2.0, 20, pre_mix),
316
+ )
317
+
318
+
319
def sinkhorn_normalize_ref(x: torch.Tensor, repeat: int, eps: float) -> torch.Tensor:
320
x = x.softmax(-1) + eps
321
x = x / (x.sum(-2, keepdim=True) + eps)
@@ -203,7 +487,7 @@ def test_hc_prenorm_gemm_tilelang(num_tokens, hidden_size):
487
_tilelang_hc_prenorm_gemm(x, fn, out, sqrsum, hidden_size, hc_mult)
488
489
torch.testing.assert_close(out, out_ref, atol=1e-5, rtol=1e-4)
206
- torch.testing.assert_close(sqrsum, sqrsum_ref, atol=8.0, rtol=5e-4)
490
+ torch.testing.assert_close(sqrsum, sqrsum_ref, atol=1e-2, rtol=1e-6)
491
492
493
@pytest.mark.skipif(
tests/model_executor/layers/test_mla_short_prefill_indexer.py
+110
@@ -290,3 +290,113 @@ def test_deepseek_v32_dispatches_selected_mha(
290
expected_args[1:],
291
)
292
)
293
+
294
+
295
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
296
+def test_select_candidate_blocks_tolerates_empty_rows():
297
+ """Full-cudagraph decode pads the batch with seq_len-0 rows. The newest
298
+ block pin must not index -1 for them (device-side assert); they select
299
+ no candidate blocks and real rows still pin their newest block."""
300
+ block_size, topk_blocks = 8, 3
301
+ logits = torch.zeros(3, 64, device="cuda")
302
+ logits[0, 3] = 5.0 # block 0 scores highest for row 0
303
+ logits[2, 9] = 5.0 # block 1 scores highest for row 2
304
+ row_ks = torch.zeros(3, dtype=torch.int64, device="cuda")
305
+ row_ke = torch.tensor([40, 0, 17], device="cuda")
306
+ out = torch.empty(3, topk_blocks, dtype=torch.int32, device="cuda")
307
+
308
+ sparse_indexer._select_candidate_blocks(
309
+ logits, row_ks, row_ke, topk_blocks, block_size, out
310
+ )
311
+
312
+ assert out[1].tolist() == [-1, -1, -1]
313
+ assert out[0, 0].item() == 4 and 0 in out[0].tolist() # newest block pinned
314
+ assert out[2, 0].item() == 2 and 1 in out[2].tolist()
315
+ assert (out[0] >= 0).all() and (out[2, :2] >= 0).all()
316
+
317
+
318
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
319
+@pytest.mark.parametrize(
320
+ "width,block_size,k,decode",
321
+ [
322
+ (73, 8, 16, False),
323
+ (97, 3, 7, False),
324
+ (32768, 8, 2048, False),
325
+ (32768, 8, 2048, True),
326
+ ],
327
+)
328
+def test_candidate_kernels_preserve_packed_bounds_and_padding(
329
+ width, block_size, k, decode
330
+):
331
+ """Preserve top-k ties, newest blocks, empty rows and candidate clamping."""
332
+ from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
333
+ apply_candidate_mask,
334
+ select_candidate_blocks,
335
+ )
336
+
337
+ torch.manual_seed(42)
338
+ rows = 6
339
+ logits = torch.randn(rows, width * 2, device="cuda")[:, ::2]
340
+ logits[0, :16] = 0
341
+ logits[2, 5] = float("nan")
342
+ starts = torch.tensor([0, 5, 3, 0, 7, 1], device="cuda", dtype=torch.int32)
343
+ ends = torch.tensor([0, width, width - 1, 1, 11, width], device="cuda")
344
+ repeat = 1
345
+ if decode:
346
+ starts = None
347
+ ends = torch.tensor([0, 1, width], device="cuda", dtype=torch.int32)
348
+ repeat = 2
349
+ ks = torch.zeros(rows, device="cuda", dtype=torch.int64) if decode else starts
350
+ ke = ends.repeat_interleave(repeat)
351
+ cols = torch.arange(width, device="cuda")
352
+ valid = (cols >= ks[:, None]) & (cols < ke[:, None])
353
+ scores = logits.masked_fill(~valid, -torch.inf)
354
+ blocks = ((cols - ks[:, None]) // block_size).clamp(min=0).long()
355
+ nblocks = (width + block_size - 1) // block_size
356
+ reduced = logits.new_full((rows, nblocks), -torch.inf)
357
+ reduced.scatter_reduce_(1, blocks, scores, reduce="amax", include_self=True)
358
+ lengths = ke - ks
359
+ last = ((lengths - 1) // block_size).clamp(min=0).long()
360
+ reduced.scatter_(
361
+ 1, last[:, None], torch.where(lengths > 0, torch.inf, -torch.inf)[:, None]
362
+ )
363
+ top = reduced.topk(min(k, nblocks), dim=-1)
364
+ expected = torch.full((rows, k), -1, device="cuda", dtype=torch.int32)
365
+ expected[:, : top.indices.shape[1]] = torch.where(
366
+ top.values > -torch.inf, top.indices, -1
367
+ ).int()
368
+ actual = torch.empty(rows, k * 2, device="cuda", dtype=torch.int32)[:, ::2]
369
+ select_candidate_blocks(logits, starts, ends, k, block_size, actual, repeat)
370
+ torch.testing.assert_close(actual, expected, rtol=0, atol=0)
371
+
372
+ candidates = actual.clone()
373
+ candidates[:, 0] = nblocks + 10
374
+ candidates[:, 1] = 0
375
+ candidates[:, 2] = 0
376
+ positions = ks[:, None, None] + candidates.long()[:, :, None] * block_size
377
+ positions = positions + torch.arange(block_size, device="cuda")
378
+ keep = torch.zeros(rows, width, device="cuda", dtype=torch.int8)
379
+ keep.scatter_reduce_(
380
+ 1,
381
+ positions.clamp(0, width - 1).reshape(rows, -1),
382
+ (candidates >= 0)[:, :, None]
383
+ .expand(-1, -1, block_size)
384
+ .reshape(rows, -1)
385
+ .to(torch.int8),
386
+ reduce="amax",
387
+ include_self=True,
388
+ )
389
+ reference = logits.masked_fill((keep == 0) | ~valid, -torch.inf)
390
+ apply_candidate_mask(logits, starts, ends, candidates, block_size, repeat)
391
+ torch.testing.assert_close(logits, reference, rtol=0, atol=0, equal_nan=True)
392
+ graph = torch.cuda.CUDAGraph()
393
+ with torch.cuda.graph(graph):
394
+ apply_candidate_mask(logits, starts, ends, candidates, block_size, repeat)
395
+ for block in (1, -1):
396
+ candidates.fill_(-1)
397
+ candidates[:, 0] = block
398
+ logits.fill_(3.0)
399
+ graph.replay()
400
+ keep = valid & (block >= 0) & ((cols - ks[:, None]) // block_size == block)
401
+ reference = torch.where(keep, 3.0, -torch.inf)
402
+ torch.testing.assert_close(logits, reference, rtol=0, atol=0)
tests/models/multimodal/generation/test_maverick.py
+2
@@ -616,6 +616,8 @@ def test_dummy_maverick(
616
) -> None:
617
# Disable multiprocessing allows us to access model executor from LLM engine
618
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
619
+ # The forked test process may inherit CUDA state from test collection.
620
+ monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
621
622
model_path = create_reduced_maverick_model(
623
original_model_name=original_model_name,
tests/models/registry.py
+9
@@ -858,6 +858,9 @@ _MULTIMODAL_EXAMPLE_MODELS = {
858
"DeepseekV4ForConditionalGeneration": _HfExamplesInfo(
859
"deepseek-ai/DeepSeek-V4-Flash-Vision-Exp",
860
),
861
+ "DeepseekV41ForCausalLM": _HfExamplesInfo(
862
+ "deepseek-ai/DeepSeek-V4.1-Flash",
863
+ ),
864
"Dots3NoteForCausalLM": _HfExamplesInfo(
865
"dots-studio/dots3-note-prev",
866
is_available_online=False,
@@ -1537,6 +1540,12 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1540
is_available_online=False,
1541
use_original_num_layers=True, # DSpark has >1 draft block
1542
),
1543
+ "DSparkV41DraftModel": _HfExamplesInfo(
1544
+ "deepseek-ai/DeepSeek-V4.1-Flash",
1545
+ speculative_model="deepseek-ai/DeepSeek-V4.1-Flash", # draft in mtp.*
1546
+ is_available_online=False,
1547
+ use_original_num_layers=True, # DSpark has >1 draft block
1548
+ ),
1549
"K3DSparkModel": _HfExamplesInfo(
1550
"moonshotai/Kimi-K3",
1551
speculative_model="Inferact/Kimi-K3-DSpark",
tests/models/test_deepseek_v4_mega_moe.py
+211
-4
@@ -18,7 +18,13 @@ from vllm.models.deepseek_v4.nvidia.model import (
18
)
19
from vllm.models.deepseek_v4.nvidia.mtp import DeepSeekV4MTP
20
from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
21
+from vllm.models.deepseek_v4_1.common.mm_preprocess import (
22
+ IMAGE_PAD_ID,
23
+ IMAGE_SENTINEL_BASE_ID,
24
+)
25
+from vllm.models.deepseek_v4_1.nvidia.model import DeepseekV4MoE as DeepseekV41MoE
26
from vllm.platforms import current_platform
27
+from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config
28
29
pytestmark = pytest.mark.skipif(
30
not current_platform.is_cuda(),
@@ -26,6 +32,133 @@ pytestmark = pytest.mark.skipif(
32
)
33
34
35
+@pytest.fixture
36
+def v41_moe_config(dist_init):
37
+ return SimpleNamespace(
38
+ model_config=SimpleNamespace(
39
+ hf_config=DeepseekV41Config(
40
+ text_config=dict(
41
+ hidden_size=128,
42
+ num_hidden_layers=2,
43
+ n_routed_experts=8,
44
+ num_experts_per_tok=2,
45
+ dspark_n_routed_experts=4,
46
+ dspark_num_experts_per_tok=3,
47
+ n_shared_experts=1,
48
+ moe_intermediate_size=128,
49
+ hidden_act="silu",
50
+ swiglu_limit=10.0,
51
+ norm_topk_prob=True,
52
+ topk_method="noaux_tc",
53
+ routed_scaling_factor=1.5,
54
+ ),
55
+ ),
56
+ ),
57
+ quant_config=None,
58
+ kernel_config=SimpleNamespace(moe_backend="deep_gemm_mega_moe"),
59
+ parallel_config=SimpleNamespace(
60
+ enable_expert_parallel=True,
61
+ enable_eplb=False,
62
+ eplb_config=SimpleNamespace(num_redundant_experts=0),
63
+ ),
64
+ scheduler_config=SimpleNamespace(max_num_batched_tokens=4),
65
+ compilation_config=SimpleNamespace(static_forward_context={}),
66
+ )
67
+
68
+
69
+@pytest.mark.parametrize("vision", [False, True])
70
+@pytest.mark.parametrize("layer_id,num_experts,top_k", [(0, 8, 2), (2, 4, 3)])
71
+def test_deepseek_v41_moe_routes_without_hash_table(
72
+ v41_moe_config, monkeypatch, vision, layer_id, num_experts, top_k
73
+):
74
+ """Main and draft layers select experts by score and preserve image routing."""
75
+ config = v41_moe_config.model_config.hf_config
76
+ config.vision_n_layers = int(vision)
77
+ with torch.device("cuda"):
78
+ moe = DeepseekV41MoE(v41_moe_config, prefix=f"model.layers.{layer_id}.ffn")
79
+ hidden_states = torch.randn(4, config.hidden_size)
80
+ input_ids = (
81
+ torch.tensor([42, IMAGE_SENTINEL_BASE_ID, IMAGE_PAD_ID, 129257])
82
+ if vision
83
+ else None
84
+ )
85
+ assert moe.gate.tid2eid is None
86
+ assert moe.gate.weight.shape == (num_experts, config.hidden_size)
87
+ assert isinstance(moe.experts, DeepseekV4MegaMoEExperts)
88
+ assert moe.experts.w13_weight.shape == (num_experts, 256, 64)
89
+ assert moe.experts.top_k == top_k
90
+ assert (config.n_routed_experts, config.num_experts_per_tok) == (8, 2)
91
+
92
+ with torch.no_grad():
93
+ moe.gate.weight.normal_(std=0.01)
94
+ moe.gate.e_score_correction_bias.copy_(torch.arange(num_experts, device="cuda"))
95
+ if vision:
96
+ moe.gate.bias_vl.copy_(-moe.gate.e_score_correction_bias)
97
+
98
+ scores = torch.nn.functional.softplus(
99
+ torch.nn.functional.linear(hidden_states, moe.gate.weight)
100
+ ).sqrt()
101
+ bias = moe.gate.e_score_correction_bias
102
+ if vision:
103
+ image_mask = (input_ids == IMAGE_SENTINEL_BASE_ID) | (input_ids == IMAGE_PAD_ID)
104
+ bias = torch.where(image_mask[:, None], moe.gate.bias_vl, bias)
105
+ expected_ids = (scores + bias).topk(top_k, dim=-1).indices
106
+ expected_weights = scores.gather(1, expected_ids)
107
+ expected_weights *= config.routed_scaling_factor / expected_weights.sum(
108
+ dim=-1, keepdim=True
109
+ )
110
+
111
+ def check_routing(x, weights, ids, *, activation_clamp):
112
+ torch.testing.assert_close(ids, expected_ids)
113
+ torch.testing.assert_close(weights, expected_weights)
114
+ assert activation_clamp == config.swiglu_limit
115
+ return x.clone()
116
+
117
+ monkeypatch.setattr(moe.experts, "forward", check_routing)
118
+ monkeypatch.setattr(moe.shared_experts, "forward", lambda x: 2 * x)
119
+ torch.testing.assert_close(moe(hidden_states, input_ids), 3 * hidden_states)
120
+
121
+
122
+@pytest.mark.parametrize(
123
+ "draft_experts,draft_top_k,expected", [(0, 0, (8, 2)), (4, 3, (4, 3))]
124
+)
125
+def test_deepseek_v41_fused_moe_uses_draft_counts_or_main_defaults(
126
+ v41_moe_config, monkeypatch, draft_experts, draft_top_k, expected
127
+):
128
+ config = v41_moe_config.model_config.hf_config
129
+ config.dspark_n_routed_experts = draft_experts
130
+ config.dspark_num_experts_per_tok = draft_top_k
131
+ v41_moe_config.kernel_config.moe_backend = "auto"
132
+ captured = {}
133
+
134
+ def make_experts(**kwargs):
135
+ captured.update(kwargs)
136
+ return torch.nn.Identity()
137
+
138
+ monkeypatch.setattr(
139
+ "vllm.models.deepseek_v4.nvidia.model.FusedMoEFactory", make_experts
140
+ )
141
+ moe = DeepseekV41MoE(v41_moe_config, prefix="model.layers.2.ffn")
142
+ assert moe.gate.weight.shape == (expected[0], 128)
143
+ assert (captured["num_experts"], captured["top_k"]) == expected
144
+ assert captured["hash_indices_table"] is None
145
+
146
+
147
+def test_deepseek_v4_moe_preserves_configured_hash_layers(v41_moe_config):
148
+ config = v41_moe_config.model_config.hf_config
149
+ config.num_hash_layers = 1
150
+ config.vocab_size = 32
151
+ moe = DeepseekV4MoE(
152
+ v41_moe_config,
153
+ prefix="model.layers.0.ffn",
154
+ num_hash_layers=config.num_hash_layers,
155
+ )
156
+ assert moe.gate.tid2eid.shape == (32, 2)
157
+ assert moe.gate.e_score_correction_bias is None
158
+ with pytest.raises(ValueError, match="hash MoE routing requires input_ids"):
159
+ moe(torch.zeros(1, 128))
160
+
161
+
162
def test_deepseek_v4_mega_moe_expert_mapping():
163
mapping = make_deepseek_v4_expert_params_mapping(2)
164
@@ -169,15 +302,79 @@ def test_deepseek_v4_mega_moe_weight_loader_uses_ep_expert_ownership():
302
assert torch.count_nonzero(experts.w13_weight[1]) == 0
303
304
305
+@pytest.mark.parametrize("intermediate_size", [512, 2304])
306
+def test_deepseek_v4_mega_moe_padding_preserves_weights(monkeypatch, intermediate_size):
307
+ """Pad gate/up separately and zero every added down-projection column."""
308
+ experts = DeepseekV4MegaMoEExperts(
309
+ SimpleNamespace(
310
+ scheduler_config=SimpleNamespace(max_num_batched_tokens=4),
311
+ compilation_config=SimpleNamespace(static_forward_context={}),
312
+ ),
313
+ num_experts=2,
314
+ num_local_experts=2,
315
+ experts_start_idx=0,
316
+ top_k=1,
317
+ hidden_size=128,
318
+ intermediate_size=intermediate_size,
319
+ )
320
+ originals = []
321
+ for param in (
322
+ experts.w13_weight,
323
+ experts.w13_weight_scale,
324
+ experts.w2_weight,
325
+ experts.w2_weight_scale,
326
+ ):
327
+ param.data.random_(1, 128)
328
+ originals.append(param.data.clone())
329
+ monkeypatch.setattr(experts, "_check_runtime_supported", lambda: None)
330
+ monkeypatch.setattr(
331
+ "vllm.utils.deep_gemm._import_deep_gemm",
332
+ lambda: SimpleNamespace(
333
+ transform_sf_into_required_layout=lambda sf, *args: sf,
334
+ transform_weights_for_mega_moe=lambda l1, l2: (l1, l2),
335
+ ),
336
+ )
337
+
338
+ experts.finalize_weights()
339
+ padded_size = experts.intermediate_size
340
+ assert padded_size == (intermediate_size + 511) // 512 * 512
341
+ for actual, original in zip(experts._transformed_l1_weights, originals[:2]):
342
+ expected = (
343
+ original
344
+ if actual.dtype == torch.int8
345
+ else (experts._ue8m0_uint8_to_float(original))
346
+ )
347
+ actual = actual.view(expected.dtype).unflatten(1, (2, padded_size))
348
+ assert torch.equal(
349
+ actual[:, :, :intermediate_size],
350
+ expected.unflatten(1, (2, intermediate_size)),
351
+ )
352
+ assert torch.count_nonzero(actual[:, :, intermediate_size:]) == 0
353
+ for actual, original in zip(experts._transformed_l2_weights, originals[2:]):
354
+ expected = (
355
+ original
356
+ if actual.dtype == torch.int8
357
+ else (experts._ue8m0_uint8_to_float(original))
358
+ )
359
+ actual = actual.view(expected.dtype)
360
+ assert torch.equal(actual[..., : original.shape[-1]], expected)
361
+ assert torch.count_nonzero(actual[..., original.shape[-1] :]) == 0
362
+
363
+ transformed_l1 = experts._transformed_l1_weights
364
+ experts.finalize_weights()
365
+ assert experts._transformed_l1_weights is transformed_l1
366
+
367
+
368
@pytest.mark.parametrize(
173
- "hidden_size,intermediate_size,block_size,fused",
369
+ "hidden_size,intermediate_size,block_size,fused,mxfp8",
370
[
175
- pytest.param(128, 512, 128, True, id="aligned-block128"),
176
- pytest.param(5120, 2304, 32, False, id="deepseek-v41-flash"),
371
+ pytest.param(128, 512, 128, True, False, id="aligned-block128"),
372
+ pytest.param(128, 512, 128, True, True, id="aligned-block128-mxfp8"),
373
+ pytest.param(5120, 2304, 32, False, False, id="deepseek-v41-flash"),
374
],
375
)
376
def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(
180
- monkeypatch, hidden_size, intermediate_size, block_size, fused
377
+ monkeypatch, hidden_size, intermediate_size, block_size, fused, mxfp8
378
):
379
"""V4.1-Flash's routed padding must preserve the serial shared-expert fallback."""
380
@@ -261,6 +458,16 @@ def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(
458
),
459
),
460
)
461
+ if mxfp8:
462
+ for linear in (shared_experts.gate_up_proj, shared_experts.down_proj):
463
+ linear.weight_block_size = (1, 32)
464
+ linear.weight_scale = torch.nn.Parameter(
465
+ linear.weight_scale_inv.view(torch.uint8)
466
+ .repeat_interleave(128, dim=0)
467
+ .repeat_interleave(4, dim=1),
468
+ requires_grad=False,
469
+ )
470
+ del linear.weight_scale_inv
471
monkeypatch.setattr("vllm.utils.deep_gemm._import_deep_gemm", lambda: FakeDeepGemm)
472
473
original_gate_up_ptr = shared_experts.gate_up_proj.weight.data_ptr()
tests/models/test_dspark_mla.py
+76
@@ -186,3 +186,79 @@ def test_context_kv_weights_are_loaded_as_merged_linear_shards():
186
assert [weight.shard_id for _, weight in mapped] == [1, 0, 1, 1]
187
assert mapped[0][1].data_ptr() == mapped[1][1].data_ptr()
188
assert mapped[2][1].data_ptr() == mapped[3][1].data_ptr()
189
+
190
+
191
+@pytest.mark.cpu_test
192
+@pytest.mark.parametrize(
193
+ "scale_dtype", [torch.uint8, torch.float8_e8m0fnu, torch.float32]
194
+)
195
+@pytest.mark.parametrize(
196
+ ("checkpoint_name", "runtime_module", "shard_id"),
197
+ [
198
+ ("mtp.0.attn.wq_a.scale", "model.layers.0.attn.fused_wqa_wkv", 0),
199
+ ("mtp.0.attn.wkv.scale", "model.layers.0.attn.fused_wqa_wkv", 1),
200
+ ("mtp.0.main_proj.scale", "model.main_proj", None),
201
+ (
202
+ "mtp.0.ffn.shared_experts.w1.scale",
203
+ "model.layers.0.ffn.shared_experts.gate_up_proj",
204
+ 0,
205
+ ),
206
+ (
207
+ "mtp.0.ffn.shared_experts.w2.scale",
208
+ "model.layers.0.ffn.shared_experts.down_proj",
209
+ None,
210
+ ),
211
+ ],
212
+)
213
+def test_v41_dspark_loads_linear_scales(
214
+ monkeypatch, scale_dtype, checkpoint_name, runtime_module, shard_id
215
+):
216
+ """Checkpoint ``.scale`` maps to the quant method's scale parameter and
217
+ loads untouched. MXFP8 block-scale expansion lives in the KMxfp8Static
218
+ loader (see tests/quantization/test_modelopt.py), not in load_weights."""
219
+ from vllm.models.deepseek_v4_1.nvidia import dspark
220
+
221
+ mxfp8 = scale_dtype != torch.float32
222
+ scale_name = "weight_scale" if mxfp8 else "weight_scale_inv"
223
+ runtime_name = f"{runtime_module}.{scale_name}"
224
+ raw = torch.tensor([[120, 127], [128, 130]], dtype=torch.uint8)
225
+ checkpoint_scale = raw.view(scale_dtype) if mxfp8 else raw.float()
226
+ param = nn.Parameter(torch.empty_like(checkpoint_scale), requires_grad=False)
227
+ shards = []
228
+
229
+ def load_scale(param, weight, *args):
230
+ shards.append(args)
231
+ assert weight.dtype == checkpoint_scale.dtype
232
+ param.copy_(weight)
233
+
234
+ param.weight_loader = load_scale
235
+ draft = SimpleNamespace(
236
+ config=SimpleNamespace(num_attention_heads=4, n_routed_experts=1),
237
+ quant_config=SimpleNamespace(
238
+ weight_block_size=[32, 32] if mxfp8 else [128, 128]
239
+ ),
240
+ linear_scale_name=scale_name,
241
+ pad_shared_expert=False,
242
+ model=SimpleNamespace(
243
+ layers=[SimpleNamespace(ffn=SimpleNamespace(use_mega_moe=False))],
244
+ confidence_head=None,
245
+ ),
246
+ named_parameters=lambda: [(runtime_name, param)],
247
+ process_weights_after_loading=lambda: None,
248
+ )
249
+ draft._remap_dspark_name = lambda name: (
250
+ dspark.DSparkDeepseekV4ForCausalLM._remap_dspark_name(draft, name)
251
+ )
252
+ monkeypatch.setattr(dspark, "get_tensor_model_parallel_world_size", lambda: 4)
253
+ monkeypatch.setattr(dspark, "get_tensor_model_parallel_rank", lambda: 0)
254
+ monkeypatch.setattr(
255
+ dspark, "fused_moe_make_expert_params_mapping", lambda *a, **kw: []
256
+ )
257
+
258
+ loaded = dspark.DSparkDeepseekV4ForCausalLM.load_weights(
259
+ draft, [(checkpoint_name, checkpoint_scale)]
260
+ )
261
+
262
+ assert loaded == {runtime_name}
263
+ assert shards == [() if shard_id is None else (shard_id,)]
264
+ torch.testing.assert_close(param, checkpoint_scale)
tests/models/test_registry.py
+1
@@ -93,6 +93,7 @@ def test_registry_imports(model_arch):
93
[
94
("LlamaForCausalLM", False, False, "bi-encoder"),
95
("LlavaForConditionalGeneration", True, True, "bi-encoder"),
96
+ ("DeepseekV41ForCausalLM", True, False, "bi-encoder"),
97
("BertForSequenceClassification", False, False, "cross-encoder"),
98
("RobertaForSequenceClassification", False, False, "cross-encoder"),
99
("XLMRobertaForSequenceClassification", False, False, "cross-encoder"),
tests/quantization/test_fp8.py
+378
@@ -37,6 +37,7 @@ from vllm.model_executor.layers.quantization.fp8 import (
37
Fp8MoEMethod,
38
)
39
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
40
+from vllm.model_executor.layers.quantization.modelopt import ModelOptLinearMethod
41
from vllm.model_executor.layers.quantization.online.fp8 import (
42
Fp8PerTensorOnlineLinearMethod,
43
)
@@ -63,6 +64,383 @@ MODELS = [
64
]
65
66
67
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="DSV4.1 requires CUDA")
68
+@pytest.mark.parametrize("tp_rank", [0, 1])
69
+@pytest.mark.parametrize("scale_dtype", [torch.uint8, torch.float8_e8m0fnu])
70
+def test_deepseek_v41_mxfp8_scale_loading(
71
+ dist_init, default_vllm_config, monkeypatch, tp_rank, scale_dtype
72
+):
73
+ """Expand linear scales before TP slicing, including padded shared experts."""
74
+ from vllm.model_executor.layers import linear as linear_module
75
+ from vllm.model_executor.layers.linear import (
76
+ ColumnParallelLinear,
77
+ MergedColumnParallelLinear,
78
+ RowParallelLinear,
79
+ )
80
+ from vllm.models.deepseek_v4_1 import quant_config as quant_module
81
+ from vllm.models.deepseek_v4_1.nvidia import model as model_module
82
+
83
+ default_vllm_config.model_config = SimpleNamespace(dtype=torch.bfloat16)
84
+ quant_config = quant_module.DeepseekV4FP8Config(
85
+ is_checkpoint_fp8_serialized=True, weight_block_size=[32, 32]
86
+ )
87
+ quant_config._resolved_expert_dtype = "fp4"
88
+ for module in (model_module, linear_module):
89
+ monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 2)
90
+ monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: tp_rank)
91
+ tp_args = dict(quant_config=quant_config, bias=False)
92
+ block = torch.nn.Module()
93
+ block.attn = torch.nn.Module()
94
+ block.attn.wq_b = ColumnParallelLinear(128, 128, **tp_args)
95
+ block.attn.fused_wqa_wkv = MergedColumnParallelLinear(128, [128, 64], **tp_args)
96
+ block.ffn = torch.nn.Module()
97
+ block.ffn.shared_experts = torch.nn.Module()
98
+ block.ffn.shared_experts.gate_up_proj = MergedColumnParallelLinear(
99
+ 128, [128, 128], **tp_args
100
+ )
101
+ block.ffn.shared_experts.down_proj = RowParallelLinear(128, 128, **tp_args)
102
+ block.ffn.experts = torch.nn.Module()
103
+ block.ffn.experts.register_parameter(
104
+ "weight_scale", torch.nn.Parameter(torch.empty(3, 4, dtype=torch.uint8), False)
105
+ )
106
+ block.engram = torch.nn.Module()
107
+ block.engram.register_parameter(
108
+ "weight_scale_inv",
109
+ torch.nn.Parameter(torch.empty(3, 4, dtype=scale_dtype), False),
110
+ )
111
+
112
+ def load_expert_scale(param, weight, *args, **kwargs):
113
+ param.data.copy_(weight)
114
+ return True
115
+
116
+ block.ffn.experts.weight_scale.weight_loader = load_expert_scale
117
+ model = model_module.DeepseekV4Model.__new__(model_module.DeepseekV4Model)
118
+ torch.nn.Module.__init__(model)
119
+ model.layers = torch.nn.ModuleList([block])
120
+ model.config = SimpleNamespace(num_attention_heads=2)
121
+ model.quant_config = quant_config
122
+ model.use_sequence_parallel = False
123
+ model.get_expert_mapping = lambda: [
124
+ ("experts.weight_scale", "experts.0.w1.weight_scale", 0, "w1")
125
+ ]
126
+
127
+ checkpoints = [
128
+ ("attn.wq_b", "attn.wq_b", 128, 128, 0, None),
129
+ ("attn.wq_a", "attn.fused_wqa_wkv", 128, 128, 0, slice(0, 64)),
130
+ ("attn.wkv", "attn.fused_wqa_wkv", 64, 128, 0, slice(64, 96)),
131
+ (
132
+ "ffn.shared_experts.w1",
133
+ "ffn.shared_experts.gate_up_proj",
134
+ 96,
135
+ 128,
136
+ 0,
137
+ slice(0, 64),
138
+ ),
139
+ (
140
+ "ffn.shared_experts.w3",
141
+ "ffn.shared_experts.gate_up_proj",
142
+ 96,
143
+ 128,
144
+ 0,
145
+ slice(64, 128),
146
+ ),
147
+ (
148
+ "ffn.shared_experts.down_proj",
149
+ "ffn.shared_experts.down_proj",
150
+ 128,
151
+ 96,
152
+ 1,
153
+ None,
154
+ ),
155
+ ]
156
+ weights = []
157
+ expected = []
158
+ for source, target, n, k, axis, shard in checkpoints:
159
+ weight = torch.randint(-4, 5, (n, k)).to(torch.float8_e4m3fn)
160
+ scale_bytes = torch.randint(124, 131, (n // 32, k // 32), dtype=torch.uint8)
161
+ weights.extend(
162
+ [
163
+ (f"layers.0.{source}.weight", weight),
164
+ (f"layers.0.{source}.scale", scale_bytes.view(scale_dtype)),
165
+ ]
166
+ )
167
+ dequant = (
168
+ weight.float().reshape(n // 32, 32, k // 32, 32)
169
+ * torch.exp2(scale_bytes.float() - 127)[:, None, :, None]
170
+ ).reshape(n, k)
171
+ if "shared_experts" in source:
172
+ padded = torch.zeros(128, 128)
173
+ padded[:n, :k] = dequant
174
+ dequant = padded
175
+ expected.append((target, shard, dequant.chunk(2, dim=axis)[tp_rank]))
176
+
177
+ expert_scale = torch.full((3, 4), 125, dtype=torch.uint8)
178
+ weights.append(
179
+ ("layers.0.ffn.experts.0.w1.weight_scale", expert_scale.view(scale_dtype))
180
+ )
181
+ weights.append(("layers.0.engram.weight_scale_inv", expert_scale.view(scale_dtype)))
182
+ mapper = model_module._make_deepseek_v4_weights_mapper("fp4", "weight_scale")
183
+ loaded = model.load_weights(
184
+ (name.removeprefix("model."), weight) for name, weight in mapper.apply(weights)
185
+ )
186
+ assert "layers.0.ffn.experts.weight_scale" in loaded
187
+ assert torch.equal(block.ffn.experts.weight_scale, expert_scale)
188
+ assert torch.equal(block.engram.weight_scale_inv.view(torch.uint8), expert_scale)
189
+ for target, shard, reference in expected:
190
+ linear = block.get_submodule(target)
191
+ assert isinstance(linear.quant_method, ModelOptLinearMethod)
192
+ assert f"layers.0.{target}.weight_scale" in loaded
193
+ weight = linear.weight if shard is None else linear.weight[shard]
194
+ scale = linear.weight_scale if shard is None else linear.weight_scale[shard]
195
+ actual = (
196
+ weight.float().unflatten(-1, (-1, 32))
197
+ * torch.exp2(scale.float() - 127).unsqueeze(-1)
198
+ ).flatten(-2)
199
+ torch.testing.assert_close(actual, reference, rtol=0, atol=0)
200
+
201
+
202
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="DSV4.1 requires CUDA")
203
+@pytest.mark.parametrize(
204
+ "weight_block_size,expert_dtype,scale_name",
205
+ [
206
+ ([32, 32], "fp4", "weight_scale"),
207
+ ([128, 128], "fp4", "weight_scale_inv"),
208
+ ([32, 32], "fp8", "weight_scale_inv"),
209
+ ],
210
+)
211
+def test_deepseek_v41_vl_mapper_routes_linear_scales(
212
+ weight_block_size, expert_dtype, scale_name
213
+):
214
+ """The VL wrapper must map ``.scale`` keys to the parameter the linear
215
+ quant method registers, as the text model does. A hardcoded
216
+ ``weight_scale_inv`` raised KeyError for native MXFP8 checkpoints."""
217
+ from vllm.models.deepseek_v4_1.nvidia import model as model_module
218
+ from vllm.models.deepseek_v4_1.nvidia import vl_model as vl_module
219
+
220
+ vllm_config = SimpleNamespace(
221
+ quant_config=SimpleNamespace(weight_block_size=weight_block_size)
222
+ )
223
+ resolved = model_module._linear_scale_param_name(vllm_config, expert_dtype)
224
+ assert resolved == scale_name
225
+
226
+ mapper = vl_module._make_deepseek_v4_vl_weights_mapper(expert_dtype, resolved)
227
+ weight = torch.empty(0)
228
+ mapped = [
229
+ name
230
+ for name, _ in mapper.apply(
231
+ [("layers.0.attn.wq_a.scale", weight), ("layers.0.attn.wkv.weight", weight)]
232
+ )
233
+ ]
234
+ assert mapped == [
235
+ f"language_model.model.layers.0.attn.wq_a.{scale_name}",
236
+ "language_model.model.layers.0.attn.wkv.weight",
237
+ ]
238
+
239
+
240
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="DeepGEMM requires CUDA")
241
+@pytest.mark.parametrize("scale_dtype", [torch.uint8, torch.float8_e8m0fnu])
242
+@pytest.mark.parametrize(
243
+ "weight_shape,is_bmm",
244
+ [
245
+ ((129, 160), False),
246
+ ((387, 160), True),
247
+ ((3, 129, 160), False),
248
+ ((3, 128, 512), False),
249
+ ],
250
+)
251
+def test_deepgemm_mxfp8_preserves_weight_and_scale_values(
252
+ scale_dtype, weight_shape, is_bmm
253
+):
254
+ """DeepGEMM layout conversion preserves native weights and E8M0 scales."""
255
+ from vllm.model_executor.layers.quantization.utils import fp8_utils
256
+ from vllm.utils.deep_gemm import is_deep_gemm_supported
257
+
258
+ if not is_deep_gemm_supported() or not current_platform.is_device_capability_family(
259
+ 100
260
+ ):
261
+ pytest.skip("DeepGEMM MXFP8 requires Blackwell")
262
+
263
+ weight = torch.randn(weight_shape, device="cuda").to(torch.float8_e4m3fn)
264
+ scales = (
265
+ torch.randint(
266
+ 1,
267
+ 255,
268
+ (weight.numel() // 32 + 1,),
269
+ device="cuda",
270
+ dtype=torch.uint8,
271
+ )[1:]
272
+ .view(*weight_shape[:-1], weight_shape[-1] // 32)
273
+ .view(scale_dtype)
274
+ )
275
+ original_weight = weight.view(torch.uint8).clone()
276
+ original_scales = scales.view(torch.uint8).clone()
277
+ processed_weight, packed = fp8_utils.deepgemm_post_process_fp8_weight_block(
278
+ weight,
279
+ scales,
280
+ quant_block_shape=(1, 32),
281
+ use_e8m0=True,
282
+ is_bmm=is_bmm,
283
+ bmm_batch_size=3 if is_bmm else 0,
284
+ )
285
+ assert packed.dtype == torch.int32
286
+ unpacked = torch.stack(
287
+ [(packed >> (8 * i)) & 0xFF for i in range(4)], dim=-1
288
+ ).flatten(-2)
289
+ unpacked = unpacked[..., : scales.shape[-1]].to(torch.uint8)
290
+ torch.testing.assert_close(
291
+ unpacked.reshape_as(original_scales), original_scales, rtol=0, atol=0
292
+ )
293
+ torch.testing.assert_close(
294
+ processed_weight.view(torch.uint8).reshape_as(original_weight),
295
+ original_weight,
296
+ rtol=0,
297
+ atol=0,
298
+ )
299
+
300
+
301
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="DeepGEMM requires CUDA")
302
+@pytest.mark.parametrize("prequantized", [False, True])
303
+@pytest.mark.parametrize("config_source", ["deepseek", "mxfp8"])
304
+@pytest.mark.parametrize("num_tokens", [1, 7, 128])
305
+def test_mxfp8_bmm_loads_and_projects_grouped_weights(
306
+ dist_init, default_vllm_config, prequantized, config_source, num_tokens
307
+):
308
+ """BMM metadata set after construction selects grouped weight processing."""
309
+ from vllm.model_executor.kernels.linear.mxfp8.deep_gemm import (
310
+ DeepGemmMxfp8BmmLinearKernel,
311
+ )
312
+ from vllm.model_executor.layers.linear import ColumnParallelLinear
313
+ from vllm.model_executor.layers.quantization import get_quantization_config
314
+ from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import (
315
+ fused_inv_rope_fp8_quant,
316
+ )
317
+ from vllm.models.deepseek_v4.nvidia.ops.o_proj import (
318
+ compute_fp8_einsum_recipe,
319
+ deep_gemm_fp8_o_proj,
320
+ )
321
+ from vllm.models.deepseek_v4_1.nvidia.model import DeepseekV4Model
322
+ from vllm.models.deepseek_v4_1.quant_config import DeepseekV4FP8Config
323
+ from vllm.utils.deep_gemm import is_deep_gemm_supported
324
+
325
+ if not is_deep_gemm_supported() or not current_platform.is_device_capability_family(
326
+ 100
327
+ ):
328
+ pytest.skip("DeepGEMM MXFP8 BMM requires Blackwell")
329
+
330
+ default_vllm_config.model_config = SimpleNamespace(dtype=torch.bfloat16)
331
+ quant_config = DeepseekV4FP8Config(
332
+ is_checkpoint_fp8_serialized=True, weight_block_size=[32, 32]
333
+ )
334
+ quant_config._resolved_expert_dtype = "fp4"
335
+ if config_source == "mxfp8":
336
+ quant_config = get_quantization_config("mxfp8").from_config(
337
+ {"quant_method": "mxfp8"}
338
+ )
339
+ with torch.device("cuda"):
340
+ linear = ColumnParallelLinear(
341
+ 512,
342
+ 256,
343
+ bias=False,
344
+ quant_config=quant_config,
345
+ return_bias=False,
346
+ )
347
+ linear.is_bmm = True
348
+ linear.bmm_batch_size = 2
349
+ weight = torch.randn(256, 512, device="cuda").to(torch.float8_e4m3fn)
350
+ scales = torch.randint(124, 131, (8, 16), device="cuda", dtype=torch.uint8)
351
+ checkpoint_scales = (
352
+ scales.view(torch.float8_e8m0fnu)
353
+ if config_source == "deepseek"
354
+ else scales.repeat_interleave(32, dim=0)
355
+ )
356
+ model = DeepseekV4Model.__new__(DeepseekV4Model)
357
+ torch.nn.Module.__init__(model)
358
+ block = torch.nn.Module()
359
+ block.attn = torch.nn.Module()
360
+ block.attn.wo_a = linear
361
+ model.layers = torch.nn.ModuleList([block])
362
+ model.config = SimpleNamespace(num_attention_heads=2)
363
+ model.quant_config = quant_config
364
+ model.use_sequence_parallel = False
365
+ model.get_expert_mapping = lambda: []
366
+ model.load_weights(
367
+ [
368
+ ("layers.0.attn.wo_a.weight", weight),
369
+ ("layers.0.attn.wo_a.weight_scale", checkpoint_scales),
370
+ ]
371
+ )
372
+ reference_weight = (
373
+ weight.float().reshape(8, 32, 16, 32)
374
+ * torch.exp2(scales.float() - 127)[:, None, :, None]
375
+ ).reshape(2, 128, 512)
376
+ linear.quant_method.process_weights_after_loading(linear)
377
+ linear.quant_method.process_weights_after_loading(linear)
378
+ assert isinstance(linear.quant_method.kernel, DeepGemmMxfp8BmmLinearKernel)
379
+ assert linear.weight.shape == (2, 128, 512)
380
+ torch.testing.assert_close(
381
+ linear.weight.view(torch.uint8).flatten(),
382
+ weight.view(torch.uint8).flatten(),
383
+ rtol=0,
384
+ atol=0,
385
+ )
386
+ unpacked_scales = torch.stack(
387
+ [(linear.weight_scale >> (8 * i)) & 0xFF for i in range(4)], dim=-1
388
+ ).flatten(-2)
389
+ torch.testing.assert_close(
390
+ unpacked_scales.to(torch.uint8).reshape(256, 16),
391
+ scales.repeat_interleave(32, dim=0),
392
+ rtol=0,
393
+ atol=0,
394
+ )
395
+
396
+ x = torch.randn(2, num_tokens, 512, device="cuda", dtype=torch.bfloat16)
397
+ reference = torch.einsum("gmk,gnk->mgn", x.float(), reference_weight)
398
+ inputs = x.transpose(0, 1)
399
+ if prequantized:
400
+ cache = torch.cat(
401
+ (
402
+ torch.ones(num_tokens, 32, device="cuda"),
403
+ torch.zeros(num_tokens, 32, device="cuda"),
404
+ ),
405
+ dim=1,
406
+ )
407
+ inputs = fused_inv_rope_fp8_quant(
408
+ x.permute(1, 0, 2),
409
+ torch.arange(num_tokens, device="cuda"),
410
+ cache,
411
+ n_groups=2,
412
+ heads_per_group=1,
413
+ nope_dim=448,
414
+ rope_dim=64,
415
+ quant_group_size=32,
416
+ tma_aligned_scales=current_platform.has_device_capability(100),
417
+ )
418
+ with torch.no_grad():
419
+ output = linear(inputs)
420
+ assert output.shape == reference.shape
421
+ assert (output.float() - reference).norm() / reference.norm() < 0.06
422
+ if prequantized:
423
+ recipe, tma_aligned_scales = compute_fp8_einsum_recipe(block_size=32)
424
+ projected = deep_gemm_fp8_o_proj(
425
+ x.permute(1, 0, 2),
426
+ torch.arange(num_tokens, device="cuda"),
427
+ cache,
428
+ linear,
429
+ torch.nn.Identity(),
430
+ n_groups=2,
431
+ heads_per_group=1,
432
+ nope_dim=448,
433
+ rope_dim=64,
434
+ o_lora_rank=128,
435
+ einsum_recipe=recipe,
436
+ tma_aligned_scales=tma_aligned_scales,
437
+ )
438
+ torch.testing.assert_close(projected, output.flatten(1), rtol=0, atol=0)
439
+ compiled = torch.compile(linear, backend="eager", fullgraph=True)
440
+ with torch.no_grad():
441
+ torch.testing.assert_close(compiled(inputs), output, rtol=0, atol=0)
442
+
443
+
444
def test_prepare_gated_trtllm_fp8_moe_weights_pads_each_projection(monkeypatch):
445
monkeypatch.setattr(
446
flashinfer_utils,
tests/quantization/test_modelopt.py
+17
@@ -127,6 +127,23 @@ def test_modelopt_nvfp4_quantizes_parallel_lm_head():
127
assert method.spec.activation is kNvfp4Dynamic
128
129
130
+def test_modelopt_mxfp8_preserves_per_row_checkpoint_scales(dist_init, monkeypatch):
131
+ """Standard MXFP8 checkpoints already have one scale row per weight row."""
132
+ from vllm.model_executor.layers.linear import ReplicatedLinear
133
+
134
+ kernel = Mock()
135
+ kernel.input_quant_key.return_value = None
136
+ monkeypatch.setattr(
137
+ "vllm.model_executor.layers.quantization.modelopt.init_mxfp8_linear_kernel",
138
+ lambda **kwargs: kernel,
139
+ )
140
+ config = ModelOptMxFp8Config.from_config({"quant_method": "mxfp8"})
141
+ linear = ReplicatedLinear(64, 64, bias=False, quant_config=config)
142
+ scales = torch.arange(128, dtype=torch.uint8).reshape(64, 2)
143
+ linear.weight_scale.weight_loader(linear.weight_scale, scales)
144
+ assert torch.equal(linear.weight_scale, scales)
145
+
146
+
147
def test_modelopt_fp8_updates_weight_dims_after_transpose():
148
"""Humming reads weight.input_dim/output_dim. Swapping the
149
ModelWeightParameter for a plain Parameter drops them, so the per-tensor
tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py
+43
-1
@@ -7,19 +7,61 @@ import pytest
7
import torch
8
9
from tests.v1.attention.utils import create_vllm_config
10
+from vllm.models.deepseek_v4.sparse_mla import DeepseekV4SparseMLABackend
11
+from vllm.models.deepseek_v4_1.sparse_mla import (
12
+ DeepseekV4SparseMLABackend as DeepseekV41SparseMLABackend,
13
+)
14
from vllm.v1.attention.backend import CommonAttentionMetadata
15
from vllm.v1.attention.backends.mla.compressor_utils import (
16
CompressedSlotMappingKernel,
17
)
18
from vllm.v1.attention.backends.mla.indexer import (
19
BuildPrefillChunkMetadataKernel,
20
+ DeepseekV4IndexerBackend,
21
DeepseekV32IndexerMetadataBuilder,
22
+ DeepseekV41IndexerBackend,
23
)
24
from vllm.v1.attention.backends.mla.sparse_utils import (
25
ConvertReqIndexToGlobalIndexKernel,
26
)
21
-from vllm.v1.kv_cache_interface import MLAAttentionSpec
27
+from vllm.v1.kv_cache_interface import (
28
+ MLAAttentionSpec,
29
+ compute_layer_kv_cache_shape_bytes,
30
+)
31
from vllm.v1.worker.block_table import get_block_table_width
32
+from vllm.v1.worker.utils import select_common_block_size
33
+
34
+
35
+def test_indexer_shares_uncompressed_block_size_with_deepseek_v4_mla():
36
+ """Packed MLA/indexer groups must retain 64 compressed rows per page."""
37
+ kernel_block_size = select_common_block_size(
38
+ 256, [DeepseekV4SparseMLABackend, DeepseekV4IndexerBackend]
39
+ )
40
+ spec = MLAAttentionSpec(
41
+ block_size=256,
42
+ num_kv_heads=1,
43
+ head_size=132,
44
+ dtype=torch.uint8,
45
+ tokens_per_state=4,
46
+ )
47
+ assert compute_layer_kv_cache_shape_bytes(spec, 2, kernel_block_size) == (
48
+ 2,
49
+ 1,
50
+ 64,
51
+ 132,
52
+ )
53
+
54
+
55
+def test_indexer_preserves_deepseek_v41_mla_block_size():
56
+ """V4.1 retains its smaller pages independently of V4's C4 indexer."""
57
+ block_size = DeepseekV41SparseMLABackend.get_supported_kernel_block_sizes()[0]
58
+ assert isinstance(block_size, int)
59
+ assert (
60
+ select_common_block_size(
61
+ block_size, [DeepseekV41SparseMLABackend, DeepseekV41IndexerBackend]
62
+ )
63
+ == block_size
64
+ )
65
66
67
def test_indexer_warmup_normalizes_zero_compress_ratios():
tests/v1/core/test_contiguous_kv_packing.py
+102
@@ -398,6 +398,56 @@ class TestCSALinearGrouping:
398
assert len(gdn) == 10
399
400
401
+class TestSlidingWindowBucketCap:
402
+ def test_sliding_window_bucket_is_capped_at_the_main_page(self):
403
+ """DeepSeek-V4.1 shape: 43 SlidingWindowMLASpec SWA caches beside an
404
+ unbalanced 8-layer paged MLA bucket. Left whole, the SWA bucket would
405
+ set the block width and the MLA group would fill under a third of
406
+ every block; capped, it splits into groups no wider than the MLA
407
+ page."""
408
+ config = _shared_layout_config()
409
+ specs: dict[str, KVCacheSpec] = {}
410
+ for layer, ratio in ((2, 2), (8, 2), (14, 2), (20, 1)):
411
+ specs[f"layers.{layer}.attn"] = MLAAttentionSpec(
412
+ block_size=128,
413
+ num_kv_heads=1,
414
+ head_size=584,
415
+ dtype=torch.uint8,
416
+ tokens_per_state=ratio,
417
+ alignment=576,
418
+ )
419
+ specs[f"layers.{layer}.idx"] = MLAAttentionSpec(
420
+ block_size=128,
421
+ num_kv_heads=1,
422
+ head_size=132,
423
+ dtype=torch.uint8,
424
+ tokens_per_state=ratio,
425
+ alignment=576,
426
+ )
427
+ for layer in range(43):
428
+ specs[f"layers.{layer}.swa"] = SlidingWindowMLASpec(
429
+ block_size=64,
430
+ num_kv_heads=1,
431
+ head_size=584,
432
+ dtype=torch.uint8,
433
+ sliding_window=128,
434
+ alignment=576,
435
+ )
436
+ groups = get_kv_cache_groups(config, specs)
437
+ pages = _pages(groups)
438
+ main = next(g for g in groups if "layers.2.attn" in g.layer_names)
439
+ main_bytes = sum(pages[n] for n in main.layer_names)
440
+ assert _get_kv_cache_bytes_per_block(groups) == main_bytes
441
+ swa_groups = [g for g in groups if g.layer_names[0].endswith(".swa")]
442
+ per_group = main_bytes // pages["layers.0.swa"]
443
+ assert len(swa_groups) == -(-43 // per_group)
444
+ for g in swa_groups:
445
+ assert sum(pages[n] for n in g.layer_names) <= main_bytes
446
+ assert sorted(n for g in swa_groups for n in g.layer_names) == sorted(
447
+ f"layers.{i}.swa" for i in range(43)
448
+ )
449
+
450
+
451
class TestDensePacking:
452
def test_bytes_per_block_is_largest_group(self):
453
groups, g1, g2 = _mixed_page_groups()
@@ -548,3 +598,55 @@ class TestDensePacking:
598
599
if __name__ == "__main__":
600
pytest.main([__file__, "-v"])
601
+
602
+
603
+class TestCompressorRingGroup:
604
+ def test_ring_beside_paged_groups_keeps_the_kv_block_as_scheduler_block(self):
605
+ """DeepSeek-V4.1 compressor state: a CircularBufferSpec ring (capacity
606
+ 8) beside 128-token paged MLA groups and 64-token SWA groups. The ring
607
+ claims one block per request, never hashes, and does not disturb the
608
+ scheduler/hash block sizes the paged groups imply."""
609
+ config = _shared_layout_config()
610
+ config.cache_config.enable_prefix_caching = True
611
+ specs: dict[str, KVCacheSpec] = {
612
+ "layers.2.attn": MLAAttentionSpec(
613
+ block_size=128,
614
+ num_kv_heads=1,
615
+ head_size=584,
616
+ dtype=torch.uint8,
617
+ tokens_per_state=2,
618
+ alignment=576,
619
+ ),
620
+ "layers.2.attn.compressor.state_cache": CircularBufferSpec(
621
+ block_size=8,
622
+ num_kv_heads=1,
623
+ head_size=1024,
624
+ head_size_v=0,
625
+ dtype=torch.float32,
626
+ ),
627
+ }
628
+ for layer in range(4):
629
+ specs[f"layers.{layer}.attn.swa_cache"] = SlidingWindowMLASpec(
630
+ block_size=64,
631
+ num_kv_heads=1,
632
+ head_size=584,
633
+ dtype=torch.uint8,
634
+ sliding_window=128,
635
+ alignment=576,
636
+ )
637
+ groups = get_kv_cache_groups(config, specs)
638
+ ring_groups = [g for g in groups if "state_cache" in g.layer_names[0]]
639
+ assert len(ring_groups) == 1
640
+ assert not ring_groups[0].kv_cache_spec.prefix_cacheable
641
+ kv_cache_config = get_kv_cache_config_from_groups(
642
+ config, groups, available_memory=64 * _get_kv_cache_bytes_per_block(groups)
643
+ )
644
+ assert resolve_kv_cache_block_sizes(kv_cache_config, config) == (128, 64)
645
+ manager = KVCacheManager(
646
+ generate_scheduler_kv_cache_config([kv_cache_config]),
647
+ max_model_len=8192,
648
+ enable_caching=True,
649
+ hash_block_size=64,
650
+ scheduler_block_size=128,
651
+ )
652
+ assert len(manager.coordinator.single_type_managers) == len(groups)
tests/v1/cudagraph/test_breakable_cudagraph.py
+51
@@ -333,6 +333,57 @@ def test_decorator_breaks_when_invoked_inside_capture(cuda_capture_stream):
333
assert torch.equal(x, torch.full((4,), 15.0, device="cuda"))
334
335
336
+@pytest.mark.parametrize("keyword", [False, True])
337
+def test_quantized_activation_replay_does_not_retain_tensors(
338
+ cuda_capture_stream, keyword
339
+):
340
+ """Release the producer tensors while preserving both buffers on replay."""
341
+ import gc
342
+ import weakref
343
+
344
+ from vllm.compilation.breakable_cudagraph import (
345
+ BreakableCUDAGraphCapture,
346
+ eager_break_during_capture,
347
+ )
348
+ from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
349
+ from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic
350
+
351
+ @eager_break_during_capture
352
+ def consume(x, data_out, scale_out):
353
+ data_out.copy_(x.data)
354
+ scale_out.copy_(x.scale)
355
+
356
+ source = torch.ones((128, 128), device="cuda", dtype=torch.bfloat16)
357
+ data_out = source.to(torch.float8_e4m3fn)
358
+ scale_out = source[:, ::32].to(torch.uint8)
359
+ output = torch.empty_like(source)
360
+ cap = BreakableCUDAGraphCapture()
361
+ with cap:
362
+ data = source.to(torch.float8_e4m3fn)
363
+ scale = source[:, ::32].to(torch.uint8)
364
+ activation = QuantizedActivation(
365
+ data, scale, source.dtype, source.shape, kMxfp8Dynamic
366
+ )
367
+ data_ref, scale_ref = weakref.ref(data), weakref.ref(scale)
368
+ if keyword:
369
+ consume(x=activation, data_out=data_out, scale_out=scale_out)
370
+ else:
371
+ consume(activation, data_out, scale_out)
372
+ output.copy_(data_out)
373
+
374
+ del data, scale, activation
375
+ gc.collect()
376
+ assert data_ref() is None
377
+ assert scale_ref() is None
378
+
379
+ for value in (1, 2, 3):
380
+ source.fill_(value)
381
+ cap.replay()
382
+ cuda_capture_stream.synchronize()
383
+ torch.testing.assert_close(output, source, rtol=0, atol=0)
384
+ torch.testing.assert_close(scale_out, source[:, ::32].to(torch.uint8))
385
+
386
+
387
def test_eager_attention_inside_multistream_overlap(cuda_capture_stream):
388
"""Handle an eager attention break inside a multi-stream overlap region."""
389
from vllm.compilation.breakable_cudagraph import (
tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py
+99
@@ -29,6 +29,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.worker import (
29
)
30
from vllm.v1.core.kv_cache_utils import BlockHash
31
from vllm.v1.kv_cache_interface import (
32
+ CircularBufferSpec,
33
FullAttentionSpec,
34
KpoolTailSpec,
35
KVCacheConfig,
@@ -798,3 +799,101 @@ def test_worker_setup_tolerates_finer_scratch_group():
799
# The scratch group's namespace never enters the store.
800
scratch_prefix = worker.token_dbs[2].key_for(hs[0]).rsplit("@", 1)[0]
801
assert not any(key.startswith(scratch_prefix) for key in store._data)
802
+
803
+
804
+def test_ring_scratch_group_is_never_stored_and_does_not_block_hits():
805
+ """A per-request ring group (CircularBufferSpec, capacity as block_size)
806
+ sits beside the paged group: worker setup tolerates its block size (the
807
+ DeepSeek-V4.1 compressor ring is 8 rows), saving a request stores paged
808
+ blocks only, and lookup hits on the paged group alone."""
809
+ full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)
810
+ ring = CircularBufferSpec(
811
+ block_size=8, num_kv_heads=1, head_size=64, head_size_v=0, dtype=torch.uint8
812
+ )
813
+ cfg = KVCacheConfig(
814
+ num_blocks=4,
815
+ kv_cache_tensors=[
816
+ KVCacheTensor(
817
+ size=4 * full.page_size_bytes,
818
+ layers=["L0"],
819
+ layer_stride=4 * full.page_size_bytes,
820
+ block_stride=full.page_size_bytes,
821
+ ),
822
+ KVCacheTensor(
823
+ size=4 * ring.page_size_bytes,
824
+ layers=["L1"],
825
+ layer_stride=4 * ring.page_size_bytes,
826
+ block_stride=ring.page_size_bytes,
827
+ ),
828
+ ],
829
+ kv_cache_groups=[
830
+ KVCacheGroupSpec(["L0"], full),
831
+ KVCacheGroupSpec(["L1"], ring),
832
+ ],
833
+ )
834
+ vllm_config = _minimal_vllm_config(cache_block_size=16)
835
+ store = _DictStore()
836
+
837
+ worker = _build_worker_with_dict_store(vllm_config, cfg, store)
838
+ worker.tp_size = 1
839
+ worker.pp_size = 1
840
+ worker.num_kv_head = 8
841
+ assert worker.token_dbs[1].hash_block_size == 8
842
+
843
+ raw_full = torch.zeros(4 * full.page_size_bytes, dtype=torch.int8)
844
+ raw_ring = torch.zeros(4 * ring.page_size_bytes, dtype=torch.int8)
845
+ kv_caches = {
846
+ "L0": dense_kv_cache_views(raw_full, full, 4, 1, KVCacheLayout.LBNHC)[0],
847
+ "L1": dense_kv_cache_views(raw_ring, ring, 4, 1, KVCacheLayout.LBNHC)[0],
848
+ }
849
+
850
+ def _fake_thread_init(*args, **kwargs):
851
+ for v in list(args) + list(kwargs.values()):
852
+ if isinstance(v, threading.Event):
853
+ v.set()
854
+ m = MagicMock()
855
+ m.start = lambda: None
856
+ return m
857
+
858
+ with (
859
+ patch.object(
860
+ mooncake_store_worker,
861
+ "KVCacheStoreSendingThread",
862
+ side_effect=_fake_thread_init,
863
+ ),
864
+ patch.object(
865
+ mooncake_store_worker,
866
+ "KVCacheStoreRecvingThread",
867
+ side_effect=_fake_thread_init,
868
+ ),
869
+ ):
870
+ worker.register_kv_caches(kv_caches)
871
+
872
+ send_thread = KVCacheStoreSendingThread(
873
+ store=store,
874
+ token_databases=worker.token_dbs,
875
+ block_size=worker.block_size,
876
+ coord=worker.coord,
877
+ tp_rank=worker.tp_rank,
878
+ group_put_steps=worker._group_tp_replication_factors,
879
+ kv_role=worker.kv_role,
880
+ ready_event=threading.Event(),
881
+ enable_kv_event=False,
882
+ group_participates=[True, False],
883
+ )
884
+ hs = [BlockHash(bytes([i + 1]) * 4) for i in range(4)]
885
+ save_req = ReqMeta(
886
+ req_id="r0",
887
+ token_len_chunk=64,
888
+ # The ring group holds one block for the request's lifetime.
889
+ block_ids=([1, 2, 3, 4], [1]),
890
+ block_hashes=hs,
891
+ can_save=True,
892
+ store_job_id=1,
893
+ )
894
+ send_thread.add_request(save_req)
895
+ send_thread._handle_request(send_thread.request_queue.get())
896
+ worker.store = store
897
+
898
+ assert worker.lookup(num_tokens=65, block_hashes=hs).hit_length == 64
899
+ assert not any("@group:1" in key for key in store._data)
tests/v1/kv_connector/unit/test_nixl_desc_geometry.py
+109
@@ -1251,3 +1251,112 @@ def test_csa_linear_remote_ple_is_copied_whole():
1251
tp_ratio=-2,
1252
transfer_info=SimpleNamespace(remote_physical_blocks_per_logical=1),
1253
)
1254
+
1255
+
1256
+def _make_ring_worker():
1257
+ """Paged MLA group plus a per-layer ring group, no Mamba: the shape of a
1258
+ model whose sliding-window KV lives in per-request rings."""
1259
+ from unittest.mock import MagicMock
1260
+
1261
+ from vllm.config import set_current_vllm_config
1262
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw
1263
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
1264
+ NixlConnectorWorker,
1265
+ )
1266
+ from vllm.v1.kv_cache_interface import (
1267
+ CircularBufferSpec,
1268
+ KVCacheConfig,
1269
+ KVCacheGroupSpec,
1270
+ KVCacheTensor,
1271
+ MLAAttentionSpec,
1272
+ UniformTypeKVCacheSpecs,
1273
+ )
1274
+
1275
+ paged_specs = {
1276
+ f"paged.{index}": MLAAttentionSpec(
1277
+ block_size=4, num_kv_heads=1, head_size=64, dtype=torch.float16
1278
+ )
1279
+ for index in range(2)
1280
+ }
1281
+ ring_specs = {
1282
+ f"ring.{index}": CircularBufferSpec(
1283
+ block_size=6,
1284
+ num_kv_heads=1,
1285
+ head_size=32,
1286
+ head_size_v=0,
1287
+ dtype=torch.float16,
1288
+ )
1289
+ for index in range(2)
1290
+ }
1291
+
1292
+ def group(specs):
1293
+ uniform = UniformTypeKVCacheSpecs.from_specs(specs)
1294
+ assert uniform is not None
1295
+ return KVCacheGroupSpec(list(specs), uniform)
1296
+
1297
+ # Both rings overlay the second paged region, so the scratch group must
1298
+ # address only that region while the paged group spans both.
1299
+ tensor_regions = (("paged.0",), ("paged.1", "ring.0", "ring.1"))
1300
+ region_size, page_size = 512, 256
1301
+ kv_cache_config = KVCacheConfig(
1302
+ num_blocks=2,
1303
+ kv_cache_tensors=[
1304
+ KVCacheTensor(
1305
+ size=len(tensor_regions) * region_size,
1306
+ layers=[layer_name],
1307
+ layer_stride=region_size,
1308
+ block_stride=page_size,
1309
+ offset=region_index * region_size,
1310
+ )
1311
+ for region_index, layer_names in enumerate(tensor_regions)
1312
+ for layer_name in layer_names
1313
+ ],
1314
+ kv_cache_groups=[group(paged_specs), group(ring_specs)],
1315
+ )
1316
+
1317
+ vllm_config = create_vllm_config(block_size=4)
1318
+ vllm_config.cache_config.enable_prefix_caching = False
1319
+ vllm_config.kv_transfer_config.kv_buffer_device = "cuda"
1320
+ fake_backend = MagicMock()
1321
+ fake_backend.get_supported_kernel_block_sizes.return_value = [4]
1322
+ fake_backend.get_name.return_value = "TEST_ATTN"
1323
+ fake_backend.full_cls_name.return_value = "test.AttentionBackend"
1324
+ fake_platform = MagicMock()
1325
+ fake_platform.device_type = "cuda"
1326
+ fake_platform.get_nixl_memory_type.return_value = "VRAM"
1327
+
1328
+ with (
1329
+ patch.object(bw, "NixlWrapper", _RecordingNixl),
1330
+ patch.object(bw, "get_tensor_model_parallel_rank", return_value=0),
1331
+ patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1),
1332
+ patch.object(bw, "get_current_attn_backends", return_value=[fake_backend]),
1333
+ patch.object(bw, "current_platform", fake_platform),
1334
+ set_current_vllm_config(vllm_config),
1335
+ ):
1336
+ worker = NixlConnectorWorker(vllm_config, "local-engine", kv_cache_config)
1337
+ tensors = [torch.zeros((2, 256), dtype=torch.uint8) for _ in range(2)]
1338
+ worker.register_kv_caches(
1339
+ {
1340
+ layer_name: tensors[region_index]
1341
+ for region_index, layer_names in enumerate(tensor_regions)
1342
+ for layer_name in layer_names
1343
+ }
1344
+ )
1345
+ return worker
1346
+
1347
+
1348
+@pytest.mark.cpu_test
1349
+def test_ring_scratch_without_mamba_registers_and_addresses_its_own_regions():
1350
+ worker = _make_ring_worker()
1351
+
1352
+ assert worker._ple_group_index is None
1353
+ assert worker._scratch_region_indices == [1]
1354
+
1355
+ desc_ids = worker._compute_desc_ids(
1356
+ block_ids=([1], [0]),
1357
+ dst_num_blocks=2,
1358
+ block_size_ratio=None,
1359
+ physical_blocks_per_logical=1,
1360
+ )
1361
+ # Paged block 1 in both regions; ring block 0 only in its scratch region.
1362
+ assert sorted(desc_ids.tolist()) == sorted([0 * 2 + 1, 1 * 2 + 1, 1 * 2 + 0])
tests/v1/kv_connector/unit/test_simple_cpu_offload_connector.py
+57
@@ -19,6 +19,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.simple_cpu_offload_connector i
19
SimpleCPUOffloadConnector,
20
)
21
from vllm.v1.kv_cache_interface import (
22
+ CircularBufferSpec,
23
FullAttentionSpec,
24
KVCacheConfig,
25
KVCacheGroupSpec,
@@ -90,3 +91,59 @@ def test_disk_mode_rejects_non_positive_capacity() -> None:
91
"disk_capacity_bytes": 0,
92
}
93
)
94
+
95
+
96
+def _make_ring_kv_cache_config(num_blocks: int = 16) -> KVCacheConfig:
97
+ """A paged full-attention group beside a per-request ring group."""
98
+ full = FullAttentionSpec(
99
+ block_size=BLOCK_SIZE,
100
+ num_kv_heads=NUM_KV_HEADS,
101
+ head_size=HEAD_SIZE,
102
+ dtype=DTYPE,
103
+ )
104
+ ring = CircularBufferSpec(
105
+ block_size=6, num_kv_heads=1, head_size=32, head_size_v=0, dtype=DTYPE
106
+ )
107
+ return KVCacheConfig(
108
+ num_blocks=num_blocks,
109
+ kv_cache_tensors=[
110
+ KVCacheTensor(
111
+ size=_BYTES_PER_BLOCK * num_blocks,
112
+ layers=["attn"],
113
+ layer_stride=_BYTES_PER_BLOCK * num_blocks,
114
+ block_stride=_BYTES_PER_BLOCK,
115
+ ),
116
+ KVCacheTensor(
117
+ size=ring.page_size_bytes * num_blocks,
118
+ layers=["ring"],
119
+ layer_stride=ring.page_size_bytes * num_blocks,
120
+ block_stride=ring.page_size_bytes,
121
+ ),
122
+ ],
123
+ kv_cache_groups=[
124
+ KVCacheGroupSpec(["attn"], full),
125
+ KVCacheGroupSpec(["ring"], ring),
126
+ ],
127
+ )
128
+
129
+
130
+def test_ring_scratch_group_is_accepted_and_never_offloaded() -> None:
131
+ """The ring group's capacity is not a token granularity, it holds no
132
+ hashed blocks, and it never participates in CPU hits."""
133
+ vllm_config = create_vllm_config(
134
+ kv_connector="SimpleCPUOffloadConnector",
135
+ kv_role="kv_both",
136
+ kv_connector_extra_config={"cpu_bytes_to_use": _BYTES_PER_BLOCK * 8},
137
+ )
138
+ connector = SimpleCPUOffloadConnector(
139
+ vllm_config=vllm_config,
140
+ role=KVConnectorRole.SCHEDULER,
141
+ kv_cache_config=_make_ring_kv_cache_config(),
142
+ )
143
+ scheduler = connector.scheduler_manager
144
+ assert scheduler.fa_gidx == 0
145
+ cpu_groups = scheduler.cpu_kv_cache_config.kv_cache_groups
146
+ assert [g.kv_cache_spec.prefix_cacheable for g in cpu_groups] == [True, False]
147
+ # Only the paged group takes part in CPU-side hit lookup.
148
+ attention_groups = scheduler.cpu_coordinator.attention_groups
149
+ assert [g.group_ids for g in attention_groups] == [[0]]
tests/v1/simple_kv_offload/test_scheduler.py
+54
-11
@@ -1319,12 +1319,10 @@ def test_inflight_finish_deferred_cleanup() -> None:
1319
# Test 8: Null GPU blocks are skipped in store and load transfer pairs
1320
# ---------------------------------------------------------------------------
1321
def test_multi_group_null_blocks_skipped() -> None:
1322
- """Null GPU blocks (no block_hash) must not appear in store or load pairs.
1322
+ """The null block id never appears in store or load pairs.
1323
1324
- In eager store mode, _prepare_eager_store_specs skips blocks whose
1325
- block_hash is None (null blocks have no hash). We verify this by mixing
1326
- real hashed blocks with unhashed (null-like) blocks in a single group and
1327
- checking that only real blocks appear in the store list.
1324
+ A null block-table slot whose hash is still cached on the GPU is stored
1325
+ from the cached block instead; the null block itself is never copied.
1326
"""
1327
fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, num_groups=1, lazy=False)
1328
sched = fix.scheduler
@@ -1356,9 +1354,11 @@ def test_multi_group_null_blocks_skipped() -> None:
1354
f"Null block id {null_block_id} should not appear in store transfer pairs"
1355
)
1356
1359
- # Only real block should be scheduled for store
1360
- assert len(meta.store_gpu_blocks) == 1
1357
+ # The nulled slot's hash is still cached on the GPU, so that block is
1358
+ # recovered from the pool and stored alongside the real one.
1359
+ assert len(meta.store_gpu_blocks) == 2
1360
assert gpu_blocks[0].block_id in meta.store_gpu_blocks
1361
+ assert gpu_blocks[1].block_id in meta.store_gpu_blocks
1362
1363
# Complete the store
1364
assert meta.store_event >= 0
@@ -1374,13 +1374,12 @@ def test_multi_group_null_blocks_skipped() -> None:
1374
block_hasher=req._block_hasher,
1375
)
1376
hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
1377
- # Only 1 block was stored (the real one)
1378
- assert hit_tokens == BLOCK_SIZE
1377
+ assert hit_tokens == 2 * BLOCK_SIZE
1378
assert is_async is True
1379
1380
# Allocate new GPU blocks for the load
1382
- gpu_blocks2 = gpu_pool.get_new_blocks(1)
1383
- kv_blocks2 = KVCacheBlocks(blocks=([gpu_blocks2[0], null_block],))
1381
+ gpu_blocks2 = gpu_pool.get_new_blocks(2)
1382
+ kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,))
1383
sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens)
1384
1385
sched_out2 = make_scheduler_output({req2.request_id: 1})
@@ -2800,3 +2799,47 @@ def test_boundary_handoff_keeps_block_meta_index_parallel() -> None:
2799
assert req_ids == [req.request_id]
2800
assert block_meta is not None
2801
assert len(block_meta) == len(gpu_ids) == len(cpu_ids)
2802
+
2803
+
2804
+def test_finished_eager_store_recovers_nulled_window_tail() -> None:
2805
+ """A sliding-window group nulls pages that left the window before the
2806
+ connector sees the finished request's block table, while the GPU keeps
2807
+ them hashed in its free queue. The eager store must copy those pages, or
2808
+ the CPU lookup (which needs the whole window) misses every time."""
2809
+ fix = make_scheduler(num_cpu_blocks=16, num_gpu_blocks=32, num_groups=2, lazy=False)
2810
+ sched = fix.scheduler
2811
+ gpu_pool = fix.gpu_block_pool
2812
+ null_block = gpu_pool.null_block
2813
+
2814
+ num_blocks = 4 # == sliding window of the test SWA group
2815
+ req = make_request(num_blocks=num_blocks)
2816
+ fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0)
2817
+ swa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=1)
2818
+ # remove_skipped_blocks() ran: the two oldest window pages are freed (still
2819
+ # hashed) and their slots hold the null block.
2820
+ gpu_pool.free_blocks(swa_blocks[:2])
2821
+ swa_table = [null_block, null_block, *swa_blocks[2:]]
2822
+ kv_blocks = KVCacheBlocks(blocks=(fa_blocks, swa_table))
2823
+ req.num_computed_tokens = num_blocks * BLOCK_SIZE
2824
+ sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
2825
+
2826
+ block_ids = kv_blocks.get_block_ids()
2827
+ sched.request_finished_all_groups(req, block_ids)
2828
+ meta = sched.build_connector_meta(make_scheduler_output({}))
2829
+
2830
+ assert null_block.block_id not in meta.store_gpu_blocks
2831
+ assert set(meta.store_gpu_blocks) == {b.block_id for b in (*fa_blocks, *swa_blocks)}
2832
+ assert meta.store_event >= 0
2833
+ simulate_store_completion(sched, meta.store_event)
2834
+
2835
+ req2 = Request(
2836
+ request_id="req-window-tail-load",
2837
+ prompt_token_ids=req.prompt_token_ids,
2838
+ sampling_params=req.sampling_params,
2839
+ pooling_params=None,
2840
+ mm_features=None,
2841
+ block_hasher=req._block_hasher,
2842
+ )
2843
+ hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
2844
+ assert hit_tokens == num_blocks * BLOCK_SIZE
2845
+ assert is_async is True
vllm/_aiter_ops.py
+6
-4
@@ -2143,13 +2143,15 @@ class rocm_aiter_ops:
2143
if not current_platform.is_rocm():
2144
return
2145
2146
- from vllm.platforms.rocm import on_gfx11
2146
+ from vllm.platforms.rocm import on_gfx11, on_gfx950
2147
2148
- if on_gfx11() and not _OPS_REGISTERED:
2148
+ # This op has self-contained Triton/C++ implementations on gfx11
2149
+ # and gfx950. Only its optional top-k fast path comes from aiter.
2150
+ if (on_gfx11() or on_gfx950()) and not _OPS_REGISTERED:
2151
direct_register_custom_op(
2152
op_name="rocm_aiter_sparse_attn_indexer",
2153
op_func=rocm_aiter_sparse_attn_indexer,
2152
- mutates_args=["topk_indices_buffer"],
2154
+ mutates_args=["topk_indices_buffer", "candidate_blocks"],
2155
fake_impl=rocm_aiter_sparse_attn_indexer_fake,
2156
dispatch_key=current_platform.dispatch_key,
2157
)
@@ -2327,7 +2329,7 @@ class rocm_aiter_ops:
2329
direct_register_custom_op(
2330
op_name="rocm_aiter_sparse_attn_indexer",
2331
op_func=rocm_aiter_sparse_attn_indexer,
2330
- mutates_args=["topk_indices_buffer"],
2332
+ mutates_args=["topk_indices_buffer", "candidate_blocks"],
2333
fake_impl=rocm_aiter_sparse_attn_indexer_fake,
2334
dispatch_key=current_platform.dispatch_key,
2335
)
vllm/compilation/breakable_cudagraph.py
+9
-7
@@ -42,6 +42,7 @@ from vllm.forward_context import (
42
is_forward_context_available,
43
)
44
from vllm.logger import init_logger
45
+from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
46
from vllm.model_executor.offloader.base import get_offloader
47
from vllm.platforms import current_platform
48
from vllm.utils.torch_utils import weak_ref_tensor, weak_ref_tensors
@@ -56,6 +57,12 @@ def is_breakable_cudagraph_enabled() -> bool:
57
F = TypeVar("F", bound=Callable[..., Any])
58
59
60
+def _weak_ref_capture_arg(arg: Any) -> Any:
61
+ if isinstance(arg, QuantizedActivation):
62
+ return arg.weak_ref()
63
+ return weak_ref_tensor(arg)
64
+
65
+
66
def eager_break_during_capture(fn: F) -> F:
67
"""Decorator that turns a custom-op Python kernel into a "break point"
68
for the breakable cudagraph capture.
@@ -105,13 +112,8 @@ def eager_break_during_capture(fn: F) -> F:
112
# Weak-ref args: strong refs in the replay lambda pin cudagraph-pool
113
# slots across batch descriptors. cudagraph owns the slot, so the
114
# weak_ref is safe to deref on replay.
108
- weak_args = tuple(
109
- weak_ref_tensor(a) if isinstance(a, torch.Tensor) else a for a in args
110
- )
111
- weak_kwargs = {
112
- k: weak_ref_tensor(v) if isinstance(v, torch.Tensor) else v
113
- for k, v in kwargs.items()
114
- }
115
+ weak_args = tuple(_weak_ref_capture_arg(a) for a in args)
116
+ weak_kwargs = {k: _weak_ref_capture_arg(v) for k, v in kwargs.items()}
117
return capture.add_eager(lambda: fn(*weak_args, **weak_kwargs))
118
119
return wrapper # type: ignore[return-value]
vllm/config/engram.py
+19
-9
@@ -15,6 +15,14 @@ if TYPE_CHECKING:
15
16
logger = init_logger(__name__)
17
18
+# Architecture -> the hf_text_config field naming its n-gram layers. A model is
19
+# only configurable here if it actually has such layers to store.
20
+_NGRAM_LAYER_FIELDS = {
21
+ "DeepseekV41ForCausalLM": "engram_layer_ids",
22
+ "Qwen4ExpForCausalLM": "ple_layer_ids",
23
+ "Qwen4ExpForConditionalGeneration": "ple_layer_ids",
24
+}
25
+
26
27
def _default_cpu_offload() -> bool:
28
"""Honor the legacy environment variable only when the field is omitted."""
@@ -40,23 +48,25 @@ class EngramConfig:
48
Otherwise, each DP rank has a separate TP-sharded embedding replica."""
49
50
def verify_model_config(self, model_config: "ModelConfig | None") -> None:
43
- """Reject Engram configuration for models without supported embeddings."""
51
+ """Reject Engram configuration for models without n-gram embeddings."""
52
from vllm.platforms import current_platform
53
46
- supported_architectures = {
47
- "Qwen4ExpForCausalLM",
48
- "Qwen4ExpForConditionalGeneration",
49
- }
54
+ field = (
55
+ _NGRAM_LAYER_FIELDS.get(model_config.architecture)
56
+ if model_config is not None
57
+ else None
58
+ )
59
if (
60
model_config is None
52
- or model_config.architecture not in supported_architectures
61
+ or field is None
62
or not current_platform.is_cuda()
54
- or not getattr(model_config.hf_text_config, "ple_layer_ids", None)
63
+ or not getattr(model_config.hf_text_config, field, None)
64
):
65
raise ValueError(
66
"EngramConfig requires a model with supported Engram "
58
- "embeddings. Currently only the CUDA Qwen4Exp implementation "
59
- "with non-empty ple_layer_ids is supported."
67
+ "embeddings. Currently only the CUDA Qwen4Exp and DeepSeek "
68
+ "V4.1 implementations with non-empty n-gram layer ids are "
69
+ "supported."
70
)
71
72
def verify_parallel_config(self, parallel_config: "ParallelConfig") -> None:
vllm/config/kernel.py
+1
@@ -165,6 +165,7 @@ FLASHINFER_MOE_EP_ARCHITECTURES = frozenset(
165
{
166
"DeepseekV4ForCausalLM",
167
"DeepSeekV4MTPModel",
168
+ "DeepseekV41ForCausalLM",
169
}
170
)
171
vllm/config/speculative.py
+40
-8
@@ -679,12 +679,18 @@ class SpeculativeConfig:
679
],
680
}
681
)
682
- if hf_config.model_type == "deepseek_v4":
682
+ if hf_config.model_type in ("deepseek_v4", "deepseek_v41"):
683
+ # V4.1 has no classic-MTP draft: its checkpoints ship DSpark stages
684
+ # under ``mtp.*``, so only V4 gets an MTP architecture here. The
685
+ # DSpark path rewrites ``architectures`` itself and needs only
686
+ # ``n_predict``; ``method="mtp"`` on V4.1 is rejected below.
687
+ is_v41 = hf_config.model_type == "deepseek_v41"
688
hf_config.model_type = "deepseek_mtp"
689
n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
685
- hf_config.update(
686
- {"n_predict": n_predict, "architectures": ["DeepSeekV4MTPModel"]}
687
- )
690
+ overrides = {"n_predict": n_predict}
691
+ if not is_v41:
692
+ overrides["architectures"] = ["DeepSeekV4MTPModel"]
693
+ hf_config.update(overrides)
694
if hf_config.model_type in ("pangu_ultra_moe"):
695
hf_config.model_type = "pangu_ultra_moe_mtp"
696
if hf_config.model_type == "pangu_ultra_moe_mtp":
@@ -1335,6 +1341,18 @@ class SpeculativeConfig:
1341
MTPModelTypes
1342
):
1343
self.method = "mtp"
1344
+ if (
1345
+ self.target_model_config is not None
1346
+ and self.target_model_config.hf_config.model_type
1347
+ == "deepseek_v41"
1348
+ ):
1349
+ raise ValueError(
1350
+ "DeepSeek V4.1 has no classic-MTP draft: its "
1351
+ "checkpoints ship DSpark stages under mtp.* "
1352
+ "(main_proj/markov_head/confidence_head) and carry "
1353
+ "no e_proj/h_proj/enorm/hnorm/hc_head weights. Use "
1354
+ "speculative method 'dspark' instead of 'mtp'."
1355
+ )
1356
if (
1357
self.num_speculative_tokens > 1
1358
and self.draft_model_config.hf_config.model_type
@@ -1398,12 +1416,26 @@ class SpeculativeConfig:
1416
and "Gemma4DSparkModel" not in self.draft_model_config.architectures
1417
and "K3DSparkModel" not in self.draft_model_config.architectures
1418
):
1401
- # DeepSeek-V4 DSpark reuses the full DeepSeek-V4 config
1419
+ # DeepSeek-V4(.1) DSpark reuses the full target config
1420
# and its weights ship in the target checkpoint.
1403
- self.draft_model_config.hf_config.model_type = "deepseek_v4"
1404
- self.draft_model_config.hf_config.architectures = [
1405
- "DSparkDraftModel"
1421
+ is_v41 = (
1422
+ self.target_model_config.hf_config.model_type == "deepseek_v41"
1423
+ )
1424
+ draft_hf_config = self.draft_model_config.hf_config
1425
+ draft_hf_config.model_type = (
1426
+ "deepseek_v41" if is_v41 else "deepseek_v4"
1427
+ )
1428
+ draft_hf_config.architectures = [
1429
+ "DSparkV41DraftModel" if is_v41 else "DSparkDraftModel"
1430
]
1431
+ if is_v41:
1432
+ # hf_config_override set n_predict to the number of
1433
+ # MTP stages (3), but one DSpark round drafts
1434
+ # dspark_block_size tokens; num_speculative_tokens
1435
+ # divisibility is checked against n_predict below.
1436
+ draft_hf_config.n_predict = getattr(
1437
+ draft_hf_config, "dspark_block_size", None
1438
+ ) or getattr(draft_hf_config, "n_predict", None)
1439
self.draft_model_config.quantization = (
1440
self.target_model_config.quantization
1441
)
vllm/config/vllm.py
+1
@@ -81,6 +81,7 @@ DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = frozenset(
81
"DeepseekV4ForCausalLM",
82
"DeepseekV4ForConditionalGeneration",
83
"DeepSeekV4MTPModel",
84
+ "DeepseekV41ForCausalLM",
85
"Dots3NoteForCausalLM",
86
"Dots3NoteMTPModel",
87
"Glm5NextForCausalLM",
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py
+1
@@ -91,6 +91,7 @@ class MooncakeStoreCoordinator:
91
assert all(
92
scheduler_block_size % g.kv_cache_spec.block_size == 0
93
for g in kv_cache_groups
94
+ if g.kv_cache_spec.prefix_cacheable
95
), "scheduler_block_size must be a multiple of each group's block_size"
96
self.kv_cache_groups = kv_cache_groups
97
self.mamba_group_ids = {
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+5
-3
@@ -460,11 +460,13 @@ class NixlBaseConnectorWorker:
460
for spec in iter_layer_specs(group.kv_cache_spec)
461
)
462
]
463
- if len(ple_groups) != 1 or len(ple_groups[0][1].layer_names) != 1:
463
+ if len(ple_groups) > 1 or any(
464
+ len(group.layer_names) != 1 for _, group in ple_groups
465
+ ):
466
raise ValueError(
465
- "CSA-linear NIXL requires exactly one PLE cache owner."
467
+ "CSA-linear NIXL requires at most one PLE cache owner."
468
)
467
- self._ple_group_index = ple_groups[0][0]
469
+ self._ple_group_index = ple_groups[0][0] if ple_groups else None
470
471
if self._has_mamba:
472
assert self._is_hma_required
vllm/model_executor/layers/fusion/quant_activation.py
+11
-1
@@ -8,7 +8,7 @@ expose_input_quant_key; the kernel validates and reads the activation via
8
as_quantized_activation.
9
"""
10
11
-from dataclasses import dataclass
11
+from dataclasses import dataclass, replace
12
13
import torch
14
@@ -34,6 +34,16 @@ class QuantizedActivation:
34
orig_shape: torch.Size
35
quant_key: QuantKey
36
37
+ def weak_ref(self) -> "QuantizedActivation":
38
+ """Return a copy with non-owning tensor references for CUDA graph replay."""
39
+ from vllm.utils.torch_utils import weak_ref_tensor
40
+
41
+ return replace(
42
+ self,
43
+ data=weak_ref_tensor(self.data),
44
+ scale=weak_ref_tensor(self.scale),
45
+ )
46
+
47
48
def expose_input_quant_key(layer: torch.nn.Module, kernel) -> None:
49
"""Advertise the kernel's pre-quantized input key on the layer, if any.
vllm/model_executor/layers/quantization/__init__.py
+15
-2
@@ -115,7 +115,20 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
115
# lazy import to avoid triggering `torch.compile` too early
116
from vllm.config.quantization import _ONLINE_SHORTHANDS
117
from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
118
- from vllm.models.deepseek_v4 import DeepseekV4FP8Config
118
+
119
+ if current_platform.is_cuda() or current_platform.is_rocm():
120
+ # The v4.1 class is the v4 one extended to accept model_type
121
+ # "deepseek_v41" and its 32x32 MXFP8 linear layout. V4.1 has
122
+ # platform-specific implementations for both CUDA and ROCm.
123
+ from vllm.models.deepseek_v4_1 import (
124
+ DeepseekV4FP8Config as DeepseekV41FP8Config,
125
+ )
126
+
127
+ deepseek_config: type[QuantizationConfig] = DeepseekV41FP8Config
128
+ else:
129
+ from vllm.models.deepseek_v4 import DeepseekV4FP8Config
130
+
131
+ deepseek_config = DeepseekV4FP8Config
132
133
from .auto_awq import AutoAWQConfig
134
from .auto_gptq import AutoGPTQConfig
@@ -161,7 +174,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
174
"inc": INCConfig,
175
"mxfp4": Mxfp4Config,
176
"gpt_oss_mxfp4": GptOssMxfp4Config,
164
- "deepseek_v4_fp8": DeepseekV4FP8Config,
177
+ "deepseek_v4_fp8": deepseek_config,
178
"humming": HummingConfig,
179
"online": OnlineQuantizationConfig,
180
# MiniMax-style checkpoints tag `quant_method: "mxfp8"`; load with the
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+9
-17
@@ -937,14 +937,13 @@ def w8a8_triton_block_scaled_mm(
937
assert len(block_size) == 2
938
block_n, block_k = block_size[0], block_size[1]
939
940
- # Triton cannot currently bind E8M0 scale tensors directly. On ROCm,
941
- # DeepSeek-V4 checkpoints store block scales in exponent-only E8M0 format,
942
- # so decode them to fp32 before launching the kernel.
943
- if current_platform.is_rocm() or current_platform.is_xpu():
944
- if As.dtype == torch.float8_e8m0fnu:
945
- As = _upcast_e8m0_to_fp32(As).contiguous()
946
- if Bs.dtype == torch.float8_e8m0fnu:
947
- Bs = _upcast_e8m0_to_fp32(Bs).contiguous()
940
+ # Triton cannot currently bind E8M0 scale tensors directly. Checkpoints
941
+ # with exponent-only UE8M0 scales (e.g. DeepSeek-V4) are decoded to fp32
942
+ # before launching the kernel.
943
+ if As.dtype == torch.float8_e8m0fnu:
944
+ As = _upcast_e8m0_to_fp32(As).contiguous()
945
+ if Bs.dtype == torch.float8_e8m0fnu:
946
+ Bs = _upcast_e8m0_to_fp32(Bs).contiguous()
947
948
assert A.shape[-1] == B.shape[-1]
949
assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous()
@@ -1089,8 +1088,6 @@ def deepgemm_post_process_weight_scale_block(
1088
is_sfa: bool = False,
1089
) -> torch.Tensor:
1090
if ws.dtype in (torch.float8_e8m0fnu, torch.uint8):
1092
- # Scales already in E8M0 from checkpoint; upcast to fp32 and let
1093
- # DeepGEMM pack the layout expected by the target architecture.
1091
ws = _upcast_e8m0_to_fp32(ws)
1092
else:
1093
assert ws.dtype == torch.float32, (
@@ -1121,12 +1118,7 @@ def deepgemm_post_process_fp8_weight_block(
1118
f"to be torch.float8_e4m3fn, got {wq.dtype} instead."
1119
)
1120
1124
- if ws.dtype in (torch.float8_e8m0fnu, torch.uint8):
1125
- # Scales already in E8M0 from checkpoint (float8_e8m0fnu, or raw E8M0
1126
- # bits as uint8 for MXFP8) - upcast to fp32 and skip requantization
1127
- # (weights already have power-of-two scales).
1128
- ws = _upcast_e8m0_to_fp32(ws)
1129
- else:
1121
+ if ws.dtype not in (torch.float8_e8m0fnu, torch.uint8):
1122
assert ws.dtype == torch.float32, (
1123
f"Expected tensor scales dtype to be torch.float32 or "
1124
f"torch.float8_e8m0fnu or torch.uint8, got {ws.dtype} instead"
@@ -1137,7 +1129,7 @@ def deepgemm_post_process_fp8_weight_block(
1129
if is_bmm:
1130
# Reshape 2D weight/scale to 3D for grouped BMM (einsum):
1131
# wq: (g*r, d) -> (g, r, d)
1140
- # ws: (g*r/128, d/128) -> (g, r/128, d/128)
1132
+ # Scale dimensions follow quant_block_shape.
1133
g = bmm_batch_size
1134
assert wq.ndim == 2 and ws.ndim == 2
1135
d = wq.size(1)
vllm/model_executor/layers/quantization/utils/mxfp8_utils.py
+8
-2
@@ -126,8 +126,14 @@ def _mxfp8_quant_triton_kernel():
126
amax = tl.maximum(tl.max(tl.abs(x), axis=1), TINY) # [BLOCK_M]
127
sb = tl.ceil(tl.log2(amax / FP8_MAX)) + 127.0
128
sb = tl.minimum(tl.maximum(sb, 0.0), 254.0)
129
- descale = tl.exp2(sb - 127.0)
130
- xq = (x / descale[:, None]).to(xq_ptr.dtype.element_ty)
129
+ # Scale by the reciprocal rather than dividing by ``exp2(sb - 127)``:
130
+ # sb == 0 (an all-zero / denormal block) makes that divisor 2**-127,
131
+ # which is subnormal in fp32 and flushes to zero on CDNA, turning the
132
+ # block's zeros into 0/0 == NaN. The reciprocal is normal for every
133
+ # reachable sb, and both forms are exact powers of two, so nothing
134
+ # else changes.
135
+ rescale = tl.exp2(127.0 - sb)
136
+ xq = (x * rescale[:, None]).to(xq_ptr.dtype.element_ty)
137
tl.store(
138
xq_ptr + offs_m[:, None] * sqm + offs_k[None, :] * sqk,
139
xq,
vllm/model_executor/models/config.py
+3
-2
@@ -330,7 +330,7 @@ class DeepseekV4ForCausalLMConfig(VerifyAndUpdateConfig):
330
quant_config = getattr(model_config.hf_config, "quantization_config", None)
331
if quant_config is not None and quant_config.get("quant_method") == "fp8":
332
model_type = getattr(model_config.hf_config, "model_type", None)
333
- if model_type == "deepseek_v4":
333
+ if model_type in ("deepseek_v4", "deepseek_v41"):
334
model_config.hf_config.quantization_config["quant_method"] = (
335
"deepseek_v4_fp8"
336
)
@@ -343,7 +343,7 @@ class DeepseekV4ForCausalLMConfig(VerifyAndUpdateConfig):
343
and hf_text_quant_config.get("quant_method") == "fp8"
344
):
345
model_type = getattr(model_config.hf_text_config, "model_type", None)
346
- if model_type == "deepseek_v4":
346
+ if model_type in ("deepseek_v4", "deepseek_v41"):
347
model_config.hf_text_config.quantization_config["quant_method"] = (
348
"deepseek_v4_fp8"
349
)
@@ -1003,6 +1003,7 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
1003
"ColQwen3_5": ColQwen3_5Config,
1004
"DeepseekV4ForCausalLM": DeepseekV4ForCausalLMConfig,
1005
"DeepseekV4ForConditionalGeneration": DeepseekV4ForCausalLMConfig,
1006
+ "DeepseekV41ForCausalLM": DeepseekV4ForCausalLMConfig,
1007
"DeepseekV32ForCausalLM": DeepseekV32ForCausalLM,
1008
"DiffusionGemmaForBlockDiffusion": DiffusionGemmaModelForBlockDiffusionConfig, # noqa: E501
1009
"Ernie4_5_VLMoeForConditionalGeneration": Ernie4_5_VLMoeForConditionalGenerationConfig, # noqa: E501
vllm/model_executor/models/registry.py
+8
@@ -376,6 +376,10 @@ _MULTIMODAL_MODELS = {
376
"vllm.models.deepseek_v4",
377
"DeepseekV4ForConditionalGeneration",
378
),
379
+ "DeepseekV41ForCausalLM": (
380
+ "vllm.models.deepseek_v4_1",
381
+ "DeepseekV41ForCausalLM",
382
+ ),
383
"Dots3NoteForCausalLM": (
384
"vllm.models.dots3_note",
385
"Dots3NoteForCausalLM",
@@ -640,6 +644,10 @@ _SPECULATIVE_DECODING_MODELS = {
644
"MuseGlimmerAssistantModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"),
645
"DFlashMuseGlimmerAssistantModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"),
646
"DSparkDraftModel": ("vllm.models.deepseek_v4", "DSparkDeepseekV4ForCausalLM"),
647
+ "DSparkV41DraftModel": (
648
+ "vllm.models.deepseek_v4_1",
649
+ "DSparkDeepseekV4ForCausalLM",
650
+ ),
651
"Qwen3DSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"),
652
"Qwen3OmniDSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"),
653
"K3DSparkModel": (
vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py
+4
-1
@@ -173,7 +173,10 @@ def deepseek_v4_mhc_warmup(
173
# large checkpoints; bail out for any model that is not DeepSeek V4.
174
config = getattr(model, "config", None)
175
model_type = getattr(config, "model_type", None) if config is not None else None
176
- if model_type is not None and model_type != "deepseek_v4":
176
+ if model_type is not None and model_type not in (
177
+ "deepseek_v4",
178
+ "deepseek_v41",
179
+ ):
180
return
181
182
layer = _find_first_mhc_layer(model)
vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py
+6
-1
@@ -28,15 +28,20 @@ _DEEPSEEK_V4_SPARSE_MLA_BACKENDS = frozenset(
28
"FLASHMLA_SPARSE_DSV4",
29
"FLASHINFER_MLA_SPARSE_DSV4",
30
"ROCM_FLASHMLA_SPARSE_DSV4",
31
+ "FLASHMLA_SPARSE_DSV41",
32
+ "FLASHINFER_MLA_SPARSE_DSV41",
33
"DEEPSEEK_SPARSE_SWA",
34
}
35
)
36
_FLASHINFER_MLA_SPARSE_BACKENDS = frozenset({"FLASHINFER_MLA_SPARSE_SM120"})
35
-_DEEPSEEK_V4_FLASHINFER_MLA_SPARSE_BACKENDS = frozenset({"FLASHINFER_MLA_SPARSE_DSV4"})
37
+_DEEPSEEK_V4_FLASHINFER_MLA_SPARSE_BACKENDS = frozenset(
38
+ {"FLASHINFER_MLA_SPARSE_DSV4", "FLASHINFER_MLA_SPARSE_DSV41"}
39
+)
40
41
_FLASHINFER_SM120_SPARSE_MLA_DECODE_LABELS = {
42
"FLASHINFER_MLA_SPARSE_SM120": "DSv3.2",
43
"FLASHINFER_MLA_SPARSE_DSV4": "DSv4",
44
+ "FLASHINFER_MLA_SPARSE_DSV41": "DSv4.1",
45
}
46
47
_SPARSE_MLA_MIXED_WARMUP_TOKENS = 16
vllm/models/deepseek_v4_1/attention.py
+2
-2
@@ -58,7 +58,7 @@ from vllm.utils.multi_stream_utils import (
58
)
59
from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
60
from vllm.v1.attention.backends.mla.indexer import (
61
- DeepseekV4IndexerBackend,
61
+ DeepseekV41IndexerBackend,
62
dsa_indexer_uses_fp4,
63
get_max_prefill_buffer_size,
64
)
@@ -1005,7 +1005,7 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase):
1005
def forward(self): ...
1006
1007
def get_attn_backend(self) -> type[AttentionBackend]:
1008
- return DeepseekV4IndexerBackend
1008
+ return DeepseekV41IndexerBackend
1009
1010
1011
class DeepseekV4Indexer(nn.Module):
vllm/v1/attention/backends/mla/indexer.py
+14
-3
@@ -245,9 +245,20 @@ class DeepseekV4IndexerBackend(DeepseekV32IndexerBackend):
245
246
@staticmethod
247
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
248
+ # Block sizes count uncompressed tokens: C4 indexer pages hold 64 rows.
249
return [256]
250
251
252
+class DeepseekV41IndexerBackend(DeepseekV4IndexerBackend):
253
+ @staticmethod
254
+ def get_name() -> str:
255
+ return "DEEPSEEK_V41_INDEXER"
256
+
257
+ @staticmethod
258
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
259
+ return [64 if current_platform.is_device_capability_family(90) else 128]
260
+
261
+
262
@dataclass
263
class DeepseekV32IndexerPrefillChunkMetadata:
264
block_table: torch.Tensor
@@ -1239,9 +1250,9 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
1250
)
1251
block_table = self.indexer_decode_block_table_buffer[:rows, :cols]
1252
1242
- seq_lens_is_buffer_view = (use_native and next_n > 1) or (
1243
- not use_native and max_decode_len > 1
1244
- )
1253
+ # Flattening always returns a buffer view, including single-token
1254
+ # batches. Keep its address stable across varlen graph replays.
1255
+ seq_lens_is_buffer_view = not use_native or next_n > 1
1256
1257
# DCP: localize the now-expanded per-token global bounds to this
1258
# rank's owned KV. Done here (after expansion) so each token's global
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+107
-5
@@ -11,7 +11,7 @@ import torch.nn.functional as F
11
12
import vllm.envs as envs
13
from vllm.compilation.breakable_cudagraph import eager_break_during_capture
14
-from vllm.config import CUDAGraphMode
14
+from vllm.config import CUDAGraphMode, get_current_vllm_config
15
from vllm.forward_context import get_forward_context
16
from vllm.logger import init_logger
17
from vllm.platforms import current_platform
@@ -819,6 +819,32 @@ def rocm_fp8_mqa_logits(
819
return fp8_mqa_logits_torch(q, kv, weights, cu_seqlen_ks, cu_seqlen_ke)
820
821
822
+def _max_decode_logits_rows(num_batched_tokens: int) -> int:
823
+ """Upper bound on decode rows the paged-MQA logits buffer can ever hold.
824
+
825
+ ``rocm_fp8_paged_mqa_logits`` sizes its workspace as
826
+ ``(batch_size * next_n, max_model_len)``. ``batch_size`` is bounded by
827
+ ``max_num_seqs`` and ``next_n`` by ``1 + num_speculative_tokens``, which is
828
+ far tighter than ``max_num_batched_tokens`` -- 192 vs 16384 for a typical
829
+ 32-seq DSpark-5 deployment. The loose bound is harmless at short contexts
830
+ but scales with ``max_model_len``, so at the model's full context it asks
831
+ for tens of TiB and the engine cannot start. Take whichever valid bound is
832
+ smaller; the workspace is locked after profiling, so it must not be under-
833
+ estimated.
834
+ """
835
+ try:
836
+ vllm_config = get_current_vllm_config()
837
+ except Exception:
838
+ return num_batched_tokens
839
+ scheduler_config = getattr(vllm_config, "scheduler_config", None)
840
+ max_num_seqs = getattr(scheduler_config, "max_num_seqs", None)
841
+ if not max_num_seqs:
842
+ return num_batched_tokens
843
+ speculative_config = getattr(vllm_config, "speculative_config", None)
844
+ num_spec = getattr(speculative_config, "num_speculative_tokens", 0) or 0
845
+ return min(num_batched_tokens, max_num_seqs * (1 + num_spec))
846
+
847
+
848
def rocm_aiter_sparse_attn_indexer_fake(
849
hidden_states: torch.Tensor,
850
k_cache_prefix: LayerNameType,
@@ -835,6 +861,9 @@ def rocm_aiter_sparse_attn_indexer_fake(
861
topk_indices_buffer: torch.Tensor | None,
862
skip_k_cache_insert: bool = False,
863
compress_ratio: int = 1,
864
+ candidate_blocks: torch.Tensor | None = None,
865
+ candidate_block_size: int = 0,
866
+ candidate_write: bool = False,
867
) -> torch.Tensor:
868
return topk_indices_buffer
869
@@ -856,6 +885,9 @@ def rocm_aiter_sparse_attn_indexer(
885
topk_indices_buffer: torch.Tensor | None,
886
skip_k_cache_insert: bool = False,
887
compress_ratio: int = 1,
888
+ candidate_blocks: torch.Tensor | None = None,
889
+ candidate_block_size: int = 0,
890
+ candidate_write: bool = False,
891
) -> torch.Tensor:
892
# careful! this will be None in dummy run
893
forward_context = get_forward_context()
@@ -881,15 +913,15 @@ def rocm_aiter_sparse_attn_indexer(
913
)
914
915
# Decode logits buffer, used by rocm_fp8_paged_mqa_logits.
884
- # batch_size * next_n <= hidden_states.shape[0] == max_num_batched_tokens
916
+ decode_rows = _max_decode_logits_rows(hidden_states.shape[0])
917
if _ON_GFX942 or _ON_GFX950:
918
workspace_manager.get_simultaneous(
887
- ((hidden_states.shape[0], max_model_len), torch.float32),
919
+ ((decode_rows, max_model_len), torch.float32),
920
)
921
else:
922
workspace_manager.get_simultaneous(
923
(
892
- (q_fp8.shape[1], hidden_states.shape[0], max_model_len),
924
+ (q_fp8.shape[1], decode_rows, max_model_len),
925
torch.float32,
926
),
927
)
@@ -919,6 +951,9 @@ def rocm_aiter_sparse_attn_indexer(
951
topk_indices_buffer,
952
skip_k_cache_insert,
953
compress_ratio,
954
+ candidate_blocks,
955
+ candidate_block_size,
956
+ candidate_write,
957
)
958
layer_attn_metadata = attn_metadata[k_cache_prefix]
959
assert isinstance(layer_attn_metadata, DeepseekV32IndexerMetadata)
@@ -974,6 +1009,30 @@ def rocm_aiter_sparse_attn_indexer(
1009
chunk.cu_seqlen_ks,
1010
chunk.cu_seqlen_ke,
1011
)
1012
+ if candidate_blocks is not None:
1013
+ from vllm.model_executor.layers.sparse_attn_indexer import (
1014
+ _apply_candidate_mask,
1015
+ _select_candidate_blocks,
1016
+ )
1017
+
1018
+ chunk_candidates = candidate_blocks[chunk.token_start : chunk.token_end]
1019
+ if candidate_write:
1020
+ _select_candidate_blocks(
1021
+ logits,
1022
+ chunk.cu_seqlen_ks,
1023
+ chunk.cu_seqlen_ke,
1024
+ chunk_candidates.shape[1],
1025
+ candidate_block_size,
1026
+ chunk_candidates,
1027
+ )
1028
+ else:
1029
+ _apply_candidate_mask(
1030
+ logits,
1031
+ chunk.cu_seqlen_ks,
1032
+ chunk.cu_seqlen_ke,
1033
+ chunk_candidates,
1034
+ candidate_block_size,
1035
+ )
1036
topk_indices = topk_indices_buffer[
1037
chunk.token_start : chunk.token_end, :topk_tokens
1038
]
@@ -1041,6 +1100,37 @@ def rocm_aiter_sparse_attn_indexer(
1100
max_model_len=max_model_len,
1101
)
1102
1103
+ if candidate_blocks is not None:
1104
+ from vllm.model_executor.layers.sparse_attn_indexer import (
1105
+ _apply_candidate_mask,
1106
+ _select_candidate_blocks,
1107
+ )
1108
+
1109
+ num_rows = logits.shape[0]
1110
+ visible = decode_metadata.seq_lens.reshape(-1)
1111
+ if visible.numel() != num_rows:
1112
+ visible = visible.repeat_interleave(next_n)
1113
+ visible = visible[:num_rows].to(torch.int64)
1114
+ row_starts = torch.zeros_like(visible)
1115
+ decode_candidates = candidate_blocks[:num_rows]
1116
+ if candidate_write:
1117
+ _select_candidate_blocks(
1118
+ logits,
1119
+ row_starts,
1120
+ visible,
1121
+ decode_candidates.shape[1],
1122
+ candidate_block_size,
1123
+ decode_candidates,
1124
+ )
1125
+ else:
1126
+ _apply_candidate_mask(
1127
+ logits,
1128
+ row_starts,
1129
+ visible,
1130
+ decode_candidates,
1131
+ candidate_block_size,
1132
+ )
1133
+
1134
topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
1135
num_rows = logits.shape[0]
1136
@@ -1097,6 +1187,10 @@ def _decode_e8m0_scales(scale: torch.Tensor) -> torch.Tensor:
1187
)
1188
1189
return _upcast_e8m0_to_fp32(scale).contiguous()
1190
+ if scale.dtype == torch.uint8:
1191
+ # MXFP8 parameters preserve E8M0 scales as their raw exponent bytes.
1192
+ # They are biased exponents, not numeric uint8 scale values.
1193
+ return torch.exp2(scale.to(torch.int16).to(torch.float32) - 127.0)
1194
return scale.to(torch.float32)
1195
1196
@@ -1226,7 +1320,15 @@ def _get_cached_wo_a_bf16(
1320
)
1321
1322
wo_a_scale_param = get_fp8_block_weight_scale(wo_a)
1229
- if wo_a_scale_param is not None:
1323
+ if wo_a_scale_param is None:
1324
+ # ModelOpt MXFP8 stores the multiplicative E8M0 scale without the
1325
+ # historical ``_inv`` suffix.
1326
+ wo_a_scale_param = getattr(wo_a, "weight_scale", None)
1327
+ # Emulated MXFP8 kernels can replace the original one-byte weight with an
1328
+ # already-dequantized BF16 tensor while retaining the scale attribute for
1329
+ # metadata. Applying that retained scale again would double-dequantize the
1330
+ # weight. Block scaling is only valid while the one-byte FP8 storage remains.
1331
+ if wo_a_scale_param is not None and wo_a.weight.element_size() == 1:
1332
wo_a_weight = wo_a.weight.view(n_local_groups, o_lora_rank, hidden_dim).to(
1333
torch.float32
1334
)
vllm/v1/core/kv_cache_utils.py
+23
-14
@@ -22,6 +22,7 @@ from vllm.utils.torch_utils import get_dtype_size
22
from vllm.v1.kv_cache_interface import (
23
AttentionSpec,
24
ChunkedLocalAttentionSpec,
25
+ CircularBufferSpec,
26
FullAttentionSpec,
27
HiddenStateCacheSpec,
28
KpoolTailSpec,
@@ -2013,15 +2014,22 @@ def _get_packed_kv_cache_groups(
2014
cdiv(len(names), n) * page for page, names in page_size_layers.items()
2015
)
2016
2016
- # Bytes a block must hold however the mamba buckets end up split: a mamba
2017
- # bucket can go down to one state per group, every other bucket's split is
2018
- # already fixed by the repeat pattern.
2017
+ # Bytes a block must hold however the state buckets end up split: mamba,
2018
+ # circular-buffer and unsplit sliding-window buckets can go down to one
2019
+ # state per group, every other bucket's split is fixed by the repeat pattern.
2020
+ def is_state_bucket(spec: UniformTypeKVCacheSpecs) -> bool:
2021
+ if isinstance(spec.first_spec, (MambaSpec, CircularBufferSpec)):
2022
+ return True
2023
+ return repeats_per_group is None and isinstance(
2024
+ spec.first_spec, SlidingWindowSpec
2025
+ )
2026
+
2027
anchor_bytes = max(
2028
(
2029
widest_group_bytes(
2030
page_size_layers,
2031
len(spec.kv_cache_specs)
2024
- if isinstance(spec.first_spec, MambaSpec)
2032
+ if is_state_bucket(spec)
2033
else num_groups_for(spec, balanced),
2034
)
2035
for spec, page_size_layers, balanced in bucketed
@@ -2032,10 +2040,9 @@ def _get_packed_kv_cache_groups(
2040
groups = []
2041
for spec, page_size_layers, balanced in bucketed:
2042
num_groups = num_groups_for(spec, balanced)
2035
- # `_align_hybrid_block_size` pads a mamba state up to one attention
2036
- # page, so cap a mamba group at the states a block already fits rather
2037
- # than let it widen the block.
2038
- if anchor_bytes and isinstance(spec.first_spec, MambaSpec):
2043
+ # Cap a state group at the states a block already fits rather than let
2044
+ # it widen the block.
2045
+ if anchor_bytes and is_state_bucket(spec):
2046
states_per_block = max(anchor_bytes // spec.first_spec.page_size_bytes, 1)
2047
num_groups = max(
2048
num_groups, cdiv(len(spec.kv_cache_specs), states_per_block)
@@ -2071,8 +2078,9 @@ def _is_deepseek_v4_eagle(vllm_config: VllmConfig) -> bool:
2078
if spec_config is None or not spec_config.use_eagle():
2079
return False
2080
model_config = vllm_config.model_config
2074
- return (
2075
- model_config is not None and model_config.hf_config.model_type == "deepseek_v4"
2081
+ return model_config is not None and model_config.hf_config.model_type in (
2082
+ "deepseek_v4",
2083
+ "deepseek_v41",
2084
)
2085
2086
@@ -2093,9 +2101,9 @@ def _annotate_eagle_groups(
2101
spec merging, wherever grouping happens to land. It is sufficient but
2102
not necessary: a drafter whose spec is indistinguishable from the
2103
target's cannot be found this way.
2096
- 2. Model-scoped positional fallback for DeepseekV4, whose MTP block reuses
2097
- the target's own decoder layer and so carries no spec marker. Its draft
2098
- attention layer is always the last registered layer, so flag whichever
2104
+ 2. Model-scoped positional fallback for DeepseekV4/V4.1, whose MTP block
2105
+ reuses the target's own decoder layer and so carries no spec marker. Its
2106
+ draft attention layer is always the last registered layer, so flag whichever
2107
group holds it. This rule is only valid where the groups partition
2108
exactly the layers of ``kv_cache_spec``, which is true on the packed
2109
grouping path and not in general; other callers must leave
@@ -2108,7 +2116,8 @@ def _annotate_eagle_groups(
2116
kv_cache_spec: The kv cache spec of each attention layer, in layer
2117
registration order. Only read by rule 2.
2118
kv_cache_groups: Groups to annotate in place.
2111
- use_deepseek_v4_fallback: Enable rule 2 for a DeepseekV4 packed group.
2119
+ use_deepseek_v4_fallback: Enable rule 2 for a DeepseekV4/V4.1 packed
2120
+ group.
2121
"""
2122
spec_config = vllm_config.speculative_config
2123
if spec_config is None or not spec_config.use_eagle_block_drop():
vllm/v1/simple_kv_offload/manager.py
+26
-8
@@ -51,7 +51,7 @@ from vllm.v1.simple_kv_offload.metadata import (
51
52
if TYPE_CHECKING:
53
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
54
- from vllm.v1.core.kv_cache_utils import KVCacheBlock
54
+ from vllm.v1.core.kv_cache_utils import BlockHashList, KVCacheBlock
55
from vllm.v1.kv_cache_interface import KVCacheConfig
56
from vllm.v1.request import Request
57
@@ -783,6 +783,18 @@ class SimpleCPUOffloadScheduler:
783
req_ids = list(dict.fromkeys(req_ids))
784
return merged_gpu_block_ids, merged_cpu_block_ids, req_ids, merged_block_meta
785
786
+ def _cached_gpu_block(
787
+ self, resolved_hashes: "BlockHashList", block_idx: int, group_id: int
788
+ ) -> "KVCacheBlock | None":
789
+ """Return the GPU block still cached under this group's block hash."""
790
+ if block_idx >= len(resolved_hashes):
791
+ return None
792
+ assert self._gpu_block_pool is not None
793
+ blocks = self._gpu_block_pool.get_cached_block(
794
+ resolved_hashes[block_idx], [group_id]
795
+ )
796
+ return blocks[0] if blocks else None
797
+
798
def _select_eager_blocks_to_store(
799
self,
800
state: StoreRequestState,
@@ -827,18 +839,25 @@ class SimpleCPUOffloadScheduler:
839
# the middle of the request.
840
group_size = self.group_block_sizes[g]
841
ready = min(len(group_gpu_ids), aligned_tokens // group_size)
830
- resolved_hashes = (
831
- resolve_block_hashes(
832
- request.block_hashes, self.hash_block_size, group_size
833
- )
834
- if block_meta is not None
835
- else None
842
+ resolved_hashes = resolve_block_hashes(
843
+ request.block_hashes, self.hash_block_size, group_size
844
)
845
curr_mm_idx = 0
846
secondary_mm_idx = 0
847
start = state.num_stored_blocks[g]
848
for i, gpu_block_id in enumerate(group_gpu_ids[start:ready], start=start):
849
gpu_block = self._gpu_block_pool.blocks[gpu_block_id]
850
+ if gpu_block.is_null:
851
+ # Sliding-window groups null pages that left the window
852
+ # before the connector sees the block table, but the
853
+ # retained prefix-cache tail stays hashed in the GPU free
854
+ # queue; store it from there.
855
+ recovered = self._cached_gpu_block(resolved_hashes, i, g)
856
+ if recovered is None:
857
+ advanced_per_group[g] += 1
858
+ continue
859
+ gpu_block = recovered
860
+ gpu_block_id = gpu_block.block_id
861
if (
862
self._classify_store_candidate(gpu_block)
863
is not _StoreAdmission.READY
@@ -854,7 +873,6 @@ class SimpleCPUOffloadScheduler:
873
if block_meta is not None:
874
token_start = i * group_size
875
token_end = token_start + group_size
857
- assert resolved_hashes is not None
876
parent_hash = (
877
None
878
if i == 0
vllm/v1/worker/gpu/spec_decode/dspark/utils.py
+2
-2
@@ -17,9 +17,9 @@ def _resolve_dspark_attention_backend(
17
) -> AttentionBackendEnum | None:
18
if draft_backend is not None:
19
return draft_backend
20
- # DeepSeek-V4 draft layers share the target's KV-cache layout. Other
20
+ # DeepSeek-V4(.1) draft layers share the target's KV-cache layout. Other
21
# DSpark architectures may use a different attention kind.
22
- if draft_model_config.hf_config.model_type == "deepseek_v4":
22
+ if draft_model_config.hf_config.model_type in ("deepseek_v4", "deepseek_v41"):
23
if target_backend is not None:
24
logger.info_once(
25
"Using the target model's %s attention backend for the "
vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py
+8
-1
@@ -112,7 +112,14 @@ def get_eagle3_aux_layers_from_config(
112
if not layer_ids:
113
dspark_layer_ids = getattr(hf_config, "dspark_target_layer_ids", None)
114
if dspark_layer_ids:
115
- layer_ids = [i + 1 for i in dspark_layer_ids]
115
+ if getattr(hf_config, "model_type", None) == "deepseek_v41":
116
+ # v4.1 reads the attention *inputs* of its target layers, and
117
+ # the target model captures the entry stream of layer L when
118
+ # idx+1 == L, so the ids are used as-is. (v4's ids are in
119
+ # capture-after semantics and keep the +1.)
120
+ layer_ids = list(dspark_layer_ids)
121
+ else:
122
+ layer_ids = [i + 1 for i in dspark_layer_ids]
123
if not layer_ids:
124
# Dense DSpark (e.g. Qwen3) also uses different aux layer semantics.
125
target_layer_ids = getattr(hf_config, "target_layer_ids", None)
vllm/v1/worker/gpu_model_runner.py
+38
-3
@@ -813,6 +813,9 @@ class GPUModelRunner(
813
814
# Persistent buffers for CUDA graphs.
815
self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32)
816
+ # [max_num_reqs, depth] prompt token ids preceding each request's chunk
817
+ # start, for models exposing `token_lookback_depth`. Sized after load.
818
+ self.lookback_token_ids: CpuGpuBuffer | None = None
819
self.positions = torch.zeros(
820
self.max_num_tokens, dtype=torch.int64, device=self.device
821
)
@@ -1065,9 +1068,36 @@ class GPUModelRunner(
1068
)
1069
return self._mamba_bufs
1070
1068
- def _init_model_kwargs(self):
1071
+ def _prepare_lookback_token_ids(self, num_reqs: int) -> torch.Tensor:
1072
+ """Gather, per request, the `depth` prompt token ids preceding its
1073
+ first scheduled token (column j is position start - 1 - j); -1 where
1074
+ the position is before the prompt or already past it. Generated
1075
+ positions are left to the model: under async scheduling the CPU token
1076
+ table holds placeholders for them."""
1077
+ buf = self.lookback_token_ids
1078
+ assert buf is not None
1079
+ buf.np.fill(-1)
1080
+ if num_reqs > 0:
1081
+ depth = buf.np.shape[1]
1082
+ starts = self.input_batch.num_computed_tokens_cpu[:num_reqs, None]
1083
+ pos = starts - np.arange(1, depth + 1)
1084
+ num_prompt = self.input_batch.num_prompt_tokens[:num_reqs, None]
1085
+ valid = (pos >= 0) & (pos < num_prompt)
1086
+ rows = np.arange(num_reqs)[:, None]
1087
+ ids = self.input_batch.token_ids_cpu[rows, np.clip(pos, 0, None)]
1088
+ buf.np[:num_reqs] = np.where(valid, ids, -1)
1089
+ return buf.copy_to_gpu()
1090
+
1091
+ def _init_model_kwargs(self, num_reqs: int | None = None):
1092
model_kwargs = dict[str, Any]()
1093
1094
+ if self.lookback_token_ids is not None:
1095
+ if num_reqs is None:
1096
+ num_reqs = self.input_batch.num_reqs
1097
+ model_kwargs["lookback_token_ids"] = self._prepare_lookback_token_ids(
1098
+ num_reqs
1099
+ )
1100
+
1101
if not self.is_pooling_model:
1102
return model_kwargs
1103
@@ -5314,6 +5344,11 @@ class GPUModelRunner(
5344
self.model = model_loader.load_model(
5345
vllm_config=self.vllm_config, model_config=self.model_config
5346
)
5347
+ lookback_depth = getattr(self.model, "token_lookback_depth", 0)
5348
+ if lookback_depth > 0:
5349
+ self.lookback_token_ids = self._make_buffer(
5350
+ self.max_num_reqs, lookback_depth, dtype=torch.int32
5351
+ )
5352
if self.lora_config:
5353
self.model = self.load_lora_model(
5354
self.model, self.vllm_config, self.device
@@ -6076,7 +6111,7 @@ class GPUModelRunner(
6111
):
6112
# Make sure padding doesn't exceed max_num_tokens
6113
assert num_tokens_padded <= self.max_num_tokens
6079
- model_kwargs = self._init_model_kwargs()
6114
+ model_kwargs = self._init_model_kwargs(num_reqs=0)
6115
if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:
6116
input_ids, inputs_embeds = self._prepare_mm_inputs(num_tokens_padded)
6117
@@ -6087,7 +6122,7 @@ class GPUModelRunner(
6122
elif self.enable_prompt_embeds:
6123
input_ids = None
6124
inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded]
6090
- model_kwargs = self._init_model_kwargs()
6125
+ model_kwargs = self._init_model_kwargs(num_reqs=0)
6126
else:
6127
input_ids = self.input_ids.gpu[:num_tokens_padded]
6128
inputs_embeds = None