@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 89876b0c54

[GLM5] Implement op fusion for GLM5/DSV3.2 (#46876)

Woosuk Kwon committed Jun 28, 2026 at 08:17 UTC 89876b0c548afdd932d41d5ef81b14347edb9ff7
9 files changed +1530 -91
tests/kernels/test_fused_deepseek_v32_norm_rope.py new
+423
@@ -0,0 +1,423 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for the horizontally-fused deepseek_v32 (NVIDIA SM100) Triton
4 +kernels used by the specialized DSA model:
5 +
6 + fused_norm_rope
7 + - q : q_lora RMSNorm
8 + - kv : kv_lora RMSNorm + (interleaved) RoPE on k_pe + MLA cache insert
9 + (bf16 or per-tensor fp8)
10 + - idx: indexer-K LayerNorm + RoPE (interleaved or NeoX) + UE8M0 fp8 quant +
11 + packed indexer cache insert; plus the top-k buffer (-1) fill
12 + fused_q
13 + - mqa: ql_nope + (interleaved) RoPE'd q_pe, concat-quantized to the fp8 MQA
14 + query
15 + - idx: indexer-Q RoPE (interleaved or NeoX) + UE8M0 fp8 quant + folded
16 + index weights
17 + fused_eh_norm (MTP): zero-at-pos-0 + enorm RMSNorm(embeds) + hnorm
18 + RMSNorm(prev), concatenated side-by-side
19 +
20 +Each kernel is compared against a PyTorch reference. The kernel keeps the whole
21 +pipeline in fp32 and rounds once, so it can land on the opposite side of a
22 +round-to-nearest tie from the reference for a few elements: deterministic fp8
23 +outputs are checked within 1 representable-step (ULP); bf16 norm/RoPE outputs use
24 +rtol/atol=1e-2 (the tolerance the sibling deepseek_v4 fused-kernel test uses).
25 +"""
26 +
27 +import pytest
28 +import torch
29 +
30 +from vllm.models.deepseek_v32.nvidia import kernels as K
31 +from vllm.platforms import current_platform
32 +
33 +FP8 = torch.float8_e4m3fn
34 +FP8_MAX = 448.0
35 +
36 +# GLM-5.2 / DeepSeek-V3.2 shapes (TP8 local heads).
37 +Q_LORA = 2048
38 +KV_LORA = 512
39 +ROPE_DIM = 64
40 +NUM_HEADS = 8
41 +INDEX_HEADS = 32
42 +INDEX_HEAD_DIM = 128
43 +HIDDEN = 6144
44 +EPS = 1e-6
45 +
46 +pytestmark = pytest.mark.skipif(
47 + not current_platform.is_cuda() or not current_platform.has_device_capability(89),
48 + reason="deepseek_v32 fused kernels require CUDA with fp8 (SM89+)",
49 +)
50 +
51 +
52 +# ── reference helpers ────────────────────────────────────────────────────────
53 +
54 +
55 +def make_cos_sin(max_pos: int, rot_dim: int, device) -> torch.Tensor:
56 + """cos||sin cache: row[pos] = [cos(theta)(rot/2), sin(theta)(rot/2)]."""
57 + half = rot_dim // 2
58 + inv_freq = 1.0 / (
59 + 10000.0 ** (torch.arange(0, half, dtype=torch.float32, device=device) / half)
60 + )
61 + t = torch.arange(max_pos, dtype=torch.float32, device=device)
62 + freqs = torch.einsum("i,j->ij", t, inv_freq)
63 + return torch.cat([freqs.cos(), freqs.sin()], dim=-1)
64 +
65 +
66 +def rms_norm(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
67 + """RMSNorm matching kernels._rms_norm (fp32, eps inside rsqrt). Returns fp32."""
68 + xf = x.float()
69 + ms = xf.pow(2).mean(dim=-1, keepdim=True)
70 + return xf * torch.rsqrt(ms + EPS) * w.float()
71 +
72 +
73 +def layer_norm(x: torch.Tensor, w: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
74 + xf = x.float()
75 + mean = xf.mean(dim=-1, keepdim=True)
76 + var = (xf - mean).pow(2).mean(dim=-1, keepdim=True)
77 + return (xf - mean) * torch.rsqrt(var + EPS) * w.float() + b.float()
78 +
79 +
80 +def rope(
81 + x: torch.Tensor, pos: torch.Tensor, cos_sin: torch.Tensor, interleave: bool
82 +) -> torch.Tensor:
83 + """Apply RoPE to the first ``rot_dim`` elements of x's last dim.
84 +
85 + x: [..., head_dim] fp32. ``cos_sin`` is [max_pos, rot_dim]. ``interleave``
86 + selects adjacent-pair (GLM) vs split-half NeoX (DeepSeek-V3.2) layout.
87 + """
88 + rot = cos_sin.shape[-1]
89 + half = rot // 2
90 + cs = cos_sin[pos.long()]
91 + cos, sin = cs[..., :half], cs[..., half:]
92 + out = x.float().clone()
93 + r = out[..., :rot]
94 + if interleave:
95 + x1, x2 = r[..., 0::2].clone(), r[..., 1::2].clone()
96 + r[..., 0::2] = x1 * cos - x2 * sin
97 + r[..., 1::2] = x2 * cos + x1 * sin
98 + else:
99 + x1, x2 = r[..., :half].clone(), r[..., half:].clone()
100 + r[..., :half] = x1 * cos - x2 * sin
101 + r[..., half:] = x2 * cos + x1 * sin
102 + return out
103 +
104 +
105 +def ue8m0_quant(vals: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
106 + """Per-row (last dim) UE8M0 fp8 quant matching kernels._fp8_ue8m0_quantize."""
107 + amax = vals.float().abs().amax(dim=-1, keepdim=True)
108 + scale = torch.clamp(amax, min=1e-4) / FP8_MAX
109 + scale = torch.exp2(torch.ceil(torch.log2(scale)))
110 + q = (vals.float() / scale).to(FP8)
111 + return q, scale.squeeze(-1)
112 +
113 +
114 +def _bf16_ulp(a: torch.Tensor, b: torch.Tensor) -> int:
115 + def key(t):
116 + u = t.contiguous().view(torch.int16).to(torch.int64) & 0xFFFF
117 + return torch.where(u >= 0x8000, 0xFFFF - u, u + 0x8000)
118 +
119 + return int((key(a) - key(b)).abs().max().item())
120 +
121 +
122 +def _fp8_ulp(a: torch.Tensor, b: torch.Tensor) -> int:
123 + def key(t):
124 + u = t.contiguous().view(torch.uint8).to(torch.int64)
125 + return torch.where(u >= 0x80, 0xFF - u, u + 0x80)
126 +
127 + return int((key(a) - key(b)).abs().max().item())
128 +
129 +
130 +def assert_bf16(got: torch.Tensor, ref_fp32: torch.Tensor, msg: str):
131 + # Kernel keeps RMSNorm/RoPE in fp32 and rounds to bf16 once; the fp32
132 + # reduction/FMA order differs from torch, so a few elements land on the
133 + # opposite side of a round-to-nearest tie. Use the same tolerance the
134 + # sibling deepseek_v4 fused-kernel test uses for this bf16 norm+rope class.
135 + torch.testing.assert_close(
136 + got.float(), ref_fp32.float(), rtol=1e-2, atol=1e-2, msg=lambda m: f"{msg}: {m}"
137 + )
138 +
139 +
140 +def assert_fp8(got: torch.Tensor, ref: torch.Tensor, msg: str):
141 + assert _fp8_ulp(got, ref) <= 1, f"{msg}: >1 fp8 ULP"
142 +
143 +
144 +# ── fused_norm_rope ──────────────────────────────────────────────────────────
145 +
146 +
147 +@pytest.mark.parametrize("num_tokens", [1, 4, 17, 512, 4096])
148 +@pytest.mark.parametrize("index_interleave", [True, False])
149 +@pytest.mark.parametrize("mla_fp8", [False, True])
150 +def test_fused_norm_rope(num_tokens: int, index_interleave: bool, mla_fp8: bool):
151 + torch.manual_seed(0)
152 + dev = "cuda"
153 + max_pos = 8192
154 + pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos
155 +
156 + q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
157 + kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
158 + k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
159 + qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
160 + kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)
161 + ik = torch.randn(num_tokens, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16)
162 + ikw = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
163 + ikb = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
164 +
165 + mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev) # MLA k_pe: interleaved
166 + idx_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
167 +
168 + bs = max_pos # single block covering all tokens
169 + mla_dim = KV_LORA + ROPE_DIM
170 + if mla_fp8:
171 + mla_cache = torch.zeros(1, bs, mla_dim, device=dev, dtype=torch.uint8)
172 + mla_dtype = "fp8"
173 + mla_k_scale = torch.tensor([0.3], device=dev, dtype=torch.float32)
174 + else:
175 + mla_cache = torch.zeros(1, bs, mla_dim, device=dev, dtype=torch.bfloat16)
176 + mla_dtype = "auto"
177 + mla_k_scale = None
178 + idx_row = INDEX_HEAD_DIM + INDEX_HEAD_DIM // 128 * 4 # 132
179 + idx_cache = torch.zeros(1, bs, idx_row, device=dev, dtype=torch.uint8)
180 + slot = torch.arange(num_tokens, device=dev, dtype=torch.int64)
181 + topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32)
182 +
183 + q_out = K.fused_norm_rope(
184 + pos,
185 + q_c,
186 + qw,
187 + EPS,
188 + kv_c,
189 + kvw,
190 + EPS,
191 + k_pe,
192 + mla_cos_sin,
193 + ik,
194 + ikw,
195 + ikb,
196 + EPS,
197 + idx_cos_sin,
198 + topk,
199 + slot_mapping=slot,
200 + indexer_k_cache=idx_cache,
201 + mla_kv_cache=mla_cache,
202 + mla_kv_cache_dtype=mla_dtype,
203 + mla_k_scale=mla_k_scale,
204 + has_indexer=True,
205 + index_rope_interleave=index_interleave,
206 + )
207 +
208 + # q_lora RMSNorm
209 + assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm")
210 +
211 + # MLA cache: [kv_c_normed | k_pe_roped(interleaved)]
212 + kv_ref = rms_norm(kv_c, kvw)
213 + kpe_ref = rope(k_pe.float(), pos, mla_cos_sin, interleave=True)
214 + if mla_fp8:
215 + cache = mla_cache.view(FP8)[0, :num_tokens]
216 + s = mla_k_scale.item()
217 + assert_fp8(cache[:, :KV_LORA], (kv_ref / s).to(FP8), "MLA kv fp8")
218 + assert_fp8(cache[:, KV_LORA:], (kpe_ref / s).to(FP8), "MLA k_pe fp8")
219 + else:
220 + cache = mla_cache[0, :num_tokens]
221 + assert_bf16(cache[:, :KV_LORA], kv_ref, "MLA kv bf16")
222 + assert_bf16(cache[:, KV_LORA:], kpe_ref, "MLA k_pe bf16")
223 +
224 + # Indexer-K cache (packed [bs*head_dim fp8 | bs*4 fp32 scale]).
225 + ik_ref = layer_norm(ik, ikw, ikb)
226 + ik_ref = rope(ik_ref, pos, idx_cos_sin, interleave=index_interleave)
227 + q_ref, s_ref = ue8m0_quant(ik_ref)
228 + flat = idx_cache[0].reshape(-1)
229 + vals = flat[: bs * INDEX_HEAD_DIM].view(FP8).reshape(bs, INDEX_HEAD_DIM)
230 + scales = flat[bs * INDEX_HEAD_DIM :].view(torch.float32)
231 + assert_fp8(vals[:num_tokens], q_ref, "indexer-K fp8")
232 + torch.testing.assert_close(scales[:num_tokens], s_ref, rtol=0, atol=0)
233 +
234 + # Top-k buffer cleared to -1 on indexer layers.
235 + assert (topk == -1).all(), "topk buffer not cleared on indexer layer"
236 +
237 +
238 +@pytest.mark.parametrize("num_tokens", [1, 17, 512])
239 +def test_fused_norm_rope_no_indexer(num_tokens: int):
240 + """Shared (no-indexer) layer: q + kv/MLA only; top-k buffer untouched."""
241 + torch.manual_seed(1)
242 + dev = "cuda"
243 + max_pos = 8192
244 + pos = torch.arange(num_tokens, device=dev, dtype=torch.int64)
245 +
246 + q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
247 + kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
248 + k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
249 + qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
250 + kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)
251 + mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
252 +
253 + bs = max_pos
254 + mla_cache = torch.zeros(1, bs, KV_LORA + ROPE_DIM, device=dev, dtype=torch.bfloat16)
255 + slot = torch.arange(num_tokens, device=dev, dtype=torch.int64)
256 + topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32)
257 +
258 + q_out = K.fused_norm_rope(
259 + pos,
260 + q_c,
261 + qw,
262 + EPS,
263 + kv_c,
264 + kvw,
265 + EPS,
266 + k_pe,
267 + mla_cos_sin,
268 + None,
269 + None,
270 + None,
271 + EPS,
272 + None,
273 + topk,
274 + slot_mapping=slot,
275 + indexer_k_cache=None,
276 + mla_kv_cache=mla_cache,
277 + mla_kv_cache_dtype="auto",
278 + mla_k_scale=None,
279 + has_indexer=False,
280 + index_rope_interleave=False,
281 + )
282 +
283 + assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm (no-indexer)")
284 + cache = mla_cache[0, :num_tokens]
285 + assert_bf16(cache[:, :KV_LORA], rms_norm(kv_c, kvw), "MLA kv (no-indexer)")
286 + assert_bf16(
287 + cache[:, KV_LORA:],
288 + rope(k_pe.float(), pos, mla_cos_sin, interleave=True),
289 + "MLA k_pe (no-indexer)",
290 + )
291 + # Shared layers reuse the previous indexer's top-k: buffer must be untouched.
292 + assert (topk == 7).all(), "topk buffer should be untouched on shared layer"
293 +
294 +
295 +# ── fused_q ──────────────────────────────────────────────────────────────────
296 +
297 +
298 +@pytest.mark.parametrize("num_tokens", [1, 4, 17, 512, 4096])
299 +@pytest.mark.parametrize("index_interleave", [True, False])
300 +def test_fused_q(num_tokens: int, index_interleave: bool):
301 + torch.manual_seed(2)
302 + dev = "cuda"
303 + max_pos = 8192
304 + pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos
305 +
306 + q_pe = torch.randn(
307 + num_tokens, NUM_HEADS, ROPE_DIM, device=dev, dtype=torch.bfloat16
308 + )
309 + ql_nope = torch.randn(
310 + num_tokens, NUM_HEADS, KV_LORA, device=dev, dtype=torch.bfloat16
311 + )
312 + index_q = torch.randn(
313 + num_tokens, INDEX_HEADS, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16
314 + )
315 + index_w = torch.randn(num_tokens, INDEX_HEADS, device=dev, dtype=torch.float32)
316 + q_scale = torch.tensor([0.37], device=dev, dtype=torch.float32)
317 + softmax_scale = INDEX_HEAD_DIM**-0.5
318 + head_scale = INDEX_HEADS**-0.5
319 + q_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev) # q_pe: interleaved
320 + idx_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
321 +
322 + iq_fp8, iw_out, mqa = K.fused_q(
323 + pos,
324 + q_pe,
325 + q_cos_sin,
326 + index_q,
327 + idx_cos_sin,
328 + ql_nope,
329 + q_scale,
330 + index_w,
331 + softmax_scale,
332 + head_scale,
333 + has_indexer=True,
334 + index_rope_interleave=index_interleave,
335 + )
336 +
337 + s = q_scale.item()
338 + # MQA query: [ql_nope | q_pe RoPE'd (interleaved)], per-tensor fp8.
339 + mqa_nope_ref = (ql_nope.float() / s).to(FP8)
340 + qpe_ref = rope(
341 + q_pe.float(),
342 + pos.unsqueeze(-1).expand(num_tokens, NUM_HEADS),
343 + q_cos_sin,
344 + interleave=True,
345 + )
346 + mqa_pe_ref = (qpe_ref / s).to(FP8)
347 + assert_fp8(mqa[:, :, :KV_LORA], mqa_nope_ref, "mqa ql_nope")
348 + assert_fp8(mqa[:, :, KV_LORA:], mqa_pe_ref, "mqa q_pe")
349 +
350 + # Indexer-Q: RoPE + UE8M0 fp8 quant; index weights fold in q-scale.
351 + iq_ref = rope(
352 + index_q.float(),
353 + pos.unsqueeze(-1).expand(num_tokens, INDEX_HEADS),
354 + idx_cos_sin,
355 + interleave=index_interleave,
356 + )
357 + q_ref, scale_ref = ue8m0_quant(iq_ref)
358 + assert_fp8(iq_fp8, q_ref, "indexer-Q fp8")
359 + iw_ref = index_w * scale_ref * softmax_scale * head_scale
360 + torch.testing.assert_close(iw_out, iw_ref, rtol=1e-3, atol=1e-3)
361 +
362 +
363 +@pytest.mark.parametrize("num_tokens", [1, 17, 512])
364 +def test_fused_q_no_indexer(num_tokens: int):
365 + torch.manual_seed(3)
366 + dev = "cuda"
367 + max_pos = 8192
368 + pos = torch.arange(num_tokens, device=dev, dtype=torch.int64)
369 + q_pe = torch.randn(
370 + num_tokens, NUM_HEADS, ROPE_DIM, device=dev, dtype=torch.bfloat16
371 + )
372 + ql_nope = torch.randn(
373 + num_tokens, NUM_HEADS, KV_LORA, device=dev, dtype=torch.bfloat16
374 + )
375 + q_scale = torch.tensor([0.5], device=dev, dtype=torch.float32)
376 + q_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
377 +
378 + _, _, mqa = K.fused_q(
379 + pos,
380 + q_pe,
381 + q_cos_sin,
382 + None,
383 + None,
384 + ql_nope,
385 + q_scale,
386 + None,
387 + 0.0,
388 + 0.0,
389 + has_indexer=False,
390 + index_rope_interleave=False,
391 + )
392 + s = q_scale.item()
393 + assert_fp8(mqa[:, :, :KV_LORA], (ql_nope.float() / s).to(FP8), "mqa ql_nope")
394 + qpe_ref = rope(
395 + q_pe.float(),
396 + pos.unsqueeze(-1).expand(num_tokens, NUM_HEADS),
397 + q_cos_sin,
398 + interleave=True,
399 + )
400 + assert_fp8(mqa[:, :, KV_LORA:], (qpe_ref / s).to(FP8), "mqa q_pe")
401 +
402 +
403 +# ── fused_eh_norm (MTP) ──────────────────────────────────────────────────────
404 +
405 +
406 +@pytest.mark.parametrize("num_tokens", [1, 4, 17, 512])
407 +def test_fused_eh_norm(num_tokens: int):
408 + torch.manual_seed(4)
409 + dev = "cuda"
410 + # Mix in a position-0 token to exercise the embeds-zeroing branch.
411 + pos = torch.arange(num_tokens, device=dev, dtype=torch.int64)
412 + pos[0] = 0
413 + embeds = torch.randn(num_tokens, HIDDEN, device=dev, dtype=torch.bfloat16)
414 + prev = torch.randn(num_tokens, HIDDEN, device=dev, dtype=torch.bfloat16)
415 + ew = torch.randn(HIDDEN, device=dev, dtype=torch.bfloat16)
416 + hw = torch.randn(HIDDEN, device=dev, dtype=torch.bfloat16)
417 +
418 + out = K.fused_eh_norm(pos, embeds, prev, ew, hw, EPS)
419 +
420 + masked = torch.where(pos.unsqueeze(-1) == 0, torch.zeros_like(embeds), embeds)
421 + ref = torch.cat([rms_norm(masked, ew), rms_norm(prev, hw)], dim=-1)
422 + assert out.shape == (num_tokens, 2 * HIDDEN)
423 + assert_bf16(out, ref, "eh_norm")
vllm/model_executor/layers/fused_moe/config.py
+6
@@ -1287,6 +1287,12 @@ class FusedMoEConfig:
1287 has_bias: bool = False
1288 is_lora_enabled: bool = False
1289
1290 + # When True, the MoE skips its final cross-rank all-reduce (and the separate
1291 + # shared-expert reduce), returning the partial per-rank sum. The caller is
1292 + # then responsible for the reduction (e.g. fusing it into the next RMSNorm).
1293 + # Only honored on the non-reduced (late-AR) TP path. Default False.
1294 + skip_final_all_reduce: bool = False
1295 +
1296 # SwiGLU clamp limit. When set, backends that do not implement the clamp
1297 # are filtered out by `FusedMoEExperts.is_supported_config` so the oracle
1298 # cannot silently select one and drop the clamp.
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
+2
@@ -423,6 +423,7 @@ class MoERunner(MoERunnerInterface):
423 if (
424 shared_output is not None
425 and not self.moe_config.is_sequence_parallel
426 + and not self.moe_config.skip_final_all_reduce
427 and self._fused_output_is_reduced
428 ):
429 shared_output = tensor_model_parallel_all_reduce(shared_output)
@@ -445,6 +446,7 @@ class MoERunner(MoERunnerInterface):
446 # - The MK already reduced the fused output itself.
447 if (
448 not self.moe_config.is_sequence_parallel
449 + and not self.moe_config.skip_final_all_reduce
450 and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1)
451 and not self._fused_output_is_reduced
452 ):
vllm/model_executor/layers/sparse_attn_indexer.py
+9 -1
@@ -227,6 +227,7 @@ def sparse_attn_indexer(
227 topk_indices_buffer: torch.Tensor,
228 skip_k_cache_insert: bool,
229 use_fp4_cache: bool = False,
230 + skip_topk_buffer_clear: bool = False,
231 ) -> torch.Tensor:
232 # careful! this will be None in dummy run
233 attn_metadata = get_forward_context().attn_metadata
@@ -303,7 +304,13 @@ def sparse_attn_indexer(
304 scale_fmt,
305 )
306
306 - topk_indices_buffer[: hidden_states.shape[0]] = -1
307 + # The buffer must be pre-filled with -1 (the "no token" sentinel) before the
308 + # top-k kernels scatter valid indices into it. On the fused deepseek_v32
309 + # nvidia path, _fused_norm_rope_kernel already cleared the same
310 + # [:num_tokens, :topk] region earlier in this forward, so skip the redundant
311 + # fill.
312 + if not skip_topk_buffer_clear:
313 + topk_indices_buffer[: hidden_states.shape[0]] = -1
314 if has_prefill:
315 prefill_metadata = attn_metadata_narrowed.prefill
316 assert prefill_metadata is not None
@@ -546,6 +553,7 @@ def sparse_attn_indexer_fake(
553 topk_indices_buffer: torch.Tensor | None,
554 skip_k_cache_insert: bool,
555 use_fp4_cache: bool = False,
556 + skip_topk_buffer_clear: bool = False,
557 ) -> torch.Tensor:
558 return topk_indices_buffer
559
vllm/models/deepseek_v32/nvidia/attention.py
+168 -81
@@ -1,7 +1,5 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from typing import TYPE_CHECKING
4 -
3 import torch
4 import torch.nn as nn
5 from transformers import DeepseekV2Config, DeepseekV3Config
@@ -23,7 +21,10 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
21 per_token_group_quant_fp8,
22 )
23 from vllm.model_executor.layers.rotary_embedding import get_rope
26 -from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
24 +from vllm.model_executor.layers.sparse_attn_indexer import (
25 + SparseAttnIndexer,
26 + sparse_attn_indexer,
27 +)
28 from vllm.model_executor.models.deepseek_v2 import (
29 DeepSeekV2FusedQkvAProjLinear,
30 DeepseekV32IndexerCache,
@@ -32,10 +33,7 @@ from vllm.model_executor.models.deepseek_v2 import (
33 from vllm.model_executor.models.utils import extract_layer_index
34 from vllm.utils.torch_utils import is_quantized_kv_cache
35
35 -if TYPE_CHECKING:
36 - from vllm.model_executor.layers.attention.mla_attention import (
37 - MLACommonMetadata,
38 - )
36 +from .kernels import fused_norm_rope, fused_q
37
38
39 class DeepseekV32Indexer(nn.Module):
@@ -160,6 +158,10 @@ class DeepseekV32Indexer(nn.Module):
158
159
160 class DeepseekV32Attention(MLAAttention):
161 + # Narrow the base's broadly-typed `indexer` to the concrete type so the
162 + # `if self.indexer is not None` guards below type-check its attributes.
163 + indexer: "DeepseekV32Indexer | None"
164 +
165 def __init__(
166 self,
167 vllm_config: VllmConfig,
@@ -263,21 +265,27 @@ class DeepseekV32Attention(MLAAttention):
265 self.num_local_heads = num_local_heads
266 self.qk_head_dim = qk_head_dim
267 self.indexer = indexer
268 + self.topk_indices_buffer = topk_indices_buffer
269 # Runtime toggle for index_share_for_mtp_iteration: MTP draft step 0
270 # computes the top-k, steps 1+ set this True to reuse it.
271 self.skip_topk = False
269 - # Whether the paged KV cache must be viewed as fp8 before the attention
270 - # (per-tensor fp8; the fp8_ds_mla layout is read as uint8).
271 - self._fp8_kv_needs_view = (
272 - is_quantized_kv_cache(self.kv_cache_dtype)
273 - and self.kv_cache_dtype != "fp8_ds_mla"
274 - )
275 - # Whether the backend takes an fp8-quantized query (FlashInfer sparse)
276 - # vs the (ql_nope, q_pe) tuple (FlashMLA sparse).
277 - self._use_concat_quant = (
272 + # Single fused fp8 path: Triton fused norm/rope/cache + fused-q write a
273 + # single fp8 MQA query and the contiguous [kv_c; k_pe] MLA cache layout.
274 + # This requires an fp8 KV cache and a sparse MLA backend that accepts a
275 + # quantized query (FlashInfer sparse on SM100).
276 + assert (
277 is_quantized_kv_cache(self.kv_cache_dtype)
278 and self.impl.supports_quant_query_input
279 + ), (
280 + "deepseek_v32 (nvidia) requires an fp8 KV cache served by the "
281 + "FlashInfer sparse MLA backend (which accepts a quantized query). "
282 + "Launch with --kv-cache-dtype fp8."
283 )
284 + # The paged KV cache is stored as uint8 and viewed as fp8 for the decode
285 + # (per-tensor fp8; never the fp8_ds_mla layout on this path).
286 + self._fp8_kv_needs_view = self.kv_cache_dtype != "fp8_ds_mla"
287 + # GLM-5.2 uses interleaved indexer RoPE; DeepSeek-V3.2 uses NeoX.
288 + self._index_rope_interleave = getattr(config, "indexer_rope_interleave", False)
289
290 # Remaining MLA projections (registered on this module).
291 self.fused_qkv_a_proj = DeepSeekV2FusedQkvAProjLinear(
@@ -295,10 +303,14 @@ class DeepseekV32Attention(MLAAttention):
303 prefix=f"{prefix}.q_b_proj",
304 )
305 self.kv_a_layernorm = RMSNorm(kv_lora_rank, eps=config.rms_norm_eps)
306 + # reduce_results=False: the attention all-reduce is fused with the
307 + # following post_attention_layernorm in the decoder layer via
308 + # fused_allreduce_rms_norm.
309 self.o_proj = RowParallelLinear(
310 num_heads * v_head_dim,
311 hidden_size,
312 bias=False,
313 + reduce_results=False,
314 quant_config=quant_config,
315 prefix=f"{prefix}.o_proj",
316 )
@@ -322,102 +334,177 @@ class DeepseekV32Attention(MLAAttention):
334 positions: torch.Tensor,
335 hidden_states: torch.Tensor,
336 ) -> torch.Tensor:
337 + # Captured: A-projections (+ indexer A-GEMM on indexer layers).
338 qkv_lora = self.fused_qkv_a_proj(hidden_states)[0]
326 - q_c, kv_lora = qkv_lora.split(
327 - [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1
328 - )
329 - q_c = self.q_a_layernorm(q_c)
330 - q = self.q_b_proj(q_c)[0]
331 -
332 - kv_c, k_pe = kv_lora.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
333 - kv_c_normed = self.kv_a_layernorm(kv_c)
334 -
335 - q = q.view(-1, self.num_local_heads, self.qk_head_dim)
336 - k_pe = k_pe.unsqueeze(1)
337 - q[..., self.qk_nope_head_dim :], k_pe = self.rotary_emb(
338 - positions, q[..., self.qk_nope_head_dim :], k_pe
339 + q_c, kv_c, k_pe = qkv_lora.split(
340 + [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
341 )
342
341 - num_tokens = hidden_states.shape[0]
342 - q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
343 - q_nope = q_nope.transpose(0, 1) # (N, B, P)
344 - ql_nope = torch.bmm(q_nope, self.W_UK_T).transpose(0, 1) # (B, N, L)
345 -
346 - # Lightning indexer writes the top-k indices into the shared buffer.
347 - # "Shared" layers (indexer is None) reuse the top-k from the previous
348 - # indexer layer already sitting in the buffer.
343 if self.indexer is not None and not self.skip_topk:
350 - self.indexer(hidden_states, q_c, positions, self.indexer_rope_emb) # type: ignore[operator]
351 -
352 - attn_latent = torch.empty(
353 - (num_tokens, self.num_local_heads, self.kv_lora_rank),
354 - dtype=q.dtype,
355 - device=q.device,
356 - )
357 - self._sparse_attention(kv_c_normed, k_pe, ql_nope, q_pe, attn_latent)
344 + kw = self.indexer.wk_weights_proj(hidden_states)[0]
345 + index_k = kw[:, : self.indexer.head_dim]
346 + index_weights = kw[:, self.indexer.head_dim :]
347 + else:
348 + index_k = None
349 + index_weights = None
350
359 - # V up-projection + output projection are metadata-independent GEMMs and
360 - # stay captured.
351 + num_tokens = hidden_states.shape[0]
352 output = torch.empty(
353 (num_tokens, self.num_local_heads * self.v_head_dim),
363 - dtype=q.dtype,
364 - device=q.device,
354 + dtype=hidden_states.dtype,
355 + device=hidden_states.device,
356 + )
357 + self._fused_attention(
358 + positions, q_c, kv_c, k_pe, index_k, index_weights, output
359 )
366 - self._v_up_proj(attn_latent, out=output)
360 return self.o_proj(output)[0]
361
362 @eager_break_during_capture
370 - def _sparse_attention(
363 + def _fused_attention(
364 self,
372 - kv_c_normed: torch.Tensor,
365 + positions: torch.Tensor,
366 + q_c: torch.Tensor,
367 + kv_c: torch.Tensor,
368 k_pe: torch.Tensor,
374 - ql_nope: torch.Tensor,
375 - q_pe: torch.Tensor,
376 - attn_latent: torch.Tensor,
369 + index_k: torch.Tensor | None,
370 + index_weights: torch.Tensor | None,
371 + output: torch.Tensor,
372 ) -> None:
373 + # One eager break for the whole attention. In FULL cudagraph mode (pure
374 + # decode) this decorator is a no-op, so everything here is captured; in
375 + # PIECEWISE (prefill) it runs eagerly. The cache writes, sparse indexer,
376 + # and forward_mqa all depend on per-step metadata and must not be split
377 + # out (PIECEWISE capture would otherwise miss them).
378 forward_context = get_forward_context()
379 attn_metadata_raw = forward_context.attn_metadata
380 - attn_metadata: MLACommonMetadata | None
380 if isinstance(attn_metadata_raw, dict):
382 - attn_metadata = attn_metadata_raw[self.layer_name] # type: ignore[assignment]
381 + attn_metadata = attn_metadata_raw.get(self.layer_name)
382 elif isinstance(attn_metadata_raw, list):
384 - # Speculative decoding: [0] is the base-model metadata dict.
385 - attn_metadata = attn_metadata_raw[0][self.layer_name] # type: ignore[assignment]
383 + attn_metadata = attn_metadata_raw[0].get(self.layer_name)
384 else:
385 attn_metadata = attn_metadata_raw
386
387 slot_mapping = forward_context.slot_mapping
388 assert isinstance(slot_mapping, dict)
391 - self.impl.do_kv_cache_update( # type: ignore[attr-defined]
392 - kv_c_normed,
389 + mla_slot = slot_mapping.get(self.layer_name)
390 +
391 + if self.indexer is not None:
392 + has_indexer = True
393 + indexer_k_norm_w = self.indexer.k_norm.weight
394 + indexer_k_norm_bias = self.indexer.k_norm.bias
395 + indexer_k_norm_eps = self.indexer.k_norm.variance_epsilon
396 + indexer_k_rope_cos_sin_cache = self.indexer_rope_emb.cos_sin_cache
397 + indexer_k_cache = self.indexer.k_cache.kv_cache
398 + indexer_softmax_scale = self.indexer.softmax_scale
399 + indexer_n_head_scale = self.indexer.n_head**-0.5
400 + else:
401 + has_indexer = False
402 + indexer_k_norm_w = None
403 + indexer_k_norm_bias = None
404 + indexer_k_norm_eps = 1e-6
405 + indexer_k_rope_cos_sin_cache = None
406 + indexer_k_cache = None
407 + indexer_softmax_scale = 0.0
408 + indexer_n_head_scale = 0.0
409 +
410 + if attn_metadata is None:
411 + mla_kv_cache = None
412 + mla_k_scale = None
413 + indexer_k_cache = None
414 + mla_slot = None
415 + else:
416 + mla_kv_cache = self.kv_cache
417 + mla_k_scale = self._k_scale
418 +
419 + q_c = fused_norm_rope(
420 + positions,
421 + q_c,
422 + self.q_a_layernorm.weight,
423 + self.q_a_layernorm.variance_epsilon,
424 + kv_c,
425 + self.kv_a_layernorm.weight,
426 + self.kv_a_layernorm.variance_epsilon,
427 k_pe,
394 - self.kv_cache,
395 - slot_mapping.get(self.layer_name),
396 - self.kv_cache_dtype,
397 - self._k_scale,
428 + self.rotary_emb.cos_sin_cache,
429 + index_k,
430 + indexer_k_norm_w,
431 + indexer_k_norm_bias,
432 + indexer_k_norm_eps,
433 + indexer_k_rope_cos_sin_cache,
434 + self.topk_indices_buffer,
435 + slot_mapping=mla_slot,
436 + indexer_k_cache=indexer_k_cache,
437 + mla_kv_cache=mla_kv_cache,
438 + mla_kv_cache_dtype=self.kv_cache_dtype,
439 + mla_k_scale=mla_k_scale,
440 + has_indexer=has_indexer,
441 + index_rope_interleave=self._index_rope_interleave,
442 + )
443 +
444 + q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim)
445 + q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
446 + q_nope = q_nope.transpose(0, 1)
447 + ql_nope = torch.bmm(q_nope, self.W_UK_T).transpose(0, 1)
448 +
449 + if self.indexer is not None:
450 + index_q = self.indexer.wq_b(q_c)[0]
451 + index_q = index_q.view(-1, self.indexer.n_head, self.indexer.head_dim)
452 + else:
453 + index_q = None
454 +
455 + index_q_fp8, index_weights_out, mqa_q = fused_q(
456 + positions,
457 + q_pe,
458 + self.rotary_emb.cos_sin_cache,
459 + index_q,
460 + self.indexer_rope_emb.cos_sin_cache if has_indexer else None,
461 + ql_nope,
462 + self._q_scale,
463 + index_weights,
464 + indexer_softmax_scale,
465 + indexer_n_head_scale,
466 + has_indexer=has_indexer,
467 + index_rope_interleave=self._index_rope_interleave,
468 )
469
470 if attn_metadata is None:
401 - # Profile / warmup: zero-fill for DP+EP determinism.
402 - attn_latent.zero_()
471 + output.zero_()
472 return
473
405 - num_actual = attn_metadata.num_actual_tokens
474 + if self.indexer is not None:
475 + sparse_attn_indexer(
476 + q_c,
477 + self.indexer.k_cache.prefix,
478 + self.indexer.k_cache.kv_cache,
479 + index_q_fp8,
480 + None, # q_scale folded into weights on the fp8 path
481 + None, # k unused when skip_k_cache_insert=True
482 + index_weights_out,
483 + self.indexer.quant_block_size,
484 + self.indexer.scale_fmt,
485 + self.indexer.topk_tokens,
486 + self.indexer.head_dim,
487 + self.indexer.max_model_len,
488 + self.indexer.max_total_seq_len,
489 + self.topk_indices_buffer,
490 + True, # skip_k_cache_insert
491 + False, # use_fp4_cache
492 + True, # skip_topk_buffer_clear (fused_norm_rope already did it)
493 + )
494 +
495 + num_actual = attn_metadata.num_actual_tokens # type: ignore[attr-defined]
496 kv_cache = self.kv_cache
497 if self._fp8_kv_needs_view:
498 kv_cache = kv_cache.view(torch.float8_e4m3fn)
409 -
410 - ql_nope = ql_nope[:num_actual]
411 - q_pe = q_pe[:num_actual]
412 - # FlashInfer sparse takes a single fp8-quantized query; FlashMLA sparse
413 - # takes the (ql_nope, q_pe) tuple and concatenates internally.
414 - mqa_q: torch.Tensor | tuple[torch.Tensor, torch.Tensor]
415 - if self._use_concat_quant:
416 - mqa_q = self._decode_concat_quant_fp8_op(ql_nope, q_pe, self._q_scale)
417 - else:
418 - mqa_q = (ql_nope, q_pe)
419 -
420 - attn_out, _ = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self) # type: ignore[attr-defined]
421 - attn_latent[:num_actual] = attn_out.view(
499 + attn_out, _ = self.impl.forward_mqa( # type: ignore[attr-defined]
500 + mqa_q[:num_actual], kv_cache, attn_metadata, self
501 + )
502 + x = attn_out.view(
503 num_actual, self.num_local_heads, self.kv_lora_rank
504 + ).transpose(0, 1)
505 + out = (
506 + output[:num_actual]
507 + .view(num_actual, self.num_local_heads, self.v_head_dim)
508 + .transpose(0, 1)
509 )
510 + torch.bmm(x, self.W_UV, out=out)
vllm/models/deepseek_v32/nvidia/fused_ops.py new
+63
@@ -0,0 +1,63 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Fused ops for deepseek_v32 (eager / breakable-cudagraph path).
4 +
5 +These recover fusions that vLLM's torch.compile passes would normally do but
6 +that don't fire when running eager under the breakable CUDA graph.
7 +"""
8 +
9 +import torch
10 +
11 +from vllm.distributed import (
12 + get_tensor_model_parallel_world_size,
13 + tensor_model_parallel_all_reduce,
14 +)
15 +from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import (
16 + _AR_RESIDUAL_RMS_NORM,
17 + _can_use_flashinfer,
18 + flashinfer_trtllm_fused_allreduce_norm,
19 +)
20 +from vllm.model_executor.layers.layernorm import RMSNorm
21 +
22 +
23 +def fused_allreduce_rms_norm(
24 + hidden_states: torch.Tensor,
25 + residual: torch.Tensor,
26 + norm: RMSNorm,
27 +) -> tuple[torch.Tensor, torch.Tensor]:
28 + """All-reduce + add residual + (standard) RMSNorm, fused via flashinfer.
29 +
30 + ``hidden_states`` is the per-rank *partial* output of a row-parallel linear
31 + run with ``reduce_results=False``; ``norm`` is the RMSNorm applied right
32 + after. Returns ``(normed_output, new_residual)``, equivalent to
33 + ``norm(all_reduce(hidden_states), residual)``. Falls back to an explicit
34 + all-reduce + RMSNorm when the flashinfer fast path is unavailable.
35 + """
36 + tp_size = get_tensor_model_parallel_world_size()
37 + if tp_size == 1:
38 + return norm(hidden_states, residual)
39 +
40 + if flashinfer_trtllm_fused_allreduce_norm is not None:
41 + ok, max_token_num = _can_use_flashinfer(hidden_states, tp_size)
42 + if ok:
43 + norm_out = torch.empty_like(hidden_states)
44 + # With norm_out provided, the kernel writes the new residual
45 + # (all_reduce(hidden_states) + residual) into the hidden_states
46 + # buffer and the normalized result into norm_out.
47 + flashinfer_trtllm_fused_allreduce_norm(
48 + allreduce_in=hidden_states,
49 + residual=residual,
50 + rms_gamma=norm.weight,
51 + rms_eps=norm.variance_epsilon,
52 + world_size=tp_size,
53 + weight_bias=0.0, # standard RMSNorm (Gemma would use 1.0)
54 + launch_with_pdl=True,
55 + fp32_acc=True,
56 + max_token_num=max_token_num,
57 + pattern_code=_AR_RESIDUAL_RMS_NORM,
58 + norm_out=norm_out,
59 + )
60 + return norm_out, hidden_states
61 +
62 + reduced = tensor_model_parallel_all_reduce(hidden_states)
63 + return norm(reduced, residual)
vllm/models/deepseek_v32/nvidia/kernels.py new
+823
@@ -0,0 +1,823 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import torch
4 +
5 +from vllm.triton_utils import tl, triton
6 +
7 +# Cache of tiny 1-element dummy tensors (per device, dtype) reused by the
8 +# has_indexer=False path so the indexer args don't allocate every call.
9 +_DUMMY_CACHE: dict[tuple, torch.Tensor] = {}
10 +
11 +
12 +def _dummy(shape: tuple, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
13 + key = (shape, dtype, device)
14 + t = _DUMMY_CACHE.get(key)
15 + if t is None:
16 + t = torch.empty(shape, dtype=dtype, device=device)
17 + _DUMMY_CACHE[key] = t
18 + return t
19 +
20 +
21 +@triton.jit
22 +def _rms_norm(x, w, eps, HIDDEN_SIZE: tl.constexpr):
23 + x = x.to(tl.float32)
24 + mean_sq = tl.sum(x * x, axis=0) / HIDDEN_SIZE
25 + rrms = tl.rsqrt(mean_sq + eps)
26 + w = w.to(tl.float32)
27 + return (x * rrms) * w
28 +
29 +
30 +@triton.jit
31 +def _get_cos_sin(
32 + cos_sin_cache_ptr,
33 + cos_sin_cache_stride,
34 + pos,
35 + HALF_ROT_DIM: tl.constexpr,
36 +):
37 + block = tl.arange(0, HALF_ROT_DIM)
38 + cos = tl.load(cos_sin_cache_ptr + pos * cos_sin_cache_stride + block)
39 + cos = cos.to(tl.float32)
40 + sin = tl.load(cos_sin_cache_ptr + pos * cos_sin_cache_stride + block + HALF_ROT_DIM)
41 + sin = sin.to(tl.float32)
42 + return cos, sin
43 +
44 +
45 +@triton.jit
46 +def _fp8_ue8m0_quantize(vals):
47 + """Quantize float32 values to FP8 E4M3 with a ue8m0 (power-of-2) scale.
48 +
49 + Returns (fp8_vals, scale) so the caller can store them or reuse the scale.
50 + """
51 + vals = vals.to(tl.float32)
52 + amax = tl.max(tl.abs(vals))
53 + scale = tl.div_rn(tl.maximum(amax, 1e-4), 448.0)
54 + scale = tl.math.exp2(tl.math.ceil(tl.math.log2(scale)))
55 + fp8_vals = tl.div_rn(vals, scale).to(tl.float8e4nv)
56 + return fp8_vals, scale
57 +
58 +
59 +@triton.jit
60 +def _fp8_quant_and_cache_write(
61 + vals,
62 + mask,
63 + slot_idx,
64 + kv_cache_ptr,
65 + kv_cache_scale_ptr,
66 + cache_block_size,
67 + cache_stride,
68 + offsets,
69 + HEAD_DIM: tl.constexpr,
70 +):
71 + k_fp8, scale = _fp8_ue8m0_quantize(vals)
72 +
73 + block_idx = slot_idx // cache_block_size
74 + block_offset = slot_idx % cache_block_size
75 + block_start = block_idx * cache_block_size * cache_stride
76 +
77 + tl.store(
78 + kv_cache_ptr + block_start + block_offset * HEAD_DIM + offsets,
79 + k_fp8,
80 + mask=mask,
81 + )
82 + scale_byte_off = block_start + cache_block_size * HEAD_DIM + block_offset * 4
83 + tl.store(kv_cache_scale_ptr + scale_byte_off // 4, scale)
84 +
85 +
86 +@triton.jit
87 +def _fused_norm_rope_kernel(
88 + pos_ptr,
89 + # Q RMS norm
90 + q_c_ptr,
91 + q_c_stride,
92 + q_rms_norm_w_ptr,
93 + q_rms_eps,
94 + q_c_out_ptr,
95 + q_c_out_stride,
96 + Q_DIM: tl.constexpr,
97 + Q_BLOCK_SIZE: tl.constexpr,
98 + # KV RMS norm
99 + kv_ptr,
100 + kv_stride,
101 + kv_rms_norm_w_ptr,
102 + kv_rms_eps,
103 + KV_DIM: tl.constexpr,
104 + # KV RoPE
105 + kpe_ptr,
106 + kpe_stride,
107 + kpe_rope_cos_sin_cache_ptr,
108 + kpe_rope_cos_sin_cache_stride,
109 + KPE_HALF_ROT_DIM: tl.constexpr,
110 + # Index K layer norm
111 + index_k_ptr,
112 + index_k_stride,
113 + index_k_layer_norm_w_ptr,
114 + index_k_layer_norm_bias_ptr,
115 + index_k_layer_norm_eps,
116 + INDEX_K_DIM: tl.constexpr,
117 + INDEX_K_BLOCK_SIZE: tl.constexpr,
118 + # Index K RoPE
119 + index_k_rope_cos_sin_cache_ptr,
120 + index_k_rope_cos_sin_cache_stride,
121 + INDEX_K_HALF_ROT_DIM: tl.constexpr,
122 + # Cache params (shared by indexer K and MLA)
123 + slot_mapping_ptr,
124 + # Index K FP8 cache
125 + indexer_cache_ptr,
126 + indexer_cache_scale_ptr,
127 + indexer_cache_block_size,
128 + indexer_cache_stride,
129 + # MLA KV cache (concat kv_c_normed + k_pe_roped, uses slot_mapping_ptr)
130 + mla_cache_ptr,
131 + mla_cache_block_stride,
132 + mla_cache_entry_stride,
133 + MLA_CACHE_FP8: tl.constexpr,
134 + mla_cache_scale_ptr,
135 + # Top k indices
136 + topk_indices_ptr,
137 + topk_indices_stride,
138 + TOPK: tl.constexpr,
139 + TOPK_BLOCK_SIZE: tl.constexpr,
140 + HAS_INDEXER: tl.constexpr,
141 + INDEX_ROPE_INTERLEAVE: tl.constexpr,
142 +):
143 + pid = tl.program_id(0)
144 + tok_idx = tl.program_id(1)
145 + if pid == 3:
146 + if not HAS_INDEXER:
147 + # Shared layer: reuse the previous indexer layer's top-k; do not
148 + # clear the buffer.
149 + return
150 + # Fill top k indices buffer with -1
151 + for i in range(0, TOPK, TOPK_BLOCK_SIZE):
152 + offset = i + tl.arange(0, TOPK_BLOCK_SIZE)
153 + mask = offset < TOPK
154 + tl.store(
155 + topk_indices_ptr + tok_idx * topk_indices_stride + offset,
156 + -1,
157 + mask=mask,
158 + )
159 + return
160 +
161 + if slot_mapping_ptr is None:
162 + # Memory profiling run.
163 + return
164 + slot_idx = tl.load(slot_mapping_ptr + tok_idx)
165 + if slot_idx < 0:
166 + # Padding
167 + return
168 +
169 + if pid == 2:
170 + # Q RMS norm
171 + q_block = tl.arange(0, Q_BLOCK_SIZE)
172 + q_mask = q_block < Q_DIM
173 + q_c = tl.load(q_c_ptr + tok_idx * q_c_stride + q_block, mask=q_mask, other=0.0)
174 + q_c_rms_w = tl.load(q_rms_norm_w_ptr + q_block, mask=q_mask)
175 + q_c = _rms_norm(q_c, q_c_rms_w, q_rms_eps, Q_DIM)
176 + tl.store(q_c_out_ptr + tok_idx * q_c_out_stride + q_block, q_c, mask=q_mask)
177 + elif pid == 1:
178 + # KV RMS Norm + KV RoPE + MLA concat_and_cache.
179 + # Merged so the normed kv_c and RoPE'd k_pe can be written
180 + # to the MLA KV cache directly without a separate kernel.
181 +
182 + # KV RMS Norm (result stays in registers for MLA cache write)
183 + kv_block = tl.arange(0, KV_DIM)
184 + kv_c = tl.load(kv_ptr + tok_idx * kv_stride + kv_block)
185 + kv_c_rms_w = tl.load(kv_rms_norm_w_ptr + kv_block)
186 + kv_c = _rms_norm(kv_c, kv_c_rms_w, kv_rms_eps, KV_DIM)
187 +
188 + # KV RoPE (interleaved) on k_pe — in registers only.
189 + # k_pe is not needed after the cache write (MLA decode reads
190 + # from kv_cache), so we skip writing back to kpe_ptr.
191 + pos = tl.load(pos_ptr + tok_idx)
192 + cos, sin = _get_cos_sin(
193 + kpe_rope_cos_sin_cache_ptr,
194 + kpe_rope_cos_sin_cache_stride,
195 + pos,
196 + KPE_HALF_ROT_DIM,
197 + )
198 + dim_off = tl.arange(0, KPE_HALF_ROT_DIM)
199 + kpe_base = kpe_ptr + tok_idx * kpe_stride
200 + x1 = tl.load(kpe_base + dim_off * 2).to(tl.float32)
201 + x2 = tl.load(kpe_base + dim_off * 2 + 1).to(tl.float32)
202 + r1 = x1 * cos - x2 * sin
203 + r2 = x2 * cos + x1 * sin
204 +
205 + # MLA concat_and_cache: write [kv_c_normed, k_pe_roped] to cache.
206 + if mla_cache_entry_stride == 0:
207 + return
208 +
209 + mla_block_size = mla_cache_block_stride // mla_cache_entry_stride
210 + mla_block_idx = slot_idx // mla_block_size
211 + mla_block_off = slot_idx % mla_block_size
212 + dst = (
213 + mla_cache_ptr
214 + + mla_block_idx * mla_cache_block_stride
215 + + mla_block_off * mla_cache_entry_stride
216 + )
217 + # kv_c_normed (KV_DIM elements)
218 + if MLA_CACHE_FP8:
219 + scale = tl.load(mla_cache_scale_ptr)
220 + kv_c_fp8 = (kv_c.to(tl.float32) / scale).to(tl.float8e4nv)
221 + tl.store(dst + kv_block, kv_c_fp8)
222 + else:
223 + tl.store(dst + kv_block, kv_c)
224 + # k_pe_roped (from registers, interleaved layout)
225 + if MLA_CACHE_FP8:
226 + tl.store(dst + KV_DIM + dim_off * 2, (r1 / scale).to(tl.float8e4nv))
227 + tl.store(dst + KV_DIM + dim_off * 2 + 1, (r2 / scale).to(tl.float8e4nv))
228 + else:
229 + tl.store(dst + KV_DIM + dim_off * 2, r1)
230 + tl.store(dst + KV_DIM + dim_off * 2 + 1, r2)
231 + elif pid == 0:
232 + if not HAS_INDEXER:
233 + # Shared layer: no indexer K to process.
234 + return
235 + # Fused: Index K LayerNorm + RoPE + FP8 quant + cache write.
236 + # Eliminates the separate indexer_k_quant_and_cache kernel launch.
237 +
238 + index_k_block = tl.arange(0, INDEX_K_BLOCK_SIZE)
239 + index_k_mask = index_k_block < INDEX_K_DIM
240 + index_k = tl.load(
241 + index_k_ptr + tok_idx * index_k_stride + index_k_block,
242 + mask=index_k_mask,
243 + other=0.0,
244 + ).to(tl.float32)
245 + index_k_w = tl.load(
246 + index_k_layer_norm_w_ptr + index_k_block, mask=index_k_mask
247 + ).to(tl.float32)
248 + index_k_b = tl.load(
249 + index_k_layer_norm_bias_ptr + index_k_block, mask=index_k_mask
250 + ).to(tl.float32)
251 +
252 + # 1. LayerNorm. Keep (mean, rstd) so the RoPE rotation partner can be
253 + # re-normalized in registers below, avoiding a global scratch buffer.
254 + mean = tl.sum(index_k, axis=0) / INDEX_K_DIM
255 + diff = tl.where(index_k_mask, index_k - mean, 0.0)
256 + var = tl.sum(diff * diff, axis=0) / INDEX_K_DIM
257 + rstd = tl.rsqrt(var + index_k_layer_norm_eps)
258 + normed = (index_k - mean) * rstd * index_k_w + index_k_b
259 +
260 + # 2. RoPE on the rotation region. Supports both interleaved (adjacent
261 + # pairs, e.g. GLM-5.2) and NeoX (split-half, e.g. DeepSeek-V3.2). The
262 + # rotation partner is gathered from the read-only inputs and
263 + # re-normalized with the same (mean, rstd) — no scratch, no atomics.
264 + pos = tl.load(pos_ptr + tok_idx)
265 + in_rope = index_k_block < 2 * INDEX_K_HALF_ROT_DIM
266 + if INDEX_ROPE_INTERLEAVE:
267 + # pair i = block // 2; partner = block ^ 1; even -> -sin, odd -> +sin.
268 + cos_idx = index_k_block // 2
269 + partner_offs = tl.where(in_rope, index_k_block ^ 1, index_k_block)
270 + sign = tl.where(index_k_block % 2 == 0, -1.0, 1.0)
271 + else:
272 + # NeoX: pair across halves; partner = block ^ HALF.
273 + cos_idx = index_k_block % INDEX_K_HALF_ROT_DIM
274 + partner_offs = tl.where(
275 + in_rope, index_k_block ^ INDEX_K_HALF_ROT_DIM, index_k_block
276 + )
277 + sign = tl.where(index_k_block < INDEX_K_HALF_ROT_DIM, -1.0, 1.0)
278 + cos_full = tl.load(
279 + index_k_rope_cos_sin_cache_ptr
280 + + pos * index_k_rope_cos_sin_cache_stride
281 + + cos_idx,
282 + mask=in_rope,
283 + other=1.0,
284 + ).to(tl.float32)
285 + sin_full = tl.load(
286 + index_k_rope_cos_sin_cache_ptr
287 + + pos * index_k_rope_cos_sin_cache_stride
288 + + INDEX_K_HALF_ROT_DIM
289 + + cos_idx,
290 + mask=in_rope,
291 + other=0.0,
292 + ).to(tl.float32)
293 + # normed[partner_offs] == (raw_partner - mean) * rstd * w_partner +
294 + # b_partner: gather the raw partner and its norm affine (read-only
295 + # loads), then apply the same per-token mean/rstd.
296 + raw_partner = tl.load(
297 + index_k_ptr + tok_idx * index_k_stride + partner_offs,
298 + mask=index_k_mask,
299 + other=0.0,
300 + ).to(tl.float32)
301 + w_partner = tl.load(
302 + index_k_layer_norm_w_ptr + partner_offs, mask=index_k_mask
303 + ).to(tl.float32)
304 + b_partner = tl.load(
305 + index_k_layer_norm_bias_ptr + partner_offs, mask=index_k_mask
306 + ).to(tl.float32)
307 + normed_partner = (raw_partner - mean) * rstd * w_partner + b_partner
308 + roped = normed * cos_full + sign * normed_partner * sin_full
309 + result = tl.where(in_rope, roped, normed)
310 +
311 + # 3. FP8 quantize + cache write from registers.
312 + # No need to write back to index_k_ptr — the only consumer
313 + # (sparse_attn_indexer) reads from the cache, not index_k.
314 + _fp8_quant_and_cache_write(
315 + result,
316 + index_k_mask,
317 + slot_idx,
318 + indexer_cache_ptr,
319 + indexer_cache_scale_ptr,
320 + indexer_cache_block_size,
321 + indexer_cache_stride,
322 + index_k_block,
323 + INDEX_K_DIM,
324 + )
325 +
326 +
327 +def fused_norm_rope(
328 + positions: torch.Tensor,
329 + q_c: torch.Tensor,
330 + q_rms_norm_w: torch.Tensor,
331 + q_rms_eps: float,
332 + kv_c: torch.Tensor,
333 + kv_rms_norm_w: torch.Tensor,
334 + kv_rms_eps: float,
335 + k_pe: torch.Tensor,
336 + k_rope_cos_sin_cache: torch.Tensor,
337 + index_k: torch.Tensor | None,
338 + index_k_layer_norm_w: torch.Tensor | None,
339 + index_k_layer_norm_bias: torch.Tensor | None,
340 + index_k_layer_norm_eps: float,
341 + index_k_rope_cos_sin_cache: torch.Tensor | None,
342 + topk_indices_buffer: torch.Tensor,
343 + # Cache params for fused writes (single slot_mapping for both caches)
344 + slot_mapping: torch.Tensor | None = None,
345 + indexer_k_cache: torch.Tensor | None = None,
346 + mla_kv_cache: torch.Tensor | None = None,
347 + mla_kv_cache_dtype: str = "auto",
348 + mla_k_scale: torch.Tensor | None = None,
349 + has_indexer: bool = True,
350 + index_rope_interleave: bool = False,
351 + q_c_out: torch.Tensor | None = None,
352 +) -> torch.Tensor:
353 + assert positions.ndim == 1
354 + assert q_c.ndim == 2
355 + assert kv_c.ndim == 2
356 + assert k_pe.ndim == 2
357 + assert topk_indices_buffer.ndim == 2
358 +
359 + num_tokens = positions.shape[0]
360 + q_dim = q_c.shape[-1]
361 + kv_dim = kv_c.shape[-1]
362 + device = positions.device
363 +
364 + # Shared (no-indexer) layers: substitute cached 1-element dummies so the
365 + # kernel launches cleanly; pid 0/3 (indexer + topk fill) skipped by
366 + # HAS_INDEXER and never dereference them.
367 + if not has_indexer:
368 + indexer_k_cache = None
369 + index_k = _dummy((1, 1), q_c.dtype, device)
370 + index_k_layer_norm_w = _dummy((1,), torch.float32, device)
371 + index_k_layer_norm_bias = _dummy((1,), torch.float32, device)
372 + index_k_rope_cos_sin_cache = k_rope_cos_sin_cache
373 + assert index_k is not None
374 + assert index_k_rope_cos_sin_cache is not None
375 + index_k_dim = index_k.shape[-1]
376 + topk = topk_indices_buffer.shape[-1]
377 +
378 + # --- Indexer K cache setup ---
379 + if indexer_k_cache is not None:
380 + assert slot_mapping is not None
381 + idx_cache_scale_view = indexer_k_cache.view(torch.uint8).view(torch.float32)
382 + idx_cache_block_size = indexer_k_cache.shape[1]
383 + idx_cache_stride = indexer_k_cache.shape[2]
384 + if indexer_k_cache.dtype == torch.uint8:
385 + indexer_k_cache = indexer_k_cache.view(torch.float8_e4m3fn)
386 + else:
387 + # No indexer cache (shared layer / MLA-only fusion). Use dummies but
388 + # KEEP the caller's slot_mapping so the MLA write (pid 1) still runs.
389 + idx_cache_scale_view = torch.empty(0, dtype=torch.float32, device=device)
390 + indexer_k_cache = torch.empty(0, dtype=torch.float8_e4m3fn, device=device)
391 + idx_cache_block_size = 1
392 + idx_cache_stride = 1
393 + if mla_kv_cache is None:
394 + # Pure profiling run (no caches at all): skip all per-token writes.
395 + slot_mapping = torch.full(
396 + (num_tokens,), -1, dtype=torch.int64, device=device
397 + )
398 +
399 + # --- MLA KV cache setup ---
400 + mla_cache_fp8 = mla_kv_cache_dtype != "auto"
401 + if mla_kv_cache is not None:
402 + mla_block_stride = mla_kv_cache.stride(0)
403 + mla_entry_stride = mla_kv_cache.stride(1)
404 + if mla_cache_fp8 and mla_kv_cache.dtype == torch.uint8:
405 + mla_kv_cache = mla_kv_cache.view(torch.float8_e4m3fn)
406 + if mla_k_scale is None:
407 + mla_k_scale = torch.ones(1, dtype=torch.float32, device=device)
408 + else:
409 + # Dummy values — pid 2 will skip the MLA cache write because
410 + # slot_mapping is all -1.
411 + mla_kv_cache = torch.empty(0, dtype=torch.bfloat16, device=device)
412 + mla_block_stride = 0
413 + mla_entry_stride = 0
414 + mla_k_scale = torch.ones(1, dtype=torch.float32, device=device)
415 +
416 + if q_c_out is None:
417 + q_c_out = torch.empty_like(q_c)
418 + _fused_norm_rope_kernel[(4, num_tokens)](
419 + positions,
420 + # Q RMS norm
421 + q_c,
422 + q_c.stride(0),
423 + q_rms_norm_w,
424 + q_rms_eps,
425 + q_c_out,
426 + q_c_out.stride(0),
427 + q_dim,
428 + triton.next_power_of_2(q_dim),
429 + # KV RMS norm
430 + kv_c,
431 + kv_c.stride(0),
432 + kv_rms_norm_w,
433 + kv_rms_eps,
434 + kv_dim,
435 + # KV RoPE
436 + k_pe,
437 + k_pe.stride(0),
438 + k_rope_cos_sin_cache,
439 + k_rope_cos_sin_cache.stride(0),
440 + k_rope_cos_sin_cache.shape[-1] // 2,
441 + # Index K layer norm + RoPE + FP8 quant
442 + index_k,
443 + index_k.stride(0),
444 + index_k_layer_norm_w,
445 + index_k_layer_norm_bias,
446 + index_k_layer_norm_eps,
447 + index_k_dim,
448 + triton.next_power_of_2(index_k_dim),
449 + index_k_rope_cos_sin_cache,
450 + index_k_rope_cos_sin_cache.stride(0),
451 + index_k_rope_cos_sin_cache.shape[-1] // 2,
452 + # Cache params
453 + slot_mapping,
454 + indexer_k_cache,
455 + idx_cache_scale_view,
456 + idx_cache_block_size,
457 + idx_cache_stride,
458 + # MLA KV cache (uses same slot_mapping)
459 + mla_kv_cache,
460 + mla_block_stride,
461 + mla_entry_stride,
462 + mla_cache_fp8,
463 + mla_k_scale,
464 + # Top k indices buffer
465 + topk_indices_buffer,
466 + topk_indices_buffer.stride(0),
467 + topk,
468 + TOPK_BLOCK_SIZE=1024,
469 + HAS_INDEXER=has_indexer,
470 + INDEX_ROPE_INTERLEAVE=index_rope_interleave,
471 + )
472 + return q_c_out
473 +
474 +
475 +@triton.jit
476 +def _fused_q_kernel(
477 + pos_ptr,
478 + # MQA query PE: RoPE + FP8 pack into output tail
479 + q_pe_ptr,
480 + q_pe_stride0,
481 + q_pe_stride1,
482 + NUM_Q_HEADS: tl.constexpr,
483 + q_pe_cos_sin_ptr,
484 + q_pe_cos_sin_stride,
485 + Q_PE_HALF_ROT_DIM: tl.constexpr,
486 + # Index Q RoPE
487 + index_q_ptr,
488 + index_q_stride0,
489 + index_q_stride1,
490 + NUM_INDEX_Q_HEADS: tl.constexpr,
491 + index_q_cos_sin_ptr,
492 + index_q_cos_sin_stride,
493 + INDEX_Q_HALF_ROT_DIM: tl.constexpr,
494 + # Index Q Quantize
495 + index_q_fp8_ptr,
496 + index_q_fp8_stride0,
497 + index_q_fp8_stride1,
498 + INDEX_Q_HEAD_DIM: tl.constexpr,
499 + # MQA query pack: quantize ql_nope and RoPE+quantize q_pe into mqa_q_fp8
500 + ql_nope_ptr,
501 + ql_nope_stride0,
502 + ql_nope_stride1,
503 + mqa_q_fp8_ptr,
504 + mqa_q_fp8_stride0,
505 + mqa_q_fp8_stride1,
506 + q_scale_ptr,
507 + QL_NOPE_DIM: tl.constexpr,
508 + QL_NOPE_BLOCK: tl.constexpr,
509 + # Index weights
510 + index_weights_ptr,
511 + index_weights_stride,
512 + index_weights_softmax_scale,
513 + index_weights_head_scale,
514 + index_weights_out_ptr,
515 + index_weights_out_stride,
516 + HAS_INDEXER: tl.constexpr,
517 + INDEX_ROPE_INTERLEAVE: tl.constexpr,
518 +):
519 + pid = tl.program_id(0)
520 + tok_idx = tl.program_id(1)
521 + head_idx = tl.program_id(2)
522 +
523 + if pid == 2:
524 + # ql_nope quantize + pack into the front of mqa_q_fp8.
525 + if 2 * head_idx >= NUM_Q_HEADS:
526 + return
527 +
528 + scale = tl.load(q_scale_ptr)
529 + for local_head in range(2):
530 + q_head_idx = head_idx * 2 + local_head
531 + if q_head_idx < NUM_Q_HEADS:
532 + ql_nope_off = tl.arange(0, QL_NOPE_BLOCK)
533 + ql_nope_mask = ql_nope_off < QL_NOPE_DIM
534 + ql_nope = tl.load(
535 + ql_nope_ptr
536 + + tok_idx * ql_nope_stride0
537 + + q_head_idx * ql_nope_stride1
538 + + ql_nope_off,
539 + mask=ql_nope_mask,
540 + ).to(tl.float32)
541 + ql_nope_fp8 = (ql_nope / scale).to(tl.float8e4nv)
542 + tl.store(
543 + mqa_q_fp8_ptr
544 + + tok_idx * mqa_q_fp8_stride0
545 + + q_head_idx * mqa_q_fp8_stride1
546 + + ql_nope_off,
547 + ql_nope_fp8,
548 + mask=ql_nope_mask,
549 + )
550 + return
551 + elif pid == 0:
552 + # q_pe RoPE + quantize + pack into the tail of mqa_q_fp8.
553 + if 2 * head_idx >= NUM_Q_HEADS:
554 + return
555 +
556 + pos = tl.load(pos_ptr + tok_idx)
557 + cos, sin = _get_cos_sin(
558 + q_pe_cos_sin_ptr,
559 + q_pe_cos_sin_stride,
560 + pos,
561 + Q_PE_HALF_ROT_DIM,
562 + )
563 +
564 + scale = tl.load(q_scale_ptr)
565 + for local_head in range(2):
566 + q_head_idx = head_idx * 2 + local_head
567 + if q_head_idx < NUM_Q_HEADS:
568 + rot_off = tl.arange(0, Q_PE_HALF_ROT_DIM)
569 + x1 = tl.load(
570 + q_pe_ptr
571 + + tok_idx * q_pe_stride0
572 + + q_head_idx * q_pe_stride1
573 + + rot_off * 2,
574 + ).to(tl.float32)
575 + x2 = tl.load(
576 + q_pe_ptr
577 + + tok_idx * q_pe_stride0
578 + + q_head_idx * q_pe_stride1
579 + + rot_off * 2
580 + + 1
581 + ).to(tl.float32)
582 + r1 = x1 * cos - x2 * sin
583 + r2 = x2 * cos + x1 * sin
584 + tl.store(
585 + mqa_q_fp8_ptr
586 + + tok_idx * mqa_q_fp8_stride0
587 + + q_head_idx * mqa_q_fp8_stride1
588 + + QL_NOPE_DIM
589 + + rot_off * 2,
590 + (r1 / scale).to(tl.float8e4nv),
591 + )
592 + tl.store(
593 + mqa_q_fp8_ptr
594 + + tok_idx * mqa_q_fp8_stride0
595 + + q_head_idx * mqa_q_fp8_stride1
596 + + QL_NOPE_DIM
597 + + rot_off * 2
598 + + 1,
599 + (r2 / scale).to(tl.float8e4nv),
600 + )
601 + return
602 + elif pid == 1:
603 + # Index Q RoPE + fp8 quant, all in registers. The roped bf16 index_q is
604 + # never consumed (only the fp8 below is), so we avoid an in-place
605 + # store-then-reload round-trip.
606 + if not HAS_INDEXER:
607 + return
608 + if head_idx >= NUM_INDEX_Q_HEADS:
609 + return
610 +
611 + pos = tl.load(pos_ptr + tok_idx)
612 + index_q_block = tl.arange(0, INDEX_Q_HEAD_DIM)
613 + iq_base = index_q_ptr + tok_idx * index_q_stride0 + head_idx * index_q_stride1
614 + index_q = tl.load(iq_base + index_q_block).to(tl.float32)
615 +
616 + # RoPE in registers (interleaved for GLM-5.2, NeoX for DeepSeek-V3.2),
617 + # gathering the rotation partner from the read-only input.
618 + in_rope = index_q_block < 2 * INDEX_Q_HALF_ROT_DIM
619 + if INDEX_ROPE_INTERLEAVE:
620 + cos_idx = index_q_block // 2
621 + partner_offs = tl.where(in_rope, index_q_block ^ 1, index_q_block)
622 + sign = tl.where(index_q_block % 2 == 0, -1.0, 1.0)
623 + else:
624 + cos_idx = index_q_block % INDEX_Q_HALF_ROT_DIM
625 + partner_offs = tl.where(
626 + in_rope, index_q_block ^ INDEX_Q_HALF_ROT_DIM, index_q_block
627 + )
628 + sign = tl.where(index_q_block < INDEX_Q_HALF_ROT_DIM, -1.0, 1.0)
629 + cos_full = tl.load(
630 + index_q_cos_sin_ptr + pos * index_q_cos_sin_stride + cos_idx,
631 + mask=in_rope,
632 + other=1.0,
633 + ).to(tl.float32)
634 + sin_full = tl.load(
635 + index_q_cos_sin_ptr
636 + + pos * index_q_cos_sin_stride
637 + + INDEX_Q_HALF_ROT_DIM
638 + + cos_idx,
639 + mask=in_rope,
640 + other=0.0,
641 + ).to(tl.float32)
642 + partner = tl.load(iq_base + partner_offs).to(tl.float32)
643 + roped = index_q * cos_full + sign * partner * sin_full
644 + index_q = tl.where(in_rope, roped, index_q)
645 +
646 + # Index Q Quantize (from registers)
647 + index_q_fp8, index_q_scale = _fp8_ue8m0_quantize(index_q)
648 + tl.store(
649 + index_q_fp8_ptr
650 + + tok_idx * index_q_fp8_stride0
651 + + head_idx * index_q_fp8_stride1
652 + + index_q_block,
653 + index_q_fp8,
654 + )
655 +
656 + # Index weights update
657 + index_weights = tl.load(
658 + index_weights_ptr + tok_idx * index_weights_stride + head_idx
659 + )
660 + index_weights = index_weights.to(tl.float32)
661 + index_weights *= index_q_scale
662 + index_weights *= index_weights_softmax_scale
663 + index_weights *= index_weights_head_scale
664 + tl.store(
665 + index_weights_out_ptr + tok_idx * index_weights_out_stride + head_idx,
666 + index_weights,
667 + )
668 +
669 +
670 +def fused_q(
671 + positions: torch.Tensor,
672 + q_pe: torch.Tensor,
673 + q_pe_cos_sin_cache: torch.Tensor,
674 + index_q: torch.Tensor | None,
675 + index_q_cos_sin_cache: torch.Tensor | None,
676 + ql_nope: torch.Tensor,
677 + q_scale: torch.Tensor,
678 + # Index weights
679 + index_weights: torch.Tensor | None,
680 + index_weights_softmax_scale: float,
681 + index_weights_head_scale: float,
682 + has_indexer: bool = True,
683 + index_rope_interleave: bool = False,
684 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
685 + assert positions.ndim == 1
686 + assert q_pe.ndim == 3
687 + assert q_pe_cos_sin_cache.ndim == 2
688 + assert ql_nope.ndim == 3
689 + assert ql_nope.shape[:2] == q_pe.shape[:2]
690 +
691 + num_tokens = positions.shape[0]
692 + num_q_heads = q_pe.shape[1]
693 + # Grid's 3rd dim must cover the MQA-pack heads (pid 0/2 iterate 2 heads
694 + # each) and, when present, the indexer heads (pid 1).
695 + mqa_grid_heads = (num_q_heads + 1) // 2
696 + if not has_indexer:
697 + # Shared layer: cached 1-element dummies; pid 1 skipped by HAS_INDEXER
698 + # and never dereferences them.
699 + index_q = _dummy((1, 1, 1), q_pe.dtype, q_pe.device)
700 + index_q_cos_sin_cache = q_pe_cos_sin_cache
701 + index_weights = _dummy((1, 1), torch.float32, q_pe.device)
702 + assert index_q is not None and index_q.ndim == 3
703 + assert index_q_cos_sin_cache is not None
704 + assert index_weights is not None
705 + num_index_q_heads = index_q.shape[1]
706 + index_q_head_dim = index_q.shape[2]
707 + grid_heads = max(mqa_grid_heads, num_index_q_heads)
708 + mqa_q_fp8 = torch.empty(
709 + q_pe.shape[0],
710 + q_pe.shape[1],
711 + ql_nope.shape[2] + q_pe.shape[2],
712 + dtype=torch.float8_e4m3fn,
713 + device=q_pe.device,
714 + )
715 +
716 + index_q_fp8 = torch.empty_like(index_q, dtype=torch.float8_e4m3fn)
717 + index_weights_out = torch.empty_like(index_weights, dtype=torch.float32)
718 + _fused_q_kernel[(3, num_tokens, grid_heads)](
719 + positions,
720 + q_pe,
721 + q_pe.stride(0),
722 + q_pe.stride(1),
723 + num_q_heads,
724 + q_pe_cos_sin_cache,
725 + q_pe_cos_sin_cache.stride(0),
726 + q_pe_cos_sin_cache.shape[-1] // 2,
727 + index_q,
728 + index_q.stride(0),
729 + index_q.stride(1),
730 + num_index_q_heads,
731 + index_q_cos_sin_cache,
732 + index_q_cos_sin_cache.stride(0),
733 + index_q_cos_sin_cache.shape[-1] // 2,
734 + index_q_fp8,
735 + index_q_fp8.stride(0),
736 + index_q_fp8.stride(1),
737 + index_q_head_dim,
738 + ql_nope,
739 + ql_nope.stride(0),
740 + ql_nope.stride(1),
741 + mqa_q_fp8,
742 + mqa_q_fp8.stride(0),
743 + mqa_q_fp8.stride(1),
744 + q_scale,
745 + ql_nope.shape[2],
746 + triton.next_power_of_2(ql_nope.shape[2]),
747 + index_weights,
748 + index_weights.stride(0),
749 + index_weights_softmax_scale,
750 + index_weights_head_scale,
751 + index_weights_out,
752 + index_weights_out.stride(0),
753 + HAS_INDEXER=has_indexer,
754 + INDEX_ROPE_INTERLEAVE=index_rope_interleave,
755 + # num_warps=1 is optimal here: each program is a single 128-element
756 + # rope+quant, so the kernel is program-count/occupancy bound, not
757 + # per-program compute bound (swept 1/2/4/8 — 1 wins or ties everywhere).
758 + num_warps=1,
759 + )
760 + return index_q_fp8, index_weights_out, mqa_q_fp8
761 +
762 +
763 +@triton.jit
764 +def _fused_eh_norm_kernel(
765 + pos_ptr,
766 + embeds_ptr,
767 + embeds_stride,
768 + prev_ptr,
769 + prev_stride,
770 + enorm_w_ptr,
771 + hnorm_w_ptr,
772 + eps,
773 + out_ptr,
774 + out_stride,
775 + H: tl.constexpr,
776 + BLOCK: tl.constexpr,
777 +):
778 + """MTP input fusion: zero embeds at position 0, RMSNorm(embeds) with enorm
779 + and RMSNorm(prev_hidden) with hnorm, written side-by-side into ``out``
780 + ([N, 2H]) ready for the eh_proj GEMM. Replaces where + 2x RMSNorm + cat."""
781 + tok = tl.program_id(0)
782 + off = tl.arange(0, BLOCK)
783 + mask = off < H
784 +
785 + pos = tl.load(pos_ptr + tok)
786 + e = tl.load(embeds_ptr + tok * embeds_stride + off, mask=mask, other=0.0)
787 + e = tl.where(pos == 0, 0.0, e.to(tl.float32))
788 + ew = tl.load(enorm_w_ptr + off, mask=mask)
789 + e_normed = _rms_norm(e, ew, eps, H)
790 + tl.store(out_ptr + tok * out_stride + off, e_normed, mask=mask)
791 +
792 + p = tl.load(prev_ptr + tok * prev_stride + off, mask=mask, other=0.0)
793 + hw = tl.load(hnorm_w_ptr + off, mask=mask)
794 + p_normed = _rms_norm(p, hw, eps, H)
795 + tl.store(out_ptr + tok * out_stride + H + off, p_normed, mask=mask)
796 +
797 +
798 +def fused_eh_norm(
799 + positions: torch.Tensor,
800 + inputs_embeds: torch.Tensor,
801 + previous_hidden: torch.Tensor,
802 + enorm_w: torch.Tensor,
803 + hnorm_w: torch.Tensor,
804 + eps: float,
805 +) -> torch.Tensor:
806 + """Returns cat([enorm(masked embeds), hnorm(prev_hidden)]) -> [N, 2H]."""
807 + n, h = inputs_embeds.shape
808 + out = torch.empty(n, 2 * h, dtype=inputs_embeds.dtype, device=inputs_embeds.device)
809 + _fused_eh_norm_kernel[(n,)](
810 + positions,
811 + inputs_embeds,
812 + inputs_embeds.stride(0),
813 + previous_hidden,
814 + previous_hidden.stride(0),
815 + enorm_w,
816 + hnorm_w,
817 + eps,
818 + out,
819 + out.stride(0),
820 + h,
821 + triton.next_power_of_2(h),
822 + )
823 + return out
vllm/models/deepseek_v32/nvidia/model.py
+21 -3
@@ -36,6 +36,7 @@ from vllm.model_executor.models.utils import (
36 from vllm.sequence import IntermediateTensors
37
38 from .attention import DeepseekV32Attention
39 +from .fused_ops import fused_allreduce_rms_norm
40
41
42 class DeepseekV32DecoderLayer(torch.nn.Module):
@@ -77,6 +78,10 @@ class DeepseekV32DecoderLayer(torch.nn.Module):
78 quant_config=quant_config,
79 prefix=f"{prefix}.mlp",
80 )
81 + # Defer the MoE cross-rank all-reduce; it is fused into the next
82 + # layer's input_layernorm (or the final norm) via
83 + # fused_allreduce_rms_norm. self.mlp.experts is the MoERunner.
84 + self.mlp.experts.moe_config.skip_final_all_reduce = True
85 else:
86 self.mlp = DeepseekV2MLP(
87 hidden_size=config.hidden_size,
@@ -84,6 +89,7 @@ class DeepseekV32DecoderLayer(torch.nn.Module):
89 hidden_act=config.hidden_act,
90 quant_config=quant_config,
91 prefix=f"{prefix}.mlp",
92 + reduce_results=False,
93 )
94 self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
95 self.post_attention_layernorm = RMSNorm(
@@ -98,12 +104,23 @@ class DeepseekV32DecoderLayer(torch.nn.Module):
104 residual: torch.Tensor | None,
105 ) -> tuple[torch.Tensor, torch.Tensor]:
106 if residual is None:
107 + # First layer: hidden_states is the (already reduced) embedding.
108 residual = hidden_states
109 hidden_states = self.input_layernorm(hidden_states)
110 else:
104 - hidden_states, residual = self.input_layernorm(hidden_states, residual)
111 + # The previous layer's MLP/MoE output is left un-reduced; fuse its
112 + # all-reduce into this input_layernorm.
113 + hidden_states, residual = fused_allreduce_rms_norm(
114 + hidden_states, residual, self.input_layernorm
115 + )
116 + # self_attn's o_proj runs reduce_results=False; fuse its all-reduce with
117 + # the post-attention RMSNorm.
118 hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states)
106 - hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
119 + hidden_states, residual = fused_allreduce_rms_norm(
120 + hidden_states, residual, self.post_attention_layernorm
121 + )
122 + # MLP/MoE runs un-reduced; its all-reduce is fused into the next layer's
123 + # input_layernorm (or the model's final norm).
124 hidden_states = self.mlp(hidden_states)
125 return hidden_states, residual
126
@@ -201,7 +218,8 @@ class DeepseekV32Model(torch.nn.Module):
218 {"hidden_states": hidden_states, "residual": residual}
219 )
220
204 - hidden_states, _ = self.norm(hidden_states, residual)
221 + # Last layer's MoE output is un-reduced; fuse its all-reduce into norm.
222 + hidden_states, _ = fused_allreduce_rms_norm(hidden_states, residual, self.norm)
223 if len(aux_hidden_states) > 0:
224 return hidden_states, aux_hidden_states
225 return hidden_states
vllm/models/deepseek_v32/nvidia/mtp.py
+15 -6
@@ -1,6 +1,5 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
3 import typing
4 from collections.abc import Callable, Iterable
5
@@ -9,6 +8,7 @@ import torch.nn as nn
8
9 from vllm._aiter_ops import rocm_aiter_ops
10 from vllm.config import VllmConfig
11 +from vllm.distributed import tensor_model_parallel_all_reduce
12 from vllm.model_executor.layers.fused_moe import (
13 fused_moe_make_expert_params_mapping,
14 )
@@ -35,6 +35,7 @@ from vllm.model_executor.models.utils import (
35 from vllm.platforms import current_platform
36 from vllm.sequence import IntermediateTensors
37
38 +from .kernels import fused_eh_norm
39 from .model import DeepseekV32DecoderLayer
40
41
@@ -75,15 +76,23 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
76 spec_step_index: int = 0,
77 ) -> torch.Tensor:
78 assert inputs_embeds is not None
78 - inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds)
79 - inputs_embeds = self.enorm(inputs_embeds)
80 - previous_hidden_states = self.hnorm(previous_hidden_states)
81 - hidden_states = self.eh_proj(
82 - torch.cat([inputs_embeds, previous_hidden_states], dim=-1)
79 + # Fused: zero pos-0 embeds + enorm(embeds) + hnorm(prev) + cat -> [N, 2H].
80 + eh_input = fused_eh_norm(
81 + positions,
82 + inputs_embeds,
83 + previous_hidden_states,
84 + self.enorm.weight,
85 + self.hnorm.weight,
86 + self.enorm.variance_epsilon,
87 )
88 + hidden_states = self.eh_proj(eh_input)
89 hidden_states, residual = self.mtp_block(
90 positions=positions, hidden_states=hidden_states, residual=None
91 )
92 + # mtp_block's MoE output is left un-reduced (skip_final_all_reduce); the
93 + # main model fuses that all-reduce into the next norm, but here the
94 + # recycle hidden is consumed directly, so reduce it now.
95 + hidden_states = tensor_model_parallel_all_reduce(hidden_states)
96 # Return the pre-final-norm recycle hidden (re-fed as the next spec
97 # step's previous_hidden_states); shared_head norm is applied in
98 # compute_logits. Matches the V2-runner / deepseek_v4 MTP contract.