[ROCm][Model][Bugfix] Enable GLM-5.2-MXFP4 on the deepseek_v32 path and fix sparse attention correctness (#51915)
Signed-off-by: Jack Hu <Jack.Hu@amd.com> Signed-off-by: Jack Hu <jack.hu@amd.com> Signed-off-by: Douglas Lehr <Doug.Lehr@amd.com> Co-authored-by: James E T Smith <jamesETsmith@users.noreply.github.com> Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> Co-authored-by: Douglas Lehr <Doug.Lehr@amd.com>
Jack Hu committed
Sep 22, 2026 at 19:08 UTC
1c0eee919db320ea3e2df37f800465c7842a650a
8 files changed
+319
-212
tests/kernels/test_fused_deepseek_v32_norm_rope.py
+187
-151
@@ -24,6 +24,8 @@ outputs are checked within 1 representable-step (ULP); bf16 norm/RoPE outputs us
24
rtol/atol=1e-2 (the tolerance the sibling deepseek_v4 fused-kernel test uses).
25
"""
26
27
+from typing import NamedTuple
28
+
29
import pytest
30
import torch
31
@@ -33,19 +35,30 @@ from vllm.platforms import current_platform
35
FP8 = torch.float8_e4m3fn
36
FP8_MAX = 448.0
37
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
38
EPS = 1e-6
39
40
+
41
+class ModelConfig(NamedTuple):
42
+ """Per-model attention dims (num_heads is the TP8 local count)."""
43
+
44
+ q_lora: int
45
+ kv_lora: int
46
+ rope_dim: int
47
+ num_heads: int
48
+ index_heads: int
49
+ index_head_dim: int
50
+ hidden: int
51
+
52
+
53
+MODEL_CONFIGS = [
54
+ ModelConfig(2048, 512, 64, 8, 32, 128, 6144), # GLM-5.2
55
+ ModelConfig(1536, 512, 64, 16, 64, 128, 7168), # DeepSeek-V3.2
56
+]
57
+MODEL_IDS = ["glm5.2", "dsv3.2"]
58
+
59
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+)",
60
+ not current_platform.supports_fp8(),
61
+ reason="deepseek_v32 fused kernels require fp8 support",
62
)
63
64
@@ -144,29 +157,32 @@ def assert_fp8(got: torch.Tensor, ref: torch.Tensor, msg: str):
157
# ── fused_norm_rope ──────────────────────────────────────────────────────────
158
159
160
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
161
@pytest.mark.parametrize("num_tokens", [1, 4, 17, 512, 4096])
162
@pytest.mark.parametrize("index_interleave", [True, False])
163
@pytest.mark.parametrize("mla_dtype", ["auto", "bfloat16", "fp8"])
150
-def test_fused_norm_rope(num_tokens: int, index_interleave: bool, mla_dtype: str):
164
+def test_fused_norm_rope(
165
+ num_tokens: int, index_interleave: bool, mla_dtype: str, cfg: ModelConfig
166
+):
167
torch.manual_seed(0)
168
dev = "cuda"
169
max_pos = 8192
170
pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos
171
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)
172
+ q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16)
173
+ kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16)
174
+ k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16)
175
+ qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16)
176
+ kvw = torch.randn(cfg.kv_lora, device=dev, dtype=torch.bfloat16)
177
+ ik = torch.randn(num_tokens, cfg.index_head_dim, device=dev, dtype=torch.bfloat16)
178
+ ikw = torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32)
179
+ ikb = torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32)
180
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)
181
+ mla_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev) # MLA k_pe: interleaved
182
+ idx_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
183
184
bs = max_pos # single block covering all tokens
169
- mla_dim = KV_LORA + ROPE_DIM
185
+ mla_dim = cfg.kv_lora + cfg.rope_dim
186
mla_fp8 = mla_dtype == "fp8"
187
if mla_fp8:
188
mla_cache = torch.zeros(1, bs, mla_dim, device=dev, dtype=torch.uint8)
@@ -174,7 +190,7 @@ def test_fused_norm_rope(num_tokens: int, index_interleave: bool, mla_dtype: str
190
else:
191
mla_cache = torch.zeros(1, bs, mla_dim, device=dev, dtype=torch.bfloat16)
192
mla_k_scale = None
177
- idx_row = INDEX_HEAD_DIM + INDEX_HEAD_DIM // 128 * 4 # 132
193
+ idx_row = cfg.index_head_dim + cfg.index_head_dim // 128 * 4 # 132
194
idx_cache = torch.zeros(1, bs, idx_row, device=dev, dtype=torch.uint8)
195
slot = torch.arange(num_tokens, device=dev, dtype=torch.int64)
196
topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32)
@@ -213,20 +229,20 @@ def test_fused_norm_rope(num_tokens: int, index_interleave: bool, mla_dtype: str
229
if mla_fp8:
230
cache = mla_cache.view(FP8)[0, :num_tokens]
231
s = mla_k_scale.item()
216
- assert_fp8(cache[:, :KV_LORA], (kv_ref / s).to(FP8), "MLA kv fp8")
217
- assert_fp8(cache[:, KV_LORA:], (kpe_ref / s).to(FP8), "MLA k_pe fp8")
232
+ assert_fp8(cache[:, : cfg.kv_lora], (kv_ref / s).to(FP8), "MLA kv fp8")
233
+ assert_fp8(cache[:, cfg.kv_lora :], (kpe_ref / s).to(FP8), "MLA k_pe fp8")
234
else:
235
cache = mla_cache[0, :num_tokens]
220
- assert_bf16(cache[:, :KV_LORA], kv_ref, "MLA kv bf16")
221
- assert_bf16(cache[:, KV_LORA:], kpe_ref, "MLA k_pe bf16")
236
+ assert_bf16(cache[:, : cfg.kv_lora], kv_ref, "MLA kv bf16")
237
+ assert_bf16(cache[:, cfg.kv_lora :], kpe_ref, "MLA k_pe bf16")
238
239
# Indexer-K cache (packed [bs*head_dim fp8 | bs*4 fp32 scale]).
240
ik_ref = layer_norm(ik, ikw, ikb)
241
ik_ref = rope(ik_ref, pos, idx_cos_sin, interleave=index_interleave)
242
q_ref, s_ref = ue8m0_quant(ik_ref)
243
flat = idx_cache[0].reshape(-1)
228
- vals = flat[: bs * INDEX_HEAD_DIM].view(FP8).reshape(bs, INDEX_HEAD_DIM)
229
- scales = flat[bs * INDEX_HEAD_DIM :].view(torch.float32)
244
+ vals = flat[: bs * cfg.index_head_dim].view(FP8).reshape(bs, cfg.index_head_dim)
245
+ scales = flat[bs * cfg.index_head_dim :].view(torch.float32)
246
assert_fp8(vals[:num_tokens], q_ref, "indexer-K fp8")
247
torch.testing.assert_close(scales[:num_tokens], s_ref, rtol=0, atol=0)
248
@@ -234,23 +250,24 @@ def test_fused_norm_rope(num_tokens: int, index_interleave: bool, mla_dtype: str
250
assert (topk == -1).all(), "topk buffer not cleared on indexer layer"
251
252
237
-def test_fused_norm_rope_packed_indexer_block_stride():
253
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
254
+def test_fused_norm_rope_packed_indexer_block_stride(cfg: ModelConfig):
255
"""Indexer writes use their HMA block stride and independent slot map."""
256
torch.manual_seed(7)
257
dev = "cuda"
258
num_tokens = block_size = 2
259
pos = torch.arange(num_tokens, device=dev, dtype=torch.int64)
243
- q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
244
- kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
245
- k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
246
- qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
247
- kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)
248
- ik = torch.randn(num_tokens, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16)
249
- ikw = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
250
- ikb = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
251
- cos_sin = make_cos_sin(32, ROPE_DIM, dev)
252
-
253
- idx_row = INDEX_HEAD_DIM + INDEX_HEAD_DIM // 128 * 4
260
+ q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16)
261
+ kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16)
262
+ k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16)
263
+ qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16)
264
+ kvw = torch.randn(cfg.kv_lora, device=dev, dtype=torch.bfloat16)
265
+ ik = torch.randn(num_tokens, cfg.index_head_dim, device=dev, dtype=torch.bfloat16)
266
+ ikw = torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32)
267
+ ikb = torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32)
268
+ cos_sin = make_cos_sin(32, cfg.rope_dim, dev)
269
+
270
+ idx_row = cfg.index_head_dim + cfg.index_head_dim // 128 * 4
271
packed_block_stride = block_size * idx_row + 64
272
backing = torch.zeros(2 * packed_block_stride, device=dev, dtype=torch.uint8)
273
idx_cache = torch.as_strided(
@@ -261,7 +278,7 @@ def test_fused_norm_rope_packed_indexer_block_stride():
278
mla_cache = torch.zeros(
279
1,
280
block_size,
264
- KV_LORA + ROPE_DIM,
281
+ cfg.kv_lora + cfg.rope_dim,
282
device=dev,
283
dtype=torch.bfloat16,
284
)
@@ -296,15 +313,20 @@ def test_fused_norm_rope_packed_indexer_block_stride():
313
q_ref, s_ref = ue8m0_quant(ik_ref)
314
packed = idx_cache[1].reshape(-1)
315
values = (
299
- packed[: block_size * INDEX_HEAD_DIM].view(FP8).view(block_size, INDEX_HEAD_DIM)
316
+ packed[: block_size * cfg.index_head_dim]
317
+ .view(FP8)
318
+ .view(block_size, cfg.index_head_dim)
319
)
301
- scales = packed[block_size * INDEX_HEAD_DIM :].view(torch.float32)
320
+ scales = packed[block_size * cfg.index_head_dim :].view(torch.float32)
321
assert_fp8(values, q_ref, "packed indexer-K fp8")
322
torch.testing.assert_close(scales, s_ref, rtol=0, atol=0)
323
assert (backing[block_size * idx_row : packed_block_stride] == 0).all()
324
325
307
-def test_fused_norm_rope_materializes_qk_without_local_cache_slots():
326
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
327
+def test_fused_norm_rope_materializes_qk_without_local_cache_slots(
328
+ cfg: ModelConfig,
329
+):
330
"""DCP non-owner ranks still need valid query shards for query AllGather."""
331
torch.manual_seed(7)
332
dev = "cuda"
@@ -312,20 +334,20 @@ def test_fused_norm_rope_materializes_qk_without_local_cache_slots():
334
max_pos = 16
335
pos = torch.arange(num_tokens, device=dev, dtype=torch.int64)
336
315
- q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
316
- kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
317
- k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
318
- qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
319
- kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)
320
- ik = torch.randn(num_tokens, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16)
321
- ikw = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
322
- ikb = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
323
- cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
337
+ q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16)
338
+ kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16)
339
+ k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16)
340
+ qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16)
341
+ kvw = torch.randn(cfg.kv_lora, device=dev, dtype=torch.bfloat16)
342
+ ik = torch.randn(num_tokens, cfg.index_head_dim, device=dev, dtype=torch.bfloat16)
343
+ ikw = torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32)
344
+ ikb = torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32)
345
+ cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
346
347
mla_cache = torch.zeros(
326
- 1, max_pos, KV_LORA + ROPE_DIM, device=dev, dtype=torch.bfloat16
348
+ 1, max_pos, cfg.kv_lora + cfg.rope_dim, device=dev, dtype=torch.bfloat16
349
)
328
- idx_row = INDEX_HEAD_DIM + INDEX_HEAD_DIM // 128 * 4
350
+ idx_row = cfg.index_head_dim + cfg.index_head_dim // 128 * 4
351
idx_cache = torch.zeros(1, max_pos, idx_row, device=dev, dtype=torch.uint8)
352
no_local_slots = torch.full((num_tokens,), -1, device=dev, dtype=torch.int64)
353
topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32)
@@ -371,23 +393,26 @@ def test_fused_norm_rope_materializes_qk_without_local_cache_slots():
393
assert (topk == -1).all(), "topk buffer not cleared on non-owner rank"
394
395
396
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
397
@pytest.mark.parametrize("num_tokens", [1, 17, 512])
375
-def test_fused_norm_rope_no_indexer(num_tokens: int):
398
+def test_fused_norm_rope_no_indexer(num_tokens: int, cfg: ModelConfig):
399
"""Shared (no-indexer) layer: q + kv/MLA only; top-k buffer untouched."""
400
torch.manual_seed(1)
401
dev = "cuda"
402
max_pos = 8192
403
pos = torch.arange(num_tokens, device=dev, dtype=torch.int64)
404
382
- q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
383
- kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
384
- k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
385
- qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
386
- kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)
387
- mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
405
+ q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16)
406
+ kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16)
407
+ k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16)
408
+ qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16)
409
+ kvw = torch.randn(cfg.kv_lora, device=dev, dtype=torch.bfloat16)
410
+ mla_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
411
412
bs = max_pos
390
- mla_cache = torch.zeros(1, bs, KV_LORA + ROPE_DIM, device=dev, dtype=torch.bfloat16)
413
+ mla_cache = torch.zeros(
414
+ 1, bs, cfg.kv_lora + cfg.rope_dim, device=dev, dtype=torch.bfloat16
415
+ )
416
slot = torch.arange(num_tokens, device=dev, dtype=torch.int64)
417
topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32)
418
@@ -418,9 +443,9 @@ def test_fused_norm_rope_no_indexer(num_tokens: int):
443
444
assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm (no-indexer)")
445
cache = mla_cache[0, :num_tokens]
421
- assert_bf16(cache[:, :KV_LORA], rms_norm(kv_c, kvw), "MLA kv (no-indexer)")
446
+ assert_bf16(cache[:, : cfg.kv_lora], rms_norm(kv_c, kvw), "MLA kv (no-indexer)")
447
assert_bf16(
423
- cache[:, KV_LORA:],
448
+ cache[:, cfg.kv_lora :],
449
rope(k_pe.float(), pos, mla_cos_sin, interleave=True),
450
"MLA k_pe (no-indexer)",
451
)
@@ -428,36 +453,39 @@ def test_fused_norm_rope_no_indexer(num_tokens: int):
453
assert (topk == 7).all(), "topk buffer should be untouched on shared layer"
454
455
456
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
457
@pytest.mark.parametrize("has_indexer", [False, True])
432
-def test_fused_norm_rope_materializes_pcp_cache_inputs(has_indexer: bool):
458
+def test_fused_norm_rope_materializes_pcp_cache_inputs(
459
+ has_indexer: bool, cfg: ModelConfig
460
+):
461
"""PCP gets local normalized/rotated K rows without direct cache writes."""
462
torch.manual_seed(6)
463
dev = "cuda"
464
num_tokens = 17
465
max_pos = 8192
466
pos = torch.arange(num_tokens, device=dev, dtype=torch.int64)
439
- q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
440
- kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
441
- k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
442
- qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
443
- kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)
467
+ q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16)
468
+ kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16)
469
+ k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16)
470
+ qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16)
471
+ kvw = torch.randn(cfg.kv_lora, device=dev, dtype=torch.bfloat16)
472
ik = (
445
- torch.randn(num_tokens, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16)
473
+ torch.randn(num_tokens, cfg.index_head_dim, device=dev, dtype=torch.bfloat16)
474
if has_indexer
475
else None
476
)
477
ikw = (
450
- torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
478
+ torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32)
479
if has_indexer
480
else None
481
)
482
ikb = (
455
- torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
483
+ torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32)
484
if has_indexer
485
else None
486
)
459
- mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
460
- idx_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev) if has_indexer else None
487
+ mla_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
488
+ idx_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev) if has_indexer else None
489
q_out = torch.empty_like(q_c)
490
kv_out = torch.empty_like(kv_c)
491
kpe_out = torch.empty_like(k_pe)
@@ -503,8 +531,9 @@ def test_fused_norm_rope_materializes_pcp_cache_inputs(has_indexer: bool):
531
assert_bf16(ik_out, ik_ref, "PCP indexer-K")
532
533
534
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
535
@pytest.mark.parametrize("num_tokens", [1, 4, 17, 512])
507
-def test_fused_norm_rope_ds_mla(num_tokens: int):
536
+def test_fused_norm_rope_ds_mla(num_tokens: int, cfg: ModelConfig):
537
"""fp8_ds_mla MLA cache layout (FlashMLA sparse, bf16-query path; SM90/SM100).
538
539
Per-token 656-byte entry: 512 fp8 NoPE (4 per-128 tiles, power-of-two scales
@@ -515,12 +544,12 @@ def test_fused_norm_rope_ds_mla(num_tokens: int):
544
max_pos = 8192
545
pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos
546
518
- q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
519
- kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
520
- k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
521
- qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
522
- kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)
523
- mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
547
+ q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16)
548
+ kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16)
549
+ k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16)
550
+ qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16)
551
+ kvw = torch.randn(cfg.kv_lora, device=dev, dtype=torch.bfloat16)
552
+ mla_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
553
554
bs = max_pos
555
mla_cache = torch.zeros(1, bs, 656, device=dev, dtype=torch.uint8)
@@ -559,13 +588,13 @@ def test_fused_norm_rope_ds_mla(num_tokens: int):
588
tiles = kv_ref.view(num_tokens, 4, 128)
589
ref_scale = torch.clamp(tiles.abs().amax(dim=-1) / FP8_MAX, min=1e-4)
590
ref_scale = torch.exp2(torch.ceil(torch.log2(ref_scale)))
562
- ref_nope = (tiles / ref_scale[..., None]).reshape(num_tokens, KV_LORA).to(FP8)
591
+ ref_nope = (tiles / ref_scale[..., None]).reshape(num_tokens, cfg.kv_lora).to(FP8)
592
593
cache = mla_cache[0, :num_tokens] # [N, 656] uint8
565
- nope = cache[:, :KV_LORA].view(FP8)
566
- scales = cache.view(torch.float32)[:, KV_LORA // 4 : KV_LORA // 4 + 4]
567
- rope_off = KV_LORA // 2 + 8
568
- rope_vals = cache.view(torch.bfloat16)[:, rope_off : rope_off + ROPE_DIM]
594
+ nope = cache[:, : cfg.kv_lora].view(FP8)
595
+ scales = cache.view(torch.float32)[:, cfg.kv_lora // 4 : cfg.kv_lora // 4 + 4]
596
+ rope_off = cfg.kv_lora // 2 + 8
597
+ rope_vals = cache.view(torch.bfloat16)[:, rope_off : rope_off + cfg.rope_dim]
598
599
torch.testing.assert_close(scales, ref_scale, rtol=0, atol=0)
600
assert_fp8(nope, ref_nope, "ds_mla NoPE fp8")
@@ -598,8 +627,9 @@ def nvfp4_sf_byte(s: torch.Tensor) -> torch.Tensor:
627
not current_platform.is_device_capability(100),
628
reason="nvfp4_ds_mla requires SM100 (Blackwell)",
629
)
630
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
631
@pytest.mark.parametrize("num_tokens", [1, 4, 17, 512])
602
-def test_fused_norm_rope_nvfp4_ds_mla(num_tokens: int):
632
+def test_fused_norm_rope_nvfp4_ds_mla(num_tokens: int, cfg: ModelConfig):
633
"""nvfp4_ds_mla MLA cache layout (FlashMLA sparse, SM100 only).
634
635
Per-token 352-byte entry: 256 B of 512 e2m1 NoPE packed 2/byte (low nibble
@@ -611,8 +641,8 @@ def test_fused_norm_rope_nvfp4_ds_mla(num_tokens: int):
641
max_pos = 8192
642
pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos
643
614
- q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
615
- kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
644
+ q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16)
645
+ kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16)
646
# Tile 0 is deliberately low-magnitude: amax/6 lands between two e4m3
647
# subnormals (spaced a flat 2^-9 there), so round-to-nearest picks a scale
648
# BELOW amax/6 and the tile's peak saturates at +-6. Asserted below, so a
@@ -622,10 +652,10 @@ def test_fused_norm_rope_nvfp4_ds_mla(num_tokens: int):
652
kv_c[:, :16] = (
653
7.5 * (2.0**-9) * torch.linspace(0.2, 1.0, 16, device=dev, dtype=torch.bfloat16)
654
)
625
- k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
626
- qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
627
- kvw = torch.ones(KV_LORA, device=dev, dtype=torch.bfloat16)
628
- mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
655
+ k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16)
656
+ qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16)
657
+ kvw = torch.ones(cfg.kv_lora, device=dev, dtype=torch.bfloat16)
658
+ mla_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
659
660
bs = max_pos
661
mla_cache = torch.zeros(1, bs, 352, device=dev, dtype=torch.uint8)
@@ -675,19 +705,19 @@ def test_fused_norm_rope_nvfp4_ds_mla(num_tokens: int):
705
)
706
707
ref_codes = quantize_to_e2m1(tiles / ref_scale.float().unsqueeze(-1))
678
- ref_codes = ref_codes.reshape(num_tokens, KV_LORA)
708
+ ref_codes = ref_codes.reshape(num_tokens, cfg.kv_lora)
709
710
cache = mla_cache[0, :num_tokens] # [N, 352] uint8
681
- packed = cache[:, : KV_LORA // 2]
711
+ packed = cache[:, : cfg.kv_lora // 2]
712
got_codes = torch.stack([packed & 0xF, packed >> 4], dim=-1)
683
- got_codes = got_codes.reshape(num_tokens, KV_LORA) # low nibble = even elem
713
+ got_codes = got_codes.reshape(num_tokens, cfg.kv_lora) # low nibble = even elem
714
torch.testing.assert_close(got_codes, ref_codes, rtol=0, atol=0)
715
686
- got_rope = cache[:, KV_LORA // 2 : KV_LORA // 2 + ROPE_DIM].view(FP8)
716
+ got_rope = cache[:, cfg.kv_lora // 2 : cfg.kv_lora // 2 + cfg.rope_dim].view(FP8)
717
assert_fp8(got_rope, kpe_ref.to(FP8), "nvfp4_ds_mla RoPE e4m3")
718
719
perm = nvfp4_sf_byte(torch.arange(32, device=dev))
690
- got_scale = cache[:, KV_LORA // 2 + ROPE_DIM :].view(FP8)[:, perm]
720
+ got_scale = cache[:, cfg.kv_lora // 2 + cfg.rope_dim :].view(FP8)[:, perm]
721
torch.testing.assert_close(got_scale.float(), ref_scale.float(), rtol=0, atol=0)
722
723
# No indexer on this call: top-k buffer must be untouched.
@@ -737,39 +767,34 @@ def test_fused_norm_rope_supports_large_token_count():
767
# ── fused_q ──────────────────────────────────────────────────────────────────
768
769
740
-@pytest.mark.parametrize(
741
- ("num_tokens", "num_q_heads"),
742
- [
743
- (1, NUM_HEADS),
744
- (4, NUM_HEADS),
745
- (17, NUM_HEADS),
746
- (512, NUM_HEADS),
747
- (4096, NUM_HEADS),
748
- (17, 64),
749
- ],
750
-)
770
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
771
+@pytest.mark.parametrize("num_tokens", [1, 4, 17, 512, 4096])
772
@pytest.mark.parametrize("index_interleave", [True, False])
752
-def test_fused_q(num_tokens: int, num_q_heads: int, index_interleave: bool):
773
+def test_fused_q(num_tokens: int, index_interleave: bool, cfg: ModelConfig):
774
torch.manual_seed(2)
775
dev = "cuda"
776
max_pos = 8192
777
pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos
778
779
q_pe = torch.randn(
759
- num_tokens, num_q_heads, ROPE_DIM, device=dev, dtype=torch.bfloat16
780
+ num_tokens, cfg.num_heads, cfg.rope_dim, device=dev, dtype=torch.bfloat16
781
)
782
ql_nope = torch.randn(
762
- num_tokens, num_q_heads, KV_LORA, device=dev, dtype=torch.bfloat16
783
+ num_tokens, cfg.num_heads, cfg.kv_lora, device=dev, dtype=torch.bfloat16
784
)
785
index_q = torch.randn(
765
- num_tokens, INDEX_HEADS, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16
786
+ num_tokens,
787
+ cfg.index_heads,
788
+ cfg.index_head_dim,
789
+ device=dev,
790
+ dtype=torch.bfloat16,
791
)
767
- index_w = torch.randn(num_tokens, INDEX_HEADS, device=dev, dtype=torch.float32)
792
+ index_w = torch.randn(num_tokens, cfg.index_heads, device=dev, dtype=torch.float32)
793
q_scale = torch.tensor([0.37], device=dev, dtype=torch.float32)
769
- softmax_scale = INDEX_HEAD_DIM**-0.5
770
- head_scale = INDEX_HEADS**-0.5
771
- q_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev) # q_pe: interleaved
772
- idx_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
794
+ softmax_scale = cfg.index_head_dim**-0.5
795
+ head_scale = cfg.index_heads**-0.5
796
+ q_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev) # q_pe: interleaved
797
+ idx_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
798
799
iq_fp8, iw_out, mqa = K.fused_q(
800
pos,
@@ -791,18 +816,18 @@ def test_fused_q(num_tokens: int, num_q_heads: int, index_interleave: bool):
816
mqa_nope_ref = (ql_nope.float() / s).to(FP8)
817
qpe_ref = rope(
818
q_pe.float(),
794
- pos.unsqueeze(-1).expand(num_tokens, num_q_heads),
819
+ pos.unsqueeze(-1).expand(num_tokens, cfg.num_heads),
820
q_cos_sin,
821
interleave=True,
822
)
823
mqa_pe_ref = (qpe_ref / s).to(FP8)
799
- assert_fp8(mqa[:, :, :KV_LORA], mqa_nope_ref, "mqa ql_nope")
800
- assert_fp8(mqa[:, :, KV_LORA:], mqa_pe_ref, "mqa q_pe")
824
+ assert_fp8(mqa[:, :, : cfg.kv_lora], mqa_nope_ref, "mqa ql_nope")
825
+ assert_fp8(mqa[:, :, cfg.kv_lora :], mqa_pe_ref, "mqa q_pe")
826
827
# Indexer-Q: RoPE + UE8M0 fp8 quant; index weights fold in q-scale.
828
iq_ref = rope(
829
index_q.float(),
805
- pos.unsqueeze(-1).expand(num_tokens, INDEX_HEADS),
830
+ pos.unsqueeze(-1).expand(num_tokens, cfg.index_heads),
831
idx_cos_sin,
832
interleave=index_interleave,
833
)
@@ -812,20 +837,21 @@ def test_fused_q(num_tokens: int, num_q_heads: int, index_interleave: bool):
837
torch.testing.assert_close(iw_out, iw_ref, rtol=1e-3, atol=1e-3)
838
839
840
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
841
@pytest.mark.parametrize("num_tokens", [1, 17, 512])
816
-def test_fused_q_no_indexer(num_tokens: int):
842
+def test_fused_q_no_indexer(num_tokens: int, cfg: ModelConfig):
843
torch.manual_seed(3)
844
dev = "cuda"
845
max_pos = 8192
846
pos = torch.arange(num_tokens, device=dev, dtype=torch.int64)
847
q_pe = torch.randn(
822
- num_tokens, NUM_HEADS, ROPE_DIM, device=dev, dtype=torch.bfloat16
848
+ num_tokens, cfg.num_heads, cfg.rope_dim, device=dev, dtype=torch.bfloat16
849
)
850
ql_nope = torch.randn(
825
- num_tokens, NUM_HEADS, KV_LORA, device=dev, dtype=torch.bfloat16
851
+ num_tokens, cfg.num_heads, cfg.kv_lora, device=dev, dtype=torch.bfloat16
852
)
853
q_scale = torch.tensor([0.5], device=dev, dtype=torch.float32)
828
- q_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
854
+ q_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
855
_, _, mqa = K.fused_q(
856
pos,
857
q_pe,
@@ -841,19 +867,20 @@ def test_fused_q_no_indexer(num_tokens: int):
867
index_rope_interleave=False,
868
)
869
s = q_scale.item()
844
- assert_fp8(mqa[:, :, :KV_LORA], (ql_nope.float() / s).to(FP8), "mqa ql_nope")
870
+ assert_fp8(mqa[:, :, : cfg.kv_lora], (ql_nope.float() / s).to(FP8), "mqa ql_nope")
871
qpe_ref = rope(
872
q_pe.float(),
847
- pos.unsqueeze(-1).expand(num_tokens, NUM_HEADS),
873
+ pos.unsqueeze(-1).expand(num_tokens, cfg.num_heads),
874
q_cos_sin,
875
interleave=True,
876
)
851
- assert_fp8(mqa[:, :, KV_LORA:], (qpe_ref / s).to(FP8), "mqa q_pe")
877
+ assert_fp8(mqa[:, :, cfg.kv_lora :], (qpe_ref / s).to(FP8), "mqa q_pe")
878
879
880
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
881
@pytest.mark.parametrize("num_tokens", [1, 17, 512])
882
@pytest.mark.parametrize("has_indexer", [True, False])
856
-def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool):
883
+def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool, cfg: ModelConfig):
884
"""bf16-query path (FlashMLA sparse, SM90/SM100): only the RoPE'd q_pe is
885
produced (bf16, unquantized); ql_nope is consumed directly by the caller."""
886
torch.manual_seed(6)
@@ -862,21 +889,27 @@ def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool):
889
pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos
890
891
q_pe = torch.randn(
865
- num_tokens, NUM_HEADS, ROPE_DIM, device=dev, dtype=torch.bfloat16
892
+ num_tokens, cfg.num_heads, cfg.rope_dim, device=dev, dtype=torch.bfloat16
893
)
894
ql_nope = torch.randn(
868
- num_tokens, NUM_HEADS, KV_LORA, device=dev, dtype=torch.bfloat16
895
+ num_tokens, cfg.num_heads, cfg.kv_lora, device=dev, dtype=torch.bfloat16
896
)
897
q_scale = torch.tensor([0.37], device=dev, dtype=torch.float32)
871
- q_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
898
+ q_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
899
900
index_q = index_w = idx_cos_sin = None
901
if has_indexer:
902
index_q = torch.randn(
876
- num_tokens, INDEX_HEADS, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16
903
+ num_tokens,
904
+ cfg.index_heads,
905
+ cfg.index_head_dim,
906
+ device=dev,
907
+ dtype=torch.bfloat16,
908
+ )
909
+ index_w = torch.randn(
910
+ num_tokens, cfg.index_heads, device=dev, dtype=torch.float32
911
)
878
- index_w = torch.randn(num_tokens, INDEX_HEADS, device=dev, dtype=torch.float32)
879
- idx_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
912
+ idx_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
913
914
iq_fp8, iw_out, q_pe_out = K.fused_q(
915
pos,
@@ -887,8 +920,8 @@ def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool):
920
ql_nope,
921
q_scale,
922
index_w,
890
- INDEX_HEAD_DIM**-0.5,
891
- INDEX_HEADS**-0.5,
923
+ cfg.index_head_dim**-0.5,
924
+ cfg.index_heads**-0.5,
925
has_indexer=has_indexer,
926
index_rope_interleave=False,
927
quantize_mqa=False,
@@ -896,10 +929,10 @@ def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool):
929
930
# MQA query: only the RoPE'd q_pe, bf16, unquantized.
931
assert q_pe_out.dtype == torch.bfloat16
899
- assert q_pe_out.shape == (num_tokens, NUM_HEADS, ROPE_DIM)
932
+ assert q_pe_out.shape == (num_tokens, cfg.num_heads, cfg.rope_dim)
933
qpe_ref = rope(
934
q_pe.float(),
902
- pos.unsqueeze(-1).expand(num_tokens, NUM_HEADS),
935
+ pos.unsqueeze(-1).expand(num_tokens, cfg.num_heads),
936
q_cos_sin,
937
interleave=True,
938
)
@@ -910,13 +943,15 @@ def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool):
943
assert index_q is not None
944
iq_ref = rope(
945
index_q.float(),
913
- pos.unsqueeze(-1).expand(num_tokens, INDEX_HEADS),
946
+ pos.unsqueeze(-1).expand(num_tokens, cfg.index_heads),
947
idx_cos_sin,
948
interleave=False,
949
)
950
q_ref, scale_ref = ue8m0_quant(iq_ref)
951
assert_fp8(iq_fp8, q_ref, "indexer-Q fp8 (bf16-query path)")
919
- iw_ref = index_w * scale_ref * (INDEX_HEAD_DIM**-0.5) * (INDEX_HEADS**-0.5)
952
+ iw_ref = (
953
+ index_w * scale_ref * (cfg.index_head_dim**-0.5) * (cfg.index_heads**-0.5)
954
+ )
955
torch.testing.assert_close(iw_out, iw_ref, rtol=1e-3, atol=1e-3)
956
957
@@ -956,21 +991,22 @@ def test_fused_q_triton_supports_large_token_count():
991
# ── fused_eh_norm (MTP) ──────────────────────────────────────────────────────
992
993
994
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
995
@pytest.mark.parametrize("num_tokens", [1, 4, 17, 512])
960
-def test_fused_eh_norm(num_tokens: int):
996
+def test_fused_eh_norm(num_tokens: int, cfg: ModelConfig):
997
torch.manual_seed(4)
998
dev = "cuda"
999
# Mix in a position-0 token to exercise the embeds-zeroing branch.
1000
pos = torch.arange(num_tokens, device=dev, dtype=torch.int64)
1001
pos[0] = 0
966
- embeds = torch.randn(num_tokens, HIDDEN, device=dev, dtype=torch.bfloat16)
967
- prev = torch.randn(num_tokens, HIDDEN, device=dev, dtype=torch.bfloat16)
968
- ew = torch.randn(HIDDEN, device=dev, dtype=torch.bfloat16)
969
- hw = torch.randn(HIDDEN, device=dev, dtype=torch.bfloat16)
1002
+ embeds = torch.randn(num_tokens, cfg.hidden, device=dev, dtype=torch.bfloat16)
1003
+ prev = torch.randn(num_tokens, cfg.hidden, device=dev, dtype=torch.bfloat16)
1004
+ ew = torch.randn(cfg.hidden, device=dev, dtype=torch.bfloat16)
1005
+ hw = torch.randn(cfg.hidden, device=dev, dtype=torch.bfloat16)
1006
1007
out = K.fused_eh_norm(pos, embeds, prev, ew, hw, EPS)
1008
1009
masked = torch.where(pos.unsqueeze(-1) == 0, torch.zeros_like(embeds), embeds)
1010
ref = torch.cat([rms_norm(masked, ew), rms_norm(prev, hw)], dim=-1)
975
- assert out.shape == (num_tokens, 2 * HIDDEN)
1011
+ assert out.shape == (num_tokens, 2 * cfg.hidden)
1012
assert_bf16(out, ref, "eh_norm")
vllm/_aiter_ops.py
+1
-1
@@ -3406,7 +3406,7 @@ class rocm_aiter_ops:
3406
X: torch.Tensor,
3407
W: torch.Tensor,
3408
w_scale: torch.Tensor,
3409
- Y: torch.Tensor,
3409
+ Y: torch.Tensor | None = None,
3410
transpose_bm: bool | None = False,
3411
prequant: bool | None = False,
3412
y_scale: torch.Tensor | None = None,
vllm/model_executor/layers/attention/mla_attention.py
+1
-3
@@ -1025,9 +1025,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
1025
mqa_q_pe = mqa_pe_padded
1026
1027
if self.is_aiter_triton_fp4_bmm_enabled:
1028
- from aiter.ops.triton.batched_gemm_a16wfp4 import batched_gemm_a16wfp4
1029
-
1030
- mqa_ql_nope = batched_gemm_a16wfp4(
1028
+ mqa_ql_nope = rocm_aiter_ops.batched_gemm_a16wfp4(
1029
mqa_q_nope,
1030
self.W_K,
1031
self.W_K_scale,
vllm/model_executor/models/deepseek_v2.py
+5
@@ -665,6 +665,11 @@ class DeepseekV32IndexerCache(torch.nn.Module, AttentionLayerBase):
665
666
def forward(self): ...
667
668
+ @property
669
+ def uses_shuffled_layout(self) -> bool:
670
+ """Whether this cache's reader expects the shuffled value layout."""
671
+ return False
672
+
673
def get_attn_backend(self) -> type[AttentionBackend]:
674
return DeepseekV32IndexerBackend
675
vllm/models/deepseek_v32/__init__.py
+8
-4
@@ -7,9 +7,9 @@ MLA + a "lightning indexer" that selects the top-k tokens for a sparse MLA
7
attend. The same model code serves any DSA checkpoint, including GLM-5.2
8
(``glm_moe_dsa``), which reuses this architecture.
9
10
-The CUDA implementation selects capability-specific kernels internally and
11
-falls back when an optimization is unavailable. Other platforms use the
12
-generic implementation by default.
10
+The CUDA and ROCm implementations select capability-specific kernels
11
+internally and fall back when an optimization is unavailable. Other platforms
12
+use the generic implementation by default.
13
"""
14
15
from vllm.platforms import current_platform
@@ -20,8 +20,12 @@ if current_platform.is_cuda():
20
from .nvidia.model import DeepseekV32ForCausalLM
21
from .nvidia.model import DeepseekV32ForCausalLM as GlmMoeDsaForCausalLM
22
from .nvidia.mtp import DeepseekV32MTP
23
+elif current_platform.is_rocm():
24
+ from .amd.model import DeepseekV32ForCausalLM
25
+ from .amd.model import DeepseekV32ForCausalLM as GlmMoeDsaForCausalLM
26
+ from .amd.mtp import DeepseekV32MTP
27
else:
24
- # ROCm, XPU, and CPU keep the generic implementation.
28
+ # XPU and CPU keep the generic implementation.
29
from vllm.model_executor.models.deepseek_mtp import DeepSeekMTP as DeepseekV32MTP
30
from vllm.model_executor.models.deepseek_v2 import (
31
DeepseekV3ForCausalLM as DeepseekV32ForCausalLM,
vllm/models/deepseek_v32/amd/rocm.py
+16
-25
@@ -3,12 +3,14 @@
3
4
import torch
5
6
+from vllm._aiter_ops import rocm_aiter_ops
7
from vllm.compilation.breakable_cudagraph import eager_break_during_capture
8
from vllm.forward_context import get_forward_context
9
from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
10
from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache
11
from vllm.models.deepseek_v32.attention import DeepseekV32Attention, DeepseekV32Indexer
12
from vllm.models.deepseek_v32.common.kernels import fused_norm_rope, fused_q
13
+from vllm.platforms import current_platform
14
from vllm.utils.torch_utils import is_quantized_kv_cache
15
from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerBackend
16
from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
@@ -16,21 +18,15 @@ from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
18
)
19
20
19
-class DeepseekV32MLASparseBackend(ROCMAiterMLASparseBackend):
20
- @staticmethod
21
- def get_supported_kernel_block_sizes() -> list:
22
- return [16, 32]
23
-
24
-
25
-class DeepseekV32ROCmIndexerBackend(DeepseekV32IndexerBackend):
26
- @staticmethod
27
- def get_supported_kernel_block_sizes() -> list:
28
- return [16, 32]
29
-
30
-
21
class DeepseekV32ROCmIndexerCache(DeepseekV32IndexerCache):
22
def get_attn_backend(self):
33
- return DeepseekV32ROCmIndexerBackend
23
+ return DeepseekV32IndexerBackend
24
+
25
+ @property
26
+ def uses_shuffled_layout(self) -> bool:
27
+ # aiter's gather/insert pair shuffles the cache above block size 1:
28
+ # [n_blocks, blk/16, head_dim/16, 16, 16] instead of [n_blocks, blk, head_dim].
29
+ return self.kv_cache.ndim == 3 and self.kv_cache.shape[1] != 1
30
31
32
class DeepseekV32ROCmIndexer(DeepseekV32Indexer):
@@ -46,7 +42,7 @@ class DeepseekV32MLAAttention(DeepseekV32Attention):
42
config,
43
prefix,
44
topk_indices_buffer,
49
- attn_backend=DeepseekV32MLASparseBackend,
45
+ attn_backend=ROCMAiterMLASparseBackend,
46
)
47
48
self.indexer_op: SparseAttnIndexer | None = None
@@ -100,14 +96,10 @@ class DeepseekV32MLAAttention(DeepseekV32Attention):
96
q_nope = q_nope.transpose(0, 1) # (N, tokens, P)
97
98
if self.is_aiter_triton_fp4_bmm_enabled:
103
- from aiter.ops.triton.batched_gemm_a16wfp4 import batched_gemm_a16wfp4
104
-
105
- ql_nope = batched_gemm_a16wfp4(
99
+ ql_nope = rocm_aiter_ops.batched_gemm_a16wfp4(
100
q_nope, self.W_K, self.W_K_scale, transpose_bm=True, prequant=True
101
)
102
elif self.is_aiter_triton_fp8_bmm_enabled:
109
- from vllm._aiter_ops import rocm_aiter_ops
110
-
103
ql_nope = rocm_aiter_ops.triton_fp8_bmm(
104
q_nope, self.W_K, self.W_K_scale, group_size=128, transpose_bm=True
105
)
@@ -152,14 +144,10 @@ class DeepseekV32MLAAttention(DeepseekV32Attention):
144
)
145
146
if self.is_aiter_triton_fp4_bmm_enabled:
155
- from aiter.ops.triton.batched_gemm_a16wfp4 import batched_gemm_a16wfp4
156
-
157
- batched_gemm_a16wfp4(
147
+ rocm_aiter_ops.batched_gemm_a16wfp4(
148
x, self.W_V, self.W_V_scale, out_view, transpose_bm=True, prequant=True
149
)
150
elif self.is_aiter_triton_fp8_bmm_enabled:
161
- from vllm._aiter_ops import rocm_aiter_ops
162
-
151
rocm_aiter_ops.triton_fp8_bmm(
152
x,
153
self.W_V,
@@ -202,6 +190,7 @@ class DeepseekV32MLAAttention(DeepseekV32Attention):
190
indexer_k_norm_eps = self.indexer.k_norm.eps
191
indexer_k_rope_cos_sin_cache = self.indexer_rope_emb.cos_sin_cache
192
indexer_k_cache = self.indexer.k_cache.kv_cache
193
+ indexer_cache_shuffled = self.indexer.k_cache.uses_shuffled_layout
194
indexer_softmax_scale = self.indexer.softmax_scale
195
indexer_n_head_scale = self.indexer.n_head**-0.5
196
else:
@@ -211,6 +200,7 @@ class DeepseekV32MLAAttention(DeepseekV32Attention):
200
indexer_k_norm_eps = 1e-6
201
indexer_k_rope_cos_sin_cache = None
202
indexer_k_cache = None
203
+ indexer_cache_shuffled = False
204
indexer_softmax_scale = 0.0
205
indexer_n_head_scale = 0.0
206
@@ -241,6 +231,7 @@ class DeepseekV32MLAAttention(DeepseekV32Attention):
231
self.topk_indices_buffer,
232
slot_mapping=mla_slot,
233
indexer_k_cache=indexer_k_cache,
234
+ indexer_cache_shuffled=indexer_cache_shuffled,
235
mla_kv_cache=mla_kv_cache,
236
mla_kv_cache_dtype=self.kv_cache_dtype,
237
mla_k_scale=mla_k_scale,
@@ -282,7 +273,7 @@ class DeepseekV32MLAAttention(DeepseekV32Attention):
273
num_actual = attn_metadata.num_actual_tokens # type: ignore[attr-defined]
274
kv_cache = self.kv_cache
275
if self._fp8_kv_needs_view:
285
- kv_cache = kv_cache.view(torch.float8_e4m3fn)
276
+ kv_cache = kv_cache.view(current_platform.fp8_dtype())
277
278
q_for_attn = self._build_q_for_attn(ql_nope, mqa_q, num_actual)
279
attn_out, _ = self.impl.forward_mqa( # type: ignore[attr-defined]
vllm/models/deepseek_v32/attention.py
+4
-1
@@ -327,6 +327,7 @@ class DeepseekV32Attention(MLAAttention):
327
indexer_k_norm_bias = self.indexer.k_norm.bias
328
indexer_k_norm_eps = self.indexer.k_norm.eps
329
indexer_k_rope_cos_sin_cache = self.indexer_rope_emb.cos_sin_cache
330
+ indexer_cache_shuffled = self.indexer.k_cache.uses_shuffled_layout
331
indexer_k_cache = None if self.use_pcp else self.indexer.k_cache.kv_cache
332
index_k_out = torch.empty_like(index_k) if self.use_pcp else None
333
indexer_softmax_scale = self.indexer.softmax_scale
@@ -338,6 +339,7 @@ class DeepseekV32Attention(MLAAttention):
339
indexer_k_norm_eps = 1e-6
340
indexer_k_rope_cos_sin_cache = None
341
indexer_k_cache = None
342
+ indexer_cache_shuffled = False
343
index_k_out = None
344
indexer_softmax_scale = 0.0
345
indexer_n_head_scale = 0.0
@@ -377,6 +379,7 @@ class DeepseekV32Attention(MLAAttention):
379
slot_mapping=mla_slot,
380
indexer_slot_mapping=indexer_slot,
381
indexer_k_cache=indexer_k_cache,
382
+ indexer_cache_shuffled=indexer_cache_shuffled,
383
mla_kv_cache=mla_kv_cache,
384
mla_kv_cache_dtype=self.kv_cache_dtype,
385
mla_k_scale=mla_k_scale,
@@ -544,7 +547,7 @@ class DeepseekV32Attention(MLAAttention):
547
return
548
549
if self._fp8_kv_needs_view:
547
- kv_cache = kv_cache.view(torch.float8_e4m3fn)
550
+ kv_cache = kv_cache.view(current_platform.fp8_dtype())
551
if self._fp8_query:
552
# FlashInfer sparse: single packed fp8 query.
553
mqa_q_arg: torch.Tensor | tuple[torch.Tensor, torch.Tensor] = mqa_q[
vllm/models/deepseek_v32/common/kernels.py
+97
-27
@@ -12,6 +12,18 @@ from vllm.utils.torch_utils import is_quantized_kv_cache
12
# has_indexer=False path so the indexer args don't allocate every call.
13
_DUMMY_CACHE: dict[tuple, torch.Tensor] = {}
14
15
+# Tile shape of the indexer-K cache's shuffled layout, used when the cache reports
16
+# uses_shuffled_layout. The block tile is a token count; the head tile is a byte
17
+# count, which the kernel converts to cache elements.
18
+_INDEXER_CACHE_BLOCK_TILE = 16
19
+_INDEXER_CACHE_HEAD_TILE_BYTES = 16
20
+
21
+# FNUZ (e4m3fnuz) on gfx942, OCP (e4m3fn) elsewhere. 224.0 keeps the ue8m0
22
+# power-of-two scale inside the FNUZ range after its ceil.
23
+_FP8_DTYPE = current_platform.fp8_dtype()
24
+_USE_FNUZ = torch.float8_e4m3fnuz == _FP8_DTYPE
25
+_FP8_MAX = 224.0 if _USE_FNUZ else 448.0
26
+
27
28
def _dummy(shape: tuple, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
29
key = (shape, dtype, device)
@@ -47,16 +59,17 @@ def _get_cos_sin(
59
60
61
@triton.jit
50
-def _fp8_ue8m0_quantize(vals):
62
+def _fp8_ue8m0_quantize(vals, FP8_MAX: tl.constexpr, USE_FNUZ: tl.constexpr):
63
"""Quantize float32 values to FP8 E4M3 with a ue8m0 (power-of-2) scale.
64
65
Returns (fp8_vals, scale) so the caller can store them or reuse the scale.
66
"""
67
vals = vals.to(tl.float32)
68
amax = tl.max(tl.abs(vals))
57
- scale = tl.div_rn(tl.maximum(amax, 1e-4), 448.0)
69
+ scale = tl.div_rn(tl.maximum(amax, 1e-4), FP8_MAX)
70
scale = tl.math.exp2(tl.math.ceil(tl.math.log2(scale)))
59
- fp8_vals = tl.div_rn(vals, scale).to(tl.float8e4nv)
71
+ fp8_dtype = tl.float8e4b8 if USE_FNUZ else tl.float8e4nv
72
+ fp8_vals = tl.div_rn(vals, scale).to(fp8_dtype)
73
return fp8_vals, scale
74
75
@@ -71,18 +84,32 @@ def _fp8_quant_and_cache_write(
84
cache_block_stride,
85
offsets,
86
HEAD_DIM: tl.constexpr,
87
+ SHUFFLE: tl.constexpr,
88
+ BLOCK_TILE: tl.constexpr,
89
+ HEAD_TILE: tl.constexpr,
90
+ FP8_MAX: tl.constexpr,
91
+ USE_FNUZ: tl.constexpr,
92
):
75
- k_fp8, scale = _fp8_ue8m0_quantize(vals)
93
+ k_fp8, scale = _fp8_ue8m0_quantize(vals, FP8_MAX, USE_FNUZ)
94
95
block_idx = slot_idx // cache_block_size
96
block_offset = slot_idx % cache_block_size
97
block_start = block_idx * cache_block_stride
98
81
- tl.store(
82
- kv_cache_ptr + block_start + block_offset * HEAD_DIM + offsets,
83
- k_fp8,
84
- mask=mask,
85
- )
99
+ # Shuffled layout: [blk/BLOCK_TILE, HEAD_DIM/HEAD_TILE, BLOCK_TILE, HEAD_TILE],
100
+ # so one contiguous run holds BLOCK_TILE tokens x HEAD_TILE bytes and the
101
+ # reader's coalesced loads land on the bytes we wrote.
102
+ if SHUFFLE:
103
+ value_off = (
104
+ block_offset // BLOCK_TILE * BLOCK_TILE * HEAD_DIM
105
+ + block_offset % BLOCK_TILE * HEAD_TILE
106
+ + offsets // HEAD_TILE * BLOCK_TILE * HEAD_TILE
107
+ + offsets % HEAD_TILE
108
+ )
109
+ else:
110
+ value_off = block_offset * HEAD_DIM + offsets
111
+
112
+ tl.store(kv_cache_ptr + block_start + value_off, k_fp8, mask=mask)
113
scale_byte_off = block_start + cache_block_size * HEAD_DIM + block_offset * 4
114
tl.store(kv_cache_scale_ptr + scale_byte_off // 4, scale)
115
@@ -137,6 +164,9 @@ def _fused_norm_rope_kernel(
164
indexer_cache_scale_ptr,
165
indexer_cache_block_size,
166
indexer_cache_stride,
167
+ INDEXER_CACHE_SHUFFLE: tl.constexpr,
168
+ INDEXER_CACHE_BLOCK_TILE: tl.constexpr,
169
+ INDEXER_CACHE_HEAD_TILE: tl.constexpr,
170
# MLA KV cache (concat kv_c_normed + k_pe_roped, uses slot_mapping_ptr)
171
mla_cache_ptr,
172
mla_cache_block_stride,
@@ -161,12 +191,15 @@ def _fused_norm_rope_kernel(
191
HAS_INDEXER: tl.constexpr,
192
INDEX_ROPE_INTERLEAVE: tl.constexpr,
193
USE_PDL: tl.constexpr,
194
+ FP8_MAX: tl.constexpr,
195
+ USE_FNUZ: tl.constexpr,
196
):
197
tok_idx = tl.program_id(0).to(tl.int64)
198
pid = tl.program_id(1)
199
if USE_PDL:
200
tl.extra.cuda.gdc_wait()
201
tl.extra.cuda.gdc_launch_dependents()
202
+ fp8_dtype = tl.float8e4b8 if USE_FNUZ else tl.float8e4nv
203
if pid == 3:
204
if not HAS_INDEXER:
205
# Shared layer: reuse the previous indexer layer's top-k; do not
@@ -327,9 +360,11 @@ def _fused_norm_rope_kernel(
360
)
361
kv_2d = tl.reshape(kv_c, (MLA_NUM_TILES, MLA_TILE_DIM))
362
tile_amax = tl.max(tl.abs(kv_2d), axis=1, keep_dims=True)
330
- tile_scale = tl.maximum(tile_amax * (1.0 / 448.0), 1e-4)
363
+ # scale = amax / FP8_MAX (fp8 e4m3 max), matching the reference
364
+ # concat_and_cache_ds_mla kernel; floored to FLT_MIN.
365
+ tile_scale = tl.maximum(tile_amax * (1.0 / FP8_MAX), 1e-4)
366
tile_scale = tl.math.exp2(tl.math.ceil(tl.math.log2(tile_scale)))
332
- kv_c_fp8 = tl.reshape((kv_2d / tile_scale).to(tl.float8e4nv), (KV_DIM,))
367
+ kv_c_fp8 = tl.reshape((kv_2d / tile_scale).to(fp8_dtype), (KV_DIM,))
368
tl.store(mla_cache_ptr + byte_base + kv_block, kv_c_fp8)
369
tile_off = tl.arange(0, MLA_NUM_TILES)
370
tl.store(
@@ -349,16 +384,16 @@ def _fused_norm_rope_kernel(
384
# kv_c_normed (KV_DIM elements)
385
if MLA_CACHE_FP8:
386
scale = tl.load(mla_cache_scale_ptr)
352
- kv_c_fp8 = (kv_c.to(tl.float32) / scale).to(tl.float8e4nv)
387
+ kv_c_fp8 = (kv_c.to(tl.float32) / scale).to(fp8_dtype)
388
tl.store(dst + kv_block, kv_c_fp8)
389
else:
390
tl.store(dst + kv_block, kv_c)
391
# k_pe_roped (from registers, interleaved layout)
392
if MLA_CACHE_FP8:
358
- tl.store(dst + KV_DIM + dim_off * 2, (r1 / scale).to(tl.float8e4nv))
393
+ tl.store(dst + KV_DIM + dim_off * 2, (r1 / scale).to(fp8_dtype))
394
tl.store(
395
dst + KV_DIM + dim_off * 2 + 1,
361
- (r2 / scale).to(tl.float8e4nv),
396
+ (r2 / scale).to(fp8_dtype),
397
)
398
else:
399
tl.store(dst + KV_DIM + dim_off * 2, r1)
@@ -463,6 +498,11 @@ def _fused_norm_rope_kernel(
498
indexer_cache_stride,
499
index_k_block,
500
INDEX_K_DIM,
501
+ INDEXER_CACHE_SHUFFLE,
502
+ INDEXER_CACHE_BLOCK_TILE,
503
+ INDEXER_CACHE_HEAD_TILE,
504
+ FP8_MAX,
505
+ USE_FNUZ,
506
)
507
508
@@ -486,6 +526,7 @@ def fused_norm_rope(
526
slot_mapping: torch.Tensor | None = None,
527
indexer_slot_mapping: torch.Tensor | None = None,
528
indexer_k_cache: torch.Tensor | None = None,
529
+ indexer_cache_shuffled: bool = False,
530
mla_kv_cache: torch.Tensor | None = None,
531
mla_kv_cache_dtype: str = "auto",
532
mla_k_scale: torch.Tensor | None = None,
@@ -532,13 +573,28 @@ def fused_norm_rope(
573
indexer_slot_mapping = slot_mapping
574
idx_cache_scale_view = indexer_k_cache.view(torch.uint8).view(torch.float32)
575
idx_cache_block_size = indexer_k_cache.shape[1]
576
+ idx_cache_stride = indexer_k_cache.shape[2]
577
+ # The caller's cache reports whether its reader expects the shuffled
578
+ # layout; see DeepseekV32IndexerCache.uses_shuffled_layout.
579
+ idx_cache_shuffle = indexer_cache_shuffled
580
+ if idx_cache_shuffle:
581
+ assert idx_cache_block_size % _INDEXER_CACHE_BLOCK_TILE == 0, (
582
+ f"indexer K cache block size {idx_cache_block_size} must be a "
583
+ f"multiple of {_INDEXER_CACHE_BLOCK_TILE} for the shuffled layout"
584
+ )
585
idx_cache_stride = indexer_k_cache.stride(0)
586
if indexer_k_cache.dtype == torch.uint8:
537
- indexer_k_cache = indexer_k_cache.view(torch.float8_e4m3fn)
587
+ indexer_k_cache = indexer_k_cache.view(_FP8_DTYPE)
588
+ # The head tile is a byte count; the kernel indexes the cache in elements.
589
+ idx_cache_head_tile = (
590
+ _INDEXER_CACHE_HEAD_TILE_BYTES // indexer_k_cache.element_size()
591
+ )
592
else:
593
idx_cache_scale_view = None
594
idx_cache_block_size = 1
595
idx_cache_stride = 0
596
+ idx_cache_shuffle = False
597
+ idx_cache_head_tile = _INDEXER_CACHE_HEAD_TILE_BYTES
598
599
# --- MLA KV cache setup ---
600
mla_cache_nvfp4 = mla_kv_cache_dtype == "nvfp4_ds_mla"
@@ -561,12 +617,12 @@ def fused_norm_rope(
617
mla_entry_stride = u8_cache.stride(1)
618
mla_ds_scale_view = u8_cache.view(torch.float32)
619
mla_ds_rope_view = u8_cache.view(torch.bfloat16)
564
- mla_kv_cache = u8_cache.view(torch.float8_e4m3fn)
620
+ mla_kv_cache = u8_cache.view(_FP8_DTYPE)
621
else:
622
mla_block_stride = mla_kv_cache.stride(0)
623
mla_entry_stride = mla_kv_cache.stride(1)
624
if mla_cache_fp8 and mla_kv_cache.dtype == torch.uint8:
569
- mla_kv_cache = mla_kv_cache.view(torch.float8_e4m3fn)
625
+ mla_kv_cache = mla_kv_cache.view(_FP8_DTYPE)
626
if mla_k_scale is None:
627
mla_k_scale = torch.ones(1, dtype=torch.float32, device=device)
628
else:
@@ -592,6 +648,9 @@ def fused_norm_rope(
648
assert index_k_out.shape == index_k.shape
649
index_k_out_stride = index_k_out.stride(0)
650
use_pdl = current_platform.is_arch_support_pdl()
651
+ pdl_kwargs = (
652
+ {"USE_PDL": use_pdl, "launch_pdl": True} if use_pdl else {"USE_PDL": use_pdl}
653
+ )
654
_fused_norm_rope_kernel[(num_tokens, 4)](
655
positions,
656
# Q RMS norm
@@ -639,6 +698,9 @@ def fused_norm_rope(
698
idx_cache_scale_view,
699
idx_cache_block_size,
700
idx_cache_stride,
701
+ idx_cache_shuffle,
702
+ _INDEXER_CACHE_BLOCK_TILE,
703
+ idx_cache_head_tile,
704
# MLA KV cache (uses same slot_mapping)
705
mla_kv_cache,
706
mla_block_stride,
@@ -659,8 +721,9 @@ def fused_norm_rope(
721
TOPK_BLOCK_SIZE=1024,
722
HAS_INDEXER=has_indexer,
723
INDEX_ROPE_INTERLEAVE=index_rope_interleave,
662
- USE_PDL=use_pdl,
663
- launch_pdl=use_pdl,
724
+ FP8_MAX=_FP8_MAX,
725
+ USE_FNUZ=_USE_FNUZ,
726
+ **pdl_kwargs,
727
)
728
return q_c_out
729
@@ -715,6 +778,8 @@ def _fused_q_kernel(
778
INDEX_ROPE_INTERLEAVE: tl.constexpr,
779
QUANTIZE_MQA: tl.constexpr,
780
USE_PDL: tl.constexpr,
781
+ FP8_MAX: tl.constexpr,
782
+ USE_FNUZ: tl.constexpr,
783
):
784
tok_idx = tl.program_id(0).to(tl.int64)
785
pid = tl.program_id(1)
@@ -722,6 +787,7 @@ def _fused_q_kernel(
787
if USE_PDL:
788
tl.extra.cuda.gdc_wait()
789
tl.extra.cuda.gdc_launch_dependents()
790
+ fp8_dtype = tl.float8e4b8 if USE_FNUZ else tl.float8e4nv
791
792
if pid == 2:
793
# ql_nope quantize + pack into the front of mqa_q_fp8. On the bf16
@@ -744,7 +810,7 @@ def _fused_q_kernel(
810
+ ql_nope_off,
811
mask=ql_nope_mask,
812
).to(tl.float32)
747
- ql_nope_fp8 = (ql_nope / scale).to(tl.float8e4nv)
813
+ ql_nope_fp8 = (ql_nope / scale).to(fp8_dtype)
814
tl.store(
815
mqa_q_fp8_ptr
816
+ tok_idx * mqa_q_fp8_stride0
@@ -794,7 +860,7 @@ def _fused_q_kernel(
860
+ q_head_idx * mqa_q_fp8_stride1
861
+ QL_NOPE_DIM
862
+ rot_off * 2,
797
- (r1 / scale).to(tl.float8e4nv),
863
+ (r1 / scale).to(fp8_dtype),
864
)
865
tl.store(
866
mqa_q_fp8_ptr
@@ -803,7 +869,7 @@ def _fused_q_kernel(
869
+ QL_NOPE_DIM
870
+ rot_off * 2
871
+ 1,
806
- (r2 / scale).to(tl.float8e4nv),
872
+ (r2 / scale).to(fp8_dtype),
873
)
874
else:
875
# bf16 query: write the RoPE'd q_pe unquantized.
@@ -861,7 +927,7 @@ def _fused_q_kernel(
927
index_q = tl.where(in_rope, roped, index_q)
928
929
# Index Q Quantize (from registers)
864
- index_q_fp8, index_q_scale = _fp8_ue8m0_quantize(index_q)
930
+ index_q_fp8, index_q_scale = _fp8_ue8m0_quantize(index_q, FP8_MAX, USE_FNUZ)
931
tl.store(
932
index_q_fp8_ptr
933
+ tok_idx * index_q_fp8_stride0
@@ -955,7 +1021,7 @@ def fused_q(
1021
q_pe.shape[0],
1022
q_pe.shape[1],
1023
ql_nope.shape[2] + q_pe.shape[2],
958
- dtype=torch.float8_e4m3fn,
1024
+ dtype=_FP8_DTYPE,
1025
device=q_pe.device,
1026
)
1027
# Placeholder; pid 0 packs q_pe into mqa_q_fp8 instead.
@@ -967,7 +1033,7 @@ def fused_q(
1033
mqa_q_fp8 = q_pe_out # unused placeholder for the fp8 pack pointer
1034
mqa_q = q_pe_out
1035
970
- index_q_fp8 = torch.empty_like(index_q, dtype=torch.float8_e4m3fn)
1036
+ index_q_fp8 = torch.empty_like(index_q, dtype=_FP8_DTYPE)
1037
index_weights_out = torch.empty_like(index_weights, dtype=torch.float32)
1038
if cutedsl_kernel is not None:
1039
cutedsl_kernel(
@@ -990,6 +1056,9 @@ def fused_q(
1056
return index_q_fp8, index_weights_out, mqa_q
1057
1058
use_pdl = current_platform.is_arch_support_pdl()
1059
+ pdl_kwargs = (
1060
+ {"USE_PDL": use_pdl, "launch_pdl": True} if use_pdl else {"USE_PDL": use_pdl}
1061
+ )
1062
_fused_q_kernel[(num_tokens, 3, grid_heads)](
1063
positions,
1064
q_pe,
@@ -1031,8 +1100,9 @@ def fused_q(
1100
HAS_INDEXER=has_indexer,
1101
INDEX_ROPE_INTERLEAVE=index_rope_interleave,
1102
QUANTIZE_MQA=quantize_mqa,
1034
- USE_PDL=use_pdl,
1035
- launch_pdl=use_pdl,
1103
+ FP8_MAX=_FP8_MAX,
1104
+ USE_FNUZ=_USE_FNUZ,
1105
+ **pdl_kwargs,
1106
# num_warps=1 is optimal here: each program is a single 128-element
1107
# rope+quant, so the kernel is program-count/occupancy bound, not
1108
# per-program compute bound (swept 1/2/4/8 — 1 wins or ties everywhere).