@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 9b959b8657

[Model] DeepSeek-V4.1-Flash Model Definitions (#56228)

Signed-off-by: Roger Wang <hey@rogerw.io> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: Kimi Code <noreply@moonshot.cn> Co-authored-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: Codex <noreply@openai.com>

Roger Wang committed Sep 10, 2026 at 05:16 UTC 9b959b86577c082c0b2bf9e2c22263255a36ad83
47 files changed +12902 -196
tests/models/test_deepseek_v4_mega_moe.py
+47 -11
@@ -169,7 +169,18 @@ def test_deepseek_v4_mega_moe_weight_loader_uses_ep_expert_ownership():
169 assert torch.count_nonzero(experts.w13_weight[1]) == 0
170
171
172 -def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(monkeypatch):
172 +@pytest.mark.parametrize(
173 + "hidden_size,intermediate_size,block_size,fused",
174 + [
175 + pytest.param(128, 512, 128, True, id="aligned-block128"),
176 + pytest.param(5120, 2304, 32, False, id="deepseek-v41-flash"),
177 + ],
178 +)
179 +def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(
180 + monkeypatch, hidden_size, intermediate_size, block_size, fused
181 +):
182 + """V4.1-Flash's routed padding must preserve the serial shared-expert fallback."""
183 +
184 class FakeDeepGemm:
185 transformed_dims: list[tuple[int, int]] = []
186 scale_inputs: list[tuple[int, ...]] = []
@@ -216,8 +227,8 @@ def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(monkeypatch
227 num_local_experts=1,
228 experts_start_idx=0,
229 top_k=1,
219 - hidden_size=128,
220 - intermediate_size=128,
230 + hidden_size=hidden_size,
231 + intermediate_size=intermediate_size,
232 num_shared_experts=1,
233 )
234 experts._check_runtime_supported = lambda: None
@@ -232,24 +243,49 @@ def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(monkeypatch
243
244 shared_experts = SimpleNamespace(
245 gate_up_proj=SimpleNamespace(
235 - weight=fp8_parameter(256, 128),
236 - weight_block_size=(128, 128),
237 - weight_scale_inv=scale_parameter(2, 1, dtype=torch.float8_e8m0fnu),
246 + weight=fp8_parameter(2 * intermediate_size, hidden_size),
247 + weight_block_size=(block_size, block_size),
248 + weight_scale_inv=scale_parameter(
249 + 2 * intermediate_size // block_size,
250 + hidden_size // block_size,
251 + dtype=torch.float8_e8m0fnu,
252 + ),
253 ),
254 down_proj=SimpleNamespace(
240 - weight=fp8_parameter(128, 128),
241 - weight_block_size=(128, 128),
242 - weight_scale_inv=scale_parameter(1, 1, dtype=torch.float8_e8m0fnu),
255 + weight=fp8_parameter(hidden_size, intermediate_size),
256 + weight_block_size=(block_size, block_size),
257 + weight_scale_inv=scale_parameter(
258 + hidden_size // block_size,
259 + intermediate_size // block_size,
260 + dtype=torch.float8_e8m0fnu,
261 + ),
262 ),
263 )
264 monkeypatch.setattr("vllm.utils.deep_gemm._import_deep_gemm", lambda: FakeDeepGemm)
265
266 original_gate_up_ptr = shared_experts.gate_up_proj.weight.data_ptr()
267 + original_down_ptr = shared_experts.down_proj.weight.data_ptr()
268 experts.finalize_weights(shared_experts)
269
270 + assert FakeDeepGemm.scale_inputs[-2:] == [
271 + (1, 2 * intermediate_size, hidden_size // 32),
272 + (1, hidden_size, intermediate_size // 32),
273 + ]
274 + assert experts.has_fused_shared_experts is fused
275 + if not fused:
276 + assert experts.intermediate_size == 2560
277 + assert FakeDeepGemm.transformed_dims == [(3, 3)]
278 + assert experts.num_shared_experts == 0
279 + assert shared_experts.gate_up_proj.weight.data_ptr() == original_gate_up_ptr
280 + assert shared_experts.down_proj.weight.data_ptr() == original_down_ptr
281 + assert shared_experts.gate_up_proj.weight.shape == (
282 + 2 * intermediate_size,
283 + hidden_size,
284 + )
285 + assert shared_experts.down_proj.weight.shape == (hidden_size, intermediate_size)
286 + return
287 +
288 assert FakeDeepGemm.transformed_dims == [(3, 3), (2, 2)]
251 - assert FakeDeepGemm.scale_inputs[-2:] == [(1, 256, 4), (1, 128, 4)]
252 - assert experts.has_fused_shared_experts
289 assert shared_experts.gate_up_proj.weight.data_ptr() != original_gate_up_ptr
290 assert (
291 experts._transformed_shared_l1_weights[0].data_ptr()
tests/parser/engine/trace_builder.py
+40 -17
@@ -31,6 +31,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
31 from vllm.parser.engine.registered_adapters import (
32 DeepSeekV4Parser,
33 DeepSeekV32Parser,
34 + DeepSeekV41Parser,
35 Gemma4Parser,
36 Glm47MoeParser,
37 InklingParser,
@@ -658,7 +659,7 @@ _DSV4_VOCAB: dict[str, int] = {
659 }
660
661
661 -def _dsv4_param_text(key: str, value: Any) -> str:
662 +def _dsv4_param_text(key: str, value: Any, dsml: str = _DSML) -> str:
663 is_string = isinstance(value, str)
664 if is_string:
665 val_str = value
@@ -670,36 +671,39 @@ def _dsv4_param_text(key: str, value: Any) -> str:
671 val_str = json.dumps(value, ensure_ascii=False)
672 string_attr = "true" if is_string else "false"
673 return (
673 - f'<{_DSML}parameter name="{key}" string="{string_attr}">'
674 - f"{val_str}</{_DSML}parameter>\n"
674 + f'<{dsml}parameter name="{key}" string="{string_attr}">'
675 + f"{val_str}</{dsml}parameter>\n"
676 )
677
678
678 -def _dsv4_tool_text(tc: ToolCallSpec) -> str:
679 - parts = [f'<{_DSML}invoke name="{tc.name}">\n']
679 +def _dsv4_tool_text(tc: ToolCallSpec, dsml: str = _DSML) -> str:
680 + parts = [f'<{dsml}invoke name="{tc.name}">\n']
681 for key, value in tc.arguments.items():
681 - parts.append(_dsv4_param_text(key, value))
682 - parts.append(f"</{_DSML}invoke>\n")
682 + parts.append(_dsv4_param_text(key, value, dsml))
683 + parts.append(f"</{dsml}invoke>\n")
684 return "".join(parts)
685
686
687 def _dsml_tool_segs(
688 scenario: Scenario,
689 tag: str,
690 + dsml: str = _DSML,
691 ) -> list[tuple[str, bool]]:
692 if not scenario.tool_calls:
693 return []
694 parts = ["\n"]
695 for tc in scenario.tool_calls:
694 - parts.append(_dsv4_tool_text(tc))
696 + parts.append(_dsv4_tool_text(tc, dsml))
697 return [
696 - (f"<{_DSML}{tag}>", True),
698 + (f"<{dsml}{tag}>", True),
699 ("".join(parts), False),
698 - (f"</{_DSML}{tag}>", True),
700 + (f"</{dsml}{tag}>", True),
701 ]
702
703
702 -def _dsv4_segments(scenario: Scenario, thinking: bool) -> list[tuple[str, bool]]:
704 +def _dsv4_segments(
705 + scenario: Scenario, thinking: bool, *, v41: bool = False
706 +) -> list[tuple[str, bool]]:
707 segs: list[tuple[str, bool]] = []
708
709 if thinking:
@@ -716,13 +720,31 @@ def _dsv4_segments(scenario: Scenario, thinking: bool) -> list[tuple[str, bool]]
720 if scenario.content is not None:
721 segs.append((scenario.content, False))
722
719 - segs.extend(_dsml_tool_segs(scenario, "tool_calls"))
723 + segs.extend(
724 + _dsml_tool_segs(scenario, "calls", f"{_DSML} ")
725 + if v41
726 + else _dsml_tool_segs(scenario, "tool_calls")
727 + )
728 return segs
729
730
723 -def _build_deepseek_v4(scenario: Scenario, validate: bool = True) -> Sample:
731 +def _build_deepseek_v4(
732 + scenario: Scenario, validate: bool = True, *, v41: bool = False
733 +) -> Sample:
734 thinking = scenario.reasoning is not None
735 chat_kwargs = {"thinking": thinking}
736 + parser_cls = DeepSeekV41Parser if v41 else DeepSeekV4Parser
737 + model = "deepseek_v41" if v41 else "deepseek_v4"
738 + vocab = (
739 + {
740 + "<think>": 128821,
741 + "</think>": 128822,
742 + f"<{_DSML} calls>": 128823,
743 + f"</{_DSML} calls>": 128824,
744 + }
745 + if v41
746 + else _DSV4_VOCAB
747 + )
748
749 if thinking:
750 expected_reasoning: str | None = scenario.reasoning or ""
@@ -730,10 +752,10 @@ def _build_deepseek_v4(scenario: Scenario, validate: bool = True) -> Sample:
752 expected_reasoning = None
753
754 sample = _make_sample(
733 - sample_id=f"deepseek_v4-{scenario.id}",
755 + sample_id=f"{model}-{scenario.id}",
756 description=scenario.description,
735 - vocab=_DSV4_VOCAB,
736 - segments=_dsv4_segments(scenario, thinking),
757 + vocab=vocab,
758 + segments=_dsv4_segments(scenario, thinking, v41=v41),
759 expected_reasoning=expected_reasoning,
760 expected_content=_qwen3_expected_content(scenario),
761 expected_tool_calls=_expected_tc(scenario),
@@ -744,7 +766,7 @@ def _build_deepseek_v4(scenario: Scenario, validate: bool = True) -> Sample:
766 kwargs = {}
767 if chat_kwargs:
768 kwargs["chat_template_kwargs"] = chat_kwargs
747 - _validate_sample(sample, DeepSeekV4Parser, **kwargs)
769 + _validate_sample(sample, parser_cls, **kwargs)
770 return sample
771
772
@@ -1025,6 +1047,7 @@ def _build_inkling(scenario: Scenario, validate: bool = True) -> Sample:
1047 _BUILDERS: dict[str, Any] = {
1048 "deepseek_v32": _build_deepseek_v32,
1049 "deepseek_v4": _build_deepseek_v4,
1050 + "deepseek_v41": functools.partial(_build_deepseek_v4, v41=True),
1051 "gemma4": _build_gemma4,
1052 "minimax_m2": _build_minimax_m2,
1053 "nemotron_v3": _build_nemotron_v3,
vllm/model_executor/kernels/attention/dsa/candidate_blocks.py new
+234
@@ -0,0 +1,234 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +from vllm.triton_utils import tl, triton
7 +
8 +
9 +@triton.jit
10 +def _max_with_nan(a, b):
11 + return tl.maximum(a, b, propagate_nan=tl.PropagateNan.ALL)
12 +
13 +
14 +@triton.jit(do_not_specialize=["width", "nblocks"])
15 +def _block_scores_kernel(
16 + logits,
17 + starts,
18 + ends,
19 + scores,
20 + stride_row,
21 + stride_col,
22 + stride_start,
23 + stride_end,
24 + width,
25 + nblocks,
26 + BLOCK_SIZE: tl.constexpr,
27 + HAS_STARTS: tl.constexpr,
28 + ROW_REPEAT: tl.constexpr,
29 + TILE: tl.constexpr,
30 +):
31 + row = tl.program_id(0).to(tl.int64)
32 + blocks = tl.program_id(1) * TILE + tl.arange(0, TILE)
33 + start = tl.load(starts + row // ROW_REPEAT * stride_start) if HAS_STARTS else 0
34 + end = tl.load(ends + row // ROW_REPEAT * stride_end)
35 + offsets = tl.arange(0, triton.next_power_of_2(BLOCK_SIZE))
36 + cols = start + blocks[:, None] * BLOCK_SIZE + offsets[None, :]
37 + values = tl.load(
38 + logits + row * stride_row + cols * stride_col,
39 + (blocks[:, None] < nblocks)
40 + & (offsets[None, :] < BLOCK_SIZE)
41 + & (cols < end)
42 + & (cols < width),
43 + other=-float("inf"),
44 + )
45 + reduced = tl.reduce(values, 1, _max_with_nan)
46 + reduced = tl.where(
47 + (end > start) & (blocks == (end - start - 1) // BLOCK_SIZE),
48 + float("inf"),
49 + reduced,
50 + )
51 + tl.store(scores + row * nblocks + blocks, reduced, blocks < nblocks)
52 +
53 +
54 +@triton.jit(do_not_specialize=["k"])
55 +def _store_candidates_kernel(
56 + values,
57 + indices,
58 + output,
59 + out_stride_row,
60 + out_stride_col,
61 + k,
62 + OUT_K: tl.constexpr,
63 + TILE: tl.constexpr,
64 +):
65 + row = tl.program_id(0).to(tl.int64)
66 + cols = tl.program_id(1) * TILE + tl.arange(0, TILE)
67 + value = tl.load(values + row * k + cols, cols < k, other=-float("inf"))
68 + index = tl.load(indices + row * k + cols, cols < k, other=-1)
69 + tl.store(
70 + output + row * out_stride_row + cols * out_stride_col,
71 + tl.where(value > -float("inf"), index, -1),
72 + cols < OUT_K,
73 + )
74 +
75 +
76 +@triton.jit(do_not_specialize=["width", "nblocks"])
77 +def _candidate_flags_kernel(
78 + candidates,
79 + starts,
80 + flags,
81 + stride_row,
82 + stride_col,
83 + stride_start,
84 + width,
85 + nblocks,
86 + BLOCK_SIZE: tl.constexpr,
87 + K: tl.constexpr,
88 + HAS_STARTS: tl.constexpr,
89 + ROW_REPEAT: tl.constexpr,
90 +):
91 + row = tl.program_id(0).to(tl.int64)
92 + offsets = tl.arange(0, 1024)
93 + for tile in range(tl.cdiv(nblocks + 1, 1024)):
94 + slots = tile * 1024 + offsets
95 + tl.store(flags + row * (nblocks + 1) + slots, 0, slots <= nblocks)
96 + start = tl.load(starts + row // ROW_REPEAT * stride_start) if HAS_STARTS else 0
97 + cols = tl.arange(0, triton.next_power_of_2(K))
98 + block = tl.load(
99 + candidates + row * stride_row + cols * stride_col, cols < K, other=-1
100 + ).to(tl.int64)
101 + # Preserve the packed-column clamp for candidates beyond the logits width.
102 + block = tl.where(start + block * BLOCK_SIZE >= width, nblocks, block)
103 + tl.debug_barrier()
104 + tl.store(flags + row * (nblocks + 1) + block, 1, (cols < K) & (block >= 0))
105 +
106 +
107 +@triton.jit(do_not_specialize=["width", "nblocks"])
108 +def _mask_candidates_kernel(
109 + logits,
110 + starts,
111 + ends,
112 + flags,
113 + stride_row,
114 + stride_col,
115 + stride_start,
116 + stride_end,
117 + width,
118 + nblocks,
119 + BLOCK_SIZE: tl.constexpr,
120 + HAS_STARTS: tl.constexpr,
121 + ROW_REPEAT: tl.constexpr,
122 + TILE: tl.constexpr,
123 +):
124 + row = tl.program_id(0).to(tl.int64)
125 + cols = tl.program_id(1) * TILE + tl.arange(0, TILE)
126 + start = tl.load(starts + row // ROW_REPEAT * stride_start) if HAS_STARTS else 0
127 + end = tl.load(ends + row // ROW_REPEAT * stride_end)
128 + valid = (cols >= start) & (cols < end) & (cols < width)
129 + block = (cols - start) // BLOCK_SIZE
130 + keep = tl.load(flags + row * (nblocks + 1) + block, valid, other=0)
131 + edge = tl.load(flags + row * (nblocks + 1) + nblocks)
132 + keep = (keep != 0) | ((cols == width - 1) & (edge != 0))
133 + tl.store(
134 + logits + row * stride_row + cols * stride_col,
135 + -float("inf"),
136 + (cols < width) & ~(valid & keep),
137 + )
138 +
139 +
140 +def select_candidate_blocks(
141 + logits: torch.Tensor,
142 + row_ks: torch.Tensor | None,
143 + row_ke: torch.Tensor,
144 + topk_blocks: int,
145 + block_size: int,
146 + out: torch.Tensor,
147 + row_repeat: int = 1,
148 +) -> None:
149 + """Select local block IDs by maximum score, pinning each row's newest block.
150 +
151 + Row bounds are in packed column space; absent starts mean zero.
152 + Decode rows share bounds in groups of ``row_repeat``. Output is -1 padded.
153 + """
154 + assert logits.is_cuda
155 + rows, width = logits.shape
156 + if not rows:
157 + return
158 + if not width:
159 + out.fill_(-1)
160 + return
161 + nblocks = triton.cdiv(width, block_size)
162 + scores = logits.new_empty((rows, nblocks))
163 + _block_scores_kernel[(rows, triton.cdiv(nblocks, 128))](
164 + logits,
165 + row_ks,
166 + row_ke,
167 + scores,
168 + *logits.stride(),
169 + row_ks.stride(0) if row_ks is not None else 0,
170 + row_ke.stride(0),
171 + width,
172 + nblocks,
173 + block_size,
174 + row_ks is not None,
175 + row_repeat,
176 + 128,
177 + )
178 + # Keep the existing top-k tie behavior.
179 + top = scores.topk(min(topk_blocks, nblocks), dim=-1)
180 + _store_candidates_kernel[(rows, triton.cdiv(topk_blocks, 256))](
181 + top.values,
182 + top.indices,
183 + out,
184 + *out.stride(),
185 + top.values.shape[1],
186 + topk_blocks,
187 + 256,
188 + )
189 +
190 +
191 +def apply_candidate_mask(
192 + logits: torch.Tensor,
193 + row_ks: torch.Tensor | None,
194 + row_ke: torch.Tensor,
195 + candidate_blocks: torch.Tensor,
196 + block_size: int,
197 + row_repeat: int = 1,
198 +) -> None:
199 + """Mask packed logits outside causal bounds and request-local candidates."""
200 + assert logits.is_cuda
201 + rows, width = logits.shape
202 + if not rows or not width:
203 + return
204 + nblocks = triton.cdiv(width, block_size)
205 + flags = torch.empty((rows, nblocks + 1), device=logits.device, dtype=torch.uint8)
206 + start_stride = row_ks.stride(0) if row_ks is not None else 0
207 + _candidate_flags_kernel[(rows,)](
208 + candidate_blocks,
209 + row_ks,
210 + flags,
211 + *candidate_blocks.stride(),
212 + start_stride,
213 + width,
214 + nblocks,
215 + block_size,
216 + candidate_blocks.shape[1],
217 + row_ks is not None,
218 + row_repeat,
219 + )
220 + _mask_candidates_kernel[(rows, triton.cdiv(width, 1024))](
221 + logits,
222 + row_ks,
223 + row_ke,
224 + flags,
225 + *logits.stride(),
226 + start_stride,
227 + row_ke.stride(0),
228 + width,
229 + nblocks,
230 + block_size,
231 + row_ks is not None,
232 + row_repeat,
233 + 1024,
234 + )
vllm/model_executor/kernels/linear/__init__.py
+15 -3
@@ -109,6 +109,9 @@ from vllm.model_executor.kernels.linear.mxfp8 import (
109 from vllm.model_executor.kernels.linear.mxfp8.b12x import (
110 B12xMxfp8LinearKernel,
111 )
112 +from vllm.model_executor.kernels.linear.mxfp8.deep_gemm import (
113 + DeepGemmMxfp8BmmLinearKernel,
114 +)
115 from vllm.model_executor.kernels.linear.mxfp8.emulation import (
116 EmulationMxfp8LinearKernel,
117 )
@@ -301,6 +304,7 @@ _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = {
304 },
305 "deep_gemm": {
306 DeepGemmFp8BlockScaledMMKernel,
307 + DeepGemmMxfp8BmmLinearKernel,
308 },
309 "torch": {
310 PerTensorTorchFP8ScaledMMLinearKernel,
@@ -851,13 +855,21 @@ def choose_mp_linear_kernel(
855 )
856
857
854 -def init_mxfp8_linear_kernel() -> Mxfp8LinearKernel:
858 +def init_mxfp8_linear_kernel(*, bmm_batch_size: int | None = None) -> Mxfp8LinearKernel:
859 """Select and instantiate the best MXFP8 linear kernel for the
860 current platform."""
857 - config = Mxfp8LinearLayerConfig()
861 + config = Mxfp8LinearLayerConfig(bmm_batch_size=bmm_batch_size)
862
863 platform = current_platform._enum
860 - possible = list(_POSSIBLE_MXFP8_KERNELS.get(platform, []))
864 + possible: list[type[Mxfp8LinearKernel]]
865 + if bmm_batch_size is not None:
866 + possible = (
867 + [DeepGemmMxfp8BmmLinearKernel, EmulationMxfp8LinearKernel]
868 + if current_platform.is_cuda()
869 + else []
870 + )
871 + else:
872 + possible = list(_POSSIBLE_MXFP8_KERNELS.get(platform, []))
873
874 # Apply --linear-backend filtering when set.
875 possible = _resolve_backend_kernels(possible, "MXFP8")
vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py
+1 -1
@@ -17,7 +17,7 @@ class Mxfp8LinearLayerConfig:
17 uint8 (E8M0) per-block scales at block size 32.
18 """
19
20 - pass
20 + bmm_batch_size: int | None = None
21
22
23 class Mxfp8LinearKernel(ABC):
vllm/model_executor/kernels/linear/mxfp8/deep_gemm.py new
+128
@@ -0,0 +1,128 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (
7 + deepgemm_post_process_fp8_weight_block,
8 + per_token_group_quant_fp8_packed_for_deepgemm,
9 +)
10 +from vllm.model_executor.utils import replace_parameter
11 +from vllm.platforms import current_platform
12 +from vllm.utils.deep_gemm import (
13 + fp8_einsum,
14 + is_deep_gemm_supported,
15 +)
16 +from vllm.utils.torch_utils import direct_register_custom_op
17 +
18 +from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
19 +
20 +
21 +class DeepGemmMxfp8BmmLinearKernel(Mxfp8LinearKernel):
22 + """Grouped MXFP8 einsum with per-row, 32-element weight scales."""
23 +
24 + def __init__(self, config: Mxfp8LinearLayerConfig):
25 + super().__init__(config)
26 + self.recipe = (1, 1, 32)
27 +
28 + @classmethod
29 + def is_supported(
30 + cls, compute_capability: int | None = None
31 + ) -> tuple[bool, str | None]:
32 + if not current_platform.is_cuda() or not is_deep_gemm_supported():
33 + return False, "DeepGEMM BMM requires a supported CUDA device."
34 + if not current_platform.is_device_capability_family(100):
35 + return False, "DeepGEMM MXFP8 BMM requires Blackwell."
36 + return True, None
37 +
38 + @classmethod
39 + def can_implement(cls, config: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
40 + if config.bmm_batch_size is None or config.bmm_batch_size <= 0:
41 + return False, "DeepGEMM BMM requires a positive batch size."
42 + return True, None
43 +
44 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
45 + weight = layer.weight.data
46 + if weight.ndim == 3:
47 + return
48 + scale = layer.weight_scale.data
49 + assert weight.dtype == torch.float8_e4m3fn and weight.ndim == 2
50 + assert scale.shape == (weight.shape[0], weight.shape[1] // 32)
51 + assert self.config.bmm_batch_size is not None
52 + weight, scale = deepgemm_post_process_fp8_weight_block(
53 + wq=weight,
54 + ws=scale,
55 + quant_block_shape=(1, 32),
56 + use_e8m0=False,
57 + is_bmm=True,
58 + bmm_batch_size=self.config.bmm_batch_size,
59 + )
60 + replace_parameter(layer, "weight", weight)
61 + replace_parameter(layer, "weight_scale", scale)
62 + layer.weight_block_size = [1, 32]
63 +
64 + def apply_weights(
65 + self,
66 + layer: torch.nn.Module,
67 + x: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
68 + bias: torch.Tensor | None = None,
69 + ) -> torch.Tensor:
70 + """Project [tokens, groups, K] input to [tokens, groups, N]."""
71 + if isinstance(x, tuple):
72 + q_input, input_scale = x
73 + else:
74 + tokens, groups, k = x.shape
75 + q_input, input_scale = per_token_group_quant_fp8_packed_for_deepgemm(
76 + x.reshape(tokens, groups * k),
77 + group_size=32,
78 + use_ue8m0=True,
79 + )
80 + q_input = q_input.view(tokens, groups, k)
81 + input_scale = input_scale.view(tokens, groups, k // 128)
82 + output = torch.ops.vllm.deepgemm_mxfp8_bmm(
83 + q_input, input_scale, layer.weight, layer.weight_scale, list(self.recipe)
84 + )
85 + if bias is not None:
86 + output = output + bias.view(output.shape[1:])
87 + return output
88 +
89 +
90 +def _deepgemm_mxfp8_bmm_fake(
91 + x: torch.Tensor,
92 + x_scale: torch.Tensor,
93 + weight: torch.Tensor,
94 + weight_scale: torch.Tensor,
95 + recipe: list[int],
96 +) -> torch.Tensor:
97 + return torch.empty(
98 + (x.shape[0], weight.shape[0], weight.shape[1]),
99 + device=x.device,
100 + dtype=torch.bfloat16,
101 + )
102 +
103 +
104 +def _deepgemm_mxfp8_bmm(
105 + x: torch.Tensor,
106 + x_scale: torch.Tensor,
107 + weight: torch.Tensor,
108 + weight_scale: torch.Tensor,
109 + recipe: list[int],
110 +) -> torch.Tensor:
111 + output = torch.empty(
112 + (x.shape[0], weight.shape[0], weight.shape[1]),
113 + device=x.device,
114 + dtype=torch.bfloat16,
115 + )
116 + fp8_einsum(
117 + "bhr,hdr->bhd",
118 + (x, x_scale),
119 + (weight, weight_scale),
120 + output,
121 + recipe=tuple(recipe),
122 + )
123 + return output
124 +
125 +
126 +direct_register_custom_op(
127 + "deepgemm_mxfp8_bmm", _deepgemm_mxfp8_bmm, fake_impl=_deepgemm_mxfp8_bmm_fake
128 +)
vllm/model_executor/kernels/linear/mxfp8/flashinfer.py
+38 -14
@@ -4,11 +4,19 @@
4 import torch
5 from torch.nn.parameter import Parameter
6
7 +from vllm.model_executor.layers.fusion.quant_activation import (
8 + QuantizedActivation,
9 + as_quantized_activation,
10 +)
11 from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
12 MXFP8_BLOCK_SIZE,
13 mxfp8_e4m3_quantize,
14 swizzle_mxfp8_scale,
15 )
16 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
17 + QuantKey,
18 + kMxfp8Dynamic,
19 +)
20 from vllm.platforms import current_platform
21 from vllm.utils import flashinfer as vllm_flashinfer
22 from vllm.utils.flashinfer import has_flashinfer, has_flashinfer_cutedsl
@@ -35,6 +43,10 @@ class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel):
43 def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
44 return True, None
45
46 + def input_quant_key(self) -> QuantKey:
47 + # Activations use FlashInfer's F8_128x4 swizzled scale layout.
48 + return kMxfp8Dynamic
49 +
50 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
51 weight = layer.weight.data # [N, K]
52 N, K = weight.shape
@@ -51,16 +63,13 @@ class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel):
63 def apply_weights(
64 self,
65 layer: torch.nn.Module,
54 - x: torch.Tensor,
66 + x: torch.Tensor | QuantizedActivation,
67 bias: torch.Tensor | None = None,
68 ) -> torch.Tensor:
69 weight = layer.weight
70 weight_scale = layer.weight_scale
59 - out_dtype = x.dtype
71 N, K = weight.shape
72
62 - input_shape = x.shape
63 - input_2d = x.view(-1, K)
73 min_dim = 128
74
75 assert min_dim <= K, (
@@ -75,9 +84,16 @@ class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel):
84 f"out_features is too small for mm_mxfp8."
85 )
86
78 - input_mxfp8, input_scale = mxfp8_e4m3_quantize(
79 - input_2d, is_sf_swizzled_layout=True
80 - )
87 + qa = as_quantized_activation(x, self.input_quant_key())
88 + if qa is not None:
89 + input_mxfp8, input_scale = qa.data, qa.scale
90 + out_dtype, input_shape = qa.orig_dtype, qa.orig_shape
91 + else:
92 + assert isinstance(x, torch.Tensor)
93 + out_dtype, input_shape = x.dtype, x.shape
94 + input_mxfp8, input_scale = mxfp8_e4m3_quantize(
95 + x.view(-1, K), is_sf_swizzled_layout=True
96 + )
97
98 if not weight.is_contiguous():
99 weight = weight.contiguous()
@@ -118,6 +134,10 @@ class FlashInferCutedslMxfp8LinearKernel(Mxfp8LinearKernel):
134 def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
135 return True, None
136
137 + def input_quant_key(self) -> QuantKey:
138 + # Activations use FlashInfer's F8_128x4 swizzled scale layout.
139 + return kMxfp8Dynamic
140 +
141 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
142 weight = layer.weight.data # [N, K]
143 N, K = weight.shape
@@ -135,16 +155,13 @@ class FlashInferCutedslMxfp8LinearKernel(Mxfp8LinearKernel):
155 def apply_weights(
156 self,
157 layer: torch.nn.Module,
138 - x: torch.Tensor,
158 + x: torch.Tensor | QuantizedActivation,
159 bias: torch.Tensor | None = None,
160 ) -> torch.Tensor:
161 weight = layer.weight # [K, N], column-major
162 weight_scale = layer.weight_scale
143 - out_dtype = x.dtype
163 K, N = weight.shape
164
146 - input_shape = x.shape
147 - input_2d = x.view(-1, K)
165 min_dim = 128
166
167 assert min_dim <= K, (
@@ -159,9 +176,16 @@ class FlashInferCutedslMxfp8LinearKernel(Mxfp8LinearKernel):
176 f"out_features is too small for mm_mxfp8."
177 )
178
162 - input_mxfp8, input_scale = mxfp8_e4m3_quantize(
163 - input_2d, is_sf_swizzled_layout=True
164 - )
179 + qa = as_quantized_activation(x, self.input_quant_key())
180 + if qa is not None:
181 + input_mxfp8, input_scale = qa.data, qa.scale
182 + out_dtype, input_shape = qa.orig_dtype, qa.orig_shape
183 + else:
184 + assert isinstance(x, torch.Tensor)
185 + out_dtype, input_shape = x.dtype, x.shape
186 + input_mxfp8, input_scale = mxfp8_e4m3_quantize(
187 + x.view(-1, K), is_sf_swizzled_layout=True
188 + )
189
190 output = vllm_flashinfer.mm_mxfp8(
191 input_mxfp8,
vllm/model_executor/kernels/mhc/tilelang.py
+200
@@ -87,6 +87,190 @@ def _tilelang_hc_prenorm_gemm(
87 )
88
89
90 +def mhc_pre_delayed_tilelang(
91 + residual: torch.Tensor,
92 + fn: torch.Tensor,
93 + hc_scale: torch.Tensor,
94 + hc_base: torch.Tensor,
95 + rms_eps: float,
96 + hc_pre_eps: float,
97 + hc_sinkhorn_eps: float,
98 + hc_post_mult_value: float,
99 + sinkhorn_repeat: int,
100 + pre_mix: torch.Tensor | None = None,
101 + x: torch.Tensor | None = None,
102 + norm_weight: torch.Tensor | None = None,
103 + norm_eps: float = 1e-6,
104 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
105 + """Run mHC pre with a carried pre-mix and return the next pre-mix.
106 +
107 + Args:
108 + residual: BF16 residual streams of shape (tokens, hc_mult, hidden_size).
109 + fn: FP32 projection of shape (hc_mult * (hc_mult + 2), input_size).
110 + hc_scale: FP32 scales of shape (3,).
111 + hc_base: FP32 bias of shape (hc_mult * (hc_mult + 2),).
112 + rms_eps: RMS normalization epsilon.
113 + hc_pre_eps: Pre-mix epsilon.
114 + hc_sinkhorn_eps: Sinkhorn epsilon.
115 + hc_post_mult_value: Post-mix multiplier.
116 + sinkhorn_repeat: Number of Sinkhorn iterations.
117 + pre_mix: FP32 coefficients from the previous sublayer, or None to
118 + select residual stream zero at model entry.
119 + x: Optional BF16 projection input of shape (tokens, input_size), for
120 + the first layer's broadcast embedding and summed projection.
121 + norm_weight: Optional BF16 RMSNorm weight for the collapsed input.
122 + norm_eps: RMSNorm epsilon for the collapsed input.
123 +
124 + Returns:
125 + Post and residual coefficients, optionally normalized BF16 layer input,
126 + and the next FP32 pre-mix, with shapes (tokens, hc_mult, 1),
127 + (tokens, hc_mult, hc_mult), (tokens, hidden_size), and (tokens, hc_mult).
128 + """
129 + from vllm.model_executor.kernels.mhc.tilelang_kernels import (
130 + mhc_pre_big_fuse_tilelang,
131 + )
132 + from vllm.model_executor.kernels.mhc.warmup import (
133 + MHC_PRE_NORM_KERNEL,
134 + compute_mhc_pre_num_splits,
135 + )
136 + from vllm.utils.deep_gemm import (
137 + is_deep_gemm_supported,
138 + tf32_hc_prenorm_gemm,
139 + )
140 +
141 + assert residual.ndim == 3 and residual.dtype == torch.bfloat16
142 + assert residual.is_contiguous()
143 + num_tokens, hc_mult, hidden_size = residual.shape
144 + if x is None:
145 + x = residual.view(num_tokens, hc_mult * hidden_size)
146 + assert x.ndim == 2 and x.dtype == torch.bfloat16 and x.is_contiguous()
147 + assert x.shape[0] == num_tokens
148 + input_size = x.shape[1]
149 + mix_size = hc_mult * (hc_mult + 2)
150 + assert fn.shape == (mix_size, input_size) and fn.dtype == torch.float32
151 + assert hc_scale.shape == (3,) and hc_scale.dtype == torch.float32
152 + assert hc_base.shape == (mix_size,) and hc_base.dtype == torch.float32
153 + if pre_mix is not None:
154 + assert pre_mix.shape == (num_tokens, hc_mult)
155 + assert pre_mix.dtype == torch.float32 and pre_mix.is_contiguous()
156 +
157 + next_pre_mix = torch.empty(
158 + num_tokens, hc_mult, dtype=torch.float32, device=residual.device
159 + )
160 + post = torch.empty_like(next_pre_mix)
161 + comb = torch.empty(
162 + num_tokens, hc_mult * hc_mult, dtype=torch.float32, device=residual.device
163 + )
164 + layer_input = torch.empty(
165 + num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device
166 + )
167 + outputs = (
168 + post.unsqueeze(-1),
169 + comb.view(num_tokens, hc_mult, hc_mult),
170 + layer_input,
171 + next_pre_mix,
172 + )
173 + if num_tokens == 0:
174 + return outputs
175 +
176 + use_deep_gemm = is_deep_gemm_supported()
177 + n_splits = (
178 + compute_mhc_pre_num_splits(input_size, num_tokens) if use_deep_gemm else 1
179 + )
180 + mixes = torch.empty(
181 + n_splits, num_tokens, mix_size, dtype=torch.float32, device=residual.device
182 + )
183 + sqrsum = torch.empty(
184 + n_splits, num_tokens, dtype=torch.float32, device=residual.device
185 + )
186 + if use_deep_gemm:
187 + tf32_hc_prenorm_gemm(x, fn, mixes, sqrsum, n_splits)
188 + else:
189 + _tilelang_hc_prenorm_gemm(x, fn, mixes, sqrsum, input_size, 1)
190 + if norm_weight is not None:
191 + assert norm_weight.shape == (hidden_size,)
192 + assert norm_weight.dtype == torch.bfloat16 and norm_weight.is_contiguous()
193 + MHC_PRE_NORM_KERNEL(
194 + mixes,
195 + sqrsum,
196 + hc_scale,
197 + hc_base,
198 + residual,
199 + post,
200 + comb,
201 + layer_input,
202 + norm_weight,
203 + pre_mix if pre_mix is not None else post,
204 + next_pre_mix,
205 + hidden_size=hidden_size,
206 + rms_eps=rms_eps,
207 + hc_pre_eps=hc_pre_eps,
208 + hc_sinkhorn_eps=hc_sinkhorn_eps,
209 + hc_post_mult_value=hc_post_mult_value,
210 + sinkhorn_repeat=sinkhorn_repeat,
211 + norm_eps=norm_eps,
212 + hc_mult=hc_mult,
213 + use_pre_mix_in=pre_mix is not None,
214 + save_pre_mix=True,
215 + rms_numel=input_size,
216 + )
217 + return outputs
218 + mhc_pre_big_fuse_tilelang(
219 + mixes,
220 + sqrsum,
221 + hc_scale,
222 + hc_base,
223 + residual,
224 + post,
225 + comb,
226 + layer_input,
227 + pre_mix if pre_mix is not None else post,
228 + next_pre_mix,
229 + hidden_size,
230 + rms_eps,
231 + hc_pre_eps,
232 + hc_sinkhorn_eps,
233 + hc_post_mult_value,
234 + sinkhorn_repeat,
235 + n_splits,
236 + hc_mult,
237 + use_pre_mix_in=pre_mix is not None,
238 + save_pre_mix=True,
239 + rms_numel=input_size,
240 + )
241 + return outputs
242 +
243 +
244 +def _mhc_pre_delayed_tilelang_fake(
245 + residual: torch.Tensor,
246 + fn: torch.Tensor,
247 + hc_scale: torch.Tensor,
248 + hc_base: torch.Tensor,
249 + rms_eps: float,
250 + hc_pre_eps: float,
251 + hc_sinkhorn_eps: float,
252 + hc_post_mult_value: float,
253 + sinkhorn_repeat: int,
254 + pre_mix: torch.Tensor | None = None,
255 + x: torch.Tensor | None = None,
256 + norm_weight: torch.Tensor | None = None,
257 + norm_eps: float = 1e-6,
258 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
259 + num_tokens, hc_mult, hidden_size = residual.shape
260 + return (
261 + torch.empty(
262 + num_tokens, hc_mult, 1, dtype=torch.float32, device=residual.device
263 + ),
264 + torch.empty(
265 + num_tokens, hc_mult, hc_mult, dtype=torch.float32, device=residual.device
266 + ),
267 + torch.empty(
268 + num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device
269 + ),
270 + torch.empty(num_tokens, hc_mult, dtype=torch.float32, device=residual.device),
271 + )
272 +
273 +
274 def mhc_pre_tilelang(
275 residual: torch.Tensor,
276 fn: torch.Tensor,
@@ -219,6 +403,9 @@ def mhc_pre_tilelang(
403 post_mix,
404 comb_mix,
405 layer_input,
406 + # Pre-mix buffers are unused in V4 mode.
407 + post_mix,
408 + post_mix,
409 hidden_size,
410 rms_eps,
411 hc_pre_eps,
@@ -239,6 +426,8 @@ def mhc_pre_tilelang(
426 comb_mix,
427 layer_input,
428 norm_weight,
429 + post_mix,
430 + post_mix,
431 hidden_size,
432 rms_eps,
433 hc_pre_eps,
@@ -622,6 +811,9 @@ def mhc_fused_post_pre_tilelang(
811 post_mix_cur,
812 comb_mix_cur,
813 layer_input_cur,
814 + # Pre-mix buffers are unused in V4 mode.
815 + post_mix_cur,
816 + post_mix_cur,
817 hidden_size,
818 rms_eps,
819 hc_pre_eps,
@@ -642,6 +834,8 @@ def mhc_fused_post_pre_tilelang(
834 comb_mix_cur,
835 layer_input_cur,
836 norm_weight,
837 + post_mix_cur,
838 + post_mix_cur,
839 hidden_size,
840 rms_eps,
841 hc_pre_eps,
@@ -762,6 +956,12 @@ def _hc_head_fused_kernel_tilelang_fake(
956 )
957
958
959 +direct_register_custom_op(
960 + op_name="mhc_pre_delayed_tilelang",
961 + op_func=mhc_pre_delayed_tilelang,
962 + mutates_args=[],
963 + fake_impl=_mhc_pre_delayed_tilelang_fake,
964 +)
965 direct_register_custom_op(
966 op_name="mhc_pre_tilelang",
967 op_func=mhc_pre_tilelang,
vllm/model_executor/kernels/mhc/tilelang_kernels.py
+67 -18
@@ -38,6 +38,8 @@ def mhc_pre_big_fuse_tilelang(
38 post_mix,
39 comb_mix,
40 layer_input,
41 + pre_mix_in,
42 + pre_mix_out,
43 hidden_size: int,
44 rms_eps: float,
45 hc_pre_eps: float,
@@ -46,11 +48,20 @@ def mhc_pre_big_fuse_tilelang(
48 sinkhorn_repeat: int,
49 n_splits: int = 16,
50 hc_mult: int = 4,
51 + use_pre_mix_in: bool = False,
52 + save_pre_mix: bool = False,
53 + rms_numel: int = 0,
54 ):
50 - """Deeply fused kernels, everything other than gemm & sqrsum in mHC pre block."""
55 + """Fuse coefficient generation and residual collapse after the projection.
56 +
57 + With save_pre_mix, store the new pre-mix and collapse with pre_mix_in,
58 + or select stream zero when use_pre_mix_in is false.
59 + """
60 num_tokens = T.dynamic("num_tokens")
61 hc_mult3 = hc_mult * (2 + hc_mult)
62 hidden_block = math.gcd(512, hidden_size)
63 + if rms_numel == 0:
64 + rms_numel = hc_mult * hidden_size
65
66 gemm_out_mul: T.Tensor[[n_splits, num_tokens, hc_mult3], T.float32] # type: ignore[no-redef, valid-type]
67 gemm_out_sqrsum: T.Tensor[[n_splits, num_tokens], T.float32] # type: ignore[no-redef, valid-type]
@@ -62,6 +73,9 @@ def mhc_pre_big_fuse_tilelang(
73 comb_mix: T.Tensor[[num_tokens, hc_mult * hc_mult], T.float32] # type: ignore[no-redef, valid-type]
74 layer_input: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type]
75
76 + pre_mix_in: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type]
77 + pre_mix_out: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type]
78 +
79 with T.Kernel(num_tokens, threads=96) as i:
80 if ENABLE_PDL:
81 T.pdl_sync()
@@ -73,7 +87,7 @@ def mhc_pre_big_fuse_tilelang(
87 rms[0] = 0
88 for i_split in T.serial(n_splits):
89 rms[0] += gemm_out_sqrsum[i_split, i]
76 - rms[0] = T.rsqrt(rms[0] / (hc_mult * hidden_size) + rms_eps)
90 + rms[0] = T.rsqrt(rms[0] / rms_numel + rms_eps)
91 for j in T.Parallel(hc_mult3):
92 mixes[j] = 0
93 for i_split in T.serial(n_splits):
@@ -87,6 +101,11 @@ def mhc_pre_big_fuse_tilelang(
101 # _pre_split_mixes_fwd (post & comb)
102 cm = T.alloc_fragment((hc_mult, hc_mult), T.float32)
103 for j in T.Parallel(hc_mult):
104 + if save_pre_mix:
105 + pre_mix_out[i, j] = (
106 + T.sigmoid(mixes_shared[j] * hc_scale[0] + hc_base[j])
107 + + hc_pre_eps
108 + )
109 post_mix[i, j] = (
110 T.sigmoid(
111 mixes_shared[j + hc_mult] * hc_scale[1] + hc_base[j + hc_mult]
@@ -137,12 +156,15 @@ def mhc_pre_big_fuse_tilelang(
156 # _pre_split_mixes_fwd (pre)
157 pre_mix_shared = T.alloc_shared(hc_mult, T.float32)
158 for j in T.Parallel(hc_mult):
140 - pre_mix_shared[j] = (
141 - T.sigmoid(
142 - mixes_shared[j] * hc_scale[0] + hc_base[j],
159 + if use_pre_mix_in:
160 + pre_mix_shared[j] = pre_mix_in[i, j]
161 + elif save_pre_mix:
162 + pre_mix_shared[j] = T.if_then_else(j == 0, 1.0, 0.0)
163 + else:
164 + pre_mix_shared[j] = (
165 + T.sigmoid(mixes_shared[j] * hc_scale[0] + hc_base[j])
166 + + hc_pre_eps
167 )
144 - + hc_pre_eps
145 - )
168 ###################################################################
169 # _pre_apply_mix_fwd
170 for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=2):
@@ -179,6 +201,8 @@ def mhc_pre_big_fuse_with_norm_tilelang(
201 comb_mix,
202 layer_input,
203 norm_weight,
204 + pre_mix_in,
205 + pre_mix_out,
206 hidden_size: int,
207 rms_eps: float,
208 hc_pre_eps: float,
@@ -189,11 +213,16 @@ def mhc_pre_big_fuse_with_norm_tilelang(
213 n_splits: int = 16,
214 hc_mult: int = 4,
215 gemm_last_dim: int = -1,
216 + use_pre_mix_in: bool = False,
217 + save_pre_mix: bool = False,
218 + rms_numel: int = 0,
219 ):
220 num_tokens = T.dynamic("num_tokens")
221 hc_mult3 = hc_mult * (2 + hc_mult)
222 if gemm_last_dim < 0:
223 gemm_last_dim = hc_mult3
224 + if rms_numel == 0:
225 + rms_numel = hc_mult * hidden_size
226 hidden_block = math.gcd(1024, hidden_size)
227
228 gemm_out_mul: T.Tensor[[n_splits, num_tokens, gemm_last_dim], T.float32] # type: ignore[no-redef, valid-type]
@@ -206,6 +235,9 @@ def mhc_pre_big_fuse_with_norm_tilelang(
235 layer_input: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type]
236 norm_weight: T.Tensor[[hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type]
237
238 + pre_mix_in: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type]
239 + pre_mix_out: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type]
240 +
241 with T.Kernel(num_tokens, threads=96) as i:
242 rms = T.alloc_fragment(1, T.float32)
243 mixes = T.alloc_fragment(hc_mult3, T.float32)
@@ -217,7 +249,7 @@ def mhc_pre_big_fuse_with_norm_tilelang(
249
250 for i_split in T.serial(n_splits):
251 rms[0] += gemm_out_sqrsum[i_split, i]
220 - rms[0] = T.rsqrt(rms[0] / (hc_mult * hidden_size) + rms_eps)
252 + rms[0] = T.rsqrt(rms[0] / rms_numel + rms_eps)
253 for j in T.Parallel(hc_mult3):
254 mixes[j] = 0
255 for i_split in T.serial(n_splits):
@@ -229,6 +261,11 @@ def mhc_pre_big_fuse_with_norm_tilelang(
261 if T.get_thread_binding() < 32:
262 cm = T.alloc_fragment((hc_mult, hc_mult), T.float32)
263 for j in T.Parallel(hc_mult):
264 + if save_pre_mix:
265 + pre_mix_out[i, j] = (
266 + T.sigmoid(mixes_shared[j] * hc_scale[0] + hc_base[j])
267 + + hc_pre_eps
268 + )
269 post_mix[i, j] = (
270 T.sigmoid(
271 mixes_shared[j + hc_mult] * hc_scale[1] + hc_base[j + hc_mult]
@@ -270,12 +307,15 @@ def mhc_pre_big_fuse_with_norm_tilelang(
307 else:
308 pre_mix_shared = T.alloc_shared(hc_mult, T.float32)
309 for j in T.Parallel(hc_mult):
273 - pre_mix_shared[j] = (
274 - T.sigmoid(
275 - mixes_shared[j] * hc_scale[0] + hc_base[j],
310 + if use_pre_mix_in:
311 + pre_mix_shared[j] = pre_mix_in[i, j]
312 + elif save_pre_mix:
313 + pre_mix_shared[j] = T.if_then_else(j == 0, 1.0, 0.0)
314 + else:
315 + pre_mix_shared[j] = (
316 + T.sigmoid(mixes_shared[j] * hc_scale[0] + hc_base[j])
317 + + hc_pre_eps
318 )
277 - + hc_pre_eps
278 - )
319
320 # Pass 1: stash unnormalized weighted-sum output in shared memory
321 # as bf16 (matches the rounding that RMSNorm would see) while
@@ -298,9 +338,18 @@ def mhc_pre_big_fuse_with_norm_tilelang(
338 for i1_h in T.Parallel(hidden_block):
339 ol[i1_h] += pre * xl[i_hc, i1_h]
340
301 - for i1_h in T.Parallel(hidden_block):
302 - sumsq_per_pos[i1_h] += ol[i1_h] * ol[i1_h]
303 - output_shared[i0_h * hidden_block + i1_h] = T.bfloat16(ol[i1_h])
341 + if save_pre_mix:
342 + # Keep the BF16 boundary before the delayed input RMSNorm.
343 + rounded = T.alloc_fragment(hidden_block, T.bfloat16)
344 + T.copy(ol, rounded)
345 + for i1_h in T.Parallel(hidden_block):
346 + value = T.float32(rounded[i1_h])
347 + sumsq_per_pos[i1_h] += value * value
348 + output_shared[i0_h * hidden_block + i1_h] = rounded[i1_h]
349 + else:
350 + for i1_h in T.Parallel(hidden_block):
351 + sumsq_per_pos[i1_h] += ol[i1_h] * ol[i1_h]
352 + output_shared[i0_h * hidden_block + i1_h] = T.bfloat16(ol[i1_h])
353
354 sumsq = T.alloc_fragment(1, T.float32)
355 T.reduce_sum(sumsq_per_pos, sumsq, dim=0)
@@ -704,7 +753,7 @@ def hc_prenorm_gemm_tilelang(
753
754 for it in T.serial(k_iters):
755 i_k = i_s * k_per_split + it * n_thr + tid
707 - x_val = x[i_n, i_k]
756 + x_val = T.cast(x[i_n, i_k], T.float32)
757 for i_o in T.unroll(tile_n):
758 out_idx = i_t * tile_n + i_o
759 if out_idx < n_out:
@@ -794,7 +843,7 @@ def hc_prenorm_gemm_block_m_tilelang(
843 for i_m in T.unroll(block_m):
844 token_idx = i_mt * block_m + i_m
845 if token_idx < num_tokens:
797 - x_val = x[token_idx, i_k]
846 + x_val = T.cast(x[token_idx, i_k], T.float32)
847 for i_o in T.unroll(tile_n):
848 acc[i_m, i_o] += x_val * fn_val[i_o]
849 if i_t == 0:
vllm/model_executor/kernels/mhc/torch.py
+42
@@ -3,6 +3,48 @@
3 import torch
4
5
6 +def mhc_pre_delayed_torch(
7 + residual: torch.Tensor,
8 + fn: torch.Tensor,
9 + hc_scale: torch.Tensor,
10 + hc_base: torch.Tensor,
11 + rms_eps: float,
12 + hc_pre_eps: float,
13 + hc_sinkhorn_eps: float,
14 + hc_post_mult_value: float,
15 + sinkhorn_repeat: int,
16 + pre_mix: torch.Tensor | None = None,
17 + x: torch.Tensor | None = None,
18 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
19 + """Reference for mHC pre using coefficients from the previous sublayer."""
20 + hc_mult = residual.shape[1]
21 + x = (residual.flatten(1) if x is None else x).float()
22 + mixes = (x @ fn.t()) * torch.rsqrt(x.square().mean(-1, keepdim=True) + rms_eps)
23 + pre = (
24 + torch.sigmoid(mixes[:, :hc_mult] * hc_scale[0] + hc_base[:hc_mult]) + hc_pre_eps
25 + )
26 + post = (
27 + torch.sigmoid(
28 + mixes[:, hc_mult : 2 * hc_mult] * hc_scale[1]
29 + + hc_base[hc_mult : 2 * hc_mult]
30 + )
31 + * hc_post_mult_value
32 + )
33 + comb = mixes[:, 2 * hc_mult :].view(-1, hc_mult, hc_mult) * hc_scale[2]
34 + comb = comb + hc_base[2 * hc_mult :].view(1, hc_mult, hc_mult)
35 + comb = torch.softmax(comb, dim=-1) + hc_sinkhorn_eps
36 + comb = comb / (comb.sum(dim=-2, keepdim=True) + hc_sinkhorn_eps)
37 + for _ in range(sinkhorn_repeat - 1):
38 + comb = comb / (comb.sum(dim=-1, keepdim=True) + hc_sinkhorn_eps)
39 + comb = comb / (comb.sum(dim=-2, keepdim=True) + hc_sinkhorn_eps)
40 + layer_input = (
41 + residual[:, 0]
42 + if pre_mix is None
43 + else (pre_mix.unsqueeze(-1) * residual.float()).sum(dim=1).to(residual.dtype)
44 + )
45 + return post.unsqueeze(-1), comb, layer_input, pre
46 +
47 +
48 def mhc_pre_torch(
49 residual: torch.Tensor,
50 fn: torch.Tensor,
vllm/model_executor/kernels/mhc/triton.py
+45 -1
@@ -78,7 +78,7 @@ def _hc_head_reduce_store_kernel(
78 out_stride_h: tl.constexpr,
79 BLOCK_H: tl.constexpr,
80 ):
81 - token_idx = tl.program_id(0)
81 + token_idx = tl.program_id(0).to(tl.int64)
82 block_idx = tl.program_id(1)
83 offsets = block_idx * BLOCK_H + tl.arange(0, BLOCK_H)
84 mask = offsets < hidden_size
@@ -105,6 +105,42 @@ def _hc_head_reduce_store_kernel(
105 )
106
107
108 +def hc_collapse_triton(x: Tensor, pre_mix: Tensor) -> Tensor:
109 + """Collapse BF16 residual streams with FP32 pre-mix coefficients."""
110 + assert x.ndim == 3 and x.dtype == torch.bfloat16
111 + num_tokens, hc_mult, hidden_size = x.shape
112 + assert pre_mix.shape == (num_tokens, hc_mult)
113 + assert pre_mix.dtype == torch.float32
114 + out = torch.empty(num_tokens, hidden_size, dtype=x.dtype, device=x.device)
115 + if num_tokens == 0:
116 + return out
117 +
118 + block_h = 1024
119 + _hc_head_reduce_store_kernel[(num_tokens, triton.cdiv(hidden_size, block_h))](
120 + pre_mix,
121 + x,
122 + out,
123 + hidden_size,
124 + hc_mult,
125 + pre_mix.stride(0),
126 + pre_mix.stride(1),
127 + x.stride(0),
128 + x.stride(1),
129 + x.stride(2),
130 + out.stride(0),
131 + out.stride(1),
132 + BLOCK_H=block_h,
133 + num_warps=4,
134 + # Preserve the separate FP32 multiply and sum in the Torch reference.
135 + enable_fp_fusion=False,
136 + )
137 + return out
138 +
139 +
140 +def _hc_collapse_triton_fake(x: Tensor, pre_mix: Tensor) -> Tensor:
141 + return torch.empty(x.shape[0], x.shape[2], dtype=x.dtype, device=x.device)
142 +
143 +
144 def hc_head_reduce_triton_kernel(
145 x: torch.Tensor,
146 hc_fn: torch.Tensor,
@@ -172,3 +208,11 @@ direct_register_custom_op(
208 op_func=_hc_head_triton,
209 mutates_args=["out"],
210 )
211 +
212 +
213 +direct_register_custom_op(
214 + op_name="hc_collapse_triton",
215 + op_func=hc_collapse_triton,
216 + mutates_args=[],
217 + fake_impl=_hc_collapse_triton_fake,
218 +)
vllm/model_executor/kernels/mhc/warmup.py new
+74
@@ -0,0 +1,74 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from dataclasses import asdict, dataclass
5 +from typing import Any
6 +
7 +from vllm.model_executor.kernels.mhc.tilelang_kernels import (
8 + compute_num_split,
9 + mhc_pre_big_fuse_with_norm_tilelang,
10 +)
11 +from vllm.model_executor.warmup.jit_warmup import VllmJitKernel
12 +from vllm.platforms import current_platform
13 +from vllm.utils.deep_gemm import is_deep_gemm_supported
14 +from vllm.utils.math_utils import cdiv
15 +
16 +
17 +def compute_mhc_pre_num_splits(input_size: int, num_tokens: int) -> int:
18 + splits = compute_num_split(64, input_size, cdiv(num_tokens, 64))
19 + # Bound both GEMM and fused-normalization specializations during startup.
20 + return 1 if splits == 1 else 4 if splits <= 4 else 16
21 +
22 +
23 +class MHCPreNormKernel(VllmJitKernel["MHCPreNormKernel.CompileKey"]):
24 + @dataclass(frozen=True)
25 + class CompileKey:
26 + hidden_size: int
27 + rms_eps: float
28 + hc_pre_eps: float
29 + hc_sinkhorn_eps: float
30 + hc_post_mult_value: float
31 + sinkhorn_repeat: int
32 + norm_eps: float
33 + n_splits: int
34 + hc_mult: int
35 + use_pre_mix_in: bool
36 + rms_numel: int
37 + save_pre_mix: bool = True
38 +
39 + kernel: Any = staticmethod(mhc_pre_big_fuse_with_norm_tilelang)
40 +
41 + def dispatch(self, *, n_splits, **fields) -> CompileKey: # type: ignore[override]
42 + return self.CompileKey(n_splits=n_splits, **fields)
43 +
44 + def get_warmup_keys(self, *, max_tokens: int, **fields) -> list[CompileKey]:
45 + # The split heuristic changes only at 64-token boundaries.
46 + splits = (
47 + sorted(
48 + {
49 + compute_mhc_pre_num_splits(fields["rms_numel"], num_tokens)
50 + for num_tokens in range(1, max_tokens + 1, 64)
51 + }
52 + )
53 + if is_deep_gemm_supported()
54 + else [1]
55 + )
56 + return self._trace_dispatch(self.dispatch)(n_splits=splits, **fields)
57 +
58 + def compile(self, compile_key: CompileKey) -> None:
59 + if compile_key not in self._compiled_cache:
60 + self._compiled_cache[compile_key] = self.kernel.compile(
61 + **asdict(compile_key)
62 + )
63 +
64 + def __call__(self, *tensors, **fields):
65 + compile_key = self.dispatch(
66 + n_splits=tensors[0].shape[0],
67 + **fields,
68 + )
69 + if not current_platform.is_cuda():
70 + return self.kernel(*tensors, **asdict(compile_key))
71 + return self._get_or_compile(compile_key)(*tensors)
72 +
73 +
74 +MHC_PRE_NORM_KERNEL = MHCPreNormKernel()
vllm/model_executor/layers/fused_moe/router/fused_moe_router.py
+4
@@ -18,6 +18,10 @@ class FusedMoERouter(ABC):
18 def __init__(self, eplb_state: EplbLayerState | None = None):
19 self._routing_replay_out: torch.Tensor | None = None
20 self.eplb_state = eplb_state
21 + # Deepseek V4.1 vision checkpoints carry a second per-expert routing
22 + # bias for image sentinel tokens; attached by model code when present.
23 + self.bias_vl: torch.Tensor | None = None
24 + self.image_sentinel_lo: int = 0
25
26 @abstractmethod
27 def set_capture_fn(
vllm/model_executor/layers/quantization/modelopt.py
+28 -2
@@ -1821,6 +1821,7 @@ class CkptCtx:
1821 """Per-checkpoint facts a QuantKey cannot carry."""
1822
1823 group_size: int | None = None
1824 + scale_block_size: tuple[int, int] | None = None
1825
1826
1827 @dataclass(frozen=True)
@@ -2153,6 +2154,23 @@ class KMxfp8Static(QuantKeyScheme):
2154
2155 key = kMxfp8Static
2156
2157 + @staticmethod
2158 + def get_scale_weight_loader(weight_loader: Callable, ctx: CkptCtx) -> Callable:
2159 + block_rows, block_cols = ctx.scale_block_size or (1, MXFP8_BLOCK_SIZE)
2160 + if block_rows < 1 or block_cols != MXFP8_BLOCK_SIZE:
2161 + raise NotImplementedError(
2162 + f"MXFP8 checkpoint scale block {ctx.scale_block_size} is unsupported"
2163 + )
2164 +
2165 + def scaled_loader(param, loaded_weight, *args, **kwargs):
2166 + assert loaded_weight.dtype in (torch.uint8, torch.float8_e8m0fnu)
2167 + loaded_weight = loaded_weight.view(torch.uint8).repeat_interleave(
2168 + block_rows, dim=0
2169 + )
2170 + return weight_loader(param, loaded_weight, *args, **kwargs)
2171 +
2172 + return scaled_loader if block_rows > 1 else weight_loader
2173 +
2174 def create_weights(self, layer, role, ctx, shapes, wl) -> None:
2175 if role is not WEIGHT:
2176 self.reject(role)
@@ -2171,6 +2189,7 @@ class KMxfp8Static(QuantKeyScheme):
2189 input_dim=1,
2190 output_dim=0,
2191 )
2192 + scale_loader = self.get_scale_weight_loader(wl, ctx)
2193 self.register_params(
2194 layer,
2195 "weight_scale",
@@ -2180,10 +2199,11 @@ class KMxfp8Static(QuantKeyScheme):
2199 ),
2200 MXFP8_SCALE_DTYPE,
2201 ModelWeightParameter,
2183 - wl,
2202 + scale_loader,
2203 input_dim=1,
2204 output_dim=0,
2205 )
2206 + layer.weight_block_size = [1, MXFP8_BLOCK_SIZE]
2207
2208 def process(self, layer, role) -> None:
2209 if role is not WEIGHT:
@@ -2194,6 +2214,8 @@ class KMxfp8Static(QuantKeyScheme):
2214 # the fp8 asserts when the weight is already >=2-byte.
2215 if layer.weight.element_size() >= 2:
2216 return
2217 + if getattr(layer, "is_bmm", False) and layer.weight.ndim == 3:
2218 + return
2219 assert layer.weight.ndim == 2 and layer.weight.dtype == MXFP8_VALUE_DTYPE
2220 assert layer.weight_scale.ndim == 2
2221 assert layer.weight_scale.dtype == MXFP8_SCALE_DTYPE
@@ -2255,7 +2277,9 @@ def select_linear_kernel(
2277 # after #50273); W4A4 → use_a16=False.
2278 return init_nvfp4_linear_kernel(use_a16=spec.activation is None)
2279 if w.scale.dtype == MXFP8_SCALE_DTYPE:
2258 - return init_mxfp8_linear_kernel()
2280 + return init_mxfp8_linear_kernel(
2281 + bmm_batch_size=getattr(layer, "bmm_batch_size", None)
2282 + )
2283 # fp8 family: init_fp8 routes block-vs-plain itself off the activation key,
2284 # and needs a real key -- weight-only fp8 is not a ModelOpt format.
2285 act = spec.activation
@@ -2482,6 +2506,8 @@ class ModelOptLinearMethod(LinearMethodBase):
2506 persistent=False,
2507 )
2508 layer._nvfp4_group_size_for_gather = self.ctx.group_size
2509 + if self.spec.weight == kMxfp8Static and getattr(layer, "is_bmm", False):
2510 + self.kernel = init_mxfp8_linear_kernel(bmm_batch_size=layer.bmm_batch_size)
2511 self.kernel.process_weights_after_loading(layer)
2512
2513 def apply(self, layer, x, bias=None):
vllm/model_executor/layers/sparse_attn_indexer.py
+97 -4
@@ -13,6 +13,12 @@ from vllm.distributed import get_dcp_group, get_pcp_group
13 from vllm.forward_context import get_forward_context
14 from vllm.logger import init_logger
15 from vllm.model_executor.custom_op import CustomOp
16 +from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
17 + apply_candidate_mask as _apply_candidate_mask,
18 +)
19 +from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
20 + select_candidate_blocks as _select_candidate_blocks,
21 +)
22 from vllm.model_executor.layers.quantization.utils.quant_utils import (
23 get_fp8_min_max,
24 )
@@ -316,6 +322,9 @@ def sparse_attn_indexer(
322 dcp_world_size: int = 1,
323 cp_kv_cache_interleave_size: int = 1,
324 skip_topk_buffer_clear: bool = False,
325 + candidate_blocks: torch.Tensor | None = None,
326 + candidate_block_size: int = 0,
327 + candidate_write: bool = False,
328 ) -> torch.Tensor:
329 # careful! this will be None in dummy run
330 forward_context = get_forward_context()
@@ -323,6 +332,14 @@ def sparse_attn_indexer(
332 fp8_dtype = current_platform.fp8_dtype()
333 k_cache_prefix = _resolve_layer_name(k_cache_prefix)
334
335 + if candidate_blocks is not None:
336 + # Candidate blocks are request-local; the DCP-sharded logits layout
337 + # would need per-rank translation that is not implemented.
338 + assert dcp_world_size == 1, (
339 + "v4.1 two-level candidate filtering is not supported with DCP."
340 + )
341 + assert candidate_block_size > 0
342 +
343 # assert isinstance(attn_metadata, dict)
344 if not isinstance(attn_metadata, dict):
345 # Reserve workspace for indexer during profiling run
@@ -361,6 +378,9 @@ def sparse_attn_indexer(
378 use_pcp,
379 dense_mha_metadata_layer_name,
380 use_fp4_cache,
381 + candidate_blocks=candidate_blocks,
382 + candidate_block_size=candidate_block_size,
383 + candidate_write=candidate_write,
384 )
385 attn_metadata_narrowed = attn_metadata[k_cache_prefix]
386 assert isinstance(attn_metadata_narrowed, DeepseekV32IndexerMetadata)
@@ -506,6 +526,30 @@ def sparse_attn_indexer(
526 clean_logits=False,
527 )
528 num_rows = logits.shape[0]
529 + if candidate_blocks is not None:
530 + # Two-level selection (v4.1): the candidate source
531 + # publishes its top blocks; later indexers mask their
532 + # scores to them. Both before the row top-k.
533 + chunk_candidates = candidate_blocks[
534 + chunk.token_start : chunk.token_end
535 + ]
536 + if candidate_write:
537 + _select_candidate_blocks(
538 + logits,
539 + cu_seqlen_ks,
540 + cu_seqlen_ke,
541 + chunk_candidates.shape[1],
542 + candidate_block_size,
543 + chunk_candidates,
544 + )
545 + else:
546 + _apply_candidate_mask(
547 + logits,
548 + cu_seqlen_ks,
549 + cu_seqlen_ke,
550 + chunk_candidates,
551 + candidate_block_size,
552 + )
553 ops.top_k_per_row_prefill(
554 logits,
555 cu_seqlen_ks,
@@ -612,6 +656,34 @@ def sparse_attn_indexer(
656 indices=decode_metadata.indices,
657 )
658 num_rows = logits.shape[0]
659 + if candidate_blocks is not None:
660 + # Two-level selection (v4.1) on the decode logits; columns are
661 + # request-local compressed positions. seq_lens is (B, next_n)
662 + # for native spec decode (per-row effective lens) and (B, 1)
663 + # otherwise.
664 + vis = seq_lens.reshape(-1)
665 + row_repeat = next_n if vis.numel() != num_rows else 1
666 + vis = vis[:num_rows]
667 + decode_candidates = candidate_blocks[:num_rows]
668 + if candidate_write:
669 + _select_candidate_blocks(
670 + logits,
671 + None,
672 + vis,
673 + decode_candidates.shape[1],
674 + candidate_block_size,
675 + decode_candidates,
676 + row_repeat,
677 + )
678 + else:
679 + _apply_candidate_mask(
680 + logits,
681 + None,
682 + vis,
683 + decode_candidates,
684 + candidate_block_size,
685 + row_repeat,
686 + )
687 topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
688
689 use_cooperative_topk = (
@@ -712,6 +784,9 @@ def sparse_attn_indexer_fake(
784 dcp_world_size: int = 1,
785 cp_kv_cache_interleave_size: int = 1,
786 skip_topk_buffer_clear: bool = False,
787 + candidate_blocks: torch.Tensor | None = None,
788 + candidate_block_size: int = 0,
789 + candidate_write: bool = False,
790 ) -> torch.Tensor:
791 return topk_indices_buffer
792
@@ -719,7 +794,7 @@ def sparse_attn_indexer_fake(
794 direct_register_custom_op(
795 op_name="sparse_attn_indexer",
796 op_func=sparse_attn_indexer,
722 - mutates_args=["topk_indices_buffer"],
797 + mutates_args=["topk_indices_buffer", "candidate_blocks"],
798 fake_impl=sparse_attn_indexer_fake,
799 dispatch_key=current_platform.dispatch_key,
800 )
@@ -751,6 +826,9 @@ class SparseAttnIndexer(CustomOp):
826 skip_k_cache_insert: bool = False,
827 use_fp4_cache: bool = False,
828 compress_ratio: int = 1,
829 + candidate_blocks: torch.Tensor | None = None,
830 + candidate_block_size: int = 0,
831 + candidate_write: bool = False,
832 ):
833 super().__init__()
834 self.k_cache = k_cache
@@ -764,6 +842,11 @@ class SparseAttnIndexer(CustomOp):
842 self.skip_k_cache_insert = skip_k_cache_insert
843 self.use_fp4_cache = use_fp4_cache
844 self.compress_ratio = compress_ratio
845 + # v4.1 two-level selection: the candidate source indexer writes the
846 + # top candidate blocks here; later indexers mask their scores with it.
847 + self.candidate_blocks = candidate_blocks
848 + self.candidate_block_size = candidate_block_size
849 + self.candidate_write = candidate_write
850 self.dense_mha_metadata_layer_name = ""
851 # DCP scalars are constant for the run; resolve them here (config is set
852 # during model construction) and pass them into the custom op, rather
@@ -871,6 +954,9 @@ class SparseAttnIndexer(CustomOp):
954 self.dcp_rank,
955 self.dcp_world_size,
956 self.cp_kv_cache_interleave_size,
957 + candidate_blocks=self.candidate_blocks,
958 + candidate_block_size=self.candidate_block_size,
959 + candidate_write=self.candidate_write,
960 )
961
962 def forward_xpu(
@@ -893,12 +979,16 @@ class SparseAttnIndexer(CustomOp):
979 assert isinstance(q_quant, torch.Tensor), (
980 "AMD sparse_attn_indexer expects a single FP8 q_quant tensor"
981 )
896 - from vllm.platforms.rocm import on_gfx11
982 + from vllm.platforms.rocm import on_gfx11, on_gfx950
983
984 if (
985 rocm_aiter_ops.is_enabled()
986 or rocm_aiter_ops.is_rdna_aiter_enabled()
987 or on_gfx11()
988 + # The so-called AITER sparse indexer op has a native gfx950 path:
989 + # its cache insert, MQA logits, and top-k fallbacks are implemented
990 + # by local Triton/C++ kernels and do not require the aiter package.
991 + or on_gfx950()
992 ):
993 return torch.ops.vllm.rocm_aiter_sparse_attn_indexer(
994 hidden_states,
@@ -916,10 +1006,13 @@ class SparseAttnIndexer(CustomOp):
1006 self.topk_indices_buffer,
1007 skip_k_cache_insert=self.skip_k_cache_insert,
1008 compress_ratio=self.compress_ratio,
1009 + candidate_blocks=self.candidate_blocks,
1010 + candidate_block_size=self.candidate_block_size,
1011 + candidate_write=self.candidate_write,
1012 )
1013 raise RuntimeError(
921 - "Sparse attention indexer ROCm path is only supported on AITER. "
922 - "Please enable aiter with VLLM_ROCM_USE_AITER=1"
1014 + "Sparse attention indexer ROCm path requires AITER or a supported "
1015 + "native architecture (gfx950/gfx11)."
1016 )
1017
1018 def forward_cpu(
vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py
+105 -70
@@ -32,37 +32,39 @@ class FusedInvRopeFP8QuantKernel(
32 fp8_max: float
33 quant_group_size: int
34 chunks_per_head: int
35 - rope_start: int
35 + nope_dim: int
36 half_rope: int
37 tma_aligned_scales: bool
38 launch_pdl: bool
39 + quantize: bool
40
41 @staticmethod
42 # scale_stride_k = align(num_tokens, 4) is only 4-aligned, so its Triton int
43 # class varies with batch size; not specialized so one warmup key covers all
43 - # (int64-cast, scalar addressing only, so perf-neutral).
44 + # (int64-cast, so perf-neutral).
45 @triton.jit(do_not_specialize=["num_tokens", "scale_stride_k"])
46 def kernel(
47 o_ptr,
48 positions_ptr,
49 cos_sin_cache_ptr,
49 - fp8_ptr,
50 + out_ptr,
51 scale_ptr,
52 num_tokens,
53 heads_per_group: tl.constexpr,
54 o_stride_token,
55 o_stride_head,
56 cache_stride_pos,
56 - fp8_stride_group,
57 - fp8_stride_token,
57 + out_stride_group,
58 + out_stride_token,
59 scale_stride_group,
60 scale_stride_k,
61 fp8_max: tl.constexpr,
62 eps: tl.constexpr,
63 QUANT_GROUP_SIZE: tl.constexpr,
64 CHUNKS_PER_HEAD: tl.constexpr,
64 - ROPE_START: tl.constexpr,
65 + NOPE_DIM: tl.constexpr,
66 HALF_ROPE: tl.constexpr,
67 + QUANTIZE: tl.constexpr,
68 TMA_ALIGNED_SCALES: tl.constexpr,
69 launch_pdl: tl.constexpr,
70 ):
@@ -74,8 +76,8 @@ class FusedInvRopeFP8QuantKernel(
76 o_stride_token = o_stride_token.to(tl.int64)
77 o_stride_head = o_stride_head.to(tl.int64)
78 cache_stride_pos = cache_stride_pos.to(tl.int64)
77 - fp8_stride_group = fp8_stride_group.to(tl.int64)
78 - fp8_stride_token = fp8_stride_token.to(tl.int64)
79 + out_stride_group = out_stride_group.to(tl.int64)
80 + out_stride_token = out_stride_token.to(tl.int64)
81 scale_stride_group = scale_stride_group.to(tl.int64)
82 scale_stride_k = scale_stride_k.to(tl.int64)
83
@@ -88,14 +90,18 @@ class FusedInvRopeFP8QuantKernel(
90 tl.extra.cuda.gdc_wait()
91 # Padding rows in the TMA-aligned scale buffer: fill with zero and skip quant.
92 if pid_token >= num_tokens:
93 + if not QUANTIZE:
94 + return
95 if TMA_ALIGNED_SCALES:
96 + packed_offsets = tl.arange(0, CHUNKS_PER_HEAD // 4)
97 scale_addr = (
98 scale_ptr
99 + g * scale_stride_group
100 + pid_token
96 - + head_in_group * scale_stride_k
101 + + (head_in_group * (CHUNKS_PER_HEAD // 4) + packed_offsets)
102 + * scale_stride_k
103 )
98 - tl.store(scale_addr, tl.zeros((), dtype=tl.int32))
104 + tl.store(scale_addr, tl.zeros((CHUNKS_PER_HEAD // 4,), dtype=tl.int32))
105 else:
106 block_offsets = tl.arange(0, CHUNKS_PER_HEAD)
107 qb_indices = qb_start + block_offsets
@@ -114,9 +120,7 @@ class FusedInvRopeFP8QuantKernel(
120 offsets = tl.arange(0, HEAD_DIM)
121 x = tl.load(input_base + offsets).to(tl.float32)
122
117 - rope_abs_start: tl.constexpr = (
118 - CHUNKS_PER_HEAD - 1
119 - ) * QUANT_GROUP_SIZE + ROPE_START
123 + rope_abs_start: tl.constexpr = NOPE_DIM
124 pos = tl.load(positions_ptr + pid_token)
125 cache_base = cos_sin_cache_ptr + pos * cache_stride_pos
126 is_rope = offsets >= rope_abs_start
@@ -134,6 +138,16 @@ class FusedInvRopeFP8QuantKernel(
138 rotated = tl.where(is_even, x_add, x_sub)
139 x = tl.where(is_rope, rotated, x)
140
141 + if not QUANTIZE:
142 + out_base = (
143 + out_ptr
144 + + g * out_stride_group
145 + + pid_token * out_stride_token
146 + + qb_start * QUANT_GROUP_SIZE
147 + )
148 + tl.store(out_base + offsets, x)
149 + return
150 +
151 x_2d = tl.reshape(tl.abs(x), (CHUNKS_PER_HEAD, QUANT_GROUP_SIZE))
152 block_absmax = tl.maximum(tl.max(x_2d, axis=1), eps)
153 scale_raw = block_absmax * (1.0 / fp8_max)
@@ -148,25 +162,31 @@ class FusedInvRopeFP8QuantKernel(
162 )
163 x_quant = tl.clamp(x / scales_exp, -fp8_max, fp8_max).to(tl.float8e4nv)
164
151 - fp8_base = (
152 - fp8_ptr
153 - + g * fp8_stride_group
154 - + pid_token * fp8_stride_token
165 + out_base = (
166 + out_ptr
167 + + g * out_stride_group
168 + + pid_token * out_stride_token
169 + qb_start * QUANT_GROUP_SIZE
170 )
157 - tl.store(fp8_base + offsets, x_quant)
171 + tl.store(out_base + offsets, x_quant)
172
173 block_offsets = tl.arange(0, CHUNKS_PER_HEAD)
174 qb_indices = qb_start + block_offsets
175 if TMA_ALIGNED_SCALES:
176 scale_bits = scales.to(tl.int32, bitcast=True)
177 ue8m0_bytes = (scale_bits >> 23) & 0xFF
164 - packed_val = tl.sum(ue8m0_bytes << (block_offsets * 8))
178 + packed_val = tl.sum(
179 + tl.reshape(ue8m0_bytes, (CHUNKS_PER_HEAD // 4, 4))
180 + << (tl.arange(0, 4)[None, :] * 8),
181 + axis=1,
182 + )
183 + packed_offsets = tl.arange(0, CHUNKS_PER_HEAD // 4)
184 scale_addr = (
185 scale_ptr
186 + g * scale_stride_group
187 + pid_token
169 - + head_in_group * scale_stride_k
188 + + (head_in_group * (CHUNKS_PER_HEAD // 4) + packed_offsets)
189 + * scale_stride_k
190 )
191 tl.store(scale_addr, packed_val)
192 else:
@@ -188,7 +208,7 @@ class FusedInvRopeFP8QuantKernel(
208 quant_group_size: int,
209 runtime_fp8_max: float | None = None,
210 runtime_chunks_per_head: int | None = None,
191 - runtime_rope_start: int | None = None,
211 + runtime_nope_dim: int | None = None,
212 runtime_half_rope: int | None = None,
213 **compile_key_fields: bool,
214 ) -> CompileKey:
@@ -202,11 +222,7 @@ class FusedInvRopeFP8QuantKernel(
222 if runtime_chunks_per_head is not None
223 else head_dim // quant_group_size
224 )
205 - rope_start = (
206 - runtime_rope_start
207 - if runtime_rope_start is not None
208 - else nope_dim % quant_group_size
209 - )
225 + nope_dim = runtime_nope_dim if runtime_nope_dim is not None else nope_dim
226 half_rope = (
227 runtime_half_rope if runtime_half_rope is not None else rope_dim // 2
228 )
@@ -216,7 +232,7 @@ class FusedInvRopeFP8QuantKernel(
232 fp8_max=fp8_max,
233 quant_group_size=quant_group_size,
234 chunks_per_head=chunks_per_head,
219 - rope_start=rope_start,
235 + nope_dim=nope_dim,
236 half_rope=half_rope,
237 )
238
@@ -250,11 +266,12 @@ class FusedInvRopeFP8QuantKernel(
266 quant_group_size=128,
267 tma_aligned_scales=capability.major >= 10,
268 launch_pdl=current_platform.is_arch_support_pdl(),
269 + quantize=True,
270 )
271
272 def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
273 head_dim = compile_key.chunks_per_head * compile_key.quant_group_size
257 - fp8_dim = compile_key.heads_per_group * head_dim
274 + out_dim = compile_key.heads_per_group * head_dim
275 scale_dtype = torch.int32 if compile_key.tma_aligned_scales else torch.float32
276 return dict(
277 o=TritonWarmupTensor(
@@ -266,10 +283,10 @@ class FusedInvRopeFP8QuantKernel(
283 torch.float32,
284 shape=(1, compile_key.half_rope * 2),
285 ),
269 - fp8_buf=TritonWarmupTensor(
270 - torch.float8_e4m3fn,
271 - shape=(compile_key.heads_per_group, 1, fp8_dim),
272 - strides=(fp8_dim, fp8_dim, 1),
286 + out_buf=TritonWarmupTensor(
287 + torch.float8_e4m3fn if compile_key.quantize else torch.bfloat16,
288 + shape=(compile_key.heads_per_group, 1, out_dim),
289 + strides=(out_dim, out_dim, 1),
290 ),
291 scale_buf=TritonWarmupTensor(
292 scale_dtype,
@@ -280,9 +297,10 @@ class FusedInvRopeFP8QuantKernel(
297 heads_per_group=compile_key.heads_per_group,
298 quant_group_size=compile_key.quant_group_size,
299 chunks_per_head=compile_key.chunks_per_head,
283 - rope_start=compile_key.rope_start,
300 + nope_dim=compile_key.nope_dim,
301 half_rope=compile_key.half_rope,
302 tma_aligned_scales=compile_key.tma_aligned_scales,
303 + quantize=compile_key.quantize,
304 fp8_max=compile_key.fp8_max,
305 launch_pdl=compile_key.launch_pdl,
306 grid=(1, compile_key.heads_per_group),
@@ -294,35 +312,37 @@ class FusedInvRopeFP8QuantKernel(
312 o: torch.Tensor,
313 positions: torch.Tensor,
314 cos_sin_cache: torch.Tensor,
297 - fp8_buf: torch.Tensor,
315 + out_buf: torch.Tensor,
316 scale_buf: torch.Tensor,
317 num_tokens: int,
318 *,
319 heads_per_group: int,
320 quant_group_size: int,
321 chunks_per_head: int,
304 - rope_start: int,
322 + nope_dim: int,
323 half_rope: int,
324 tma_aligned_scales: bool,
325 + quantize: bool,
326 fp8_max: float,
327 launch_pdl: bool,
328 grid: tuple[int, int],
329 ) -> LaunchSpec:
330 return grid, dict(
312 - fp8_ptr=fp8_buf,
331 + out_ptr=out_buf,
332 scale_ptr=scale_buf,
333 o_stride_token=o.stride(0),
334 o_stride_head=o.stride(1),
335 cache_stride_pos=cos_sin_cache.stride(0),
317 - fp8_stride_group=fp8_buf.stride(0),
318 - fp8_stride_token=fp8_buf.stride(1),
319 - scale_stride_group=scale_buf.stride(0),
320 - scale_stride_k=scale_buf.stride(2),
336 + out_stride_group=out_buf.stride(0),
337 + out_stride_token=out_buf.stride(1),
338 + scale_stride_group=scale_buf.stride(0) if quantize else 0,
339 + scale_stride_k=scale_buf.stride(2) if quantize else 0,
340 eps=1e-10,
341 QUANT_GROUP_SIZE=quant_group_size,
342 CHUNKS_PER_HEAD=chunks_per_head,
324 - ROPE_START=rope_start,
343 + NOPE_DIM=nope_dim,
344 HALF_ROPE=half_rope,
345 + QUANTIZE=quantize,
346 TMA_ALIGNED_SCALES=tma_aligned_scales,
347 launch_pdl=launch_pdl,
348 num_stages=1,
@@ -340,6 +360,7 @@ def fused_inv_rope_fp8_quant(
360 rope_dim: int = 64,
361 quant_group_size: int = 128,
362 tma_aligned_scales: bool = False,
363 + quantize: bool = True,
364 ) -> tuple[torch.Tensor, torch.Tensor]:
365 """Fused inverse RoPE + block-scaled FP8 quantization.
366
@@ -354,10 +375,11 @@ def fused_inv_rope_fp8_quant(
375 quant_group_size: FP8 quantization block size (default 128).
376 tma_aligned_scales: Output INT32 packed UE8M0 for SM100 (True)
377 or FP32 for SM90 (False).
378 + quantize: Quantize the rotated output to FP8 and return its scales.
379
380 Returns:
359 - o_fp8: [T, G, D] float8_e4m3fn, strides (D, T*D, 1).
360 - o_scale: Pre-transformed scale tensor for fp8_einsum.
381 + Rotated output in [T, G, D] and its FP8 scales. The scale tensor is
382 + empty when quantization is disabled.
383 """
384 from vllm.utils.deep_gemm import get_tma_aligned_size
385
@@ -365,7 +387,6 @@ def fused_inv_rope_fp8_quant(
387 assert num_heads == n_groups * heads_per_group
388 assert head_dim == nope_dim + rope_dim
389 assert head_dim % quant_group_size == 0
368 - assert nope_dim % quant_group_size == (quant_group_size - rope_dim)
390 assert rope_dim % 2 == 0
391 assert cos_sin_cache.shape[-1] == rope_dim
392 assert cos_sin_cache.dtype == torch.float32
@@ -377,23 +398,26 @@ def fused_inv_rope_fp8_quant(
398 fp8_dtype = torch.float8_e4m3fn
399 fp8_max = torch.finfo(fp8_dtype).max
400
380 - tma_aligned_T = get_tma_aligned_size(num_tokens, 4)
381 - if tma_aligned_scales:
401 + tma_aligned_T = get_tma_aligned_size(num_tokens, 4) if quantize else num_tokens
402 + if quantize and tma_aligned_scales:
403 + assert chunks_per_head % 4 == 0
404 packed_sf_k = (num_scale_blocks + 3) // 4
405 scale_inner = packed_sf_k
384 - else:
406 + elif quantize:
407 scale_inner = num_scale_blocks
408 + else:
409 + scale_inner = 0
410
411 # Run kernel through a custom op so inductor sees an opaque boundary.
412 # It's a pytorch bug, see https://github.com/vllm-project/vllm/issues/41106
389 - fp8_buf, scale_buf = torch.ops.vllm.fused_inv_rope_fp8_quant_kernel(
413 + out_buf, scale_buf = torch.ops.vllm.fused_inv_rope_fp8_quant_kernel(
414 o,
415 positions,
416 cos_sin_cache,
417 heads_per_group,
418 quant_group_size,
419 chunks_per_head,
396 - nope_dim % quant_group_size,
420 + nope_dim,
421 rope_dim // 2,
422 tma_aligned_scales,
423 fp8_max,
@@ -402,8 +426,11 @@ def fused_inv_rope_fp8_quant(
426 n_groups,
427 d,
428 scale_inner,
429 + quantize,
430 )
406 - return fp8_buf.transpose(0, 1), scale_buf.transpose(0, 1)
431 + output = out_buf.transpose(0, 1)
432 + scales = scale_buf.transpose(0, 1) if quantize else scale_buf
433 + return output, scales
434
435
436 def _fused_inv_rope_fp8_quant_kernel_impl(
@@ -413,7 +440,7 @@ def _fused_inv_rope_fp8_quant_kernel_impl(
440 heads_per_group: int,
441 quant_group_size: int,
442 chunks_per_head: int,
416 - rope_start: int,
443 + nope_dim: int,
444 half_rope: int,
445 tma_aligned_scales: bool,
446 fp8_max: float,
@@ -422,41 +449,46 @@ def _fused_inv_rope_fp8_quant_kernel_impl(
449 n_groups: int,
450 d: int,
451 scale_inner: int,
452 + quantize: bool,
453 ) -> tuple[torch.Tensor, torch.Tensor]:
426 - fp8_buf = torch.empty(
427 - (n_groups, num_tokens, d),
428 - dtype=torch.float8_e4m3fn,
429 - device=o.device,
430 - )
454 scale_dtype = torch.int32 if tma_aligned_scales else torch.float32
432 - scale_buf = torch.empty(
433 - n_groups * scale_inner * tma_aligned_T,
434 - dtype=scale_dtype,
455 + out_buf = torch.empty(
456 + (n_groups, num_tokens, d),
457 + dtype=torch.float8_e4m3fn if quantize else o.dtype,
458 device=o.device,
436 - ).as_strided(
437 - (n_groups, num_tokens, scale_inner),
438 - (scale_inner * tma_aligned_T, 1, tma_aligned_T),
459 )
460 + if quantize:
461 + scale_buf = torch.empty(
462 + n_groups * scale_inner * tma_aligned_T,
463 + dtype=scale_dtype,
464 + device=o.device,
465 + ).as_strided(
466 + (n_groups, num_tokens, scale_inner),
467 + (scale_inner * tma_aligned_T, 1, tma_aligned_T),
468 + )
469 + else:
470 + scale_buf = torch.empty(0, dtype=scale_dtype, device=o.device)
471 grid = (tma_aligned_T, n_groups * heads_per_group)
472 launch_pdl = current_platform.is_arch_support_pdl()
473 _FUSED_INV_ROPE_FP8_QUANT_KERNEL(
474 o,
475 positions,
476 cos_sin_cache,
446 - fp8_buf,
477 + out_buf,
478 scale_buf,
479 num_tokens,
480 heads_per_group=heads_per_group,
481 quant_group_size=quant_group_size,
482 chunks_per_head=chunks_per_head,
452 - rope_start=rope_start,
483 + nope_dim=nope_dim,
484 half_rope=half_rope,
485 tma_aligned_scales=tma_aligned_scales,
486 + quantize=quantize,
487 fp8_max=fp8_max,
488 launch_pdl=launch_pdl,
489 grid=grid,
490 )
459 - return fp8_buf, scale_buf
491 + return out_buf, scale_buf
492
493
494 def _fused_inv_rope_fp8_quant_kernel_fake(
@@ -466,7 +498,7 @@ def _fused_inv_rope_fp8_quant_kernel_fake(
498 heads_per_group: int,
499 quant_group_size: int,
500 chunks_per_head: int,
469 - rope_start: int,
501 + nope_dim: int,
502 half_rope: int,
503 tma_aligned_scales: bool,
504 fp8_max: float,
@@ -475,13 +507,16 @@ def _fused_inv_rope_fp8_quant_kernel_fake(
507 n_groups: int,
508 d: int,
509 scale_inner: int,
510 + quantize: bool,
511 ) -> tuple[torch.Tensor, torch.Tensor]:
479 - fp8_buf = torch.empty(
512 + scale_dtype = torch.int32 if tma_aligned_scales else torch.float32
513 + out_buf = torch.empty(
514 (n_groups, num_tokens, d),
481 - dtype=torch.float8_e4m3fn,
515 + dtype=torch.float8_e4m3fn if quantize else o.dtype,
516 device=o.device,
517 )
484 - scale_dtype = torch.int32 if tma_aligned_scales else torch.float32
518 + if not quantize:
519 + return out_buf, torch.empty(0, dtype=scale_dtype, device=o.device)
520 scale_buf = torch.empty(
521 n_groups * scale_inner * tma_aligned_T,
522 dtype=scale_dtype,
@@ -490,7 +525,7 @@ def _fused_inv_rope_fp8_quant_kernel_fake(
525 (n_groups, num_tokens, scale_inner),
526 (scale_inner * tma_aligned_T, 1, tma_aligned_T),
527 )
493 - return fp8_buf, scale_buf
528 + return out_buf, scale_buf
529
530
531 direct_register_custom_op(
vllm/models/deepseek_v4/common/vision.py
+93 -2
@@ -11,13 +11,18 @@ embed and norms are replicated, so the residual stream is full-width on
11 every rank.
12 """
13
14 +import itertools
15 from functools import lru_cache
16
17 import torch
18 import torch.nn.functional as F
19 from torch import nn
20
20 -from vllm.distributed import get_tensor_model_parallel_world_size
21 +from vllm.distributed import (
22 + get_tensor_model_parallel_rank,
23 + get_tensor_model_parallel_world_size,
24 + tensor_model_parallel_all_gather,
25 +)
26 from vllm.model_executor.layers.activation import SiluAndMul
27 from vllm.model_executor.layers.attention.mm_encoder_attention import (
28 MMEncoderAttention,
@@ -29,7 +34,10 @@ from vllm.model_executor.layers.linear import (
34 ReplicatedLinear,
35 RowParallelLinear,
36 )
32 -from vllm.model_executor.models.vision import is_vit_use_data_parallel
37 +from vllm.model_executor.models.vision import (
38 + get_load_balance_assignment,
39 + is_vit_use_data_parallel,
40 +)
41
42
43 @lru_cache(8)
@@ -206,6 +214,7 @@ class DeepseekV4Aligner(nn.Module):
214 super().__init__()
215 use_data_parallel = is_vit_use_data_parallel(config.vision_n_heads)
216 self.downsample_ratio = config.vision_downsample_ratio
217 + self.out_dim = config.hidden_size
218 in_dim = config.vision_dim * self.downsample_ratio**2
219 self.w1 = ColumnParallelLinear(
220 in_dim,
@@ -230,3 +239,85 @@ class DeepseekV4Aligner(nn.Module):
239 hidden, _ = self.w1(x)
240 out, _ = self.w2(F.gelu(hidden))
241 return out
242 +
243 +
244 +def run_dp_sharded_vision_tower(
245 + vision_model: DeepseekV4ViT,
246 + aligner: DeepseekV4Aligner,
247 + patches: torch.Tensor,
248 + vit_grid: list[list[int]],
249 +) -> list[torch.Tensor]:
250 + """Run the ViT + aligner with images sharded across TP ranks.
251 +
252 + Every rank holds the full tower weights (``--mm-encoder-tp-mode data``)
253 + and receives the full ``patches`` batch. Images are assigned to ranks by
254 + patch count (greedy load balancing), each rank encodes only its share,
255 + and per-image embeddings are exchanged with one padded all-gather and
256 + returned in the original image order.
257 +
258 + Args:
259 + vision_model: The (weight-replicated) ViT tower.
260 + aligner: The (weight-replicated) spatial-merge projector.
261 + patches: ``(sum(n_vit_h * n_vit_w), 3, p, p)`` patches of all images.
262 + vit_grid: ``[n_vit_h, n_vit_w]`` per image.
263 +
264 + Returns:
265 + One ``(n_aligner_rows, hidden_size)`` embedding tensor per image.
266 + """
267 + tp_size = get_tensor_model_parallel_world_size()
268 + tp_rank = get_tensor_model_parallel_rank()
269 +
270 + sizes = [h * w for h, w in vit_grid]
271 + cum_patches = [0, *itertools.accumulate(sizes)]
272 + image_to_tp_rank, gpu_sample_counts, _ = get_load_balance_assignment(sizes, tp_size)
273 + cum_sample_counts = [0, *itertools.accumulate(gpu_sample_counts)]
274 +
275 + # Rows the aligner emits per image: each grid dim is padded up to a
276 + # multiple of the merge ratio before r x r patches fold into one row.
277 + r = aligner.downsample_ratio
278 + rows_per_image = [-(-h // r) * -(-w // r) for h, w in vit_grid]
279 +
280 + def rank_image_idxs(g: int) -> list[int]:
281 + return image_to_tp_rank[cum_sample_counts[g] : cum_sample_counts[g + 1]]
282 +
283 + # All-gather needs one shape on every rank; pad each rank's packed
284 + # output to the largest per-rank row count (computable locally).
285 + max_rows = max(
286 + sum(rows_per_image[i] for i in rank_image_idxs(g)) for g in range(tp_size)
287 + )
288 +
289 + local_embeds = [
290 + aligner(
291 + vision_model(
292 + patches[cum_patches[i] : cum_patches[i + 1]],
293 + vit_grid[i][0],
294 + vit_grid[i][1],
295 + ),
296 + vit_grid[i][0],
297 + vit_grid[i][1],
298 + )
299 + for i in rank_image_idxs(tp_rank)
300 + ]
301 + if local_embeds:
302 + embeds_local = torch.cat(local_embeds, dim=0)
303 + else:
304 + embeds_local = patches.new_zeros((0, aligner.out_dim))
305 + if embeds_local.shape[0] < max_rows:
306 + embeds_local = torch.cat(
307 + [
308 + embeds_local,
309 + embeds_local.new_zeros(
310 + (max_rows - embeds_local.shape[0], embeds_local.shape[1])
311 + ),
312 + ],
313 + dim=0,
314 + )
315 + gathered = tensor_model_parallel_all_gather(embeds_local.contiguous(), dim=0)
316 +
317 + out: list[torch.Tensor] = [None] * len(vit_grid) # type: ignore[list-item]
318 + for g in range(tp_size):
319 + offset = g * max_rows
320 + for i in rank_image_idxs(g):
321 + out[i] = gathered[offset : offset + rows_per_image[i]]
322 + offset += rows_per_image[i]
323 + return out
vllm/models/deepseek_v4/nvidia/model.py
+51 -15
@@ -368,9 +368,17 @@ class DeepseekV4MegaMoEExperts(nn.Module):
368 gate_up = shared_experts.gate_up_proj
369 down = shared_experts.down_proj
370 gate_up_weight = gate_up.weight.data
371 - gate_up_scale = gate_up.weight_scale_inv.data
371 + gate_up_scale = (
372 + gate_up.weight_scale
373 + if hasattr(gate_up, "weight_scale")
374 + else gate_up.weight_scale_inv
375 + ).data
376 down_weight = down.weight.data
373 - down_scale = down.weight_scale_inv.data
377 + down_scale = (
378 + down.weight_scale
379 + if hasattr(down, "weight_scale")
380 + else down.weight_scale_inv
381 + ).data
382
383 # MegaMoE's shared FP8 MMA consumes a 1x32 scale for every weight row,
384 # while the checkpoint uses coarser block-FP8 scales (usually
@@ -507,6 +515,23 @@ class DeepseekV4MegaMoEExperts(nn.Module):
515
516 if self._transformed_l1_weights is None:
517 self._check_runtime_supported()
518 + # MegaMoE's 1x32 activation scales need 16-byte TMA rows, so
519 + # pad each gate/up half and the down projection to a multiple of 512.
520 + padded_size = (self.intermediate_size + 511) // 512 * 512
521 + padding = padded_size - self.intermediate_size
522 + if padding:
523 + for param in (self.w13_weight, self.w13_weight_scale):
524 + gate_up = param.data.unflatten(1, (2, self.intermediate_size))
525 + param.data = torch.nn.functional.pad(
526 + gate_up, (0, 0, 0, padding)
527 + ).flatten(1, 2)
528 + self.w2_weight.data = torch.nn.functional.pad(
529 + self.w2_weight.data, (0, padding // 2)
530 + )
531 + self.w2_weight_scale.data = torch.nn.functional.pad(
532 + self.w2_weight_scale.data, (0, padding // 32)
533 + )
534 + self.intermediate_size = padded_size
535 w13_scale = deep_gemm.transform_sf_into_required_layout(
536 self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(),
537 2 * self.intermediate_size,
@@ -752,6 +777,11 @@ class DeepseekV4MoE(nn.Module):
777 vllm_config: VllmConfig,
778 prefix: str = "",
779 use_sequence_parallel: bool = False,
780 + *,
781 + num_hash_layers: int,
782 + n_routed_experts: int | None = None,
783 + n_activated_experts: int | None = None,
784 + image_sentinel_lo: int = IMAGE_SENTINEL_BASE_ID,
785 ):
786 super().__init__()
787
@@ -774,8 +804,14 @@ class DeepseekV4MoE(nn.Module):
804 self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
805 self.hidden_size = config.hidden_size
806
777 - self.n_routed_experts = config.n_routed_experts
778 - self.n_activated_experts = config.num_experts_per_tok
807 + self.n_routed_experts = (
808 + config.n_routed_experts if n_routed_experts is None else n_routed_experts
809 + )
810 + self.n_activated_experts = (
811 + config.num_experts_per_tok
812 + if n_activated_experts is None
813 + else n_activated_experts
814 + )
815 self.moe_intermediate_size = config.moe_intermediate_size
816 self.swiglu_limit = config.swiglu_limit
817 self.renormalize = config.norm_topk_prob
@@ -793,7 +829,7 @@ class DeepseekV4MoE(nn.Module):
829
830 self.gate = GateLinear(
831 input_size=config.hidden_size,
796 - output_size=config.n_routed_experts,
832 + output_size=self.n_routed_experts,
833 bias=False,
834 out_dtype=torch.float32,
835 prefix=f"{prefix}.gate",
@@ -805,9 +841,9 @@ class DeepseekV4MoE(nn.Module):
841 # Image tokens borrow five consecutive reserved in-vocab ids starting
842 # at IMAGE_SENTINEL_BASE_ID; 0 disables vision routing (text model).
843 self.image_sentinel_lo = (
808 - IMAGE_SENTINEL_BASE_ID if getattr(config, "vision_n_layers", 0) > 0 else 0
844 + image_sentinel_lo if getattr(config, "vision_n_layers", 0) > 0 else 0
845 )
810 - is_hash_moe = extract_layer_index(prefix) < config.num_hash_layers
846 + is_hash_moe = extract_layer_index(prefix) < num_hash_layers
847 self.hash_indices_dtype = torch.int64 if self.use_mega_moe else torch.int32
848 if is_hash_moe:
849 # hash MoE doesn't use e_score_correction_bias
@@ -816,8 +852,8 @@ class DeepseekV4MoE(nn.Module):
852 self.gate.tid2eid = nn.Parameter(
853 torch.randint(
854 0,
819 - config.n_routed_experts,
820 - (config.vocab_size, config.num_experts_per_tok),
855 + self.n_routed_experts,
856 + (config.vocab_size, self.n_activated_experts),
857 dtype=self.hash_indices_dtype,
858 ),
859 requires_grad=False,
@@ -828,7 +864,7 @@ class DeepseekV4MoE(nn.Module):
864 # Vision checkpoints ship a gate bias on hash layers too (it is
865 # unused for routing there; image tokens use bias_vl instead).
866 self.gate.e_score_correction_bias = nn.Parameter(
831 - torch.empty(config.n_routed_experts, dtype=torch.float32),
867 + torch.empty(self.n_routed_experts, dtype=torch.float32),
868 requires_grad=False,
869 )
870
@@ -837,7 +873,7 @@ class DeepseekV4MoE(nn.Module):
873 # instead of e_score_correction_bias / the hash table. Created on
874 # every MoE layer, hash layers included.
875 self.gate.bias_vl = nn.Parameter(
840 - torch.empty(config.n_routed_experts, dtype=torch.float32),
876 + torch.empty(self.n_routed_experts, dtype=torch.float32),
877 requires_grad=False,
878 )
879
@@ -874,7 +910,6 @@ class DeepseekV4MoE(nn.Module):
910
911 eplb_config = vllm_config.parallel_config.eplb_config
912 self.n_redundant_experts = eplb_config.num_redundant_experts
877 - self.n_routed_experts = config.n_routed_experts
913 self.n_shared_experts = config.n_shared_experts or 0
914 self.n_logical_experts = self.n_routed_experts
915 self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
@@ -922,7 +957,7 @@ class DeepseekV4MoE(nn.Module):
957 num_local_experts=self.n_local_physical_experts,
958 experts_start_idx=self.physical_expert_start,
959 num_logical_experts=self.n_logical_experts,
925 - top_k=config.num_experts_per_tok,
960 + top_k=self.n_activated_experts,
961 hidden_size=config.hidden_size,
962 intermediate_size=config.moe_intermediate_size,
963 num_shared_experts=(self.n_shared_experts if fuse_shared_experts else 0),
@@ -961,8 +996,8 @@ class DeepseekV4MoE(nn.Module):
996 self.experts = FusedMoEFactory(
997 shared_experts=self.shared_experts,
998 gate=self.gate,
964 - num_experts=config.n_routed_experts,
965 - top_k=config.num_experts_per_tok,
999 + num_experts=self.n_routed_experts,
1000 + top_k=self.n_activated_experts,
1001 hidden_size=config.hidden_size,
1002 intermediate_size=config.moe_intermediate_size,
1003 renormalize=config.norm_topk_prob,
@@ -1120,6 +1155,7 @@ class DeepseekV4DecoderLayer(nn.Module):
1155 vllm_config,
1156 prefix=f"{prefix}.ffn",
1157 use_sequence_parallel=self.use_sequence_parallel,
1158 + num_hash_layers=config.num_hash_layers,
1159 )
1160
1161 self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps)
vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py
+7 -3
@@ -165,9 +165,13 @@ class IndexerQRopeQuantKernel:
165 # NOTE: token_id may exceed bounds, hence we need to add load/store guards
166 # we can't do early exit because CuteDSL doesn't support it. and we also need
167 # all threads in a warp to be active since we utilize warp shuffle later.
168 - # must_in_bounds is constexpr, True when 1 threadblock fit within 1 token
169 - # position. the compiler will remove bounds check when that happens.
170 - must_in_bounds = cutlass.const_expr(self.tb_size % self.threads_per_token == 0)
168 + # must_in_bounds is constexpr, True when a token spans a whole number of
169 + # threadblocks. Then num_tokens * threads_per_token is always a multiple of
170 + # tb_size, so the trailing block never straddles num_tokens and the compiler
171 + # removes the bounds check. Otherwise the guard has to stay: e.g. 32 heads
172 + # with coarsen=4 needs 64 threads per token, so an odd num_tokens leaves half
173 + # of the last block addressing token num_tokens.
174 + must_in_bounds = cutlass.const_expr(self.threads_per_token % self.tb_size == 0)
175 in_bounds = must_in_bounds or (token_id < num_tokens)
176
177 cp_op = cute.nvgpu.CopyUniversalOp()
vllm/models/deepseek_v4/nvidia/ops/o_proj.py
+32 -18
@@ -10,17 +10,18 @@ from vllm.platforms import current_platform
10 from vllm.utils.deep_gemm import fp8_einsum
11
12
13 -def compute_fp8_einsum_recipe() -> tuple[tuple[int, int, int], bool]:
13 +def compute_fp8_einsum_recipe(
14 + block_size: int = 128,
15 +) -> tuple[tuple[int, int, int], bool]:
16 """fp8_einsum recipe + scale layout for the current GPU arch.
17
16 - SM90: FP32 block scales stay [g, r/128, d/128] → sfb_gran_mn=128.
17 - SM100: INT32 packed scales become [g, r, ...] → sfb_gran_mn=1.
18 + SM90 keeps block-row FP32 scales. SM100 uses packed per-row E8M0 scales.
19
20 Returns ``(einsum_recipe, tma_aligned_scales)`` for ``deep_gemm_fp8_o_proj``.
21 """
22 cap = current_platform.get_device_capability()
23 assert cap is not None, "DeepseekV4 attention requires a CUDA device"
23 - einsum_recipe = (1, 128, 128) if cap.major <= 9 else (1, 1, 128)
24 + einsum_recipe = (1, 128, 128) if cap.major <= 9 else (1, 1, block_size)
25 tma_aligned_scales = cap.major >= 10
26 return einsum_recipe, tma_aligned_scales
27
@@ -40,12 +41,13 @@ def deep_gemm_fp8_o_proj(
41 einsum_recipe: tuple[int, int, int],
42 tma_aligned_scales: bool,
43 ) -> torch.Tensor:
43 - """O projection: inverse RoPE + FP8 quant + einsum + wo_b.
44 + """O projection: inverse RoPE + grouped wo_a + wo_b.
45
45 - Shared by the FlashMLA and FlashInfer CUDA backends. ``einsum_recipe`` /
46 - ``tma_aligned_scales`` come from ``compute_fp8_einsum_recipe``.
46 + Shared by the FlashMLA and FlashInfer CUDA backends. The attention
47 + layer selects the recipe at initialization.
48 """
48 - o_fp8, o_scale = fused_inv_rope_fp8_quant(
49 + use_fp8 = wo_a.weight.dtype == torch.float8_e4m3fn
50 + o_proj_input, o_scale = fused_inv_rope_fp8_quant(
51 o,
52 positions,
53 cos_sin_cache,
@@ -53,21 +55,33 @@ def deep_gemm_fp8_o_proj(
55 heads_per_group=heads_per_group,
56 nope_dim=nope_dim,
57 rope_dim=rope_dim,
58 + quant_group_size=einsum_recipe[2],
59 tma_aligned_scales=tma_aligned_scales,
60 + quantize=use_fp8,
61 )
62 z = torch.empty(
63 (o.shape[0], n_groups, o_lora_rank),
64 device=o.device,
65 dtype=torch.bfloat16,
66 )
63 - weight_scale = (
64 - wo_a.weight_scale if hasattr(wo_a, "weight_scale") else wo_a.weight_scale_inv
65 - )
66 - fp8_einsum(
67 - "bhr,hdr->bhd",
68 - (o_fp8, o_scale),
69 - (wo_a.weight, weight_scale),
70 - z,
71 - recipe=einsum_recipe,
72 - )
67 + if use_fp8:
68 + weight_scale = (
69 + wo_a.weight_scale
70 + if hasattr(wo_a, "weight_scale")
71 + else wo_a.weight_scale_inv
72 + )
73 + fp8_einsum(
74 + "bhr,hdr->bhd",
75 + (o_proj_input, o_scale),
76 + (wo_a.weight, weight_scale),
77 + z,
78 + recipe=einsum_recipe,
79 + )
80 + else:
81 + grouped_weight = wo_a.weight.view(n_groups, o_lora_rank, -1)
82 + torch.bmm(
83 + o_proj_input.transpose(0, 1),
84 + grouped_weight.transpose(1, 2),
85 + out=z.transpose(0, 1),
86 + )
87 return wo_b(z.flatten(1))
vllm/models/deepseek_v4_1/__init__.py new
+20
@@ -0,0 +1,20 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""DeepSeek V4.1 hardware-isolated model entry point."""
4 +
5 +from vllm.platforms import current_platform
6 +
7 +from .quant_config import DeepseekV4FP8Config
8 +
9 +if current_platform.is_rocm():
10 + from .amd.dspark import DSparkDeepseekV4ForCausalLM
11 + from .amd.vl_model import DeepseekV41ForCausalLM
12 +else:
13 + from .nvidia.dspark import DSparkDeepseekV4ForCausalLM
14 + from .nvidia.vl_model import DeepseekV41ForCausalLM
15 +
16 +__all__ = [
17 + "DSparkDeepseekV4ForCausalLM",
18 + "DeepseekV4FP8Config",
19 + "DeepseekV41ForCausalLM",
20 +]
vllm/models/deepseek_v4_1/amd/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/deepseek_v4_1/amd/dspark.py new
+546
@@ -0,0 +1,546 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""DSpark draft model for DeepSeek-V4.1 (semi-autoregressive speculative decoding).
4 +
5 +See: qwen3_dspark.py for base architecture. This one is specialized to the
6 +DSV4.1 DSpark, which reuses the target model's architecture similarly to MTP.
7 +Its weights ship in the target checkpoint under the ``mtp.{0,1,2}.*`` prefix.
8 +
9 +To implement non-causal attention, we leverage the sparse attention implementation to
10 +include the future query tokens in the top-k indices for each query token.
11 +"""
12 +
13 +from collections.abc import Iterable
14 +
15 +import regex as re
16 +import torch
17 +import torch.nn as nn
18 +
19 +import vllm.envs as envs
20 +from vllm.config import VllmConfig, get_current_vllm_config
21 +from vllm.distributed import (
22 + get_tensor_model_parallel_rank,
23 + get_tensor_model_parallel_world_size,
24 +)
25 +from vllm.forward_context import get_forward_context, is_forward_context_available
26 +from vllm.logger import init_logger
27 +from vllm.model_executor.kernels.mhc.torch import mhc_post_torch
28 +from vllm.model_executor.layers.fused_moe import (
29 + fused_moe_make_expert_params_mapping,
30 +)
31 +from vllm.model_executor.layers.layernorm import RMSNorm
32 +from vllm.model_executor.layers.linear import ReplicatedLinear
33 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
34 +from vllm.model_executor.layers.vocab_parallel_embedding import (
35 + ParallelLMHead,
36 + VocabParallelEmbedding,
37 +)
38 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
39 +from vllm.model_executor.models.qwen3_dspark import (
40 + DSparkConfidenceHead,
41 + DSparkMarkovHead,
42 +)
43 +from vllm.model_executor.models.utils import maybe_prefix
44 +from vllm.models.common.ops.sequence_parallel import (
45 + sp_all_gather,
46 + sp_padding_mask,
47 + sp_shard,
48 +)
49 +
50 +from .model import (
51 + DeepseekV4DecoderLayer,
52 + DeepseekV4Model,
53 + _linear_scale_param_name,
54 + _use_sequence_parallel,
55 +)
56 +
57 +logger = init_logger(__name__)
58 +
59 +# MoE expert scale suffix differs by expert dtype (mirrors deepseek_v4 loaders):
60 +# fp4 experts register ``.weight_scale``; block-fp8 experts ``.weight_scale_inv``.
61 +_EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$")
62 +
63 +
64 +class DSparkDeepseekV4Model(nn.Module):
65 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
66 + super().__init__()
67 + assert vllm_config.speculative_config is not None
68 + config = vllm_config.speculative_config.draft_model_config.hf_config
69 + self.config = config
70 + self.hidden_size = config.hidden_size
71 + self.hc_mult = config.hc_mult
72 + self.hc_eps = config.hc_eps
73 + self.rms_norm_eps = config.rms_norm_eps
74 + self.num_hidden_layers = config.num_hidden_layers
75 + self.target_layer_ids = tuple(config.dspark_target_layer_ids)
76 + self.use_sequence_parallel = _use_sequence_parallel(vllm_config)
77 +
78 + self.num_dspark_layers = (
79 + getattr(config, "n_mtp_layers", None)
80 + or getattr(config, "num_nextn_predict_layers", None)
81 + or 3
82 + )
83 +
84 + # Shared with the target (aliased by the speculator's loading utility).
85 + self.embed_tokens = VocabParallelEmbedding(
86 + config.vocab_size,
87 + config.hidden_size,
88 + prefix=maybe_prefix(prefix, "embed_tokens"),
89 + )
90 +
91 + self.main_proj = ReplicatedLinear(
92 + config.hidden_size * len(self.target_layer_ids),
93 + config.hidden_size,
94 + bias=False,
95 + return_bias=False,
96 + quant_config=vllm_config.quant_config,
97 + prefix=maybe_prefix(prefix, "main_proj"),
98 + )
99 + self.main_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
100 +
101 + self.topk_indices_buffer = torch.empty(
102 + vllm_config.scheduler_config.max_num_batched_tokens,
103 + config.index_topk,
104 + dtype=torch.int32,
105 + )
106 +
107 + current_vllm_config = get_current_vllm_config()
108 + self.layers = nn.ModuleList(
109 + [
110 + DeepseekV4DecoderLayer(
111 + current_vllm_config,
112 + prefix=maybe_prefix(prefix, f"layers.{self.num_hidden_layers + i}"),
113 + topk_indices_buffer=self.topk_indices_buffer,
114 + )
115 + for i in range(self.num_dspark_layers)
116 + ]
117 + )
118 +
119 + # Heads: final norm, and the Markov + confidence heads.
120 + # Loaded from the "final" MTP layer weights (mtp.*) in the target
121 + # checkpoint. v4.1 has no learned hc_head: the hc copies are
122 + # collapsed with a pre-mix derived from the last layer's hc_ffn
123 + # projection at the end of forward().
124 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
125 + draft_vocab_size = (
126 + getattr(config, "draft_vocab_size", None) or config.vocab_size
127 + )
128 + self.markov_head = DSparkMarkovHead(
129 + config.vocab_size,
130 + draft_vocab_size,
131 + config.dspark_markov_rank,
132 + prefix=maybe_prefix(prefix, "markov_head"),
133 + )
134 + self.confidence_head: DSparkConfidenceHead | None = None
135 + if getattr(config, "enable_confidence_head", True):
136 + self.confidence_head = DSparkConfidenceHead(
137 + config.hidden_size + config.dspark_markov_rank,
138 + prefix=maybe_prefix(prefix, "confidence_head"),
139 + )
140 +
141 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
142 + return self.embed_tokens(input_ids)
143 +
144 + def combine_hidden_states(self, aux_hidden_states: torch.Tensor) -> torch.Tensor:
145 + """main_x = main_norm(main_proj(concat of target aux hidden states)).
146 +
147 + ``aux_hidden_states`` is [T, hidden_size * len(target_layer_ids)].
148 + """
149 + return self.main_norm(self.main_proj(aux_hidden_states))
150 +
151 + @torch.inference_mode()
152 + def precompute_and_store_context_kv(
153 + self,
154 + main_x: torch.Tensor,
155 + context_positions: torch.Tensor,
156 + context_slot_mappings: list[torch.Tensor | None] | None = None,
157 + ) -> None:
158 + """Insert the sliding-window context KV for every draft layer.
159 +
160 + Mirrors the reference DSparkAttention: each layer derives its context KV
161 + from the SAME projected target hidden ``main_x``, via that layer's own
162 + ``wkv`` + ``kv_norm`` + RoPE + quant, then writes it at the
163 + layer's context slots.
164 +
165 + ``context_slot_mappings`` is a per-layer list (each entry is the context
166 + slot mapping for that layer's kv-cache group, since the hybrid manager may
167 + place draft layers in different groups). ``None`` (or a ``None`` entry)
168 + runs the projection to reserve workspace but writes nothing (profiling).
169 + """
170 + for i, layer in enumerate(self.layers):
171 + slot_mapping = (
172 + None if context_slot_mappings is None else context_slot_mappings[i]
173 + )
174 + attn = layer.attn
175 + # Optimized DSV4 MLA path: wkv part of the fused wq_a|wkv projection
176 + # (q_lora part discarded), then RoPE/quant/insert via the fused op.
177 + qr_kv, _ = attn.fused_wqa_wkv(main_x)
178 + kv = qr_kv[..., attn.q_lora_rank :]
179 + kv = attn.kv_norm(kv)
180 + if slot_mapping is None:
181 + continue
182 + _insert_context_kv(attn, kv, context_positions, slot_mapping)
183 +
184 + def forward(
185 + self,
186 + input_ids: torch.Tensor,
187 + positions: torch.Tensor,
188 + inputs_embeds: torch.Tensor | None = None,
189 + ) -> torch.Tensor:
190 + if inputs_embeds is None:
191 + inputs_embeds = self.embed_input_ids(input_ids)
192 + full_num_tokens = positions.shape[0]
193 + if self.use_sequence_parallel:
194 + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available():
195 + forward_context = get_forward_context()
196 + forward_context.is_padding = sp_padding_mask(
197 + forward_context.is_padding, inputs_embeds
198 + )
199 + inputs_embeds = sp_shard(inputs_embeds)
200 + input_ids = sp_shard(input_ids)
201 + # Expand to hc_mult copies for hyper-connections ([T, H] -> [T, hc, H]).
202 + hidden_states = inputs_embeds.unsqueeze(-2).repeat(1, self.hc_mult, 1)
203 +
204 + residual = post_mix = res_mix = pre_mix = None
205 + for layer in self.layers:
206 + hidden_states, residual, post_mix, res_mix, pre_mix = layer(
207 + hidden_states,
208 + positions,
209 + input_ids,
210 + pre_mix,
211 + post_mix,
212 + res_mix,
213 + residual,
214 + )
215 + hidden_states = mhc_post_torch(hidden_states, residual, post_mix, res_mix)
216 + if self.use_sequence_parallel:
217 + hidden_states = sp_all_gather(hidden_states)[:full_num_tokens]
218 + pre_mix = sp_all_gather(pre_mix)[:full_num_tokens]
219 + # Collapse the hc copies with the pre-mix from the last layer's FFN
220 + # mixes — the mix the reference forward_head applies (hc_pre with the
221 + # last block's ffn pre-mix). Return the PRE-norm head hidden;
222 + # compute_logits applies self.norm.
223 + hidden_states = DeepseekV4DecoderLayer._hc_collapse(hidden_states, pre_mix)
224 + return hidden_states
225 +
226 +
227 +def _insert_context_kv(
228 + attn: nn.Module,
229 + kv: torch.Tensor,
230 + positions: torch.Tensor,
231 + slot_mapping: torch.Tensor,
232 +) -> None:
233 + """RoPE + quant + paged-cache insert of (already kv_norm'd) context KV.
234 +
235 + Reuses the DSV4 fused insert ops (which also process a query; we pass a dummy
236 + query and discard it, since context tokens have no query). Mirrors
237 + ``DeepseekV4Attention._fused_qnorm_rope_kv_insert``.
238 + """
239 + swa_cache = attn.swa_cache_layer.kv_cache
240 + block_size = attn.swa_cache_layer.block_size
241 + cos_sin_cache = attn.rotary_emb.cos_sin_cache
242 + cache_dtype = swa_cache.dtype
243 + n_ctx = kv.shape[0]
244 + dummy_q = torch.zeros(
245 + (n_ctx, attn.n_local_heads, attn.head_dim),
246 + dtype=kv.dtype,
247 + device=kv.device,
248 + )
249 + if cache_dtype == torch.uint8:
250 + # fp8_ds_mla UE8M0 paged layout
251 + swa_2d = swa_cache.view(swa_cache.shape[0], -1)
252 + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
253 + dummy_q,
254 + kv,
255 + swa_2d,
256 + slot_mapping,
257 + positions,
258 + cos_sin_cache,
259 + attn.padded_heads,
260 + attn.eps,
261 + block_size,
262 + )
263 + elif cache_dtype == torch.bfloat16:
264 + swa_3d = swa_cache.view(-1, block_size, attn.head_dim)
265 + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
266 + dummy_q,
267 + kv,
268 + swa_3d,
269 + slot_mapping,
270 + positions,
271 + cos_sin_cache,
272 + attn.eps,
273 + block_size,
274 + )
275 + else: # per-tensor fp8 (torch.float8_e4m3fn)
276 + # TODO(ben): double-check if this is being dispatched correctly for FI backend
277 + swa_3d = swa_cache.view(-1, block_size, attn.head_dim)
278 + dummy_q_fp8 = torch.zeros_like(dummy_q, dtype=torch.float8_e4m3fn)
279 + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
280 + dummy_q,
281 + kv,
282 + dummy_q_fp8,
283 + swa_3d,
284 + slot_mapping,
285 + positions,
286 + cos_sin_cache,
287 + attn._flashinfer_fp8_kv_scale,
288 + attn._flashinfer_fp8_q_scale_inv,
289 + attn.eps,
290 + block_size,
291 + )
292 +
293 +
294 +class DSparkDeepseekV4ForCausalLM(nn.Module):
295 + # Draft weights ship in the target checkpoint (mtp.*) without embed/head, so
296 + # load_dspark_model always aliases the target's.
297 + has_own_embed_tokens = False
298 + has_own_lm_head = False
299 + # Full-vocab draft: draft ids are target ids, no remapping needed.
300 + draft_id_to_target_id = None
301 +
302 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
303 + super().__init__()
304 + assert vllm_config.speculative_config is not None
305 + self.draft_model_config = vllm_config.speculative_config.draft_model_config
306 + self.config = self.draft_model_config.hf_config
307 + self.quant_config = vllm_config.quant_config
308 + self.linear_scale_name = _linear_scale_param_name(
309 + vllm_config, getattr(self.config, "expert_dtype", "fp4")
310 + )
311 + self.pad_shared_expert = getattr(
312 + self.quant_config, "weight_block_size", None
313 + ) is not None and not _use_sequence_parallel(vllm_config)
314 + self.model = DSparkDeepseekV4Model(
315 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
316 + )
317 + # Shared with the target (aliased by the speculator's load utility).
318 + self.lm_head = ParallelLMHead(
319 + self.config.vocab_size,
320 + self.config.hidden_size,
321 + prefix=maybe_prefix(prefix, "lm_head"),
322 + )
323 + self.logits_processor = LogitsProcessor(self.config.vocab_size)
324 +
325 + # --- Hooks used by the speculator -------------------------------------
326 +
327 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
328 + return self.model.embed_input_ids(input_ids)
329 +
330 + def combine_hidden_states(self, aux_hidden_states: torch.Tensor) -> torch.Tensor:
331 + return self.model.combine_hidden_states(aux_hidden_states)
332 +
333 + def get_draft_kv_cache_layer_names(self) -> list[str]:
334 + # DSV4 MLA path: each draft layer's sliding-window cache is a separate
335 + # layer, named by its prefix.
336 + return [layer.attn.swa_cache_layer.prefix for layer in self.model.layers]
337 +
338 + def precompute_and_store_context_kv(
339 + self,
340 + context_states: torch.Tensor,
341 + context_positions: torch.Tensor,
342 + context_slot_mappings: list[torch.Tensor | None] | None = None,
343 + ) -> None:
344 + self.model.precompute_and_store_context_kv(
345 + context_states, context_positions, context_slot_mappings
346 + )
347 +
348 + def forward(
349 + self,
350 + input_ids: torch.Tensor,
351 + positions: torch.Tensor,
352 + inputs_embeds: torch.Tensor | None = None,
353 + ) -> torch.Tensor:
354 + # Returns the pre-norm collapsed head hidden ([T, hidden_size]).
355 + return self.model(input_ids, positions, inputs_embeds)
356 +
357 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
358 + """Base logits U_k = lm_head(norm(head_hidden))."""
359 + return self.logits_processor(self.lm_head, self.model.norm(hidden_states))
360 +
361 + def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
362 + # Full-vocab draft: base logits, no d2t scatter.
363 + return self.compute_logits(hidden_states)
364 +
365 + def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor:
366 + return draft_ids # full-vocab: draft ids are target ids
367 +
368 + def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor:
369 + return self.model.markov_head.embed(token_ids)
370 +
371 + def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor:
372 + return self.model.markov_head.bias(markov_embed, self.logits_processor)
373 +
374 + def compute_confidence(
375 + self, head_hidden: torch.Tensor, markov_embed: torch.Tensor
376 + ) -> torch.Tensor:
377 + """Per-position acceptance probability for each drafted token."""
378 + assert self.model.confidence_head is not None
379 + return torch.sigmoid(self.model.confidence_head(head_hidden, markov_embed))
380 +
381 + # --- Weight loading ----------------------------------------------------
382 +
383 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
384 + """Load the ``mtp.{0,1,2}.*`` draft weights from the target checkpoint.
385 +
386 + Non-mtp weights (embed/head/main layers) belong to the target model and
387 + are skipped here. ``embed_tokens``/``lm_head`` are aliased from the target.
388 + """
389 + # Draft MoE layers use the dspark_* expert counts, not the
390 + # backbone's (see DeepseekV4MoE and the reference
391 + # ModelArgs.get_moe_config).
392 + n_draft_experts = (
393 + getattr(self.config, "dspark_n_routed_experts", 0)
394 + or self.config.n_routed_experts
395 + )
396 + expert_mapping = fused_moe_make_expert_params_mapping(
397 + self,
398 + ckpt_gate_proj_name="w1",
399 + ckpt_down_proj_name="w2",
400 + ckpt_up_proj_name="w3",
401 + num_experts=n_draft_experts,
402 + )
403 + expert_scale_suffix = (
404 + ".weight_scale"
405 + if getattr(self.config, "expert_dtype", "fp4") == "fp4"
406 + else ".weight_scale_inv"
407 + )
408 +
409 + # (param_name, ckpt_shard_name, shard_id) for non-expert stacked params.
410 + stacked_params_mapping = [
411 + ("gate_up_proj", "w1", 0),
412 + ("gate_up_proj", "w3", 1),
413 + ("attn.fused_wqa_wkv", "attn.wq_a", 0),
414 + ("attn.fused_wqa_wkv", "attn.wkv", 1),
415 + ]
416 +
417 + params_dict = dict(self.named_parameters())
418 + loaded_params: set[str] = set()
419 + loaded_confidence_head = False
420 +
421 + tp_size = get_tensor_model_parallel_world_size()
422 + tp_rank = get_tensor_model_parallel_rank()
423 + n_local_head = self.config.num_attention_heads // tp_size
424 + head_start = n_local_head * tp_rank
425 + head_end = n_local_head * (tp_rank + 1)
426 +
427 + for name, loaded_weight in weights:
428 + mapped = self._remap_dspark_name(name)
429 + if mapped is None:
430 + continue
431 + name = mapped
432 + if "confidence_head." in name:
433 + loaded_confidence_head = True
434 +
435 + # ``.scale`` -> per-method scale suffix.
436 + if name.endswith(".scale"):
437 + suffix = (
438 + expert_scale_suffix
439 + if _EXPERT_SCALE_RE.search(name)
440 + else f".{self.linear_scale_name}"
441 + )
442 + name = name.removesuffix(".scale") + suffix
443 + if ".shared_experts.w2" in name:
444 + name = name.replace(".shared_experts.w2", ".shared_experts.down_proj")
445 + if self.pad_shared_expert and ".shared_experts." in name:
446 + loaded_weight = DeepseekV4Model._pad_shared_expert_weight(
447 + self.quant_config, name, loaded_weight
448 + )
449 +
450 + # E8M0 expert scales: keep raw exponent bytes.
451 + if ".experts." in name:
452 + if (
453 + "weight_scale" in name
454 + and loaded_weight.dtype == torch.float8_e8m0fnu
455 + ):
456 + loaded_weight = loaded_weight.view(torch.uint8)
457 + for param_name, weight_name, expert_id, shard_id in expert_mapping:
458 + if weight_name not in name:
459 + continue
460 + name_mapped = name.replace(weight_name, param_name)
461 + param = params_dict[name_mapped]
462 + success = param.weight_loader(
463 + param,
464 + loaded_weight,
465 + name_mapped,
466 + shard_id=shard_id,
467 + expert_id=expert_id,
468 + return_success=True,
469 + )
470 + if success:
471 + loaded_params.add(name_mapped)
472 + break
473 + continue
474 +
475 + # Stacked rules only apply to decoder-layer weights. Head-stack params
476 + # (main_proj/norm/markov_head/confidence_head) load directly —
477 + # otherwise e.g. "markov_w1" would collide with the "w1" shard rule.
478 + is_layer_param = name.startswith("model.layers.")
479 + for param_name, weight_name, stacked_shard_id in stacked_params_mapping:
480 + if not is_layer_param or weight_name not in name:
481 + continue
482 + name = name.replace(weight_name, param_name)
483 + param = params_dict[name]
484 + param.weight_loader(param, loaded_weight, stacked_shard_id)
485 + loaded_params.add(name)
486 + break
487 + else:
488 + if "attn_sink" in name:
489 + narrow = loaded_weight[head_start:head_end]
490 + params_dict[name][: narrow.shape[0]].copy_(narrow)
491 + loaded_params.add(name)
492 + continue
493 + if name.endswith(".ffn.gate.bias"):
494 + name = name.replace(
495 + ".ffn.gate.bias", ".ffn.gate.e_score_correction_bias"
496 + )
497 + param = params_dict[name]
498 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
499 + weight_loader(param, loaded_weight)
500 + loaded_params.add(name)
501 +
502 + if self.model.confidence_head is not None and not loaded_confidence_head:
503 + self.model.confidence_head = None
504 + self.process_weights_after_loading()
505 + logger.info_once("DSpark draft model loaded: %d params", len(loaded_params))
506 + return loaded_params
507 +
508 + def process_weights_after_loading(self) -> None:
509 + # ROCm linears and fused-MoE runners finalize their own quantized
510 + # parameters. NVIDIA-only MegaMoE/WO-A requantization is not needed.
511 + return
512 +
513 + def _remap_dspark_name(self, name: str) -> str | None:
514 + """Map a checkpoint ``mtp.{i}.*`` name to this model's parameter path.
515 +
516 + Returns None for non-mtp weights (owned by the target model).
517 + """
518 + m = re.match(r"mtp\.(\d+)\.(.*)", name)
519 + if m is None:
520 + return None
521 + stage = int(m.group(1))
522 + rest = m.group(2)
523 + if rest.startswith("confidence_head.") and self.model.confidence_head is None:
524 + return None
525 + # The checkpoint calls the Markov head's factors ``embed``/``head``;
526 + # DSparkMarkovHead registers them as ``markov_w1``/``markov_w2``.
527 + if rest.startswith("markov_head.embed."):
528 + return "model.markov_head.markov_w1." + rest.removeprefix(
529 + "markov_head.embed."
530 + )
531 + if rest.startswith("markov_head.head."):
532 + return "model.markov_head.markov_w2." + rest.removeprefix(
533 + "markov_head.head."
534 + )
535 + # Head-stack params live at model level (mtp.last), context combiner at
536 + # model level (mtp.0); everything else is a per-layer decoder block.
537 + head_prefixes = (
538 + "norm.",
539 + "markov_head.",
540 + "confidence_head.",
541 + )
542 + if rest.startswith(("main_proj.", "main_norm.")) or rest.startswith(
543 + head_prefixes
544 + ):
545 + return f"model.{rest}"
546 + return f"model.layers.{stage}.{rest}"
vllm/models/deepseek_v4_1/amd/model.py new
+1086
@@ -0,0 +1,1086 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import typing
4 +from collections.abc import Callable, Iterable
5 +from itertools import islice
6 +
7 +import regex as re
8 +import torch
9 +import torch.nn as nn
10 +
11 +import vllm.envs as envs
12 +from vllm.config import VllmConfig
13 +from vllm.distributed import (
14 + get_pp_group,
15 + get_tensor_model_parallel_rank,
16 + get_tensor_model_parallel_world_size,
17 +)
18 +from vllm.forward_context import get_forward_context, is_forward_context_available
19 +from vllm.logger import init_logger
20 +from vllm.model_executor.kernels.mhc.torch import (
21 + mhc_post_torch,
22 + mhc_pre_delayed_torch,
23 +)
24 +from vllm.model_executor.layers.fused_moe import (
25 + fused_moe_make_expert_params_mapping,
26 +)
27 +from vllm.model_executor.layers.layernorm import RMSNorm
28 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
29 +from vllm.model_executor.layers.quantization import QuantizationConfig
30 +from vllm.model_executor.layers.vocab_parallel_embedding import (
31 + ParallelLMHead,
32 + VocabParallelEmbedding,
33 +)
34 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
35 +from vllm.model_executor.models.interfaces import (
36 + EagleModelMixin,
37 + MixtureOfExperts,
38 + SupportsEagle3,
39 + SupportsLoRA,
40 + SupportsPP,
41 +)
42 +from vllm.model_executor.models.utils import (
43 + AutoWeightsLoader,
44 + PPMissingLayer,
45 + WeightsMapper,
46 + extract_layer_index,
47 + is_pp_missing_parameter,
48 + make_layers,
49 + maybe_prefix,
50 +)
51 +from vllm.models.common.ops.sequence_parallel import (
52 + sp_all_gather,
53 + sp_padding_mask,
54 + sp_reduce_scatter,
55 + sp_shard,
56 +)
57 +from vllm.models.deepseek_v4.amd.model import (
58 + DeepseekV4MoE as DeepseekV4MoEBase,
59 +)
60 +from vllm.models.deepseek_v4_1.amd.rocm import DeepseekV41ROCMAiterMLAAttention
61 +from vllm.models.deepseek_v4_1.attention import DeepseekV4Attention
62 +from vllm.sequence import IntermediateTensors
63 +from vllm.utils.math_utils import cdiv
64 +from vllm.v1.attention.backends.registry import AttentionBackendEnum
65 +from vllm.v1.worker.ubatching import dbo_current_ubatch_id
66 +
67 +from ..common.engram import Engram, EngramLayout, NgramHashState
68 +from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask
69 +
70 +if typing.TYPE_CHECKING:
71 + from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
72 +
73 +logger = init_logger(__name__)
74 +
75 +
76 +class DeepseekV4MoE(DeepseekV4MoEBase):
77 + def __init__(
78 + self,
79 + vllm_config: VllmConfig,
80 + prefix: str = "",
81 + use_sequence_parallel: bool = False,
82 + ):
83 + config = vllm_config.model_config.hf_config
84 + n_routed_experts = config.n_routed_experts
85 + n_activated_experts = config.num_experts_per_tok
86 + if extract_layer_index(prefix) >= config.num_hidden_layers:
87 + n_routed_experts = (
88 + getattr(config, "dspark_n_routed_experts", 0) or n_routed_experts
89 + )
90 + n_activated_experts = (
91 + getattr(config, "dspark_num_experts_per_tok", 0) or n_activated_experts
92 + )
93 + # The reused ROCm V4 MoE reads expert counts from the HF config. Draft
94 + # layers use V4.1's smaller DSpark expert set, so temporarily expose
95 + # the per-layer counts while constructing the module, then restore the
96 + # shared config for the remaining backbone layers.
97 + old_routed = config.n_routed_experts
98 + old_activated = config.num_experts_per_tok
99 + old_hash_layers = getattr(config, "num_hash_layers", None)
100 + config.n_routed_experts = n_routed_experts
101 + config.num_experts_per_tok = n_activated_experts
102 + config.num_hash_layers = 0
103 + try:
104 + super().__init__(vllm_config, prefix=prefix)
105 + finally:
106 + config.n_routed_experts = old_routed
107 + config.num_experts_per_tok = old_activated
108 + if old_hash_layers is None:
109 + delattr(config, "num_hash_layers")
110 + else:
111 + config.num_hash_layers = old_hash_layers
112 +
113 + # The V4.1 checkpoint contains a second per-expert routing bias for
114 + # image sentinel tokens. The reused ROCm V4 MoE predates this field,
115 + # so register it here and attach the same Parameter to the router,
116 + # matching the CUDA V4.1 implementation. Text tokens continue to use
117 + # e_score_correction_bias.
118 + self.gate.bias_vl = None
119 + self.image_sentinel_lo = 0
120 + if getattr(config, "vision_n_layers", 0) > 0:
121 + self.gate.bias_vl = nn.Parameter(
122 + torch.empty(self.n_routed_experts, dtype=torch.float32),
123 + requires_grad=False,
124 + )
125 + self.image_sentinel_lo = IMAGE_SENTINEL_BASE_ID
126 + self.experts.router.bias_vl = self.gate.bias_vl
127 + self.experts.router.image_sentinel_lo = self.image_sentinel_lo
128 + self.use_mega_moe = False
129 + # Expose the MixtureOfExperts protocol fields expected by the V4.1
130 + # outer model. The current ROCm runner has no redundant EPLB copies,
131 + # so logical and physical expert counts are identical.
132 + self.n_logical_experts = self.n_routed_experts
133 + self.n_physical_experts = self.n_routed_experts
134 + self.n_local_physical_experts = self.n_local_experts
135 + self.n_redundant_experts = 0
136 +
137 +
138 +def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]:
139 + """Pick the ROCm sparse-MLA implementation for DeepSeek V4.1."""
140 + backend = vllm_config.attention_config.backend
141 + if backend is not None and (
142 + backend is not AttentionBackendEnum.ROCM_FLASHMLA_SPARSE_DSV4
143 + ):
144 + raise ValueError(
145 + f"{backend.name} is not supported for DeepSeek V4.1 on ROCm; "
146 + "use ROCM_FLASHMLA_SPARSE_DSV4."
147 + )
148 + return DeepseekV41ROCMAiterMLAAttention
149 +
150 +
151 +def _use_sequence_parallel(vllm_config: VllmConfig) -> bool:
152 + # The ROCm fused-MoE path currently uses ordinary tensor/expert parallel.
153 + return False
154 +
155 +
156 +class DeepseekV4DecoderLayer(nn.Module):
157 + def __init__(
158 + self,
159 + vllm_config,
160 + prefix,
161 + topk_indices_buffer: torch.Tensor | None = None,
162 + aux_stream_list: list[torch.cuda.Stream] | None = None,
163 + candidate_block_buffer: torch.Tensor | None = None,
164 + engram_layout: EngramLayout | None = None,
165 + ):
166 + super().__init__()
167 + config = vllm_config.model_config.hf_config
168 + self.hidden_size = config.hidden_size
169 + self.use_sequence_parallel = _use_sequence_parallel(vllm_config)
170 +
171 + self.engram: Engram | None = None
172 + if engram_layout is not None:
173 + layer_id = extract_layer_index(prefix)
174 + if layer_id in engram_layout.layer_ids:
175 + self.engram = Engram(
176 + config,
177 + vllm_config.quant_config,
178 + engram_layout,
179 + engram_layout.layer_ids.index(layer_id),
180 + use_sequence_parallel=self.use_sequence_parallel,
181 + prefix=f"{prefix}.engram",
182 + )
183 +
184 + self.rms_norm_eps = config.rms_norm_eps
185 + self.attn = _select_dsv4_attn_cls(vllm_config)(
186 + vllm_config,
187 + prefix=f"{prefix}.attn",
188 + topk_indices_buffer=topk_indices_buffer,
189 + aux_stream_list=aux_stream_list,
190 + candidate_block_buffer=candidate_block_buffer,
191 + )
192 + if self.use_sequence_parallel:
193 + self.attn.wo_b.reduce_results = False
194 + self.ffn = DeepseekV4MoE(
195 + vllm_config,
196 + prefix=f"{prefix}.ffn",
197 + use_sequence_parallel=self.use_sequence_parallel,
198 + )
199 +
200 + self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps)
201 + self.ffn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps)
202 + self.hc_mult = config.hc_mult
203 + self.hc_sinkhorn_iters = config.hc_sinkhorn_iters
204 + self.hc_eps = config.hc_eps
205 + self.hc_post_alpha = 2.0
206 + mix_hc = (2 + self.hc_mult) * self.hc_mult
207 + hc_dim = self.hc_mult * self.hidden_size
208 + self.hc_attn_fn = nn.Parameter(
209 + torch.empty(
210 + (mix_hc, hc_dim),
211 + dtype=torch.float32,
212 + ),
213 + requires_grad=False,
214 + )
215 + self.hc_attn_fn_broadcast: torch.Tensor | None = None
216 + self.hc_ffn_fn = nn.Parameter(
217 + torch.empty(
218 + (mix_hc, hc_dim),
219 + dtype=torch.float32,
220 + ),
221 + requires_grad=False,
222 + )
223 + self.hc_attn_base = nn.Parameter(
224 + torch.empty(
225 + mix_hc,
226 + dtype=torch.float32,
227 + ),
228 + requires_grad=False,
229 + )
230 + self.hc_ffn_base = nn.Parameter(
231 + torch.empty(
232 + mix_hc,
233 + dtype=torch.float32,
234 + ),
235 + requires_grad=False,
236 + )
237 + self.hc_attn_scale = nn.Parameter(
238 + torch.empty(
239 + 3,
240 + dtype=torch.float32,
241 + ),
242 + requires_grad=False,
243 + )
244 + self.hc_ffn_scale = nn.Parameter(
245 + torch.empty(
246 + 3,
247 + dtype=torch.float32,
248 + ),
249 + requires_grad=False,
250 + )
251 +
252 + @staticmethod
253 + def _hc_collapse(x: torch.Tensor, pre_mix: torch.Tensor) -> torch.Tensor:
254 + """Reference hyper-connection collapse used by DSpark on ROCm."""
255 + return (pre_mix.unsqueeze(-1) * x.float()).sum(dim=1).to(x.dtype)
256 +
257 + def forward(
258 + self,
259 + x: torch.Tensor,
260 + positions: torch.Tensor,
261 + input_ids: torch.Tensor | None,
262 + pre_mix: torch.Tensor | None = None,
263 + post_mix: torch.Tensor | None = None,
264 + res_mix: torch.Tensor | None = None,
265 + residual: torch.Tensor | None = None,
266 + engram_hashes: torch.Tensor | None = None,
267 + engram_mask: torch.Tensor | None = None,
268 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
269 + # The reference collapses each sublayer's input with the *previous*
270 + # sublayer's pre-mix: attention uses the pre-mix carried in (identity
271 + # for the first layer), the FFN uses this layer's attention pre-mix.
272 + if residual is None:
273 + if x.dim() == 2:
274 + # First layer: the stream is the embedding broadcast to hc
275 + # copies and the identity pre-mix selects copy 0.
276 + assert self.hc_attn_fn_broadcast is not None
277 + residual = x.unsqueeze(1).expand(-1, self.hc_mult, -1).contiguous()
278 + post_mix, res_mix, x, attn_pre = mhc_pre_delayed_torch(
279 + residual,
280 + self.hc_attn_fn_broadcast,
281 + self.hc_attn_scale,
282 + self.hc_attn_base,
283 + self.rms_norm_eps,
284 + self.hc_eps,
285 + self.hc_eps,
286 + self.hc_post_alpha,
287 + self.hc_sinkhorn_iters,
288 + x=x,
289 + )
290 + else:
291 + residual = x
292 + post_mix, res_mix, x, attn_pre = mhc_pre_delayed_torch(
293 + residual,
294 + self.hc_attn_fn,
295 + self.hc_attn_scale,
296 + self.hc_attn_base,
297 + self.rms_norm_eps,
298 + self.hc_eps,
299 + self.hc_eps,
300 + self.hc_post_alpha,
301 + self.hc_sinkhorn_iters,
302 + pre_mix=pre_mix,
303 + )
304 + else:
305 + residual = mhc_post_torch(x, residual, post_mix, res_mix)
306 + if self.engram is not None and engram_hashes is not None:
307 + # Engram injection happens between the previous sublayer's
308 + # post and this block's pre, on the full hc stream, so the
309 + # mix coefficients see the injected stream.
310 + residual = self.engram(
311 + residual,
312 + engram_hashes[:, self.engram.layer_hash_index],
313 + engram_mask,
314 + )
315 + post_mix, res_mix, x, attn_pre = mhc_pre_delayed_torch(
316 + residual,
317 + self.hc_attn_fn,
318 + self.hc_attn_scale,
319 + self.hc_attn_base,
320 + self.rms_norm_eps,
321 + self.hc_eps,
322 + self.hc_eps,
323 + self.hc_post_alpha,
324 + self.hc_sinkhorn_iters,
325 + pre_mix=pre_mix,
326 + )
327 + x = self.attn_norm(x)
328 +
329 + if self.use_sequence_parallel:
330 + x = sp_all_gather(x)[: positions.shape[0]]
331 +
332 + x = self.attn(positions, x, None)
333 + if self.use_sequence_parallel:
334 + x = sp_reduce_scatter(x)
335 +
336 + residual = mhc_post_torch(x, residual, post_mix, res_mix)
337 + post_mix, res_mix, x, ffn_pre = mhc_pre_delayed_torch(
338 + residual,
339 + self.hc_ffn_fn,
340 + self.hc_ffn_scale,
341 + self.hc_ffn_base,
342 + self.rms_norm_eps,
343 + self.hc_eps,
344 + self.hc_eps,
345 + self.hc_post_alpha,
346 + self.hc_sinkhorn_iters,
347 + pre_mix=attn_pre,
348 + )
349 + x = self.ffn_norm(x)
350 + x = self.ffn(x, input_ids)
351 + return x, residual, post_mix, res_mix, ffn_pre
352 +
353 +
354 +class DeepseekV4Model(nn.Module, EagleModelMixin):
355 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
356 + super().__init__()
357 +
358 + config = vllm_config.model_config.hf_config
359 + quant_config = vllm_config.quant_config
360 + self.config = config
361 + self.quant_config = quant_config
362 + self.parallel_config = vllm_config.parallel_config
363 + self.use_mega_moe = False
364 + self.use_sequence_parallel = _use_sequence_parallel(vllm_config)
365 + if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel:
366 + raise NotImplementedError(
367 + "DeepSeek V4 MegaMoE currently requires expert parallel. "
368 + "Enable it with --enable-expert-parallel, or pick a different "
369 + "moe backend."
370 + )
371 + self.vocab_size = config.vocab_size
372 + self.hc_eps = config.hc_eps
373 + self.hc_mult = config.hc_mult
374 + self.hc_dim = self.hc_mult * config.hidden_size
375 + self.rms_norm_eps = config.rms_norm_eps
376 +
377 + # Three aux streams: one per non-default input GEMM in
378 + # DeepseekV4Attention._run_parallel_input_projections
379 + # (compressor kv_score, indexer.weights_proj). fused_wqa_wkv stays on
380 + # the default stream.
381 + aux_stream_list = [torch.cuda.Stream() for _ in range(3)]
382 +
383 + # Reserved topk indices buffer for all Indexer layers to reuse.
384 + self.topk_indices_buffer = torch.empty(
385 + vllm_config.scheduler_config.max_num_batched_tokens,
386 + config.index_topk,
387 + dtype=torch.int32,
388 + )
389 +
390 + # Two-level candidate filtering: the indexer at
391 + # candidate_source_layer_id publishes the top candidate blocks of
392 + # compressed positions here; later ratio-1 indexers (24/28/32/36)
393 + # mask their scores with it.
394 + candidate_source_layer = getattr(config, "candidate_source_layer_id", -1)
395 + candidate_topk_blocks = getattr(config, "candidate_topk_blocks", 0)
396 + if candidate_source_layer >= 0 and candidate_topk_blocks > 0:
397 + self.candidate_block_buffer = torch.empty(
398 + vllm_config.scheduler_config.max_num_batched_tokens,
399 + candidate_topk_blocks,
400 + dtype=torch.int32,
401 + )
402 + else:
403 + self.candidate_block_buffer = None
404 +
405 + if get_pp_group().is_first_rank:
406 + self.embed_tokens = VocabParallelEmbedding(
407 + config.vocab_size,
408 + config.hidden_size,
409 + quant_config=quant_config,
410 + prefix=f"{prefix}.embed_tokens",
411 + )
412 + else:
413 + self.embed_tokens = PPMissingLayer()
414 +
415 + self.engram_layout = EngramLayout.from_config(config)
416 +
417 + self.start_layer, self.end_layer, self.layers = make_layers(
418 + config.num_hidden_layers,
419 + lambda prefix: DeepseekV4DecoderLayer(
420 + vllm_config,
421 + prefix=prefix,
422 + topk_indices_buffer=self.topk_indices_buffer,
423 + aux_stream_list=aux_stream_list,
424 + candidate_block_buffer=self.candidate_block_buffer,
425 + engram_layout=self.engram_layout,
426 + ),
427 + prefix=f"{prefix}.layers",
428 + )
429 +
430 + # The n-gram hash needs a slot-keyed rolling store of compressed ids
431 + # (chunked prefill / decode lookback); key it off the first local
432 + # layer's sliding-window KV cache. Only PP ranks owning an engram
433 + # layer need it.
434 + self.engram_hash: NgramHashState | None = None
435 + self.engram_swa_prefix: str | None = None
436 + if self.engram_layout is not None:
437 + local_engram = any(
438 + isinstance(layer, DeepseekV4DecoderLayer) and layer.engram is not None
439 + for layer in islice(self.layers, self.start_layer, self.end_layer)
440 + )
441 + if local_engram:
442 + first_layer = next(
443 + iter(islice(self.layers, self.start_layer, self.end_layer))
444 + )
445 + swa_cache_module = first_layer.attn.swa_cache_layer
446 + self.engram_hash = NgramHashState(
447 + vllm_config, self.engram_layout, swa_cache_module
448 + )
449 + self.engram_swa_prefix = swa_cache_module.prefix
450 +
451 + if get_pp_group().is_last_rank:
452 + self.norm = RMSNorm(config.hidden_size, self.rms_norm_eps)
453 + else:
454 + self.norm = PPMissingLayer()
455 +
456 + spec_config = vllm_config.speculative_config
457 + needs_mtp_hidden_states = spec_config is not None and (
458 + spec_config.use_eagle() or spec_config.uses_draft_model()
459 + )
460 + if get_pp_group().is_last_rank and needs_mtp_hidden_states:
461 + self._mtp_hidden_buffer = torch.empty(
462 + vllm_config.scheduler_config.max_num_batched_tokens,
463 + self.hc_dim,
464 + dtype=vllm_config.model_config.dtype,
465 + )
466 + else:
467 + self._mtp_hidden_buffer = None
468 +
469 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
470 + return self.embed_tokens(input_ids)
471 +
472 + def make_empty_intermediate_tensors(
473 + self,
474 + batch_size: int,
475 + dtype: torch.dtype,
476 + device: torch.device,
477 + ) -> IntermediateTensors:
478 + # PP intermediate tensors carry the multi-stream hidden_states
479 + # of shape (num_tokens, hc_mult, hidden_size) — V4 expands the
480 + # token embedding to hc_mult streams before the first decoder
481 + # layer and keeps that shape until the final hc collapse — plus the
482 + # (num_tokens, hc_mult) pre-mix the next rank's first layer needs
483 + # for its attention collapse.
484 + return IntermediateTensors(
485 + {
486 + "hidden_states": torch.zeros(
487 + (batch_size, self.hc_mult, self.config.hidden_size),
488 + dtype=dtype,
489 + device=device,
490 + ),
491 + "pre_mix": torch.zeros(
492 + (batch_size, self.hc_mult),
493 + dtype=torch.float32,
494 + device=device,
495 + ),
496 + }
497 + )
498 +
499 + def forward(
500 + self,
501 + input_ids: torch.Tensor,
502 + positions: torch.Tensor,
503 + intermediate_tensors: IntermediateTensors | None,
504 + inputs_embeds: torch.Tensor | None = None,
505 + lookback_token_ids: torch.Tensor | None = None,
506 + ) -> torch.Tensor | IntermediateTensors:
507 + if get_pp_group().is_first_rank:
508 + if inputs_embeds is not None:
509 + hidden_states = inputs_embeds
510 + else:
511 + hidden_states = self.embed_input_ids(input_ids)
512 + else:
513 + assert intermediate_tensors is not None
514 + hidden_states = intermediate_tensors["hidden_states"]
515 +
516 + if self.use_mega_moe:
517 + input_ids = input_ids.to(torch.int64)
518 +
519 + # Engram n-gram hashes for the whole (flattened) batch, computed once
520 + # on the full token stream — before any sequence-parallel sharding —
521 + # and consumed by the engram layers (1 and 14) below. Skipped on
522 + # profile runs (KV cache unbound -> no slot space to key the rolling
523 + # compressed-id cache into).
524 + engram_hashes: torch.Tensor | None = None
525 + engram_mask: torch.Tensor | None = None
526 + if (
527 + self.engram_hash is not None
528 + and input_ids is not None
529 + and is_forward_context_available()
530 + ):
531 + attn_metadata = get_forward_context().attn_metadata
532 + if isinstance(attn_metadata, list):
533 + attn_metadata = attn_metadata[dbo_current_ubatch_id()]
534 + if isinstance(attn_metadata, dict) and self.engram_hash.ensure_cache():
535 + assert self.engram_swa_prefix is not None
536 + swa_metadata = typing.cast(
537 + "DeepseekSparseSWAMetadata", attn_metadata[self.engram_swa_prefix]
538 + )
539 + # Image-span tokens are dead: they break n-grams (hash op
540 + # takes True=dead) and their gate is zeroed (Engram.forward
541 + # takes True=keep).
542 + image_mask = image_sentinel_mask(input_ids)
543 + engram_mask = ~image_mask
544 + if lookback_token_ids is None:
545 + if not self.engram_hash.use_slot_cache:
546 + raise NotImplementedError(
547 + "engram needs `lookback_token_ids` from the model "
548 + "runner (the DBO/ubatch wrapper drops model kwargs)"
549 + )
550 + num_reqs = swa_metadata.num_decodes + swa_metadata.num_prefills
551 + lookback_token_ids = input_ids.new_full(
552 + (num_reqs, self.engram_hash.lookback_depth), -1
553 + )
554 + engram_hashes = self.engram_hash(
555 + input_ids,
556 + positions,
557 + swa_metadata.query_start_loc,
558 + image_mask,
559 + lookback_token_ids,
560 + image_sentinel_mask(lookback_token_ids),
561 + swa_metadata.slot_mapping,
562 + swa_metadata.block_table,
563 + )
564 + # Gather all Engram rows before entering the decoder layers.
565 + for layer in islice(self.layers, self.start_layer, self.end_layer):
566 + engram = getattr(layer, "engram", None)
567 + if engram is not None:
568 + engram.prepare_embeddings(
569 + engram_hashes[:, engram.layer_hash_index]
570 + )
571 +
572 + full_num_tokens = positions.shape[0]
573 + if self.use_sequence_parallel:
574 + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available():
575 + forward_context = get_forward_context()
576 + forward_context.is_padding = sp_padding_mask(
577 + forward_context.is_padding, hidden_states
578 + )
579 + hidden_states = sp_shard(hidden_states)
580 + input_ids = sp_shard(input_ids)
581 +
582 + residual, post_mix, res_mix = None, None, None
583 + pre_mix: torch.Tensor | None = None
584 + if not get_pp_group().is_first_rank:
585 + assert intermediate_tensors is not None
586 + pre_mix = intermediate_tensors["pre_mix"]
587 + aux_hidden_states: list[torch.Tensor] = []
588 + final_aux_recon: torch.Tensor | None = None # avoid duplicate mhc_post call
589 + for idx, layer in enumerate(
590 + islice(self.layers, self.start_layer, self.end_layer),
591 + start=self.start_layer,
592 + ):
593 + hidden_states, residual, post_mix, res_mix, pre_mix = layer(
594 + hidden_states,
595 + positions,
596 + input_ids,
597 + pre_mix,
598 + post_mix,
599 + res_mix,
600 + residual,
601 + engram_hashes,
602 + engram_mask,
603 + )
604 + if idx + 1 in self.aux_hidden_state_layers:
605 + # Reconstruct the aux hidden state for draft models
606 + aux_recon = mhc_post_torch(hidden_states, residual, post_mix, res_mix)
607 + aux_hidden_state = aux_recon.mean(dim=1)
608 + if self.use_sequence_parallel:
609 + aux_hidden_state = sp_all_gather(aux_hidden_state)[:full_num_tokens]
610 + aux_hidden_states.append(aux_hidden_state)
611 + final_aux_recon = aux_recon
612 + if layer is not None:
613 + # Reuse if the last layer was captured as an aux hidden state
614 + if self.end_layer in self.aux_hidden_state_layers:
615 + hidden_states = final_aux_recon
616 + else:
617 + hidden_states = mhc_post_torch(
618 + hidden_states, residual, post_mix, res_mix
619 + )
620 +
621 + if not get_pp_group().is_last_rank:
622 + return IntermediateTensors(
623 + {"hidden_states": hidden_states, "pre_mix": pre_mix}
624 + )
625 +
626 + if self.use_sequence_parallel:
627 + hidden_states = sp_all_gather(hidden_states)[:full_num_tokens]
628 + pre_mix = sp_all_gather(pre_mix)[:full_num_tokens]
629 +
630 + if self._mtp_hidden_buffer is not None:
631 + num_tokens = hidden_states.shape[0]
632 + self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1))
633 +
634 + # Collapse the hc copies with the pre-mix from the last layer's FFN
635 + # mixes — the mix the reference applies via
636 + # ``last_layer.hc_pre(h, pre_mix)`` (v4.1 has no learned hc_head).
637 + assert pre_mix is not None
638 + hidden_states = (
639 + (pre_mix.unsqueeze(-1) * hidden_states.float())
640 + .sum(dim=1)
641 + .to(hidden_states.dtype)
642 + )
643 + hidden_states = self.norm(hidden_states)
644 + if len(aux_hidden_states) > 0:
645 + return hidden_states, aux_hidden_states
646 + return hidden_states
647 +
648 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
649 + stacked_params_mapping = [
650 + # (param_name, shard_name, shard_id)
651 + ("gate_up_proj", "w1", 0),
652 + ("gate_up_proj", "w3", 1),
653 + ("attn.fused_wqa_wkv", "attn.wq_a", 0),
654 + ("attn.fused_wqa_wkv", "attn.wkv", 1),
655 + ("compressor.fused_wkv_wgate", "compressor.wkv", 0),
656 + ("compressor.fused_wkv_wgate", "compressor.wgate", 1),
657 + ]
658 + params_dict = dict(self.named_parameters())
659 + loaded_params: set[str] = set()
660 +
661 + def _resolve_param_name(name: str) -> str:
662 + """Resolve scale naming across ROCm FP8 and MXFP8 methods."""
663 + if name.endswith("_inv"):
664 + direct_name = name.removesuffix("_inv")
665 + if direct_name in params_dict:
666 + return direct_name
667 + inv_name = f"{name}_inv"
668 + if name not in params_dict and inv_name in params_dict:
669 + return inv_name
670 + return name
671 +
672 + # TP for attention
673 + tp_size = get_tensor_model_parallel_world_size()
674 + tp_rank = get_tensor_model_parallel_rank()
675 + n_head = self.config.num_attention_heads
676 + n_local_head = n_head // tp_size
677 + head_rank_start = n_local_head * tp_rank
678 + head_rank_end = n_local_head * (tp_rank + 1)
679 +
680 + # Pre-compute expert mapping ONCE.
681 + expert_mapping = self.get_expert_mapping()
682 +
683 + # Block-FP8 shared experts: pad the intermediate up to the TP-uniform
684 + # block count so the standard loaders below slice it evenly (trailing
685 + # ranks land on the zero pad). SP / unquantized ones need no padding.
686 + pad_shared_expert = (
687 + getattr(self.quant_config, "weight_block_size", None) is not None
688 + and not self.use_sequence_parallel
689 + )
690 +
691 + for name, loaded_weight in weights:
692 + if name.startswith(("vision.", "aligner.", "image_")):
693 + # Vision weights are loaded by the outer multimodal wrapper.
694 + logger.warning_once("Skipping non-text weight: %s", name)
695 + continue
696 + if pad_shared_expert and ".shared_experts." in name:
697 + loaded_weight = self._pad_shared_expert_weight(
698 + self.quant_config, name, loaded_weight
699 + )
700 + for param_name, weight_name, shard_id in stacked_params_mapping:
701 + # Skip non-stacked layers and experts (experts handled below).
702 + if ".experts." in name:
703 + continue
704 + if weight_name not in name:
705 + continue
706 + name = name.replace(weight_name, param_name)
707 +
708 + if is_pp_missing_parameter(name, self):
709 + break
710 + if name not in params_dict:
711 + head, _, leaf = name.rpartition(".")
712 + suffixed = f"{head}.base_layer.{leaf}"
713 + if suffixed in params_dict:
714 + name = suffixed
715 + name = _resolve_param_name(name)
716 + param = params_dict[name]
717 + weight_loader = param.weight_loader
718 + weight_loader(param, loaded_weight, shard_id)
719 + loaded_params.add(name)
720 + break
721 + else:
722 + if ".experts." in name:
723 + # E8M0 scales are stored as float8_e8m0fnu in
724 + # checkpoints but the MoE param is uint8. copy_()
725 + # would do a numeric conversion (e.g. 2^-7 → 0),
726 + # destroying the raw exponent bytes.
727 + if (
728 + "weight_scale" in name
729 + and loaded_weight.dtype == torch.float8_e8m0fnu
730 + ):
731 + loaded_weight = loaded_weight.view(torch.uint8)
732 + for mapping in expert_mapping:
733 + param_name, weight_name, expert_id, expert_shard_id = mapping
734 + if weight_name not in name:
735 + continue
736 + name_mapped = name.replace(weight_name, param_name)
737 + if is_pp_missing_parameter(name_mapped, self):
738 + continue
739 + name_mapped = _resolve_param_name(name_mapped)
740 + param = params_dict[name_mapped]
741 + # We should ask the weight loader to return success or not
742 + # here since otherwise we may skip experts with other
743 + # available replicas.
744 + weight_loader = typing.cast(
745 + Callable[..., bool], param.weight_loader
746 + )
747 + success = weight_loader(
748 + param,
749 + loaded_weight,
750 + name_mapped,
751 + shard_id=expert_shard_id,
752 + expert_id=expert_id,
753 + return_success=True,
754 + )
755 + if success:
756 + name = name_mapped
757 + break
758 + loaded_params.add(name_mapped)
759 + continue
760 + elif "attn_sink" in name:
761 + if is_pp_missing_parameter(name, self):
762 + continue
763 + narrow_weight = loaded_weight[head_rank_start:head_rank_end]
764 + n = narrow_weight.shape[0]
765 + params_dict[name][:n].copy_(narrow_weight)
766 + loaded_params.add(name)
767 + continue
768 + else:
769 + if is_pp_missing_parameter(name, self):
770 + continue
771 + # Non-LoRA params on a LoRA-wrapped module live at
772 + # ``<head>.base_layer.<leaf>``; the checkpoint is plain.
773 + if name not in params_dict:
774 + head, _, leaf = name.rpartition(".")
775 + suffixed = f"{head}.base_layer.{leaf}"
776 + if suffixed in params_dict:
777 + name = suffixed
778 + name = _resolve_param_name(name)
779 + param = params_dict[name]
780 + weight_loader = getattr(
781 + param, "weight_loader", default_weight_loader
782 + )
783 + weight_loader(param, loaded_weight)
784 + loaded_params.add(name)
785 + continue
786 +
787 + return loaded_params
788 +
789 + @staticmethod
790 + def _pad_shared_expert_weight(
791 + quant_config: QuantizationConfig | None,
792 + name: str,
793 + loaded_weight: torch.Tensor,
794 + ) -> torch.Tensor:
795 + """Zero-pad a block-FP8 shared-expert weight/scale on its intermediate
796 + axis so the standard TP loaders split it into even, block-aligned shards
797 + (trailing ranks get the zero pad). gate (w1)/up (w3) [I, H] pad dim 0;
798 + down (w2 -> down_proj) [H, I] pads dim 1.
799 + """
800 + block_size = getattr(quant_config, "weight_block_size", None)
801 + assert block_size is not None
802 + # Round the intermediate axis up to a whole number of TP shards. The axis
803 + # is in elements for weights (step = block) and in blocks for scales.
804 + step = (
805 + 1 if name.endswith(("weight_scale_inv", "weight_scale")) else block_size[0]
806 + )
807 + dim = 1 if ".down_proj." in name else 0
808 + mult = get_tensor_model_parallel_world_size() * step
809 + pad = cdiv(loaded_weight.shape[dim], mult) * mult - loaded_weight.shape[dim]
810 + if pad == 0:
811 + return loaded_weight
812 + pad_shape = list(loaded_weight.shape)
813 + pad_shape[dim] = pad
814 + return torch.cat([loaded_weight, loaded_weight.new_zeros(pad_shape)], dim=dim)
815 +
816 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
817 + # Params for weights, fp8 weight scales, fp8 activation scales
818 + # (param_name, weight_name, expert_id, shard_id)
819 + return fused_moe_make_expert_params_mapping(
820 + self,
821 + ckpt_gate_proj_name="w1",
822 + ckpt_down_proj_name="w2",
823 + ckpt_up_proj_name="w3",
824 + num_experts=self.config.n_routed_experts,
825 + )
826 +
827 + def finalize_mega_moe_weights(self) -> None:
828 + return
829 +
830 + def finalize_mhc_broadcast_weights(self) -> None:
831 + if not get_pp_group().is_first_rank or self.start_layer >= self.end_layer:
832 + return
833 + layer = self.layers[self.start_layer]
834 + if isinstance(layer, DeepseekV4DecoderLayer):
835 + broadcast = (
836 + layer.hc_attn_fn.detach()
837 + .view(-1, layer.hc_mult, layer.hidden_size)
838 + .sum(dim=1)
839 + )
840 + if layer.hc_attn_fn_broadcast is None:
841 + layer.hc_attn_fn_broadcast = broadcast
842 + else:
843 + layer.hc_attn_fn_broadcast.copy_(broadcast)
844 +
845 +
846 +def _linear_scale_param_name(vllm_config: VllmConfig, expert_dtype: str) -> str:
847 + """Parameter name the linear quant method registers for weight scales.
848 +
849 + Native MXFP8 checkpoints ([32, 32] blocks with MXFP4 experts) route linear
850 + layers through ModelOptLinearMethod, which registers ``weight_scale``;
851 + block-FP8 linear layers register ``weight_scale_inv``.
852 + """
853 + use_mxfp8 = (
854 + getattr(vllm_config.quant_config, "weight_block_size", None) == [32, 32]
855 + and expert_dtype == "fp4"
856 + )
857 + return "weight_scale" if use_mxfp8 else "weight_scale_inv"
858 +
859 +
860 +def _make_deepseek_v4_weights_mapper(
861 + expert_dtype: str, linear_scale_name: str = "weight_scale_inv"
862 +) -> WeightsMapper:
863 + if expert_dtype == "fp4":
864 + # MXFP4 experts use Mxfp4MoEMethod, which registers scales as
865 + # ``w{1,2,3}_weight_scale`` (no _inv suffix). Linear scales use the
866 + # parameter name registered by their quantization method.
867 + scale_regex = {
868 + # ``.base_layer.``-namespace variant (LoRA-wrapped experts).
869 + re.compile(
870 + r"(\.experts\.\d+\.w[123]\.base_layer)\.scale$"
871 + ): r"\1.weight_scale",
872 + re.compile(r"(\.experts\.\d+\.w[123])\.scale$"): r"\1.weight_scale",
873 + # The ``embed.weight`` -> ``embed_tokens.weight`` suffix rule
874 + # renames the engram fp8 table but not its scale; route the
875 + # scale explicitly to the same module.
876 + re.compile(r"(engram\.embed)\.scale$"): r"\1_tokens.weight_scale_inv",
877 + re.compile(r"\.scale$"): f".{linear_scale_name}",
878 + }
879 + else:
880 + # FP8 experts use Fp8MoEMethod (block_quant=True), which registers
881 + # scales as ``w{13,2}_weight_scale_inv``. Map all ``.scale`` keys
882 + # there.
883 + scale_regex = {
884 + # ``.base_layer.``-namespace variant of the above.
885 + re.compile(
886 + r"(\.experts\.\d+\.w[123]\.base_layer)\.scale$"
887 + ): r"\1.weight_scale_inv",
888 + # Same engram reroute as the fp4 branch above.
889 + re.compile(r"(engram\.embed)\.scale$"): r"\1_tokens.weight_scale_inv",
890 + re.compile(r"\.scale$"): f".{linear_scale_name}",
891 + }
892 + return WeightsMapper(
893 + orig_to_new_prefix={
894 + "layers.": "model.layers.",
895 + "embed.": "model.embed.",
896 + "norm.": "model.norm.",
897 + "mtp.": "model.mtp.",
898 + },
899 + orig_to_new_regex=scale_regex,
900 + orig_to_new_suffix={
901 + "head.weight": "lm_head.weight",
902 + "embed.weight": "embed_tokens.weight",
903 + ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias",
904 + },
905 + orig_to_new_substr={
906 + ".shared_experts.w2": ".shared_experts.down_proj",
907 + "mtp.": None,
908 + # The v4.1 checkpoint declares the VL arch even for text-only
909 + # use; the text backbone drops the vision-tower weights.
910 + "vision.": None,
911 + "aligner.": None,
912 + "image_": None,
913 + },
914 + )
915 +
916 +
917 +class DeepseekV4MixtureOfExperts(MixtureOfExperts):
918 + moe_mlp_layers: list["DeepseekV4MoE"]
919 +
920 + def extract_moe_parameters(self, example_moe: "DeepseekV4MoE | None") -> None:
921 + if example_moe is None:
922 + self.num_moe_layers = 0
923 + self.num_expert_groups = 0
924 + self.num_logical_experts = 0
925 + self.num_physical_experts = 0
926 + self.num_local_physical_experts = 0
927 + self.num_routed_experts = 0
928 + self.num_shared_experts = 0
929 + self.num_redundant_experts = 0
930 + return
931 + self.num_logical_experts = example_moe.n_logical_experts
932 + self.num_physical_experts = example_moe.n_physical_experts
933 + self.num_local_physical_experts = example_moe.n_local_physical_experts
934 + self.num_routed_experts = example_moe.n_routed_experts
935 + self.num_shared_experts = example_moe.n_shared_experts
936 + self.num_redundant_experts = example_moe.n_redundant_experts
937 +
938 + def update_physical_experts_metadata(
939 + self,
940 + num_physical_experts: int,
941 + num_local_physical_experts: int,
942 + ) -> None:
943 + assert self.num_local_physical_experts == num_local_physical_experts
944 + self.num_physical_experts = num_physical_experts
945 + self.num_local_physical_experts = num_local_physical_experts
946 + self.num_redundant_experts = num_physical_experts - self.num_logical_experts
947 + for moe in self.moe_mlp_layers:
948 + moe.n_local_physical_experts = num_local_physical_experts
949 + moe.n_physical_experts = num_physical_experts
950 + moe.n_redundant_experts = self.num_redundant_experts
951 + moe.experts.update_expert_map()
952 +
953 +
954 +class DeepseekV41LLMForCausalLM(
955 + nn.Module,
956 + SupportsPP,
957 + SupportsEagle3,
958 + SupportsLoRA,
959 + DeepseekV4MixtureOfExperts,
960 +):
961 + model_cls = DeepseekV4Model
962 +
963 + # Default mapper assumes the original FP4-expert checkpoint layout.
964 + # Overridden per-instance in __init__ when expert_dtype != "fp4".
965 + hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper("fp4")
966 +
967 + packed_modules_mapping = {
968 + "gate_up_proj": ["w1", "w3"],
969 + "fused_wqa_wkv": ["wq_a", "wkv"],
970 + "fused_wkv_wgate": ["wkv", "wgate"],
971 + }
972 +
973 + # The MTP draft head is not LoRA-adapted.
974 + lora_skip_prefixes = ["mtp."]
975 +
976 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
977 + super().__init__()
978 +
979 + config = vllm_config.model_config.hf_config
980 + self.config = config
981 + expert_dtype = getattr(config, "expert_dtype", "fp4")
982 + self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper(
983 + expert_dtype, _linear_scale_param_name(vllm_config, expert_dtype)
984 + )
985 +
986 + self.model = self.model_cls(
987 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
988 + )
989 + if get_pp_group().is_last_rank:
990 + self.lm_head = ParallelLMHead(
991 + config.vocab_size,
992 + config.hidden_size,
993 + prefix=maybe_prefix(prefix, "lm_head"),
994 + )
995 + else:
996 + self.lm_head = PPMissingLayer()
997 + self.logits_processor = LogitsProcessor(config.vocab_size)
998 + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign]
999 + self.model.make_empty_intermediate_tensors
1000 + )
1001 +
1002 + self.set_moe_parameters()
1003 +
1004 + def set_moe_parameters(self) -> None:
1005 + self.num_expert_groups = getattr(self.config, "n_group", 1)
1006 + self.num_moe_layers = self.config.num_hidden_layers
1007 + self.moe_layers: list[nn.Module] = []
1008 + self.moe_mlp_layers: list[DeepseekV4MoE] = []
1009 + example_moe: DeepseekV4MoE | None = None
1010 + for layer in self.model.layers:
1011 + if isinstance(layer, PPMissingLayer):
1012 + continue
1013 + if not isinstance(layer, DeepseekV4DecoderLayer):
1014 + continue
1015 + if isinstance(layer.ffn, DeepseekV4MoE):
1016 + example_moe = layer.ffn
1017 + self.moe_mlp_layers.append(layer.ffn)
1018 + self.moe_layers.append(layer.ffn.experts)
1019 +
1020 + self.num_moe_layers = len(self.moe_layers)
1021 + self.extract_moe_parameters(example_moe)
1022 +
1023 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1024 + return self.model.embed_input_ids(input_ids)
1025 +
1026 + def compute_logits(
1027 + self,
1028 + hidden_states: torch.Tensor,
1029 + ) -> torch.Tensor | None:
1030 + logits = self.logits_processor(self.lm_head, hidden_states)
1031 + return logits
1032 +
1033 + def compute_logits_local(
1034 + self,
1035 + hidden_states: torch.Tensor,
1036 + ) -> torch.Tensor:
1037 + return self.logits_processor(self.lm_head, hidden_states, skip_gather=True)
1038 +
1039 + @staticmethod
1040 + def get_model_state_cls():
1041 + from ..nvidia.model_state import DeepseekV41ModelState
1042 +
1043 + return DeepseekV41ModelState
1044 +
1045 + @property
1046 + def token_lookback_depth(self) -> int:
1047 + """Tokens before a chunk start the engram hash needs; the model runner
1048 + passes them as `lookback_token_ids`."""
1049 + engram_hash = self.model.engram_hash
1050 + return engram_hash.lookback_depth if engram_hash is not None else 0
1051 +
1052 + def forward(
1053 + self,
1054 + input_ids: torch.Tensor,
1055 + positions: torch.Tensor,
1056 + intermediate_tensors: IntermediateTensors | None = None,
1057 + inputs_embeds: torch.Tensor | None = None,
1058 + lookback_token_ids: torch.Tensor | None = None,
1059 + ) -> torch.Tensor | IntermediateTensors:
1060 + hidden_states = self.model(
1061 + input_ids,
1062 + positions,
1063 + intermediate_tensors,
1064 + inputs_embeds,
1065 + lookback_token_ids=lookback_token_ids,
1066 + )
1067 + return hidden_states
1068 +
1069 + def get_mtp_target_hidden_states(self) -> torch.Tensor | None:
1070 + """Pre-collapse residual stream buffer (max_num_batched_tokens,
1071 + hc_mult * hidden_size) for the MTP draft model. Populated by
1072 + forward(); valid after each target step."""
1073 + return getattr(self.model, "_mtp_hidden_buffer", None)
1074 +
1075 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1076 + loader = AutoWeightsLoader(self)
1077 + loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
1078 + self.process_weights_after_loading()
1079 + return loaded_params
1080 +
1081 + def process_weights_after_loading(self) -> None:
1082 + self.model.finalize_mega_moe_weights()
1083 + self.model.finalize_mhc_broadcast_weights()
1084 +
1085 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1086 + return self.model.get_expert_mapping()
vllm/models/deepseek_v4_1/amd/rocm.py new
+905
@@ -0,0 +1,905 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import functools
5 +from dataclasses import dataclass
6 +from typing import cast
7 +
8 +import torch
9 +
10 +from vllm.distributed import (
11 + get_tensor_model_parallel_world_size,
12 + tensor_model_parallel_all_reduce,
13 +)
14 +from vllm.forward_context import get_forward_context
15 +from vllm.logger import init_logger
16 +from vllm.models.deepseek_v4_1.attention import DeepseekV4Attention
17 +from vllm.models.deepseek_v4_1.common.ops import dequantize_and_gather_k_cache
18 +from vllm.models.deepseek_v4_1.sparse_mla import (
19 + DeepseekV4FlashMLAMetadata,
20 + DeepseekV4SparseMLABackend,
21 + DeepseekV4SparseMLAMetadataBuilder,
22 + DeepseekV41SparseSWAMetadataBuilder,
23 +)
24 +from vllm.platforms import current_platform
25 +from vllm.platforms.rocm import _ON_GFX950
26 +from vllm.triton_utils import tl, triton
27 +from vllm.v1.attention.backend import (
28 + CommonAttentionMetadata,
29 +)
30 +from vllm.v1.attention.backends.mla.sparse_swa import (
31 + DeepseekSparseSWABackend,
32 + DeepseekSparseSWAMetadata,
33 +)
34 +from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
35 + build_ragged_indices_from_dense,
36 + rocm_inv_rope_einsum,
37 + rocm_sparse_attn_decode,
38 + rocm_sparse_attn_prefill,
39 +)
40 +from vllm.v1.worker.workspace import current_workspace_manager
41 +
42 +logger = init_logger(__name__)
43 +
44 +
45 +def _trust_dsv4_extra_cache_nan_free(
46 + kv_cache_dtype: str,
47 + has_kv_transfer: bool,
48 + has_extra_cache: bool,
49 +) -> bool:
50 + return (
51 + _ON_GFX950
52 + and kv_cache_dtype == "fp8_ds_mla"
53 + and not has_kv_transfer
54 + and has_extra_cache
55 + )
56 +
57 +
58 +def _build_indptr_from_lengths(lengths: torch.Tensor) -> torch.Tensor:
59 + lengths = lengths.to(dtype=torch.int32).contiguous()
60 + indptr = torch.zeros(lengths.shape[0] + 1, dtype=torch.int32, device=lengths.device)
61 + torch.cumsum(lengths, dim=0, out=indptr[1:])
62 + return indptr
63 +
64 +
65 +def apply_pre_quantized_block_scaled_mm(
66 + linear: torch.nn.Module,
67 + x_fp8: torch.Tensor,
68 + x_scale: torch.Tensor,
69 +) -> torch.Tensor:
70 + """Block-scaled fp8 GEMM on pre-quantized activations.
71 +
72 + The fused q/kv norm kernel writes fp8 qr + per-1x128 scales; this
73 + drives the linear's block-scaled GEMM directly with them, bypassing
74 + apply_weights which would re-quantize the fp8 input. Only valid for
75 + the wq_b-style column/replicated linears: their output is the local
76 + TP shard, so no all-reduce is needed.
77 + """
78 + from vllm.model_executor.kernels.linear.scaled_mm.BlockScaledMMLinearKernel import (
79 + FP8BlockParams,
80 + )
81 +
82 + params = FP8BlockParams.from_layer(linear)
83 + weight_scale = (
84 + params.weight_scale
85 + if params.weight_scale_inv is None
86 + else params.weight_scale_inv
87 + )
88 + kernel = linear.quant_method.fp8_linear
89 + out = kernel.apply_block_scaled_mm(
90 + A=x_fp8, B=params.weight, As=x_scale, Bs=weight_scale
91 + )
92 + return out.to(dtype=kernel.config.out_dtype)
93 +
94 +
95 +# ROCm sparse prefill keeps this dense combine local so AMD-specific SWA changes
96 +# do not touch the shared DeepSeek V4 cache utilities.
97 +_SPARSE_PREFILL_TOPK_ALIGNMENT = 128
98 +
99 +
100 +@triton.jit
101 +def _combine_topk_swa_indices_kernel(
102 + combined_indices_ptr,
103 + combined_indices_stride,
104 + combined_lens_ptr,
105 + topk_indices_ptr,
106 + topk_indices_stride,
107 + query_start_loc_ptr,
108 + seq_lens_ptr,
109 + gather_lens_ptr,
110 + M,
111 + N,
112 + TOP_K: tl.constexpr,
113 + COMPRESS_RATIO: tl.constexpr,
114 + WINDOW_SIZE: tl.constexpr,
115 + TOPK_WIDTH: tl.constexpr,
116 + PADDED_TOP_K: tl.constexpr,
117 +):
118 + batch_idx = tl.program_id(0)
119 + worker_id = tl.program_id(1)
120 + num_workers = tl.num_programs(1)
121 +
122 + base = tl.load(query_start_loc_ptr)
123 + query_start = tl.load(query_start_loc_ptr + batch_idx) - base
124 + query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base
125 + query_len = query_end - query_start
126 + seq_len = tl.load(seq_lens_ptr + batch_idx)
127 + gather_len = tl.load(gather_lens_ptr + batch_idx)
128 + start_pos = seq_len - query_len
129 + gather_start = seq_len - gather_len
130 +
131 + for token_idx in range(query_start + worker_id, query_end, num_workers):
132 + token_idx_in_query = token_idx - query_start
133 + pos = start_pos + token_idx_in_query
134 + topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
135 + swa_len = tl.minimum(pos + 1, WINDOW_SIZE)
136 +
137 + topk_offset = tl.arange(0, PADDED_TOP_K)
138 + topk_mask = topk_offset < topk_len
139 + safe_topk_offset = tl.where(topk_offset < TOPK_WIDTH, topk_offset, 0)
140 + topk_indices = tl.load(
141 + topk_indices_ptr + token_idx * topk_indices_stride + safe_topk_offset,
142 + mask=topk_mask,
143 + other=-1,
144 + )
145 + valid_topk = (topk_indices >= 0) & (topk_indices < N)
146 + topk_indices = tl.where(valid_topk, topk_indices + M * batch_idx, -1)
147 + tl.store(
148 + combined_indices_ptr + token_idx * combined_indices_stride + topk_offset,
149 + topk_indices,
150 + mask=topk_mask,
151 + )
152 +
153 + swa_offset = tl.arange(0, WINDOW_SIZE)
154 + tl.store(
155 + combined_indices_ptr
156 + + token_idx * combined_indices_stride
157 + + topk_len
158 + + swa_offset,
159 + M * batch_idx + N + swa_offset + pos - swa_len + 1 - gather_start,
160 + mask=swa_offset < swa_len,
161 + )
162 +
163 + tl.store(combined_lens_ptr + token_idx, topk_len + swa_len)
164 +
165 +
166 +def combine_topk_swa_indices(
167 + topk_indices: torch.Tensor,
168 + query_start_loc: torch.Tensor,
169 + seq_lens: torch.Tensor,
170 + gather_lens: torch.Tensor,
171 + window_size: int,
172 + compress_ratio: int,
173 + topk: int,
174 + M: int,
175 + N: int,
176 +) -> tuple[torch.Tensor, torch.Tensor]:
177 + """Combine compressed-attention and sliding-window indices with Torch.
178 +
179 + The Triton implementation inherited from DeepSeek V4 launches a
180 + two-dimensional grid with 128 workers per request. On gfx950 it can issue
181 + an out-of-bounds access for V4.1's mixed prefill metadata (including the
182 + synthetic mixed-token warmup). This path is prefill-only and the tensors
183 + are small, so use ordinary Torch indexing until a gfx950-safe fused kernel
184 + is available.
185 + """
186 + topk_indices = topk_indices.reshape(topk_indices.shape[0], -1).contiguous()
187 + num_tokens = topk_indices.shape[0]
188 + combined_topk = (
189 + (topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1)
190 + // _SPARSE_PREFILL_TOPK_ALIGNMENT
191 + * _SPARSE_PREFILL_TOPK_ALIGNMENT
192 + )
193 + combined_indices = torch.full(
194 + (num_tokens, combined_topk),
195 + fill_value=-1,
196 + dtype=torch.int32,
197 + device=topk_indices.device,
198 + )
199 + combined_lens = torch.empty(
200 + num_tokens, dtype=torch.int32, device=topk_indices.device
201 + )
202 +
203 + # query_start_loc may have a non-zero base for a narrowed mixed batch.
204 + query_lens = query_start_loc[1:] - query_start_loc[:-1]
205 + req_ids = torch.repeat_interleave(
206 + torch.arange(seq_lens.shape[0], device=seq_lens.device), query_lens
207 + )
208 + query_starts = query_start_loc[:-1] - query_start_loc[0]
209 + token_offsets = torch.arange(num_tokens, device=seq_lens.device) - (
210 + torch.repeat_interleave(query_starts, query_lens)
211 + )
212 + positions = seq_lens[req_ids] - query_lens[req_ids] + token_offsets
213 +
214 + logical_topk_width = min(topk, topk_indices.shape[1])
215 + topk_lens = torch.minimum(
216 + (positions + 1) // compress_ratio,
217 + torch.full_like(positions, logical_topk_width),
218 + ).clamp_min(0)
219 + topk_offsets = torch.arange(logical_topk_width, device=seq_lens.device)
220 + topk_mask = topk_offsets[None, :] < topk_lens[:, None]
221 + topk_values = topk_indices[:, :logical_topk_width].to(torch.int32)
222 + topk_valid = topk_mask & (topk_values >= 0) & (topk_values < N)
223 + combined_indices[:, :logical_topk_width] = torch.where(
224 + topk_valid,
225 + topk_values + (M * req_ids).to(torch.int32)[:, None],
226 + -1,
227 + )
228 +
229 + swa_lens = torch.minimum(
230 + positions + 1, torch.full_like(positions, window_size)
231 + ).clamp_min(0)
232 + swa_offsets = torch.arange(window_size, device=seq_lens.device)
233 + swa_mask = swa_offsets[None, :] < swa_lens[:, None]
234 + swa_columns = topk_lens[:, None] + swa_offsets[None, :]
235 + gather_starts = seq_lens - gather_lens
236 + swa_values = (
237 + M * req_ids[:, None]
238 + + N
239 + + swa_offsets[None, :]
240 + + positions[:, None]
241 + - swa_lens[:, None]
242 + + 1
243 + - gather_starts[req_ids, None]
244 + ).to(torch.int32)
245 + rows = torch.arange(num_tokens, device=seq_lens.device)[:, None].expand_as(
246 + swa_columns
247 + )
248 + flat_dst = rows[swa_mask] * combined_topk + swa_columns[swa_mask]
249 + combined_indices.view(-1).index_copy_(0, flat_dst, swa_values[swa_mask])
250 + combined_lens.copy_((topk_lens + swa_lens).to(torch.int32))
251 + return combined_indices, combined_lens
252 +
253 +
254 +@triton.jit
255 +def _compute_topk_lens_kernel(
256 + topk_lens_ptr,
257 + topk_indices_ptr,
258 + topk_indices_stride,
259 + topk,
260 + is_valid_token_ptr,
261 + TRITON_BLOCK_SIZE: tl.constexpr,
262 +):
263 + token_idx = tl.program_id(0)
264 + is_valid_token = tl.load(is_valid_token_ptr + token_idx)
265 +
266 + count = tl.zeros((), dtype=tl.int32)
267 + for i in range(0, topk, TRITON_BLOCK_SIZE):
268 + offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
269 + mask = offset < topk
270 + local_idx = tl.load(
271 + topk_indices_ptr + token_idx * topk_indices_stride + offset,
272 + mask=mask,
273 + other=-1,
274 + )
275 + count += tl.sum((local_idx >= 0).to(tl.int32), axis=0)
276 +
277 + tl.store(topk_lens_ptr + token_idx, tl.where(is_valid_token, count, 0))
278 +
279 +
280 +@triton.jit
281 +def _pack_global_topk_ragged_kernel(
282 + global_topk_ragged_ptr,
283 + topk_indptr_ptr,
284 + topk_indices_ptr,
285 + topk_indices_stride,
286 + token_to_req_indices_ptr,
287 + block_table_ptr,
288 + block_table_stride,
289 + block_size,
290 + topk,
291 + BLOCK_SIZE: tl.constexpr,
292 +):
293 + token_idx = tl.program_id(0)
294 + block_idx = tl.program_id(1)
295 + offset = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
296 +
297 + out_start = tl.load(topk_indptr_ptr + token_idx)
298 + out_end = tl.load(topk_indptr_ptr + token_idx + 1)
299 + out_len = out_end - out_start
300 + if block_idx * BLOCK_SIZE >= out_len:
301 + return
302 +
303 + req_idx = tl.load(token_to_req_indices_ptr + token_idx)
304 + mask = (offset < out_len) & (offset < topk)
305 + local_idx = tl.load(
306 + topk_indices_ptr + token_idx * topk_indices_stride + offset,
307 + mask=mask,
308 + other=-1,
309 + )
310 + valid = mask & (local_idx >= 0)
311 + block_indices = local_idx // block_size
312 + block_numbers = tl.load(
313 + block_table_ptr + req_idx * block_table_stride + block_indices,
314 + mask=valid,
315 + other=0,
316 + )
317 + block_offsets = local_idx % block_size
318 + slot_ids = tl.where(valid, block_numbers * block_size + block_offsets, -1)
319 + tl.store(global_topk_ragged_ptr + out_start + offset, slot_ids, mask=mask)
320 +
321 +
322 +def compute_global_topk_ragged_indices_and_indptr(
323 + topk_indices: torch.Tensor,
324 + token_to_req_indices: torch.Tensor,
325 + block_table: torch.Tensor,
326 + block_size: int,
327 + is_valid_token: torch.Tensor,
328 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
329 + topk_indices = topk_indices.reshape(topk_indices.shape[0], -1).contiguous()
330 + num_tokens = topk_indices.shape[0]
331 + topk = topk_indices.shape[1]
332 +
333 + topk_lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device)
334 + _compute_topk_lens_kernel[(num_tokens,)](
335 + topk_lens,
336 + topk_indices,
337 + topk_indices.stride(0),
338 + topk,
339 + is_valid_token,
340 + TRITON_BLOCK_SIZE=1024,
341 + )
342 +
343 + topk_indptr = _build_indptr_from_lengths(topk_lens)
344 + global_topk_ragged = torch.empty(
345 + num_tokens * topk,
346 + dtype=torch.int32,
347 + device=topk_indices.device,
348 + )
349 + if global_topk_ragged.numel() > 0:
350 + block = 128
351 + _pack_global_topk_ragged_kernel[(num_tokens, triton.cdiv(topk, block))](
352 + global_topk_ragged,
353 + topk_indptr,
354 + topk_indices,
355 + topk_indices.stride(0),
356 + token_to_req_indices,
357 + block_table,
358 + block_table.stride(0),
359 + block_size,
360 + topk,
361 + BLOCK_SIZE=block,
362 + )
363 + return global_topk_ragged, topk_indptr, topk_lens
364 +
365 +
366 +def _copy_ragged_to_graph_buffers(
367 + ragged_indices: torch.Tensor,
368 + ragged_indptr: torch.Tensor,
369 + ragged_indices_buffer: torch.Tensor,
370 + ragged_indptr_buffer: torch.Tensor,
371 + num_rows: int,
372 + max_entries_per_row: int,
373 +) -> tuple[torch.Tensor, torch.Tensor]:
374 + """Copy dynamic ragged metadata into persistent CUDA graph buffers.
375 +
376 + FULL decode graphs capture kernel argument addresses. Keep the returned
377 + tensors backed by stable storage, while indptr continues to bound reads.
378 + """
379 + indptr_out = ragged_indptr_buffer[: num_rows + 1]
380 + indptr_out.copy_(ragged_indptr, non_blocking=True)
381 +
382 + max_entries = max(num_rows * max_entries_per_row, 1)
383 + ragged_out = ragged_indices_buffer[:max_entries]
384 + source_entries = ragged_indices.numel()
385 + if source_entries > 0:
386 + ragged_out[:source_entries].copy_(ragged_indices, non_blocking=True)
387 + if _ON_GFX950:
388 + # Preserve the graph-stable base pointer while exposing source capacity
389 + # to the sync-free split selector; indptr still carries the true NNZ.
390 + ragged_out = ragged_out[: max(source_entries, 1)]
391 + return ragged_out, indptr_out
392 +
393 +
394 +@dataclass
395 +class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata):
396 + decode_swa_ragged_indices: torch.Tensor | None = None
397 + decode_swa_ragged_indptr: torch.Tensor | None = None
398 +
399 +
400 +class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekV41SparseSWAMetadataBuilder):
401 + # Keep fused multi-step decode disabled until update_draft_decode_metadata()
402 + # also refreshes the ROCm-specific ragged SWA indices and indptrs.
403 + supports_draft_decode_metadata_update = False
404 +
405 + def __init__(self, *args, **kwargs):
406 + super().__init__(*args, **kwargs)
407 + max_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens
408 + # The non-causal (DSpark draft) path widens each token's SWA index list
409 + # to ``noncausal_index_width`` (>= window_size), so size the persistent
410 + # ragged buffer to the wider bound to cover both causal and non-causal.
411 + swa_index_width = max(self.window_size, self.noncausal_index_width)
412 + self.decode_swa_ragged_indices_buffer = torch.empty(
413 + max_tokens * swa_index_width,
414 + dtype=torch.int32,
415 + device=self.device,
416 + )
417 + self.decode_swa_ragged_indptr_buffer = torch.empty(
418 + max_tokens + 1,
419 + dtype=torch.int32,
420 + device=self.device,
421 + )
422 +
423 + def build(
424 + self,
425 + common_prefix_len: int,
426 + common_attn_metadata: CommonAttentionMetadata,
427 + fast_build: bool = False,
428 + ) -> DeepseekV4ROCMAiterSparseSWAMetadata:
429 + base = super().build(
430 + common_prefix_len=common_prefix_len,
431 + common_attn_metadata=common_attn_metadata,
432 + fast_build=fast_build,
433 + )
434 +
435 + ragged_indices = None
436 + ragged_indptr = None
437 + if (
438 + base.num_decode_tokens > 0
439 + and base.decode_swa_indices is not None
440 + and base.decode_swa_lens is not None
441 + ):
442 + ragged_indices, ragged_indptr = build_ragged_indices_from_dense(
443 + base.decode_swa_indices.reshape(
444 + base.num_decode_tokens, base.decode_swa_width
445 + ),
446 + base.decode_swa_lens,
447 + )
448 + ragged_indices, ragged_indptr = _copy_ragged_to_graph_buffers(
449 + ragged_indices,
450 + ragged_indptr,
451 + self.decode_swa_ragged_indices_buffer,
452 + self.decode_swa_ragged_indptr_buffer,
453 + base.num_decode_tokens,
454 + base.decode_swa_width,
455 + )
456 +
457 + return DeepseekV4ROCMAiterSparseSWAMetadata(
458 + **vars(base),
459 + decode_swa_ragged_indices=ragged_indices,
460 + decode_swa_ragged_indptr=ragged_indptr,
461 + )
462 +
463 +
464 +class DeepseekV4ROCMAiterMLASparseBackend(DeepseekV4SparseMLABackend):
465 + @staticmethod
466 + def get_name() -> str:
467 + return "ROCM_FLASHMLA_SPARSE_DSV4"
468 +
469 + @staticmethod
470 + def get_builder_cls() -> type[DeepseekV4SparseMLAMetadataBuilder]:
471 + return DeepseekV4SparseMLAMetadataBuilder
472 +
473 +
474 +class DeepseekV41ROCMAiterSparseSWABackend(DeepseekSparseSWABackend):
475 + @staticmethod
476 + def get_builder_cls() -> type["DeepseekV4ROCMAiterSparseSWAMetadataBuilder"]:
477 + return DeepseekV4ROCMAiterSparseSWAMetadataBuilder
478 +
479 +
480 +class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
481 + """ROCm sparse MLA attention layer for DeepSeek V4.1."""
482 +
483 + backend_cls = DeepseekV4ROCMAiterMLASparseBackend
484 + swa_backend_cls = DeepseekV41ROCMAiterSparseSWABackend
485 +
486 + def __init__(self, *args, **kwargs):
487 + vllm_config = args[0] if args else kwargs["vllm_config"]
488 + super().__init__(*args, **kwargs)
489 + # CUDA executes WO_A with a quantized grouped-BMM kernel. ROCm's
490 + # correctness path below dequantizes WO_A once and uses torch.einsum,
491 + # so retain the ordinary MXFP8 linear kernel for post-load processing
492 + # instead of asking for the CUDA-only BMM kernel.
493 + self.wo_a.is_bmm = False
494 + self._has_kv_transfer = vllm_config.kv_transfer_config is not None
495 + # Block scale for the preshuffled weight; None = not preshuffled.
496 + self._wqa_wkv_scale: torch.Tensor | None = None
497 + self._wo_b_scale: torch.Tensor | None = None
498 + self._fused_compressor_weight: torch.Tensor | None
499 + self.register_buffer("_fused_compressor_weight", None, persistent=False)
500 + self._fused_compressor_split_sizes: tuple[int, int] | None = None
501 +
502 + @classmethod
503 + def get_padded_num_q_heads(cls, num_heads: int) -> int:
504 + return num_heads
505 +
506 + def prepare_attn_preshuffle(self) -> None:
507 + from vllm._aiter_ops import rocm_aiter_ops
508 +
509 + if not rocm_aiter_ops.is_enabled():
510 + return
511 + from vllm.model_executor.layers.quantization.utils.fp8_utils import (
512 + _upcast_e8m0_to_fp32,
513 + )
514 + from vllm.model_executor.utils import replace_parameter
515 +
516 + def _prep(linear) -> torch.Tensor | None:
517 + w = getattr(linear, "weight", None)
518 + if w is None or w.dim() != 2:
519 + return None
520 + # K % 128 (group-128 quant) and N % 16 (shuffle_weight) must hold.
521 + if w.shape[-1] % 128 != 0 or w.shape[0] % 16 != 0:
522 + return None
523 + ws = getattr(linear, "weight_scale_inv", None) # per-block scale
524 + if ws is None:
525 + return None
526 + if ws.dtype == torch.float8_e8m0fnu:
527 + ws = _upcast_e8m0_to_fp32(ws).contiguous()
528 + # Shuffle the weight in place (single weight, no unshuffled copy).
529 + replace_parameter(
530 + linear,
531 + "weight",
532 + rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)),
533 + )
534 + return ws
535 +
536 + self._wqa_wkv_scale = _prep(self.fused_wqa_wkv)
537 + self._wo_b_scale = _prep(self.wo_b)
538 +
539 + def prepare_compressor_gemm_fusion(self) -> bool:
540 + # V4.1 derives index keys from the source compressor's emitted latent
541 + # and has no nested ``indexer.compressor``. Keep the projections
542 + # separate and use the shared linear/PyTorch correctness path.
543 + return False
544 +
545 + def _bpre_attn_gemm(
546 + self,
547 + weight: torch.Tensor,
548 + scale: torch.Tensor,
549 + x: torch.Tensor,
550 + reduce_tp: bool,
551 + ) -> torch.Tensor:
552 + from vllm._aiter_ops import rocm_aiter_ops
553 +
554 + x_fp8, x_scale = rocm_aiter_ops.group_fp8_quant(x, transpose_scale=True)
555 + out = rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle(
556 + x_fp8, weight, x_scale, scale, output_dtype=x.dtype
557 + )
558 + if reduce_tp and get_tensor_model_parallel_world_size() > 1:
559 + out = tensor_model_parallel_all_reduce(out)
560 + return out
561 +
562 + def _fused_wqa_wkv_gemm(self, hidden_states: torch.Tensor) -> torch.Tensor:
563 + if self._wqa_wkv_scale is not None and hidden_states.dim() == 2:
564 + return self._bpre_attn_gemm(
565 + self.fused_wqa_wkv.weight, self._wqa_wkv_scale, hidden_states, False
566 + )
567 + return super()._fused_wqa_wkv_gemm(hidden_states)
568 +
569 + def _run_parallel_input_projections(
570 + self, hidden_states: torch.Tensor
571 + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
572 + return super()._run_parallel_input_projections(hidden_states)
573 +
574 + @functools.cached_property
575 + def _wq_b_uses_aiter_block_scaled(self) -> bool:
576 + """True when both wq_b GEMMs run the aiter block-scaled fp8 kernel.
577 +
578 + Cached: the linear kernels and the aiter env gates are fixed once
579 + the model is built, so this is evaluated at the first forward
580 + only.
581 +
582 + The fused norm+quant path is only valid if the quant and GEMM it
583 + replaces are exactly the aiter ones; otherwise fall back to the
584 + shared path.
585 + """
586 + from vllm._aiter_ops import rocm_aiter_ops
587 + from vllm.model_executor.kernels.linear.scaled_mm import (
588 + Fp8BlockScaledMMLinearKernel,
589 + )
590 +
591 + if not rocm_aiter_ops.is_linear_fp8_enabled():
592 + return False
593 +
594 + linears = [self.wq_b]
595 + if self.indexer is not None:
596 + linears.append(self.indexer.wq_b)
597 + for linear in linears:
598 + kernel = getattr(getattr(linear, "quant_method", None), "fp8_linear", None)
599 + if not isinstance(kernel, Fp8BlockScaledMMLinearKernel):
600 + return False
601 + return True
602 +
603 + def _split_qkv_and_norm(
604 + self, qr_kv: torch.Tensor
605 + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]:
606 + """Fuse q/kv RMSNorm + per-1x128 fp8 q quant into one aiter kernel.
607 +
608 + The shared path norms q and kv in one triton kernel and the wq_b
609 + linears then re-read the bf16 qr to quantize it. The aiter kernel
610 + computes both RMSNorms (fp32 accumulate) and the fp8 group quant
611 + in a single pass, writing fp8 qr + group scales directly; both
612 + wq_b GEMMs (attention and indexer) then consume that pair and
613 + skip their own input quant. kv stays bf16: the fused insert
614 + kernel RoPE/quantizes it itself. Falls back to the shared path
615 + when the aiter linear path is not active.
616 + """
617 + qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1)
618 + if not (
619 + qr.dim() == 2
620 + and qr.shape[0] > 0
621 + and self.q_lora_rank % 128 == 0
622 + and self._wq_b_uses_aiter_block_scaled
623 + ):
624 + return super()._split_qkv_and_norm(qr_kv)
625 +
626 + from vllm._aiter_ops import rocm_aiter_ops
627 +
628 + return rocm_aiter_ops.fused_qk_rmsnorm_group_quant(
629 + q=qr,
630 + q_weight=self.q_norm.weight.data,
631 + q_epsilon=self.eps,
632 + kv=kv,
633 + kv_weight=self.kv_norm.weight.data,
634 + kv_epsilon=self.eps,
635 + group_size=128,
636 + transpose_scale=False,
637 + )
638 +
639 + def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
640 + # ROCm BF16 reference wo_a path (inverse RoPE + einsum) + wo_b.
641 + z = rocm_inv_rope_einsum(
642 + self.rotary_emb,
643 + o,
644 + positions,
645 + self.rope_head_dim,
646 + self.n_local_groups,
647 + self.o_lora_rank,
648 + self.wo_a,
649 + )
650 + zf = z.flatten(1)
651 + if self._wo_b_scale is not None and zf.dim() == 2:
652 + result = self._bpre_attn_gemm(self.wo_b.weight, self._wo_b_scale, zf, True)
653 + else:
654 + result = self.wo_b(zf)
655 + return result
656 +
657 + def forward_mqa(
658 + self,
659 + q: torch.Tensor,
660 + kv: torch.Tensor,
661 + positions: torch.Tensor,
662 + output: torch.Tensor,
663 + ) -> None:
664 + assert output.shape == q.shape, (
665 + f"output buffer shape {output.shape} must match q shape {q.shape}"
666 + )
667 + assert output.dtype == q.dtype, (
668 + f"output buffer dtype {output.dtype} must match q dtype {q.dtype}"
669 + )
670 +
671 + forward_context = get_forward_context()
672 + attn_metadata = forward_context.attn_metadata
673 +
674 + if attn_metadata is None:
675 + # Warmup dummy run: no real metadata. Reserve the same bf16
676 + # gather workspace _forward_prefill would; the dequantize / topk
677 + # / sparse_fwd kernels are skipped this step.
678 + swa_only = self.compress_ratio == 0
679 + N = (
680 + 0
681 + if swa_only
682 + else (self.max_model_len + self.compress_ratio - 1)
683 + // self.compress_ratio
684 + )
685 + M = N + self.window_size + self.max_num_batched_tokens
686 + current_workspace_manager().get_simultaneous(
687 + ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
688 + )
689 + output.zero_()
690 + return
691 +
692 + assert isinstance(attn_metadata, dict)
693 + rocm_metadata = cast(
694 + DeepseekV4FlashMLAMetadata | None,
695 + attn_metadata.get(self.compressed_cache_prefix)
696 + if self.compressed_cache_prefix is not None
697 + else None,
698 + )
699 + swa_metadata = cast(
700 + DeepseekV4ROCMAiterSparseSWAMetadata | None,
701 + attn_metadata.get(self.swa_cache_layer.prefix),
702 + )
703 + assert swa_metadata is not None
704 +
705 + swa_only = self.compress_ratio == 0
706 + self_kv_cache = None if swa_only else self._compressed_kv_cache()
707 + swa_kv_cache = self.swa_cache_layer.kv_cache
708 +
709 + num_decodes = swa_metadata.num_decodes
710 + num_prefills = swa_metadata.num_prefills
711 + num_decode_tokens = swa_metadata.num_decode_tokens
712 +
713 + if num_prefills > 0:
714 + self._forward_prefill(
715 + q=q[num_decode_tokens:],
716 + positions=positions[num_decode_tokens:],
717 + compressed_k_cache=self_kv_cache,
718 + swa_k_cache=swa_kv_cache,
719 + output=output[num_decode_tokens:],
720 + attn_metadata=rocm_metadata,
721 + swa_metadata=swa_metadata,
722 + )
723 + if num_decodes > 0:
724 + self._forward_decode(
725 + q=q[:num_decode_tokens],
726 + kv_cache=self_kv_cache,
727 + swa_metadata=swa_metadata,
728 + attn_metadata=rocm_metadata,
729 + swa_only=swa_only,
730 + output=output[:num_decode_tokens],
731 + )
732 +
733 + def _forward_decode(
734 + self,
735 + q: torch.Tensor,
736 + kv_cache: torch.Tensor | None,
737 + swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,
738 + attn_metadata: DeepseekV4FlashMLAMetadata | None,
739 + swa_only: bool,
740 + output: torch.Tensor,
741 + ) -> None:
742 + num_decodes = swa_metadata.num_decodes
743 + num_decode_tokens = swa_metadata.num_decode_tokens
744 +
745 + topk_lens = None
746 + topk_ragged_indices = None
747 + topk_ragged_indptr = None
748 + if not swa_only:
749 + # Local indices filled by the index-source layer's indexer.
750 + assert attn_metadata is not None
751 + assert swa_metadata.is_valid_token is not None
752 + assert self.topk_indices_buffer is not None
753 + block_size = attn_metadata.block_size // self.compress_ratio
754 + is_valid = swa_metadata.is_valid_token[:num_decode_tokens]
755 + (
756 + topk_ragged_indices,
757 + topk_ragged_indptr,
758 + topk_lens,
759 + ) = compute_global_topk_ragged_indices_and_indptr(
760 + self.topk_indices_buffer[:num_decode_tokens],
761 + swa_metadata.token_to_req_indices,
762 + attn_metadata.block_table[:num_decodes],
763 + block_size,
764 + is_valid,
765 + )
766 +
767 + rocm_sparse_attn_decode(
768 + q=q,
769 + kv_cache=kv_cache,
770 + swa_k_cache=self.swa_cache_layer.kv_cache,
771 + swa_only=swa_only,
772 + topk_indices=None,
773 + topk_lens=topk_lens,
774 + swa_indices=swa_metadata.decode_swa_indices,
775 + swa_lens=swa_metadata.decode_swa_lens,
776 + swa_ragged_indices=swa_metadata.decode_swa_ragged_indices,
777 + swa_ragged_indptr=swa_metadata.decode_swa_ragged_indptr,
778 + topk_ragged_indices=topk_ragged_indices,
779 + topk_ragged_indptr=topk_ragged_indptr,
780 + attn_sink=self.attn_sink,
781 + scale=self.scale,
782 + head_dim=self.head_dim,
783 + nope_head_dim=self.nope_head_dim,
784 + rope_head_dim=self.rope_head_dim,
785 + output=output,
786 + extra_cache_nan_free=_trust_dsv4_extra_cache_nan_free(
787 + self.kv_cache_dtype,
788 + self._has_kv_transfer,
789 + not swa_only and kv_cache is not None,
790 + ),
791 + )
792 +
793 + def _forward_prefill(
794 + self,
795 + q: torch.Tensor,
796 + positions: torch.Tensor,
797 + compressed_k_cache: torch.Tensor | None,
798 + swa_k_cache: torch.Tensor,
799 + output: torch.Tensor,
800 + attn_metadata: DeepseekV4FlashMLAMetadata | None,
801 + swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,
802 + ) -> None:
803 + swa_only = attn_metadata is None
804 +
805 + num_prefills = swa_metadata.num_prefills
806 + num_prefill_tokens = swa_metadata.num_prefill_tokens
807 + num_decodes = swa_metadata.num_decodes
808 + num_decode_tokens = swa_metadata.num_decode_tokens
809 +
810 + seq_lens = swa_metadata.prefill_seq_lens
811 + gather_lens = swa_metadata.prefill_gather_lens
812 + assert seq_lens is not None
813 + assert gather_lens is not None
814 +
815 + query_start_loc_cpu = swa_metadata.query_start_loc_cpu
816 + query_start_loc = swa_metadata.query_start_loc
817 + assert query_start_loc_cpu is not None
818 + assert query_start_loc is not None
819 + prefill_token_base = query_start_loc_cpu[num_decodes]
820 +
821 + # Local indices filled by the index source; SWA-only layers pass
822 + # top_k=0 and never read them.
823 + assert self.topk_indices_buffer is not None
824 + topk_indices = self.topk_indices_buffer[num_decode_tokens:]
825 + topk_indices = topk_indices[:num_prefill_tokens]
826 + if not swa_only:
827 + top_k = topk_indices.shape[-1]
828 + N = (self.max_model_len + self.compress_ratio - 1) // self.compress_ratio
829 + else:
830 + top_k = 0
831 + N = 0
832 +
833 + M = N + self.window_size + self.max_num_batched_tokens
834 + num_chunks = (num_prefills + self.PREFILL_CHUNK_SIZE - 1) // (
835 + self.PREFILL_CHUNK_SIZE
836 + )
837 +
838 + workspace_manager = current_workspace_manager()
839 + kv = workspace_manager.get_simultaneous(
840 + ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
841 + )[0]
842 + for chunk_idx in range(num_chunks):
843 + chunk_start = chunk_idx * self.PREFILL_CHUNK_SIZE
844 + chunk_end = min(chunk_start + self.PREFILL_CHUNK_SIZE, num_prefills)
845 + chunk_size = chunk_end - chunk_start
846 + if not swa_only:
847 + assert attn_metadata is not None
848 + assert compressed_k_cache is not None
849 + block_table = attn_metadata.block_table[num_decodes:]
850 + # compressed_k_cache is OCP on every platform (Triton encoder).
851 + dequantize_and_gather_k_cache(
852 + kv[:chunk_size],
853 + compressed_k_cache,
854 + seq_lens=seq_lens[chunk_start:chunk_end] // self.compress_ratio,
855 + gather_lens=None,
856 + block_table=block_table[chunk_start:chunk_end],
857 + block_size=attn_metadata.block_size // self.compress_ratio,
858 + offset=0,
859 + use_fnuz=False,
860 + )
861 +
862 + swa_block_table = swa_metadata.block_table[num_decodes:]
863 + dequantize_and_gather_k_cache(
864 + kv[:chunk_size],
865 + swa_k_cache,
866 + seq_lens=seq_lens[chunk_start:chunk_end],
867 + gather_lens=gather_lens[chunk_start:chunk_end],
868 + block_table=swa_block_table[chunk_start:chunk_end],
869 + block_size=swa_metadata.block_size,
870 + offset=N,
871 + use_fnuz=current_platform.is_fp8_fnuz(),
872 + )
873 +
874 + query_start = (
875 + query_start_loc_cpu[num_decodes + chunk_start] - prefill_token_base
876 + )
877 + query_end = (
878 + query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base
879 + )
880 +
881 + combined_indices, combined_lens = combine_topk_swa_indices(
882 + topk_indices[query_start:query_end],
883 + query_start_loc[
884 + num_decodes + chunk_start : num_decodes + chunk_end + 1
885 + ],
886 + seq_lens[chunk_start:chunk_end],
887 + gather_lens[chunk_start:chunk_end],
888 + self.window_size,
889 + self.compress_ratio,
890 + top_k,
891 + M,
892 + N,
893 + )
894 + rocm_sparse_attn_prefill(
895 + q=q[query_start:query_end],
896 + kv=kv.view(-1, 1, q.shape[-1]),
897 + indices=combined_indices,
898 + topk_length=combined_lens,
899 + scale=self.scale,
900 + head_dim=self.head_dim,
901 + nope_head_dim=self.nope_head_dim,
902 + rope_head_dim=self.rope_head_dim,
903 + attn_sink=self.attn_sink,
904 + output=output[query_start:query_end],
905 + )
vllm/models/deepseek_v4_1/amd/vl_model.py new
+357
@@ -0,0 +1,357 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""DeepSeek-V4.1 vision variant.
4 +
5 +Thin multimodal wrapper around the text-only ``DeepseekV41LLMForCausalLM``:
6 +
7 +- ``vision`` ViT + ``aligner`` produce per-image embeddings for the IMAGE
8 + positions; three learned vectors (``image_start`` / ``image_newline`` /
9 + ``image_end``) fill the delimiter positions. Every image-span position
10 + carries ``image_token_id`` (129264) in ``input_ids``; the per-position
11 + roles come from the processor's ``types`` tensor (see
12 + ``common/mm_preprocess.py``, and ``common/vision.py`` for the tower).
13 +- Merged embeddings enter the text model via ``inputs_embeds``, i.e. before
14 + its hyper-connection stream expansion. Raw ``input_ids`` still flow into
15 + the model so the MoE router can apply ``bias_vl`` to image tokens
16 + (``requires_raw_input_tokens``).
17 +"""
18 +
19 +from collections.abc import Iterable
20 +from typing import Annotated
21 +
22 +import torch
23 +from torch import nn
24 +
25 +from vllm.distributed import get_tensor_model_parallel_world_size
26 +from vllm.model_executor.models.interfaces import (
27 + MultiModalEmbeddings,
28 + SupportsEagle3,
29 + SupportsMultiModal,
30 + SupportsPP,
31 +)
32 +from vllm.model_executor.models.utils import (
33 + AutoWeightsLoader,
34 + WeightsMapper,
35 + maybe_prefix,
36 +)
37 +from vllm.model_executor.models.vision import is_vit_use_data_parallel
38 +from vllm.models.deepseek_v4.common.vision import (
39 + DeepseekV4Aligner,
40 + DeepseekV4ViT,
41 + run_dp_sharded_vision_tower,
42 +)
43 +from vllm.multimodal import MULTIMODAL_REGISTRY
44 +from vllm.utils.tensor_schema import TensorSchema, TensorShape
45 +
46 +from ..common.mm_preprocess import (
47 + IMAGE,
48 + IMAGE_END,
49 + IMAGE_NEW_LINE,
50 + IMAGE_PAD_ID,
51 + IMAGE_PLACEHOLDER,
52 + IMAGE_SENTINEL_BASE_ID,
53 + IMAGE_START,
54 + DeepseekV4VLDummyInputsBuilder,
55 + DeepseekV4VLMultiModalProcessor,
56 + DeepseekV4VLProcessingInfo,
57 +)
58 +from .model import (
59 + DeepseekV41LLMForCausalLM,
60 + _linear_scale_param_name,
61 + _make_deepseek_v4_weights_mapper,
62 +)
63 +
64 +
65 +class DeepseekV4VLImagePixelInputs(TensorSchema):
66 + """ViT patch inputs for one batched set of images.
67 +
68 + Dimensions:
69 + - np: Total ViT patches across images (sum of n_vit_h * n_vit_w)
70 + - c: Number of image channels (3)
71 + - p: ViT patch size
72 + - ni: Number of images
73 + - ns: Total image-span positions (sum of n_llm_h * (n_llm_w + 1) + 2)
74 + """
75 +
76 + patches: Annotated[
77 + torch.Tensor, TensorShape("np", 3, "p", "p", dynamic_dims={"np"})
78 + ]
79 + # [n_vit_h, n_vit_w] per image
80 + vit_grid: Annotated[torch.Tensor, TensorShape("ni", 2)]
81 + # [n_llm_h, n_llm_w] per image
82 + llm_grid: Annotated[torch.Tensor, TensorShape("ni", 2)]
83 + types: Annotated[torch.Tensor, TensorShape("ns", dynamic_dims={"ns"})]
84 +
85 +
86 +def _make_deepseek_v4_vl_weights_mapper(
87 + expert_dtype: str, linear_scale_name: str
88 +) -> WeightsMapper:
89 + """Text-checkpoint mapping rules re-rooted under ``language_model.``."""
90 + base = _make_deepseek_v4_weights_mapper(expert_dtype, linear_scale_name)
91 + orig_to_new_prefix: dict[str, str | None] = {
92 + "layers.": "language_model.model.layers.",
93 + "embed.": "language_model.model.embed.",
94 + "norm.": "language_model.model.norm.",
95 + "hc_head": "language_model.model.hc_head",
96 + "mtp.": "language_model.model.mtp.",
97 + }
98 + return WeightsMapper(
99 + orig_to_new_prefix=orig_to_new_prefix,
100 + orig_to_new_regex=base.orig_to_new_regex,
101 + orig_to_new_suffix={
102 + "head.weight": "language_model.lm_head.weight",
103 + "embed.weight": "embed_tokens.weight",
104 + ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias",
105 + },
106 + orig_to_new_substr={
107 + ".shared_experts.w2": ".shared_experts.down_proj",
108 + # The MTP/DSpark draft heads are not supported for the vision
109 + # variant; drop their weights.
110 + "mtp.": None,
111 + },
112 + )
113 +
114 +
115 +@MULTIMODAL_REGISTRY.register_processor(
116 + DeepseekV4VLMultiModalProcessor,
117 + info=DeepseekV4VLProcessingInfo,
118 + dummy_inputs=DeepseekV4VLDummyInputsBuilder,
119 +)
120 +class DeepseekV41ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3):
121 + """Multimodal entry point for DeepSeek-V4.1 checkpoints with a vision tower.
122 +
123 + ``SupportsEagle3`` (aux hidden-state plumbing for MTP/DSpark drafters)
124 + delegates through ``language_model`` via the protocol defaults.
125 + """
126 +
127 + supports_encoder_tp_data = True
128 +
129 + # The MoE router needs raw token ids to detect image-span tokens
130 + # (all carrying image_token_id, see common/mm_preprocess.py) and apply
131 + # bias_vl.
132 + requires_raw_input_tokens = True
133 +
134 + @classmethod
135 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:
136 + if modality == "image":
137 + return IMAGE_PLACEHOLDER
138 + raise ValueError(f"Unsupported modality: {modality!r}")
139 +
140 + def __init__(self, *, vllm_config, prefix: str = "") -> None:
141 + super().__init__()
142 + model_config = vllm_config.model_config
143 + config = model_config.hf_config
144 + self.config = config
145 + self.multimodal_config = model_config.multimodal_config
146 + assert self.multimodal_config is not None
147 +
148 + # The tower is always built; _mark_tower_model stubs it out
149 + # (StageMissingLayer, weights skipped) when the image limit is 0.
150 + with self._mark_tower_model(vllm_config, {"image"}):
151 + self.use_data_parallel = is_vit_use_data_parallel(config.vision_n_heads)
152 + self.vision = DeepseekV4ViT(config)
153 + self.aligner = DeepseekV4Aligner(config)
154 + self.image_start = nn.Parameter(
155 + torch.empty(config.hidden_size, dtype=torch.float32)
156 + )
157 + self.image_end = nn.Parameter(
158 + torch.empty(config.hidden_size, dtype=torch.float32)
159 + )
160 + self.image_newline = nn.Parameter(
161 + torch.empty(config.hidden_size, dtype=torch.float32)
162 + )
163 + self.vision.to(dtype=model_config.dtype)
164 + self.aligner.to(dtype=model_config.dtype)
165 +
166 + with self._mark_language_model(vllm_config):
167 + self.language_model = DeepseekV41LLMForCausalLM(
168 + vllm_config=vllm_config,
169 + prefix=maybe_prefix(prefix, "language_model"),
170 + )
171 + # The outer mapper (see load_weights) fully resolves HF names into
172 + # this wrapper's namespace before AutoWeightsLoader strips the
173 + # "language_model." prefix and delegates to the child's load_weights,
174 + # so the child's own mapper must be a no-op. Its suffix rules are not
175 + # idempotent (e.g. "lm_head.weight".endswith("head.weight") would
176 + # re-fire "head.weight" -> "lm_head.weight").
177 + self.language_model.hf_to_vllm_mapper = WeightsMapper()
178 + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign]
179 + self.language_model.make_empty_intermediate_tensors
180 + )
181 +
182 + expert_dtype = getattr(config, "expert_dtype", "fp4")
183 + self.hf_to_vllm_mapper = _make_deepseek_v4_vl_weights_mapper(
184 + expert_dtype, _linear_scale_param_name(vllm_config, expert_dtype)
185 + )
186 +
187 + def _parse_and_validate_image_input(
188 + self, **kwargs: object
189 + ) -> DeepseekV4VLImagePixelInputs | None:
190 + patches = kwargs.pop("patches", None)
191 + if patches is None:
192 + return None
193 + return DeepseekV4VLImagePixelInputs(
194 + patches=patches,
195 + vit_grid=kwargs.pop("vit_grid"),
196 + llm_grid=kwargs.pop("llm_grid"),
197 + types=kwargs.pop("types"),
198 + resolve_bindings={"p": self.config.vision_patch_size},
199 + )
200 +
201 + def _encode_image(
202 + self,
203 + patches: torch.Tensor,
204 + n_vit_h: int,
205 + n_vit_w: int,
206 + ) -> torch.Tensor:
207 + # Aligner rows in reading order, one per IMAGE slot.
208 + return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w)
209 +
210 + def _build_image_span(
211 + self, image_embeds: torch.Tensor, types: torch.Tensor
212 + ) -> torch.Tensor:
213 + """Full image span: aligner rows at IMAGE slots, the learned
214 + delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END."""
215 + types = types.to(image_embeds.device)
216 + span = image_embeds.new_empty(types.numel(), image_embeds.shape[-1])
217 + dtype = image_embeds.dtype
218 + span[types == IMAGE_START] = self.image_start.to(dtype)
219 + span[types == IMAGE_END] = self.image_end.to(dtype)
220 + span[types == IMAGE_NEW_LINE] = self.image_newline.to(dtype)
221 + span[types == IMAGE] = image_embeds
222 + return span
223 +
224 + def _process_image_input(
225 + self,
226 + image_input: DeepseekV4VLImagePixelInputs,
227 + ) -> tuple[torch.Tensor, ...]:
228 + patches = image_input.patches.to(self.aligner.w1.weight.dtype)
229 + vit_grid = image_input.vit_grid.tolist()
230 +
231 + image_embeds_list: list[torch.Tensor]
232 + if self.use_data_parallel and get_tensor_model_parallel_world_size() > 1:
233 + # Data-parallel ViT: shard images across TP ranks and all-gather
234 + # the per-image embeddings (weights are replicated on every rank).
235 + image_embeds_list = run_dp_sharded_vision_tower(
236 + self.vision, self.aligner, patches, vit_grid
237 + )
238 + else:
239 + image_embeds_list = []
240 + vit_offset = 0
241 + for n_vit_h, n_vit_w in vit_grid:
242 + n_vit = n_vit_h * n_vit_w
243 + image_embeds_list.append(
244 + self._encode_image(
245 + patches[vit_offset : vit_offset + n_vit], n_vit_h, n_vit_w
246 + )
247 + )
248 + vit_offset += n_vit
249 +
250 + embeds: list[torch.Tensor] = []
251 + span_offset = 0
252 + for image_embeds, (n_llm_h, n_llm_w) in zip(
253 + image_embeds_list, image_input.llm_grid.tolist(), strict=True
254 + ):
255 + span_len = n_llm_h * (n_llm_w + 1) + 2
256 + embeds.append(
257 + self._build_image_span(
258 + image_embeds,
259 + image_input.types[span_offset : span_offset + span_len],
260 + )
261 + )
262 + span_offset += span_len
263 + return tuple(embeds)
264 +
265 + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
266 + image_input = self._parse_and_validate_image_input(**kwargs)
267 + if image_input is None:
268 + return []
269 + return self._process_image_input(image_input)
270 +
271 + def embed_input_ids(
272 + self,
273 + input_ids: torch.Tensor,
274 + multimodal_embeddings: MultiModalEmbeddings | None = None,
275 + *,
276 + is_multimodal: torch.Tensor | None = None,
277 + ) -> torch.Tensor:
278 + from vllm.model_executor.models.utils import _merge_multimodal_embeddings
279 +
280 + # Compressor-alignment pads borrow a reserved id; embed them as the
281 + # plain image token (the checkpoint has no image_pad vector).
282 + # Branch-free: safe inside compiled/captured regions.
283 + input_ids = input_ids.masked_fill(
284 + input_ids == IMAGE_PAD_ID, IMAGE_SENTINEL_BASE_ID
285 + )
286 + inputs_embeds = self.language_model.embed_input_ids(input_ids)
287 +
288 + if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
289 + return inputs_embeds
290 +
291 + assert is_multimodal is not None
292 + return _merge_multimodal_embeddings(
293 + inputs_embeds=inputs_embeds,
294 + multimodal_embeddings=multimodal_embeddings,
295 + is_multimodal=is_multimodal,
296 + )
297 +
298 + @staticmethod
299 + def get_model_state_cls():
300 + from ..nvidia.model_state import DeepseekV41ModelState
301 +
302 + return DeepseekV41ModelState
303 +
304 + @property
305 + def token_lookback_depth(self) -> int:
306 + return self.language_model.token_lookback_depth
307 +
308 + def forward(
309 + self,
310 + input_ids: torch.Tensor,
311 + positions: torch.Tensor,
312 + intermediate_tensors=None,
313 + inputs_embeds: torch.Tensor | None = None,
314 + lookback_token_ids: torch.Tensor | None = None,
315 + **kwargs,
316 + ) -> torch.Tensor:
317 + return self.language_model(
318 + input_ids,
319 + positions,
320 + intermediate_tensors,
321 + inputs_embeds,
322 + lookback_token_ids=lookback_token_ids,
323 + )
324 +
325 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
326 + return self.language_model.compute_logits(hidden_states)
327 +
328 + def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor:
329 + return self.language_model.compute_logits_local(hidden_states)
330 +
331 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
332 + return self.language_model.get_expert_mapping()
333 +
334 + def get_mtp_target_hidden_states(self) -> torch.Tensor | None:
335 + """Pre-hc_head residual stream buffer for the MTP/DSpark draft model."""
336 + return self.language_model.get_mtp_target_hidden_states()
337 +
338 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
339 + # Map HF names into this wrapper's namespace up front and sort, so
340 + # the "language_model." group reaches the child loader as one
341 + # contiguous block (AutoWeightsLoader delegates per contiguous group,
342 + # and the child's load_weights finalizes fused expert weights, which
343 + # must not run on a partially loaded model).
344 + mapped = sorted(self.hf_to_vllm_mapper.apply(weights), key=lambda x: x[0])
345 + loader = AutoWeightsLoader(self)
346 + loaded_params = loader.load_weights(mapped)
347 + # The child's load_weights already ran its post-load finalization.
348 + self._weights_finalized = True
349 + return loaded_params
350 +
351 + def process_weights_after_loading(self) -> None:
352 + # Model-level post-load hook (called by the loader after any load
353 + # format). Under DummyModelLoader the child's load_weights — and
354 + # hence its finalize step — is bypassed, so run it here instead.
355 + if getattr(self, "_weights_finalized", False):
356 + return
357 + self.language_model.process_weights_after_loading()
vllm/models/deepseek_v4_1/attention.py new
+1225
@@ -0,0 +1,1225 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +DeepseekV4 MLA Attention Layer
5 +"""
6 +
7 +from abc import ABC, abstractmethod
8 +from collections.abc import Callable
9 +from functools import cached_property
10 +from typing import TYPE_CHECKING, Any, ClassVar, cast
11 +
12 +import regex as re
13 +import torch
14 +import torch.nn as nn
15 +import torch.nn.functional as F
16 +from transformers import DeepseekV2Config, DeepseekV3Config
17 +
18 +import vllm.envs as envs
19 +from vllm.compilation.breakable_cudagraph import eager_break_during_capture
20 +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
21 +from vllm.model_executor.layers.linear import (
22 + ColumnParallelLinear,
23 + MergedColumnParallelLinear,
24 + ReplicatedLinear,
25 + RowParallelLinear,
26 +)
27 +from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
28 +from vllm.models.common.ops import fused_q_kv_rmsnorm
29 +from vllm.models.deepseek_v4_1.common.ops import (
30 + MXFP4_BLOCK_SIZE,
31 + fused_indexer_q_rope_quant,
32 + indexer_k_norm_rope_store,
33 +)
34 +
35 +if TYPE_CHECKING:
36 + from vllm.v1.attention.backends.mla.sparse_swa import (
37 + DeepseekSparseSWAMetadata,
38 + )
39 +
40 +from vllm.config import (
41 + CacheConfig,
42 + VllmConfig,
43 + get_current_vllm_config,
44 +)
45 +from vllm.distributed import get_tensor_model_parallel_world_size
46 +from vllm.forward_context import get_forward_context
47 +from vllm.logger import init_logger
48 +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
49 +from vllm.model_executor.layers.layernorm import RMSNorm
50 +from vllm.model_executor.layers.quantization import QuantizationConfig
51 +from vllm.model_executor.models.utils import extract_layer_index
52 +from vllm.models.deepseek_v4_1.common.rope import build_deepseek_v4_rope
53 +from vllm.models.deepseek_v4_1.compressor import DeepseekCompressor
54 +from vllm.triton_utils import tl, triton
55 +from vllm.utils.multi_stream_utils import (
56 + execute_in_parallel,
57 + maybe_execute_in_parallel,
58 +)
59 +from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
60 +from vllm.v1.attention.backends.mla.indexer import (
61 + DeepseekV4IndexerBackend,
62 + dsa_indexer_uses_fp4,
63 + get_max_prefill_buffer_size,
64 +)
65 +from vllm.v1.attention.backends.mla.sparse_swa import DeepseekV4SWACache
66 +from vllm.v1.kv_cache_interface import (
67 + KVCacheSpec,
68 + MLAAttentionSpec,
69 + get_kv_quant_mode,
70 +)
71 +
72 +logger = init_logger(__name__)
73 +
74 +
75 +def _replace_layer_index(prefix: str, layer_id: int) -> str:
76 + """Swap the layer index inside a ``...layers.<idx>...`` prefix."""
77 + new_prefix, n = re.subn(r"\.layers\.\d+\.", f".layers.{layer_id}.", prefix)
78 + assert n == 1, f"Cannot locate layer index in prefix {prefix}"
79 + return new_prefix
80 +
81 +
82 +def _indexer_k_cache_head_dim(index_head_dim: int, use_fp4_kv: bool) -> int:
83 + """Per-token byte width of the paged indexer K cache row."""
84 + if use_fp4_kv:
85 + # MXFP4 stores two values per byte plus one UE8M0 byte per 32 values.
86 + # head_dim bytes = 64 packed values + 4 UE8M0 scales = 68.
87 + return index_head_dim // 2 + index_head_dim // MXFP4_BLOCK_SIZE
88 + # NOTE(yifan): FP8 indexer cache uses the same layout as V3.2:
89 + # head_dim bytes = 128 fp8 + 4 fp32 scale = 132.
90 + return index_head_dim + index_head_dim // 128 * 4
91 +
92 +
93 +@triton.jit
94 +def _fill_short_context_topk_indices(
95 + output,
96 + positions,
97 + TOP_K: tl.constexpr,
98 + COMPRESS_RATIO: tl.constexpr,
99 + PADDED_TOP_K: tl.constexpr,
100 +):
101 + # small triton kernel that selects every candidate, -1 otherwise
102 + row = tl.program_id(0)
103 + offsets = tl.arange(0, PADDED_TOP_K)
104 + num_compressed = (tl.load(positions + row) + 1) // COMPRESS_RATIO
105 + tl.store(
106 + output + row * TOP_K + offsets,
107 + tl.where(offsets < num_compressed, offsets, -1),
108 + mask=offsets < TOP_K,
109 + )
110 +
111 +
112 +def _resolve_dsv4_kv_cache_dtype(
113 + use_fp8_ds_mla_layout: bool,
114 + kv_cache_dtype: str,
115 + cache_config: CacheConfig | None,
116 +) -> tuple[str, torch.dtype]:
117 + """Map ``(layout, --kv-cache-dtype)`` to ``(cache_dtype_str, torch_dtype)``.
118 +
119 + Both layouts are paged; they differ in the per-token block format. The
120 + ``fp8_ds_mla`` format is UE8M0 block-scaled fp8 packed as ``uint8`` (the
121 + canonical ``fp8_ds_mla`` string is written back onto ``cache_config`` so the
122 + page-size specs pick the 576B per-token slot). Plain-row backends store each
123 + token's KV row in its element dtype: bf16 or per-tensor FP8 E4M3.
124 + """
125 + if use_fp8_ds_mla_layout:
126 + # fp8_ds_mla block format: UE8M0 block-scaled fp8 packed as uint8.
127 + if kv_cache_dtype == "auto":
128 + kv_cache_dtype = "fp8"
129 + if not kv_cache_dtype.startswith("fp8"):
130 + raise ValueError(
131 + "DeepseekV4 fp8_ds_mla layout only supports fp8 "
132 + f"kv-cache, got {kv_cache_dtype}. Please set "
133 + "`--kv-cache-dtype fp8` or select a backend that supports "
134 + "bfloat16 KV cache."
135 + )
136 + if kv_cache_dtype != "fp8_ds_mla":
137 + if cache_config is not None:
138 + cache_config.cache_dtype = "fp8_ds_mla"
139 + kv_cache_dtype = "fp8_ds_mla"
140 + logger.info_once("Using DeepSeek's fp8_ds_mla KV cache format.")
141 + return kv_cache_dtype, torch.uint8
142 +
143 + # Plain bf16 / per-tensor fp8 KV row (FlashInfer).
144 + if kv_cache_dtype.startswith("fp8"):
145 + return kv_cache_dtype, torch.float8_e4m3fn
146 + # auto / bfloat16 -> plain bf16 KV row.
147 + return kv_cache_dtype, torch.bfloat16
148 +
149 +
150 +class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
151 + """DeepseekV4 MLA attention layer.
152 +
153 + The platform-specific sparse-MLA forward (``forward_mqa`` /
154 + ``get_padded_num_q_heads`` / ``_o_proj`` / ``backend_cls``) is provided by a
155 + subclass — ``DeepseekV4FlashMLAAttention`` /
156 + ``DeepseekV4FlashInferSM120Attention`` /
157 + ``DeepseekV4FlashInferMLAAttention`` (CUDA) or
158 + ``DeepseekV41ROCMAiterMLAAttention`` (ROCm) — selected by the platform-specific
159 + deepseek_v4_1 model module. The base is never instantiated directly.
160 + """
161 +
162 + # Provided by the platform subclass.
163 + backend_cls: ClassVar[type[AttentionBackend]]
164 + # Backend for the SWA cache layer; None uses the default SWA backend.
165 + swa_backend_cls: ClassVar[type[AttentionBackend] | None] = None
166 + # KV-cache per-token block format (both layouts are paged). True (default)
167 + # = fp8_ds_mla (UE8M0 block-scaled fp8 packed as uint8); False = plain
168 + # bf16 / per-tensor fp8 KV row. Backends can override the instance hook when
169 + # a single attention class dispatches across arch-specific layouts.
170 + use_fp8_ds_mla_layout: ClassVar[bool] = True
171 + # Prefill is processed in fixed-size chunks; this bounds the bf16 kv-gather
172 + # workspace allocated in _forward_prefill and is also read by the dummy-run
173 + # path to pre-reserve that workspace.
174 + PREFILL_CHUNK_SIZE: ClassVar[int] = 4
175 +
176 + @classmethod
177 + @abstractmethod
178 + def get_padded_num_q_heads(cls, num_heads: int) -> int:
179 + """Q head count the q/output buffers are allocated at.
180 +
181 + The layer allocates the q/output buffers at
182 + ``[N, get_padded_num_q_heads(n_local_heads), head_dim]``. Must satisfy
183 + ``result >= num_heads``. Backends with no padding constraint return
184 + ``num_heads``.
185 + """
186 + raise NotImplementedError
187 +
188 + @abstractmethod
189 + def forward_mqa(
190 + self,
191 + q: torch.Tensor,
192 + kv: torch.Tensor,
193 + positions: torch.Tensor,
194 + output: torch.Tensor,
195 + ) -> None:
196 + """Platform-specific sparse MLA forward; writes attention into ``output``."""
197 + raise NotImplementedError
198 +
199 + @abstractmethod
200 + def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
201 + """Inverse-RoPE + wo_a + wo_b output projection (platform-specific)."""
202 + raise NotImplementedError
203 +
204 + def _uses_fp8_ds_mla_layout(self) -> bool:
205 + """Return whether this instance stores fp8 KV in fp8_ds_mla layout."""
206 + return self.use_fp8_ds_mla_layout
207 +
208 + def __init__(
209 + self,
210 + vllm_config: VllmConfig,
211 + prefix: str,
212 + topk_indices_buffer: torch.Tensor | None = None,
213 + aux_stream_list: list[torch.cuda.Stream] | None = None,
214 + candidate_block_buffer: torch.Tensor | None = None,
215 + ) -> None:
216 + super().__init__()
217 + config = vllm_config.model_config.hf_config
218 + quant_config = vllm_config.quant_config
219 + cache_config = vllm_config.cache_config
220 + tp_size = get_tensor_model_parallel_world_size()
221 + layer_id = extract_layer_index(prefix)
222 + self.layer_id = layer_id
223 +
224 + self.prefix = prefix # Alias for compatibility with compressor
225 + self.hidden_size = config.hidden_size
226 + self.n_heads = config.num_attention_heads
227 + assert self.n_heads % tp_size == 0
228 + self.n_local_heads = self.n_heads // tp_size
229 + self.q_lora_rank = config.q_lora_rank
230 + self.o_lora_rank = config.o_lora_rank
231 + self.head_dim = config.head_dim
232 + self.rope_head_dim = config.qk_rope_head_dim
233 + self.nope_head_dim = self.head_dim - self.rope_head_dim
234 + self.n_groups = config.o_groups
235 + self.n_local_groups = self.n_groups // tp_size
236 + self.window_size = config.sliding_window
237 + # Vision variant: image spans are visible bidirectionally, widening
238 + # prefill SWA index rows by up to max_image_tokens columns.
239 + self.max_image_tokens = (
240 + getattr(config, "vision_max_n_token", 0)
241 + if getattr(config, "vision_n_layers", 0) > 0
242 + else 0
243 + )
244 + # ---- v4.1 sparse-attention topology ----
245 + # compress_ratios has one entry per layer (MTP layers included):
246 + # 0 = pure sliding window, 1 = full-length compressed cache,
247 + # 2 = ratio-2 compressed. Compressors and compressed-KV caches live
248 + # only on ``kv_source_layer_ids``; indexers only on
249 + # ``index_source_layer_ids``. Consumers reuse the most recently
250 + # published source below them.
251 + compress_ratios = getattr(config, "compress_ratios", None)
252 + if compress_ratios is not None and layer_id < len(compress_ratios):
253 + self.compress_ratio = int(compress_ratios[layer_id])
254 + else:
255 + # MTP layers past the configured list are pure sliding-window.
256 + self.compress_ratio = 0
257 + if self.compress_ratio not in (0, 1, 2):
258 + raise ValueError(
259 + f"DeepSeek V4.1 layer {layer_id} has compress_ratio="
260 + f"{self.compress_ratio}; only 0 (sliding window), 1 and 2 are "
261 + "supported."
262 + )
263 + self.kv_source_layers = tuple(
264 + getattr(config, "kv_source_layer_ids", None) or ()
265 + )
266 + self.index_source_layers = tuple(
267 + getattr(config, "index_source_layer_ids", None) or ()
268 + )
269 + self.candidate_source_layer = getattr(config, "candidate_source_layer_id", -1)
270 + self.candidate_topk_blocks = getattr(config, "candidate_topk_blocks", 0)
271 + self.candidate_block_size = getattr(config, "candidate_block_size", 0)
272 +
273 + is_backbone = layer_id < config.num_hidden_layers
274 + self.is_kv_source = is_backbone and layer_id in self.kv_source_layers
275 + self.is_index_source = is_backbone and layer_id in self.index_source_layers
276 + if self.compress_ratio > 0:
277 + if not self.kv_source_layers or not self.index_source_layers:
278 + raise ValueError(
279 + "DeepSeek V4.1 requires kv_source_layer_ids / "
280 + "index_source_layer_ids in the config for compressed "
281 + f"layers (layer {layer_id} has "
282 + f"compress_ratio={self.compress_ratio})."
283 + )
284 + self.kv_source_layer_id = max(
285 + s for s in self.kv_source_layers if s <= layer_id
286 + )
287 + self.index_source_layer_id = max(
288 + s for s in self.index_source_layers if s <= layer_id
289 + )
290 + else:
291 + self.kv_source_layer_id = None
292 + self.index_source_layer_id = None
293 + self.eps = config.rms_norm_eps
294 + self.scale = self.head_dim**-0.5
295 +
296 + # Padded Q head count is dictated by the platform subclass.
297 + self.padded_heads = self.get_padded_num_q_heads(self.n_local_heads)
298 + # Sink padded to the same head count, initialized to -inf (no sink
299 + # effect). Weight loading fills the first n_local_heads slots.
300 + self.attn_sink = nn.Parameter(
301 + torch.full((self.padded_heads,), -float("inf"), dtype=torch.float32),
302 + requires_grad=False,
303 + )
304 +
305 + self.fused_wqa_wkv = MergedColumnParallelLinear(
306 + self.hidden_size,
307 + [self.q_lora_rank, self.head_dim],
308 + bias=False,
309 + quant_config=quant_config,
310 + prefix=f"{prefix}.fused_wqa_wkv",
311 + disable_tp=True, # fused ReplicatedLinear
312 + )
313 + self.q_norm = RMSNorm(self.q_lora_rank, self.eps)
314 + self.wq_b = ColumnParallelLinear(
315 + self.q_lora_rank,
316 + self.n_heads * self.head_dim,
317 + bias=False,
318 + quant_config=quant_config,
319 + return_bias=False,
320 + prefix=f"{prefix}.wq_b",
321 + )
322 +
323 + self.kv_norm = RMSNorm(self.head_dim, self.eps)
324 + self.wo_a = ColumnParallelLinear(
325 + self.n_heads * self.head_dim // self.n_groups,
326 + self.n_groups * self.o_lora_rank,
327 + bias=False,
328 + quant_config=quant_config,
329 + return_bias=False,
330 + prefix=f"{prefix}.wo_a",
331 + )
332 + self.wo_a.is_bmm = True
333 + self.wo_a.bmm_batch_size = self.n_local_groups
334 + self._o_proj_block_size = (
335 + 32 if getattr(self.wo_a, "weight_block_size", None) == [1, 32] else 128
336 + )
337 + self.wo_b = RowParallelLinear(
338 + self.n_groups * self.o_lora_rank,
339 + self.hidden_size,
340 + bias=False,
341 + quant_config=quant_config,
342 + return_bias=False,
343 + prefix=f"{prefix}.wo_b",
344 + )
345 +
346 + # Initialize rotary embedding before the indexer/compressor consume it.
347 + self.rotary_emb = build_deepseek_v4_rope(
348 + config,
349 + head_dim=self.head_dim,
350 + rope_head_dim=self.rope_head_dim,
351 + max_position_embeddings=config.max_position_embeddings,
352 + compress_ratio=self.compress_ratio,
353 + )
354 + self.indexer_rotary_emb = self.rotary_emb
355 + self.topk_indices_buffer = topk_indices_buffer
356 + self.candidate_block_buffer = candidate_block_buffer
357 +
358 + # Register with compilation context for metadata lookup. Done before
359 + # indexer/compressor creation so consumers can resolve their source
360 + # layers through it.
361 + compilation_config = vllm_config.compilation_config
362 + if prefix and prefix in compilation_config.static_forward_context:
363 + raise ValueError(f"Duplicate layer name: {prefix}")
364 + if prefix:
365 + compilation_config.static_forward_context[prefix] = self
366 + self.kv_cache = torch.tensor([])
367 + self._static_forward_context = compilation_config.static_forward_context
368 +
369 + self.indexer = None
370 + if self.is_index_source:
371 + index_k_cache: DeepseekV4IndexerCache | None
372 + # Index K cache: owned only by kv-source layers (their indexer has
373 + # wk/k_norm); non-owning index sources share the K cache of the
374 + # latest kv source below them (its indexer produces the keys).
375 + if self.is_kv_source:
376 + index_k_cache = DeepseekV4IndexerCache(
377 + head_dim=_indexer_k_cache_head_dim(
378 + config.index_head_dim, dsa_indexer_uses_fp4(vllm_config)
379 + ),
380 + dtype=torch.uint8,
381 + prefix=f"{prefix}.indexer.k_cache",
382 + cache_config=cache_config,
383 + compress_ratio=self.compress_ratio,
384 + )
385 + else:
386 + assert self.kv_source_layer_id is not None
387 + k_cache_prefix = (
388 + f"{_replace_layer_index(prefix, self.kv_source_layer_id)}"
389 + ".indexer.k_cache"
390 + )
391 + index_k_cache = self._static_forward_context.get(k_cache_prefix)
392 + if index_k_cache is None:
393 + raise NotImplementedError(
394 + f"Indexer K cache source {k_cache_prefix} not found on "
395 + "this rank; PP splits inside a v4.1 kv-sharing group "
396 + "are not supported."
397 + )
398 + is_candidate_source = layer_id == self.candidate_source_layer
399 + uses_candidates = 0 <= self.candidate_source_layer < layer_id
400 + self.indexer = DeepseekV4Indexer(
401 + vllm_config,
402 + config=config,
403 + hidden_size=self.hidden_size,
404 + q_lora_rank=self.q_lora_rank,
405 + quant_config=quant_config,
406 + cache_config=cache_config,
407 + topk_indices_buffer=topk_indices_buffer,
408 + compress_ratio=self.compress_ratio,
409 + prefix=f"{prefix}.indexer",
410 + owns_k=self.is_kv_source,
411 + k_cache=index_k_cache,
412 + main_head_dim=self.head_dim,
413 + candidate_block_buffer=(
414 + candidate_block_buffer
415 + if (is_candidate_source or uses_candidates)
416 + else None
417 + ),
418 + candidate_block_size=self.candidate_block_size,
419 + candidate_write=is_candidate_source,
420 + )
421 +
422 + self._prepare_and_attn_fn = self._prepare_and_attn
423 + if not vllm_config.use_v2_model_runner:
424 + # MRV1's piecewise capture only tolerates the wide eager region: with
425 + # the narrow one the attention input preparation stays in the captured
426 + # graph and MRV1 produces garbage (#51430).
427 + self._prepare_and_attn_fn = self._prepare_and_attn_eager
428 +
429 + # Will be None on ROCm for now.
430 + self.aux_stream_list = aux_stream_list
431 + # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events;
432 + # [1] doubles as post-GEMM event1. Reuse is safe: GEMM fully joins
433 + # before post-GEMM starts.
434 + self.ln_events = [torch.cuda.Event() for _ in range(4)]
435 +
436 + assert cache_config is not None, "DeepseekV4 attention requires cache_config"
437 + # ---- Attention / KV-cache setup ----
438 + self.max_num_batched_tokens = (
439 + vllm_config.scheduler_config.max_num_batched_tokens
440 + )
441 + self.max_model_len = vllm_config.model_config.max_model_len
442 +
443 + # Resolve the kv-cache dtype from this backend's block format. The same
444 + # resolution drives the SWA cache tensor dtype below.
445 + self.kv_cache_dtype, self.kv_cache_torch_dtype = _resolve_dsv4_kv_cache_dtype(
446 + self._uses_fp8_ds_mla_layout(), cache_config.cache_dtype, cache_config
447 + )
448 +
449 + self.swa_cache_layer = DeepseekV4SWACache(
450 + head_dim=self.head_dim,
451 + window_size=self.window_size,
452 + dtype=self.kv_cache_torch_dtype,
453 + prefix=f"{prefix}.swa_cache",
454 + cache_config=cache_config,
455 + backend_cls=self.swa_backend_cls,
456 + block_size=32,
457 + )
458 +
459 + # The attention layer itself was already registered with the
460 + # compilation context above (before indexer/compressor creation).
461 +
462 + # Compressors live only on kv-source layers; consumers read the
463 + # source's compressed cache through the forward context.
464 + self.compressor = None
465 + if self.is_kv_source:
466 + self.compressor = DeepseekCompressor(
467 + vllm_config=vllm_config,
468 + compress_ratio=self.compress_ratio,
469 + hidden_size=self.hidden_size,
470 + head_dim=self.head_dim,
471 + rotate=True,
472 + prefix=f"{prefix}.compressor",
473 + k_cache_prefix=self.prefix,
474 + )
475 + # Prefix of the attention layer owning this layer's compressed KV
476 + # cache (self for kv sources).
477 + if self.compress_ratio > 0:
478 + assert self.kv_source_layer_id is not None
479 + self.compressed_cache_prefix: str | None = _replace_layer_index(
480 + prefix, self.kv_source_layer_id
481 + )
482 + if (
483 + not self.is_kv_source
484 + and self.compressed_cache_prefix not in self._static_forward_context
485 + ):
486 + raise NotImplementedError(
487 + f"Compressed-KV source {self.compressed_cache_prefix} not "
488 + "found on this rank; PP splits inside a v4.1 kv-sharing "
489 + "group are not supported."
490 + )
491 + else:
492 + self.compressed_cache_prefix = None
493 +
494 + if vllm_config.kernel_config.enable_jit_warmup:
495 + from vllm.v1.attention.backends.mla.sparse_swa import (
496 + _COMPUTE_PREFILL_METADATA_KERNEL,
497 + _COMPUTE_SWA_INDICES_AND_LENS_KERNEL,
498 + )
499 +
500 + _COMPUTE_PREFILL_METADATA_KERNEL.register_warmup()
501 + _COMPUTE_SWA_INDICES_AND_LENS_KERNEL.register_warmup(
502 + window_size=self.window_size,
503 + block_size=self.swa_cache_layer.block_size,
504 + max_image_tokens=self.max_image_tokens,
505 + )
506 +
507 + if self.compress_ratio > 1:
508 + from vllm.v1.attention.backends.mla.compressor_utils import (
509 + _COMPRESSED_SLOT_MAPPING_KERNEL,
510 + )
511 +
512 + _COMPRESSED_SLOT_MAPPING_KERNEL.register_warmup()
513 +
514 + if self.indexer is not None:
515 + from vllm.v1.attention.backends.mla.indexer import (
516 + _BUILD_PREFILL_CHUNK_METADATA_KERNEL,
517 + _PREPARE_UNIFORM_DECODE_KERNEL,
518 + )
519 +
520 + _PREPARE_UNIFORM_DECODE_KERNEL.register_warmup()
521 + _BUILD_PREFILL_CHUNK_METADATA_KERNEL.register_warmup()
522 +
523 + spec_config = vllm_config.speculative_config
524 + if spec_config is not None and spec_config.use_dspark():
525 + from vllm.v1.attention.backends.mla.sparse_swa import (
526 + _COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL,
527 + )
528 +
529 + _COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL.register_warmup(
530 + window_size=self.window_size,
531 + num_speculative_tokens=spec_config.num_speculative_tokens,
532 + block_size=self.swa_cache_layer.block_size,
533 + )
534 +
535 + if self.backend_cls.get_name() in (
536 + "FLASHMLA_SPARSE_DSV41",
537 + "ROCM_FLASHMLA_SPARSE_DSV4",
538 + ):
539 + from vllm.models.deepseek_v4_1.common.ops.cache_utils import (
540 + _COMBINE_TOPK_SWA_INDICES_KERNEL,
541 + )
542 +
543 + _COMBINE_TOPK_SWA_INDICES_KERNEL.register_warmup()
544 +
545 + def forward(
546 + self,
547 + positions: torch.Tensor,
548 + hidden_states: torch.Tensor,
549 + llama_4_scaling: torch.Tensor | None = None,
550 + ) -> torch.Tensor:
551 + # Pre-allocate attention output with FlashMLA-padded head count.
552 + # The op writes into `o_padded`; we slice to n_local_heads after.
553 + num_tokens = hidden_states.shape[0]
554 + o_padded = torch.empty(
555 + (num_tokens, self.padded_heads, self.head_dim),
556 + dtype=hidden_states.dtype,
557 + device=hidden_states.device,
558 + )
559 +
560 + # Keep the attention input preparation in the captured graph. Only the
561 + # sparse indexer and MLA attention run in the eager break below.
562 + qr_kv, kv_score, indexer_weights = self._run_parallel_input_projections(
563 + hidden_states
564 + )
565 + qr, qr_scale, kv = self._split_qkv_and_norm(qr_kv)
566 +
567 + self._prepare_and_attn_fn(
568 + hidden_states,
569 + qr,
570 + kv,
571 + qr_scale,
572 + kv_score,
573 + indexer_weights,
574 + positions,
575 + o_padded,
576 + )
577 + o = o_padded[:, : self.n_local_heads, :]
578 +
579 + # Inverse-RoPE + wo_a + wo_b output projection (platform-specific).
580 + return self._o_proj(o, positions)
581 +
582 + @cached_property
583 + def _can_fuse_query_quant(self) -> bool:
584 + from vllm.models.deepseek_v4_1.common.ops.query_quant import (
585 + can_fuse_query_quant,
586 + )
587 +
588 + linears = [self.wq_b]
589 + if self.indexer is not None:
590 + linears.append(self.indexer.wq_b)
591 + return can_fuse_query_quant(linears)
592 +
593 + def _split_qkv_and_norm(
594 + self, qr_kv: torch.Tensor
595 + ) -> tuple[torch.Tensor | QuantizedActivation, torch.Tensor | None, torch.Tensor]:
596 + """Split the fused q-lora / kv projection and RMSNorm both halves.
597 +
598 + Compatible MXFP8 projections share the quantized Q and scales;
599 + other projection backends consume the normalized Q directly.
600 + """
601 + qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1)
602 + if self.q_lora_rank % 32 == 0 and self._can_fuse_query_quant:
603 + from vllm.models.deepseek_v4_1.common.ops.query_quant import (
604 + fused_q_kv_rmsnorm_quant,
605 + )
606 +
607 + qr_quant, kv = fused_q_kv_rmsnorm_quant(
608 + qr,
609 + kv,
610 + self.q_norm.weight.data,
611 + self.kv_norm.weight.data,
612 + self.eps,
613 + )
614 + return qr_quant, None, kv
615 + qr, kv = fused_q_kv_rmsnorm(
616 + qr,
617 + kv,
618 + self.q_norm.weight.data,
619 + self.kv_norm.weight.data,
620 + self.eps,
621 + )
622 + return qr, None, kv
623 +
624 + @eager_break_during_capture
625 + def _prepare_and_attn_eager(
626 + self,
627 + hidden_states: torch.Tensor,
628 + qr: torch.Tensor | QuantizedActivation,
629 + kv: torch.Tensor,
630 + qr_scale: torch.Tensor | None,
631 + kv_score: torch.Tensor,
632 + indexer_weights: torch.Tensor,
633 + positions: torch.Tensor,
634 + o_padded: torch.Tensor,
635 + ) -> None:
636 + """Wide eager region: the whole of ``_prepare_and_attn`` runs eagerly.
637 +
638 + The nested ``_sparse_indexer_and_attn`` break runs inline, since
639 + ``add_eager`` clears ``_capturing`` before invoking this.
640 + """
641 + self._prepare_and_attn(
642 + hidden_states,
643 + qr,
644 + kv,
645 + qr_scale,
646 + kv_score,
647 + indexer_weights,
648 + positions,
649 + o_padded,
650 + )
651 +
652 + def _prepare_and_attn(
653 + self,
654 + hidden_states: torch.Tensor,
655 + qr: torch.Tensor | QuantizedActivation,
656 + kv: torch.Tensor,
657 + qr_scale: torch.Tensor | None,
658 + kv_score: torch.Tensor,
659 + indexer_weights: torch.Tensor,
660 + positions: torch.Tensor,
661 + o_padded: torch.Tensor,
662 + ) -> None:
663 + """Attention input preparation followed by the sparse indexer and MLA.
664 +
665 + Only the latter runs in the eager break.
666 +
667 + Q/SWA preparation overlaps state saving and compression. Once the
668 + latent is ready, main-cache insertion overlaps indexer preparation;
669 + both cache writes finish before sparse attention reads either cache.
670 + """
671 + attn_metadata = get_forward_context().attn_metadata
672 + indexer = self.indexer
673 + compressor = self.compressor
674 + aux_streams = self.aux_stream_list
675 +
676 + def project_query_and_cache_kv() -> torch.Tensor:
677 + q = self._wq_b_proj(qr, qr_scale).view(
678 + -1, self.n_local_heads, self.head_dim
679 + )
680 + return self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
681 +
682 + index_q: torch.Tensor | None = None
683 + index_q_scale: torch.Tensor | None = None
684 + index_weights_out: torch.Tensor | None = None
685 + latent: torch.Tensor | None = None
686 + aux_stream = aux_streams[0] if aux_streams is not None else None
687 +
688 + if compressor is not None:
689 + # Q projection / KV insertion on the default stream overlaps the
690 + # compressor on aux stream 0 (sequential on ROCm).
691 + q, latent = maybe_execute_in_parallel(
692 + project_query_and_cache_kv,
693 + lambda: compressor(kv_score, positions),
694 + self.ln_events[0],
695 + self.ln_events[1],
696 + aux_stream,
697 + )
698 + else:
699 + q = project_query_and_cache_kv()
700 +
701 + def prepare_indexer():
702 + if indexer is None:
703 + return None, None, None
704 + return indexer(
705 + qr,
706 + latent,
707 + indexer_weights,
708 + positions,
709 + self.indexer_rotary_emb,
710 + qr_scale,
711 + )
712 +
713 + if compressor is not None:
714 + indexer_result, _ = maybe_execute_in_parallel(
715 + prepare_indexer,
716 + lambda: compressor.insert_cache(latent, positions, self.rotary_emb),
717 + self.ln_events[0],
718 + self.ln_events[1],
719 + aux_stream,
720 + )
721 + else:
722 + indexer_result = prepare_indexer()
723 + index_q, index_q_scale, index_weights_out = indexer_result
724 +
725 + self._sparse_indexer_and_attn(
726 + hidden_states,
727 + index_q,
728 + index_q_scale,
729 + index_weights_out,
730 + q,
731 + kv,
732 + positions,
733 + o_padded,
734 + )
735 +
736 + def _fused_wqa_wkv_gemm(self, hidden_states: torch.Tensor) -> torch.Tensor:
737 + # Override point: the ROCm layer preshuffles this weight in place, so
738 + # it cannot go through fused_wqa_wkv directly.
739 + # MergedColumnParallelLinear returns (output, bias); bias is None.
740 + qr_kv, _ = self.fused_wqa_wkv(hidden_states)
741 + return qr_kv
742 +
743 + def _wq_b_proj(
744 + self,
745 + qr: torch.Tensor | QuantizedActivation,
746 + qr_scale: torch.Tensor | None = None,
747 + ) -> torch.Tensor:
748 + """Project normalized Q, bypassing quantization when already fused."""
749 + assert qr_scale is None, "ROCm-only path"
750 + return self.wq_b(qr)
751 +
752 + def _run_parallel_input_projections(
753 + self, hidden_states: torch.Tensor
754 + ) -> tuple[
755 + torch.Tensor,
756 + torch.Tensor | None,
757 + torch.Tensor | None,
758 + ]:
759 + aux_streams = self.aux_stream_list
760 + if aux_streams is not None:
761 + aux_streams = aux_streams[:2]
762 +
763 + # fused_wqa_wkv (heaviest) on default; the two lighter input GEMMs on
764 + # aux streams 0/1 when their owning module exists. ln_events[0] is the
765 + # fan-out start event; ln_events[1..2] are per-aux done events. The
766 + # v4.1 indexer derives K from the kv-source compressor's latent, so
767 + # unlike v4.0 there is no indexer K GEMM over hidden_states here.
768 + aux_fns: list[Callable[[], Any] | None] = [None, None]
769 +
770 + if self.compressor is not None:
771 + # Local ref so the closure keeps a non-None type for mypy.
772 + compressor = self.compressor
773 +
774 + def compressor_kv_score() -> torch.Tensor:
775 + return torch.mm(
776 + hidden_states,
777 + compressor.fused_wkv_wgate.weight.T,
778 + out_dtype=torch.float32,
779 + )
780 +
781 + aux_fns[0] = compressor_kv_score
782 +
783 + if self.indexer is not None:
784 + indexer = self.indexer
785 +
786 + def indexer_weights_proj() -> torch.Tensor:
787 + # ReplicatedLinear returns (output, bias); bias is None.
788 + weights, _ = indexer.weights_proj(hidden_states)
789 + return weights
790 +
791 + aux_fns[1] = indexer_weights_proj
792 +
793 + qr_kv, (kv_score, indexer_weights) = execute_in_parallel(
794 + lambda: self._fused_wqa_wkv_gemm(hidden_states),
795 + aux_fns,
796 + self.ln_events[0],
797 + self.ln_events[1:3],
798 + aux_streams,
799 + enable=hidden_states.shape[0]
800 + <= envs.VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD,
801 + )
802 +
803 + return qr_kv, kv_score, indexer_weights
804 +
805 + @eager_break_during_capture
806 + def _sparse_indexer_and_attn(
807 + self,
808 + hidden_states: torch.Tensor,
809 + index_q: torch.Tensor | None,
810 + index_q_scale: torch.Tensor | None,
811 + index_weights: torch.Tensor | None,
812 + q: torch.Tensor,
813 + kv: torch.Tensor,
814 + positions: torch.Tensor,
815 + out: torch.Tensor,
816 + ) -> None:
817 + if self.indexer is not None and index_q is not None:
818 + assert index_weights is not None
819 + q_quant = (index_q, index_q_scale) if index_q_scale is not None else index_q
820 + self.indexer.indexer_op(
821 + hidden_states,
822 + q_quant,
823 + None,
824 + index_weights,
825 + )
826 +
827 + # MLA attention writes into the pre-allocated `out` buffer
828 + # ([num_tokens, padded_heads, head_dim]).
829 + self.forward_mqa(q, kv, positions, out)
830 +
831 + def _fused_qnorm_rope_kv_insert(
832 + self,
833 + q: torch.Tensor,
834 + kv: torch.Tensor,
835 + positions: torch.Tensor,
836 + attn_metadata: (
837 + dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None
838 + ),
839 + ) -> torch.Tensor:
840 + if not isinstance(attn_metadata, dict):
841 + # Profile run: kernel doesn't fire; produce a padded tensor so
842 + # downstream FlashMLA gets the right shape.
843 + if self.n_local_heads < self.padded_heads:
844 + return F.pad(
845 + q,
846 + (0, 0, 0, self.padded_heads - self.n_local_heads),
847 + value=0.0,
848 + )
849 + return q
850 +
851 + swa_metadata = cast(
852 + "DeepseekSparseSWAMetadata | None",
853 + attn_metadata.get(self.swa_cache_layer.prefix),
854 + )
855 + assert swa_metadata is not None
856 +
857 + swa_kv_cache = self.swa_cache_layer.kv_cache
858 + # The fused insert ops require int64 position_ids; the runner's positions
859 + # buffer is already int64, so no cast is needed.
860 + assert positions.dtype == torch.int64
861 + cos_sin_cache = self.rotary_emb.cos_sin_cache
862 + cache_dtype = swa_kv_cache.dtype
863 +
864 + # kv is unchanged; attention reads kv solely via swa_kv_cache.
865 + if cache_dtype == torch.uint8:
866 + # fp8_ds_mla UE8M0 paged path. Horizontally fused:
867 + # Q side: GPT-J RoPE, zero-filling the padding head slots; the
868 + # kernel allocates and returns the padded q tensor.
869 + # KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert.
870 + swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1)
871 + return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
872 + q,
873 + kv,
874 + swa_kv_cache_2d,
875 + swa_metadata.slot_mapping,
876 + positions,
877 + cos_sin_cache,
878 + self.padded_heads,
879 + self.eps,
880 + swa_metadata.block_size,
881 + False,
882 + )
883 +
884 + # Plain-row path: the [num_blocks, block_size, 512] cache stores the KV
885 + # row in its element dtype (no Q padding). bf16 rewrites q in place;
886 + # per-tensor fp8 writes a separately-allocated fp8 q and quantizes the
887 + # KV row.
888 + block_size = swa_metadata.block_size
889 + assert swa_kv_cache.shape[1:] == (block_size, self.head_dim)
890 + swa_kv_cache_3d = swa_kv_cache
891 + if cache_dtype == torch.bfloat16:
892 + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
893 + q,
894 + kv,
895 + swa_kv_cache_3d,
896 + swa_metadata.slot_mapping,
897 + positions,
898 + cos_sin_cache,
899 + self.eps,
900 + block_size,
901 + False,
902 + )
903 + return q
904 +
905 + # per-tensor fp8 (torch.float8_e4m3fn)
906 + q_fp8 = torch.empty_like(q, dtype=torch.float8_e4m3fn)
907 + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
908 + q,
909 + kv,
910 + q_fp8,
911 + swa_kv_cache_3d,
912 + swa_metadata.slot_mapping,
913 + positions,
914 + cos_sin_cache,
915 + self._flashinfer_fp8_kv_scale,
916 + self._flashinfer_fp8_q_scale_inv,
917 + self.eps,
918 + block_size,
919 + False,
920 + )
921 + return q_fp8
922 +
923 + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
924 + # [B, H=1, N, C] -> [B, N, C]
925 + self.kv_cache = kv_cache.squeeze(1)
926 +
927 + def get_attn_backend(self) -> type[AttentionBackend]:
928 + return self.backend_cls
929 +
930 + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
931 + # Only kv-source layers own a compressed-KV cache; consumers read the
932 + # source's cache through the forward context, and cr==0 layers are
933 + # pure SWA. The SWA cache is allocated separately as
934 + # DeepseekV4SWACache.
935 + if not self.is_kv_source:
936 + return None
937 + # fp8_ds_mla is a UE8M0 block-scaled uint8 layout and needs 576B
938 + # alignment; plain bf16 / per-tensor fp8 rows use natural element-size
939 + # pages.
940 + uses_fp8_ds_mla_layout = self.kv_cache_dtype == "fp8_ds_mla"
941 + return MLAAttentionSpec(
942 + block_size=vllm_config.cache_config.block_size,
943 + num_kv_heads=1,
944 + head_size=self.head_dim,
945 + dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype,
946 + tokens_per_state=self.compress_ratio,
947 + cache_dtype_str=self.kv_cache_dtype,
948 + alignment=576 if uses_fp8_ds_mla_layout else 512,
949 + model_version="deepseek_v4",
950 + kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
951 + # DeepseekV4: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B per token;
952 + # head_size stays semantic (512).
953 + state_content_bytes=584 if uses_fp8_ds_mla_layout else None,
954 + )
955 +
956 + def _compressed_kv_cache(self) -> torch.Tensor:
957 + """The compressed-KV cache tensor of this layer's kv source (own
958 + cache for kv-source layers)."""
959 + if self.is_kv_source:
960 + return self.kv_cache
961 + assert self.compressed_cache_prefix is not None
962 + source = self._static_forward_context[self.compressed_cache_prefix]
963 + return source.kv_cache
964 +
965 +
966 +class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase):
967 + def __init__(
968 + self,
969 + head_dim: int,
970 + dtype: torch.dtype,
971 + prefix: str,
972 + cache_config: CacheConfig,
973 + compress_ratio: int = 1,
974 + ):
975 + super().__init__()
976 + self.kv_cache = torch.tensor([])
977 + self.head_dim = head_dim
978 + self.prefix = prefix
979 + self.cache_config = cache_config
980 + self.dtype = dtype
981 + self.compress_ratio = compress_ratio
982 + compilation_config = get_current_vllm_config().compilation_config
983 + if prefix in compilation_config.static_forward_context:
984 + raise ValueError(f"Duplicate layer name: {prefix}")
985 + compilation_config.static_forward_context[prefix] = self
986 +
987 + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
988 + # [B, H=1, N, C] -> [B, N, C]
989 + self.kv_cache = kv_cache.squeeze(1)
990 +
991 + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
992 + # head_dim already carries the fp8 scale padding
993 + # tokens_per_state=1 for V3.2, >1 for DeepseekV4; same cache layout.
994 + uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype == "fp8_ds_mla"
995 + return MLAAttentionSpec(
996 + block_size=self.cache_config.block_size,
997 + num_kv_heads=1,
998 + head_size=self.head_dim,
999 + dtype=self.dtype,
1000 + tokens_per_state=self.compress_ratio,
1001 + # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577).
1002 + alignment=576 if uses_fp8_ds_mla_layout else 512,
1003 + )
1004 +
1005 + def forward(self): ...
1006 +
1007 + def get_attn_backend(self) -> type[AttentionBackend]:
1008 + return DeepseekV4IndexerBackend
1009 +
1010 +
1011 +class DeepseekV4Indexer(nn.Module):
1012 + """DeepSeek V4.1 sparse-attention indexer.
1013 +
1014 + Exists only on ``index_source_layer_ids``; consumers reuse the topk
1015 + indices it publishes into the shared ``topk_indices_buffer``. Unlike v4.0
1016 + the index key is derived from the kv-source layer's compressor latent
1017 + (``k = k_norm(wk(latent))``, ``owns_k``) instead of an indexer-local
1018 + compressor over hidden states, so there is no hidden-state K GEMM here.
1019 + Non-owning index sources share the kv source's paged K cache.
1020 +
1021 + Two-level candidate filtering: the indexer at ``candidate_source_layer_id``
1022 + additionally publishes the top ``candidate_topk_blocks`` blocks of
1023 + ``candidate_block_size`` compressed positions (``candidate_write``);
1024 + later indexers mask their scores to those blocks before their own top-k.
1025 + """
1026 +
1027 + def __init__(
1028 + self,
1029 + vllm_config: VllmConfig,
1030 + config: DeepseekV2Config | DeepseekV3Config,
1031 + hidden_size: int,
1032 + q_lora_rank: int,
1033 + quant_config: QuantizationConfig | None,
1034 + cache_config: CacheConfig | None,
1035 + topk_indices_buffer: torch.Tensor | None,
1036 + compress_ratio: int,
1037 + prefix: str,
1038 + *,
1039 + owns_k: bool,
1040 + k_cache: DeepseekV4IndexerCache,
1041 + main_head_dim: int,
1042 + candidate_block_buffer: torch.Tensor | None = None,
1043 + candidate_block_size: int = 0,
1044 + candidate_write: bool = False,
1045 + ):
1046 + super().__init__()
1047 + self.vllm_config = vllm_config
1048 + self.config = config
1049 + self.quant_config = quant_config
1050 + self.topk_tokens = config.index_topk
1051 + self.n_head = config.index_n_heads # 32
1052 + self.head_dim = config.index_head_dim # 128
1053 + self.rope_dim = config.qk_rope_head_dim # 64
1054 + self.q_lora_rank = q_lora_rank # 1280
1055 + self.compress_ratio = compress_ratio
1056 + self.owns_k = owns_k
1057 + self.use_fp4_kv = dsa_indexer_uses_fp4(vllm_config)
1058 + logger.info_once(
1059 + "Using %s indexer cache for Lightning Indexer.",
1060 + "MXFP4" if self.use_fp4_kv else "FP8",
1061 + )
1062 +
1063 + # no tensor parallel, just replicated
1064 + self.wq_b = ReplicatedLinear(
1065 + self.q_lora_rank,
1066 + self.head_dim * self.n_head,
1067 + bias=False,
1068 + quant_config=quant_config,
1069 + prefix=f"{prefix}.wq_b",
1070 + )
1071 + self.weights_proj = ReplicatedLinear(
1072 + hidden_size,
1073 + self.n_head,
1074 + bias=False,
1075 + quant_config=None,
1076 + prefix=f"{prefix}.weights_proj",
1077 + )
1078 + self.softmax_scale = self.head_dim**-0.5
1079 +
1080 + self.scale_fmt = "ue8m0"
1081 + self.quant_block_size = 128 # TODO: get from config
1082 + self.topk_indices_buffer = topk_indices_buffer
1083 +
1084 + self.max_model_len = (
1085 + vllm_config.model_config.max_model_len // self.compress_ratio
1086 + )
1087 + self.prefix = prefix
1088 +
1089 + self.max_total_seq_len = (
1090 + get_max_prefill_buffer_size(vllm_config) // self.compress_ratio
1091 + )
1092 +
1093 + assert cache_config is not None, "Deepseek V4 indexer requires cache_config"
1094 + if owns_k:
1095 + # wk maps the main compressor's (pre-RoPE) latent to index keys.
1096 + # The checkpoint stores it in bf16 with no quantization scales.
1097 + self.wk = ReplicatedLinear(
1098 + main_head_dim,
1099 + self.head_dim,
1100 + bias=False,
1101 + quant_config=None,
1102 + prefix=f"{prefix}.wk",
1103 + )
1104 + self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps)
1105 + self.k_cache = k_cache
1106 +
1107 + self.indexer_op = SparseAttnIndexer(
1108 + self.k_cache,
1109 + self.quant_block_size,
1110 + self.scale_fmt,
1111 + self.topk_tokens,
1112 + self.head_dim,
1113 + self.max_model_len,
1114 + self.max_total_seq_len,
1115 + self.topk_indices_buffer,
1116 + skip_k_cache_insert=True,
1117 + use_fp4_cache=self.use_fp4_kv,
1118 + compress_ratio=self.compress_ratio,
1119 + candidate_blocks=candidate_block_buffer,
1120 + candidate_block_size=candidate_block_size,
1121 + candidate_write=candidate_write,
1122 + )
1123 +
1124 + def _produce_k(
1125 + self,
1126 + latent: torch.Tensor | None,
1127 + positions: torch.Tensor,
1128 + rotary_emb: nn.Module,
1129 + ) -> None:
1130 + """Turn the compressor's pre-RoPE latent into paged index keys.
1131 +
1132 + ``k_norm(wk(latent))`` at group-boundary tokens, RoPE'd at the group
1133 + position and MXFP4/FP8-quantized into the indexer K cache.
1134 + """
1135 + attn_metadata = get_forward_context().attn_metadata
1136 + if not isinstance(attn_metadata, dict) or latent is None:
1137 + # Profile run: the indexer K cache is not bound yet and the
1138 + # compressor skipped its latent output.
1139 + return
1140 + assert self.owns_k
1141 + indexer_metadata = cast(Any, attn_metadata[self.k_cache.prefix])
1142 + # ReplicatedLinear returns (output, bias); bias is None. Rows at
1143 + # non-boundary tokens hold garbage latent and are skipped by the
1144 + # store kernel.
1145 + k_pre, _ = self.wk(latent)
1146 + indexer_k_norm_rope_store(
1147 + k_pre,
1148 + positions,
1149 + rotary_emb.cos_sin_cache,
1150 + self.k_norm.weight,
1151 + self.k_norm.variance_epsilon,
1152 + self.k_cache.kv_cache,
1153 + indexer_metadata.slot_mapping,
1154 + self.compress_ratio,
1155 + self.use_fp4_kv,
1156 + )
1157 +
1158 + def forward(
1159 + self,
1160 + qr: torch.Tensor | QuantizedActivation,
1161 + latent: torch.Tensor | None,
1162 + indexer_weights: torch.Tensor,
1163 + positions: torch.Tensor,
1164 + rotary_emb: nn.Module,
1165 + qr_scale: torch.Tensor | None = None,
1166 + ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
1167 + attn_metadata = get_forward_context().attn_metadata
1168 + if isinstance(attn_metadata, dict):
1169 + indexer_metadata = cast(Any, attn_metadata[self.k_cache.prefix])
1170 + if (
1171 + indexer_metadata.max_seq_len // self.compress_ratio <= self.topk_tokens
1172 + and not torch.cuda.is_current_stream_capturing()
1173 + ):
1174 + # candidates num smaller than topk, every candidate is selected
1175 + # but we still need to build k cache
1176 + if self.owns_k:
1177 + self._produce_k(latent, positions, rotary_emb)
1178 + assert self.topk_indices_buffer is not None
1179 + num_tokens = (
1180 + indexer_metadata.num_decode_tokens
1181 + + indexer_metadata.num_prefill_tokens
1182 + )
1183 + if num_tokens > 0:
1184 + _fill_short_context_topk_indices[(num_tokens,)](
1185 + self.topk_indices_buffer,
1186 + positions,
1187 + TOP_K=self.topk_tokens,
1188 + COMPRESS_RATIO=self.compress_ratio,
1189 + PADDED_TOP_K=triton.next_power_of_2(self.topk_tokens),
1190 + num_warps=8,
1191 + )
1192 + return None, None, None
1193 +
1194 + if self.owns_k:
1195 + # K write must land before indexer_op reads the cache
1196 + # (skip_k_cache_insert=True).
1197 + self._produce_k(latent, positions, rotary_emb)
1198 +
1199 + q = self._wq_b_proj(qr, qr_scale)
1200 + q = q.view(-1, self.n_head, self.head_dim)
1201 + q_quant, weights = fused_indexer_q_rope_quant(
1202 + positions,
1203 + q,
1204 + rotary_emb.cos_sin_cache,
1205 + indexer_weights,
1206 + self.softmax_scale,
1207 + self.n_head**-0.5,
1208 + use_fp4=self.use_fp4_kv,
1209 + )
1210 + if isinstance(q_quant, tuple):
1211 + q, q_scale = q_quant
1212 + else:
1213 + q, q_scale = q_quant, None
1214 + return q, q_scale, weights
1215 +
1216 + def _wq_b_proj(
1217 + self,
1218 + qr: torch.Tensor | QuantizedActivation,
1219 + qr_scale: torch.Tensor | None = None,
1220 + ) -> torch.Tensor:
1221 + """Project normalized Q with the main attention's shared quantization."""
1222 + assert qr_scale is None, "ROCm-only path"
1223 + # ReplicatedLinear returns (output, bias); bias is None.
1224 + q, _ = self.wq_b(qr)
1225 + return q
vllm/models/deepseek_v4_1/common/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/deepseek_v4_1/common/engram.py new
+1033
@@ -0,0 +1,1033 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Engram: n-gram hash lookups gated into the hyper-connection stream.
4 +
5 +Port of the reference ``inference/engram.py`` + ``Engram`` /
6 +``ParallelEngramEmbedding`` from ``inference/model.py`` (DeepSeek V4.1
7 +checkpoint layout). Engram modules live on the backbone layers listed in
8 +``engram_layer_ids`` only.
9 +
10 +Two pieces of cross-forward state are needed because vLLM streams tokens
11 +chunk-by-chunk while an n-gram at position ``p`` needs the token ids at
12 +``p-1..p-3``:
13 +
14 +- ``token_map``: token id -> compressed vocab id, built once from the
15 + model's tokenizer at init (deterministic; asserted against
16 + ``engram_compressed_vocab_size``).
17 +- ``hash_cache``: one int32 slot per KV slot of the first local layer's
18 + sliding-window cache, holding the compressed id (or DEAD) of the token
19 + last written to that slot. Slots are stable per (request, position) —
20 + the block table pins a position to a physical slot, prefix-cache hits
21 + reuse both the physical blocks and the identical token ids, and
22 + spec-decode rollbacks rewrite the same slots — so lookbacks read back
23 + exactly what the owning request wrote. Lookback depth (3) is far inside
24 + the sliding window (128), so window eviction never frees a block a
25 + live lookback still needs.
26 +
27 + Slots are not part of the KV cache, so KV loaded from another instance
28 + (P/D, offload connectors) leaves them unwritten. The runner therefore
29 + passes ``lookback_token_ids``, the ids just before each request's chunk
30 + start, which take precedence over the slots. The V2 runner reads them
31 + from its device-resident token history and needs no slot cache; the V1
32 + runner's CPU token table holds placeholders for generated tokens under
33 + async scheduling, so it passes prompt positions only and keeps the slot
34 + cache for the rest.
35 +"""
36 +
37 +import weakref
38 +
39 +import numpy as np
40 +import torch
41 +from torch import nn
42 +
43 +from vllm.config import VllmConfig, get_current_vllm_config
44 +from vllm.distributed import (
45 + get_tensor_model_parallel_rank,
46 + get_tensor_model_parallel_world_size,
47 + tensor_model_parallel_all_gather,
48 +)
49 +from vllm.logger import init_logger
50 +from vllm.model_executor.layers.linear import ReplicatedLinear
51 +from vllm.model_executor.layers.quantization import QuantizationConfig
52 +from vllm.model_executor.utils import set_weight_attrs
53 +from vllm.triton_utils import tl, triton
54 +from vllm.utils.platform_utils import is_uva_available
55 +from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
56 +
57 +logger = init_logger(__name__)
58 +
59 +# Cache value for tokens that take no part in an n-gram (image spans).
60 +DEAD_ID = -1
61 +
62 +
63 +def _is_prime(n: int) -> bool:
64 + """Deterministic Miller-Rabin for n < 2**32 (avoids a sympy import)."""
65 + if n < 2:
66 + return False
67 + for p in (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37):
68 + if n % p == 0:
69 + return n == p
70 + d = n - 1
71 + r = 0
72 + while d % 2 == 0:
73 + d //= 2
74 + r += 1
75 + for a in (2, 7, 61):
76 + x = pow(a, d, n)
77 + if x in (1, n - 1):
78 + continue
79 + for _ in range(r - 1):
80 + x = x * x % n
81 + if x == n - 1:
82 + break
83 + else:
84 + return False
85 + return True
86 +
87 +
88 +def find_next_prime(start: int, seen_primes: set[int]) -> int:
89 + """The smallest prime above `start` that has not been handed out yet."""
90 + candidate = start + 1
91 + while not _is_prime(candidate) or candidate in seen_primes:
92 + candidate += 1
93 + return candidate
94 +
95 +
96 +def build_compressed_token_map(tokenizer) -> tuple[list[int], int]:
97 + """Map every token id onto a smaller id space where tokens that normalize
98 + alike collapse together.
99 +
100 + N-grams are hashed over these compressed ids, so " The", "the" and "THE"
101 + all hash the same way. The compressed size matters beyond bounds checking:
102 + every hash multiplier is derived from it.
103 + """
104 + from tokenizers import Regex, normalizers
105 +
106 + # A private-use char, so a token that is exactly one space survives
107 + # Strip() instead of collapsing to the empty string and merging with
108 + # unrelated tokens.
109 + sentinel = "\ue000"
110 + normalizer = normalizers.Sequence(
111 + [
112 + normalizers.NFKC(),
113 + normalizers.NFD(),
114 + normalizers.StripAccents(),
115 + normalizers.Lowercase(),
116 + normalizers.Replace(Regex(r"[ \t\r\n]+"), " "),
117 + normalizers.Replace(Regex(r"^ $"), sentinel),
118 + normalizers.Strip(),
119 + normalizers.Replace(sentinel, " "),
120 + ]
121 + )
122 +
123 + # The raw Rust tokenizer, matching what training decodes with
124 + # (no clean_up_tokenization_spaces).
125 + backend = tokenizer.backend_tokenizer
126 + key_to_new: dict[str, int] = {}
127 + lookup = [0] * len(tokenizer)
128 + for token_id in range(len(tokenizer)):
129 + text = backend.decode([token_id], skip_special_tokens=False)
130 + if "\ufffd" in text:
131 + # A partial UTF-8 byte token: nothing to normalize, so key it
132 + # by its raw form.
133 + key = backend.id_to_token(token_id)
134 + else:
135 + normalized = normalizer.normalize_str(text)
136 + key = normalized if normalized else text
137 +
138 + new_id = key_to_new.get(key)
139 + if new_id is None:
140 + new_id = len(key_to_new)
141 + key_to_new[key] = new_id
142 + lookup[token_id] = new_id
143 +
144 + return lookup, len(key_to_new)
145 +
146 +
147 +def compute_hash_multipliers(
148 + layer_ids: tuple[int, ...], max_ngram_size: int, compressed_vocab_size: int
149 +) -> torch.Tensor:
150 + """One multiplier per (layer, lookback), from a per-layer RNG so layers
151 + hash differently. Kept odd and bounded so `token_id * multiplier` cannot
152 + overflow int64.
153 + """
154 + max_long = np.iinfo(np.int64).max
155 + multiplier_bound = max(1, (max_long // compressed_vocab_size) // 2)
156 + rows = []
157 + for layer_id in layer_ids:
158 + generator = np.random.default_rng(10007 * layer_id)
159 + values = generator.integers(
160 + low=0,
161 + high=multiplier_bound,
162 + size=(max_ngram_size,),
163 + dtype=np.int64,
164 + )
165 + rows.append(torch.tensor(values * 2 + 1))
166 + return torch.stack(rows)
167 +
168 +
169 +class EngramLayout:
170 + """Bucket layout of the n-gram hash tables.
171 +
172 + A position is hashed as `max_ngram_size - 1` n-grams (2-gram .. max), each
173 + split over `n_heads` heads. Every (n-gram size, head) pair owns its own
174 + prime-sized bucket range in the layer's table; the primes are drawn in
175 + order and never reused, which keeps the ranges disjoint.
176 + """
177 +
178 + def __init__(self, config) -> None:
179 + self.layer_ids: tuple[int, ...] = tuple(config.engram_layer_ids)
180 + self.num_embeddings: tuple[int, ...] = tuple(config.engram_num_embeddings)
181 + self.max_ngram_size: int = config.engram_max_ngram_size
182 + self.n_heads: int = config.engram_n_heads
183 + self.head_dim: int = config.engram_head_dim
184 + self.compressed_vocab_size: int = config.engram_compressed_vocab_size
185 + self.pad_token_id: int = config.engram_pad_token_id
186 + assert len(self.layer_ids) == len(self.num_embeddings)
187 +
188 + primes = []
189 + seen: set[int] = set()
190 + for _ in self.layer_ids:
191 + per_ngram = []
192 + for _ in range(self.max_ngram_size - 1):
193 + sizes, current = [], config.engram_vocab_size - 1
194 + for _ in range(self.n_heads):
195 + current = find_next_prime(current, seen)
196 + seen.add(current)
197 + sizes.append(current)
198 + per_ngram.append(tuple(sizes))
199 + primes.append(tuple(per_ngram))
200 + self.primes: tuple[tuple[tuple[int, ...], ...], ...] = tuple(primes)
201 + self.n_hash_cols = (self.max_ngram_size - 1) * self.n_heads
202 + flat = [[p for per_ngram in layer for p in per_ngram] for layer in primes]
203 + offsets = [np.cumsum([0, *sizes[:-1]]) for sizes in flat]
204 + self.offsets = torch.tensor(np.array(offsets)) # [n_layers, n_hash_cols]
205 +
206 + @classmethod
207 + def from_config(cls, config) -> "EngramLayout | None":
208 + if not getattr(config, "engram_layer_ids", None):
209 + return None
210 + return cls(config)
211 +
212 +
213 +@triton.jit(do_not_specialize=["num_tokens"])
214 +def _write_hash_cache_kernel(
215 + input_ids,
216 + token_map,
217 + dead_mask,
218 + slot_mapping,
219 + cache,
220 + num_tokens,
221 + input_stride,
222 + mask_stride,
223 + slot_stride,
224 + BLOCK_SIZE: tl.constexpr,
225 + dead_id,
226 +):
227 + token_idx = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
228 + slot = tl.load(
229 + slot_mapping + token_idx * slot_stride, token_idx < num_tokens, other=-1
230 + ).to(tl.int64)
231 + valid = (token_idx < num_tokens) & (slot >= 0)
232 + token = tl.load(input_ids + token_idx * input_stride, valid, other=0)
233 + value = tl.load(token_map + token, valid, other=0)
234 + dead = tl.load(dead_mask + token_idx * mask_stride, valid, other=False)
235 + value = tl.where(dead, dead_id, value)
236 + tl.store(cache + slot, value, valid)
237 +
238 +
239 +@triton.jit(
240 + do_not_specialize=[
241 + "num_tokens",
242 + "num_slots",
243 + "num_query_rows",
244 + "num_table_rows",
245 + "max_blocks",
246 + ]
247 +)
248 +def _hash_ids_kernel(
249 + input_ids,
250 + token_map,
251 + dead_mask,
252 + positions,
253 + block_table,
254 + query_start_loc,
255 + multipliers,
256 + primes,
257 + offsets,
258 + cache,
259 + lookback_token_ids,
260 + lookback_dead_mask,
261 + output,
262 + num_tokens,
263 + num_slots,
264 + pad_id,
265 + input_stride,
266 + mask_stride,
267 + position_stride,
268 + table_stride,
269 + table_col_stride,
270 + query_stride,
271 + num_query_rows,
272 + num_table_rows,
273 + max_blocks,
274 + cache_block_size,
275 + MAX_NGRAM: tl.constexpr,
276 + num_heads,
277 + BLOCK_T: tl.constexpr,
278 + BLOCK_H: tl.constexpr,
279 + dead_id,
280 + lookback_depth,
281 + lookback_row_stride,
282 + lookback_col_stride,
283 + lookback_mask_row_stride,
284 + lookback_mask_col_stride,
285 +):
286 + token = tl.program_id(0) * BLOCK_T + tl.arange(0, BLOCK_T)
287 + layer = tl.program_id(1)
288 + num_layers = tl.num_programs(1)
289 + valid = token < num_tokens
290 + # Upper bound in query_start_loc[1:], including repeated padding boundaries.
291 + lo = tl.full((BLOCK_T,), 0, tl.int32)
292 + hi = tl.full((BLOCK_T,), num_query_rows, tl.int32)
293 + while tl.sum((lo < hi).to(tl.int32), 0) > 0:
294 + mid = (lo + hi) // 2
295 + end = tl.load(
296 + query_start_loc + (mid + 1) * query_stride,
297 + lo < hi,
298 + other=0,
299 + )
300 + right = token >= end
301 + active = lo < hi
302 + lo = tl.where(active & right, mid + 1, lo)
303 + hi = tl.where(active & ~right, mid, hi)
304 + req = tl.minimum(lo, num_query_rows - 1).to(tl.int64)
305 + chunk_idx = tl.load(query_start_loc + req * query_stride)
306 + chunk_idx = tl.minimum(chunk_idx, num_tokens - 1).to(tl.int64)
307 + chunk_start = tl.load(positions + chunk_idx * position_stride)
308 + position = tl.load(positions + token * position_stride, valid, other=0).to(tl.int64)
309 + head = tl.arange(0, BLOCK_H)
310 + blocked = tl.full((BLOCK_T,), False, tl.int1)
311 + rolling = tl.full((BLOCK_T,), 0, tl.int64)
312 + for shift in tl.static_range(MAX_NGRAM):
313 + lookback = position - shift
314 + in_batch = lookback >= chunk_start
315 + batch_idx = tl.maximum(token - shift, 0)
316 + batch_token = tl.load(
317 + input_ids + batch_idx * input_stride, valid & in_batch, other=0
318 + )
319 + batch_source = tl.load(token_map + batch_token, valid & in_batch, other=0)
320 + batch_dead = tl.load(
321 + dead_mask + batch_idx * mask_stride, valid & in_batch, other=False
322 + )
323 + batch_source = tl.where(batch_dead, dead_id, batch_source)
324 +
325 + col = chunk_start - 1 - lookback
326 + in_window = valid & ~in_batch & (col >= 0) & (col < lookback_depth)
327 + col = tl.minimum(tl.maximum(col, 0), lookback_depth - 1)
328 + window_token = tl.load(
329 + lookback_token_ids + req * lookback_row_stride + col * lookback_col_stride,
330 + in_window,
331 + other=-1,
332 + )
333 + known = in_window & (window_token >= 0)
334 + window_source = tl.load(token_map + window_token, known, other=0)
335 + window_dead = tl.load(
336 + lookback_dead_mask
337 + + req * lookback_mask_row_stride
338 + + col * lookback_mask_col_stride,
339 + known,
340 + other=False,
341 + )
342 + window_source = tl.where(window_dead, dead_id, window_source)
343 +
344 + if cache is not None:
345 + clamped = tl.minimum(
346 + tl.maximum(lookback, 0), max_blocks * cache_block_size - 1
347 + )
348 + block_row = tl.minimum(req, num_table_rows - 1)
349 + needs_cache = valid & ~in_batch & ~known
350 + block = tl.load(
351 + block_table
352 + + block_row * table_stride
353 + + (clamped // cache_block_size) * table_col_stride,
354 + needs_cache,
355 + other=0,
356 + ).to(tl.int64)
357 + slot = tl.minimum(
358 + tl.maximum(block * cache_block_size + clamped % cache_block_size, 0),
359 + num_slots - 1,
360 + )
361 + fallback = tl.load(cache + slot, needs_cache, other=0)
362 + else:
363 + fallback = tl.full((BLOCK_T,), pad_id, tl.int32)
364 + source = tl.where(
365 + in_batch, batch_source, tl.where(known, window_source, fallback)
366 + ).to(tl.int64)
367 + blocked |= (lookback < 0) | (source == dead_id)
368 + value = tl.where(blocked, pad_id, source)
369 + multiplier = tl.load(multipliers + layer * MAX_NGRAM + shift)
370 + rolling ^= value * multiplier
371 + if shift > 0:
372 + col = (shift - 1) * num_heads + head
373 + param_offset = layer * (MAX_NGRAM - 1) * num_heads + col
374 + prime = tl.load(primes + param_offset, head < num_heads, other=1)
375 + offset = tl.load(offsets + param_offset, head < num_heads, other=0)
376 + hashed = rolling[:, None] % prime[None, :] + offset[None, :]
377 + out_offset = (token.to(tl.int64) * num_layers + layer)[:, None] * (
378 + (MAX_NGRAM - 1) * num_heads
379 + ) + col[None, :]
380 + tl.store(output + out_offset, hashed, valid[:, None] & (head < num_heads))
381 +
382 +
383 +class NgramHashState(nn.Module):
384 + """Maps each position to the hash ids of the n-grams ending there.
385 +
386 + Stateless on the V2 runner, which supplies every lookback token id. On
387 + the V1 runner it also keeps `hash_cache`, the slot-keyed rolling store
388 + of compressed ids (see module docstring), for generated tokens.
389 + """
390 +
391 + def __init__(
392 + self,
393 + vllm_config: VllmConfig,
394 + layout: EngramLayout,
395 + swa_cache_module: nn.Module,
396 + ) -> None:
397 + super().__init__()
398 + self.layout = layout
399 + self.swa_cache_module = swa_cache_module
400 + self.block_size: int = swa_cache_module.block_size
401 + self.lookback_depth: int = layout.max_ngram_size - 1
402 + self.use_slot_cache: bool = not vllm_config.use_v2_model_runner
403 + self._cache: torch.Tensor | None = None
404 + self._kv_cache_ref: weakref.ReferenceType[torch.Tensor] | None = None
405 +
406 + model_config = vllm_config.model_config
407 + from transformers import AutoTokenizer
408 +
409 + tokenizer = AutoTokenizer.from_pretrained(
410 + model_config.tokenizer,
411 + trust_remote_code=model_config.trust_remote_code,
412 + revision=model_config.revision,
413 + )
414 + token_map, vocab_size = build_compressed_token_map(tokenizer)
415 + if vocab_size != layout.compressed_vocab_size:
416 + raise ValueError(
417 + f"Compressed vocab size mismatch: built {vocab_size} from the "
418 + f"tokenizer, config expects {layout.compressed_vocab_size}; "
419 + "every hash multiplier derives from it, so the engram tables "
420 + "would be silently rehashed."
421 + )
422 + self.pad_id = token_map[layout.pad_token_id]
423 + multipliers = compute_hash_multipliers(
424 + layout.layer_ids, layout.max_ngram_size, vocab_size
425 + )
426 + self.register_buffer(
427 + "token_map", torch.tensor(token_map, dtype=torch.int32), persistent=False
428 + )
429 + self.register_buffer("primes", torch.tensor(layout.primes), persistent=False)
430 + self.register_buffer("offsets", layout.offsets, persistent=False)
431 + self.register_buffer("multipliers", multipliers, persistent=False)
432 + logger.info(
433 + "Built engram token map (%d -> %d ids) for layers %s",
434 + len(token_map),
435 + vocab_size,
436 + layout.layer_ids,
437 + )
438 +
439 + def ensure_cache(self) -> bool:
440 + """Lazily size the slot-keyed cache from the bound SWA KV cache.
441 +
442 + Returns False while the KV cache is unbound (profile run); the caller
443 + skips engram hashing then. Without the slot cache only that check
444 + remains.
445 + """
446 + kv_cache = self.swa_cache_module.kv_cache
447 + if kv_cache.numel() == 0:
448 + self._cache = None
449 + self._kv_cache_ref = None
450 + return False
451 + if not self.use_slot_cache:
452 + return True
453 + if self._kv_cache_ref is not None and self._kv_cache_ref() is kv_cache:
454 + return True
455 + # Graph memory profiling binds a temporary, smaller KV cache first.
456 + # Rebinding must discard its hash history without retaining KV storage.
457 + self._cache = torch.zeros(
458 + kv_cache.shape[0] * self.block_size,
459 + dtype=torch.int32,
460 + device=kv_cache.device,
461 + )
462 + self._kv_cache_ref = weakref.ref(kv_cache)
463 + return True
464 +
465 + def forward(
466 + self,
467 + input_ids: torch.Tensor,
468 + positions: torch.Tensor,
469 + query_start_loc: torch.Tensor,
470 + dead_mask: torch.Tensor,
471 + lookback_token_ids: torch.Tensor,
472 + lookback_dead_mask: torch.Tensor,
473 + slot_mapping: torch.Tensor | None,
474 + block_table: torch.Tensor | None,
475 + ) -> torch.Tensor:
476 + """Compute [tokens, layers, hash columns] int32 n-gram hashes.
477 +
478 + History comes from the current chunk, then the runner's lookback
479 + window, then the optional V1 slot cache. V2 needs only one launch.
480 + """
481 + cache = self._cache if self.use_slot_cache else None
482 + num_tokens = input_ids.shape[0]
483 + num_layers, max_ngram = self.multipliers.shape
484 + num_heads = self.primes.shape[-1]
485 + output = input_ids.new_empty(
486 + (num_tokens, num_layers, (max_ngram - 1) * num_heads), dtype=torch.int32
487 + )
488 + if num_tokens == 0:
489 + return output
490 + if self.use_slot_cache:
491 + assert cache is not None and slot_mapping is not None
492 + assert block_table is not None
493 + # Finish writes before other thread blocks read fallback history.
494 + _write_hash_cache_kernel[(triton.cdiv(num_tokens, 256),)](
495 + input_ids,
496 + self.token_map,
497 + dead_mask,
498 + slot_mapping,
499 + cache,
500 + num_tokens,
501 + input_ids.stride(0),
502 + dead_mask.stride(0),
503 + slot_mapping.stride(0),
504 + 256,
505 + DEAD_ID,
506 + )
507 + _hash_ids_kernel[(triton.cdiv(num_tokens, 32), num_layers)](
508 + input_ids,
509 + self.token_map,
510 + dead_mask,
511 + positions,
512 + block_table,
513 + query_start_loc,
514 + self.multipliers,
515 + self.primes,
516 + self.offsets,
517 + cache,
518 + lookback_token_ids,
519 + lookback_dead_mask,
520 + output,
521 + num_tokens,
522 + cache.shape[0] if cache is not None else 0,
523 + self.pad_id,
524 + input_stride=input_ids.stride(0),
525 + mask_stride=dead_mask.stride(0),
526 + position_stride=positions.stride(0),
527 + table_stride=block_table.stride(0) if block_table is not None else 0,
528 + table_col_stride=block_table.stride(1) if block_table is not None else 0,
529 + query_stride=query_start_loc.stride(0),
530 + num_query_rows=query_start_loc.numel() - 1,
531 + num_table_rows=block_table.shape[0] if block_table is not None else 0,
532 + max_blocks=block_table.shape[1] if block_table is not None else 0,
533 + cache_block_size=self.block_size,
534 + MAX_NGRAM=max_ngram,
535 + num_heads=num_heads,
536 + BLOCK_T=32,
537 + BLOCK_H=triton.next_power_of_2(num_heads),
538 + dead_id=DEAD_ID,
539 + lookback_depth=lookback_token_ids.shape[1],
540 + lookback_row_stride=lookback_token_ids.stride(0),
541 + lookback_col_stride=lookback_token_ids.stride(1),
542 + lookback_mask_row_stride=lookback_dead_mask.stride(0),
543 + lookback_mask_col_stride=lookback_dead_mask.stride(1),
544 + num_warps=4,
545 + )
546 + return output
547 +
548 +
549 +def _engram_head_shard_weight_loader(
550 + param: torch.nn.Parameter, loaded_weight: torch.Tensor
551 +) -> None:
552 + """Load this rank's complete head buckets. ue8m0 scales arrive as
553 + float8_e8m0fnu; keep the raw bytes (the param stores uint8)."""
554 + part_rows = param.shape[0]
555 + if loaded_weight.dtype == torch.float8_e8m0fnu:
556 + loaded_weight = loaded_weight.view(torch.uint8)
557 + shard = loaded_weight.narrow(0, param.engram_vocab_start, part_rows)
558 + assert shard.shape == param.shape, (
559 + f"engram shard {tuple(shard.shape)} does not fit param {tuple(param.shape)}"
560 + )
561 + param.data.copy_(shard)
562 +
563 +
564 +@triton.jit
565 +def _engram_lookup_kernel(
566 + weight,
567 + scales,
568 + ids,
569 + out,
570 + vocab_start,
571 + vocab_end,
572 + num_rows,
573 + ids_stride_t,
574 + ids_stride_h,
575 + HEAD_START: tl.constexpr,
576 + LOCAL_HEADS: tl.constexpr,
577 + TOTAL_HEADS: tl.constexpr,
578 + DIM: tl.constexpr,
579 + QUANT_BLOCK: tl.constexpr,
580 + BLOCK_R: tl.constexpr,
581 + GRID: tl.constexpr,
582 +):
583 + """Gather fp8 rows, apply their ue8m0 block scales, write bf16.
584 +
585 + Only this rank's heads are read; padded heads write zeros for all-gather.
586 + `weight`/`scales` may address pinned host memory through UVA.
587 + """
588 + cols = tl.arange(0, DIM)
589 + scale_cols = cols // QUANT_BLOCK
590 + for base in tl.range(tl.program_id(0) * BLOCK_R, num_rows, GRID * BLOCK_R):
591 + rows = base + tl.arange(0, BLOCK_R)
592 + valid = rows < num_rows
593 + head = HEAD_START + rows % LOCAL_HEADS
594 + token = (rows // LOCAL_HEADS).to(tl.int64)
595 + index = tl.load(
596 + ids + token * ids_stride_t + head * ids_stride_h,
597 + mask=valid & (head < TOTAL_HEADS),
598 + other=-1,
599 + ).to(tl.int64)
600 + owned = valid & (head < TOTAL_HEADS)
601 + owned &= (index >= vocab_start) & (index < vocab_end)
602 + local = tl.where(owned, index - vocab_start, 0)
603 + values = tl.load(
604 + weight + local[:, None] * DIM + cols[None, :],
605 + mask=owned[:, None],
606 + other=0.0,
607 + )
608 + scale = tl.load(
609 + scales + local[:, None] * (DIM // QUANT_BLOCK) + scale_cols[None, :],
610 + mask=owned[:, None],
611 + other=0,
612 + )
613 + # ue8m0 is a power of two, so its byte *is* the fp32 exponent field.
614 + scale = (scale.to(tl.int32) << 23).to(tl.float32, bitcast=True)
615 + tl.store(
616 + out + rows[:, None] * DIM + cols[None, :],
617 + (values.to(tl.float32) * scale).to(tl.bfloat16),
618 + mask=valid[:, None],
619 + )
620 +
621 +
622 +class ParallelEngramEmbedding(nn.Module):
623 + """The n-gram hash table, sharded by complete hash heads over TP ranks.
624 + Rows stay fp8 and are dequantized with ue8m0 per-32 scales on lookup.
625 +
626 + With `cpu_offload` the shard lives in pinned host memory and is read over
627 + UVA instead of HBM; the TP sharding is unchanged either way.
628 + """
629 +
630 + def __init__(
631 + self,
632 + num_embeddings: int,
633 + dim: int,
634 + head_sizes: tuple[int, ...],
635 + block_size: int = 32,
636 + cpu_offload: bool = False,
637 + ):
638 + super().__init__()
639 + tp_size = get_tensor_model_parallel_world_size()
640 + tp_rank = get_tensor_model_parallel_rank()
641 + assert head_sizes and all(size > 0 for size in head_sizes)
642 + assert sum(head_sizes) <= num_embeddings
643 + if cpu_offload and not is_uva_available():
644 + raise RuntimeError("Engram CPU offload requires UVA support")
645 + self.num_embeddings = num_embeddings
646 + self.dim = dim
647 + self.block_size = block_size
648 + self.n_hash_cols = len(head_sizes)
649 + self.part_n_hash_cols = triton.cdiv(self.n_hash_cols, tp_size)
650 + self.head_start = tp_rank * self.part_n_hash_cols
651 + head_end = self.head_start + self.part_n_hash_cols
652 + self.vocab_start_idx = sum(head_sizes[: self.head_start])
653 + self.vocab_end_idx = sum(head_sizes[:head_end])
654 + self.part_num_embeddings = self.vocab_end_idx - self.vocab_start_idx
655 + self.tp_size = tp_size
656 + self.cpu_offload = cpu_offload
657 + self._views: tuple[torch.Tensor, torch.Tensor] | None = None
658 + self._view_src: tuple[int, int] | None = None
659 + self._num_sms = torch.cuda.get_device_properties(
660 + torch.accelerator.current_device_index()
661 + ).multi_processor_count
662 +
663 + # Explicit device: model init runs under a `torch.device("cuda")`
664 + # context, which would otherwise put the shard in HBM.
665 + kwargs = {"device": "cpu", "pin_memory": True} if cpu_offload else {}
666 + self.weight = nn.Parameter(
667 + torch.empty(
668 + self.part_num_embeddings, dim, dtype=torch.float8_e4m3fn, **kwargs
669 + ),
670 + requires_grad=False,
671 + )
672 + self.weight_scale_inv = nn.Parameter(
673 + torch.empty(
674 + self.part_num_embeddings,
675 + dim // block_size,
676 + dtype=torch.uint8,
677 + **kwargs,
678 + ),
679 + requires_grad=False,
680 + )
681 + for param in (self.weight, self.weight_scale_inv):
682 + set_weight_attrs(
683 + param,
684 + {
685 + "weight_loader": _engram_head_shard_weight_loader,
686 + "engram_vocab_start": self.vocab_start_idx,
687 + },
688 + )
689 + if cpu_offload:
690 + logger.info(
691 + "Engram table offloaded to pinned host memory: %d rows x %d, "
692 + "%.2f GiB per rank",
693 + self.part_num_embeddings,
694 + dim,
695 + self.part_num_embeddings * (dim + dim // block_size) / 1024**3,
696 + )
697 +
698 + def _storage(self) -> tuple[torch.Tensor, torch.Tensor]:
699 + """Parameters when resident, else cached UVA views of the pinned shard.
700 +
701 + Rebuilt if anything swaps `.data`, so a stale device pointer cannot
702 + survive silently.
703 + """
704 + if not self.cpu_offload:
705 + return self.weight.data, self.weight_scale_inv.data
706 + src = (self.weight.data_ptr(), self.weight_scale_inv.data_ptr())
707 + if self._view_src != src:
708 + self._views = (
709 + get_accelerator_view_from_cpu_tensor(self.weight.data),
710 + get_accelerator_view_from_cpu_tensor(self.weight_scale_inv.data),
711 + )
712 + self._view_src = src
713 + assert self._views is not None
714 + return self._views
715 +
716 + def lookup(
717 + self, indices: torch.Tensor, out: torch.Tensor, background: bool = False
718 + ) -> None:
719 + """Look up local heads of [T, heads] into [T, local_heads, dim] bf16.
720 +
721 + `background` limits the grid to leave SMs for concurrent work.
722 + """
723 + rows = indices.shape[0] * self.part_n_hash_cols
724 + if not rows:
725 + return
726 + weight, scales = self._storage()
727 + # The table dwarfs TLB reach, so a persistent grid near the SM count
728 + # beats one program per row; halve it to leave SMs for the main stream.
729 + tiles = triton.cdiv(rows, 16)
730 + grid = min(tiles, self._num_sms // 2 if background else self._num_sms)
731 + _engram_lookup_kernel[(grid,)](
732 + weight,
733 + scales,
734 + indices,
735 + out,
736 + self.vocab_start_idx,
737 + self.vocab_end_idx,
738 + rows,
739 + indices.stride(0),
740 + indices.stride(1),
741 + HEAD_START=self.head_start,
742 + LOCAL_HEADS=self.part_n_hash_cols,
743 + TOTAL_HEADS=self.n_hash_cols,
744 + DIM=self.dim,
745 + QUANT_BLOCK=self.block_size,
746 + BLOCK_R=16,
747 + GRID=grid,
748 + )
749 +
750 + def forward(self, indices: torch.Tensor) -> torch.Tensor:
751 + """indices: [num_tokens, n_hash_cols] -> [num_tokens, n_hash_cols, dim]
752 + bf16, gathered from all TP shards."""
753 + out = torch.empty(
754 + (indices.shape[0], self.part_n_hash_cols, self.dim),
755 + dtype=torch.bfloat16,
756 + device=indices.device,
757 + )
758 + self.lookup(indices, out)
759 + if self.tp_size > 1:
760 + out = tensor_model_parallel_all_gather(out, dim=1)
761 + out = out[:, : self.n_hash_cols]
762 + return out
763 +
764 +
765 +@triton.jit(do_not_specialize=["num_kv_tokens"])
766 +def _fused_engram_post_wkv_kernel(
767 + hidden_states,
768 + kv,
769 + q_weight,
770 + k_weight,
771 + token_mask,
772 + output,
773 + num_kv_tokens,
774 + hidden_stride_t,
775 + hidden_stride_h,
776 + hidden_stride_d,
777 + kv_stride_t,
778 + kv_stride_d,
779 + q_stride_h,
780 + q_stride_d,
781 + k_stride_h,
782 + k_stride_d,
783 + mask_stride,
784 + output_stride_t,
785 + output_stride_h,
786 + output_stride_d,
787 + eps,
788 + clamp_value,
789 + DIM: tl.constexpr,
790 + HC_MULT: tl.constexpr,
791 + BLOCK_SIZE: tl.constexpr,
792 + HAS_MASK: tl.constexpr,
793 +):
794 + program_idx = tl.program_id(0)
795 + token_idx = program_idx // HC_MULT
796 + hc_idx = program_idx % HC_MULT
797 + token_idx = token_idx.to(tl.int64)
798 + source_idx = token_idx
799 + source_valid = source_idx < num_kv_tokens
800 +
801 + dim_offsets = tl.arange(0, BLOCK_SIZE)
802 + dim_valid = dim_offsets < DIM
803 + hidden = tl.load(
804 + hidden_states
805 + + token_idx * hidden_stride_t
806 + + hc_idx * hidden_stride_h
807 + + dim_offsets * hidden_stride_d,
808 + mask=dim_valid,
809 + other=0.0,
810 + ).to(tl.float32)
811 + key = tl.load(
812 + kv + source_idx * kv_stride_t + (hc_idx * DIM + dim_offsets) * kv_stride_d,
813 + mask=source_valid & dim_valid,
814 + other=0.0,
815 + ).to(tl.float32)
816 + q = tl.load(
817 + q_weight + hc_idx * q_stride_h + dim_offsets * q_stride_d,
818 + mask=dim_valid,
819 + other=0.0,
820 + ).to(tl.float32)
821 + k = tl.load(
822 + k_weight + hc_idx * k_stride_h + dim_offsets * k_stride_d,
823 + mask=dim_valid,
824 + other=0.0,
825 + ).to(tl.float32)
826 +
827 + hidden_rms = tl.rsqrt(tl.sum(hidden * hidden, axis=0) / DIM + eps)
828 + key_rms = tl.rsqrt(tl.sum(key * key, axis=0) / DIM + eps)
829 + dot = tl.sum(hidden * q * k * key, axis=0)
830 + dot *= hidden_rms * key_rms * tl.rsqrt(DIM * 1.0)
831 + gate_input = tl.sqrt(tl.maximum(tl.abs(dot), clamp_value))
832 + gate_input = tl.where(dot < 0.0, -gate_input, gate_input)
833 + gate = tl.sigmoid(gate_input)
834 + if HAS_MASK:
835 + active = tl.load(
836 + token_mask + source_idx * mask_stride,
837 + mask=source_valid,
838 + other=0,
839 + )
840 + gate = tl.where(active, gate, 0.0)
841 +
842 + value = tl.load(
843 + kv + source_idx * kv_stride_t + (HC_MULT * DIM + dim_offsets) * kv_stride_d,
844 + mask=source_valid & dim_valid,
845 + other=0.0,
846 + ).to(tl.float32)
847 + tl.store(
848 + output
849 + + token_idx * output_stride_t
850 + + hc_idx * output_stride_h
851 + + dim_offsets * output_stride_d,
852 + hidden + gate * value,
853 + mask=dim_valid,
854 + )
855 +
856 +
857 +@triton.jit(do_not_specialize=["num_tokens", "token_start", "num_elements"])
858 +def _engram_sp_rows_kernel(
859 + gathered,
860 + output,
861 + num_tokens,
862 + token_start,
863 + num_elements,
864 + LOCAL_WIDTH: tl.constexpr,
865 + WIDTH: tl.constexpr,
866 + BLOCK_SIZE: tl.constexpr,
867 +):
868 + offsets = tl.program_id(0).to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
869 + tokens = token_start + offsets // WIDTH
870 + cols = offsets % WIDTH
871 + source = (cols // LOCAL_WIDTH * num_tokens + tokens) * LOCAL_WIDTH
872 + source += cols % LOCAL_WIDTH
873 + values = tl.load(
874 + gathered + source, (offsets < num_elements) & (tokens < num_tokens), other=0
875 + )
876 + tl.store(output + offsets, values, offsets < num_elements)
877 +
878 +
879 +class Engram(nn.Module):
880 + """Writes an n-gram lookup into the residual stream, gated by how well it
881 + matches that stream.
882 +
883 + The hash ids fetch `n_hash_cols` rows; `wkv` turns them into one key per
884 + hc copy plus a shared value. The gate is a normalized dot product of the
885 + stream against the key, signed-sqrt'ed before the sigmoid (matching the
886 + training kernel).
887 + """
888 +
889 + def __init__(
890 + self,
891 + config,
892 + quant_config: QuantizationConfig | None,
893 + layout: EngramLayout,
894 + layer_hash_index: int,
895 + use_sequence_parallel: bool,
896 + prefix: str,
897 + ) -> None:
898 + super().__init__()
899 + self.layer_hash_index = layer_hash_index
900 + self.dim = config.hidden_size
901 + self.hc_mult = config.hc_mult
902 + self.eps = config.rms_norm_eps
903 + self.clamp_value = 1e-6
904 + self.use_sequence_parallel = use_sequence_parallel
905 +
906 + # Named ``embed_tokens`` so the checkpoint's ``engram.embed.weight``
907 + # survives the mapper's ``embed.weight`` -> ``embed_tokens.weight``
908 + # suffix rule.
909 + engram_config = get_current_vllm_config().engram_config
910 + self.embed_tokens = ParallelEngramEmbedding(
911 + layout.num_embeddings[layer_hash_index],
912 + layout.head_dim,
913 + tuple(size for order in layout.primes[layer_hash_index] for size in order),
914 + cpu_offload=engram_config.cpu_offload if engram_config else True,
915 + )
916 + n_hash_cols = (layout.max_ngram_size - 1) * layout.n_heads
917 + self.wkv = ReplicatedLinear(
918 + n_hash_cols * layout.head_dim,
919 + self.dim * (self.hc_mult + 1),
920 + bias=False,
921 + quant_config=quant_config,
922 + return_bias=False,
923 + prefix=f"{prefix}.wkv",
924 + )
925 + self.q_weight = nn.Parameter(
926 + torch.empty(self.hc_mult, self.dim, dtype=torch.bfloat16),
927 + requires_grad=False,
928 + )
929 + self.k_weight = nn.Parameter(
930 + torch.empty(self.hc_mult, self.dim, dtype=torch.bfloat16),
931 + requires_grad=False,
932 + )
933 +
934 + max_tokens = get_current_vllm_config().scheduler_config.max_num_batched_tokens
935 + # Keep lookup results alive across breakable graph segments.
936 + self.staged_rows = torch.empty(
937 + max_tokens,
938 + self.embed_tokens.part_n_hash_cols,
939 + layout.head_dim,
940 + dtype=torch.bfloat16,
941 + )
942 +
943 + def prepare_embeddings(self, hash_ids: torch.Tensor) -> None:
944 + """Gather this layer's rows on the main stream before decoder layers."""
945 + self.embed_tokens.lookup(hash_ids, self.staged_rows[: hash_ids.shape[0]])
946 +
947 + def embed(self, hash_ids: torch.Tensor) -> torch.Tensor:
948 + """Gather heads, returning only local tokens when SP is enabled."""
949 + rows = self.staged_rows[: hash_ids.shape[0]]
950 + if self.embed_tokens.tp_size == 1:
951 + return rows
952 + if self.use_sequence_parallel:
953 + tp_size = self.embed_tokens.tp_size
954 + num_tokens, local_heads, dim = rows.shape
955 + gathered = tensor_model_parallel_all_gather(rows, dim=0)
956 + chunk = (num_tokens + tp_size - 1) // tp_size
957 + rows = rows.new_empty((chunk, self.embed_tokens.n_hash_cols, dim))
958 + _engram_sp_rows_kernel[(triton.cdiv(rows.numel(), 1024),)](
959 + gathered,
960 + rows,
961 + num_tokens,
962 + get_tensor_model_parallel_rank() * chunk,
963 + rows.numel(),
964 + local_heads * dim,
965 + self.embed_tokens.n_hash_cols * dim,
966 + BLOCK_SIZE=1024,
967 + )
968 + return rows
969 + rows = tensor_model_parallel_all_gather(rows, dim=1)
970 + return rows[:, : self.embed_tokens.n_hash_cols]
971 +
972 + def forward(
973 + self,
974 + hidden_states: torch.Tensor,
975 + hash_ids: torch.Tensor,
976 + token_mask: torch.Tensor | None = None,
977 + ) -> torch.Tensor:
978 + """hidden_states: [T, hc_mult, dim]; hash_ids: [T, n_hash_cols] (all
979 + tokens, pre sequence-parallel shard); token_mask: [T], False shuts
980 + the gate so those positions pass through untouched."""
981 + kv = self.wkv(self.embed(hash_ids).flatten(-2))
982 + num_kv_tokens = hash_ids.shape[0]
983 + assert token_mask is None or token_mask.shape == (num_kv_tokens,)
984 + if self.use_sequence_parallel:
985 + tp_size = get_tensor_model_parallel_world_size()
986 + tp_rank = get_tensor_model_parallel_rank()
987 + shard_size = (num_kv_tokens + tp_size - 1) // tp_size
988 + assert hidden_states.shape[0] == shard_size
989 + start = min(tp_rank * shard_size, num_kv_tokens)
990 + num_kv_tokens = min(shard_size, num_kv_tokens - start)
991 + if token_mask is not None:
992 + token_mask = token_mask[start : start + num_kv_tokens]
993 +
994 + num_tokens, hc_mult, dim = hidden_states.shape
995 + assert hc_mult == self.hc_mult and dim == self.dim
996 + assert kv.ndim == 2 and kv.shape[1] == (hc_mult + 1) * dim
997 + output = torch.empty_like(hidden_states)
998 + if num_tokens == 0:
999 + return output
1000 +
1001 + block_size = triton.next_power_of_2(dim)
1002 + num_warps = 8 if block_size >= 2048 else 4
1003 + mask = token_mask if token_mask is not None else hidden_states
1004 + _fused_engram_post_wkv_kernel[(num_tokens * hc_mult,)](
1005 + hidden_states,
1006 + kv,
1007 + self.q_weight,
1008 + self.k_weight,
1009 + mask,
1010 + output,
1011 + num_kv_tokens,
1012 + hidden_states.stride(0),
1013 + hidden_states.stride(1),
1014 + hidden_states.stride(2),
1015 + kv.stride(0),
1016 + kv.stride(1),
1017 + self.q_weight.stride(0),
1018 + self.q_weight.stride(1),
1019 + self.k_weight.stride(0),
1020 + self.k_weight.stride(1),
1021 + token_mask.stride(0) if token_mask is not None else 0,
1022 + output.stride(0),
1023 + output.stride(1),
1024 + output.stride(2),
1025 + self.eps,
1026 + self.clamp_value,
1027 + DIM=dim,
1028 + HC_MULT=hc_mult,
1029 + BLOCK_SIZE=block_size,
1030 + HAS_MASK=token_mask is not None,
1031 + num_warps=num_warps,
1032 + )
1033 + return output
vllm/models/deepseek_v4_1/common/ops/__init__.py new
+42
@@ -0,0 +1,42 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""DeepSeek V4.1 fused ops.
4 +
5 +Single bridging point for the kernels V4.1 shares verbatim with V4: the
6 +re-exports below are the only place this package reaches into
7 +``vllm.models.deepseek_v4``, so submodules import them from ``.`` instead of
8 +naming the V4 path themselves. Keep these imports above the ``.``-relative
9 +ones — ``indexer_k_store`` reads them off this partially-initialized module.
10 +"""
11 +
12 +from vllm.models.deepseek_v4.common.ops.fused_indexer_q import (
13 + MXFP4_BLOCK_SIZE,
14 + fused_indexer_q_rope_quant,
15 +)
16 +from vllm.models.deepseek_v4.common.ops.fused_indexer_q import (
17 + _fp32x2_to_fp4x2 as _fp32x2_to_fp4x2, # re-export for .indexer_k_store
18 +)
19 +from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import (
20 + fused_inv_rope_fp8_quant,
21 +)
22 +
23 +from .cache_utils import (
24 + build_flashinfer_mixed_sparse_indices,
25 + combine_topk_swa_indices,
26 + compute_global_topk_indices_and_lens,
27 + dequantize_and_gather_k_cache,
28 + quantize_and_insert_k_cache,
29 +)
30 +from .indexer_k_store import indexer_k_norm_rope_store
31 +
32 +__all__ = [
33 + "MXFP4_BLOCK_SIZE",
34 + "build_flashinfer_mixed_sparse_indices",
35 + "combine_topk_swa_indices",
36 + "compute_global_topk_indices_and_lens",
37 + "dequantize_and_gather_k_cache",
38 + "fused_indexer_q_rope_quant",
39 + "fused_inv_rope_fp8_quant",
40 + "indexer_k_norm_rope_store",
41 + "quantize_and_insert_k_cache",
42 +]
vllm/models/deepseek_v4_1/common/ops/cache_utils.py new
+1254
@@ -0,0 +1,1254 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Triton kernels for DeepseekV4 paged K-cache management and sparse-attention index
5 +preparation.
6 +
7 +- quantize_and_insert_k_cache: quantize bf16 K to UE8M0 FP8 and insert into
8 + the paged cache.
9 +- dequantize_and_gather_k_cache: gather and dequantize FP8 K from the paged
10 + cache for sparse/SWA prefill.
11 +- compute_global_topk_indices_and_lens: map local topk indices to global KV
12 + cache slots and count valid entries.
13 +- combine_topk_swa_indices: concatenate topk compressed indices with SWA
14 + window indices for sparse prefill.
15 +"""
16 +
17 +from dataclasses import dataclass
18 +from typing import Any
19 +
20 +import torch
21 +
22 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
23 + get_fp8_min_max,
24 +)
25 +from vllm.model_executor.warmup.jit_warmup import (
26 + VllmJitKernel,
27 + _WarmupInputRows,
28 + zip_inputs,
29 +)
30 +from vllm.model_executor.warmup.jit_warmup_triton_helper import (
31 + TritonPointerInputVariant,
32 + TritonWarmupTensor,
33 +)
34 +from vllm.platforms import current_platform
35 +from vllm.triton_utils import tl, triton
36 +from vllm.utils.import_utils import has_cutedsl
37 +from vllm.utils.math_utils import next_power_of_2
38 +
39 +
40 +@triton.jit
41 +def quantize_and_insert_k_kernel(
42 + # Input tensors
43 + k_ptr, # [num_tokens, 512] bf16
44 + slot_mapping_ptr, # [num_tokens] int64
45 + # Output tensor
46 + k_cache_ptr, # [num_blocks, block_bytes] as uint8 (flattened view)
47 + # Dimensions
48 + num_tokens,
49 + input_dim: tl.constexpr, # 512
50 + fp8_dim: tl.constexpr, # 448
51 + bf16_dim: tl.constexpr, # 64
52 + scale_dim: tl.constexpr, # 8
53 + quant_block: tl.constexpr, # 64 (quantization block size)
54 + cache_block_size: tl.constexpr, # 64 (paged cache block size)
55 + token_data_size: tl.constexpr, # 576 bytes per token data
56 + block_stride: tl.constexpr, # total bytes per block (padded)
57 + fp8_max: tl.constexpr,
58 + n_quant_blocks: tl.constexpr, # 8 (7 real + 1 padding)
59 + use_fnuz: tl.constexpr = False,
60 +):
61 + """
62 + Quantize K tensor and insert into paged K cache.
63 +
64 + K Cache block layout (block_size=64 tokens):
65 + - [0, 64*576): Token data, each token has 448 fp8 + 128 bf16
66 + - [64*576, 64*576 + 64*8): Scales, each token has 8 uint8 scales
67 + - [64*576 + 64*8, block_stride): Padding
68 +
69 + One program per token.
70 +
71 + ``use_fnuz=True`` selects FNUZ (``tl.float8e4b8``); default OCP
72 + (``tl.float8e4nv``) matches every production caller.
73 + """
74 + pid = tl.program_id(0)
75 +
76 + if pid >= num_tokens:
77 + return
78 +
79 + # Get slot mapping
80 + slot_idx = tl.load(slot_mapping_ptr + pid)
81 + if slot_idx == -1:
82 + return
83 +
84 + block_idx = slot_idx // cache_block_size
85 + pos_in_block = slot_idx % cache_block_size
86 +
87 + # Input pointer for this token
88 + input_row_ptr = k_ptr + pid * input_dim
89 +
90 + # int64: block_idx * block_stride can exceed 2^31 with many KV-cache blocks
91 + # (e.g. >= 57K at block_stride ~37K). Matches gather path below.
92 + cache_block_ptr = k_cache_ptr + block_idx.to(tl.int64) * block_stride
93 +
94 + # Token data pointer: token data is stored contiguously at start of block
95 + # Each token's data is at offset pos_in_block * token_data_size
96 + token_data_ptr = cache_block_ptr + pos_in_block * token_data_size
97 +
98 + # Scale pointer: scales are stored after ALL token data in the block
99 + # Scale for this token is at offset (64 * 576) + pos_in_block * 8
100 + token_scale_ptr = (
101 + cache_block_ptr + cache_block_size * token_data_size + pos_in_block * scale_dim
102 + )
103 +
104 + # Token data layout: [0:448] fp8, [448:576] bf16
105 + token_fp8_ptr = token_data_ptr
106 + token_bf16_ptr = token_data_ptr + fp8_dim
107 +
108 + # ========== Quantize and store FP8 portion (first 448 elements) ==========
109 + # Using UE8M0 quantization strategy (scale is power of 2, stored as uint8 exponent)
110 + for qblock_idx in tl.static_range(n_quant_blocks):
111 + qblock_start = qblock_idx * quant_block
112 +
113 + if qblock_start < fp8_dim:
114 + offsets = qblock_start + tl.arange(0, quant_block)
115 + mask = offsets < fp8_dim
116 +
117 + # Load bf16 input
118 + x = tl.load(input_row_ptr + offsets, mask=mask, other=0.0)
119 +
120 + # Compute absmax scale (same as CUDA kernel)
121 + abs_x = tl.abs(x)
122 + block_max = tl.max(abs_x, axis=0)
123 + block_max = tl.maximum(block_max, 1e-4) # Match CUDA: fmaxf(amax, 1e-4)
124 +
125 + # UE8M0: Round scale UP to next power of 2
126 + # scale = 2^ceil(log2(block_max / fp8_max))
127 + raw_scale = block_max / fp8_max
128 + log_scale = tl.log2(raw_scale)
129 + exponent = tl.ceil(log_scale) # Round UP to next integer exponent
130 + scale = tl.exp2(exponent) # scale = 2^exponent (power of 2)
131 +
132 + # Quantize to fp8: fp8_value = bf16_value / scale
133 + x_scaled = x / scale
134 + x_clamped = tl.clamp(x_scaled, -fp8_max, fp8_max)
135 +
136 + # Convert to fp8 (FNUZ on gfx942, OCP elsewhere), then bitcast to uint8.
137 + if use_fnuz:
138 + x_fp8 = x_clamped.to(tl.float8e4b8)
139 + else:
140 + x_fp8 = x_clamped.to(tl.float8e4nv)
141 + x_uint8 = x_fp8.to(tl.uint8, bitcast=True)
142 +
143 + # Store as uint8 (1 byte each)
144 + tl.store(token_fp8_ptr + offsets, x_uint8, mask=mask)
145 +
146 + # UE8M0 scale encoding: stored_value = exponent + 127 (bias)
147 + # During dequant: scale = 2^(stored_value - 127)
148 + encoded_scale = exponent + 127.0
149 + encoded_scale = tl.maximum(tl.minimum(encoded_scale, 255.0), 0.0)
150 + tl.store(token_scale_ptr + qblock_idx, encoded_scale.to(tl.uint8))
151 +
152 + # Padding scale at index 7
153 + tl.store(token_scale_ptr + 7, tl.zeros((), dtype=tl.uint8))
154 +
155 + # ========== Store BF16 portion (last 64 elements, no quantization) ==========
156 + bf16_input_offset = fp8_dim
157 +
158 + # Process bf16 in chunks of 16
159 + bf16_out_ptr = token_bf16_ptr.to(tl.pointer_type(tl.bfloat16))
160 + for i in tl.static_range(bf16_dim // 16):
161 + chunk_offsets = i * 16 + tl.arange(0, 16)
162 + bf16_vals = tl.load(input_row_ptr + bf16_input_offset + chunk_offsets)
163 + tl.store(bf16_out_ptr + chunk_offsets, bf16_vals)
164 +
165 +
166 +def quantize_and_insert_k_cache(
167 + k: torch.Tensor, # [num_tokens, 512] bf16
168 + k_cache: torch.Tensor, # [num_blocks, block_bytes] uint8
169 + slot_mapping: torch.Tensor, # [num_tokens] int64
170 + block_size: int = 64,
171 + is_ue8m0: bool = True,
172 + use_fnuz: bool = False,
173 +):
174 + """
175 + Quantize K tensor and insert into paged K cache.
176 +
177 + K Cache block layout (block_size=64 tokens):
178 + - First 64 * 576 = 36864 bytes: Token data
179 + - Each token: 448 bytes (fp8) + 128 bytes (bf16)
180 + - Next 64 * 8 = 512 bytes: Scales
181 + - Each token: 8 bytes (uint8 scales, 7 real + 1 padding)
182 + - Padded to multiple of 576
183 +
184 + ``use_fnuz=True`` selects FNUZ E4M3 cache encoding and is only valid on
185 + platforms whose FP8 format is FNUZ. ``use_fnuz=False`` selects OCP E4M3,
186 + which is used by OCP-encoded caches even on gfx942.
187 + """
188 + assert k.dim() == 2 and k.shape[1] == 512, (
189 + f"K must be [num_tokens, 512], got {k.shape}"
190 + )
191 + assert k.dtype == torch.bfloat16, f"K must be bf16, got {k.dtype}"
192 + assert is_ue8m0, "Only support ue8m0 quantization."
193 +
194 + # NOTE: When using DP, slot_mapping.shape[0] can be less than k.shape[0] due to
195 + # padding. Always use slot_mapping.shape[0] as the token count.
196 + num_tokens = slot_mapping.shape[0]
197 + block_stride = k_cache.stride(0) # bytes per block
198 +
199 + TOKEN_FP8_DIM = 448
200 + TOKEN_BF16_DIM = 64
201 + TOKEN_SCALE_DIM = 8
202 + QUANT_BLOCK_SIZE = 64
203 + if use_fnuz:
204 + if not current_platform.is_fp8_fnuz():
205 + raise ValueError("use_fnuz=True requires a platform using FNUZ FP8")
206 + _, FP8_MAX = get_fp8_min_max()
207 + else:
208 + FP8_MAX = torch.finfo(torch.float8_e4m3fn).max
209 + TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2
210 +
211 + grid = (num_tokens,)
212 +
213 + quantize_and_insert_k_kernel[grid](
214 + k,
215 + slot_mapping,
216 + k_cache,
217 + num_tokens,
218 + input_dim=512,
219 + fp8_dim=TOKEN_FP8_DIM,
220 + bf16_dim=TOKEN_BF16_DIM,
221 + scale_dim=TOKEN_SCALE_DIM,
222 + quant_block=QUANT_BLOCK_SIZE,
223 + cache_block_size=block_size,
224 + token_data_size=TOKEN_DATA_SIZE,
225 + block_stride=block_stride,
226 + fp8_max=FP8_MAX,
227 + n_quant_blocks=8,
228 + use_fnuz=use_fnuz,
229 + )
230 +
231 +
232 +@triton.jit
233 +def _dequantize_and_gather_k_kernel(
234 + out_ptr,
235 + out_stride0,
236 + out_stride1,
237 + k_cache_ptr,
238 + seq_lens_ptr,
239 + block_table_ptr,
240 + offset,
241 + gather_lens_ptr,
242 + # Constants
243 + max_blocks_per_seq: tl.constexpr,
244 + fp8_dim: tl.constexpr, # 448
245 + bf16_dim: tl.constexpr, # 64
246 + scale_dim: tl.constexpr, # 8
247 + quant_block: tl.constexpr, # 64 (quantization block size)
248 + cache_block_size: tl.constexpr, # 64 or 128 (paged cache block size)
249 + token_data_size: tl.constexpr, # 576 bytes per token data
250 + block_stride: tl.constexpr, # total bytes per block (padded) int32
251 + output_dim: tl.constexpr, # 512
252 + fp8_max: tl.constexpr,
253 + n_quant_blocks: tl.constexpr, # 7 real blocks
254 + use_fnuz: tl.constexpr = False,
255 +):
256 + batch_idx = tl.program_id(0)
257 + worker_id = tl.program_id(1)
258 + num_workers = tl.num_programs(1)
259 +
260 + seq_len = tl.load(seq_lens_ptr + batch_idx)
261 + if gather_lens_ptr is not None: # noqa: SIM108
262 + gather_len = tl.load(gather_lens_ptr + batch_idx)
263 + else:
264 + # Gather all tokens
265 + gather_len = seq_len
266 + start_pos = seq_len - gather_len
267 +
268 + for i in range(worker_id, gather_len, num_workers):
269 + # Calculate the actual token index in the sequence
270 + pos = start_pos + i
271 +
272 + # Calculate which block and position within block
273 + block_in_seq = pos // cache_block_size
274 + pos_in_block = pos % cache_block_size
275 +
276 + # Get physical block index from block table
277 + block_table_row_ptr = block_table_ptr + batch_idx * max_blocks_per_seq
278 + physical_block_idx = tl.load(block_table_row_ptr + block_in_seq) # int32
279 +
280 + # int64: physical_block_idx * block_stride can exceed 2^31 with many
281 + # KV-cache blocks (e.g. >= 57K at block_stride ~37K).
282 + cache_block_ptr = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride
283 +
284 + # Token data pointer
285 + token_data_ptr = cache_block_ptr + pos_in_block * token_data_size
286 +
287 + # Scale pointer: after all token data
288 + token_scale_ptr = (
289 + cache_block_ptr
290 + + cache_block_size * token_data_size
291 + + pos_in_block * scale_dim
292 + )
293 +
294 + # Token data layout: [0:448] fp8, [448:576] bf16
295 + token_fp8_ptr = token_data_ptr
296 + token_bf16_ptr = token_data_ptr + fp8_dim
297 +
298 + # Output pointer for this token (flattened)
299 + output_row_ptr = out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1
300 +
301 + # ========== Dequantize FP8 portion using UE8M0 ==========
302 + for qblock_idx in tl.static_range(n_quant_blocks):
303 + qblock_start = qblock_idx * quant_block
304 +
305 + if qblock_start < fp8_dim:
306 + offsets = qblock_start + tl.arange(0, quant_block)
307 + mask = offsets < fp8_dim
308 +
309 + # Load quantized fp8 values (stored as uint8)
310 + x_uint8 = tl.load(token_fp8_ptr + offsets, mask=mask, other=0)
311 +
312 + # Bitcast uint8 back to fp8 (FNUZ on gfx942, OCP elsewhere).
313 + if use_fnuz:
314 + x_fp8 = x_uint8.to(tl.float8e4b8, bitcast=True)
315 + else:
316 + x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)
317 +
318 + # Convert fp8 to float32 for computation
319 + x_float = x_fp8.to(tl.float32)
320 +
321 + # Load and decode UE8M0 scale
322 + # UE8M0: scale = 2^(stored_value - 127)
323 + encoded_scale = tl.load(token_scale_ptr + qblock_idx)
324 + exponent = encoded_scale.to(tl.float32) - 127.0
325 + scale = tl.exp2(exponent)
326 +
327 + # Dequantize: bf16_value = fp8_value * scale
328 + x_dequant = x_float * scale
329 +
330 + # Store as bf16
331 + tl.store(output_row_ptr + offsets, x_dequant.to(tl.bfloat16), mask=mask)
332 +
333 + # ========== Copy BF16 portion directly ==========
334 + bf16_output_offset = fp8_dim # After 448 elements in output
335 +
336 + # Read bf16 from cache
337 + bf16_cache_ptr = token_bf16_ptr.to(tl.pointer_type(tl.bfloat16))
338 +
339 + # Process in chunks of 16
340 + for j in tl.static_range(bf16_dim // 16):
341 + chunk_offsets = j * 16 + tl.arange(0, 16)
342 + bf16_vals = tl.load(bf16_cache_ptr + chunk_offsets)
343 + tl.store(output_row_ptr + bf16_output_offset + chunk_offsets, bf16_vals)
344 +
345 +
346 +def dequantize_and_gather_k_cache_triton(
347 + # [num_reqs, max_num_tokens, head_size]
348 + out: torch.Tensor,
349 + # [num_blocks, block_size, head_bytes]
350 + k_cache: torch.Tensor,
351 + # [num_reqs]
352 + seq_lens: torch.Tensor,
353 + # [num_reqs]
354 + gather_lens: torch.Tensor | None,
355 + # [num_reqs, max_blocks_per_seq]
356 + block_table: torch.Tensor,
357 + block_size: int,
358 + offset: int,
359 + use_fnuz: bool = False,
360 +) -> None:
361 + TOKEN_FP8_DIM = 448
362 + TOKEN_BF16_DIM = 64
363 + TOKEN_SCALE_DIM = 8
364 + QUANT_BLOCK_SIZE = 64
365 + FP8_MAX = 448.0
366 + TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2
367 +
368 + num_reqs = seq_lens.shape[0]
369 + NUM_WORKERS = 128
370 + _dequantize_and_gather_k_kernel[(num_reqs, NUM_WORKERS)](
371 + out,
372 + out.stride(0),
373 + out.stride(1),
374 + k_cache,
375 + seq_lens,
376 + block_table,
377 + offset,
378 + gather_lens,
379 + max_blocks_per_seq=block_table.shape[-1],
380 + fp8_dim=TOKEN_FP8_DIM,
381 + bf16_dim=TOKEN_BF16_DIM,
382 + scale_dim=TOKEN_SCALE_DIM,
383 + quant_block=QUANT_BLOCK_SIZE,
384 + cache_block_size=block_size,
385 + token_data_size=TOKEN_DATA_SIZE,
386 + block_stride=k_cache.stride(0),
387 + output_dim=512,
388 + fp8_max=FP8_MAX,
389 + n_quant_blocks=7,
390 + use_fnuz=use_fnuz,
391 + )
392 +
393 +
394 +def dequantize_and_gather_k_cache(
395 + # [num_reqs, max_num_tokens, head_size]
396 + out: torch.Tensor,
397 + # [num_blocks, block_size, head_bytes]
398 + k_cache: torch.Tensor,
399 + # [num_reqs]
400 + seq_lens: torch.Tensor,
401 + # [num_reqs]
402 + gather_lens: torch.Tensor | None,
403 + # [num_reqs, max_blocks_per_seq]
404 + block_table: torch.Tensor,
405 + block_size: int,
406 + offset: int,
407 + use_fnuz: bool = False,
408 +) -> None:
409 + """Dequantize and gather a paged DSv4 K cache.
410 +
411 + ``use_fnuz`` MUST match the encoder of the specific cache being read:
412 + ``False`` for ``compressed_k_cache`` (Triton encoder is OCP everywhere),
413 + ``current_platform.is_fp8_fnuz()`` for ``swa_k_cache`` (C++ encoder
414 + writes FNUZ on gfx942 and OCP on gfx950).
415 + """
416 + if has_cutedsl():
417 + # lazily import, otherwise some tests fail due to CUDA driver init failure.
418 + from vllm.models.deepseek_v4.nvidia.ops.dequant_gather_k_cutedsl import (
419 + dequantize_and_gather_k_cache_cutedsl,
420 + )
421 +
422 + dequantize_and_gather_k_cache_cutedsl(
423 + out, k_cache, seq_lens, gather_lens, block_table, block_size, offset
424 + )
425 + return
426 +
427 + dequantize_and_gather_k_cache_triton(
428 + out,
429 + k_cache,
430 + seq_lens,
431 + gather_lens,
432 + block_table,
433 + block_size,
434 + offset,
435 + use_fnuz=use_fnuz,
436 + )
437 +
438 +
439 +def compute_global_topk_indices_and_lens(
440 + topk_indices: torch.Tensor,
441 + token_to_req_indices: torch.Tensor,
442 + block_table: torch.Tensor,
443 + block_size: int,
444 + is_valid_token: torch.Tensor,
445 +) -> tuple[torch.Tensor, torch.Tensor]:
446 + """Map local topk indices to global KV cache slots and count valid entries.
447 +
448 + Fuses three operations into a single kernel:
449 + 1. Block-table lookup (local index → global slot id)
450 + 2. Valid-entry counting (topk_lens per token)
451 + 3. Masking padding tokens to length 0
452 + """
453 + num_tokens = topk_indices.shape[0]
454 + global_topk_indices = torch.empty_like(topk_indices)
455 + topk_lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device)
456 + _compute_global_topk_indices_and_lens_kernel[(num_tokens,)](
457 + global_topk_indices,
458 + global_topk_indices.stride(0),
459 + topk_lens,
460 + topk_indices,
461 + topk_indices.stride(0),
462 + topk_indices.shape[-1],
463 + token_to_req_indices,
464 + block_table,
465 + block_table.stride(0),
466 + block_size,
467 + is_valid_token,
468 + TRITON_BLOCK_SIZE=1024,
469 + )
470 + return global_topk_indices, topk_lens
471 +
472 +
473 +@triton.jit
474 +def _compute_global_topk_indices_and_lens_kernel(
475 + global_topk_indices_ptr,
476 + global_topk_indices_stride: tl.constexpr,
477 + topk_lens_ptr,
478 + topk_indices_ptr,
479 + topk_indices_stride: tl.constexpr,
480 + topk: tl.constexpr,
481 + token_to_req_indices_ptr,
482 + block_table_ptr,
483 + block_table_stride: tl.constexpr,
484 + block_size: tl.constexpr,
485 + is_valid_token_ptr,
486 + TRITON_BLOCK_SIZE: tl.constexpr,
487 +):
488 + token_idx = tl.program_id(0)
489 + is_valid_token = tl.load(is_valid_token_ptr + token_idx)
490 + req_idx = tl.load(token_to_req_indices_ptr + token_idx)
491 +
492 + count = tl.zeros((), dtype=tl.int32)
493 + for i in range(0, topk, TRITON_BLOCK_SIZE):
494 + offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
495 + mask = offset < topk
496 +
497 + local_idx = tl.load(
498 + topk_indices_ptr + token_idx * topk_indices_stride + offset,
499 + mask=mask,
500 + other=-1,
501 + )
502 + is_valid = local_idx >= 0
503 +
504 + block_indices = local_idx // block_size
505 + block_numbers = tl.load(
506 + block_table_ptr + req_idx * block_table_stride + block_indices,
507 + mask=mask & is_valid,
508 + )
509 + block_offsets = local_idx % block_size
510 +
511 + slot_ids = block_numbers * block_size + block_offsets
512 + slot_ids = tl.where(is_valid, slot_ids, -1)
513 + tl.store(
514 + global_topk_indices_ptr + token_idx * global_topk_indices_stride + offset,
515 + slot_ids,
516 + mask=mask,
517 + )
518 + count += tl.sum(is_valid.to(tl.int32), axis=0)
519 +
520 + # Zero out length for padding tokens.
521 + tl.store(topk_lens_ptr + token_idx, tl.where(is_valid_token, count, 0))
522 +
523 +
524 +# FlashMLA sparse prefill asserts `params.topk % B_TOPK == 0` (see
525 +# flashmla/csrc/sm100/prefill/sparse/fwd/head{64,128}/phase1.cuh). B_TOPK is
526 +# 64 for the h_q=64 kernel and 128 for h_q=128; pad to 128 to satisfy both.
527 +# The extra slots stay as -1 sentinels and `combined_lens` caps the valid
528 +# range via `topk_length`, so padding is a no-op at kernel level.
529 +_SPARSE_PREFILL_TOPK_ALIGNMENT = 128
530 +
531 +
532 +def combine_topk_swa_indices(
533 + topk_indices: torch.Tensor,
534 + query_start_loc: torch.Tensor,
535 + seq_lens: torch.Tensor,
536 + gather_lens: torch.Tensor,
537 + window_size: int,
538 + compress_ratio: int,
539 + topk: int,
540 + M: int,
541 + N: int,
542 + out: tuple[torch.Tensor, torch.Tensor] | None = None,
543 + left_visible: torch.Tensor | None = None,
544 + right_visible: torch.Tensor | None = None,
545 + max_image_tokens: int = 0,
546 +) -> tuple[torch.Tensor, torch.Tensor]:
547 + num_tokens = topk_indices.shape[0]
548 + # max_image_tokens widens the SWA column region for in-image
549 + # bidirectional visibility; the width is fixed per model so the
550 + # caller-provided workspace stays valid on image-free batches.
551 + combined_topk = (
552 + (topk + window_size + max_image_tokens + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1)
553 + // _SPARSE_PREFILL_TOPK_ALIGNMENT
554 + * _SPARSE_PREFILL_TOPK_ALIGNMENT
555 + )
556 + if out is None:
557 + combined_indices = torch.full(
558 + (num_tokens, combined_topk),
559 + fill_value=-1,
560 + dtype=torch.int32,
561 + device=topk_indices.device,
562 + )
563 + combined_lens = torch.empty(
564 + num_tokens, dtype=torch.int32, device=topk_indices.device
565 + )
566 + else:
567 + combined_indices, combined_lens = out
568 +
569 + _COMBINE_TOPK_SWA_INDICES_KERNEL(
570 + combined_indices,
571 + combined_lens,
572 + topk_indices,
573 + query_start_loc,
574 + seq_lens,
575 + gather_lens,
576 + M,
577 + N,
578 + TOP_K=topk,
579 + COMPRESS_RATIO=compress_ratio,
580 + WINDOW_SIZE=window_size,
581 + left_visible=left_visible,
582 + right_visible=right_visible,
583 + max_image_tokens=max_image_tokens,
584 + )
585 + return combined_indices, combined_lens
586 +
587 +
588 +_COMBINE_TOPK_SWA_NUM_WORKERS = 256
589 +
590 +
591 +# Representative pointer alignment variants for Triton pointer specialization.
592 +_COMBINE_TOPK_SWA_POINTER_INPUTS = zip_inputs(
593 + dict(
594 + topk_indices=True,
595 + query_start_loc=True,
596 + seq_lens=True,
597 + gather_lens=True,
598 + ),
599 + dict(
600 + topk_indices=True,
601 + query_start_loc=False,
602 + seq_lens=False,
603 + gather_lens=True,
604 + ),
605 + dict(
606 + topk_indices=False,
607 + query_start_loc=False,
608 + seq_lens=False,
609 + gather_lens=False,
610 + ),
611 +)
612 +
613 +
614 +def _hf_config_int(vllm_config: Any, name: str, default: int) -> int:
615 + model_config = getattr(vllm_config, "model_config", None)
616 + hf_config = getattr(model_config, "hf_config", None)
617 + return int(getattr(hf_config, name, default) or default)
618 +
619 +
620 +def _combine_topk_swa_warmup_inputs(vllm_config: Any) -> _WarmupInputRows:
621 + """One warmup row per v4.1 layer type (compress ratio) in the model.
622 +
623 + Every layer slices the shared ``index_topk``-wide indices buffer, so the
624 + padded width is the same for all rows; SWA-only layers (ratio 0) pass
625 + ``TOP_K=0``.
626 + """
627 + model_config = getattr(vllm_config, "model_config", None)
628 + hf_config = getattr(model_config, "hf_config", None)
629 + compress_ratios = getattr(hf_config, "compress_ratios", None) or [0]
630 + topk = _hf_config_int(vllm_config, "index_topk", 0)
631 + return zip_inputs(
632 + *(
633 + dict(compress_ratio=ratio, topk=topk if ratio > 0 else 0, topk_width=topk)
634 + for ratio in sorted({int(ratio) for ratio in compress_ratios})
635 + )
636 + )
637 +
638 +
639 +def _scheduler_config_int(vllm_config: Any, name: str, default: int) -> int:
640 + scheduler_config = getattr(vllm_config, "scheduler_config", None)
641 + return int(getattr(scheduler_config, name, default) or default)
642 +
643 +
644 +class CombineTopkSwaIndicesKernel(
645 + VllmJitKernel["CombineTopkSwaIndicesKernel.CompileKey"]
646 +):
647 + @dataclass(frozen=True)
648 + class CompileKey:
649 + TOP_K: int
650 + COMPRESS_RATIO: int
651 + WINDOW_SIZE: int
652 + IMAGE_WIDTH: int
653 + PADDED_TOP_K: int
654 + input_variant: TritonPointerInputVariant
655 +
656 + @staticmethod
657 + @triton.jit(
658 + do_not_specialize=[
659 + "combined_indices_stride",
660 + "topk_indices_stride",
661 + "M",
662 + "N",
663 + ]
664 + )
665 + def kernel(
666 + combined_indices_ptr,
667 + combined_indices_stride,
668 + combined_lens_ptr,
669 + topk_indices_ptr,
670 + topk_indices_stride,
671 + query_start_loc_ptr,
672 + seq_lens_ptr,
673 + gather_lens_ptr,
674 + left_visible_ptr,
675 + right_visible_ptr,
676 + M,
677 + N,
678 + TOP_K: tl.constexpr,
679 + COMPRESS_RATIO: tl.constexpr,
680 + WINDOW_SIZE: tl.constexpr,
681 + IMAGE_WIDTH: tl.constexpr,
682 + PADDED_TOP_K: tl.constexpr,
683 + ):
684 + batch_idx = tl.program_id(0)
685 + worker_id = tl.program_id(1)
686 + num_workers = tl.num_programs(1)
687 +
688 + # query_start_loc is a global tensor; rebase to chunk-local offsets
689 + # by subtracting the chunk's starting value.
690 + base = tl.load(query_start_loc_ptr)
691 + query_start = tl.load(query_start_loc_ptr + batch_idx) - base
692 + query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base
693 + query_len = query_end - query_start
694 + seq_len = tl.load(seq_lens_ptr + batch_idx)
695 + gather_len = tl.load(gather_lens_ptr + batch_idx)
696 + start_pos = seq_len - query_len
697 + # The SWA portion of the gathered buffer starts from position
698 + # (seq_len - gather_len), not position 0. We need this offset
699 + # to correctly index into the gathered buffer.
700 + gather_start = seq_len - gather_len
701 +
702 + for token_idx in range(query_start + worker_id, query_end, num_workers):
703 + # topk_len is fully determined by the query token's absolute position:
704 + # the indexer emits min((pos + 1) // compress_ratio, topk_tokens)
705 + # valid entries. Caller passes TOP_K=0 for SWA-only layers to zero
706 + # this out.
707 + token_idx_in_query = token_idx - query_start
708 + pos = start_pos + token_idx_in_query
709 + if COMPRESS_RATIO > 0:
710 + topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
711 + else:
712 + # SWA-only layers pass TOP_K=0; skip the division entirely
713 + # (integer div by constexpr 0 is UB and yields garbage, which
714 + # turns the store index below negative).
715 + topk_len = 0
716 + if IMAGE_WIDTH > 0:
717 + # In-image bidirectional visibility: the window starts up to
718 + # max(left - (window - 1), 0) positions earlier and extends
719 + # `right` positions past the query token (both 0 outside
720 + # image spans, reducing to the plain causal window).
721 + left = tl.load(left_visible_ptr + token_idx)
722 + right = tl.load(right_visible_ptr + token_idx)
723 + else:
724 + left = 0
725 + right = 0
726 + left_add = tl.maximum(left - (WINDOW_SIZE - 1), 0)
727 + swa_start = tl.maximum(pos - (WINDOW_SIZE - 1) - left_add, 0)
728 + swa_len = pos + right - swa_start + 1
729 +
730 + offset = tl.arange(0, PADDED_TOP_K)
731 + mask = offset < topk_len
732 + topk_indices = tl.load(
733 + topk_indices_ptr + token_idx * topk_indices_stride + offset,
734 + mask=mask,
735 + )
736 + tl.store(
737 + combined_indices_ptr + token_idx * combined_indices_stride + offset,
738 + topk_indices + M * batch_idx,
739 + mask=mask,
740 + )
741 + # Index into gathered buffer: N + (position - gather_start)
742 + # For positions [swa_start, pos + right], the buffer indices are:
743 + # [N + swa_start - gather_start, N + pos + right - gather_start]
744 + for i in range(0, WINDOW_SIZE + IMAGE_WIDTH, WINDOW_SIZE):
745 + swa_offset = i + tl.arange(0, WINDOW_SIZE)
746 + tl.store(
747 + combined_indices_ptr
748 + + token_idx * combined_indices_stride
749 + + topk_len
750 + + swa_offset,
751 + M * batch_idx + N + swa_start + swa_offset - gather_start,
752 + mask=swa_offset < swa_len,
753 + )
754 +
755 + combined_len = topk_len + swa_len
756 + tl.store(combined_lens_ptr + token_idx, combined_len)
757 +
758 + def dispatch( # type: ignore[override]
759 + self,
760 + *,
761 + topk_width: int,
762 + topk_indices: bool,
763 + query_start_loc: bool,
764 + seq_lens: bool,
765 + gather_lens: bool,
766 + topk: int,
767 + compress_ratio: int,
768 + WINDOW_SIZE: int,
769 + image_width: int = 0,
770 + ) -> CompileKey:
771 + padded_topk = next_power_of_2(topk_width)
772 + input_variant = TritonPointerInputVariant.from_alignment(
773 + topk_indices=topk_indices,
774 + query_start_loc=query_start_loc,
775 + seq_lens=seq_lens,
776 + gather_lens=gather_lens,
777 + )
778 + return self.CompileKey(
779 + TOP_K=topk,
780 + COMPRESS_RATIO=compress_ratio,
781 + WINDOW_SIZE=WINDOW_SIZE,
782 + IMAGE_WIDTH=image_width,
783 + PADDED_TOP_K=padded_topk,
784 + input_variant=input_variant,
785 + )
786 +
787 + def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
788 + if _scheduler_config_int(vllm_config, "max_num_batched_tokens", 0) <= 0:
789 + return []
790 +
791 + window_size = _hf_config_int(vllm_config, "sliding_window", 128)
792 + image_width = (
793 + _hf_config_int(vllm_config, "vision_max_n_token", 0)
794 + if _hf_config_int(vllm_config, "vision_n_layers", 0) > 0
795 + else 0
796 + )
797 + # Warm both the plain-window variant (batches without image spans)
798 + # and the in-image bidirectional variant.
799 + image_widths = [0, image_width] if image_width > 0 else [0]
800 + return self._trace_dispatch(self.dispatch)(
801 + _combine_topk_swa_warmup_inputs(vllm_config),
802 + _COMBINE_TOPK_SWA_POINTER_INPUTS,
803 + WINDOW_SIZE=window_size,
804 + image_width=image_widths,
805 + )
806 +
807 + def compile(self, compile_key: CompileKey) -> None:
808 + warmup = getattr(self.kernel, "warmup", None)
809 + assert warmup is not None
810 + int32_ptr = TritonWarmupTensor(torch.int32)
811 + input_variant = compile_key.input_variant
812 + warmup(
813 + int32_ptr,
814 + 1, # do not specialize combined_indices_stride
815 + int32_ptr,
816 + input_variant.pointer("topk_indices", torch.int32),
817 + 1, # do not specialize topk_indices_stride
818 + input_variant.pointer("query_start_loc", torch.int32),
819 + input_variant.pointer("seq_lens", torch.int32),
820 + input_variant.pointer("gather_lens", torch.int32),
821 + int32_ptr,
822 + int32_ptr,
823 + 1, # do not specialize M
824 + 1, # do not specialize N
825 + TOP_K=compile_key.TOP_K,
826 + COMPRESS_RATIO=compile_key.COMPRESS_RATIO,
827 + WINDOW_SIZE=compile_key.WINDOW_SIZE,
828 + IMAGE_WIDTH=compile_key.IMAGE_WIDTH,
829 + PADDED_TOP_K=compile_key.PADDED_TOP_K,
830 + grid=(1, _COMBINE_TOPK_SWA_NUM_WORKERS),
831 + )
832 +
833 + def __call__(
834 + self,
835 + combined_indices: torch.Tensor,
836 + combined_lens: torch.Tensor,
837 + topk_indices: torch.Tensor,
838 + query_start_loc: torch.Tensor,
839 + seq_lens: torch.Tensor,
840 + gather_lens: torch.Tensor,
841 + M: int,
842 + N: int,
843 + *,
844 + TOP_K: int,
845 + COMPRESS_RATIO: int,
846 + WINDOW_SIZE: int,
847 + left_visible: torch.Tensor | None = None,
848 + right_visible: torch.Tensor | None = None,
849 + max_image_tokens: int = 0,
850 + ) -> None:
851 + num_reqs = seq_lens.shape[0]
852 + has_image = left_visible is not None
853 + image_width = max_image_tokens if has_image else 0
854 + self.kernel[(num_reqs, _COMBINE_TOPK_SWA_NUM_WORKERS)](
855 + combined_indices,
856 + combined_indices.stride(0),
857 + combined_lens,
858 + topk_indices,
859 + topk_indices.stride(0),
860 + query_start_loc,
861 + seq_lens,
862 + gather_lens,
863 + left_visible if has_image else topk_indices,
864 + right_visible if has_image else topk_indices,
865 + M,
866 + N,
867 + TOP_K=TOP_K,
868 + COMPRESS_RATIO=COMPRESS_RATIO,
869 + WINDOW_SIZE=WINDOW_SIZE,
870 + IMAGE_WIDTH=image_width,
871 + PADDED_TOP_K=next_power_of_2(topk_indices.shape[-1]),
872 + )
873 +
874 +
875 +_COMBINE_TOPK_SWA_INDICES_KERNEL = CombineTopkSwaIndicesKernel()
876 +
877 +
878 +def build_flashinfer_mixed_sparse_indices(
879 + decode_swa_indices: torch.Tensor,
880 + decode_compressed_indices: torch.Tensor | None,
881 + decode_compressed_topk_lens: torch.Tensor | None,
882 + prefill_topk_indices: torch.Tensor,
883 + query_start_loc: torch.Tensor,
884 + seq_lens: torch.Tensor,
885 + token_to_req_indices: torch.Tensor,
886 + swa_block_table: torch.Tensor,
887 + swa_block_size: int,
888 + compressed_block_table: torch.Tensor | None,
889 + compressed_block_size: int,
890 + window_size: int,
891 + compress_ratio: int,
892 + topk: int,
893 + decode_compressed_indices_are_local: bool = False,
894 + decode_is_valid_token: torch.Tensor | None = None,
895 + swa_block_span: int | None = None,
896 + compressed_block_span: int | None = None,
897 + prefill_left_visible: torch.Tensor | None = None,
898 + prefill_right_visible: torch.Tensor | None = None,
899 + max_image_tokens: int = 0,
900 +) -> tuple[torch.Tensor, torch.Tensor]:
901 + """Build the FlashInfer DSV4 sparse-index matrix for decode-first batches.
902 +
903 + Produces ``sparse_indices`` of shape ``[num_tokens, swa_total_width +
904 + padded_topk]`` (the first ``swa_total_width`` columns are SWA slot ids, the
905 + rest are compressed/top-k slot ids) and ``sparse_topk_lens`` (active length
906 + per token). Decode tokens read precomputed SWA/compressed indices; prefill
907 + tokens derive their SWA window from the position and translate local
908 + compressed indices to global slots via the block tables.
909 +
910 + When ``prefill_left_visible``/``prefill_right_visible`` are given (vision
911 + variant), the SWA column region widens by ``max_image_tokens`` and prefill
912 + tokens inside an image span get a bidirectionally widened window; decode
913 + rows are padded with -1 across the extra columns.
914 + """
915 + assert decode_swa_indices.dtype == torch.int32
916 + assert decode_swa_indices.dim() == 2
917 + swa_index_width = decode_swa_indices.shape[-1]
918 + assert swa_index_width >= window_size
919 + has_image = prefill_left_visible is not None
920 + image_width = max_image_tokens if has_image else 0
921 + swa_total_width = swa_index_width + image_width
922 + if has_image:
923 + assert prefill_right_visible is not None
924 + if decode_compressed_topk_lens is not None:
925 + assert decode_compressed_topk_lens.dtype == torch.int32
926 + assert prefill_topk_indices.dtype == torch.int32
927 + assert prefill_topk_indices.dim() == 2
928 + assert query_start_loc.dtype == torch.int32
929 + assert seq_lens.dtype == torch.int32
930 + assert token_to_req_indices.dtype == torch.int32
931 + assert swa_block_table.dtype == torch.int32
932 +
933 + num_decode_tokens = decode_swa_indices.shape[0]
934 + num_prefill_tokens = prefill_topk_indices.shape[0]
935 + num_tokens = num_decode_tokens + num_prefill_tokens
936 + assert token_to_req_indices.shape[0] >= num_tokens
937 + if decode_compressed_topk_lens is not None:
938 + assert decode_compressed_topk_lens.shape[0] >= num_decode_tokens
939 +
940 + decode_compressed_topk = 0
941 + if decode_compressed_indices is None:
942 + decode_compressed_indices = prefill_topk_indices
943 + else:
944 + assert decode_compressed_indices.dtype == torch.int32
945 + assert decode_compressed_indices.dim() == 2
946 + assert decode_compressed_indices.shape[0] == num_decode_tokens
947 + decode_compressed_topk = decode_compressed_indices.shape[-1]
948 + if decode_compressed_topk > 0 and decode_compressed_indices_are_local:
949 + assert decode_is_valid_token is not None
950 + assert decode_is_valid_token.dtype == torch.bool
951 + assert decode_is_valid_token.shape[0] >= num_decode_tokens
952 + else:
953 + decode_is_valid_token = token_to_req_indices
954 +
955 + if compressed_block_table is None:
956 + compressed_block_table = swa_block_table
957 + assert compressed_block_table.dtype == torch.int32
958 + has_decode_compressed_lens = decode_compressed_topk_lens is not None
959 + if decode_compressed_topk_lens is None:
960 + decode_compressed_topk_lens = token_to_req_indices
961 +
962 + # The FlashInfer TRTLLM-gen sparse-MLA kernels require every per-token topk
963 + # index row to start on a 16-byte boundary: the kernel loads the compressed
964 + # indices with 128-bit (16-byte) vectorized loads, so a misaligned row would
965 + # fault or read across rows. 16 bytes = 4 int32 indices, so round the topk
966 + # width (and hence the row stride, since the SWA columns are fixed-width) up
967 + # to a multiple of 4. The extra columns are filled with -1 (invalid) and bounded
968 + # by ``sparse_topk_lens``, so padding never changes the attention result.
969 + padded_topk = max(topk, decode_compressed_topk)
970 + padded_topk = (padded_topk + 3) // 4 * 4
971 + sparse_indices = torch.empty(
972 + (num_tokens, swa_total_width + padded_topk),
973 + dtype=torch.int32,
974 + device=decode_swa_indices.device,
975 + )
976 + sparse_topk_lens = torch.empty(
977 + num_tokens, dtype=torch.int32, device=decode_swa_indices.device
978 + )
979 + if num_tokens == 0:
980 + return sparse_indices, sparse_topk_lens
981 +
982 + window_block_size = triton.next_power_of_2(max(swa_index_width, 1))
983 + topk_block_size = triton.next_power_of_2(max(padded_topk, 1))
984 + max_block_size = max(window_block_size, topk_block_size)
985 + num_warps = 4 if max_block_size >= 256 else 1
986 +
987 + # block_span = page_stride / token_stride; == block_size (no-op) for unpacked KV.
988 + swa_span = swa_block_size if swa_block_span is None else swa_block_span
989 + compressed_span = (
990 + compressed_block_size
991 + if compressed_block_span is None
992 + else compressed_block_span
993 + )
994 + _build_flashinfer_mixed_sparse_indices_kernel[(num_tokens,)](
995 + sparse_indices,
996 + sparse_indices.stride(0),
997 + sparse_topk_lens,
998 + decode_swa_indices,
999 + decode_swa_indices.stride(0),
1000 + decode_compressed_indices,
1001 + decode_compressed_indices.stride(0),
1002 + decode_compressed_topk_lens,
1003 + decode_is_valid_token,
1004 + prefill_topk_indices,
1005 + prefill_topk_indices.stride(0),
1006 + prefill_left_visible if has_image else token_to_req_indices,
1007 + prefill_right_visible if has_image else token_to_req_indices,
1008 + query_start_loc,
1009 + seq_lens,
1010 + token_to_req_indices,
1011 + swa_block_table,
1012 + swa_block_table.stride(0),
1013 + swa_block_size,
1014 + swa_span,
1015 + compressed_block_table,
1016 + compressed_block_table.stride(0),
1017 + compressed_block_size,
1018 + compressed_span,
1019 + NUM_DECODE_TOKENS=num_decode_tokens,
1020 + WINDOW_SIZE=window_size,
1021 + SWA_INDEX_WIDTH=swa_index_width,
1022 + IMAGE_WIDTH=image_width,
1023 + COMPRESS_RATIO=compress_ratio,
1024 + TOP_K=topk,
1025 + PADDED_TOP_K=padded_topk,
1026 + PREFILL_TOPK_STRIDE=prefill_topk_indices.shape[-1],
1027 + DECODE_COMPRESSED_TOPK=decode_compressed_topk,
1028 + DECODE_COMPRESSED_INDICES_ARE_LOCAL=decode_compressed_indices_are_local,
1029 + HAS_DECODE_COMPRESSED_LENS=has_decode_compressed_lens,
1030 + WINDOW_BLOCK_SIZE=window_block_size,
1031 + TOPK_BLOCK_SIZE=topk_block_size,
1032 + num_warps=num_warps,
1033 + )
1034 + return sparse_indices, sparse_topk_lens
1035 +
1036 +
1037 +@triton.jit
1038 +def _remap_flashinfer_index(values, block_size, block_span):
1039 + # FlashInfer's DSv4 kernel indexes sparse KV by physical token stride, so
1040 + # packed pages (#44577) need block*block_size+off -> block*block_span+off.
1041 + # TODO: remove once flashinfer-ai/flashinfer#3856 is fixed.
1042 + is_valid = values >= 0
1043 + safe_values = tl.where(is_valid, values, 0)
1044 + values = (safe_values // block_size) * block_span
1045 + values += safe_values % block_size
1046 + return tl.where(is_valid, values, -1)
1047 +
1048 +
1049 +@triton.jit(
1050 + do_not_specialize=[
1051 + "sparse_indices_stride",
1052 + "decode_swa_stride",
1053 + "decode_compressed_stride",
1054 + "prefill_topk_stride",
1055 + "swa_block_table_stride",
1056 + "swa_block_size",
1057 + "swa_block_span",
1058 + "compressed_block_table_stride",
1059 + "compressed_block_size",
1060 + "compressed_block_span",
1061 + "NUM_DECODE_TOKENS",
1062 + "PREFILL_TOPK_STRIDE",
1063 + ]
1064 +)
1065 +def _build_flashinfer_mixed_sparse_indices_kernel(
1066 + sparse_indices_ptr,
1067 + sparse_indices_stride,
1068 + sparse_topk_lens_ptr,
1069 + decode_swa_indices_ptr,
1070 + decode_swa_stride,
1071 + decode_compressed_indices_ptr,
1072 + decode_compressed_stride,
1073 + decode_compressed_topk_lens_ptr,
1074 + decode_is_valid_token_ptr,
1075 + prefill_topk_indices_ptr,
1076 + prefill_topk_stride,
1077 + left_visible_ptr,
1078 + right_visible_ptr,
1079 + query_start_loc_ptr,
1080 + seq_lens_ptr,
1081 + token_to_req_indices_ptr,
1082 + swa_block_table_ptr,
1083 + swa_block_table_stride,
1084 + swa_block_size,
1085 + swa_block_span,
1086 + compressed_block_table_ptr,
1087 + compressed_block_table_stride,
1088 + compressed_block_size,
1089 + compressed_block_span,
1090 + NUM_DECODE_TOKENS,
1091 + WINDOW_SIZE: tl.constexpr,
1092 + SWA_INDEX_WIDTH: tl.constexpr,
1093 + IMAGE_WIDTH: tl.constexpr,
1094 + COMPRESS_RATIO: tl.constexpr,
1095 + TOP_K: tl.constexpr,
1096 + PADDED_TOP_K: tl.constexpr,
1097 + PREFILL_TOPK_STRIDE,
1098 + DECODE_COMPRESSED_TOPK: tl.constexpr,
1099 + DECODE_COMPRESSED_INDICES_ARE_LOCAL: tl.constexpr,
1100 + HAS_DECODE_COMPRESSED_LENS: tl.constexpr,
1101 + WINDOW_BLOCK_SIZE: tl.constexpr,
1102 + TOPK_BLOCK_SIZE: tl.constexpr,
1103 +):
1104 + token_idx = tl.program_id(0)
1105 + # Total SWA column count; > SWA_INDEX_WIDTH only for in-image
1106 + # bidirectional visibility (vision variant), where the extra columns are
1107 + # -1-padded for decode rows.
1108 + SWA_TOTAL_WIDTH: tl.constexpr = SWA_INDEX_WIDTH + IMAGE_WIDTH
1109 +
1110 + if token_idx < NUM_DECODE_TOKENS:
1111 + for i in range(0, SWA_TOTAL_WIDTH, WINDOW_BLOCK_SIZE):
1112 + offset = i + tl.arange(0, WINDOW_BLOCK_SIZE)
1113 + mask = offset < SWA_TOTAL_WIDTH
1114 + values = tl.load(
1115 + decode_swa_indices_ptr + token_idx * decode_swa_stride + offset,
1116 + mask=offset < SWA_INDEX_WIDTH,
1117 + other=-1,
1118 + )
1119 + values = _remap_flashinfer_index(values, swa_block_size, swa_block_span)
1120 + tl.store(
1121 + sparse_indices_ptr + token_idx * sparse_indices_stride + offset,
1122 + values,
1123 + mask=mask,
1124 + )
1125 +
1126 + compressed_len = tl.zeros((), dtype=tl.int32)
1127 + for i in range(0, PADDED_TOP_K, TOPK_BLOCK_SIZE):
1128 + offset = i + tl.arange(0, TOPK_BLOCK_SIZE)
1129 + mask = offset < PADDED_TOP_K
1130 + values = tl.load(
1131 + decode_compressed_indices_ptr
1132 + + token_idx * decode_compressed_stride
1133 + + offset,
1134 + mask=offset < DECODE_COMPRESSED_TOPK,
1135 + other=-1,
1136 + )
1137 + if DECODE_COMPRESSED_INDICES_ARE_LOCAL:
1138 + token_valid = tl.load(decode_is_valid_token_ptr + token_idx)
1139 + is_valid = values >= 0
1140 + req_idx = tl.load(token_to_req_indices_ptr + token_idx)
1141 + block_indices = values // compressed_block_size
1142 + block_numbers = tl.load(
1143 + compressed_block_table_ptr
1144 + + req_idx * compressed_block_table_stride
1145 + + block_indices,
1146 + mask=mask & is_valid,
1147 + other=-1,
1148 + )
1149 + block_offsets = values % compressed_block_size
1150 + values = block_numbers * compressed_block_size + block_offsets
1151 + values = tl.where(is_valid, values, -1)
1152 + compressed_len += tl.sum((is_valid & token_valid).to(tl.int32), axis=0)
1153 + values = _remap_flashinfer_index(
1154 + values, compressed_block_size, compressed_block_span
1155 + )
1156 + tl.store(
1157 + sparse_indices_ptr
1158 + + token_idx * sparse_indices_stride
1159 + + SWA_TOTAL_WIDTH
1160 + + offset,
1161 + values,
1162 + mask=mask,
1163 + )
1164 +
1165 + if DECODE_COMPRESSED_TOPK == 0:
1166 + compressed_len = tl.zeros((), dtype=tl.int32)
1167 + elif not DECODE_COMPRESSED_INDICES_ARE_LOCAL:
1168 + if HAS_DECODE_COMPRESSED_LENS:
1169 + compressed_len = tl.load(decode_compressed_topk_lens_ptr + token_idx)
1170 + else:
1171 + compressed_len = tl.full((), DECODE_COMPRESSED_TOPK, dtype=tl.int32)
1172 +
1173 + tl.store(sparse_topk_lens_ptr + token_idx, SWA_TOTAL_WIDTH + compressed_len)
1174 + return
1175 +
1176 + prefill_idx = token_idx - NUM_DECODE_TOKENS
1177 + req_idx = tl.load(token_to_req_indices_ptr + token_idx)
1178 + query_start = tl.load(query_start_loc_ptr + req_idx)
1179 + query_end = tl.load(query_start_loc_ptr + req_idx + 1)
1180 + query_len = query_end - query_start
1181 + seq_len = tl.load(seq_lens_ptr + req_idx)
1182 + start_pos = seq_len - query_len
1183 + token_idx_in_query = token_idx - query_start
1184 + pos = start_pos + token_idx_in_query
1185 + if IMAGE_WIDTH > 0:
1186 + # In-image bidirectional visibility: window starts up to
1187 + # max(left - (window - 1), 0) earlier and extends `right` past pos.
1188 + left = tl.load(left_visible_ptr + token_idx)
1189 + right = tl.load(right_visible_ptr + token_idx)
1190 + else:
1191 + left = 0
1192 + right = 0
1193 + left_add = tl.maximum(left - (WINDOW_SIZE - 1), 0)
1194 + swa_start_pos = tl.maximum(pos - (WINDOW_SIZE - 1) - left_add, 0)
1195 + swa_len = pos + right - swa_start_pos + 1
1196 + if COMPRESS_RATIO > 0:
1197 + topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
1198 + else:
1199 + # SWA-only layers (ratio 0): no indexer topk; div by constexpr 0 is UB.
1200 + topk_len = 0
1201 +
1202 + for i in range(0, SWA_TOTAL_WIDTH, WINDOW_BLOCK_SIZE):
1203 + offset = i + tl.arange(0, WINDOW_BLOCK_SIZE)
1204 + mask = offset < SWA_TOTAL_WIDTH
1205 + pos_offset = swa_start_pos + offset
1206 + block_indices = pos_offset // swa_block_size
1207 + block_numbers = tl.load(
1208 + swa_block_table_ptr + req_idx * swa_block_table_stride + block_indices,
1209 + mask=mask & (offset < swa_len),
1210 + other=-1,
1211 + )
1212 + block_offsets = pos_offset % swa_block_size
1213 + slot_ids = block_numbers * swa_block_size + block_offsets
1214 + slot_ids = tl.where(offset < swa_len, slot_ids, -1)
1215 + slot_ids = _remap_flashinfer_index(slot_ids, swa_block_size, swa_block_span)
1216 + tl.store(
1217 + sparse_indices_ptr + token_idx * sparse_indices_stride + offset,
1218 + slot_ids,
1219 + mask=mask,
1220 + )
1221 +
1222 + for i in range(0, PADDED_TOP_K, TOPK_BLOCK_SIZE):
1223 + offset = i + tl.arange(0, TOPK_BLOCK_SIZE)
1224 + mask = offset < PADDED_TOP_K
1225 + local_idx = tl.load(
1226 + prefill_topk_indices_ptr + prefill_idx * prefill_topk_stride + offset,
1227 + mask=(offset < PREFILL_TOPK_STRIDE) & (offset < topk_len),
1228 + other=-1,
1229 + )
1230 + is_valid = local_idx >= 0
1231 + block_indices = local_idx // compressed_block_size
1232 + block_numbers = tl.load(
1233 + compressed_block_table_ptr
1234 + + req_idx * compressed_block_table_stride
1235 + + block_indices,
1236 + mask=mask & is_valid,
1237 + other=-1,
1238 + )
1239 + block_offsets = local_idx % compressed_block_size
1240 + slot_ids = block_numbers * compressed_block_size + block_offsets
1241 + slot_ids = tl.where((offset < topk_len) & is_valid, slot_ids, -1)
1242 + slot_ids = _remap_flashinfer_index(
1243 + slot_ids, compressed_block_size, compressed_block_span
1244 + )
1245 + tl.store(
1246 + sparse_indices_ptr
1247 + + token_idx * sparse_indices_stride
1248 + + SWA_TOTAL_WIDTH
1249 + + offset,
1250 + slot_ids,
1251 + mask=mask,
1252 + )
1253 +
1254 + tl.store(sparse_topk_lens_ptr + token_idx, SWA_TOTAL_WIDTH + topk_len)
vllm/models/deepseek_v4_1/common/ops/fused_compress_quant_cache.py new
+383
@@ -0,0 +1,383 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""V4.1 state saving/compression and independently schedulable cache insertion."""
4 +
5 +import torch
6 +
7 +from vllm.platforms import current_platform
8 +from vllm.triton_utils import tl, triton
9 +
10 +if current_platform.is_rocm():
11 + from vllm.platforms.rocm import _ON_GFX950
12 +else:
13 + _ON_GFX950 = False
14 +
15 +# The ring and raw rows a ratio-2 request program pools are not adjacent, so
16 +# reading them as one tile means joining two pointers. ROCm's Triton fails to
17 +# legalize `tt.join` on pointers, so there the two rows are loaded separately.
18 +_JOIN_ROW_PTRS = not current_platform.is_rocm()
19 +
20 +
21 +def fused_save_compress_norm(
22 + kv_score: torch.Tensor,
23 + positions: torch.Tensor,
24 + state_cache: torch.Tensor | None,
25 + slot_mapping: torch.Tensor,
26 + query_start_loc: torch.Tensor | None,
27 + token_to_req_indices: torch.Tensor | None,
28 + rms_norm_weight: torch.Tensor,
29 + rms_norm_eps: float,
30 + compress_ratio: int,
31 + latent_out: torch.Tensor,
32 +) -> None:
33 + """Pool each closed group into a normalized BF16 latent; save FP32 states.
34 +
35 + The latent feeds the main-cache insert and the indexer K path, which the
36 + attention layer schedules on separate streams.
37 +
38 + Ratio 2 keeps one ring block per request holding the open group's rows:
39 + position ``p`` lives in row ``p % capacity`` and ``slot_mapping`` encodes
40 + ``block * capacity + p % capacity``. The grid has one program per request
41 + followed by one per pair of packed tokens. A request program handles the
42 + group that the chunk's first token closes with its predecessor's ring row,
43 + then stores the chunk's last ``capacity`` rows to the ring; because the
44 + same program does both, ring reads and writes never race. A pair program
45 + handles the group that ends inside its pair, reading both rows from the
46 + raw input. Ratio 1 has no ring and one program per token; ``slot_mapping``
47 + then only marks valid tokens.
48 +
49 + Args:
50 + kv_score: FP32 [tokens, 512] for CR1, [tokens, 1024] for CR2.
51 + positions: Absolute positions of the packed request tokens.
52 + state_cache: Ring FP32 [blocks, capacity, 1024] KV/score states (CR2).
53 + slot_mapping: Ring slots (CR2) or main-cache slots (CR1).
54 + query_start_loc: [num_reqs + 1] token offsets of each request's chunk.
55 + token_to_req_indices: Request indices for the packed token rows.
56 + rms_norm_weight: BF16 [512] normalization weight.
57 + rms_norm_eps: RMSNorm epsilon.
58 + compress_ratio: Group size, either 1 or 2.
59 + latent_out: BF16 [tokens, 512], written only at valid group boundaries.
60 + """
61 + assert compress_ratio in (1, 2)
62 + assert kv_score.dtype == torch.float32
63 + assert kv_score.shape[1] == 512 * compress_ratio and kv_score.stride(1) == 1
64 + assert latent_out.shape == (kv_score.shape[0], 512)
65 + assert latent_out.is_contiguous() and latent_out.dtype == torch.bfloat16
66 + assert positions.is_contiguous() and slot_mapping.is_contiguous()
67 + # Rows stay 64-byte aligned, as the kernel's tl.multiple_of hint promises.
68 + assert kv_score.stride(0) % 16 == 0
69 + if compress_ratio == 2:
70 + assert state_cache is not None and query_start_loc is not None
71 + assert token_to_req_indices is not None
72 + assert query_start_loc.is_contiguous()
73 + assert token_to_req_indices.is_contiguous()
74 + assert state_cache.dtype == torch.float32
75 + assert state_cache.shape[2] == 1024 and state_cache.stride(2) == 1
76 + assert state_cache.stride(1) % 16 == 0
77 + state_stride, state_row_stride, state_block = (
78 + state_cache.stride(0),
79 + state_cache.stride(1),
80 + state_cache.shape[1],
81 + )
82 + num_reqs = query_start_loc.numel() - 1
83 + else:
84 + state_cache = query_start_loc = token_to_req_indices = None
85 + state_stride = state_row_stride = state_block = 1
86 + num_reqs = 0
87 + num_tokens = slot_mapping.numel()
88 + assert num_tokens <= min(kv_score.shape[0], positions.numel())
89 + if num_tokens == 0:
90 + return
91 + grid = num_reqs + triton.cdiv(num_tokens, compress_ratio)
92 + _fused_save_compress_norm_kernel[(grid,)](
93 + kv_score,
94 + positions,
95 + state_cache,
96 + slot_mapping,
97 + query_start_loc,
98 + token_to_req_indices,
99 + rms_norm_weight,
100 + latent_out,
101 + num_tokens,
102 + num_reqs,
103 + RAW_STRIDE=kv_score.stride(0),
104 + STATE_STRIDE=state_stride,
105 + STATE_ROW_STRIDE=state_row_stride,
106 + STATE_BLOCK=state_block,
107 + COMPRESS_RATIO=compress_ratio,
108 + EPS=rms_norm_eps,
109 + JOIN_ROW_PTRS=_JOIN_ROW_PTRS,
110 + num_warps=4,
111 + **({"launch_pdl": False} if current_platform.is_cuda() else {}),
112 + )
113 +
114 +
115 +@triton.jit(do_not_specialize=["num_tokens", "num_reqs"])
116 +def _fused_save_compress_norm_kernel(
117 + raw,
118 + positions,
119 + state,
120 + state_slots,
121 + query_start_loc,
122 + req_ids,
123 + norm_weight,
124 + latent,
125 + num_tokens,
126 + num_reqs,
127 + RAW_STRIDE: tl.constexpr,
128 + STATE_STRIDE: tl.constexpr,
129 + STATE_ROW_STRIDE: tl.constexpr,
130 + STATE_BLOCK: tl.constexpr,
131 + COMPRESS_RATIO: tl.constexpr,
132 + EPS: tl.constexpr,
133 + JOIN_ROW_PTRS: tl.constexpr,
134 +):
135 + pid = tl.program_id(0)
136 + d = tl.arange(0, 512)
137 +
138 + if COMPRESS_RATIO == 2: # noqa: SIM102 (constexpr branch, then runtime)
139 + if pid < num_reqs:
140 + # Request program: the group closed by the chunk's first token,
141 + # then the ring tail store.
142 + start = tl.load(query_start_loc + pid)
143 + end = tl.load(query_start_loc + pid + 1)
144 + if start >= end:
145 + return
146 + state_slot = tl.load(state_slots + start)
147 + if state_slot < 0:
148 + return
149 + position = tl.load(positions + start)
150 + if (position + 1) % 2 == 0:
151 + ring = state + (state_slot // STATE_BLOCK).to(tl.int64) * STATE_STRIDE
152 + prev = ring + ((position - 1) % STATE_BLOCK) * STATE_ROW_STRIDE
153 + current = raw + start.to(tl.int64) * RAW_STRIDE
154 + if JOIN_ROW_PTRS:
155 + # Joining pointers hides row alignment from Triton; restate it.
156 + rows = tl.multiple_of(tl.join(prev, current)[:, None], (16, 16))
157 + kv = tl.load(rows + d[None, :])
158 + score = tl.load(rows + 512 + d[None, :])
159 + # Every lane has read the ring before the tail store below
160 + # writes it.
161 + tl.debug_barrier()
162 + pooled = tl.sum(kv * tl.softmax(score, 0), 0)
163 + else:
164 + kv_prev = tl.load(prev + d)
165 + score_prev = tl.load(prev + 512 + d)
166 + kv_current = tl.load(current + d)
167 + score_current = tl.load(current + 512 + d)
168 + # Every lane has read the ring before the tail store below
169 + # writes it.
170 + tl.debug_barrier()
171 + peak = tl.maximum(score_prev, score_current)
172 + weight_prev = tl.exp(score_prev - peak)
173 + weight_current = tl.exp(score_current - peak)
174 + pooled = (kv_prev * weight_prev + kv_current * weight_current) / (
175 + weight_prev + weight_current
176 + )
177 + _store_latent(pooled, start, norm_weight, latent, EPS)
178 + num_rows = tl.minimum(end - start, STATE_BLOCK)
179 + for k in tl.range(0, num_rows):
180 + token = end - num_rows + k
181 + slot = tl.load(state_slots + token)
182 + if slot >= 0:
183 + row = (
184 + state
185 + + (slot // STATE_BLOCK).to(tl.int64) * STATE_STRIDE
186 + + (slot % STATE_BLOCK) * STATE_ROW_STRIDE
187 + )
188 + src = raw + token.to(tl.int64) * RAW_STRIDE
189 + tl.store(row + d, tl.load(src + d))
190 + tl.store(row + 512 + d, tl.load(src + 512 + d))
191 + return
192 +
193 + # Group program: one token (ratio 1) or the token of the pair that ends
194 + # a group (ratio 2). A chunk's first token belongs to its request program.
195 + t = (pid - num_reqs) * COMPRESS_RATIO
196 + if COMPRESS_RATIO == 2:
197 + t += (tl.load(positions + t) + 1) % 2
198 + state_slot = tl.load(state_slots + t, t < num_tokens, other=-1)
199 + if state_slot < 0:
200 + return
201 + if COMPRESS_RATIO == 1:
202 + pooled = tl.load(raw + t.to(tl.int64) * RAW_STRIDE + d)
203 + else:
204 + if t == tl.load(query_start_loc + tl.load(req_ids + t)):
205 + return
206 + rows = raw + (t.to(tl.int64) - 1 + tl.arange(0, 2)[:, None]) * RAW_STRIDE
207 + kv = tl.load(rows + d[None, :])
208 + score = tl.load(rows + 512 + d[None, :])
209 + pooled = tl.sum(kv * tl.softmax(score, 0), 0)
210 + _store_latent(pooled, t, norm_weight, latent, EPS)
211 +
212 +
213 +@triton.jit
214 +def _store_latent(pooled, t, norm_weight, latent, EPS: tl.constexpr):
215 + d = tl.arange(0, 512)
216 + weight = tl.load(norm_weight + d).to(tl.float32)
217 + variance = tl.sum(pooled * pooled, 0) / 512
218 + normed = pooled * tl.rsqrt(variance + EPS) * weight
219 + tl.store(latent + t.to(tl.int64) * 512 + d, normed.to(tl.bfloat16))
220 +
221 +
222 +def rope_quant_insert(
223 + latent: torch.Tensor,
224 + positions: torch.Tensor,
225 + cos_sin_cache: torch.Tensor,
226 + kv_cache: torch.Tensor,
227 + slot_mapping: torch.Tensor,
228 + compress_ratio: int,
229 + fp8_scale: torch.Tensor | None = None,
230 +) -> None:
231 + """Apply GPT-J RoPE and publish a latent to the compressed KV cache.
232 +
233 + The BF16 latent supplies both NoPE quantization and RoPE input. It is read
234 + only for valid slots at group boundaries. The cache dtype selects the
235 + layout: ``uint8`` is the fp8_ds_mla paged layout (576 value bytes and eight
236 + segregated UE8M0 scale bytes per token, including one zero padding scale);
237 + ``bfloat16`` and ``float8_e4m3fn`` are the plain [448 NoPE | 64 RoPE] rows
238 + read by FlashInfer, the latter scaled by the per-tensor ``fp8_scale``.
239 + """
240 + assert compress_ratio in (1, 2)
241 + assert latent.shape[1] == 512 and latent.dtype == torch.bfloat16
242 + assert latent.is_contiguous()
243 + num_tokens = slot_mapping.numel()
244 + assert num_tokens <= min(latent.shape[0], positions.numel())
245 + if num_tokens == 0:
246 + return
247 + launch_kwargs = {"launch_pdl": False} if current_platform.is_cuda() else {}
248 + if kv_cache.dtype == torch.uint8:
249 + assert kv_cache.shape[-1] == 584
250 + _rope_quant_insert_kernel[(num_tokens,)](
251 + latent,
252 + positions,
253 + cos_sin_cache,
254 + kv_cache,
255 + slot_mapping,
256 + COS_STRIDE=cos_sin_cache.stride(0),
257 + CACHE_STRIDE=kv_cache.stride(0),
258 + CACHE_BLOCK=kv_cache.shape[1],
259 + COMPRESS_RATIO=compress_ratio,
260 + SANITIZE_CACHE_NANS=_ON_GFX950,
261 + num_warps=4,
262 + **launch_kwargs,
263 + )
264 + return
265 +
266 + assert kv_cache.dtype in (torch.bfloat16, torch.float8_e4m3fn)
267 + assert kv_cache.shape[-1] == 512 and kv_cache.stride(-1) == 1
268 + store_fp8 = kv_cache.dtype == torch.float8_e4m3fn
269 + if store_fp8:
270 + assert fp8_scale is not None and fp8_scale.numel() == 1
271 + assert fp8_scale.dtype == torch.float32
272 + _rope_plain_insert_kernel[(num_tokens,)](
273 + latent,
274 + positions,
275 + cos_sin_cache,
276 + kv_cache,
277 + slot_mapping,
278 + fp8_scale if store_fp8 else None,
279 + COS_STRIDE=cos_sin_cache.stride(0),
280 + CACHE_STRIDE=kv_cache.stride(0),
281 + ROW_STRIDE=kv_cache.stride(1),
282 + CACHE_BLOCK=kv_cache.shape[1],
283 + COMPRESS_RATIO=compress_ratio,
284 + STORE_FP8=store_fp8,
285 + num_warps=4,
286 + **launch_kwargs,
287 + )
288 +
289 +
290 +@triton.jit
291 +def _rope_quant_insert_kernel(
292 + latent,
293 + positions,
294 + cos_sin,
295 + cache,
296 + cache_slots,
297 + COS_STRIDE: tl.constexpr,
298 + CACHE_STRIDE: tl.constexpr,
299 + CACHE_BLOCK: tl.constexpr,
300 + COMPRESS_RATIO: tl.constexpr,
301 + SANITIZE_CACHE_NANS: tl.constexpr,
302 +):
303 + t = tl.program_id(0)
304 + slot = tl.load(cache_slots + t)
305 + if slot < 0:
306 + return
307 + position = tl.load(positions + t)
308 + if (position + 1) % COMPRESS_RATIO != 0:
309 + return
310 + d = tl.arange(0, 512)
311 + normed = tl.load(latent + t.to(tl.int64) * 512 + d).to(tl.float32)
312 + page = cache + (slot // CACHE_BLOCK).to(tl.int64) * CACHE_STRIDE
313 + values = page + (slot % CACHE_BLOCK) * 576
314 + scales = page + CACHE_BLOCK * 576 + (slot % CACHE_BLOCK) * 8
315 +
316 + quant = tl.reshape(normed, (8, 64))
317 + amax = tl.maximum(tl.max(tl.abs(quant), 1), 1e-4)
318 + exponent = tl.ceil(tl.log2(amax * (1.0 / 448.0)))
319 + scaled = quant * tl.reshape(tl.exp2(-exponent), (8, 1))
320 + fp8 = tl.clamp(scaled, -448.0, 448.0).to(tl.float8e4nv)
321 + packed = tl.reshape(fp8.to(tl.uint8, bitcast=True), (512,))
322 + tl.store(values + d, packed, d < 448)
323 + s = tl.arange(0, 8)
324 + max_encoded: tl.constexpr = 254.0 if SANITIZE_CACHE_NANS else 255.0
325 + encoded = tl.minimum(tl.maximum(exponent + 127.0, 0.0), max_encoded)
326 + tl.store(scales + s, encoded.to(tl.uint8), s < 7)
327 + tl.store(scales + 7, tl.full((), 0, tl.uint8))
328 +
329 + even, odd = tl.split(tl.reshape(normed, (256, 2)))
330 + pair = tl.arange(0, 256) - 224
331 + cs = cos_sin + (position // COMPRESS_RATIO * COMPRESS_RATIO) * COS_STRIDE
332 + c = tl.load(cs + tl.maximum(pair, 0), pair >= 0, other=1.0).to(tl.float32)
333 + s = tl.load(cs + 32 + tl.maximum(pair, 0), pair >= 0, other=0.0).to(tl.float32)
334 + rotated = tl.interleave(even * c - odd * s, odd * c + even * s)
335 + if SANITIZE_CACHE_NANS:
336 + rotated = tl.where(rotated == rotated, rotated, 0.0)
337 + rope_dst = (values + 448).to(tl.pointer_type(tl.bfloat16))
338 + tl.store(rope_dst + d - 448, rotated.to(tl.bfloat16), d >= 448)
339 +
340 +
341 +@triton.jit
342 +def _rope_plain_insert_kernel(
343 + latent,
344 + positions,
345 + cos_sin,
346 + cache,
347 + cache_slots,
348 + fp8_scale,
349 + COS_STRIDE: tl.constexpr,
350 + CACHE_STRIDE: tl.constexpr,
351 + ROW_STRIDE: tl.constexpr,
352 + CACHE_BLOCK: tl.constexpr,
353 + COMPRESS_RATIO: tl.constexpr,
354 + STORE_FP8: tl.constexpr,
355 +):
356 + t = tl.program_id(0)
357 + slot = tl.load(cache_slots + t)
358 + if slot < 0:
359 + return
360 + position = tl.load(positions + t)
361 + if (position + 1) % COMPRESS_RATIO != 0:
362 + return
363 + d = tl.arange(0, 512)
364 + normed = tl.load(latent + t.to(tl.int64) * 512 + d).to(tl.float32)
365 +
366 + # NoPE pairs load (cos, sin) = (1, 0), so the rotation is the identity there.
367 + even, odd = tl.split(tl.reshape(normed, (256, 2)))
368 + pair = tl.arange(0, 256) - 224
369 + cs = cos_sin + (position // COMPRESS_RATIO * COMPRESS_RATIO) * COS_STRIDE
370 + c = tl.load(cs + tl.maximum(pair, 0), pair >= 0, other=1.0).to(tl.float32)
371 + s = tl.load(cs + 32 + tl.maximum(pair, 0), pair >= 0, other=0.0).to(tl.float32)
372 + row = tl.interleave(even * c - odd * s, odd * c + even * s).to(tl.bfloat16)
373 +
374 + dst = (
375 + cache
376 + + (slot // CACHE_BLOCK).to(tl.int64) * CACHE_STRIDE
377 + + (slot % CACHE_BLOCK) * ROW_STRIDE
378 + )
379 + if STORE_FP8:
380 + scaled = row.to(tl.float32) * (1.0 / tl.load(fp8_scale))
381 + tl.store(dst + d, tl.clamp(scaled, -448.0, 448.0).to(tl.float8e4nv))
382 + else:
383 + tl.store(dst + d, row)
vllm/models/deepseek_v4_1/common/ops/indexer_k_store.py new
+253
@@ -0,0 +1,253 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Indexer K production for DeepSeek V4.1 kv-source layers.
4 +
5 +In v4.1 the index key is derived from the *main* compressor's latent:
6 +``k = k_norm(wk(latent))`` (reference model.py Indexer.forward), then RoPE'd
7 +at the group's first-token position and MXFP4/FP8-quantized into the paged
8 +indexer K cache. The ``wk(latent)`` GEMM runs in torch; this kernel fuses the
9 +remaining k_norm → RoPE → quant → paged store, one program per token.
10 +
11 +Unlike the legacy (v4.0) indexer path there is no per-token pooling from a
12 +compressor state cache: the latent already stands for a whole group, so only
13 +group-boundary tokens ``(position + 1) % compress_ratio == 0`` produce a key.
14 +"""
15 +
16 +import torch
17 +
18 +from vllm.platforms import current_platform
19 +from vllm.triton_utils import tl, triton
20 +
21 +from . import MXFP4_BLOCK_SIZE, _fp32x2_to_fp4x2
22 +
23 +# ROCm tiled ("SHUFFLE") indexer K value layout, mirroring
24 +# indexer_k_quant_and_cache_triton's defaults: 16 positions x 16 bytes.
25 +_BLOCK_TILE_SIZE = 16
26 +_HEAD_TILE_SIZE = 16
27 +
28 +
29 +def indexer_k_norm_rope_store(
30 + k_pre: torch.Tensor,
31 + positions: torch.Tensor,
32 + cos_sin_cache: torch.Tensor,
33 + rms_norm_weight: torch.Tensor,
34 + rms_norm_eps: float,
35 + k_cache: torch.Tensor,
36 + kv_slot_mapping: torch.Tensor,
37 + compress_ratio: int,
38 + use_fp4_cache: bool,
39 +) -> None:
40 + """k_norm → RoPE → quant → paged store for indexer keys.
41 +
42 + Args:
43 + k_pre: [num_tokens, 128] bf16, the ``wk(latent)`` projection. Only
44 + group-boundary rows are read.
45 + positions: [num_tokens] int64 token positions.
46 + cos_sin_cache: [max_pos, rope_head_dim] GPT-J layout (cos half, then
47 + sin half), from the layer's compress-RoPE instance.
48 + rms_norm_weight: [128] k_norm weight.
49 + k_cache: uint8 paged indexer cache [num_blocks, block_size, row_bytes].
50 + kv_slot_mapping: [num_tokens] slots in the indexer cache (-1 = skip).
51 + compress_ratio: group size; keys are emitted at group boundaries.
52 + use_fp4_cache: MXFP4 (2 nibbles/byte + ue8m0 per 32) when True, else
53 + per-token FP8 with a single fp32 scale.
54 + """
55 + num_tokens = kv_slot_mapping.numel()
56 + assert k_pre.ndim == 2 and k_pre.shape[1] == 128
57 + assert k_pre.dtype == torch.bfloat16 and k_pre.stride(1) == 1
58 + assert num_tokens <= k_pre.shape[0] and num_tokens <= positions.numel()
59 + assert compress_ratio in (1, 2)
60 + if num_tokens == 0:
61 + return
62 +
63 + head_dim = k_pre.shape[1]
64 + if use_fp4_cache:
65 + token_stride = head_dim // 2
66 + scale_dim = head_dim // MXFP4_BLOCK_SIZE
67 + else:
68 + token_stride = head_dim
69 + scale_dim = 4 # single float32 scale
70 +
71 + # ROCm reads this cache back with the 16x16-tiled ("SHUFFLE") value
72 + # layout: cp_gather_indexer_k_quant_cache_triton on prefill and aiter's
73 + # deepgemm_fp8_paged_mqa_logits(Preshuffle=True) on decode both select it
74 + # from ``block_size > 1``, matching the v4.0 writer
75 + # (indexer_k_quant_and_cache_triton). Writing row-major here would hand
76 + # both readers permuted key bytes.
77 + block_size = k_cache.shape[1]
78 + shuffle = current_platform.is_rocm() and block_size > 1
79 + if shuffle:
80 + if use_fp4_cache:
81 + raise NotImplementedError(
82 + "MXFP4 indexer K cache has no tiled ROCm layout; "
83 + "the ROCm readers only implement the FP8 one."
84 + )
85 + if block_size % _BLOCK_TILE_SIZE != 0 or head_dim % _HEAD_TILE_SIZE != 0:
86 + raise ValueError(
87 + f"ROCm tiled indexer K cache needs block_size "
88 + f"({block_size}) % {_BLOCK_TILE_SIZE} == 0 and head_dim "
89 + f"({head_dim}) % {_HEAD_TILE_SIZE} == 0."
90 + )
91 +
92 + launch_kwargs = {"launch_pdl": False} if current_platform.is_cuda() else {}
93 + _indexer_k_norm_rope_quant_store_kernel[(num_tokens,)](
94 + k_pre,
95 + k_pre.stride(0),
96 + positions,
97 + rms_norm_weight,
98 + rms_norm_eps,
99 + cos_sin_cache,
100 + cos_sin_cache.stride(0),
101 + k_cache,
102 + kv_slot_mapping,
103 + k_cache.shape[1],
104 + HEAD_SIZE=head_dim,
105 + ROPE_HEAD_DIM=64,
106 + COMPRESS_RATIO=compress_ratio,
107 + TOKEN_STRIDE=token_stride,
108 + SCALE_DIM=scale_dim,
109 + KV_BLOCK_STRIDE=k_cache.stride(0),
110 + FP8_MAX=448.0,
111 + USE_FP4=use_fp4_cache,
112 + SHUFFLE=shuffle,
113 + BLOCK_TILE_SIZE=_BLOCK_TILE_SIZE,
114 + HEAD_TILE_SIZE=_HEAD_TILE_SIZE,
115 + num_warps=1,
116 + **launch_kwargs,
117 + )
118 +
119 +
120 +@triton.jit
121 +def _indexer_k_norm_rope_quant_store_kernel(
122 + k_pre_ptr,
123 + k_pre_stride,
124 + positions_ptr,
125 + rms_norm_weight_ptr,
126 + rms_norm_eps,
127 + cos_sin_cache_ptr,
128 + cos_sin_stride,
129 + k_cache_ptr,
130 + kv_slot_mapping_ptr,
131 + kv_cache_block_size,
132 + HEAD_SIZE: tl.constexpr,
133 + ROPE_HEAD_DIM: tl.constexpr,
134 + COMPRESS_RATIO: tl.constexpr,
135 + TOKEN_STRIDE: tl.constexpr,
136 + SCALE_DIM: tl.constexpr,
137 + KV_BLOCK_STRIDE: tl.constexpr,
138 + FP8_MAX: tl.constexpr,
139 + USE_FP4: tl.constexpr,
140 + SHUFFLE: tl.constexpr,
141 + BLOCK_TILE_SIZE: tl.constexpr,
142 + HEAD_TILE_SIZE: tl.constexpr,
143 +):
144 + token_idx = tl.program_id(0)
145 +
146 + kv_slot_idx = tl.load(kv_slot_mapping_ptr + token_idx)
147 + if kv_slot_idx < 0:
148 + return
149 + position = tl.load(positions_ptr + token_idx)
150 + # Only the last token of a group publishes that group's index key.
151 + if (position + 1) % COMPRESS_RATIO != 0:
152 + return
153 +
154 + block = tl.arange(0, HEAD_SIZE)
155 +
156 + # ── k_norm (fp32 throughout, bf16 roundtrip like the reference) ────
157 + k = tl.load(k_pre_ptr + token_idx * k_pre_stride + block).to(tl.float32)
158 + rms_w = tl.load(rms_norm_weight_ptr + block).to(tl.float32)
159 + variance = tl.sum(k * k, axis=0) / HEAD_SIZE
160 + k = (k * tl.rsqrt(variance + rms_norm_eps) * rms_w).to(tl.bfloat16)
161 + k = k.to(tl.float32)
162 +
163 + # ── Register-based GPT-J forward RoPE in fp32 ─────────────────────
164 + # A latent stands for the first token of its group, so group j takes
165 + # position j * compress_ratio.
166 + NUM_PAIRS: tl.constexpr = HEAD_SIZE // 2
167 + NOPE_HEAD_DIM: tl.constexpr = HEAD_SIZE - ROPE_HEAD_DIM
168 + NOPE_PAIRS: tl.constexpr = NOPE_HEAD_DIM // 2
169 + HALF_ROPE: tl.constexpr = ROPE_HEAD_DIM // 2
170 +
171 + even, odd = tl.split(tl.reshape(k, (NUM_PAIRS, 2))) # each [NUM_PAIRS]
172 + pair_idx = tl.arange(0, NUM_PAIRS)
173 + rope_pair_local = pair_idx - NOPE_PAIRS
174 + is_rope_pair = rope_pair_local >= 0
175 + cs_idx = tl.maximum(rope_pair_local, 0)
176 +
177 + compressed_pos = (position // COMPRESS_RATIO) * COMPRESS_RATIO
178 + cache_base = cos_sin_cache_ptr + compressed_pos * cos_sin_stride
179 + cos_v = tl.load(cache_base + cs_idx, mask=is_rope_pair, other=1.0)
180 + sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope_pair, other=0.0)
181 +
182 + new_even = even * cos_v - odd * sin_v
183 + new_odd = odd * cos_v + even * sin_v
184 +
185 + # bf16 roundtrip for parity with the reference / Q-side kernel numerics.
186 + new_even = new_even.to(tl.bfloat16).to(tl.float32)
187 + new_odd = new_odd.to(tl.bfloat16).to(tl.float32)
188 +
189 + # ── Paged cache pointers (segregated: values first, then scales) ──
190 + kv_block_idx = kv_slot_idx // kv_cache_block_size
191 + kv_pos_in_block = kv_slot_idx % kv_cache_block_size
192 + cache_block_ptr = k_cache_ptr + kv_block_idx.to(tl.int64) * KV_BLOCK_STRIDE
193 + if SHUFFLE:
194 + # [block_size // BT, HEAD_SIZE // HT, BT, HT] value tiles.
195 + val_ptr = (
196 + cache_block_ptr
197 + + (kv_pos_in_block // BLOCK_TILE_SIZE) * BLOCK_TILE_SIZE * TOKEN_STRIDE
198 + + (kv_pos_in_block % BLOCK_TILE_SIZE) * HEAD_TILE_SIZE
199 + )
200 + else:
201 + val_ptr = cache_block_ptr + kv_pos_in_block * TOKEN_STRIDE
202 + scale_ptr = (
203 + cache_block_ptr
204 + + kv_cache_block_size * TOKEN_STRIDE
205 + + kv_pos_in_block * SCALE_DIM
206 + )
207 +
208 + if USE_FP4:
209 + # MXFP4: each 32-element block = 16 consecutive even/odd pairs, so
210 + # tiling the halves into (N_BLOCKS, 16) lands one block per row.
211 + N_QUANT_BLOCKS: tl.constexpr = HEAD_SIZE // 32
212 + HALF_BLOCK: tl.constexpr = 16
213 + even_2d = tl.reshape(new_even, (N_QUANT_BLOCKS, HALF_BLOCK))
214 + odd_2d = tl.reshape(new_odd, (N_QUANT_BLOCKS, HALF_BLOCK))
215 +
216 + amax = tl.maximum(
217 + tl.max(tl.abs(even_2d), axis=1),
218 + tl.max(tl.abs(odd_2d), axis=1),
219 + )
220 + amax = tl.maximum(amax, 6.0 * (2**-126))
221 + # ue8m0 block scale: 2^ceil(log2(amax / 6.0)), stored (exp + 127).
222 + log2_ratio = tl.ceil(tl.log2(amax * (1.0 / 6.0)))
223 + log2_ratio = tl.minimum(tl.maximum(log2_ratio, -127.0), 127.0)
224 + inv_scale = tl.exp2(-log2_ratio)
225 + ue8m0 = (log2_ratio + 127.0).to(tl.uint8) # [N_QUANT_BLOCKS]
226 +
227 + inv_scale_col = tl.reshape(inv_scale, (N_QUANT_BLOCKS, 1))
228 + packed = _fp32x2_to_fp4x2(
229 + even_2d * inv_scale_col, odd_2d * inv_scale_col
230 + ) # (N_BLOCKS, HALF_BLOCK) uint8
231 + packed_flat = tl.reshape(packed, (TOKEN_STRIDE,))
232 +
233 + tl.store(val_ptr + tl.arange(0, TOKEN_STRIDE), packed_flat)
234 + tl.store(scale_ptr + tl.arange(0, SCALE_DIM), ue8m0)
235 + else:
236 + # Per-token FP8 (single 128-wide block) with one float32 scale.
237 + result = tl.interleave(new_even, new_odd) # [HEAD_SIZE] fp32
238 + result_bf16 = result.to(tl.bfloat16).to(tl.float32)
239 + INV_FP8_MAX: tl.constexpr = 1.0 / FP8_MAX
240 + absmax = tl.maximum(tl.max(tl.abs(result_bf16), axis=0), 1e-4)
241 + exponent = tl.ceil(tl.log2(absmax * INV_FP8_MAX))
242 + inv_scale = tl.exp2(-exponent)
243 + x_clamped = tl.clamp(result_bf16 * inv_scale, -FP8_MAX, FP8_MAX)
244 + x_uint8 = x_clamped.to(tl.float8e4nv).to(tl.uint8, bitcast=True)
245 + if SHUFFLE:
246 + tiled = (
247 + block // HEAD_TILE_SIZE * BLOCK_TILE_SIZE * HEAD_TILE_SIZE
248 + + block % HEAD_TILE_SIZE
249 + )
250 + tl.store(val_ptr + tiled, x_uint8)
251 + else:
252 + tl.store(val_ptr + block, x_uint8)
253 + tl.store(scale_ptr.to(tl.pointer_type(tl.float32)), tl.exp2(exponent))
vllm/models/deepseek_v4_1/common/ops/query_quant.py new
+142
@@ -0,0 +1,142 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
7 +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic
8 +from vllm.platforms import current_platform
9 +from vllm.triton_utils import tl, triton
10 +
11 +
12 +@triton.jit
13 +def _norm_row(x, weight, row, stride, eps, SIZE: tl.constexpr, BLOCK: tl.constexpr):
14 + cols = tl.arange(0, BLOCK)
15 + w = tl.load(weight + cols, cols < SIZE, 0).to(tl.float32)
16 + values = tl.load(x + row * stride + cols, cols < SIZE, 0).to(tl.float32)
17 + rrms = tl.rsqrt(tl.sum(values * values, 0) / SIZE + eps)
18 + return values * rrms * w
19 +
20 +
21 +@triton.jit(do_not_specialize=["num_tokens"])
22 +def _q_kv_norm_quant_kernel(
23 + q,
24 + kv,
25 + qw,
26 + kvw,
27 + qo,
28 + kvo,
29 + scales,
30 + num_tokens,
31 + q_stride,
32 + kv_stride,
33 + eps,
34 + Q_SIZE: tl.constexpr,
35 + KV_SIZE: tl.constexpr,
36 + BLOCK: tl.constexpr,
37 + LAUNCH_PDL: tl.constexpr,
38 +):
39 + row = tl.program_id(0).to(tl.int64)
40 + cols = tl.arange(0, BLOCK)
41 + if LAUNCH_PDL:
42 + tl.extra.cuda.gdc_wait()
43 + tl.extra.cuda.gdc_launch_dependents()
44 + if tl.program_id(1) == 0:
45 + groups = tl.arange(0, BLOCK // 32)
46 + if row < num_tokens:
47 + y = _norm_row(q, qw, row, q_stride, eps, Q_SIZE, BLOCK)
48 + # Preserve the materialized normalization output's rounding boundary.
49 + y = y.to(q.dtype.element_ty).to(tl.float32)
50 + grouped = tl.reshape(y, (BLOCK // 32, 32))
51 + amax = tl.max(tl.abs(grouped), 1)
52 + normalized = amax * (1.0 / 448.0)
53 + bits = normalized.to(tl.uint32, bitcast=True)
54 + exponent = (bits >> 23) & 255
55 + mantissa = bits & 0x7FFFFF
56 + bump = (mantissa != 0) & ~((exponent == 0) & (mantissa <= 0x400000))
57 + sf = tl.minimum(exponent + bump, 254)
58 + sf = tl.where(normalized <= 0, 0, sf)
59 + # Match FlashInfer's UE8M0 conversion, including zero/subnormal scales.
60 + inv_bits = tl.where(sf == 0, 0, (254 - sf) << 23)
61 + inv_scale = inv_bits.to(tl.float32, bitcast=True)
62 + quantized = tl.reshape(grouped * inv_scale[:, None], (BLOCK,))
63 + tl.store(qo + row * Q_SIZE + cols, quantized, cols < Q_SIZE)
64 + else:
65 + sf = tl.full((BLOCK // 32,), 0, tl.uint32)
66 + padded_groups: tl.constexpr = triton.cdiv(Q_SIZE // 32, 4) * 4
67 + sf = tl.where(groups < Q_SIZE // 32, sf, 0)
68 + # F8_128x4: [row/128, group/4, row%32, row%128/32, group%4].
69 + offsets = (
70 + row // 128 * (128 * padded_groups)
71 + + groups // 4 * 512
72 + + row % 32 * 16
73 + + row % 128 // 32 * 4
74 + + groups % 4
75 + )
76 + tl.store(scales + offsets, sf, groups < padded_groups)
77 + elif row < num_tokens:
78 + y = _norm_row(kv, kvw, row, kv_stride, eps, KV_SIZE, BLOCK)
79 + tl.store(kvo + row * KV_SIZE + cols, y, cols < KV_SIZE)
80 +
81 +
82 +def fused_q_kv_rmsnorm_quant(
83 + qr: torch.Tensor,
84 + kv: torch.Tensor,
85 + q_weight: torch.Tensor,
86 + kv_weight: torch.Tensor,
87 + eps: float,
88 +) -> tuple[QuantizedActivation, torch.Tensor]:
89 + """Normalize Q/KV and quantize Q with FlashInfer's swizzled MXFP8 scales."""
90 + assert qr.ndim == kv.ndim == 2 and qr.shape[0] == kv.shape[0]
91 + assert qr.stride(-1) == kv.stride(-1) == 1
92 + assert q_weight.is_contiguous() and kv_weight.is_contiguous()
93 + assert qr.shape[1] % 32 == 0
94 + tokens, q_size = qr.shape
95 + kv_size = kv.shape[1]
96 + qo = torch.empty(qr.shape, dtype=torch.float8_e4m3fn, device=qr.device)
97 + kvo = torch.empty(kv.shape, dtype=kv.dtype, device=kv.device)
98 + padded_tokens = triton.cdiv(tokens, 128) * 128
99 + padded_groups = triton.cdiv(q_size // 32, 4) * 4
100 + scales = torch.empty(
101 + padded_tokens * padded_groups, dtype=torch.uint8, device=qr.device
102 + )
103 + if tokens:
104 + block = triton.next_power_of_2(max(q_size, kv_size))
105 + _q_kv_norm_quant_kernel[(padded_tokens, 2)](
106 + qr,
107 + kv,
108 + q_weight,
109 + kv_weight,
110 + qo,
111 + kvo,
112 + scales,
113 + tokens,
114 + qr.stride(0),
115 + kv.stride(0),
116 + eps,
117 + q_size,
118 + kv_size,
119 + block,
120 + current_platform.is_arch_support_pdl(),
121 + num_warps=8 if block >= 2048 else 4,
122 + )
123 + return QuantizedActivation(qo, scales, qr.dtype, qr.shape, kMxfp8Dynamic), kvo
124 +
125 +
126 +def can_fuse_query_quant(linears: list[torch.nn.Module]) -> bool:
127 + """Require both local projections to use the same MXFP8 activation ABI."""
128 + if not current_platform.is_cuda():
129 + return False
130 + from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
131 + FlashInferCutedslMxfp8LinearKernel,
132 + FlashInferCutlassMxfp8LinearKernel,
133 + )
134 +
135 + # QuantKey does not encode scale layout; restrict this producer to the
136 + # consumers that accept F8_128x4 swizzled scales.
137 + return all(
138 + getattr(linear, "input_quant_key", None) == kMxfp8Dynamic
139 + and type(getattr(getattr(linear, "quant_method", None), "kernel", None))
140 + in (FlashInferCutedslMxfp8LinearKernel, FlashInferCutlassMxfp8LinearKernel)
141 + for linear in linears
142 + )
vllm/models/deepseek_v4_1/common/rope.py new
+53
@@ -0,0 +1,53 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""DeepseekV4 rotary embedding initialization."""
4 +
5 +from vllm.model_executor.layers.rotary_embedding import get_rope
6 +from vllm.model_executor.layers.rotary_embedding.base import RotaryEmbedding
7 +
8 +
9 +def build_deepseek_v4_rope(
10 + config,
11 + *,
12 + head_dim: int,
13 + rope_head_dim: int,
14 + max_position_embeddings: int,
15 + compress_ratio: int,
16 +) -> RotaryEmbedding:
17 + rope_parameters = config.rope_parameters
18 + # Newer checkpoints nest per-layer-type rope dicts ({"main", "compress"});
19 + # older ones ship a single flat dict shared by all layer types.
20 + if isinstance(rope_parameters.get("main"), dict) and isinstance(
21 + rope_parameters.get("compress"), dict
22 + ):
23 + key = "compress" if compress_ratio > 0 else "main"
24 + rope_parameters = dict(rope_parameters[key])
25 + else:
26 + rope_parameters = dict(rope_parameters)
27 +
28 + rope_parameters["rope_theta"] = (
29 + config.compress_rope_theta if compress_ratio > 0 else config.rope_theta
30 + )
31 + if compress_ratio > 0 and rope_parameters["rope_type"] != "default":
32 + # YaRN applies to all compressed-KV layers (compress_ratio > 0); only
33 + # pure sliding-window layers (compress_ratio == 0) use plain RoPE.
34 + rope_parameters["rope_type"] = (
35 + "deepseek_yarn"
36 + if rope_parameters.get("apply_yarn_scaling", True)
37 + else "deepseek_llama_scaling"
38 + )
39 + else:
40 + # Sliding-window layers use plain RoPE (theta=rope_theta, no YaRN).
41 + rope_parameters["rope_type"] = "deepseek_yarn"
42 + rope_parameters["factor"] = 1.0
43 + rope_parameters["original_max_position_embeddings"] = max_position_embeddings
44 + rope_parameters["mscale"] = 0 # Disable mscale
45 + rope_parameters["mscale_all_dim"] = 0 # Disable mscale
46 + rope_parameters["is_deepseek_v4"] = True
47 + rope_parameters["rope_dim"] = rope_head_dim
48 + return get_rope(
49 + head_dim,
50 + max_position=max_position_embeddings,
51 + rope_parameters=rope_parameters,
52 + is_neox_style=False,
53 + )
vllm/models/deepseek_v4_1/compressor.py new
+323
@@ -0,0 +1,323 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from dataclasses import dataclass
5 +from typing import Any, ClassVar, cast
6 +
7 +import torch
8 +from torch import nn
9 +
10 +from vllm.config import VllmConfig, get_current_vllm_config
11 +from vllm.forward_context import get_forward_context
12 +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
13 +from vllm.model_executor.layers.layernorm import RMSNorm
14 +from vllm.model_executor.layers.linear import MergedColumnParallelLinear
15 +from vllm.models.deepseek_v4_1.common.ops.fused_compress_quant_cache import (
16 + fused_save_compress_norm,
17 + rope_quant_insert,
18 +)
19 +from vllm.platforms import current_platform
20 +from vllm.triton_utils import tl, triton
21 +from vllm.v1.attention.backend import (
22 + AttentionBackend,
23 + AttentionCGSupport,
24 + AttentionMetadataBuilder,
25 + CommonAttentionMetadata,
26 + MultipleOf,
27 +)
28 +from vllm.v1.kv_cache_interface import CircularBufferSpec, KVCacheSpec
29 +
30 +
31 +class CompressorBackend(AttentionBackend):
32 + def __init__(self):
33 + super().__init__()
34 +
35 + @staticmethod
36 + def get_name() -> str:
37 + return "CompressorBackend"
38 +
39 + @staticmethod
40 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
41 + return [MultipleOf(1)]
42 +
43 + @classmethod
44 + def get_supported_head_sizes(cls) -> list[int]:
45 + return [512, 1024]
46 +
47 + @staticmethod
48 + def get_builder_cls() -> type["CompressorMetadataBuilder"]:
49 + return CompressorMetadataBuilder
50 +
51 +
52 +@dataclass
53 +class CompressorMetadata:
54 + # [num_tokens] ring slot of every token: block * capacity + pos % capacity.
55 + slot_mapping: torch.Tensor
56 + query_start_loc: torch.Tensor # [num_reqs + 1]
57 + token_to_req_indices: torch.Tensor # [num_tokens]
58 +
59 +
60 +@triton.jit
61 +def _ring_slot_mapping_kernel(
62 + slot_mapping_ptr,
63 + block_table_ptr,
64 + block_table_stride,
65 + token_to_req_ptr,
66 + positions_ptr,
67 + num_actual_tokens,
68 + num_tokens,
69 + CAPACITY: tl.constexpr,
70 + BLOCK: tl.constexpr,
71 +):
72 + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
73 + valid = offsets < num_actual_tokens
74 + req = tl.load(token_to_req_ptr + offsets, mask=valid, other=0).to(tl.int64)
75 + block = tl.load(block_table_ptr + req * block_table_stride, mask=valid, other=0)
76 + pos = tl.load(positions_ptr + offsets, mask=valid, other=0)
77 + slot = block.to(tl.int64) * CAPACITY + pos % CAPACITY
78 + tl.store(
79 + slot_mapping_ptr + offsets, tl.where(valid, slot, -1), mask=offsets < num_tokens
80 + )
81 +
82 +
83 +class CompressorMetadataBuilder(AttentionMetadataBuilder):
84 + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
85 +
86 + def __init__(self, *args, **kwargs):
87 + super().__init__(*args, **kwargs)
88 + assert isinstance(self.kv_cache_spec, CircularBufferSpec)
89 + self.capacity = self.kv_cache_spec.block_size
90 + max_num_batched_tokens = (
91 + self.vllm_config.scheduler_config.max_num_batched_tokens
92 + )
93 + self.token_to_req_indices = torch.zeros(
94 + max_num_batched_tokens, dtype=torch.int32, device=self.device
95 + )
96 + self.slot_mapping_buffer = torch.empty(
97 + max_num_batched_tokens, dtype=torch.int64, device=self.device
98 + )
99 +
100 + def build(
101 + self,
102 + common_prefix_len: int,
103 + common_attn_metadata: CommonAttentionMetadata,
104 + fast_build: bool = False,
105 + ) -> CompressorMetadata:
106 + num_tokens = common_attn_metadata.slot_mapping.numel()
107 + positions = common_attn_metadata.positions
108 + assert positions is not None
109 + token_to_req_indices = common_attn_metadata.token_to_req_indices(
110 + self.token_to_req_indices
111 + )
112 + slot_mapping = self.slot_mapping_buffer[:num_tokens]
113 + block_table = common_attn_metadata.block_table_tensor
114 + _ring_slot_mapping_kernel[(triton.cdiv(num_tokens, 256),)](
115 + slot_mapping,
116 + block_table,
117 + block_table.stride(0),
118 + token_to_req_indices,
119 + positions,
120 + common_attn_metadata.num_actual_tokens,
121 + num_tokens,
122 + CAPACITY=self.capacity,
123 + BLOCK=256,
124 + )
125 + return CompressorMetadata(
126 + slot_mapping=slot_mapping,
127 + query_start_loc=common_attn_metadata.query_start_loc,
128 + token_to_req_indices=token_to_req_indices,
129 + )
130 +
131 +
132 +class CompressorStateCache(torch.nn.Module, AttentionLayerBase):
133 + """Per-request ring of the open group's [kv, score] rows (ratio > 1)."""
134 +
135 + def __init__(self, state_dim: int, dtype: torch.dtype, prefix: str):
136 + super().__init__()
137 + self.state_dim = state_dim
138 + self.dtype = dtype
139 + self.prefix = prefix
140 + self.kv_cache = torch.tensor([])
141 + compilation_config = get_current_vllm_config().compilation_config
142 + if prefix in compilation_config.static_forward_context:
143 + raise ValueError(f"Duplicate layer name: {prefix}")
144 + compilation_config.static_forward_context[prefix] = self
145 +
146 + assert self.dtype == torch.float32
147 + # Drafts + bonus token + the previous row, as a power of two >= 8.
148 + rows_per_step = get_current_vllm_config().num_speculative_tokens + 2
149 + self.block_size = max(8, 1 << (rows_per_step - 1).bit_length())
150 +
151 + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
152 + # [B, H=1, N, C] -> [B, N, C]
153 + self.kv_cache = kv_cache.squeeze(1)
154 +
155 + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
156 + del vllm_config
157 + return CircularBufferSpec(
158 + block_size=self.block_size,
159 + num_kv_heads=1,
160 + head_size=self.state_dim,
161 + head_size_v=0,
162 + dtype=self.dtype,
163 + )
164 +
165 + def forward(self): ...
166 +
167 + def get_attn_backend(self) -> type[AttentionBackend]:
168 + return CompressorBackend
169 +
170 +
171 +class DeepseekCompressor(nn.Module):
172 + """DeepSeek V4.1 KV/score compressor.
173 +
174 + Pools ``compress_ratio`` consecutive tokens into one KV latent with a
175 + learned softmax gate (ratio 1 has no gate and no pooling). Owns the
176 + linear, norm and state cache. State saving, compression and RMSNorm share
177 + one Triton kernel. The emitted latent feeds independent main-cache and
178 + indexer K kernels, which the attention layer can schedule concurrently.
179 + """
180 +
181 + def __init__(
182 + self,
183 + vllm_config: VllmConfig,
184 + compress_ratio: int,
185 + hidden_size: int,
186 + head_dim: int,
187 + rotate: bool = False,
188 + prefix: str = "",
189 + k_cache_prefix="",
190 + ):
191 + super().__init__()
192 + if compress_ratio not in (1, 2):
193 + raise NotImplementedError(
194 + "DeepSeek V4.1 compressor supports compress_ratio 1 (full-length "
195 + f"compressed cache) and 2; got {compress_ratio}. The ratio-4/128 "
196 + "CuTe-DSL kernels are v4.0-specific and not wired here."
197 + )
198 + self.compress_ratio = compress_ratio
199 + self.hidden_size = hidden_size
200 + self.head_dim = head_dim
201 + self.rotate = rotate
202 + self.prefix = prefix
203 + self.k_cache_prefix = k_cache_prefix
204 + # Ratio 1 pools single tokens, so the checkpoint carries no gate.
205 + self.has_gate = compress_ratio > 1
206 +
207 + config = vllm_config.model_config.hf_config
208 + self.rope_head_dim = config.qk_rope_head_dim
209 + self.nope_head_dim = self.head_dim - self.rope_head_dim
210 + assert self.head_dim == 512 and self.rope_head_dim == 64
211 + self.rms_norm_eps = config.rms_norm_eps
212 + self.device = current_platform.device_type
213 + self.max_num_reqs = vllm_config.scheduler_config.max_num_seqs
214 + self.max_model_len = vllm_config.model_config.max_model_len
215 +
216 + wkv_wgate_sizes = (
217 + [self.head_dim, self.head_dim] if self.has_gate else [self.head_dim]
218 + )
219 + self.fused_wkv_wgate = MergedColumnParallelLinear(
220 + self.hidden_size,
221 + wkv_wgate_sizes,
222 + bias=False,
223 + return_bias=False,
224 + quant_config=None,
225 + disable_tp=True,
226 + prefix=f"{prefix}.fused_wkv_wgate",
227 + )
228 + self.norm = RMSNorm(self.head_dim, self.rms_norm_eps)
229 +
230 + self.state_cache = (
231 + CompressorStateCache(
232 + state_dim=2 * self.head_dim, # kv_state + score_state
233 + dtype=torch.float32,
234 + prefix=f"{prefix}.state_cache",
235 + )
236 + if compress_ratio > 1
237 + else None
238 + )
239 +
240 + # Save reference to static_forward_context for forward-time KV cache lookup.
241 + # get_current_vllm_config() is only available during __init__, not forward.
242 + self._static_forward_context = (
243 + vllm_config.compilation_config.static_forward_context
244 + )
245 +
246 + def forward(
247 + self,
248 + # [num_tokens, (2 if has_gate else 1) * self.head_dim]
249 + kv_score: torch.Tensor,
250 + # [num_tokens]
251 + positions: torch.Tensor,
252 + ) -> torch.Tensor | None:
253 + """Save states and return the BF16 latent for cache insertion and indexing.
254 +
255 + Only valid group-boundary rows are written.
256 + """
257 + attn_metadata = get_forward_context().attn_metadata
258 + if not isinstance(attn_metadata, dict):
259 + return None
260 +
261 + if self.state_cache is None:
262 + state_cache = query_start_loc = token_to_req_indices = None
263 + slot_mapping = cast(Any, attn_metadata[self.k_cache_prefix]).slot_mapping
264 + else:
265 + state_metadata = cast(
266 + CompressorMetadata, attn_metadata[self.state_cache.prefix]
267 + )
268 + state_cache = self.state_cache.kv_cache
269 + slot_mapping = state_metadata.slot_mapping
270 + query_start_loc = state_metadata.query_start_loc
271 + token_to_req_indices = state_metadata.token_to_req_indices
272 +
273 + latent = torch.empty(
274 + kv_score.shape[0],
275 + self.head_dim,
276 + dtype=torch.bfloat16,
277 + device=kv_score.device,
278 + )
279 + fused_save_compress_norm(
280 + kv_score,
281 + positions,
282 + state_cache,
283 + slot_mapping,
284 + query_start_loc,
285 + token_to_req_indices,
286 + self.norm.weight,
287 + self.rms_norm_eps,
288 + self.compress_ratio,
289 + latent,
290 + )
291 + return latent
292 +
293 + def insert_cache(
294 + self,
295 + latent: torch.Tensor | None,
296 + positions: torch.Tensor,
297 + rotary_emb,
298 + ) -> None:
299 + """Publish compressed main-cache rows after the latent becomes ready."""
300 + if latent is None:
301 + return
302 + attn_metadata = get_forward_context().attn_metadata
303 + assert isinstance(attn_metadata, dict)
304 + k_cache_metadata = cast(Any, attn_metadata[self.k_cache_prefix])
305 + k_cache_layer = self._static_forward_context[self.k_cache_prefix]
306 + kv_cache = k_cache_layer.kv_cache
307 + # Plain-row per-tensor fp8 caches (FlashInfer) carry the layer's scale;
308 + # fp8_ds_mla and bf16 rows need none.
309 + fp8_scale = (
310 + getattr(k_cache_layer, "_flashinfer_fp8_kv_scale", None)
311 + if kv_cache.dtype == torch.float8_e4m3fn
312 + else None
313 + )
314 +
315 + rope_quant_insert(
316 + latent,
317 + positions,
318 + rotary_emb.cos_sin_cache,
319 + kv_cache,
320 + k_cache_metadata.slot_mapping,
321 + self.compress_ratio,
322 + fp8_scale=fp8_scale,
323 + )
vllm/models/deepseek_v4_1/nvidia/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/deepseek_v4_1/nvidia/dspark.py new
+558
@@ -0,0 +1,558 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""DSpark draft model for DeepSeek-V4.1 (semi-autoregressive speculative decoding).
4 +
5 +See: qwen3_dspark.py for base architecture. This one is specialized to the
6 +DSV4.1 DSpark, which reuses the target model's architecture similarly to MTP.
7 +Its weights ship in the target checkpoint under the ``mtp.{0,1,2}.*`` prefix.
8 +
9 +To implement non-causal attention, we leverage the sparse attention implementation to
10 +include the future query tokens in the top-k indices for each query token.
11 +"""
12 +
13 +from collections.abc import Iterable
14 +
15 +import regex as re
16 +import torch
17 +import torch.nn as nn
18 +
19 +import vllm.envs as envs
20 +from vllm.config import VllmConfig, get_current_vllm_config
21 +from vllm.distributed import (
22 + get_tensor_model_parallel_rank,
23 + get_tensor_model_parallel_world_size,
24 +)
25 +from vllm.forward_context import get_forward_context, is_forward_context_available
26 +from vllm.logger import init_logger
27 +from vllm.model_executor.kernels.mhc.tilelang import (
28 + mhc_post_tilelang,
29 +)
30 +from vllm.model_executor.kernels.mhc.triton import hc_collapse_triton
31 +from vllm.model_executor.layers.fused_moe import (
32 + fused_moe_make_expert_params_mapping,
33 +)
34 +from vllm.model_executor.layers.layernorm import RMSNorm
35 +from vllm.model_executor.layers.linear import ReplicatedLinear
36 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
37 +from vllm.model_executor.layers.vocab_parallel_embedding import (
38 + ParallelLMHead,
39 + VocabParallelEmbedding,
40 +)
41 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
42 +from vllm.model_executor.models.qwen3_dspark import (
43 + DSparkConfidenceHead,
44 + DSparkMarkovHead,
45 +)
46 +from vllm.model_executor.models.utils import maybe_prefix
47 +from vllm.models.common.ops.sequence_parallel import (
48 + sp_all_gather,
49 + sp_padding_mask,
50 + sp_shard,
51 +)
52 +
53 +from .model import (
54 + DeepseekV4DecoderLayer,
55 + DeepseekV4Model,
56 + _linear_scale_param_name,
57 + _use_sequence_parallel,
58 + make_deepseek_v4_expert_params_mapping,
59 +)
60 +
61 +logger = init_logger(__name__)
62 +
63 +# MoE expert scale suffix differs by expert dtype (mirrors deepseek_v4 loaders):
64 +# fp4 experts register ``.weight_scale``; block-fp8 experts ``.weight_scale_inv``.
65 +_EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$")
66 +
67 +
68 +class DSparkDeepseekV4Model(nn.Module):
69 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
70 + super().__init__()
71 + assert vllm_config.speculative_config is not None
72 + config = vllm_config.speculative_config.draft_model_config.hf_config
73 + self.config = config
74 + self.hidden_size = config.hidden_size
75 + self.hc_mult = config.hc_mult
76 + self.hc_eps = config.hc_eps
77 + self.rms_norm_eps = config.rms_norm_eps
78 + self.num_hidden_layers = config.num_hidden_layers
79 + self.target_layer_ids = tuple(config.dspark_target_layer_ids)
80 + self.use_sequence_parallel = _use_sequence_parallel(vllm_config)
81 +
82 + self.num_dspark_layers = (
83 + getattr(config, "n_mtp_layers", None)
84 + or getattr(config, "num_nextn_predict_layers", None)
85 + or 3
86 + )
87 +
88 + # Shared with the target (aliased by the speculator's loading utility).
89 + self.embed_tokens = VocabParallelEmbedding(
90 + config.vocab_size,
91 + config.hidden_size,
92 + prefix=maybe_prefix(prefix, "embed_tokens"),
93 + )
94 +
95 + self.main_proj = ReplicatedLinear(
96 + config.hidden_size * len(self.target_layer_ids),
97 + config.hidden_size,
98 + bias=False,
99 + return_bias=False,
100 + quant_config=vllm_config.quant_config,
101 + prefix=maybe_prefix(prefix, "main_proj"),
102 + )
103 + self.main_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
104 +
105 + self.topk_indices_buffer = torch.empty(
106 + vllm_config.scheduler_config.max_num_batched_tokens,
107 + config.index_topk,
108 + dtype=torch.int32,
109 + )
110 +
111 + current_vllm_config = get_current_vllm_config()
112 + self.layers = nn.ModuleList(
113 + [
114 + DeepseekV4DecoderLayer(
115 + current_vllm_config,
116 + prefix=maybe_prefix(prefix, f"layers.{self.num_hidden_layers + i}"),
117 + topk_indices_buffer=self.topk_indices_buffer,
118 + )
119 + for i in range(self.num_dspark_layers)
120 + ]
121 + )
122 +
123 + # Heads: final norm, and the Markov + confidence heads.
124 + # Loaded from the "final" MTP layer weights (mtp.*) in the target
125 + # checkpoint. v4.1 has no learned hc_head: the hc copies are
126 + # collapsed with a pre-mix derived from the last layer's hc_ffn
127 + # projection at the end of forward().
128 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
129 + draft_vocab_size = (
130 + getattr(config, "draft_vocab_size", None) or config.vocab_size
131 + )
132 + self.markov_head = DSparkMarkovHead(
133 + config.vocab_size,
134 + draft_vocab_size,
135 + config.dspark_markov_rank,
136 + prefix=maybe_prefix(prefix, "markov_head"),
137 + )
138 + self.confidence_head: DSparkConfidenceHead | None = None
139 + if getattr(config, "enable_confidence_head", True):
140 + self.confidence_head = DSparkConfidenceHead(
141 + config.hidden_size + config.dspark_markov_rank,
142 + prefix=maybe_prefix(prefix, "confidence_head"),
143 + )
144 +
145 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
146 + return self.embed_tokens(input_ids)
147 +
148 + def combine_hidden_states(self, aux_hidden_states: torch.Tensor) -> torch.Tensor:
149 + """main_x = main_norm(main_proj(concat of target aux hidden states)).
150 +
151 + ``aux_hidden_states`` is [T, hidden_size * len(target_layer_ids)].
152 + """
153 + return self.main_norm(self.main_proj(aux_hidden_states))
154 +
155 + @torch.inference_mode()
156 + def precompute_and_store_context_kv(
157 + self,
158 + main_x: torch.Tensor,
159 + context_positions: torch.Tensor,
160 + context_slot_mappings: list[torch.Tensor | None] | None = None,
161 + ) -> None:
162 + """Insert the sliding-window context KV for every draft layer.
163 +
164 + Mirrors the reference DSparkAttention: each layer derives its context KV
165 + from the SAME projected target hidden ``main_x``, via that layer's own
166 + ``wkv`` + ``kv_norm`` + RoPE + quant, then writes it at the
167 + layer's context slots.
168 +
169 + ``context_slot_mappings`` is a per-layer list (each entry is the context
170 + slot mapping for that layer's kv-cache group, since the hybrid manager may
171 + place draft layers in different groups). ``None`` (or a ``None`` entry)
172 + runs the projection to reserve workspace but writes nothing (profiling).
173 + """
174 + for i, layer in enumerate(self.layers):
175 + slot_mapping = (
176 + None if context_slot_mappings is None else context_slot_mappings[i]
177 + )
178 + attn = layer.attn
179 + # Optimized DSV4 MLA path: wkv part of the fused wq_a|wkv projection
180 + # (q_lora part discarded), then RoPE/quant/insert via the fused op.
181 + qr_kv, _ = attn.fused_wqa_wkv(main_x)
182 + kv = qr_kv[..., attn.q_lora_rank :]
183 + kv = attn.kv_norm(kv)
184 + if slot_mapping is None:
185 + continue
186 + _insert_context_kv(attn, kv, context_positions, slot_mapping)
187 +
188 + def forward(
189 + self,
190 + input_ids: torch.Tensor,
191 + positions: torch.Tensor,
192 + inputs_embeds: torch.Tensor | None = None,
193 + ) -> torch.Tensor:
194 + if inputs_embeds is None:
195 + inputs_embeds = self.embed_input_ids(input_ids)
196 + full_num_tokens = positions.shape[0]
197 + if self.use_sequence_parallel:
198 + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available():
199 + forward_context = get_forward_context()
200 + forward_context.is_padding = sp_padding_mask(
201 + forward_context.is_padding, inputs_embeds
202 + )
203 + inputs_embeds = sp_shard(inputs_embeds)
204 + input_ids = sp_shard(input_ids)
205 + # Expand to hc_mult copies for hyper-connections ([T, H] -> [T, hc, H]).
206 + hidden_states = inputs_embeds.unsqueeze(-2).repeat(1, self.hc_mult, 1)
207 +
208 + residual = post_mix = res_mix = pre_mix = None
209 + for layer in self.layers:
210 + hidden_states, residual, post_mix, res_mix, pre_mix = layer(
211 + hidden_states,
212 + positions,
213 + input_ids,
214 + pre_mix,
215 + post_mix,
216 + res_mix,
217 + residual,
218 + )
219 + hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix)
220 + if self.use_sequence_parallel:
221 + hidden_states = sp_all_gather(hidden_states)[:full_num_tokens]
222 + pre_mix = sp_all_gather(pre_mix)[:full_num_tokens]
223 + # Collapse the hc copies with the pre-mix from the last layer's FFN
224 + # mixes — the mix the reference forward_head applies (hc_pre with the
225 + # last block's ffn pre-mix). Return the PRE-norm head hidden;
226 + # compute_logits applies self.norm.
227 + assert pre_mix is not None
228 + hidden_states = hc_collapse_triton(hidden_states, pre_mix)
229 + return hidden_states
230 +
231 +
232 +def _insert_context_kv(
233 + attn: nn.Module,
234 + kv: torch.Tensor,
235 + positions: torch.Tensor,
236 + slot_mapping: torch.Tensor,
237 +) -> None:
238 + """RoPE + quant + paged-cache insert of (already kv_norm'd) context KV.
239 +
240 + Reuses the DSV4 fused insert ops (which also process a query; we pass a dummy
241 + query and discard it, since context tokens have no query). Mirrors
242 + ``DeepseekV4Attention._fused_qnorm_rope_kv_insert``.
243 + """
244 + swa_cache = attn.swa_cache_layer.kv_cache
245 + block_size = attn.swa_cache_layer.block_size
246 + cos_sin_cache = attn.rotary_emb.cos_sin_cache
247 + cache_dtype = swa_cache.dtype
248 + n_ctx = kv.shape[0]
249 + dummy_q = torch.zeros(
250 + (n_ctx, attn.n_local_heads, attn.head_dim),
251 + dtype=kv.dtype,
252 + device=kv.device,
253 + )
254 + if cache_dtype == torch.uint8:
255 + # fp8_ds_mla UE8M0 paged layout
256 + swa_2d = swa_cache.view(swa_cache.shape[0], -1)
257 + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
258 + dummy_q,
259 + kv,
260 + swa_2d,
261 + slot_mapping,
262 + positions,
263 + cos_sin_cache,
264 + attn.padded_heads,
265 + attn.eps,
266 + block_size,
267 + )
268 + elif cache_dtype == torch.bfloat16:
269 + swa_3d = swa_cache.view(-1, block_size, attn.head_dim)
270 + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
271 + dummy_q,
272 + kv,
273 + swa_3d,
274 + slot_mapping,
275 + positions,
276 + cos_sin_cache,
277 + attn.eps,
278 + block_size,
279 + )
280 + else: # per-tensor fp8 (torch.float8_e4m3fn)
281 + # TODO(ben): double-check if this is being dispatched correctly for FI backend
282 + swa_3d = swa_cache.view(-1, block_size, attn.head_dim)
283 + dummy_q_fp8 = torch.zeros_like(dummy_q, dtype=torch.float8_e4m3fn)
284 + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
285 + dummy_q,
286 + kv,
287 + dummy_q_fp8,
288 + swa_3d,
289 + slot_mapping,
290 + positions,
291 + cos_sin_cache,
292 + attn._flashinfer_fp8_kv_scale,
293 + attn._flashinfer_fp8_q_scale_inv,
294 + attn.eps,
295 + block_size,
296 + )
297 +
298 +
299 +class DSparkDeepseekV4ForCausalLM(nn.Module):
300 + # Draft weights ship in the target checkpoint (mtp.*) without embed/head, so
301 + # load_dspark_model always aliases the target's.
302 + has_own_embed_tokens = False
303 + has_own_lm_head = False
304 + # Full-vocab draft: draft ids are target ids, no remapping needed.
305 + draft_id_to_target_id = None
306 +
307 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
308 + super().__init__()
309 + assert vllm_config.speculative_config is not None
310 + self.draft_model_config = vllm_config.speculative_config.draft_model_config
311 + self.config = self.draft_model_config.hf_config
312 + self.quant_config = vllm_config.quant_config
313 + self.linear_scale_name = _linear_scale_param_name(
314 + vllm_config, getattr(self.config, "expert_dtype", "fp4")
315 + )
316 + self.pad_shared_expert = getattr(
317 + self.quant_config, "weight_block_size", None
318 + ) is not None and not _use_sequence_parallel(vllm_config)
319 + self.model = DSparkDeepseekV4Model(
320 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
321 + )
322 + # Shared with the target (aliased by the speculator's load utility).
323 + self.lm_head = ParallelLMHead(
324 + self.config.vocab_size,
325 + self.config.hidden_size,
326 + prefix=maybe_prefix(prefix, "lm_head"),
327 + )
328 + self.logits_processor = LogitsProcessor(self.config.vocab_size)
329 +
330 + # --- Hooks used by the speculator -------------------------------------
331 +
332 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
333 + return self.model.embed_input_ids(input_ids)
334 +
335 + def combine_hidden_states(self, aux_hidden_states: torch.Tensor) -> torch.Tensor:
336 + return self.model.combine_hidden_states(aux_hidden_states)
337 +
338 + def get_draft_kv_cache_layer_names(self) -> list[str]:
339 + # DSV4 MLA path: each draft layer's sliding-window cache is a separate
340 + # layer, named by its prefix.
341 + return [layer.attn.swa_cache_layer.prefix for layer in self.model.layers]
342 +
343 + def precompute_and_store_context_kv(
344 + self,
345 + context_states: torch.Tensor,
346 + context_positions: torch.Tensor,
347 + context_slot_mappings: list[torch.Tensor | None] | None = None,
348 + ) -> None:
349 + self.model.precompute_and_store_context_kv(
350 + context_states, context_positions, context_slot_mappings
351 + )
352 +
353 + def forward(
354 + self,
355 + input_ids: torch.Tensor,
356 + positions: torch.Tensor,
357 + inputs_embeds: torch.Tensor | None = None,
358 + ) -> torch.Tensor:
359 + # Returns the pre-norm collapsed head hidden ([T, hidden_size]).
360 + return self.model(input_ids, positions, inputs_embeds)
361 +
362 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
363 + """Base logits U_k = lm_head(norm(head_hidden))."""
364 + return self.logits_processor(self.lm_head, self.model.norm(hidden_states))
365 +
366 + def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
367 + # Full-vocab draft: base logits, no d2t scatter.
368 + return self.compute_logits(hidden_states)
369 +
370 + def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor:
371 + return draft_ids # full-vocab: draft ids are target ids
372 +
373 + def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor:
374 + return self.model.markov_head.embed(token_ids)
375 +
376 + def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor:
377 + return self.model.markov_head.bias(markov_embed, self.logits_processor)
378 +
379 + def compute_confidence(
380 + self, head_hidden: torch.Tensor, markov_embed: torch.Tensor
381 + ) -> torch.Tensor:
382 + """Per-position acceptance probability for each drafted token."""
383 + assert self.model.confidence_head is not None
384 + return torch.sigmoid(self.model.confidence_head(head_hidden, markov_embed))
385 +
386 + # --- Weight loading ----------------------------------------------------
387 +
388 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
389 + """Load the ``mtp.{0,1,2}.*`` draft weights from the target checkpoint.
390 +
391 + Non-mtp weights (embed/head/main layers) belong to the target model and
392 + are skipped here. ``embed_tokens``/``lm_head`` are aliased from the target.
393 + """
394 + first_layer = self.model.layers[0]
395 + use_mega_moe = first_layer.ffn.use_mega_moe
396 + # Draft MoE layers use the dspark_* expert counts, not the
397 + # backbone's (see DeepseekV4MoE and the reference
398 + # ModelArgs.get_moe_config).
399 + n_draft_experts = (
400 + getattr(self.config, "dspark_n_routed_experts", 0)
401 + or self.config.n_routed_experts
402 + )
403 + if use_mega_moe:
404 + expert_mapping = make_deepseek_v4_expert_params_mapping(n_draft_experts)
405 + else:
406 + expert_mapping = fused_moe_make_expert_params_mapping(
407 + self,
408 + ckpt_gate_proj_name="w1",
409 + ckpt_down_proj_name="w2",
410 + ckpt_up_proj_name="w3",
411 + num_experts=n_draft_experts,
412 + )
413 + expert_scale_suffix = (
414 + ".weight_scale"
415 + if getattr(self.config, "expert_dtype", "fp4") == "fp4"
416 + else ".weight_scale_inv"
417 + )
418 +
419 + # (param_name, ckpt_shard_name, shard_id) for non-expert stacked params.
420 + stacked_params_mapping = [
421 + ("gate_up_proj", "w1", 0),
422 + ("gate_up_proj", "w3", 1),
423 + ("attn.fused_wqa_wkv", "attn.wq_a", 0),
424 + ("attn.fused_wqa_wkv", "attn.wkv", 1),
425 + ]
426 +
427 + params_dict = dict(self.named_parameters())
428 + loaded_params: set[str] = set()
429 + loaded_confidence_head = False
430 +
431 + tp_size = get_tensor_model_parallel_world_size()
432 + tp_rank = get_tensor_model_parallel_rank()
433 + n_local_head = self.config.num_attention_heads // tp_size
434 + head_start = n_local_head * tp_rank
435 + head_end = n_local_head * (tp_rank + 1)
436 +
437 + for name, loaded_weight in weights:
438 + mapped = self._remap_dspark_name(name)
439 + if mapped is None:
440 + continue
441 + name = mapped
442 + if "confidence_head." in name:
443 + loaded_confidence_head = True
444 +
445 + # ``.scale`` -> per-method scale suffix.
446 + if name.endswith(".scale"):
447 + suffix = (
448 + expert_scale_suffix
449 + if _EXPERT_SCALE_RE.search(name)
450 + else f".{self.linear_scale_name}"
451 + )
452 + name = name.removesuffix(".scale") + suffix
453 + if ".shared_experts.w2" in name:
454 + name = name.replace(".shared_experts.w2", ".shared_experts.down_proj")
455 + if self.pad_shared_expert and ".shared_experts." in name:
456 + loaded_weight = DeepseekV4Model._pad_shared_expert_weight(
457 + self.quant_config, name, loaded_weight
458 + )
459 +
460 + # E8M0 expert scales: keep raw exponent bytes.
461 + if ".experts." in name:
462 + if (
463 + "weight_scale" in name
464 + and loaded_weight.dtype == torch.float8_e8m0fnu
465 + ):
466 + loaded_weight = loaded_weight.view(torch.uint8)
467 + for param_name, weight_name, expert_id, shard_id in expert_mapping:
468 + if weight_name not in name:
469 + continue
470 + name_mapped = name.replace(weight_name, param_name)
471 + param = params_dict[name_mapped]
472 + success = param.weight_loader(
473 + param,
474 + loaded_weight,
475 + name_mapped,
476 + shard_id=shard_id,
477 + expert_id=expert_id,
478 + return_success=True,
479 + )
480 + if success:
481 + loaded_params.add(name_mapped)
482 + break
483 + continue
484 +
485 + # Stacked rules only apply to decoder-layer weights. Head-stack params
486 + # (main_proj/norm/markov_head/confidence_head) load directly —
487 + # otherwise e.g. "markov_w1" would collide with the "w1" shard rule.
488 + is_layer_param = name.startswith("model.layers.")
489 + for param_name, weight_name, stacked_shard_id in stacked_params_mapping:
490 + if not is_layer_param or weight_name not in name:
491 + continue
492 + name = name.replace(weight_name, param_name)
493 + param = params_dict[name]
494 + param.weight_loader(param, loaded_weight, stacked_shard_id)
495 + loaded_params.add(name)
496 + break
497 + else:
498 + if "attn_sink" in name:
499 + narrow = loaded_weight[head_start:head_end]
500 + params_dict[name][: narrow.shape[0]].copy_(narrow)
501 + loaded_params.add(name)
502 + continue
503 + if name.endswith(".ffn.gate.bias"):
504 + name = name.replace(
505 + ".ffn.gate.bias", ".ffn.gate.e_score_correction_bias"
506 + )
507 + param = params_dict[name]
508 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
509 + weight_loader(param, loaded_weight)
510 + loaded_params.add(name)
511 +
512 + if self.model.confidence_head is not None and not loaded_confidence_head:
513 + self.model.confidence_head = None
514 + self.process_weights_after_loading()
515 + logger.info_once("DSpark draft model loaded: %d params", len(loaded_params))
516 + return loaded_params
517 +
518 + def _finalize_moe(self) -> None:
519 + for layer in self.model.layers:
520 + layer.ffn.finalize_mega_moe_weights()
521 +
522 + def process_weights_after_loading(self) -> None:
523 + self._finalize_moe()
524 +
525 + def _remap_dspark_name(self, name: str) -> str | None:
526 + """Map a checkpoint ``mtp.{i}.*`` name to this model's parameter path.
527 +
528 + Returns None for non-mtp weights (owned by the target model).
529 + """
530 + m = re.match(r"mtp\.(\d+)\.(.*)", name)
531 + if m is None:
532 + return None
533 + stage = int(m.group(1))
534 + rest = m.group(2)
535 + if rest.startswith("confidence_head.") and self.model.confidence_head is None:
536 + return None
537 + # The checkpoint calls the Markov head's factors ``embed``/``head``;
538 + # DSparkMarkovHead registers them as ``markov_w1``/``markov_w2``.
539 + if rest.startswith("markov_head.embed."):
540 + return "model.markov_head.markov_w1." + rest.removeprefix(
541 + "markov_head.embed."
542 + )
543 + if rest.startswith("markov_head.head."):
544 + return "model.markov_head.markov_w2." + rest.removeprefix(
545 + "markov_head.head."
546 + )
547 + # Head-stack params live at model level (mtp.last), context combiner at
548 + # model level (mtp.0); everything else is a per-layer decoder block.
549 + head_prefixes = (
550 + "norm.",
551 + "markov_head.",
552 + "confidence_head.",
553 + )
554 + if rest.startswith(("main_proj.", "main_norm.")) or rest.startswith(
555 + head_prefixes
556 + ):
557 + return f"model.{rest}"
558 + return f"model.layers.{stage}.{rest}"
vllm/models/deepseek_v4_1/nvidia/flashinfer_sparse.py new
+908
@@ -0,0 +1,908 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""DeepSeek V4 FlashInfer sparse MLA backend."""
4 +
5 +from typing import TYPE_CHECKING, ClassVar, cast
6 +
7 +import torch
8 +
9 +from vllm.config import VllmConfig
10 +from vllm.config.cache import CacheDType
11 +from vllm.forward_context import get_forward_context
12 +from vllm.models.deepseek_v4.nvidia.ops.o_proj import (
13 + compute_fp8_einsum_recipe,
14 + deep_gemm_fp8_o_proj,
15 +)
16 +from vllm.models.deepseek_v4_1.attention import DeepseekV4Attention
17 +from vllm.models.deepseek_v4_1.common.ops import (
18 + build_flashinfer_mixed_sparse_indices,
19 + compute_global_topk_indices_and_lens,
20 +)
21 +from vllm.models.deepseek_v4_1.sparse_mla import (
22 + DeepseekV4FlashMLAMetadata,
23 + DeepseekV4SparseMLABackend,
24 + DeepseekV4SparseMLAMetadataBuilder,
25 + DeepseekV41SparseSWAMetadataBuilder,
26 +)
27 +from vllm.platforms.interface import DeviceCapability
28 +from vllm.utils.flashinfer import flashinfer_trtllm_batch_decode_sparse_mla_dsv4
29 +from vllm.v1.attention.backend import AttentionCGSupport, MultipleOf
30 +from vllm.v1.attention.backends.mla.compressor_utils import (
31 + get_dspark_swa_index_width,
32 +)
33 +from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWABackend
34 +
35 +if TYPE_CHECKING:
36 + from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
37 +
38 +_FLASHINFER_DSV4_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
39 +_flashinfer_dsv4_workspace_by_device: dict[torch.device, torch.Tensor] = {}
40 +
41 +
42 +def _get_flashinfer_dsv4_workspace(device: torch.device) -> torch.Tensor:
43 + workspace = _flashinfer_dsv4_workspace_by_device.get(device)
44 + if workspace is None:
45 + workspace = torch.zeros(
46 + _FLASHINFER_DSV4_WORKSPACE_BUFFER_SIZE,
47 + dtype=torch.uint8,
48 + device=device,
49 + )
50 + _flashinfer_dsv4_workspace_by_device[device] = workspace
51 + return workspace
52 +
53 +
54 +def _packed_block_span(pool: torch.Tensor) -> int:
55 + """Per-block stride of ``pool`` in tokens (``stride(0)//stride(-2)``): ==
56 + block_size for unpacked KV, larger when packed (#44577). Raises if not
57 + token-aligned."""
58 + block_stride = pool.stride(0)
59 + token_stride = pool.stride(-2)
60 + if block_stride % token_stride != 0:
61 + raise NotImplementedError(
62 + "FLASHINFER_MLA_SPARSE_DSV4 packed KV requires the per-block stride "
63 + f"({block_stride}) to be a multiple of the per-token stride "
64 + f"({token_stride}); this layout is not supported yet."
65 + )
66 + return block_stride // token_stride
67 +
68 +
69 +# Sparse MLA h_q counts accepted natively (flashinfer>=0.6.14, #3545).
70 +_SPARSE_MLA_SUPPORTED_Q_HEADS = (8, 16, 32, 64, 128)
71 +
72 +
73 +def _pad_to_supported_q_heads(num_heads: int) -> int:
74 + for supported in _SPARSE_MLA_SUPPORTED_Q_HEADS:
75 + if num_heads <= supported:
76 + return supported
77 + raise ValueError(
78 + f"DeepseekV4 FlashInfer MLA Sparse does not support {num_heads} heads "
79 + "(sparse MLA kernel requires h_q in {8, 16, 32, 64, 128})."
80 + )
81 +
82 +
83 +def _required_sm120_sparse_topk(vllm_config: VllmConfig, window_size: int) -> int:
84 + """Return the SM120 DSV4 SWA specialization needed by this model."""
85 + if not vllm_config.attention_config.use_non_causal:
86 + return window_size
87 + speculative_config = vllm_config.speculative_config
88 + if speculative_config is None:
89 + return window_size
90 + return get_dspark_swa_index_width(
91 + window_size,
92 + speculative_config.num_speculative_tokens,
93 + )
94 +
95 +
96 +class DeepseekV4FlashInferMLASparseBackend(DeepseekV4SparseMLABackend):
97 + """FlashInfer backend using the DSv4 sparse metadata/cache layout.
98 +
99 + Inherits the base and backend reuses its``DeepseekV4SparseMLAMetadataBuilder``
100 + """
101 +
102 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
103 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
104 + "auto",
105 + "bfloat16",
106 + "fp8",
107 + "fp8_e4m3",
108 + "fp8_ds_mla",
109 + ]
110 +
111 + @staticmethod
112 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
113 + return [128]
114 +
115 + @staticmethod
116 + def get_name() -> str:
117 + return "FLASHINFER_MLA_SPARSE_DSV41"
118 +
119 + @classmethod
120 + def get_supported_head_sizes(cls) -> list[int]:
121 + return [512]
122 +
123 + @classmethod
124 + def supports_sink(cls) -> bool:
125 + return True
126 +
127 + @classmethod
128 + def is_sparse(cls) -> bool:
129 + return True
130 +
131 + @classmethod
132 + def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
133 + return capability.major in [10, 12]
134 +
135 + @classmethod
136 + def supports_combination(
137 + cls,
138 + head_size: int,
139 + dtype: torch.dtype,
140 + kv_cache_dtype: CacheDType | None,
141 + block_size: int | None,
142 + use_mla: bool,
143 + has_sink: bool,
144 + use_sparse: bool,
145 + use_mm_prefix: bool,
146 + device_capability: DeviceCapability,
147 + ) -> str | None:
148 + if device_capability.major == 10:
149 + if kv_cache_dtype == "fp8_ds_mla":
150 + return (
151 + "FLASHINFER_MLA_SPARSE_DSV4 SM10x uses the plain "
152 + "per-tensor FP8 KV layout, not fp8_ds_mla"
153 + )
154 + if kv_cache_dtype not in (None, "auto", "bfloat16", "fp8", "fp8_e4m3"):
155 + return "kv_cache_dtype not supported"
156 + return None
157 + if device_capability.major == 12:
158 + if kv_cache_dtype not in ("fp8", "fp8_e4m3", "fp8_ds_mla"):
159 + return "kv_cache_dtype not supported"
160 + from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120
161 +
162 + if not has_flashinfer_sparse_mla_sm120():
163 + return (
164 + "FLASHINFER_MLA_SPARSE_DSV4 SM120 requires FlashInfer's "
165 + "sparse MLA decode API"
166 + )
167 + return None
168 + return "FLASHINFER_MLA_SPARSE_DSV4 requires SM10x or SM12x"
169 +
170 + @staticmethod
171 + def get_builder_cls() -> type["DeepseekV4FlashInferSparseMLAMetadataBuilder"]:
172 + return DeepseekV4FlashInferSparseMLAMetadataBuilder
173 +
174 +
175 +class DeepseekV4FlashInferSparseMLAMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder):
176 + """Varlen-capable metadata builder for the FlashInfer sparse MLA backend."""
177 +
178 + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
179 +
180 +
181 +class DeepseekSparseSWAFlashInferMetadataBuilder(DeepseekV41SparseSWAMetadataBuilder):
182 + """SWA metadata for the FlashInfer sparse decode path (varlen decode)."""
183 +
184 + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
185 +
186 +
187 +class DeepseekSparseSWAFlashInferBackend(DeepseekSparseSWABackend):
188 + @staticmethod
189 + def get_builder_cls() -> type[DeepseekSparseSWAFlashInferMetadataBuilder]:
190 + return DeepseekSparseSWAFlashInferMetadataBuilder
191 +
192 +
193 +class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
194 + """FlashInfer TRTLLM-gen sparse MLA attention layer for SM100 DeepSeek V4."""
195 +
196 + backend_cls = DeepseekV4FlashInferMLASparseBackend
197 + swa_backend_cls = DeepseekSparseSWAFlashInferBackend
198 + use_fp8_ds_mla_layout: ClassVar[bool] = False
199 +
200 + @classmethod
201 + def get_padded_num_q_heads(cls, num_heads: int) -> int:
202 + return _pad_to_supported_q_heads(num_heads)
203 +
204 + def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
205 + return deep_gemm_fp8_o_proj(
206 + o,
207 + positions,
208 + self.rotary_emb.cos_sin_cache,
209 + self.wo_a,
210 + self.wo_b,
211 + n_groups=self.n_local_groups,
212 + heads_per_group=self.n_local_heads // self.n_local_groups,
213 + nope_dim=self.nope_head_dim,
214 + rope_dim=self.rope_head_dim,
215 + o_lora_rank=self.o_lora_rank,
216 + einsum_recipe=self._einsum_recipe,
217 + tma_aligned_scales=self._tma_aligned_scales,
218 + )
219 +
220 + def __init__(self, *args, **kwargs) -> None:
221 + super().__init__(*args, **kwargs)
222 + self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe(
223 + self._o_proj_block_size
224 + )
225 + # Per-tensor FP8 scale buffers + precomputed scalar BMM scales. Only the
226 + # per-tensor FP8 cache path consumes these; bf16 reads ``self.scale``.
227 + if self.kv_cache_torch_dtype != torch.float8_e4m3fn:
228 + return
229 + fp8_q_scale = 1.0
230 + fp8_kv_scale = 1.0
231 + self.register_buffer(
232 + "_flashinfer_fp8_q_scale",
233 + torch.tensor([fp8_q_scale], dtype=torch.float32),
234 + persistent=False,
235 + )
236 + self.register_buffer(
237 + "_flashinfer_fp8_q_scale_inv",
238 + torch.tensor([1.0 / fp8_q_scale], dtype=torch.float32),
239 + persistent=False,
240 + )
241 + self.register_buffer(
242 + "_flashinfer_fp8_kv_scale",
243 + torch.tensor([fp8_kv_scale], dtype=torch.float32),
244 + persistent=False,
245 + )
246 + # TRTLLM-gen takes scalar scale args on a distinct C++ path vs
247 + # one-element tensors, so these are Python floats.
248 + self._flashinfer_fp8_bmm1_scale = self.scale * fp8_q_scale * fp8_kv_scale
249 + self._flashinfer_fp8_bmm2_scale = fp8_kv_scale
250 +
251 + def forward_mqa(
252 + self,
253 + q: torch.Tensor,
254 + kv: torch.Tensor,
255 + positions: torch.Tensor,
256 + output: torch.Tensor,
257 + ) -> None:
258 + # The TRTLLM-gen kernel requires h_q in {64, 128}, so the output buffer
259 + # is allocated at the padded head count while q arrives at the local
260 + # head count; _forward pads q to match before the launcher.
261 + assert output.shape[0] == q.shape[0] and output.shape[-1] == q.shape[-1], (
262 + f"output buffer shape {output.shape} incompatible with q shape {q.shape}"
263 + )
264 + assert output.shape[1] >= q.shape[1], (
265 + f"output heads {output.shape[1]} must be >= q heads {q.shape[1]}"
266 + )
267 + # Per-tensor FP8 q produces a bf16 attention output.
268 + expected_output_dtype = (
269 + torch.bfloat16 if q.dtype == torch.float8_e4m3fn else q.dtype
270 + )
271 + assert output.dtype == expected_output_dtype, (
272 + f"output dtype {output.dtype} must match expected {expected_output_dtype} "
273 + f"for q dtype {q.dtype}"
274 + )
275 +
276 + forward_context = get_forward_context()
277 + attn_metadata = forward_context.attn_metadata
278 + if attn_metadata is None:
279 + # Warmup dummy run: FlashInfer reads the cache directly and lazily
280 + # allocates its workspace, so nothing to reserve here.
281 + output.zero_()
282 + return
283 +
284 + assert isinstance(attn_metadata, dict)
285 + # Compressed-cache metadata lives on the kv-source layer's prefix;
286 + # consumers share that cache and its block table.
287 + flashmla_metadata = cast(
288 + DeepseekV4FlashMLAMetadata | None,
289 + attn_metadata.get(self.compressed_cache_prefix)
290 + if self.compressed_cache_prefix is not None
291 + else None,
292 + )
293 + swa_metadata = cast(
294 + "DeepseekSparseSWAMetadata | None",
295 + attn_metadata.get(self.swa_cache_layer.prefix),
296 + )
297 + assert swa_metadata is not None
298 +
299 + swa_only = self.compress_ratio == 0
300 + # SWA-only layers have no compressed KV cache; consumers read the kv
301 + # source's cache.
302 + self_kv_cache = None if swa_only else self._compressed_kv_cache()
303 + swa_kv_cache = self.swa_cache_layer.kv_cache
304 +
305 + self._forward(
306 + q=q,
307 + kv_cache=self_kv_cache,
308 + swa_k_cache=swa_kv_cache,
309 + swa_metadata=swa_metadata,
310 + attn_metadata=flashmla_metadata,
311 + swa_only=swa_only,
312 + output=output,
313 + )
314 +
315 + def _build_sparse_index_metadata(
316 + self,
317 + kv_cache: torch.Tensor | None,
318 + swa_k_cache: torch.Tensor,
319 + swa_metadata: "DeepseekSparseSWAMetadata",
320 + attn_metadata: DeepseekV4FlashMLAMetadata | None,
321 + swa_only: bool,
322 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
323 + """Build the combined sparse-index tensors for the mixed batch.
324 +
325 + Returns ``(compressed_kv_cache, seq_lens, sparse_indices,
326 + sparse_topk_lens)``.
327 + """
328 + num_decodes = swa_metadata.num_decodes
329 + num_prefills = swa_metadata.num_prefills
330 + num_decode_tokens = swa_metadata.num_decode_tokens
331 + num_prefill_tokens = swa_metadata.num_prefill_tokens
332 + num_reqs = num_decodes + num_prefills
333 + num_tokens = num_decode_tokens + num_prefill_tokens
334 +
335 + assert swa_metadata.seq_lens is not None
336 + assert swa_metadata.query_start_loc is not None
337 + assert swa_metadata.token_to_req_indices is not None
338 + assert swa_metadata.decode_swa_indices is not None
339 + assert swa_metadata.block_table is not None
340 +
341 + decode_swa_indices = swa_metadata.decode_swa_indices.reshape(
342 + num_decode_tokens, swa_metadata.decode_swa_width
343 + )
344 + decode_compressed_topk_lens = None
345 + decode_compressed_indices_are_local = False
346 + decode_is_valid_token = None
347 +
348 + if swa_only:
349 + assert self.topk_indices_buffer is not None
350 + compressed_kv_cache = swa_k_cache
351 + decode_compressed_indices = None
352 + prefill_topk_indices = self.topk_indices_buffer[
353 + num_decode_tokens:num_tokens, :0
354 + ]
355 + compressed_block_table = None
356 + compressed_block_size = swa_metadata.block_size
357 + top_k = 0
358 + else:
359 + assert kv_cache is not None
360 + assert attn_metadata is not None
361 + assert self.topk_indices_buffer is not None
362 + assert swa_metadata.is_valid_token is not None
363 + compressed_kv_cache = kv_cache
364 + compressed_block_table = attn_metadata.block_table[:num_reqs]
365 + compressed_block_size = attn_metadata.block_size // self.compress_ratio
366 +
367 + # Local indices filled by the index-source layer's indexer.
368 + if num_prefill_tokens > 0:
369 + prefill_topk_indices = self.topk_indices_buffer[
370 + num_decode_tokens:num_tokens
371 + ]
372 + top_k = prefill_topk_indices.shape[-1]
373 + else:
374 + prefill_topk_indices = self.topk_indices_buffer[:0, :0]
375 + top_k = 0
376 +
377 + decode_compressed_indices_are_local = True
378 + decode_is_valid_token = swa_metadata.is_valid_token[:num_decode_tokens]
379 + if num_decode_tokens > 0:
380 + decode_compressed_indices = self.topk_indices_buffer[:num_decode_tokens]
381 + else:
382 + # Keep the logical width aligned with the mixed-batch case so
383 + # pure-prefill steps reuse the same Triton specialization.
384 + decode_compressed_indices = prefill_topk_indices[:0]
385 +
386 + query_start_loc = swa_metadata.query_start_loc[: num_reqs + 1]
387 + seq_lens = swa_metadata.seq_lens[:num_reqs]
388 + assert seq_lens.dtype == torch.int32
389 + # SWA-only layers all build the same mixed sparse indices, so the first
390 + # one caches them for the step; indexer layers depend on their own topk
391 + # indices and stay uncached.
392 + cached_sparse = (
393 + swa_metadata.flashinfer_sparse_index_cache.get("swa_only")
394 + if swa_only
395 + else None
396 + )
397 + if cached_sparse is None:
398 + swa_block_span = _packed_block_span(swa_k_cache)
399 + compressed_block_span = _packed_block_span(compressed_kv_cache)
400 + sparse_indices, sparse_topk_lens = build_flashinfer_mixed_sparse_indices(
401 + decode_swa_indices,
402 + decode_compressed_indices,
403 + decode_compressed_topk_lens,
404 + prefill_topk_indices[:num_prefill_tokens],
405 + query_start_loc,
406 + seq_lens,
407 + swa_metadata.token_to_req_indices[:num_tokens],
408 + swa_metadata.block_table[:num_reqs],
409 + swa_metadata.block_size,
410 + compressed_block_table,
411 + compressed_block_size,
412 + self.window_size,
413 + self.compress_ratio,
414 + top_k,
415 + decode_compressed_indices_are_local=decode_compressed_indices_are_local,
416 + decode_is_valid_token=decode_is_valid_token,
417 + swa_block_span=swa_block_span,
418 + compressed_block_span=compressed_block_span,
419 + prefill_left_visible=swa_metadata.prefill_left_visible,
420 + prefill_right_visible=swa_metadata.prefill_right_visible,
421 + # getattr for tests that bypass __init__ via object.__new__.
422 + max_image_tokens=getattr(self, "max_image_tokens", 0),
423 + )
424 + if swa_only:
425 + swa_metadata.flashinfer_sparse_index_cache["swa_only"] = (
426 + sparse_indices,
427 + sparse_topk_lens,
428 + )
429 + else:
430 + sparse_indices, sparse_topk_lens = cached_sparse
431 + return compressed_kv_cache, seq_lens, sparse_indices, sparse_topk_lens
432 +
433 + def _forward(
434 + self,
435 + q: torch.Tensor,
436 + kv_cache: torch.Tensor | None,
437 + swa_k_cache: torch.Tensor,
438 + swa_metadata: "DeepseekSparseSWAMetadata",
439 + attn_metadata: DeepseekV4FlashMLAMetadata | None,
440 + swa_only: bool,
441 + output: torch.Tensor,
442 + ) -> None:
443 + assert self.kv_cache_torch_dtype in (torch.bfloat16, torch.float8_e4m3fn)
444 + num_decodes = swa_metadata.num_decodes
445 + num_prefills = swa_metadata.num_prefills
446 + num_decode_tokens = swa_metadata.num_decode_tokens
447 + num_prefill_tokens = swa_metadata.num_prefill_tokens
448 + num_reqs = num_decodes + num_prefills
449 + num_tokens = num_decode_tokens + num_prefill_tokens
450 + if num_tokens == 0:
451 + return
452 +
453 + (
454 + compressed_kv_cache,
455 + seq_lens,
456 + sparse_indices,
457 + sparse_topk_lens,
458 + ) = self._build_sparse_index_metadata(
459 + kv_cache=kv_cache,
460 + swa_k_cache=swa_k_cache,
461 + swa_metadata=swa_metadata,
462 + attn_metadata=attn_metadata,
463 + swa_only=swa_only,
464 + )
465 +
466 + # CUDA graph execution can pad q/output past the scheduled token count;
467 + # restrict to the real tokens (the launcher validates sparse indices).
468 + query = q[:num_tokens]
469 + output = output[:num_tokens]
470 + bmm1_scale: float | torch.Tensor = self.scale
471 + bmm2_scale: float | torch.Tensor = 1.0
472 + if self.kv_cache_torch_dtype == torch.float8_e4m3fn:
473 + assert query.dtype == torch.float8_e4m3fn
474 + bmm1_scale = self._flashinfer_fp8_bmm1_scale
475 + bmm2_scale = self._flashinfer_fp8_bmm2_scale
476 + else:
477 + assert query.dtype == torch.bfloat16
478 + query = query.contiguous()
479 +
480 + # The TRTLLM-gen sparse-MLA kernel requires h_q in {64, 128}; zero-pad
481 + # the query heads to the allocated output head count. Padded heads attend
482 + # to the shared KV and are sliced off downstream (output is padded too).
483 + padded_heads = output.shape[1]
484 + if query.shape[1] < padded_heads:
485 + padded_query = query.new_zeros(
486 + (query.shape[0], padded_heads, query.shape[2])
487 + )
488 + padded_query[:, : query.shape[1], :] = query
489 + query = padded_query
490 +
491 + workspace = _get_flashinfer_dsv4_workspace(q.device)
492 + query_start_loc = swa_metadata.query_start_loc
493 + query_start_loc_cpu = swa_metadata.query_start_loc_cpu
494 + assert query_start_loc is not None and query_start_loc_cpu is not None
495 +
496 + # Keep the TRTLLM-gen decode/prefill split: the launcher is tuned for
497 + # uniform-q batches, and this avoids flattening mixed batches into one call.
498 + if num_decode_tokens > 0:
499 + decode_cu = query_start_loc[: num_decodes + 1]
500 + flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
501 + query=query[:num_decode_tokens],
502 + swa_kv_cache=swa_k_cache,
503 + workspace_buffer=workspace,
504 + sparse_indices=sparse_indices[:num_decode_tokens],
505 + compressed_kv_cache=compressed_kv_cache,
506 + sparse_topk_lens=sparse_topk_lens[:num_decode_tokens],
507 + seq_lens=seq_lens[:num_decodes],
508 + out=output[:num_decode_tokens],
509 + bmm1_scale=bmm1_scale,
510 + bmm2_scale=bmm2_scale,
511 + sinks=self.attn_sink,
512 + cum_seq_lens_q=decode_cu,
513 + max_q_len=swa_metadata.max_decode_query_len,
514 + )
515 +
516 + if num_prefill_tokens > 0:
517 + # The prefill query view re-anchors at offset 0, so rebase the
518 + # cumulative query offsets to start at 0.
519 + prefill_cu = (
520 + query_start_loc[num_decodes : num_reqs + 1]
521 + - query_start_loc[num_decodes]
522 + )
523 + prefill_cu_cpu = query_start_loc_cpu[num_decodes : num_reqs + 1]
524 + prefill_lens_cpu = prefill_cu_cpu[1:] - prefill_cu_cpu[:-1]
525 + flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
526 + query=query[num_decode_tokens:num_tokens],
527 + swa_kv_cache=swa_k_cache,
528 + workspace_buffer=workspace,
529 + sparse_indices=sparse_indices[num_decode_tokens:num_tokens],
530 + compressed_kv_cache=compressed_kv_cache,
531 + sparse_topk_lens=sparse_topk_lens[num_decode_tokens:num_tokens],
532 + seq_lens=seq_lens[num_decodes:num_reqs],
533 + out=output[num_decode_tokens:num_tokens],
534 + bmm1_scale=bmm1_scale,
535 + bmm2_scale=bmm2_scale,
536 + sinks=self.attn_sink,
537 + cum_seq_lens_q=prefill_cu,
538 + max_q_len=int(prefill_lens_cpu.max().item()),
539 + )
540 +
541 +
542 +class DeepseekV4FlashInferSM120Attention(DeepseekV4Attention):
543 + """DeepSeek V4 sparse MLA attention through FlashInfer's SM120 kernels."""
544 +
545 + backend_cls = DeepseekV4FlashInferMLASparseBackend
546 + swa_backend_cls = DeepseekSparseSWAFlashInferBackend
547 + use_fp8_ds_mla_layout: ClassVar[bool] = True
548 +
549 + @staticmethod
550 + def _get_workspace(device: torch.device) -> torch.Tensor:
551 + return _get_flashinfer_dsv4_workspace(device)
552 +
553 + @staticmethod
554 + def _as_sparse_cache(kv_cache: torch.Tensor) -> torch.Tensor:
555 + if kv_cache.dtype == torch.float8_e4m3fn:
556 + kv_cache = kv_cache.view(torch.uint8)
557 + if kv_cache.dim() == 4:
558 + return kv_cache
559 + return kv_cache.unsqueeze(-2)
560 +
561 + @classmethod
562 + def get_padded_num_q_heads(cls, num_heads: int) -> int:
563 + return _pad_to_supported_q_heads(num_heads)
564 +
565 + def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
566 + return deep_gemm_fp8_o_proj(
567 + o,
568 + positions,
569 + self.rotary_emb.cos_sin_cache,
570 + self.wo_a,
571 + self.wo_b,
572 + n_groups=self.n_local_groups,
573 + heads_per_group=self.n_local_heads // self.n_local_groups,
574 + nope_dim=self.nope_head_dim,
575 + rope_dim=self.rope_head_dim,
576 + o_lora_rank=self.o_lora_rank,
577 + einsum_recipe=self._einsum_recipe,
578 + tma_aligned_scales=self._tma_aligned_scales,
579 + )
580 +
581 + def __init__(self, vllm_config: VllmConfig, *args, **kwargs) -> None:
582 + super().__init__(vllm_config, *args, **kwargs)
583 + from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120_config
584 +
585 + required_topk = _required_sm120_sparse_topk(vllm_config, self.window_size)
586 + if not has_flashinfer_sparse_mla_sm120_config(self.padded_heads, required_topk):
587 + raise RuntimeError(
588 + "FLASHINFER_MLA_SPARSE_DSV4 on SM120 requires a FlashInfer "
589 + "DSV4 sparse MLA decode specialization for "
590 + f"(num_q_heads={self.padded_heads}, top_k={required_topk}). "
591 + "Install a FlashInfer build containing "
592 + "flashinfer-ai/flashinfer#4380."
593 + )
594 + self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe(
595 + self._o_proj_block_size
596 + )
597 + # Per-tensor FP8 cache path scales.
598 + if self.kv_cache_torch_dtype != torch.float8_e4m3fn:
599 + return
600 + fp8_q_scale = 1.0
601 + fp8_kv_scale = 1.0
602 + self.register_buffer(
603 + "_flashinfer_fp8_q_scale",
604 + torch.tensor([fp8_q_scale], dtype=torch.float32),
605 + persistent=False,
606 + )
607 + self.register_buffer(
608 + "_flashinfer_fp8_q_scale_inv",
609 + torch.tensor([1.0 / fp8_q_scale], dtype=torch.float32),
610 + persistent=False,
611 + )
612 + self.register_buffer(
613 + "_flashinfer_fp8_kv_scale",
614 + torch.tensor([fp8_kv_scale], dtype=torch.float32),
615 + persistent=False,
616 + )
617 + # FlashInfer expects scalar scale arguments for this path.
618 + self._flashinfer_fp8_bmm1_scale = self.scale * fp8_q_scale * fp8_kv_scale
619 + self._flashinfer_fp8_bmm2_scale = fp8_kv_scale
620 +
621 + def _reserve_empty_forward_workspace(self) -> None:
622 + self._get_workspace(
623 + torch.device("cuda", torch.accelerator.current_device_index())
624 + )
625 +
626 + def _forward_sparse_impl(
627 + self,
628 + q: torch.Tensor,
629 + output: torch.Tensor,
630 + flashmla_metadata: DeepseekV4FlashMLAMetadata | None,
631 + swa_metadata: "DeepseekSparseSWAMetadata",
632 + self_kv_cache: torch.Tensor | None,
633 + swa_kv_cache: torch.Tensor,
634 + swa_only: bool,
635 + ) -> None:
636 + num_decode_tokens = swa_metadata.num_decode_tokens
637 + if swa_metadata.num_prefills > 0:
638 + self._forward_prefill(
639 + q=q[num_decode_tokens:],
640 + compressed_k_cache=self_kv_cache,
641 + swa_k_cache=swa_kv_cache,
642 + output=output[num_decode_tokens:],
643 + attn_metadata=flashmla_metadata,
644 + swa_metadata=swa_metadata,
645 + )
646 + if swa_metadata.num_decodes > 0:
647 + self._forward_decode(
648 + q=q[:num_decode_tokens],
649 + kv_cache=self_kv_cache,
650 + swa_metadata=swa_metadata,
651 + attn_metadata=flashmla_metadata,
652 + swa_only=swa_only,
653 + output=output[:num_decode_tokens],
654 + )
655 +
656 + def forward_mqa(
657 + self,
658 + q: torch.Tensor,
659 + kv: torch.Tensor,
660 + positions: torch.Tensor,
661 + output: torch.Tensor,
662 + ) -> None:
663 + # Output may be padded to backend-supported head counts.
664 + assert output.shape[0] == q.shape[0] and output.shape[-1] == q.shape[-1], (
665 + f"output buffer shape {output.shape} incompatible with q shape {q.shape}"
666 + )
667 + assert output.shape[1] >= q.shape[1], (
668 + f"output heads {output.shape[1]} must be >= q heads {q.shape[1]}"
669 + )
670 + # Per-tensor FP8 q produces a bf16 attention output.
671 + expected_output_dtype = (
672 + torch.bfloat16 if q.dtype == torch.float8_e4m3fn else q.dtype
673 + )
674 + assert output.dtype == expected_output_dtype, (
675 + f"output dtype {output.dtype} must match expected {expected_output_dtype} "
676 + f"for q dtype {q.dtype}"
677 + )
678 +
679 + forward_context = get_forward_context()
680 + attn_metadata = forward_context.attn_metadata
681 + if attn_metadata is None:
682 + self._reserve_empty_forward_workspace()
683 + output.zero_()
684 + return
685 +
686 + assert isinstance(attn_metadata, dict)
687 + # Compressed-cache metadata lives on the kv-source layer's prefix;
688 + # consumers share that cache and its block table.
689 + flashmla_metadata = cast(
690 + DeepseekV4FlashMLAMetadata | None,
691 + attn_metadata.get(self.compressed_cache_prefix)
692 + if self.compressed_cache_prefix is not None
693 + else None,
694 + )
695 + swa_metadata = cast(
696 + "DeepseekSparseSWAMetadata | None",
697 + attn_metadata.get(self.swa_cache_layer.prefix),
698 + )
699 + assert swa_metadata is not None
700 +
701 + swa_only = self.compress_ratio == 0
702 + # SWA-only layers have no compressed KV cache; consumers read the kv
703 + # source's cache.
704 + self_kv_cache = None if swa_only else self._compressed_kv_cache()
705 + swa_kv_cache = self.swa_cache_layer.kv_cache
706 +
707 + self._forward_sparse_impl(
708 + q=q,
709 + output=output,
710 + flashmla_metadata=flashmla_metadata,
711 + swa_metadata=swa_metadata,
712 + self_kv_cache=self_kv_cache,
713 + swa_kv_cache=swa_kv_cache,
714 + swa_only=swa_only,
715 + )
716 +
717 + def _prepare_query(self, q: torch.Tensor, output: torch.Tensor) -> torch.Tensor:
718 + if self.kv_cache_torch_dtype == torch.float8_e4m3fn:
719 + assert q.dtype == torch.float8_e4m3fn
720 + q = q.to(torch.bfloat16)
721 + else:
722 + assert q.dtype == torch.bfloat16
723 + padded_heads = output.shape[1]
724 + if q.shape[1] < padded_heads:
725 + padded_query = q.new_zeros((q.shape[0], padded_heads, q.shape[2]))
726 + padded_query[:, : q.shape[1], :] = q
727 + q = padded_query
728 + return q.contiguous()
729 +
730 + def _forward_decode(
731 + self,
732 + q: torch.Tensor,
733 + kv_cache: torch.Tensor | None,
734 + swa_metadata: "DeepseekSparseSWAMetadata",
735 + attn_metadata: DeepseekV4FlashMLAMetadata | None,
736 + swa_only: bool,
737 + output: torch.Tensor,
738 + ) -> None:
739 + num_decodes = swa_metadata.num_decodes
740 + num_decode_tokens = swa_metadata.num_decode_tokens
741 +
742 + extra_sparse_indices = None
743 + extra_sparse_lengths = None
744 + if not swa_only:
745 + if attn_metadata is None:
746 + raise RuntimeError(
747 + "Sparse MLA metadata is required for compressed layers."
748 + )
749 + if swa_metadata.is_valid_token is None:
750 + raise RuntimeError(
751 + "SWA validity metadata is required for compressed layers."
752 + )
753 + if self.topk_indices_buffer is None:
754 + raise RuntimeError(
755 + "Compressed-layer decode requires top-k indices from the indexer."
756 + )
757 + # Local indices filled by the index-source layer's indexer.
758 + is_valid = swa_metadata.is_valid_token[:num_decode_tokens]
759 + block_size = attn_metadata.block_size // self.compress_ratio
760 + global_indices, extra_sparse_lengths = compute_global_topk_indices_and_lens(
761 + self.topk_indices_buffer[:num_decode_tokens],
762 + swa_metadata.token_to_req_indices,
763 + attn_metadata.block_table[:num_decodes],
764 + block_size,
765 + is_valid,
766 + )
767 + extra_sparse_indices = global_indices.view(num_decode_tokens, 1, -1)
768 +
769 + swa_indices = swa_metadata.decode_swa_indices
770 + swa_lens = swa_metadata.decode_swa_lens
771 + assert swa_indices is not None
772 + assert swa_lens is not None
773 + q = self._prepare_query(q, output)
774 + swa_cache = self._as_sparse_cache(self.swa_cache_layer.kv_cache)
775 + extra_cache = self._as_sparse_cache(kv_cache) if kv_cache is not None else None
776 + if extra_cache is not None and extra_sparse_indices is None:
777 + raise RuntimeError(
778 + "Compressed sparse MLA decode requires compressed sparse indices."
779 + )
780 + flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
781 + query=q,
782 + swa_kv_cache=swa_cache,
783 + workspace_buffer=self._get_workspace(q.device),
784 + sparse_indices=swa_indices,
785 + compressed_kv_cache=extra_cache,
786 + out=output,
787 + bmm1_scale=self.scale,
788 + sinks=self.attn_sink,
789 + kv_layout="NHD",
790 + swa_topk_lens=swa_lens,
791 + extra_sparse_indices=extra_sparse_indices,
792 + extra_sparse_topk_lens=extra_sparse_lengths,
793 + )
794 +
795 + def _forward_prefill(
796 + self,
797 + q: torch.Tensor,
798 + compressed_k_cache: torch.Tensor | None,
799 + swa_k_cache: torch.Tensor,
800 + output: torch.Tensor,
801 + attn_metadata: DeepseekV4FlashMLAMetadata | None,
802 + swa_metadata: "DeepseekSparseSWAMetadata",
803 + ) -> None:
804 + swa_only = self.compress_ratio == 0
805 +
806 + num_prefills = swa_metadata.num_prefills
807 + num_decodes = swa_metadata.num_decodes
808 + num_decode_tokens = swa_metadata.num_decode_tokens
809 + num_prefill_tokens = swa_metadata.num_prefill_tokens
810 +
811 + query_start_loc_cpu = swa_metadata.query_start_loc_cpu
812 + assert query_start_loc_cpu is not None
813 + prefill_token_base = query_start_loc_cpu[num_decodes]
814 +
815 + extra_sparse_indices: torch.Tensor | None = None
816 + extra_sparse_lengths: torch.Tensor | None = None
817 + if not swa_only:
818 + if self.topk_indices_buffer is None:
819 + raise RuntimeError(
820 + "Compressed-layer prefill requires top-k indices from the indexer."
821 + )
822 + if attn_metadata is None:
823 + raise RuntimeError("Compressed-layer prefill metadata is missing.")
824 + if swa_metadata.token_to_req_indices is None:
825 + raise RuntimeError(
826 + "Compressed-layer prefill request mapping is missing."
827 + )
828 + if swa_metadata.is_valid_token is None:
829 + raise RuntimeError(
830 + "Compressed-layer prefill validity metadata is missing."
831 + )
832 + # Local indices filled by the index-source layer's indexer.
833 + local_topk_indices = self.topk_indices_buffer[
834 + num_decode_tokens : num_decode_tokens + num_prefill_tokens
835 + ]
836 + prefill_token_slice = slice(
837 + num_decode_tokens, num_decode_tokens + num_prefill_tokens
838 + )
839 + block_size = attn_metadata.block_size // self.compress_ratio
840 + extra_sparse_indices, extra_sparse_lengths = (
841 + compute_global_topk_indices_and_lens(
842 + local_topk_indices,
843 + swa_metadata.token_to_req_indices[prefill_token_slice],
844 + attn_metadata.block_table,
845 + block_size,
846 + swa_metadata.is_valid_token[prefill_token_slice],
847 + )
848 + )
849 +
850 + assert swa_metadata.prefill_swa_indices is not None
851 + assert swa_metadata.prefill_swa_lens is not None
852 +
853 + q = self._prepare_query(q, output)
854 + swa_kv_paged = self._as_sparse_cache(swa_k_cache)
855 + if swa_only:
856 + extra_kv_paged = None
857 + else:
858 + if compressed_k_cache is None:
859 + raise RuntimeError(
860 + "Compressed sparse MLA layers require their compressed KV cache."
861 + )
862 + extra_kv_paged = self._as_sparse_cache(compressed_k_cache)
863 +
864 + num_chunks = (
865 + num_prefills + self.PREFILL_CHUNK_SIZE - 1
866 + ) // self.PREFILL_CHUNK_SIZE
867 + for chunk_idx in range(num_chunks):
868 + chunk_start = chunk_idx * self.PREFILL_CHUNK_SIZE
869 + chunk_end = min(chunk_start + self.PREFILL_CHUNK_SIZE, num_prefills)
870 + query_start = (
871 + query_start_loc_cpu[num_decodes + chunk_start] - prefill_token_base
872 + )
873 + query_end = (
874 + query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base
875 + )
876 +
877 + extra_sparse_indices_chunk = (
878 + extra_sparse_indices[query_start:query_end]
879 + if extra_sparse_indices is not None
880 + else None
881 + )
882 + extra_sparse_lengths_chunk = (
883 + extra_sparse_lengths[query_start:query_end]
884 + if extra_sparse_lengths is not None
885 + else None
886 + )
887 +
888 + q_chunk = q[query_start:query_end]
889 + swa_indices_chunk = swa_metadata.prefill_swa_indices[query_start:query_end]
890 + swa_lens_chunk = swa_metadata.prefill_swa_lens[query_start:query_end]
891 + if extra_kv_paged is not None and extra_sparse_indices_chunk is None:
892 + raise RuntimeError(
893 + "Compressed sparse MLA prefill requires compressed sparse indices."
894 + )
895 + flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
896 + query=q_chunk,
897 + swa_kv_cache=swa_kv_paged,
898 + workspace_buffer=self._get_workspace(q.device),
899 + sparse_indices=swa_indices_chunk,
900 + compressed_kv_cache=extra_kv_paged,
901 + out=output[query_start:query_end],
902 + bmm1_scale=self.scale,
903 + sinks=self.attn_sink,
904 + kv_layout="NHD",
905 + swa_topk_lens=swa_lens_chunk,
906 + extra_sparse_indices=extra_sparse_indices_chunk,
907 + extra_sparse_topk_lens=extra_sparse_lengths_chunk,
908 + )
vllm/models/deepseek_v4_1/nvidia/flashmla.py new
+387
@@ -0,0 +1,387 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from typing import TYPE_CHECKING, ClassVar, cast
5 +
6 +import torch
7 +
8 +from vllm.forward_context import get_forward_context
9 +from vllm.models.deepseek_v4.nvidia.ops.o_proj import (
10 + compute_fp8_einsum_recipe,
11 + deep_gemm_fp8_o_proj,
12 +)
13 +from vllm.models.deepseek_v4_1.attention import DeepseekV4Attention
14 +from vllm.models.deepseek_v4_1.common.ops import (
15 + combine_topk_swa_indices,
16 + compute_global_topk_indices_and_lens,
17 + dequantize_and_gather_k_cache,
18 +)
19 +from vllm.models.deepseek_v4_1.sparse_mla import (
20 + DeepseekV4FlashMLABackend,
21 + DeepseekV4FlashMLAMetadata,
22 + DeepseekV41SparseSWAMetadataBuilder,
23 +)
24 +from vllm.utils.math_utils import round_up
25 +from vllm.v1.attention.backend import AttentionCGSupport
26 +from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWABackend
27 +from vllm.v1.attention.ops.flashmla import (
28 + flash_mla_sparse_fwd,
29 + flash_mla_with_kvcache,
30 +)
31 +from vllm.v1.worker.workspace import current_workspace_manager
32 +
33 +if TYPE_CHECKING:
34 + from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
35 +
36 +
37 +class DeepseekSparseSWAFlashMLAMetadataBuilder(DeepseekV41SparseSWAMetadataBuilder):
38 + """SWA metadata for the FlashMLA decode path, which allows varlen decode."""
39 +
40 + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
41 +
42 +
43 +class DeepseekSparseSWAFlashMLABackend(DeepseekSparseSWABackend):
44 + @staticmethod
45 + def get_builder_cls() -> type[DeepseekSparseSWAFlashMLAMetadataBuilder]:
46 + return DeepseekSparseSWAFlashMLAMetadataBuilder
47 +
48 +
49 +class DeepseekV4FlashMLAAttention(DeepseekV4Attention):
50 + """FlashMLA sparse MLA attention layer for DeepSeek V4.1 (CUDA)."""
51 +
52 + backend_cls = DeepseekV4FlashMLABackend
53 + swa_backend_cls = DeepseekSparseSWAFlashMLABackend
54 +
55 + def __init__(self, *args, **kwargs) -> None:
56 + super().__init__(*args, **kwargs)
57 + self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe(
58 + self._o_proj_block_size
59 + )
60 +
61 + def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
62 + return deep_gemm_fp8_o_proj(
63 + o,
64 + positions,
65 + self.rotary_emb.cos_sin_cache,
66 + self.wo_a,
67 + self.wo_b,
68 + n_groups=self.n_local_groups,
69 + heads_per_group=self.n_local_heads // self.n_local_groups,
70 + nope_dim=self.nope_head_dim,
71 + rope_dim=self.rope_head_dim,
72 + o_lora_rank=self.o_lora_rank,
73 + einsum_recipe=self._einsum_recipe,
74 + tma_aligned_scales=self._tma_aligned_scales,
75 + )
76 +
77 + @classmethod
78 + def get_padded_num_q_heads(cls, num_heads: int) -> int:
79 + # FP8 decode kernel only supports h_q = 64 or 128.
80 + if num_heads > 128:
81 + raise ValueError(
82 + f"DeepseekV4 FlashMLA does not support {num_heads} heads "
83 + "(FP8 decode kernel requires h_q in {64, 128})."
84 + )
85 + return 64 if num_heads <= 64 else 128
86 +
87 + def forward_mqa(
88 + self,
89 + q: torch.Tensor,
90 + kv: torch.Tensor,
91 + positions: torch.Tensor,
92 + output: torch.Tensor,
93 + ) -> None:
94 + assert output.shape == q.shape, (
95 + f"output buffer shape {output.shape} must match q shape {q.shape}"
96 + )
97 + assert output.dtype == q.dtype, (
98 + f"output buffer dtype {output.dtype} must match q dtype {q.dtype}"
99 + )
100 +
101 + # Get SWA and indexer metadata from forward context
102 + forward_context = get_forward_context()
103 + attn_metadata = forward_context.attn_metadata
104 +
105 + if attn_metadata is None:
106 + # Warmup dummy run: no real metadata. Reserve the same bf16
107 + # gather workspace _forward_prefill would; the dequantize / topk
108 + # / sparse_fwd kernels are skipped this step.
109 + swa_only = self.compress_ratio == 0
110 + N = (
111 + 0
112 + if swa_only
113 + else (self.max_model_len + self.compress_ratio - 1)
114 + // self.compress_ratio
115 + )
116 + M = N + self.window_size + self.max_num_batched_tokens
117 + if swa_only:
118 + top_k = 0
119 + else:
120 + assert self.topk_indices_buffer is not None
121 + top_k = self.topk_indices_buffer.shape[-1]
122 + combined_topk = round_up(
123 + top_k + self.window_size + self.max_image_tokens, 128
124 + )
125 + current_workspace_manager().get_simultaneous(
126 + ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
127 + ((self.max_num_batched_tokens, combined_topk), torch.int32),
128 + ((self.max_num_batched_tokens,), torch.int32),
129 + )
130 + output.zero_()
131 + return
132 +
133 + assert isinstance(attn_metadata, dict)
134 + # Compressed-cache metadata lives on the kv-source layer's prefix;
135 + # consumers share that cache and its block table.
136 + flashmla_metadata = cast(
137 + DeepseekV4FlashMLAMetadata | None,
138 + attn_metadata.get(self.compressed_cache_prefix)
139 + if self.compressed_cache_prefix is not None
140 + else None,
141 + )
142 + swa_metadata = cast(
143 + "DeepseekSparseSWAMetadata | None",
144 + attn_metadata.get(self.swa_cache_layer.prefix),
145 + )
146 + assert swa_metadata is not None
147 +
148 + swa_only = self.compress_ratio == 0
149 + # SWA-only layers (compress_ratio == 0) don't have their own KV cache
150 + # allocation; consumers read the kv source's cache, which is bound on
151 + # the source layer.
152 + self_kv_cache = None if swa_only else self._compressed_kv_cache()
153 + swa_kv_cache = self.swa_cache_layer.kv_cache
154 +
155 + # Split prefill and decode
156 + num_decodes = swa_metadata.num_decodes
157 + num_prefills = swa_metadata.num_prefills
158 + num_decode_tokens = swa_metadata.num_decode_tokens
159 +
160 + if num_prefills > 0:
161 + self._forward_prefill(
162 + q=q[num_decode_tokens:],
163 + positions=positions[num_decode_tokens:],
164 + compressed_k_cache=self_kv_cache,
165 + swa_k_cache=swa_kv_cache,
166 + output=output[num_decode_tokens:],
167 + attn_metadata=flashmla_metadata,
168 + swa_metadata=swa_metadata,
169 + )
170 + if num_decodes > 0:
171 + self._forward_decode(
172 + q=q[:num_decode_tokens],
173 + kv_cache=self_kv_cache,
174 + swa_metadata=swa_metadata,
175 + attn_metadata=flashmla_metadata,
176 + swa_only=swa_only,
177 + output=output[:num_decode_tokens],
178 + )
179 +
180 + def _forward_decode(
181 + self,
182 + q: torch.Tensor,
183 + kv_cache: torch.Tensor | None, # None for SWA-only layers
184 + swa_metadata: "DeepseekSparseSWAMetadata",
185 + attn_metadata: DeepseekV4FlashMLAMetadata | None,
186 + swa_only: bool,
187 + output: torch.Tensor,
188 + ) -> None:
189 + num_decodes = swa_metadata.num_decodes
190 + num_decode_tokens = swa_metadata.num_decode_tokens
191 +
192 + topk_indices = None
193 + topk_lens = None
194 + if not swa_only:
195 + # Local indices filled by the index-source layer's indexer.
196 + assert attn_metadata is not None
197 + assert swa_metadata.is_valid_token is not None
198 + assert self.topk_indices_buffer is not None
199 + block_size = attn_metadata.block_size // self.compress_ratio
200 + is_valid = swa_metadata.is_valid_token[:num_decode_tokens]
201 + global_indices, topk_lens = compute_global_topk_indices_and_lens(
202 + self.topk_indices_buffer[:num_decode_tokens],
203 + swa_metadata.token_to_req_indices,
204 + attn_metadata.block_table[:num_decodes],
205 + block_size,
206 + is_valid,
207 + )
208 + topk_indices = global_indices.view(num_decode_tokens, 1, -1)
209 +
210 + swa_indices = swa_metadata.decode_swa_indices
211 + swa_lens = swa_metadata.decode_swa_lens
212 +
213 + # We treat queries in the same seq as different queries
214 + # and later we only attend by generated indices.
215 + # q arrives pre-padded to self.padded_heads by the outer wrapper.
216 + q = q.unsqueeze(1)
217 +
218 + # Prepare SWA cache (num_blocks, swa_block_size, 1, head_bytes)
219 + # Use unsqueeze to preserve strides (handles padded blocks correctly)
220 + swa_cache = self.swa_cache_layer.kv_cache.unsqueeze(-2)
221 + # Reshape KV cache to (num_blocks, block_size, 1, head_bytes)
222 + if kv_cache is not None:
223 + kv_cache = kv_cache.unsqueeze(-2)
224 +
225 + # One FlashMLASchedMeta per layer type, shared across all same-type
226 + # layers within this decode step. The first forward call per type
227 + # triggers the in-kernel planner (allocating tile_scheduler_metadata
228 + # and num_splits via PyTorch's graph-aware allocator so CUDA graph
229 + # capture reuses the same addresses on replay); subsequent same-type
230 + # layers see have_initialized=True and skip the planner.
231 + tile_metadata = {
232 + 0: swa_metadata.tile_sched_swaonly,
233 + 1: swa_metadata.tile_sched_c1a,
234 + 2: swa_metadata.tile_sched_c2a,
235 + }[self.compress_ratio]
236 + assert tile_metadata is not None, (
237 + "swa_metadata missing tile_sched entry for "
238 + f"compress_ratio={self.compress_ratio}; "
239 + "DeepseekSparseSWAMetadataBuilder.build_tile_scheduler did not "
240 + "allocate one for this layer type."
241 + )
242 +
243 + out, _ = flash_mla_with_kvcache(
244 + q=q,
245 + k_cache=swa_cache,
246 + block_table=None,
247 + head_dim_v=512,
248 + tile_scheduler_metadata=tile_metadata,
249 + cache_seqlens=None,
250 + is_fp8_kvcache=True,
251 + indices=swa_indices,
252 + topk_length=swa_lens,
253 + softmax_scale=self.scale,
254 + attn_sink=self.attn_sink,
255 + extra_k_cache=kv_cache if not swa_only else None,
256 + extra_indices_in_kvcache=topk_indices,
257 + extra_topk_length=topk_lens,
258 + out=output.unsqueeze(1),
259 + )
260 +
261 + def _forward_prefill(
262 + self,
263 + q: torch.Tensor,
264 + positions: torch.Tensor,
265 + compressed_k_cache: torch.Tensor | None, # None for SWA-only layers
266 + swa_k_cache: torch.Tensor,
267 + output: torch.Tensor,
268 + attn_metadata: DeepseekV4FlashMLAMetadata | None,
269 + swa_metadata: "DeepseekSparseSWAMetadata",
270 + ) -> None:
271 + swa_only = self.compress_ratio == 0
272 +
273 + num_prefill_tokens = swa_metadata.num_prefill_tokens
274 + num_decodes = swa_metadata.num_decodes
275 + num_decode_tokens = swa_metadata.num_decode_tokens
276 +
277 + # Use pre-computed prefill metadata.
278 + seq_lens = swa_metadata.prefill_seq_lens
279 + gather_lens = swa_metadata.prefill_gather_lens
280 + assert seq_lens is not None
281 + assert gather_lens is not None
282 +
283 + # Derive prefill-local token offsets from the full query_start_loc_cpu.
284 + query_start_loc_cpu = swa_metadata.query_start_loc_cpu
285 + query_start_loc = swa_metadata.query_start_loc
286 + assert query_start_loc_cpu is not None
287 + assert query_start_loc is not None
288 + prefill_token_base = query_start_loc_cpu[num_decodes]
289 +
290 + # Local indices filled by the index source; SWA-only layers pass
291 + # top_k=0 and never read them.
292 + assert self.topk_indices_buffer is not None
293 + topk_indices = self.topk_indices_buffer[num_decode_tokens:]
294 + topk_indices = topk_indices[:num_prefill_tokens]
295 + top_k = 0 if swa_only else topk_indices.shape[-1]
296 + chunk_plan = swa_metadata.get_prefill_chunk_plan(
297 + compress_ratio=self.compress_ratio,
298 + prefill_chunk_size=self.PREFILL_CHUNK_SIZE,
299 + # v4.1: every cr>0 layer (including cr==1) gathers a full-length
300 + # compressed region; only cr==0 layers are SWA-only.
301 + has_compressed=not swa_only,
302 + )
303 + assert chunk_plan, "prefill chunk plan must be non-empty when num_prefills > 0"
304 + workspace_manager = current_workspace_manager()
305 + combined_topk = round_up(top_k + self.window_size + self.max_image_tokens, 128)
306 + for chunk_start, chunk_end, chunk_N, chunk_M in chunk_plan:
307 + chunk_size = chunk_end - chunk_start
308 + workspace = workspace_manager.get_simultaneous(
309 + ((chunk_size, chunk_M, q.shape[-1]), torch.bfloat16),
310 + ((self.max_num_batched_tokens, combined_topk), torch.int32),
311 + ((self.max_num_batched_tokens,), torch.int32),
312 + )
313 + kv, combined_indices_out, combined_lens_out = workspace
314 + if not swa_only:
315 + # Gather compressed KV
316 + assert attn_metadata is not None
317 + block_table = attn_metadata.block_table[num_decodes:]
318 + dequantize_and_gather_k_cache(
319 + kv[:chunk_size],
320 + compressed_k_cache,
321 + seq_lens=seq_lens[chunk_start:chunk_end] // self.compress_ratio,
322 + gather_lens=None,
323 + block_table=block_table[chunk_start:chunk_end],
324 + block_size=attn_metadata.block_size // self.compress_ratio,
325 + offset=0,
326 + )
327 +
328 + # Gather SWA KV
329 + swa_block_table = swa_metadata.block_table[num_decodes:]
330 + dequantize_and_gather_k_cache(
331 + kv[:chunk_size],
332 + swa_k_cache,
333 + seq_lens=seq_lens[chunk_start:chunk_end],
334 + gather_lens=gather_lens[chunk_start:chunk_end],
335 + block_table=swa_block_table[chunk_start:chunk_end],
336 + block_size=swa_metadata.block_size,
337 + offset=chunk_N,
338 + )
339 +
340 + # Combine the topk indices and SWA indices for gathered KV cache
341 + query_start = (
342 + query_start_loc_cpu[num_decodes + chunk_start] - prefill_token_base
343 + )
344 + query_end = (
345 + query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base
346 + )
347 + combined_indices_out = combined_indices_out[: query_end - query_start]
348 + combined_lens_out = combined_lens_out[: query_end - query_start]
349 +
350 + combined_indices, combined_lens = combine_topk_swa_indices(
351 + topk_indices[query_start:query_end],
352 + query_start_loc[
353 + num_decodes + chunk_start : num_decodes + chunk_end + 1
354 + ],
355 + seq_lens[chunk_start:chunk_end],
356 + gather_lens[chunk_start:chunk_end],
357 + self.window_size,
358 + self.compress_ratio,
359 + top_k,
360 + chunk_M,
361 + chunk_N,
362 + out=(combined_indices_out, combined_lens_out),
363 + left_visible=(
364 + swa_metadata.prefill_left_visible[
365 + num_decode_tokens + query_start : num_decode_tokens + query_end
366 + ]
367 + if swa_metadata.prefill_left_visible is not None
368 + else None
369 + ),
370 + right_visible=(
371 + swa_metadata.prefill_right_visible[
372 + num_decode_tokens + query_start : num_decode_tokens + query_end
373 + ]
374 + if swa_metadata.prefill_right_visible is not None
375 + else None
376 + ),
377 + max_image_tokens=self.max_image_tokens,
378 + )
379 + flash_mla_sparse_fwd(
380 + q=q[query_start:query_end],
381 + kv=kv.view(-1, 1, q.shape[-1]),
382 + indices=combined_indices.unsqueeze(1),
383 + sm_scale=self.scale,
384 + attn_sink=self.attn_sink,
385 + topk_length=combined_lens,
386 + out=output[query_start:query_end],
387 + )
vllm/models/deepseek_v4_1/nvidia/model.py new
+1111
@@ -0,0 +1,1111 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import typing
4 +from collections.abc import Callable, Iterable
5 +from itertools import islice
6 +
7 +import regex as re
8 +import torch
9 +import torch.nn as nn
10 +
11 +import vllm.envs as envs
12 +from vllm.config import VllmConfig
13 +from vllm.config.kernel import MEGA_MOE_BACKENDS
14 +from vllm.distributed import (
15 + get_pp_group,
16 + get_tensor_model_parallel_rank,
17 + get_tensor_model_parallel_world_size,
18 +)
19 +from vllm.forward_context import get_forward_context, is_forward_context_available
20 +from vllm.logger import init_logger
21 +from vllm.model_executor.kernels.mhc.tilelang import (
22 + mhc_post_tilelang,
23 + mhc_pre_delayed_tilelang,
24 +)
25 +from vllm.model_executor.kernels.mhc.triton import hc_collapse_triton
26 +from vllm.model_executor.layers.fused_moe import (
27 + fused_moe_make_expert_params_mapping,
28 +)
29 +from vllm.model_executor.layers.layernorm import RMSNorm
30 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
31 +from vllm.model_executor.layers.quantization import QuantizationConfig
32 +from vllm.model_executor.layers.vocab_parallel_embedding import (
33 + ParallelLMHead,
34 + VocabParallelEmbedding,
35 +)
36 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
37 +from vllm.model_executor.models.interfaces import (
38 + EagleModelMixin,
39 + MixtureOfExperts,
40 + SupportsEagle3,
41 + SupportsLoRA,
42 + SupportsPP,
43 +)
44 +from vllm.model_executor.models.utils import (
45 + AutoWeightsLoader,
46 + PPMissingLayer,
47 + WeightsMapper,
48 + extract_layer_index,
49 + is_pp_missing_parameter,
50 + make_layers,
51 + maybe_prefix,
52 +)
53 +from vllm.models.common.ops.sequence_parallel import (
54 + sp_all_gather,
55 + sp_padding_mask,
56 + sp_reduce_scatter,
57 + sp_shard,
58 +)
59 +from vllm.models.deepseek_v4.nvidia.model import (
60 + DeepseekV4MoE as DeepseekV4MoEBase,
61 +)
62 +from vllm.models.deepseek_v4.nvidia.model import (
63 + make_deepseek_v4_expert_params_mapping,
64 +)
65 +from vllm.models.deepseek_v4_1.attention import DeepseekV4Attention
66 +from vllm.models.deepseek_v4_1.nvidia.flashinfer_sparse import (
67 + DeepseekV4FlashInferMLAAttention,
68 + DeepseekV4FlashInferSM120Attention,
69 +)
70 +from vllm.models.deepseek_v4_1.nvidia.flashmla import DeepseekV4FlashMLAAttention
71 +from vllm.platforms import current_platform
72 +from vllm.sequence import IntermediateTensors
73 +from vllm.utils.math_utils import cdiv
74 +from vllm.v1.attention.backends.registry import AttentionBackendEnum
75 +from vllm.v1.worker.ubatching import dbo_current_ubatch_id
76 +
77 +from ..common.engram import Engram, EngramLayout, NgramHashState
78 +from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask
79 +
80 +if typing.TYPE_CHECKING:
81 + from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
82 +
83 +logger = init_logger(__name__)
84 +
85 +
86 +class DeepseekV4MoE(DeepseekV4MoEBase):
87 + def __init__(
88 + self,
89 + vllm_config: VllmConfig,
90 + prefix: str = "",
91 + use_sequence_parallel: bool = False,
92 + ):
93 + config = vllm_config.model_config.hf_config
94 + n_routed_experts = config.n_routed_experts
95 + n_activated_experts = config.num_experts_per_tok
96 + if extract_layer_index(prefix) >= config.num_hidden_layers:
97 + n_routed_experts = (
98 + getattr(config, "dspark_n_routed_experts", 0) or n_routed_experts
99 + )
100 + n_activated_experts = (
101 + getattr(config, "dspark_num_experts_per_tok", 0) or n_activated_experts
102 + )
103 + super().__init__(
104 + vllm_config,
105 + prefix=prefix,
106 + use_sequence_parallel=use_sequence_parallel,
107 + n_routed_experts=n_routed_experts,
108 + n_activated_experts=n_activated_experts,
109 + num_hash_layers=0,
110 + image_sentinel_lo=IMAGE_SENTINEL_BASE_ID,
111 + )
112 +
113 +
114 +def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]:
115 + """Pick the CUDA sparse-MLA attention class for the configured backend.
116 +
117 + The generic CUDA backend selector does not instantiate DSv4 layers directly,
118 + so map generic sparse-MLA choices to the DSv4-specialized attention class.
119 + Without an explicit backend, SM12 defaults to FlashInfer while the other
120 + CUDA arches keep the FlashMLA path.
121 + """
122 + backend = vllm_config.attention_config.backend
123 + device_capability = current_platform.get_device_capability()
124 + if backend in (
125 + AttentionBackendEnum.FLASHINFER_MLA_SPARSE,
126 + AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM120,
127 + ):
128 + raise ValueError(
129 + f"{backend.name} is not a DeepSeek V4.1 attention backend. "
130 + "Use FLASHINFER_MLA_SPARSE_DSV41 for DeepSeek V4.1 FlashInfer "
131 + "sparse MLA."
132 + )
133 + if backend in (
134 + AttentionBackendEnum.FLASHINFER_MLA_SPARSE_DSV4,
135 + AttentionBackendEnum.FLASHINFER_MLA_SPARSE_DSV41,
136 + ):
137 + if device_capability is not None and device_capability.major == 12:
138 + return DeepseekV4FlashInferSM120Attention
139 + return DeepseekV4FlashInferMLAAttention
140 + if backend in (
141 + AttentionBackendEnum.FLASHMLA_SPARSE,
142 + AttentionBackendEnum.FLASHMLA_SPARSE_DSV4,
143 + AttentionBackendEnum.FLASHMLA_SPARSE_DSV41,
144 + ):
145 + return DeepseekV4FlashMLAAttention
146 +
147 + if device_capability is not None and device_capability.major == 12:
148 + return DeepseekV4FlashInferSM120Attention
149 + return DeepseekV4FlashMLAAttention
150 +
151 +
152 +def _use_sequence_parallel(vllm_config: VllmConfig) -> bool:
153 + parallel_config = vllm_config.parallel_config
154 + use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS
155 + return (
156 + parallel_config.pipeline_parallel_size == 1
157 + and parallel_config.enable_expert_parallel
158 + and parallel_config.tensor_parallel_size > 1
159 + and (use_mega_moe or parallel_config.data_parallel_size > 1)
160 + )
161 +
162 +
163 +class DeepseekV4DecoderLayer(nn.Module):
164 + def __init__(
165 + self,
166 + vllm_config,
167 + prefix,
168 + topk_indices_buffer: torch.Tensor | None = None,
169 + aux_stream_list: list[torch.cuda.Stream] | None = None,
170 + candidate_block_buffer: torch.Tensor | None = None,
171 + engram_layout: EngramLayout | None = None,
172 + ):
173 + super().__init__()
174 +
175 + config = vllm_config.model_config.hf_config
176 + self.hidden_size = config.hidden_size
177 + self.use_sequence_parallel = _use_sequence_parallel(vllm_config)
178 +
179 + self.engram: Engram | None = None
180 + if engram_layout is not None:
181 + layer_id = extract_layer_index(prefix)
182 + if layer_id in engram_layout.layer_ids:
183 + self.engram = Engram(
184 + config,
185 + vllm_config.quant_config,
186 + engram_layout,
187 + engram_layout.layer_ids.index(layer_id),
188 + use_sequence_parallel=self.use_sequence_parallel,
189 + prefix=f"{prefix}.engram",
190 + )
191 +
192 + self.rms_norm_eps = config.rms_norm_eps
193 + self.attn = _select_dsv4_attn_cls(vllm_config)(
194 + vllm_config,
195 + prefix=f"{prefix}.attn",
196 + topk_indices_buffer=topk_indices_buffer,
197 + aux_stream_list=aux_stream_list,
198 + candidate_block_buffer=candidate_block_buffer,
199 + )
200 + if self.use_sequence_parallel:
201 + self.attn.wo_b.reduce_results = False
202 + self.ffn = DeepseekV4MoE(
203 + vllm_config,
204 + prefix=f"{prefix}.ffn",
205 + use_sequence_parallel=self.use_sequence_parallel,
206 + )
207 +
208 + self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps)
209 + self.ffn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps)
210 + self.hc_mult = config.hc_mult
211 + self.hc_sinkhorn_iters = config.hc_sinkhorn_iters
212 + self.hc_eps = config.hc_eps
213 + self.hc_post_alpha = 2.0
214 + if vllm_config.kernel_config.enable_jit_warmup and current_platform.is_cuda():
215 + from vllm.model_executor.kernels.mhc.warmup import MHC_PRE_NORM_KERNEL
216 +
217 + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
218 + if self.use_sequence_parallel:
219 + tp_size = vllm_config.parallel_config.tensor_parallel_size
220 + max_tokens = (max_tokens + tp_size - 1) // tp_size
221 + for input_size, use_pre_mix in (
222 + (self.hidden_size, False),
223 + (self.hc_mult * self.hidden_size, False),
224 + (self.hc_mult * self.hidden_size, True),
225 + ):
226 + MHC_PRE_NORM_KERNEL.register_warmup(
227 + max_tokens=max_tokens,
228 + hidden_size=self.hidden_size,
229 + rms_eps=self.rms_norm_eps,
230 + hc_pre_eps=self.hc_eps,
231 + hc_sinkhorn_eps=self.hc_eps,
232 + hc_post_mult_value=self.hc_post_alpha,
233 + sinkhorn_repeat=self.hc_sinkhorn_iters,
234 + norm_eps=self.rms_norm_eps,
235 + hc_mult=self.hc_mult,
236 + use_pre_mix_in=use_pre_mix,
237 + rms_numel=input_size,
238 + )
239 + mix_hc = (2 + self.hc_mult) * self.hc_mult
240 + hc_dim = self.hc_mult * self.hidden_size
241 + self.hc_attn_fn = nn.Parameter(
242 + torch.empty(
243 + (mix_hc, hc_dim),
244 + dtype=torch.float32,
245 + ),
246 + requires_grad=False,
247 + )
248 + self.hc_attn_fn_broadcast: torch.Tensor | None = None
249 + self.hc_ffn_fn = nn.Parameter(
250 + torch.empty(
251 + (mix_hc, hc_dim),
252 + dtype=torch.float32,
253 + ),
254 + requires_grad=False,
255 + )
256 + self.hc_attn_base = nn.Parameter(
257 + torch.empty(
258 + mix_hc,
259 + dtype=torch.float32,
260 + ),
261 + requires_grad=False,
262 + )
263 + self.hc_ffn_base = nn.Parameter(
264 + torch.empty(
265 + mix_hc,
266 + dtype=torch.float32,
267 + ),
268 + requires_grad=False,
269 + )
270 + self.hc_attn_scale = nn.Parameter(
271 + torch.empty(
272 + 3,
273 + dtype=torch.float32,
274 + ),
275 + requires_grad=False,
276 + )
277 + self.hc_ffn_scale = nn.Parameter(
278 + torch.empty(
279 + 3,
280 + dtype=torch.float32,
281 + ),
282 + requires_grad=False,
283 + )
284 +
285 + def forward(
286 + self,
287 + x: torch.Tensor,
288 + positions: torch.Tensor,
289 + input_ids: torch.Tensor | None,
290 + pre_mix: torch.Tensor | None = None,
291 + post_mix: torch.Tensor | None = None,
292 + res_mix: torch.Tensor | None = None,
293 + residual: torch.Tensor | None = None,
294 + engram_hashes: torch.Tensor | None = None,
295 + engram_mask: torch.Tensor | None = None,
296 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
297 + # The reference collapses each sublayer's input with the *previous*
298 + # sublayer's pre-mix: attention uses the pre-mix carried in (identity
299 + # for the first layer), the FFN uses this layer's attention pre-mix.
300 + if residual is None:
301 + if x.dim() == 2:
302 + # First layer: the stream is the embedding broadcast to hc
303 + # copies and the identity pre-mix selects copy 0.
304 + assert self.hc_attn_fn_broadcast is not None
305 + residual = x.unsqueeze(1).expand(-1, self.hc_mult, -1).contiguous()
306 + post_mix, res_mix, x, attn_pre = mhc_pre_delayed_tilelang(
307 + residual,
308 + self.hc_attn_fn_broadcast,
309 + self.hc_attn_scale,
310 + self.hc_attn_base,
311 + self.rms_norm_eps,
312 + self.hc_eps,
313 + self.hc_eps,
314 + self.hc_post_alpha,
315 + self.hc_sinkhorn_iters,
316 + x=x,
317 + norm_weight=self.attn_norm.weight,
318 + norm_eps=self.attn_norm.variance_epsilon,
319 + )
320 + else:
321 + residual = x
322 + post_mix, res_mix, x, attn_pre = mhc_pre_delayed_tilelang(
323 + residual,
324 + self.hc_attn_fn,
325 + self.hc_attn_scale,
326 + self.hc_attn_base,
327 + self.rms_norm_eps,
328 + self.hc_eps,
329 + self.hc_eps,
330 + self.hc_post_alpha,
331 + self.hc_sinkhorn_iters,
332 + pre_mix=pre_mix,
333 + norm_weight=self.attn_norm.weight,
334 + norm_eps=self.attn_norm.variance_epsilon,
335 + )
336 + else:
337 + residual = mhc_post_tilelang(x, residual, post_mix, res_mix)
338 + if self.engram is not None and engram_hashes is not None:
339 + # Engram injection happens between the previous sublayer's
340 + # post and this block's pre, on the full hc stream, so the
341 + # mix coefficients see the injected stream.
342 + residual = self.engram(
343 + residual,
344 + engram_hashes[:, self.engram.layer_hash_index],
345 + engram_mask,
346 + )
347 + post_mix, res_mix, x, attn_pre = mhc_pre_delayed_tilelang(
348 + residual,
349 + self.hc_attn_fn,
350 + self.hc_attn_scale,
351 + self.hc_attn_base,
352 + self.rms_norm_eps,
353 + self.hc_eps,
354 + self.hc_eps,
355 + self.hc_post_alpha,
356 + self.hc_sinkhorn_iters,
357 + pre_mix=pre_mix,
358 + norm_weight=self.attn_norm.weight,
359 + norm_eps=self.attn_norm.variance_epsilon,
360 + )
361 +
362 + if self.use_sequence_parallel:
363 + x = sp_all_gather(x)[: positions.shape[0]]
364 +
365 + x = self.attn(positions, x, None)
366 + if self.use_sequence_parallel:
367 + x = sp_reduce_scatter(x)
368 +
369 + residual = mhc_post_tilelang(x, residual, post_mix, res_mix)
370 + post_mix, res_mix, x, ffn_pre = mhc_pre_delayed_tilelang(
371 + residual,
372 + self.hc_ffn_fn,
373 + self.hc_ffn_scale,
374 + self.hc_ffn_base,
375 + self.rms_norm_eps,
376 + self.hc_eps,
377 + self.hc_eps,
378 + self.hc_post_alpha,
379 + self.hc_sinkhorn_iters,
380 + pre_mix=attn_pre,
381 + norm_weight=self.ffn_norm.weight,
382 + norm_eps=self.ffn_norm.variance_epsilon,
383 + )
384 + x = self.ffn(x, input_ids)
385 + return x, residual, post_mix, res_mix, ffn_pre
386 +
387 +
388 +class DeepseekV4Model(nn.Module, EagleModelMixin):
389 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
390 + super().__init__()
391 +
392 + config = vllm_config.model_config.hf_config
393 + quant_config = vllm_config.quant_config
394 + self.config = config
395 + self.quant_config = quant_config
396 + self.parallel_config = vllm_config.parallel_config
397 + self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS
398 + self.use_sequence_parallel = _use_sequence_parallel(vllm_config)
399 + if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel:
400 + raise NotImplementedError(
401 + "DeepSeek V4 MegaMoE currently requires expert parallel. "
402 + "Enable it with --enable-expert-parallel, or pick a different "
403 + "moe backend."
404 + )
405 + self.vocab_size = config.vocab_size
406 + self.hc_eps = config.hc_eps
407 + self.hc_mult = config.hc_mult
408 + self.hc_dim = self.hc_mult * config.hidden_size
409 + self.rms_norm_eps = config.rms_norm_eps
410 +
411 + # Three aux streams: one per non-default input GEMM in
412 + # DeepseekV4Attention._run_parallel_input_projections
413 + # (compressor kv_score, indexer.weights_proj). fused_wqa_wkv stays on
414 + # the default stream.
415 + aux_stream_list = [torch.cuda.Stream() for _ in range(3)]
416 +
417 + # Reserved topk indices buffer for all Indexer layers to reuse.
418 + self.topk_indices_buffer = torch.empty(
419 + vllm_config.scheduler_config.max_num_batched_tokens,
420 + config.index_topk,
421 + dtype=torch.int32,
422 + )
423 +
424 + # Two-level candidate filtering: the indexer at
425 + # candidate_source_layer_id publishes the top candidate blocks of
426 + # compressed positions here; later ratio-1 indexers (24/28/32/36)
427 + # mask their scores with it.
428 + candidate_source_layer = getattr(config, "candidate_source_layer_id", -1)
429 + candidate_topk_blocks = getattr(config, "candidate_topk_blocks", 0)
430 + if candidate_source_layer >= 0 and candidate_topk_blocks > 0:
431 + self.candidate_block_buffer = torch.empty(
432 + vllm_config.scheduler_config.max_num_batched_tokens,
433 + candidate_topk_blocks,
434 + dtype=torch.int32,
435 + )
436 + else:
437 + self.candidate_block_buffer = None
438 +
439 + if get_pp_group().is_first_rank:
440 + self.embed_tokens = VocabParallelEmbedding(
441 + config.vocab_size,
442 + config.hidden_size,
443 + quant_config=quant_config,
444 + prefix=f"{prefix}.embed_tokens",
445 + )
446 + else:
447 + self.embed_tokens = PPMissingLayer()
448 +
449 + self.engram_layout = EngramLayout.from_config(config)
450 +
451 + self.start_layer, self.end_layer, self.layers = make_layers(
452 + config.num_hidden_layers,
453 + lambda prefix: DeepseekV4DecoderLayer(
454 + vllm_config,
455 + prefix=prefix,
456 + topk_indices_buffer=self.topk_indices_buffer,
457 + aux_stream_list=aux_stream_list,
458 + candidate_block_buffer=self.candidate_block_buffer,
459 + engram_layout=self.engram_layout,
460 + ),
461 + prefix=f"{prefix}.layers",
462 + )
463 +
464 + # The n-gram hash needs a slot-keyed rolling store of compressed ids
465 + # (chunked prefill / decode lookback); key it off the first local
466 + # layer's sliding-window KV cache. Only PP ranks owning an engram
467 + # layer need it.
468 + self.engram_hash: NgramHashState | None = None
469 + self.engram_swa_prefix: str | None = None
470 + if self.engram_layout is not None:
471 + local_engram = any(
472 + isinstance(layer, DeepseekV4DecoderLayer) and layer.engram is not None
473 + for layer in islice(self.layers, self.start_layer, self.end_layer)
474 + )
475 + if local_engram:
476 + first_layer = next(
477 + iter(islice(self.layers, self.start_layer, self.end_layer))
478 + )
479 + swa_cache_module = first_layer.attn.swa_cache_layer
480 + self.engram_hash = NgramHashState(
481 + vllm_config, self.engram_layout, swa_cache_module
482 + )
483 + self.engram_swa_prefix = swa_cache_module.prefix
484 +
485 + if get_pp_group().is_last_rank:
486 + self.norm = RMSNorm(config.hidden_size, self.rms_norm_eps)
487 + else:
488 + self.norm = PPMissingLayer()
489 +
490 + spec_config = vllm_config.speculative_config
491 + needs_mtp_hidden_states = spec_config is not None and (
492 + spec_config.use_eagle() or spec_config.uses_draft_model()
493 + )
494 + if get_pp_group().is_last_rank and needs_mtp_hidden_states:
495 + self._mtp_hidden_buffer = torch.empty(
496 + vllm_config.scheduler_config.max_num_batched_tokens,
497 + self.hc_dim,
498 + dtype=vllm_config.model_config.dtype,
499 + )
500 + else:
501 + self._mtp_hidden_buffer = None
502 +
503 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
504 + return self.embed_tokens(input_ids)
505 +
506 + def make_empty_intermediate_tensors(
507 + self,
508 + batch_size: int,
509 + dtype: torch.dtype,
510 + device: torch.device,
511 + ) -> IntermediateTensors:
512 + # PP intermediate tensors carry the multi-stream hidden_states
513 + # of shape (num_tokens, hc_mult, hidden_size) — V4 expands the
514 + # token embedding to hc_mult streams before the first decoder
515 + # layer and keeps that shape until the final hc collapse — plus the
516 + # (num_tokens, hc_mult) pre-mix the next rank's first layer needs
517 + # for its attention collapse.
518 + return IntermediateTensors(
519 + {
520 + "hidden_states": torch.zeros(
521 + (batch_size, self.hc_mult, self.config.hidden_size),
522 + dtype=dtype,
523 + device=device,
524 + ),
525 + "pre_mix": torch.zeros(
526 + (batch_size, self.hc_mult),
527 + dtype=torch.float32,
528 + device=device,
529 + ),
530 + }
531 + )
532 +
533 + def forward(
534 + self,
535 + input_ids: torch.Tensor,
536 + positions: torch.Tensor,
537 + intermediate_tensors: IntermediateTensors | None,
538 + inputs_embeds: torch.Tensor | None = None,
539 + lookback_token_ids: torch.Tensor | None = None,
540 + ) -> torch.Tensor | IntermediateTensors:
541 + if get_pp_group().is_first_rank:
542 + if inputs_embeds is not None:
543 + hidden_states = inputs_embeds
544 + else:
545 + hidden_states = self.embed_input_ids(input_ids)
546 + else:
547 + assert intermediate_tensors is not None
548 + hidden_states = intermediate_tensors["hidden_states"]
549 +
550 + if self.use_mega_moe:
551 + input_ids = input_ids.to(torch.int64)
552 +
553 + # Engram n-gram hashes for the whole (flattened) batch, computed once
554 + # on the full token stream — before any sequence-parallel sharding —
555 + # and consumed by the engram layers (1 and 14) below. Skipped on
556 + # profile runs (KV cache unbound).
557 + engram_hashes: torch.Tensor | None = None
558 + engram_mask: torch.Tensor | None = None
559 + if (
560 + self.engram_hash is not None
561 + and input_ids is not None
562 + and is_forward_context_available()
563 + ):
564 + attn_metadata = get_forward_context().attn_metadata
565 + if isinstance(attn_metadata, list):
566 + attn_metadata = attn_metadata[dbo_current_ubatch_id()]
567 + if isinstance(attn_metadata, dict) and self.engram_hash.ensure_cache():
568 + assert self.engram_swa_prefix is not None
569 + swa_metadata = typing.cast(
570 + "DeepseekSparseSWAMetadata", attn_metadata[self.engram_swa_prefix]
571 + )
572 + # Image-span tokens are dead: they break n-grams (hash op
573 + # takes True=dead) and their gate is zeroed (Engram.forward
574 + # takes True=keep).
575 + image_mask = image_sentinel_mask(input_ids)
576 + engram_mask = ~image_mask
577 + if lookback_token_ids is None:
578 + if not self.engram_hash.use_slot_cache:
579 + raise NotImplementedError(
580 + "engram needs `lookback_token_ids` from the model "
581 + "runner (the DBO/ubatch wrapper drops model kwargs)"
582 + )
583 + num_reqs = swa_metadata.num_decodes + swa_metadata.num_prefills
584 + lookback_token_ids = input_ids.new_full(
585 + (num_reqs, self.engram_hash.lookback_depth), -1
586 + )
587 + engram_hashes = self.engram_hash(
588 + input_ids,
589 + positions,
590 + swa_metadata.query_start_loc,
591 + image_mask,
592 + lookback_token_ids,
593 + image_sentinel_mask(lookback_token_ids),
594 + swa_metadata.slot_mapping,
595 + swa_metadata.block_table,
596 + )
597 + # Gather all Engram rows before entering the decoder layers.
598 + for layer in islice(self.layers, self.start_layer, self.end_layer):
599 + engram = getattr(layer, "engram", None)
600 + if engram is not None:
601 + engram.prepare_embeddings(
602 + engram_hashes[:, engram.layer_hash_index]
603 + )
604 +
605 + full_num_tokens = positions.shape[0]
606 + if self.use_sequence_parallel:
607 + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available():
608 + forward_context = get_forward_context()
609 + forward_context.is_padding = sp_padding_mask(
610 + forward_context.is_padding, hidden_states
611 + )
612 + hidden_states = sp_shard(hidden_states)
613 + input_ids = sp_shard(input_ids)
614 +
615 + residual, post_mix, res_mix = None, None, None
616 + pre_mix: torch.Tensor | None = None
617 + if not get_pp_group().is_first_rank:
618 + assert intermediate_tensors is not None
619 + pre_mix = intermediate_tensors["pre_mix"]
620 + aux_hidden_states: list[torch.Tensor] = []
621 + final_aux_recon: torch.Tensor | None = None # avoid duplicate mhc_post call
622 + for idx, layer in enumerate(
623 + islice(self.layers, self.start_layer, self.end_layer),
624 + start=self.start_layer,
625 + ):
626 + hidden_states, residual, post_mix, res_mix, pre_mix = layer(
627 + hidden_states,
628 + positions,
629 + input_ids,
630 + pre_mix,
631 + post_mix,
632 + res_mix,
633 + residual,
634 + engram_hashes,
635 + engram_mask,
636 + )
637 + if idx + 1 in self.aux_hidden_state_layers:
638 + # Reconstruct the aux hidden state for draft models
639 + aux_recon = mhc_post_tilelang(
640 + hidden_states, residual, post_mix, res_mix
641 + )
642 + aux_hidden_state = aux_recon.mean(dim=1)
643 + if self.use_sequence_parallel:
644 + aux_hidden_state = sp_all_gather(aux_hidden_state)[:full_num_tokens]
645 + aux_hidden_states.append(aux_hidden_state)
646 + final_aux_recon = aux_recon
647 + if layer is not None:
648 + # Reuse if the last layer was captured as an aux hidden state
649 + if self.end_layer in self.aux_hidden_state_layers:
650 + hidden_states = final_aux_recon
651 + else:
652 + hidden_states = mhc_post_tilelang(
653 + hidden_states, residual, post_mix, res_mix
654 + )
655 +
656 + if not get_pp_group().is_last_rank:
657 + return IntermediateTensors(
658 + {"hidden_states": hidden_states, "pre_mix": pre_mix}
659 + )
660 +
661 + # MTP needs full HC states; otherwise collapse and normalize locally
662 + # before gathering to reduce communication.
663 + if self._mtp_hidden_buffer is not None:
664 + if self.use_sequence_parallel:
665 + hidden_states = sp_all_gather(hidden_states)[:full_num_tokens]
666 + pre_mix = sp_all_gather(pre_mix)[:full_num_tokens]
667 + num_tokens = hidden_states.shape[0]
668 + self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1))
669 +
670 + # Collapse the hc copies with the pre-mix from the last layer's FFN
671 + # mixes — the mix the reference applies via
672 + # ``last_layer.hc_pre(h, pre_mix)`` (v4.1 has no learned hc_head).
673 + assert pre_mix is not None
674 + hidden_states = hc_collapse_triton(hidden_states, pre_mix)
675 + hidden_states = self.norm(hidden_states)
676 + if self.use_sequence_parallel and self._mtp_hidden_buffer is None:
677 + # Without MTP, gather only the collapsed and normalized hidden states.
678 + hidden_states = sp_all_gather(hidden_states)[:full_num_tokens]
679 + if len(aux_hidden_states) > 0:
680 + return hidden_states, aux_hidden_states
681 + return hidden_states
682 +
683 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
684 + stacked_params_mapping = [
685 + # (param_name, shard_name, shard_id)
686 + ("gate_up_proj", "w1", 0),
687 + ("gate_up_proj", "w3", 1),
688 + ("attn.fused_wqa_wkv", "attn.wq_a", 0),
689 + ("attn.fused_wqa_wkv", "attn.wkv", 1),
690 + ("compressor.fused_wkv_wgate", "compressor.wkv", 0),
691 + ("compressor.fused_wkv_wgate", "compressor.wgate", 1),
692 + ]
693 + params_dict = dict(self.named_parameters())
694 + loaded_params: set[str] = set()
695 +
696 + # TP for attention
697 + tp_size = get_tensor_model_parallel_world_size()
698 + tp_rank = get_tensor_model_parallel_rank()
699 + n_head = self.config.num_attention_heads
700 + n_local_head = n_head // tp_size
701 + head_rank_start = n_local_head * tp_rank
702 + head_rank_end = n_local_head * (tp_rank + 1)
703 +
704 + # Pre-compute expert mapping ONCE.
705 + expert_mapping = self.get_expert_mapping()
706 +
707 + # Block-FP8 shared experts: pad the intermediate up to the TP-uniform
708 + # block count so the standard loaders below slice it evenly (trailing
709 + # ranks land on the zero pad). SP / unquantized ones need no padding.
710 + pad_shared_expert = (
711 + getattr(self.quant_config, "weight_block_size", None) is not None
712 + and not self.use_sequence_parallel
713 + )
714 +
715 + for name, loaded_weight in weights:
716 + if name.startswith(("vision.", "aligner.", "image_")):
717 + # Vision weights are loaded by the outer multimodal wrapper.
718 + logger.warning_once("Skipping non-text weight: %s", name)
719 + continue
720 + if pad_shared_expert and ".shared_experts." in name:
721 + loaded_weight = self._pad_shared_expert_weight(
722 + self.quant_config, name, loaded_weight
723 + )
724 + for param_name, weight_name, shard_id in stacked_params_mapping:
725 + # Skip non-stacked layers and experts (experts handled below).
726 + if ".experts." in name:
727 + continue
728 + if weight_name not in name:
729 + continue
730 + name = name.replace(weight_name, param_name)
731 +
732 + if is_pp_missing_parameter(name, self):
733 + break
734 + if name not in params_dict:
735 + head, _, leaf = name.rpartition(".")
736 + suffixed = f"{head}.base_layer.{leaf}"
737 + if suffixed in params_dict:
738 + name = suffixed
739 + param = params_dict[name]
740 + weight_loader = param.weight_loader
741 + weight_loader(param, loaded_weight, shard_id)
742 + loaded_params.add(name)
743 + break
744 + else:
745 + if ".experts." in name:
746 + # E8M0 scales are stored as float8_e8m0fnu in
747 + # checkpoints but the MoE param is uint8. copy_()
748 + # would do a numeric conversion (e.g. 2^-7 → 0),
749 + # destroying the raw exponent bytes.
750 + if (
751 + "weight_scale" in name
752 + and loaded_weight.dtype == torch.float8_e8m0fnu
753 + ):
754 + loaded_weight = loaded_weight.view(torch.uint8)
755 + for mapping in expert_mapping:
756 + param_name, weight_name, expert_id, expert_shard_id = mapping
757 + if weight_name not in name:
758 + continue
759 + name_mapped = name.replace(weight_name, param_name)
760 + if is_pp_missing_parameter(name_mapped, self):
761 + continue
762 + param = params_dict[name_mapped]
763 + # We should ask the weight loader to return success or not
764 + # here since otherwise we may skip experts with other
765 + # available replicas.
766 + weight_loader = typing.cast(
767 + Callable[..., bool], param.weight_loader
768 + )
769 + success = weight_loader(
770 + param,
771 + loaded_weight,
772 + name_mapped,
773 + shard_id=expert_shard_id,
774 + expert_id=expert_id,
775 + return_success=True,
776 + )
777 + if success:
778 + name = name_mapped
779 + break
780 + loaded_params.add(name_mapped)
781 + continue
782 + elif "attn_sink" in name:
783 + if is_pp_missing_parameter(name, self):
784 + continue
785 + narrow_weight = loaded_weight[head_rank_start:head_rank_end]
786 + n = narrow_weight.shape[0]
787 + params_dict[name][:n].copy_(narrow_weight)
788 + loaded_params.add(name)
789 + continue
790 + else:
791 + if is_pp_missing_parameter(name, self):
792 + continue
793 + # Non-LoRA params on a LoRA-wrapped module live at
794 + # ``<head>.base_layer.<leaf>``; the checkpoint is plain.
795 + if name not in params_dict:
796 + head, _, leaf = name.rpartition(".")
797 + suffixed = f"{head}.base_layer.{leaf}"
798 + if suffixed in params_dict:
799 + name = suffixed
800 + param = params_dict[name]
801 + weight_loader = getattr(
802 + param, "weight_loader", default_weight_loader
803 + )
804 + weight_loader(param, loaded_weight)
805 + loaded_params.add(name)
806 + continue
807 +
808 + return loaded_params
809 +
810 + @staticmethod
811 + def _pad_shared_expert_weight(
812 + quant_config: QuantizationConfig | None,
813 + name: str,
814 + loaded_weight: torch.Tensor,
815 + ) -> torch.Tensor:
816 + """Zero-pad a block-FP8 shared-expert weight/scale on its intermediate
817 + axis so the standard TP loaders split it into even, block-aligned shards
818 + (trailing ranks get the zero pad). gate (w1)/up (w3) [I, H] pad dim 0;
819 + down (w2 -> down_proj) [H, I] pads dim 1.
820 + """
821 + block_size = getattr(quant_config, "weight_block_size", None)
822 + assert block_size is not None
823 + # Round the intermediate axis up to a whole number of TP shards. The axis
824 + # is in elements for weights (step = block) and in blocks for scales.
825 + step = (
826 + 1 if name.endswith(("weight_scale_inv", "weight_scale")) else block_size[0]
827 + )
828 + dim = 1 if ".down_proj." in name else 0
829 + mult = get_tensor_model_parallel_world_size() * step
830 + pad = cdiv(loaded_weight.shape[dim], mult) * mult - loaded_weight.shape[dim]
831 + if pad == 0:
832 + return loaded_weight
833 + pad_shape = list(loaded_weight.shape)
834 + pad_shape[dim] = pad
835 + return torch.cat([loaded_weight, loaded_weight.new_zeros(pad_shape)], dim=dim)
836 +
837 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
838 + first_layer = next(iter(islice(self.layers, self.start_layer, self.end_layer)))
839 + if first_layer.ffn.use_mega_moe:
840 + return make_deepseek_v4_expert_params_mapping(self.config.n_routed_experts)
841 + # Params for weights, fp8 weight scales, fp8 activation scales
842 + # (param_name, weight_name, expert_id, shard_id)
843 + return fused_moe_make_expert_params_mapping(
844 + self,
845 + ckpt_gate_proj_name="w1",
846 + ckpt_down_proj_name="w2",
847 + ckpt_up_proj_name="w3",
848 + num_experts=self.config.n_routed_experts,
849 + )
850 +
851 + def finalize_mega_moe_weights(self) -> None:
852 + for layer in islice(self.layers, self.start_layer, self.end_layer):
853 + layer.ffn.finalize_mega_moe_weights()
854 +
855 + def finalize_mhc_broadcast_weights(self) -> None:
856 + if not get_pp_group().is_first_rank or self.start_layer >= self.end_layer:
857 + return
858 + layer = self.layers[self.start_layer]
859 + if isinstance(layer, DeepseekV4DecoderLayer):
860 + broadcast = (
861 + layer.hc_attn_fn.detach()
862 + .view(-1, layer.hc_mult, layer.hidden_size)
863 + .sum(dim=1)
864 + )
865 + if layer.hc_attn_fn_broadcast is None:
866 + layer.hc_attn_fn_broadcast = broadcast
867 + else:
868 + layer.hc_attn_fn_broadcast.copy_(broadcast)
869 +
870 +
871 +def _linear_scale_param_name(vllm_config: VllmConfig, expert_dtype: str) -> str:
872 + """Parameter name the linear quant method registers for weight scales.
873 +
874 + Native MXFP8 checkpoints ([32, 32] blocks with MXFP4 experts) route linear
875 + layers through ModelOptLinearMethod, which registers ``weight_scale``;
876 + block-FP8 linear layers register ``weight_scale_inv``.
877 + """
878 + use_mxfp8 = (
879 + getattr(vllm_config.quant_config, "weight_block_size", None) == [32, 32]
880 + and expert_dtype == "fp4"
881 + )
882 + return "weight_scale" if use_mxfp8 else "weight_scale_inv"
883 +
884 +
885 +def _make_deepseek_v4_weights_mapper(
886 + expert_dtype: str, linear_scale_name: str = "weight_scale_inv"
887 +) -> WeightsMapper:
888 + if expert_dtype == "fp4":
889 + # MXFP4 experts use Mxfp4MoEMethod, which registers scales as
890 + # ``w{1,2,3}_weight_scale`` (no _inv suffix). Linear scales use the
891 + # parameter name registered by their quantization method.
892 + scale_regex = {
893 + # ``.base_layer.``-namespace variant (LoRA-wrapped experts).
894 + re.compile(
895 + r"(\.experts\.\d+\.w[123]\.base_layer)\.scale$"
896 + ): r"\1.weight_scale",
897 + re.compile(r"(\.experts\.\d+\.w[123])\.scale$"): r"\1.weight_scale",
898 + # The ``embed.weight`` -> ``embed_tokens.weight`` suffix rule
899 + # renames the engram fp8 table but not its scale; route the
900 + # scale explicitly to the same module.
901 + re.compile(r"(engram\.embed)\.scale$"): r"\1_tokens.weight_scale_inv",
902 + re.compile(r"\.scale$"): f".{linear_scale_name}",
903 + }
904 + else:
905 + # FP8 experts use Fp8MoEMethod (block_quant=True), which registers
906 + # scales as ``w{13,2}_weight_scale_inv``. Map all ``.scale`` keys
907 + # there.
908 + scale_regex = {
909 + # ``.base_layer.``-namespace variant of the above.
910 + re.compile(
911 + r"(\.experts\.\d+\.w[123]\.base_layer)\.scale$"
912 + ): r"\1.weight_scale_inv",
913 + # Same engram reroute as the fp4 branch above.
914 + re.compile(r"(engram\.embed)\.scale$"): r"\1_tokens.weight_scale_inv",
915 + re.compile(r"\.scale$"): f".{linear_scale_name}",
916 + }
917 + return WeightsMapper(
918 + orig_to_new_prefix={
919 + "layers.": "model.layers.",
920 + "embed.": "model.embed.",
921 + "norm.": "model.norm.",
922 + "mtp.": "model.mtp.",
923 + },
924 + orig_to_new_regex=scale_regex,
925 + orig_to_new_suffix={
926 + "head.weight": "lm_head.weight",
927 + "embed.weight": "embed_tokens.weight",
928 + ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias",
929 + },
930 + orig_to_new_substr={
931 + ".shared_experts.w2": ".shared_experts.down_proj",
932 + "mtp.": None,
933 + # The v4.1 checkpoint declares the VL arch even for text-only
934 + # use; the text backbone drops the vision-tower weights.
935 + "vision.": None,
936 + "aligner.": None,
937 + "image_": None,
938 + },
939 + )
940 +
941 +
942 +class DeepseekV4MixtureOfExperts(MixtureOfExperts):
943 + moe_mlp_layers: list["DeepseekV4MoE"]
944 +
945 + def extract_moe_parameters(self, example_moe: "DeepseekV4MoE | None") -> None:
946 + if example_moe is None:
947 + self.num_moe_layers = 0
948 + self.num_expert_groups = 0
949 + self.num_logical_experts = 0
950 + self.num_physical_experts = 0
951 + self.num_local_physical_experts = 0
952 + self.num_routed_experts = 0
953 + self.num_shared_experts = 0
954 + self.num_redundant_experts = 0
955 + return
956 + self.num_logical_experts = example_moe.n_logical_experts
957 + self.num_physical_experts = example_moe.n_physical_experts
958 + self.num_local_physical_experts = example_moe.n_local_physical_experts
959 + self.num_routed_experts = example_moe.n_routed_experts
960 + self.num_shared_experts = example_moe.n_shared_experts
961 + self.num_redundant_experts = example_moe.n_redundant_experts
962 +
963 + def update_physical_experts_metadata(
964 + self,
965 + num_physical_experts: int,
966 + num_local_physical_experts: int,
967 + ) -> None:
968 + assert self.num_local_physical_experts == num_local_physical_experts
969 + self.num_physical_experts = num_physical_experts
970 + self.num_local_physical_experts = num_local_physical_experts
971 + self.num_redundant_experts = num_physical_experts - self.num_logical_experts
972 + for moe in self.moe_mlp_layers:
973 + moe.n_local_physical_experts = num_local_physical_experts
974 + moe.n_physical_experts = num_physical_experts
975 + moe.n_redundant_experts = self.num_redundant_experts
976 + moe.experts.update_expert_map()
977 +
978 +
979 +class DeepseekV41LLMForCausalLM(
980 + nn.Module,
981 + SupportsPP,
982 + SupportsEagle3,
983 + SupportsLoRA,
984 + DeepseekV4MixtureOfExperts,
985 +):
986 + model_cls = DeepseekV4Model
987 +
988 + # Default mapper assumes the original FP4-expert checkpoint layout.
989 + # Overridden per-instance in __init__ when expert_dtype != "fp4".
990 + hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper("fp4")
991 +
992 + packed_modules_mapping = {
993 + "gate_up_proj": ["w1", "w3"],
994 + "fused_wqa_wkv": ["wq_a", "wkv"],
995 + "fused_wkv_wgate": ["wkv", "wgate"],
996 + }
997 +
998 + # The MTP draft head is not LoRA-adapted.
999 + lora_skip_prefixes = ["mtp."]
1000 +
1001 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1002 + super().__init__()
1003 +
1004 + config = vllm_config.model_config.hf_config
1005 + self.config = config
1006 + expert_dtype = getattr(config, "expert_dtype", "fp4")
1007 + self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper(
1008 + expert_dtype, _linear_scale_param_name(vllm_config, expert_dtype)
1009 + )
1010 +
1011 + self.model = self.model_cls(
1012 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
1013 + )
1014 + if get_pp_group().is_last_rank:
1015 + self.lm_head = ParallelLMHead(
1016 + config.vocab_size,
1017 + config.hidden_size,
1018 + prefix=maybe_prefix(prefix, "lm_head"),
1019 + )
1020 + else:
1021 + self.lm_head = PPMissingLayer()
1022 + self.logits_processor = LogitsProcessor(config.vocab_size)
1023 + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign]
1024 + self.model.make_empty_intermediate_tensors
1025 + )
1026 +
1027 + self.set_moe_parameters()
1028 +
1029 + def set_moe_parameters(self) -> None:
1030 + self.num_expert_groups = getattr(self.config, "n_group", 1)
1031 + self.num_moe_layers = self.config.num_hidden_layers
1032 + self.moe_layers: list[nn.Module] = []
1033 + self.moe_mlp_layers: list[DeepseekV4MoE] = []
1034 + example_moe: DeepseekV4MoE | None = None
1035 + for layer in self.model.layers:
1036 + if isinstance(layer, PPMissingLayer):
1037 + continue
1038 + if not isinstance(layer, DeepseekV4DecoderLayer):
1039 + continue
1040 + if isinstance(layer.ffn, DeepseekV4MoE):
1041 + example_moe = layer.ffn
1042 + self.moe_mlp_layers.append(layer.ffn)
1043 + self.moe_layers.append(layer.ffn.experts)
1044 +
1045 + self.num_moe_layers = len(self.moe_layers)
1046 + self.extract_moe_parameters(example_moe)
1047 +
1048 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1049 + return self.model.embed_input_ids(input_ids)
1050 +
1051 + def compute_logits(
1052 + self,
1053 + hidden_states: torch.Tensor,
1054 + ) -> torch.Tensor | None:
1055 + logits = self.logits_processor(self.lm_head, hidden_states)
1056 + return logits
1057 +
1058 + def compute_logits_local(
1059 + self,
1060 + hidden_states: torch.Tensor,
1061 + ) -> torch.Tensor:
1062 + return self.logits_processor(self.lm_head, hidden_states, skip_gather=True)
1063 +
1064 + @staticmethod
1065 + def get_model_state_cls():
1066 + from .model_state import DeepseekV41ModelState
1067 +
1068 + return DeepseekV41ModelState
1069 +
1070 + @property
1071 + def token_lookback_depth(self) -> int:
1072 + """Tokens before a chunk start the engram hash needs; the model runner
1073 + passes them as `lookback_token_ids`."""
1074 + engram_hash = self.model.engram_hash
1075 + return engram_hash.lookback_depth if engram_hash is not None else 0
1076 +
1077 + def forward(
1078 + self,
1079 + input_ids: torch.Tensor,
1080 + positions: torch.Tensor,
1081 + intermediate_tensors: IntermediateTensors | None = None,
1082 + inputs_embeds: torch.Tensor | None = None,
1083 + lookback_token_ids: torch.Tensor | None = None,
1084 + ) -> torch.Tensor | IntermediateTensors:
1085 + hidden_states = self.model(
1086 + input_ids,
1087 + positions,
1088 + intermediate_tensors,
1089 + inputs_embeds,
1090 + lookback_token_ids=lookback_token_ids,
1091 + )
1092 + return hidden_states
1093 +
1094 + def get_mtp_target_hidden_states(self) -> torch.Tensor | None:
1095 + """Pre-collapse residual stream buffer (max_num_batched_tokens,
1096 + hc_mult * hidden_size) for the MTP draft model. Populated by
1097 + forward(); valid after each target step."""
1098 + return getattr(self.model, "_mtp_hidden_buffer", None)
1099 +
1100 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1101 + loader = AutoWeightsLoader(self)
1102 + loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
1103 + self.process_weights_after_loading()
1104 + return loaded_params
1105 +
1106 + def process_weights_after_loading(self) -> None:
1107 + self.model.finalize_mega_moe_weights()
1108 + self.model.finalize_mhc_broadcast_weights()
1109 +
1110 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1111 + return self.model.get_expert_mapping()
vllm/models/deepseek_v4_1/nvidia/model_state.py new
+98
@@ -0,0 +1,98 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import Any
4 +
5 +import torch
6 +import torch.nn as nn
7 +
8 +from vllm.config import VllmConfig
9 +from vllm.triton_utils import tl, triton
10 +from vllm.v1.worker.gpu.input_batch import InputBatch
11 +from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
12 +from vllm.v1.worker.gpu.model_states.default import DefaultModelState
13 +from vllm.v1.worker.gpu.states import RequestState
14 +
15 +
16 +@triton.jit
17 +def _gather_lookback_kernel(
18 + lookback_ptr,
19 + idx_mapping_ptr,
20 + num_computed_tokens_ptr,
21 + all_token_ids_ptr,
22 + all_token_ids_stride,
23 + num_reqs,
24 + DEPTH: tl.constexpr,
25 + BLOCK_DEPTH: tl.constexpr,
26 +):
27 + # One program per lookback row; rows past the batch are filled with -1.
28 + batch_idx = tl.program_id(0)
29 + in_batch = batch_idx < num_reqs
30 + req_state_idx = tl.load(idx_mapping_ptr + batch_idx, mask=in_batch, other=0)
31 + num_computed = tl.load(num_computed_tokens_ptr + req_state_idx)
32 +
33 + offs = tl.arange(0, BLOCK_DEPTH)
34 + pos = num_computed - 1 - offs
35 + valid = in_batch & (offs < DEPTH) & (pos >= 0)
36 + ids = tl.load(
37 + all_token_ids_ptr + req_state_idx * all_token_ids_stride + pos,
38 + mask=valid,
39 + other=-1,
40 + )
41 + tl.store(lookback_ptr + batch_idx * DEPTH + offs, ids, mask=offs < DEPTH)
42 +
43 +
44 +class DeepseekV41ModelState(DefaultModelState):
45 + """DefaultModelState plus the engram lookback window.
46 +
47 + The engram n-gram hash needs the ids of the ``depth`` tokens preceding
48 + each request's chunk start (see ``common/engram.py``). The runner keeps
49 + the full token history on device, so the window is gathered there every
50 + step: exact for prompt and generated tokens alike, whatever instance
51 + produced their KV.
52 + """
53 +
54 + def __init__(
55 + self,
56 + vllm_config: VllmConfig,
57 + model: nn.Module,
58 + encoder_cache: EncoderCache | None,
59 + device: torch.device,
60 + ):
61 + super().__init__(vllm_config, model, encoder_cache, device)
62 + depth = model.token_lookback_depth
63 + self.lookback_token_ids: torch.Tensor | None = None
64 + if depth > 0:
65 + # Persistent so a captured graph can read it on replay.
66 + self.lookback_token_ids = torch.full(
67 + (self.max_num_reqs, depth), -1, dtype=torch.int32, device=device
68 + )
69 +
70 + def prepare_inputs(
71 + self, input_batch: InputBatch, req_states: RequestState
72 + ) -> dict[str, torch.Tensor | None]:
73 + model_inputs = super().prepare_inputs(input_batch, req_states)
74 + window = self.lookback_token_ids
75 + if window is None:
76 + return model_inputs
77 + all_token_ids = req_states.all_token_ids.gpu
78 + depth = window.shape[1]
79 + _gather_lookback_kernel[(window.shape[0],)](
80 + window,
81 + input_batch.idx_mapping,
82 + req_states.num_computed_tokens.gpu,
83 + all_token_ids,
84 + all_token_ids.stride(0),
85 + input_batch.idx_mapping.shape[0],
86 + DEPTH=depth,
87 + BLOCK_DEPTH=triton.next_power_of_2(depth),
88 + )
89 + model_inputs["lookback_token_ids"] = window
90 + return model_inputs
91 +
92 + def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]:
93 + model_inputs = super().prepare_dummy_inputs(num_reqs, num_tokens)
94 + if self.lookback_token_ids is not None:
95 + # The captured graph reads this buffer; replays refill it in place.
96 + self.lookback_token_ids.fill_(-1)
97 + model_inputs["lookback_token_ids"] = self.lookback_token_ids
98 + return model_inputs
vllm/models/deepseek_v4_1/nvidia/vl_model.py new
+357
@@ -0,0 +1,357 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""DeepSeek-V4.1 vision variant.
4 +
5 +Thin multimodal wrapper around the text-only ``DeepseekV41LLMForCausalLM``:
6 +
7 +- ``vision`` ViT + ``aligner`` produce per-image embeddings for the IMAGE
8 + positions; three learned vectors (``image_start`` / ``image_newline`` /
9 + ``image_end``) fill the delimiter positions. Every image-span position
10 + carries ``image_token_id`` (129264) in ``input_ids``; the per-position
11 + roles come from the processor's ``types`` tensor (see
12 + ``common/mm_preprocess.py``, and ``common/vision.py`` for the tower).
13 +- Merged embeddings enter the text model via ``inputs_embeds``, i.e. before
14 + its hyper-connection stream expansion. Raw ``input_ids`` still flow into
15 + the model so the MoE router can apply ``bias_vl`` to image tokens
16 + (``requires_raw_input_tokens``).
17 +"""
18 +
19 +from collections.abc import Iterable
20 +from typing import Annotated
21 +
22 +import torch
23 +from torch import nn
24 +
25 +from vllm.distributed import get_tensor_model_parallel_world_size
26 +from vllm.model_executor.models.interfaces import (
27 + MultiModalEmbeddings,
28 + SupportsEagle3,
29 + SupportsMultiModal,
30 + SupportsPP,
31 +)
32 +from vllm.model_executor.models.utils import (
33 + AutoWeightsLoader,
34 + WeightsMapper,
35 + maybe_prefix,
36 +)
37 +from vllm.model_executor.models.vision import is_vit_use_data_parallel
38 +from vllm.models.deepseek_v4.common.vision import (
39 + DeepseekV4Aligner,
40 + DeepseekV4ViT,
41 + run_dp_sharded_vision_tower,
42 +)
43 +from vllm.multimodal import MULTIMODAL_REGISTRY
44 +from vllm.utils.tensor_schema import TensorSchema, TensorShape
45 +
46 +from ..common.mm_preprocess import (
47 + IMAGE,
48 + IMAGE_END,
49 + IMAGE_NEW_LINE,
50 + IMAGE_PAD_ID,
51 + IMAGE_PLACEHOLDER,
52 + IMAGE_SENTINEL_BASE_ID,
53 + IMAGE_START,
54 + DeepseekV4VLDummyInputsBuilder,
55 + DeepseekV4VLMultiModalProcessor,
56 + DeepseekV4VLProcessingInfo,
57 +)
58 +from .model import (
59 + DeepseekV41LLMForCausalLM,
60 + _linear_scale_param_name,
61 + _make_deepseek_v4_weights_mapper,
62 +)
63 +
64 +
65 +class DeepseekV4VLImagePixelInputs(TensorSchema):
66 + """ViT patch inputs for one batched set of images.
67 +
68 + Dimensions:
69 + - np: Total ViT patches across images (sum of n_vit_h * n_vit_w)
70 + - c: Number of image channels (3)
71 + - p: ViT patch size
72 + - ni: Number of images
73 + - ns: Total image-span positions (sum of n_llm_h * (n_llm_w + 1) + 2)
74 + """
75 +
76 + patches: Annotated[
77 + torch.Tensor, TensorShape("np", 3, "p", "p", dynamic_dims={"np"})
78 + ]
79 + # [n_vit_h, n_vit_w] per image
80 + vit_grid: Annotated[torch.Tensor, TensorShape("ni", 2)]
81 + # [n_llm_h, n_llm_w] per image
82 + llm_grid: Annotated[torch.Tensor, TensorShape("ni", 2)]
83 + types: Annotated[torch.Tensor, TensorShape("ns", dynamic_dims={"ns"})]
84 +
85 +
86 +def _make_deepseek_v4_vl_weights_mapper(
87 + expert_dtype: str, linear_scale_name: str
88 +) -> WeightsMapper:
89 + """Text-checkpoint mapping rules re-rooted under ``language_model.``."""
90 + base = _make_deepseek_v4_weights_mapper(expert_dtype, linear_scale_name)
91 + orig_to_new_prefix: dict[str, str | None] = {
92 + "layers.": "language_model.model.layers.",
93 + "embed.": "language_model.model.embed.",
94 + "norm.": "language_model.model.norm.",
95 + "hc_head": "language_model.model.hc_head",
96 + "mtp.": "language_model.model.mtp.",
97 + }
98 + return WeightsMapper(
99 + orig_to_new_prefix=orig_to_new_prefix,
100 + orig_to_new_regex=base.orig_to_new_regex,
101 + orig_to_new_suffix={
102 + "head.weight": "language_model.lm_head.weight",
103 + "embed.weight": "embed_tokens.weight",
104 + ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias",
105 + },
106 + orig_to_new_substr={
107 + ".shared_experts.w2": ".shared_experts.down_proj",
108 + # The MTP/DSpark draft heads are not supported for the vision
109 + # variant; drop their weights.
110 + "mtp.": None,
111 + },
112 + )
113 +
114 +
115 +@MULTIMODAL_REGISTRY.register_processor(
116 + DeepseekV4VLMultiModalProcessor,
117 + info=DeepseekV4VLProcessingInfo,
118 + dummy_inputs=DeepseekV4VLDummyInputsBuilder,
119 +)
120 +class DeepseekV41ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3):
121 + """Multimodal entry point for DeepSeek-V4.1 checkpoints with a vision tower.
122 +
123 + ``SupportsEagle3`` (aux hidden-state plumbing for MTP/DSpark drafters)
124 + delegates through ``language_model`` via the protocol defaults.
125 + """
126 +
127 + supports_encoder_tp_data = True
128 +
129 + # The MoE router needs raw token ids to detect image-span tokens
130 + # (all carrying image_token_id, see common/mm_preprocess.py) and apply
131 + # bias_vl.
132 + requires_raw_input_tokens = True
133 +
134 + @classmethod
135 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:
136 + if modality == "image":
137 + return IMAGE_PLACEHOLDER
138 + raise ValueError(f"Unsupported modality: {modality!r}")
139 +
140 + def __init__(self, *, vllm_config, prefix: str = "") -> None:
141 + super().__init__()
142 + model_config = vllm_config.model_config
143 + config = model_config.hf_config
144 + self.config = config
145 + self.multimodal_config = model_config.multimodal_config
146 + assert self.multimodal_config is not None
147 +
148 + # The tower is always built; _mark_tower_model stubs it out
149 + # (StageMissingLayer, weights skipped) when the image limit is 0.
150 + with self._mark_tower_model(vllm_config, {"image"}):
151 + self.use_data_parallel = is_vit_use_data_parallel(config.vision_n_heads)
152 + self.vision = DeepseekV4ViT(config)
153 + self.aligner = DeepseekV4Aligner(config)
154 + self.image_start = nn.Parameter(
155 + torch.empty(config.hidden_size, dtype=torch.float32)
156 + )
157 + self.image_end = nn.Parameter(
158 + torch.empty(config.hidden_size, dtype=torch.float32)
159 + )
160 + self.image_newline = nn.Parameter(
161 + torch.empty(config.hidden_size, dtype=torch.float32)
162 + )
163 + self.vision.to(dtype=model_config.dtype)
164 + self.aligner.to(dtype=model_config.dtype)
165 +
166 + with self._mark_language_model(vllm_config):
167 + self.language_model = DeepseekV41LLMForCausalLM(
168 + vllm_config=vllm_config,
169 + prefix=maybe_prefix(prefix, "language_model"),
170 + )
171 + # The outer mapper (see load_weights) fully resolves HF names into
172 + # this wrapper's namespace before AutoWeightsLoader strips the
173 + # "language_model." prefix and delegates to the child's load_weights,
174 + # so the child's own mapper must be a no-op. Its suffix rules are not
175 + # idempotent (e.g. "lm_head.weight".endswith("head.weight") would
176 + # re-fire "head.weight" -> "lm_head.weight").
177 + self.language_model.hf_to_vllm_mapper = WeightsMapper()
178 + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign]
179 + self.language_model.make_empty_intermediate_tensors
180 + )
181 +
182 + expert_dtype = getattr(config, "expert_dtype", "fp4")
183 + self.hf_to_vllm_mapper = _make_deepseek_v4_vl_weights_mapper(
184 + expert_dtype, _linear_scale_param_name(vllm_config, expert_dtype)
185 + )
186 +
187 + def _parse_and_validate_image_input(
188 + self, **kwargs: object
189 + ) -> DeepseekV4VLImagePixelInputs | None:
190 + patches = kwargs.pop("patches", None)
191 + if patches is None:
192 + return None
193 + return DeepseekV4VLImagePixelInputs(
194 + patches=patches,
195 + vit_grid=kwargs.pop("vit_grid"),
196 + llm_grid=kwargs.pop("llm_grid"),
197 + types=kwargs.pop("types"),
198 + resolve_bindings={"p": self.config.vision_patch_size},
199 + )
200 +
201 + def _encode_image(
202 + self,
203 + patches: torch.Tensor,
204 + n_vit_h: int,
205 + n_vit_w: int,
206 + ) -> torch.Tensor:
207 + # Aligner rows in reading order, one per IMAGE slot.
208 + return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w)
209 +
210 + def _build_image_span(
211 + self, image_embeds: torch.Tensor, types: torch.Tensor
212 + ) -> torch.Tensor:
213 + """Full image span: aligner rows at IMAGE slots, the learned
214 + delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END."""
215 + types = types.to(image_embeds.device)
216 + span = image_embeds.new_empty(types.numel(), image_embeds.shape[-1])
217 + dtype = image_embeds.dtype
218 + span[types == IMAGE_START] = self.image_start.to(dtype)
219 + span[types == IMAGE_END] = self.image_end.to(dtype)
220 + span[types == IMAGE_NEW_LINE] = self.image_newline.to(dtype)
221 + span[types == IMAGE] = image_embeds
222 + return span
223 +
224 + def _process_image_input(
225 + self,
226 + image_input: DeepseekV4VLImagePixelInputs,
227 + ) -> tuple[torch.Tensor, ...]:
228 + patches = image_input.patches.to(self.aligner.w1.weight.dtype)
229 + vit_grid = image_input.vit_grid.tolist()
230 +
231 + image_embeds_list: list[torch.Tensor]
232 + if self.use_data_parallel and get_tensor_model_parallel_world_size() > 1:
233 + # Data-parallel ViT: shard images across TP ranks and all-gather
234 + # the per-image embeddings (weights are replicated on every rank).
235 + image_embeds_list = run_dp_sharded_vision_tower(
236 + self.vision, self.aligner, patches, vit_grid
237 + )
238 + else:
239 + image_embeds_list = []
240 + vit_offset = 0
241 + for n_vit_h, n_vit_w in vit_grid:
242 + n_vit = n_vit_h * n_vit_w
243 + image_embeds_list.append(
244 + self._encode_image(
245 + patches[vit_offset : vit_offset + n_vit], n_vit_h, n_vit_w
246 + )
247 + )
248 + vit_offset += n_vit
249 +
250 + embeds: list[torch.Tensor] = []
251 + span_offset = 0
252 + for image_embeds, (n_llm_h, n_llm_w) in zip(
253 + image_embeds_list, image_input.llm_grid.tolist(), strict=True
254 + ):
255 + span_len = n_llm_h * (n_llm_w + 1) + 2
256 + embeds.append(
257 + self._build_image_span(
258 + image_embeds,
259 + image_input.types[span_offset : span_offset + span_len],
260 + )
261 + )
262 + span_offset += span_len
263 + return tuple(embeds)
264 +
265 + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
266 + image_input = self._parse_and_validate_image_input(**kwargs)
267 + if image_input is None:
268 + return []
269 + return self._process_image_input(image_input)
270 +
271 + def embed_input_ids(
272 + self,
273 + input_ids: torch.Tensor,
274 + multimodal_embeddings: MultiModalEmbeddings | None = None,
275 + *,
276 + is_multimodal: torch.Tensor | None = None,
277 + ) -> torch.Tensor:
278 + from vllm.model_executor.models.utils import _merge_multimodal_embeddings
279 +
280 + # Compressor-alignment pads borrow a reserved id; embed them as the
281 + # plain image token (the checkpoint has no image_pad vector).
282 + # Branch-free: safe inside compiled/captured regions.
283 + input_ids = input_ids.masked_fill(
284 + input_ids == IMAGE_PAD_ID, IMAGE_SENTINEL_BASE_ID
285 + )
286 + inputs_embeds = self.language_model.embed_input_ids(input_ids)
287 +
288 + if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
289 + return inputs_embeds
290 +
291 + assert is_multimodal is not None
292 + return _merge_multimodal_embeddings(
293 + inputs_embeds=inputs_embeds,
294 + multimodal_embeddings=multimodal_embeddings,
295 + is_multimodal=is_multimodal,
296 + )
297 +
298 + @staticmethod
299 + def get_model_state_cls():
300 + from .model_state import DeepseekV41ModelState
301 +
302 + return DeepseekV41ModelState
303 +
304 + @property
305 + def token_lookback_depth(self) -> int:
306 + return self.language_model.token_lookback_depth
307 +
308 + def forward(
309 + self,
310 + input_ids: torch.Tensor,
311 + positions: torch.Tensor,
312 + intermediate_tensors=None,
313 + inputs_embeds: torch.Tensor | None = None,
314 + lookback_token_ids: torch.Tensor | None = None,
315 + **kwargs,
316 + ) -> torch.Tensor:
317 + return self.language_model(
318 + input_ids,
319 + positions,
320 + intermediate_tensors,
321 + inputs_embeds,
322 + lookback_token_ids=lookback_token_ids,
323 + )
324 +
325 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
326 + return self.language_model.compute_logits(hidden_states)
327 +
328 + def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor:
329 + return self.language_model.compute_logits_local(hidden_states)
330 +
331 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
332 + return self.language_model.get_expert_mapping()
333 +
334 + def get_mtp_target_hidden_states(self) -> torch.Tensor | None:
335 + """Pre-hc_head residual stream buffer for the MTP/DSpark draft model."""
336 + return self.language_model.get_mtp_target_hidden_states()
337 +
338 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
339 + # Map HF names into this wrapper's namespace up front and sort, so
340 + # the "language_model." group reaches the child loader as one
341 + # contiguous block (AutoWeightsLoader delegates per contiguous group,
342 + # and the child's load_weights finalizes fused expert weights, which
343 + # must not run on a partially loaded model).
344 + mapped = sorted(self.hf_to_vllm_mapper.apply(weights), key=lambda x: x[0])
345 + loader = AutoWeightsLoader(self)
346 + loaded_params = loader.load_weights(mapped)
347 + # The child's load_weights already ran its post-load finalization.
348 + self._weights_finalized = True
349 + return loaded_params
350 +
351 + def process_weights_after_loading(self) -> None:
352 + # Model-level post-load hook (called by the loader after any load
353 + # format). Under DummyModelLoader the child's load_weights — and
354 + # hence its finalize step — is bypassed, so run it here instead.
355 + if getattr(self, "_weights_finalized", False):
356 + return
357 + self.language_model.process_weights_after_loading()
vllm/models/deepseek_v4_1/quant_config.py new
+229
@@ -0,0 +1,229 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Quantization config for DeepSeek V4."""
4 +
5 +from __future__ import annotations
6 +
7 +from typing import TYPE_CHECKING, cast
8 +
9 +from vllm.config import get_current_vllm_config
10 +from vllm.config.quantization import QuantSpec
11 +from vllm.model_executor.layers.fused_moe import (
12 + RoutedExperts,
13 + UnquantizedFusedMoEMethod,
14 +)
15 +from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
16 +from vllm.model_executor.layers.quantization import QuantizationMethods
17 +from vllm.model_executor.layers.quantization.fp8 import Fp8Config
18 +from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod
19 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
20 + is_layer_skipped,
21 + kMxfp8Dynamic,
22 + kMxfp8Static,
23 +)
24 +
25 +_DEEPSEEK_V4_EXPERT_DTYPES = ("fp4", "fp8")
26 +
27 +if TYPE_CHECKING:
28 + from vllm.model_executor.layers.quantization.modelopt import (
29 + ModelOptNvFp4Config,
30 + )
31 +
32 +
33 +class DeepseekV4FP8Config(Fp8Config):
34 + """FP8 config for DeepSeek V4 with expert-dtype-aware MoE dispatch.
35 +
36 + DeepSeek V4 checkpoints always use FP8 block quantization for
37 + linear/attention layers. The MoE expert weights vary by checkpoint:
38 + - ``expert_dtype="fp4"`` (e.g. DeepSeek-V4-Flash): MXFP4 experts
39 + with ue8m0 (e8m0fnu) FP8 linear scales.
40 + - ``expert_dtype="fp8"`` (e.g. DeepSeek-V4-Flash-Base): FP8 block
41 + experts with float32 FP8 linear scales.
42 +
43 + The dispatch and the linear scale dtype are both keyed off
44 + ``expert_dtype`` from the model's hf_config; missing values default
45 + to ``"fp4"`` so existing FP4 checkpoints stay unchanged.
46 +
47 + NOTE: ``expert_dtype`` is resolved lazily because this config is
48 + constructed during VllmConfig setup, before ``set_current_vllm_config``
49 + is active. Reading hf_config eagerly in ``__init__`` would always see
50 + the default ``"fp4"`` and silently misroute Flash-Base checkpoints.
51 + """
52 +
53 + def __init__(self, *args, **kwargs):
54 + super().__init__(*args, **kwargs)
55 + self._resolved_expert_dtype: str | None = None
56 + self._resolved_moe_quant_algo: str | None = None
57 + self._nvfp4_config: ModelOptNvFp4Config | None = None
58 + # ``is_scale_e8m0`` is a property that resolves on first read,
59 + # by which time the current vllm_config has been set.
60 +
61 + @property
62 + def expert_dtype(self) -> str:
63 + if self._resolved_expert_dtype is None:
64 + try:
65 + hf_config = get_current_vllm_config().model_config.hf_config
66 + except Exception:
67 + # vllm_config not yet set; defer the decision until a
68 + # later call lands inside set_current_vllm_config.
69 + return "fp4"
70 + expert_dtype = getattr(hf_config, "expert_dtype", "fp4")
71 + if expert_dtype not in _DEEPSEEK_V4_EXPERT_DTYPES:
72 + raise ValueError(
73 + f"Unsupported DeepSeek V4 expert_dtype={expert_dtype!r}; "
74 + f"expected one of {_DEEPSEEK_V4_EXPERT_DTYPES}."
75 + )
76 + self._resolved_expert_dtype = expert_dtype
77 + from vllm.logger import init_logger
78 +
79 + init_logger(__name__).info_once(
80 + "DeepSeek V4 expert_dtype resolved to %r", expert_dtype
81 + )
82 + return self._resolved_expert_dtype
83 +
84 + @property
85 + def is_scale_e8m0(self) -> bool:
86 + # FP4 checkpoints store FP8 linear scales as e8m0fnu; FP8 expert
87 + # checkpoints (Flash-Base) store them as float32.
88 + return self.expert_dtype == "fp4"
89 +
90 + def _resolve_moe_overrides(self) -> None:
91 + if self._resolved_moe_quant_algo is not None:
92 + return
93 + try:
94 + hf_config = get_current_vllm_config().model_config.hf_config
95 + except Exception:
96 + return
97 + quant_cfg = getattr(hf_config, "quantization_config", None) or {}
98 + algo = (quant_cfg.get("moe_quant_algo") or "").upper() or None
99 + self._resolved_moe_quant_algo = algo or ""
100 +
101 + @property
102 + def moe_quant_algo(self) -> str:
103 + self._resolve_moe_overrides()
104 + return self._resolved_moe_quant_algo or ""
105 +
106 + def _get_nvfp4_config(self) -> ModelOptNvFp4Config:
107 + if self._nvfp4_config is None:
108 + from vllm.model_executor.layers.quantization.modelopt import (
109 + ModelOptNvFp4Config,
110 + )
111 +
112 + self._nvfp4_config = ModelOptNvFp4Config(
113 + is_checkpoint_nvfp4_serialized=True,
114 + kv_cache_quant_algo=None,
115 + exclude_modules=[],
116 + group_size=16,
117 + )
118 + return self._nvfp4_config
119 +
120 + @classmethod
121 + def get_name(cls) -> QuantizationMethods:
122 + return "deepseek_v4_fp8"
123 +
124 + @staticmethod
125 + def _is_quark_mxfp4_ocp(hf_quant_cfg: dict) -> bool:
126 + """True for AMD-Quark exports whose global scheme is MXFP4."""
127 + weight = (hf_quant_cfg.get("global_quant_config") or {}).get("weight")
128 + # A non-dict weight (e.g. a list of multiple specs) means not an OCP
129 + # MXFP4 scheme (e.g. NVFP4 with 2-level scale).
130 + if not isinstance(weight, dict):
131 + return False
132 + return (
133 + weight.get("dtype") == "fp4"
134 + and weight.get("qscheme") == "per_group"
135 + and weight.get("group_size") == 32
136 + )
137 +
138 + @classmethod
139 + def override_quantization_method(
140 + cls, hf_quant_cfg, user_quant, hf_config=None
141 + ) -> QuantizationMethods | None:
142 + if not (
143 + isinstance(hf_quant_cfg, dict)
144 + and (
145 + hf_quant_cfg.get("quant_method") in ("fp8", "deepseek_v4_fp8")
146 + or (
147 + hf_quant_cfg.get("quant_method") == "quark"
148 + and cls._is_quark_mxfp4_ocp(hf_quant_cfg)
149 + )
150 + )
151 + ):
152 + return None
153 + model_type = getattr(hf_config, "model_type", None)
154 + if (
155 + model_type
156 + in (
157 + "deepseek_v4",
158 + "deepseek_v4_text",
159 + "deepseek_v41",
160 + "deepseek_v41_text",
161 + )
162 + or user_quant == "deepseek_v4_fp8"
163 + ):
164 + return "deepseek_v4_fp8"
165 + return None
166 +
167 + @classmethod
168 + def from_config(cls, config: dict) -> DeepseekV4FP8Config:
169 + # Reroute AMD-Quark fused shared expert MXFP4 checkpoints onto the fp8
170 + # path: the runtime layout matches the DeepSeek-native fp8 checkpoint,
171 + # so translate the schema into format Fp8Config.from_config expects.
172 + if config.get("quant_method") == "quark":
173 + quark_exclude = config.get("exclude") or []
174 + config = {
175 + "quant_method": "fp8",
176 + "activation_scheme": "dynamic",
177 + "fmt": "e4m3",
178 + "scale_fmt": "ue8m0",
179 + "weight_block_size": [128, 128],
180 + "ignored_layers": [
181 + name for name in quark_exclude if isinstance(name, str)
182 + ],
183 + }
184 + return cast("DeepseekV4FP8Config", super().from_config(config))
185 +
186 + def get_quant_method(self, layer, prefix):
187 + if (
188 + isinstance(layer, LinearBase)
189 + and self.weight_block_size == [32, 32]
190 + and self.is_scale_e8m0
191 + ):
192 + if is_layer_skipped(
193 + prefix=prefix,
194 + ignored_layers=self.ignored_layers,
195 + fused_mapping=self.packed_modules_mapping,
196 + match_mode=self.ignored_layers_match_mode,
197 + ):
198 + return UnquantizedLinearMethod()
199 + from vllm.model_executor.layers.quantization.modelopt import (
200 + CkptCtx,
201 + ModelOptLinearMethod,
202 + )
203 +
204 + rows, cols = self.weight_block_size
205 + return ModelOptLinearMethod(
206 + QuantSpec(weight=kMxfp8Static, activation=kMxfp8Dynamic),
207 + CkptCtx(scale_block_size=(rows, cols)),
208 + )
209 + if isinstance(layer, RoutedExperts):
210 + if is_layer_skipped(
211 + prefix=prefix,
212 + ignored_layers=self.ignored_layers,
213 + fused_mapping=self.packed_modules_mapping,
214 + ):
215 + return UnquantizedFusedMoEMethod(layer.moe_config)
216 + if self.expert_dtype == "fp4":
217 + if self.moe_quant_algo == "NVFP4":
218 + from vllm.model_executor.layers.quantization.modelopt import (
219 + ModelOptNvFp4FusedMoE,
220 + )
221 +
222 + return ModelOptNvFp4FusedMoE(
223 + quant_config=self._get_nvfp4_config(),
224 + moe_config=layer.moe_config,
225 + )
226 + return Mxfp4MoEMethod(layer.moe_config)
227 + # expert_dtype == "fp8": fall through to Fp8Config which
228 + # returns Fp8MoEMethod with block-wise float32 scales.
229 + return super().get_quant_method(layer, prefix)
vllm/models/deepseek_v4_1/sparse_mla.py new
+231
@@ -0,0 +1,231 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""DeepSeek-V4.1 FlashMLA sparse backend, metadata, and metadata builders."""
4 +
5 +from dataclasses import dataclass
6 +from typing import Any, ClassVar
7 +
8 +import torch
9 +
10 +from vllm.config import VllmConfig
11 +from vllm.config.cache import CacheDType
12 +from vllm.platforms import current_platform
13 +from vllm.platforms.interface import DeviceCapability
14 +from vllm.v1.attention.backend import (
15 + AttentionBackend,
16 + AttentionCGSupport,
17 + AttentionMetadata,
18 + AttentionMetadataBuilder,
19 + CommonAttentionMetadata,
20 + MultipleOf,
21 +)
22 +from vllm.v1.attention.backends.mla.compressor_utils import get_compressed_slot_mapping
23 +from vllm.v1.attention.backends.mla.sparse_swa import (
24 + _LAYER_TYPE_C1A,
25 + _LAYER_TYPE_C2A,
26 + _LAYER_TYPE_SWAONLY,
27 + DeepseekSparseSWAMetadataBuilder,
28 +)
29 +from vllm.v1.kv_cache_interface import AttentionSpec
30 +
31 +# v4.1 per-layer compress ratios: 0 = pure sliding window, 1 = full-length
32 +# compressed cache, 2 = ratio-2 compressed cache. Ratio-1 and ratio-2 layers
33 +# both attend over indexer topk indices into a shared compressed cache but
34 +# differ in compressed page block size (block_size // ratio), so each needs its
35 +# own FlashMLA tile-scheduler plan.
36 +_V41_LAYER_TYPES: dict[int, str] = {
37 + 0: _LAYER_TYPE_SWAONLY,
38 + 1: _LAYER_TYPE_C1A,
39 + 2: _LAYER_TYPE_C2A,
40 +}
41 +
42 +
43 +def deepseek_v41_layer_type(compress_ratio: int) -> str:
44 + layer_type = _V41_LAYER_TYPES.get(compress_ratio)
45 + if layer_type is None:
46 + raise ValueError(
47 + f"Unsupported DeepSeek V4.1 compress_ratio={compress_ratio}; "
48 + "expected 0, 1, or 2."
49 + )
50 + return layer_type
51 +
52 +
53 +class DeepseekV41SparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder):
54 + """SWA metadata builder base for v4.1.
55 +
56 + The shared builder classifies decode layer types by the v4.0 ratios
57 + (1 = SWA-only, 4, 128). v4.1 uses 0 / 1 / 2, so recompute the set of
58 + tile-scheduler plans from the v4.1 topology.
59 + """
60 +
61 + def __init__(self, *args, **kwargs):
62 + super().__init__(*args, **kwargs)
63 + compress_ratios = getattr(
64 + self.vllm_config.model_config.hf_config, "compress_ratios", None
65 + ) or [0]
66 + self._layer_types = {
67 + deepseek_v41_layer_type(int(ratio)) for ratio in compress_ratios
68 + }
69 +
70 +
71 +class DeepseekV4SparseMLABackend(AttentionBackend):
72 + """DeepSeek-V4.1 sparse-MLA backend base.
73 +
74 + Subclasses ``AttentionBackend`` directly (not the V3.2
75 + ``FlashMLASparseBackend``): DeepSeek-V4.1 runs its own attention layer
76 + (``DeepseekV4Attention``), so it does not reuse the V3.2 builder or impl, and
77 + only needs to declare its own metadata builder, KV-cache layout, and the
78 + sparse-MLA capability flags.
79 + """
80 +
81 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
82 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
83 + "auto",
84 + "fp8_ds_mla",
85 + "fp8", # alias for fp8_ds_mla
86 + ]
87 +
88 + @staticmethod
89 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
90 + return [64 if current_platform.is_device_capability_family(90) else 128]
91 +
92 + @staticmethod
93 + def get_builder_cls() -> type["DeepseekV4SparseMLAMetadataBuilder"]:
94 + return DeepseekV4SparseMLAMetadataBuilder
95 +
96 + @staticmethod
97 + def get_impl_cls() -> type[Any]:
98 + # DeepSeek-V4.1 runs its attention through ``DeepseekV4Attention.forward``,
99 + # not the generic ``Attention``/``MLAAttention`` layer, so the backend's
100 + # impl class is never instantiated.
101 + raise NotImplementedError(
102 + "DeepseekV4SparseMLABackend has no separate impl class; DeepSeek-V4.1 "
103 + "attention runs through DeepseekV4Attention."
104 + )
105 +
106 + @classmethod
107 + def get_supported_head_sizes(cls) -> list[int]:
108 + # DeepSeek V4 layout: 448 NoPE + 64 RoPE = 512.
109 + return [512]
110 +
111 + @classmethod
112 + def is_mla(cls) -> bool:
113 + return True
114 +
115 + @classmethod
116 + def is_sparse(cls) -> bool:
117 + return True
118 +
119 + @classmethod
120 + def supports_sink(cls) -> bool:
121 + return True
122 +
123 + @classmethod
124 + def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
125 + return capability.major in [9, 10]
126 +
127 +
128 +@dataclass
129 +class DeepseekV4FlashMLAMetadata(AttentionMetadata):
130 + num_reqs: int
131 + max_query_len: int
132 + max_seq_len: int
133 +
134 + num_actual_tokens: int # Number of tokens excluding padding.
135 + query_start_loc: torch.Tensor
136 + slot_mapping: torch.Tensor
137 +
138 + block_table: torch.Tensor
139 + req_id_per_token: torch.Tensor
140 + block_size: int
141 + topk_tokens: int
142 +
143 +
144 +class DeepseekV4SparseMLAMetadataBuilder(
145 + AttentionMetadataBuilder[DeepseekV4FlashMLAMetadata]
146 +):
147 + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
148 +
149 + def __init__(
150 + self,
151 + kv_cache_spec: AttentionSpec,
152 + layer_names: list[str],
153 + vllm_config: VllmConfig,
154 + device: torch.device,
155 + ) -> None:
156 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
157 + self.model_config = vllm_config.model_config
158 + # Classify single-token queries (plus num_speculative_tokens via
159 + # supports_spec_as_decode=True) as decodes; longer queries go to prefill.
160 + self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
161 + self.topk_tokens = self.model_config.hf_config.index_topk
162 +
163 + max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
164 + self.req_id_per_token_buffer = torch.empty(
165 + (max_num_batched_tokens,), dtype=torch.int32, device=device
166 + )
167 +
168 + assert isinstance(self.kv_cache_spec.tokens_per_state, int)
169 + self.compress_ratio = self.kv_cache_spec.tokens_per_state
170 + # Only kv-source layers own a compressed cache, and v4.1 compresses at
171 + # ratio 1 or 2.
172 + if self.compress_ratio not in (1, 2):
173 + raise ValueError(
174 + "DeepSeek V4.1 compressed-KV caches use compress_ratio 1 or 2; "
175 + f"got {self.compress_ratio}."
176 + )
177 +
178 + # Pre-allocate compressed slot mapping buffer for CUDA graph address
179 + # stability when compress_ratio > 1.
180 + if self.compress_ratio > 1:
181 + self.compressed_slot_mapping_buffer = torch.empty(
182 + max_num_batched_tokens, dtype=torch.int64, device=device
183 + )
184 +
185 + def build(
186 + self,
187 + common_prefix_len: int,
188 + common_attn_metadata: CommonAttentionMetadata,
189 + fast_build: bool = False,
190 + ) -> DeepseekV4FlashMLAMetadata:
191 + cm = common_attn_metadata
192 + req_id_per_token = cm.token_to_req_indices(self.req_id_per_token_buffer)
193 +
194 + slot_mapping = cm.slot_mapping
195 + if self.compress_ratio > 1:
196 + slot_mapping = get_compressed_slot_mapping(
197 + cm.num_actual_tokens,
198 + cm.query_start_loc,
199 + cm.seq_lens,
200 + cm.block_table_tensor.clamp_(min=0),
201 + int(self.kv_cache_spec.num_states),
202 + self.compress_ratio,
203 + out=self.compressed_slot_mapping_buffer,
204 + )
205 +
206 + return DeepseekV4FlashMLAMetadata(
207 + num_reqs=cm.num_reqs,
208 + max_query_len=cm.max_query_len,
209 + max_seq_len=cm.max_seq_len,
210 + num_actual_tokens=cm.num_actual_tokens,
211 + query_start_loc=cm.query_start_loc,
212 + slot_mapping=slot_mapping,
213 + block_table=cm.block_table_tensor,
214 + req_id_per_token=req_id_per_token,
215 + block_size=self.kv_cache_spec.block_size,
216 + topk_tokens=self.topk_tokens,
217 + )
218 +
219 +
220 +class DeepseekV4FlashMLAMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder):
221 + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
222 +
223 +
224 +class DeepseekV4FlashMLABackend(DeepseekV4SparseMLABackend):
225 + @staticmethod
226 + def get_name() -> str:
227 + return "FLASHMLA_SPARSE_DSV41"
228 +
229 + @staticmethod
230 + def get_builder_cls() -> type[DeepseekV4FlashMLAMetadataBuilder]:
231 + return DeepseekV4FlashMLAMetadataBuilder
vllm/v1/attention/backends/mla/sparse_swa.py
+37 -17
@@ -46,18 +46,26 @@ from vllm.v1.kv_cache_interface import (
46 _LAYER_TYPE_SWAONLY = "swaonly"
47 _LAYER_TYPE_C4A = "c4a"
48 _LAYER_TYPE_C128A = "c128a"
49 +# v4.1 ratio-1 / ratio-2 indexer layers: same indexer-path shape as C4A, but
50 +# their compressed page block is block_size // ratio, so each gets its own plan.
51 +# v4.1 builders classify layer types themselves (ratio 0 is SWA-only there); see
52 +# deepseek_v4_1/sparse_mla.py.
53 +_LAYER_TYPE_C1A = "c1a"
54 +_LAYER_TYPE_C2A = "c2a"
55
56
57 def _layer_type_for(compress_ratio: int) -> str:
58 if compress_ratio <= 1:
59 return _LAYER_TYPE_SWAONLY
60 + if compress_ratio == 2:
61 + return _LAYER_TYPE_C2A
62 if compress_ratio == 4:
63 return _LAYER_TYPE_C4A
64 if compress_ratio == 128:
65 return _LAYER_TYPE_C128A
66 raise ValueError(
67 f"Unsupported DeepseekV4 compress_ratio={compress_ratio}; "
60 - "expected 1, 4, or 128."
68 + "expected 1, 2, 4, or 128."
69 )
70
71
@@ -70,6 +78,7 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):
78 prefix: str,
79 cache_config: CacheConfig,
80 backend_cls: "type[AttentionBackend] | None" = None,
81 + block_size: int = 64,
82 ):
83 super().__init__()
84 self.backend_cls = backend_cls or DeepseekSparseSWABackend
@@ -84,12 +93,9 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):
93 raise ValueError(f"Duplicate layer name: {prefix}")
94 compilation_config.static_forward_context[prefix] = self
95
87 - # Block size is constrained by tensor sharing between SWA and C4A KV blocks.
88 - # Since both block types share the same physical tensor, they must use the
89 - # same page size. The C4A KV block shape [256//4, head_dim] = [64, head_dim]
90 - # determines the SWA block size of 64 tokens per block.
91 - # TODO(yifan): make SWA block size automatically determined and configurable.
92 - self.block_size = 64
96 + # Any multiple of 32; the sparse decode kernels take the page size at
97 + # runtime.
98 + self.block_size = block_size
99 # uint8: fp8_ds_mla UE8M0 paged layout. bfloat16 / float8_e4m3fn:
100 # contiguous full-cache layout.
101 assert self.dtype in (torch.uint8, torch.bfloat16, torch.float8_e4m3fn)
@@ -130,7 +136,7 @@ class DeepseekSparseSWABackend(AttentionBackend):
136
137 @staticmethod
138 def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
133 - return [MultipleOf(64)]
139 + return [MultipleOf(32)]
140
141 @classmethod
142 def get_preferred_block_size(cls, default_block_size: int) -> int:
@@ -213,12 +219,17 @@ class DeepseekSparseSWAMetadata:
219 tile_sched_swaonly: "FlashMLASchedMeta | None" = None
220 tile_sched_c4a: "FlashMLASchedMeta | None" = None
221 tile_sched_c128a: "FlashMLASchedMeta | None" = None
222 + tile_sched_c1a: "FlashMLASchedMeta | None" = None
223 + tile_sched_c2a: "FlashMLASchedMeta | None" = None
224 flashinfer_sparse_index_cache: dict[str, tuple[torch.Tensor, torch.Tensor]] = field(
225 default_factory=dict
226 )
227
228 def get_prefill_chunk_plan(
221 - self, compress_ratio: int, prefill_chunk_size: int
229 + self,
230 + compress_ratio: int,
231 + prefill_chunk_size: int,
232 + has_compressed: bool | None = None,
233 ) -> list[tuple[int, int, int, int]]:
234 if self.num_prefills == 0:
235 return []
@@ -226,17 +237,20 @@ class DeepseekSparseSWAMetadata:
237 assert self.prefill_seq_lens_cpu is not None
238 assert self.prefill_query_lens_cpu is not None
239
240 + # Whether the layer gathers a compressed-KV region into the prefill
241 + # workspace. v4.0 callers keep the legacy default (ratio <= 1 means
242 + # SWA-only); v4.1 passes has_compressed explicitly because its
243 + # compress_ratio==1 layers DO have a full-length compressed cache.
244 + if has_compressed is None:
245 + has_compressed = compress_ratio > 1
246 +
247 # query_len <= max_num_batched_tokens and
248 # gather_len = query_len + min(prefix_len, window_size - 1), so the
249 # worst-case gathered width is bounded by
250 # max_num_batched_tokens + window_size - 1. The compressed prefix pool
251 # is bounded by ceil(max_model_len / compress_ratio).
252 max_workspace_area = prefill_chunk_size * (
235 - (
236 - 0
237 - if compress_ratio <= 1
238 - else cdiv(self.prefill_max_model_len, compress_ratio)
239 - )
253 + (cdiv(self.prefill_max_model_len, compress_ratio) if has_compressed else 0)
254 + self.prefill_window_size
255 + self.prefill_max_num_batched_tokens
256 )
@@ -245,13 +259,13 @@ class DeepseekSparseSWAMetadata:
259 prefix_lens_cpu, min=0, max=self.prefill_window_size - 1
260 )
261 compressed_lens_cpu = (
248 - torch.zeros_like(self.prefill_seq_lens_cpu)
249 - if compress_ratio <= 1
250 - else torch.div(
262 + torch.div(
263 self.prefill_seq_lens_cpu,
264 compress_ratio,
265 rounding_mode="floor",
266 )
267 + if has_compressed
268 + else torch.zeros_like(self.prefill_seq_lens_cpu)
269 )
270
271 chunk_plan: list[tuple[int, int, int, int]] = []
@@ -715,6 +729,8 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
729 tile_sched_swaonly=tile_sched[_LAYER_TYPE_SWAONLY],
730 tile_sched_c4a=tile_sched[_LAYER_TYPE_C4A],
731 tile_sched_c128a=tile_sched[_LAYER_TYPE_C128A],
732 + tile_sched_c1a=tile_sched[_LAYER_TYPE_C1A],
733 + tile_sched_c2a=tile_sched[_LAYER_TYPE_C2A],
734 **deepseek_v4_fields, # type: ignore[arg-type]
735 )
736
@@ -795,6 +811,8 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
811 metadata.tile_sched_swaonly = tile_sched[_LAYER_TYPE_SWAONLY]
812 metadata.tile_sched_c4a = tile_sched[_LAYER_TYPE_C4A]
813 metadata.tile_sched_c128a = tile_sched[_LAYER_TYPE_C128A]
814 + metadata.tile_sched_c1a = tile_sched[_LAYER_TYPE_C1A]
815 + metadata.tile_sched_c2a = tile_sched[_LAYER_TYPE_C2A]
816 metadata.flashinfer_sparse_index_cache.clear()
817
818 def build_tile_scheduler(
@@ -815,6 +833,8 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
833 _LAYER_TYPE_SWAONLY: None,
834 _LAYER_TYPE_C4A: None,
835 _LAYER_TYPE_C128A: None,
836 + _LAYER_TYPE_C1A: None,
837 + _LAYER_TYPE_C2A: None,
838 }
839 if (
840 num_decode_tokens == 0
vllm/v1/attention/backends/registry.py
+10
@@ -98,6 +98,16 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
98 ROCM_FLASHMLA_SPARSE_DSV4 = (
99 "vllm.models.deepseek_v4.amd.rocm.DeepseekV4ROCMAiterMLASparseBackend"
100 )
101 + # DeepSeek V4.1 sparse MLA backends (model-driven; selected via the V4.1
102 + # layer). Separate names from DSV4 so a V4.1 model never resolves the
103 + # V4.0 backend classes through this enum.
104 + FLASHMLA_SPARSE_DSV41 = (
105 + "vllm.models.deepseek_v4_1.sparse_mla.DeepseekV4FlashMLABackend"
106 + )
107 + FLASHINFER_MLA_SPARSE_DSV41 = (
108 + "vllm.models.deepseek_v4_1.nvidia.flashinfer_sparse."
109 + "DeepseekV4FlashInferMLASparseBackend"
110 + )
111 B12X = "vllm.v1.attention.backends.b12x.B12xPagedAttentionBackend"
112 FLASH_ATTN_MLA = "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend"
113 FLASH_ATTN_MLA_SPARSE = (