From 1c0eee919db320ea3e2df37f800465c7842a650a Mon Sep 17 00:00:00 2001
From: Jack Hu <jack.hu@amd.com>
Date: Tue, 22 Sep 2026 19:08:51 -0400
Subject: [PATCH] [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>
---
 .../test_fused_deepseek_v32_norm_rope.py      | 336 ++++++++++--------
 vllm/_aiter_ops.py                            |   2 +-
 .../layers/attention/mla_attention.py         |   4 +-
 vllm/model_executor/models/deepseek_v2.py     |   5 +
 vllm/models/deepseek_v32/__init__.py          |  12 +-
 vllm/models/deepseek_v32/amd/rocm.py          |  41 +--
 vllm/models/deepseek_v32/attention.py         |   5 +-
 vllm/models/deepseek_v32/common/kernels.py    | 124 +++++--
 8 files changed, 318 insertions(+), 211 deletions(-)

diff --git a/tests/kernels/test_fused_deepseek_v32_norm_rope.py b/tests/kernels/test_fused_deepseek_v32_norm_rope.py
index 61f977ddbd..df9d45a99d 100644
--- a/tests/kernels/test_fused_deepseek_v32_norm_rope.py
+++ b/tests/kernels/test_fused_deepseek_v32_norm_rope.py
@@ -24,6 +24,8 @@ outputs are checked within 1 representable-step (ULP); bf16 norm/RoPE outputs us
 rtol/atol=1e-2 (the tolerance the sibling deepseek_v4 fused-kernel test uses).
 """
 
+from typing import NamedTuple
+
 import pytest
 import torch
 
@@ -33,19 +35,30 @@ from vllm.platforms import current_platform
 FP8 = torch.float8_e4m3fn
 FP8_MAX = 448.0
 
-# GLM-5.2 / DeepSeek-V3.2 shapes (TP8 local heads).
-Q_LORA = 2048
-KV_LORA = 512
-ROPE_DIM = 64
-NUM_HEADS = 8
-INDEX_HEADS = 32
-INDEX_HEAD_DIM = 128
-HIDDEN = 6144
 EPS = 1e-6
 
+
+class ModelConfig(NamedTuple):
+    """Per-model attention dims (num_heads is the TP8 local count)."""
+
+    q_lora: int
+    kv_lora: int
+    rope_dim: int
+    num_heads: int
+    index_heads: int
+    index_head_dim: int
+    hidden: int
+
+
+MODEL_CONFIGS = [
+    ModelConfig(2048, 512, 64, 8, 32, 128, 6144),  # GLM-5.2
+    ModelConfig(1536, 512, 64, 16, 64, 128, 7168),  # DeepSeek-V3.2
+]
+MODEL_IDS = ["glm5.2", "dsv3.2"]
+
 pytestmark = pytest.mark.skipif(
-    not current_platform.is_cuda() or not current_platform.has_device_capability(89),
-    reason="deepseek_v32 fused kernels require CUDA with fp8 (SM89+)",
+    not current_platform.supports_fp8(),
+    reason="deepseek_v32 fused kernels require fp8 support",
 )
 
 
@@ -144,29 +157,32 @@ def assert_fp8(got: torch.Tensor, ref: torch.Tensor, msg: str):
 # ── fused_norm_rope ──────────────────────────────────────────────────────────
 
 
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
 @pytest.mark.parametrize("num_tokens", [1, 4, 17, 512, 4096])
 @pytest.mark.parametrize("index_interleave", [True, False])
 @pytest.mark.parametrize("mla_dtype", ["auto", "bfloat16", "fp8"])
-def test_fused_norm_rope(num_tokens: int, index_interleave: bool, mla_dtype: str):
+def test_fused_norm_rope(
+    num_tokens: int, index_interleave: bool, mla_dtype: str, cfg: ModelConfig
+):
     torch.manual_seed(0)
     dev = "cuda"
     max_pos = 8192
     pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos
 
-    q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
-    kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
-    k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
-    qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
-    kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)
-    ik = torch.randn(num_tokens, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16)
-    ikw = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
-    ikb = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
+    q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16)
+    k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16)
+    qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kvw = torch.randn(cfg.kv_lora, device=dev, dtype=torch.bfloat16)
+    ik = torch.randn(num_tokens, cfg.index_head_dim, device=dev, dtype=torch.bfloat16)
+    ikw = torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32)
+    ikb = torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32)
 
-    mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)  # MLA k_pe: interleaved
-    idx_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
+    mla_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)  # MLA k_pe: interleaved
+    idx_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
 
     bs = max_pos  # single block covering all tokens
-    mla_dim = KV_LORA + ROPE_DIM
+    mla_dim = cfg.kv_lora + cfg.rope_dim
     mla_fp8 = mla_dtype == "fp8"
     if mla_fp8:
         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
     else:
         mla_cache = torch.zeros(1, bs, mla_dim, device=dev, dtype=torch.bfloat16)
         mla_k_scale = None
-    idx_row = INDEX_HEAD_DIM + INDEX_HEAD_DIM // 128 * 4  # 132
+    idx_row = cfg.index_head_dim + cfg.index_head_dim // 128 * 4  # 132
     idx_cache = torch.zeros(1, bs, idx_row, device=dev, dtype=torch.uint8)
     slot = torch.arange(num_tokens, device=dev, dtype=torch.int64)
     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
     if mla_fp8:
         cache = mla_cache.view(FP8)[0, :num_tokens]
         s = mla_k_scale.item()
-        assert_fp8(cache[:, :KV_LORA], (kv_ref / s).to(FP8), "MLA kv fp8")
-        assert_fp8(cache[:, KV_LORA:], (kpe_ref / s).to(FP8), "MLA k_pe fp8")
+        assert_fp8(cache[:, : cfg.kv_lora], (kv_ref / s).to(FP8), "MLA kv fp8")
+        assert_fp8(cache[:, cfg.kv_lora :], (kpe_ref / s).to(FP8), "MLA k_pe fp8")
     else:
         cache = mla_cache[0, :num_tokens]
-        assert_bf16(cache[:, :KV_LORA], kv_ref, "MLA kv bf16")
-        assert_bf16(cache[:, KV_LORA:], kpe_ref, "MLA k_pe bf16")
+        assert_bf16(cache[:, : cfg.kv_lora], kv_ref, "MLA kv bf16")
+        assert_bf16(cache[:, cfg.kv_lora :], kpe_ref, "MLA k_pe bf16")
 
     # Indexer-K cache (packed [bs*head_dim fp8 | bs*4 fp32 scale]).
     ik_ref = layer_norm(ik, ikw, ikb)
     ik_ref = rope(ik_ref, pos, idx_cos_sin, interleave=index_interleave)
     q_ref, s_ref = ue8m0_quant(ik_ref)
     flat = idx_cache[0].reshape(-1)
-    vals = flat[: bs * INDEX_HEAD_DIM].view(FP8).reshape(bs, INDEX_HEAD_DIM)
-    scales = flat[bs * INDEX_HEAD_DIM :].view(torch.float32)
+    vals = flat[: bs * cfg.index_head_dim].view(FP8).reshape(bs, cfg.index_head_dim)
+    scales = flat[bs * cfg.index_head_dim :].view(torch.float32)
     assert_fp8(vals[:num_tokens], q_ref, "indexer-K fp8")
     torch.testing.assert_close(scales[:num_tokens], s_ref, rtol=0, atol=0)
 
@@ -234,23 +250,24 @@ def test_fused_norm_rope(num_tokens: int, index_interleave: bool, mla_dtype: str
     assert (topk == -1).all(), "topk buffer not cleared on indexer layer"
 
 
-def test_fused_norm_rope_packed_indexer_block_stride():
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
+def test_fused_norm_rope_packed_indexer_block_stride(cfg: ModelConfig):
     """Indexer writes use their HMA block stride and independent slot map."""
     torch.manual_seed(7)
     dev = "cuda"
     num_tokens = block_size = 2
     pos = torch.arange(num_tokens, device=dev, dtype=torch.int64)
-    q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
-    kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
-    k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
-    qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
-    kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)
-    ik = torch.randn(num_tokens, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16)
-    ikw = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
-    ikb = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
-    cos_sin = make_cos_sin(32, ROPE_DIM, dev)
+    q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16)
+    k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16)
+    qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kvw = torch.randn(cfg.kv_lora, device=dev, dtype=torch.bfloat16)
+    ik = torch.randn(num_tokens, cfg.index_head_dim, device=dev, dtype=torch.bfloat16)
+    ikw = torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32)
+    ikb = torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32)
+    cos_sin = make_cos_sin(32, cfg.rope_dim, dev)
 
-    idx_row = INDEX_HEAD_DIM + INDEX_HEAD_DIM // 128 * 4
+    idx_row = cfg.index_head_dim + cfg.index_head_dim // 128 * 4
     packed_block_stride = block_size * idx_row + 64
     backing = torch.zeros(2 * packed_block_stride, device=dev, dtype=torch.uint8)
     idx_cache = torch.as_strided(
@@ -261,7 +278,7 @@ def test_fused_norm_rope_packed_indexer_block_stride():
     mla_cache = torch.zeros(
         1,
         block_size,
-        KV_LORA + ROPE_DIM,
+        cfg.kv_lora + cfg.rope_dim,
         device=dev,
         dtype=torch.bfloat16,
     )
@@ -296,15 +313,20 @@ def test_fused_norm_rope_packed_indexer_block_stride():
     q_ref, s_ref = ue8m0_quant(ik_ref)
     packed = idx_cache[1].reshape(-1)
     values = (
-        packed[: block_size * INDEX_HEAD_DIM].view(FP8).view(block_size, INDEX_HEAD_DIM)
+        packed[: block_size * cfg.index_head_dim]
+        .view(FP8)
+        .view(block_size, cfg.index_head_dim)
     )
-    scales = packed[block_size * INDEX_HEAD_DIM :].view(torch.float32)
+    scales = packed[block_size * cfg.index_head_dim :].view(torch.float32)
     assert_fp8(values, q_ref, "packed indexer-K fp8")
     torch.testing.assert_close(scales, s_ref, rtol=0, atol=0)
     assert (backing[block_size * idx_row : packed_block_stride] == 0).all()
 
 
-def test_fused_norm_rope_materializes_qk_without_local_cache_slots():
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
+def test_fused_norm_rope_materializes_qk_without_local_cache_slots(
+    cfg: ModelConfig,
+):
     """DCP non-owner ranks still need valid query shards for query AllGather."""
     torch.manual_seed(7)
     dev = "cuda"
@@ -312,20 +334,20 @@ def test_fused_norm_rope_materializes_qk_without_local_cache_slots():
     max_pos = 16
     pos = torch.arange(num_tokens, device=dev, dtype=torch.int64)
 
-    q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
-    kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
-    k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
-    qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
-    kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)
-    ik = torch.randn(num_tokens, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16)
-    ikw = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
-    ikb = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
-    cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
+    q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16)
+    k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16)
+    qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kvw = torch.randn(cfg.kv_lora, device=dev, dtype=torch.bfloat16)
+    ik = torch.randn(num_tokens, cfg.index_head_dim, device=dev, dtype=torch.bfloat16)
+    ikw = torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32)
+    ikb = torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32)
+    cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
 
     mla_cache = torch.zeros(
-        1, max_pos, KV_LORA + ROPE_DIM, device=dev, dtype=torch.bfloat16
+        1, max_pos, cfg.kv_lora + cfg.rope_dim, device=dev, dtype=torch.bfloat16
     )
-    idx_row = INDEX_HEAD_DIM + INDEX_HEAD_DIM // 128 * 4
+    idx_row = cfg.index_head_dim + cfg.index_head_dim // 128 * 4
     idx_cache = torch.zeros(1, max_pos, idx_row, device=dev, dtype=torch.uint8)
     no_local_slots = torch.full((num_tokens,), -1, device=dev, dtype=torch.int64)
     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():
     assert (topk == -1).all(), "topk buffer not cleared on non-owner rank"
 
 
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
 @pytest.mark.parametrize("num_tokens", [1, 17, 512])
-def test_fused_norm_rope_no_indexer(num_tokens: int):
+def test_fused_norm_rope_no_indexer(num_tokens: int, cfg: ModelConfig):
     """Shared (no-indexer) layer: q + kv/MLA only; top-k buffer untouched."""
     torch.manual_seed(1)
     dev = "cuda"
     max_pos = 8192
     pos = torch.arange(num_tokens, device=dev, dtype=torch.int64)
 
-    q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
-    kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
-    k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
-    qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
-    kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)
-    mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
+    q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16)
+    k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16)
+    qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kvw = torch.randn(cfg.kv_lora, device=dev, dtype=torch.bfloat16)
+    mla_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
 
     bs = max_pos
-    mla_cache = torch.zeros(1, bs, KV_LORA + ROPE_DIM, device=dev, dtype=torch.bfloat16)
+    mla_cache = torch.zeros(
+        1, bs, cfg.kv_lora + cfg.rope_dim, device=dev, dtype=torch.bfloat16
+    )
     slot = torch.arange(num_tokens, device=dev, dtype=torch.int64)
     topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32)
 
@@ -418,9 +443,9 @@ def test_fused_norm_rope_no_indexer(num_tokens: int):
 
     assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm (no-indexer)")
     cache = mla_cache[0, :num_tokens]
-    assert_bf16(cache[:, :KV_LORA], rms_norm(kv_c, kvw), "MLA kv (no-indexer)")
+    assert_bf16(cache[:, : cfg.kv_lora], rms_norm(kv_c, kvw), "MLA kv (no-indexer)")
     assert_bf16(
-        cache[:, KV_LORA:],
+        cache[:, cfg.kv_lora :],
         rope(k_pe.float(), pos, mla_cos_sin, interleave=True),
         "MLA k_pe (no-indexer)",
     )
@@ -428,36 +453,39 @@ def test_fused_norm_rope_no_indexer(num_tokens: int):
     assert (topk == 7).all(), "topk buffer should be untouched on shared layer"
 
 
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
 @pytest.mark.parametrize("has_indexer", [False, True])
-def test_fused_norm_rope_materializes_pcp_cache_inputs(has_indexer: bool):
+def test_fused_norm_rope_materializes_pcp_cache_inputs(
+    has_indexer: bool, cfg: ModelConfig
+):
     """PCP gets local normalized/rotated K rows without direct cache writes."""
     torch.manual_seed(6)
     dev = "cuda"
     num_tokens = 17
     max_pos = 8192
     pos = torch.arange(num_tokens, device=dev, dtype=torch.int64)
-    q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
-    kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
-    k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
-    qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
-    kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)
+    q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16)
+    k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16)
+    qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kvw = torch.randn(cfg.kv_lora, device=dev, dtype=torch.bfloat16)
     ik = (
-        torch.randn(num_tokens, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16)
+        torch.randn(num_tokens, cfg.index_head_dim, device=dev, dtype=torch.bfloat16)
         if has_indexer
         else None
     )
     ikw = (
-        torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
+        torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32)
         if has_indexer
         else None
     )
     ikb = (
-        torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
+        torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32)
         if has_indexer
         else None
     )
-    mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
-    idx_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev) if has_indexer else None
+    mla_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
+    idx_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev) if has_indexer else None
     q_out = torch.empty_like(q_c)
     kv_out = torch.empty_like(kv_c)
     kpe_out = torch.empty_like(k_pe)
@@ -503,8 +531,9 @@ def test_fused_norm_rope_materializes_pcp_cache_inputs(has_indexer: bool):
         assert_bf16(ik_out, ik_ref, "PCP indexer-K")
 
 
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
 @pytest.mark.parametrize("num_tokens", [1, 4, 17, 512])
-def test_fused_norm_rope_ds_mla(num_tokens: int):
+def test_fused_norm_rope_ds_mla(num_tokens: int, cfg: ModelConfig):
     """fp8_ds_mla MLA cache layout (FlashMLA sparse, bf16-query path; SM90/SM100).
 
     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):
     max_pos = 8192
     pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos
 
-    q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
-    kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
-    k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
-    qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
-    kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)
-    mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
+    q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16)
+    k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16)
+    qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kvw = torch.randn(cfg.kv_lora, device=dev, dtype=torch.bfloat16)
+    mla_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
 
     bs = max_pos
     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):
     tiles = kv_ref.view(num_tokens, 4, 128)
     ref_scale = torch.clamp(tiles.abs().amax(dim=-1) / FP8_MAX, min=1e-4)
     ref_scale = torch.exp2(torch.ceil(torch.log2(ref_scale)))
-    ref_nope = (tiles / ref_scale[..., None]).reshape(num_tokens, KV_LORA).to(FP8)
+    ref_nope = (tiles / ref_scale[..., None]).reshape(num_tokens, cfg.kv_lora).to(FP8)
 
     cache = mla_cache[0, :num_tokens]  # [N, 656] uint8
-    nope = cache[:, :KV_LORA].view(FP8)
-    scales = cache.view(torch.float32)[:, KV_LORA // 4 : KV_LORA // 4 + 4]
-    rope_off = KV_LORA // 2 + 8
-    rope_vals = cache.view(torch.bfloat16)[:, rope_off : rope_off + ROPE_DIM]
+    nope = cache[:, : cfg.kv_lora].view(FP8)
+    scales = cache.view(torch.float32)[:, cfg.kv_lora // 4 : cfg.kv_lora // 4 + 4]
+    rope_off = cfg.kv_lora // 2 + 8
+    rope_vals = cache.view(torch.bfloat16)[:, rope_off : rope_off + cfg.rope_dim]
 
     torch.testing.assert_close(scales, ref_scale, rtol=0, atol=0)
     assert_fp8(nope, ref_nope, "ds_mla NoPE fp8")
@@ -598,8 +627,9 @@ def nvfp4_sf_byte(s: torch.Tensor) -> torch.Tensor:
     not current_platform.is_device_capability(100),
     reason="nvfp4_ds_mla requires SM100 (Blackwell)",
 )
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
 @pytest.mark.parametrize("num_tokens", [1, 4, 17, 512])
-def test_fused_norm_rope_nvfp4_ds_mla(num_tokens: int):
+def test_fused_norm_rope_nvfp4_ds_mla(num_tokens: int, cfg: ModelConfig):
     """nvfp4_ds_mla MLA cache layout (FlashMLA sparse, SM100 only).
 
     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):
     max_pos = 8192
     pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos
 
-    q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
-    kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
+    q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16)
     # Tile 0 is deliberately low-magnitude: amax/6 lands between two e4m3
     # subnormals (spaced a flat 2^-9 there), so round-to-nearest picks a scale
     # 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):
     kv_c[:, :16] = (
         7.5 * (2.0**-9) * torch.linspace(0.2, 1.0, 16, device=dev, dtype=torch.bfloat16)
     )
-    k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
-    qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
-    kvw = torch.ones(KV_LORA, device=dev, dtype=torch.bfloat16)
-    mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
+    k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16)
+    qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16)
+    kvw = torch.ones(cfg.kv_lora, device=dev, dtype=torch.bfloat16)
+    mla_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
 
     bs = max_pos
     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):
     )
 
     ref_codes = quantize_to_e2m1(tiles / ref_scale.float().unsqueeze(-1))
-    ref_codes = ref_codes.reshape(num_tokens, KV_LORA)
+    ref_codes = ref_codes.reshape(num_tokens, cfg.kv_lora)
 
     cache = mla_cache[0, :num_tokens]  # [N, 352] uint8
-    packed = cache[:, : KV_LORA // 2]
+    packed = cache[:, : cfg.kv_lora // 2]
     got_codes = torch.stack([packed & 0xF, packed >> 4], dim=-1)
-    got_codes = got_codes.reshape(num_tokens, KV_LORA)  # low nibble = even elem
+    got_codes = got_codes.reshape(num_tokens, cfg.kv_lora)  # low nibble = even elem
     torch.testing.assert_close(got_codes, ref_codes, rtol=0, atol=0)
 
-    got_rope = cache[:, KV_LORA // 2 : KV_LORA // 2 + ROPE_DIM].view(FP8)
+    got_rope = cache[:, cfg.kv_lora // 2 : cfg.kv_lora // 2 + cfg.rope_dim].view(FP8)
     assert_fp8(got_rope, kpe_ref.to(FP8), "nvfp4_ds_mla RoPE e4m3")
 
     perm = nvfp4_sf_byte(torch.arange(32, device=dev))
-    got_scale = cache[:, KV_LORA // 2 + ROPE_DIM :].view(FP8)[:, perm]
+    got_scale = cache[:, cfg.kv_lora // 2 + cfg.rope_dim :].view(FP8)[:, perm]
     torch.testing.assert_close(got_scale.float(), ref_scale.float(), rtol=0, atol=0)
 
     # No indexer on this call: top-k buffer must be untouched.
@@ -737,39 +767,34 @@ def test_fused_norm_rope_supports_large_token_count():
 # ── fused_q ──────────────────────────────────────────────────────────────────
 
 
-@pytest.mark.parametrize(
-    ("num_tokens", "num_q_heads"),
-    [
-        (1, NUM_HEADS),
-        (4, NUM_HEADS),
-        (17, NUM_HEADS),
-        (512, NUM_HEADS),
-        (4096, NUM_HEADS),
-        (17, 64),
-    ],
-)
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
+@pytest.mark.parametrize("num_tokens", [1, 4, 17, 512, 4096])
 @pytest.mark.parametrize("index_interleave", [True, False])
-def test_fused_q(num_tokens: int, num_q_heads: int, index_interleave: bool):
+def test_fused_q(num_tokens: int, index_interleave: bool, cfg: ModelConfig):
     torch.manual_seed(2)
     dev = "cuda"
     max_pos = 8192
     pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos
 
     q_pe = torch.randn(
-        num_tokens, num_q_heads, ROPE_DIM, device=dev, dtype=torch.bfloat16
+        num_tokens, cfg.num_heads, cfg.rope_dim, device=dev, dtype=torch.bfloat16
     )
     ql_nope = torch.randn(
-        num_tokens, num_q_heads, KV_LORA, device=dev, dtype=torch.bfloat16
+        num_tokens, cfg.num_heads, cfg.kv_lora, device=dev, dtype=torch.bfloat16
     )
     index_q = torch.randn(
-        num_tokens, INDEX_HEADS, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16
+        num_tokens,
+        cfg.index_heads,
+        cfg.index_head_dim,
+        device=dev,
+        dtype=torch.bfloat16,
     )
-    index_w = torch.randn(num_tokens, INDEX_HEADS, device=dev, dtype=torch.float32)
+    index_w = torch.randn(num_tokens, cfg.index_heads, device=dev, dtype=torch.float32)
     q_scale = torch.tensor([0.37], device=dev, dtype=torch.float32)
-    softmax_scale = INDEX_HEAD_DIM**-0.5
-    head_scale = INDEX_HEADS**-0.5
-    q_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)  # q_pe: interleaved
-    idx_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
+    softmax_scale = cfg.index_head_dim**-0.5
+    head_scale = cfg.index_heads**-0.5
+    q_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)  # q_pe: interleaved
+    idx_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
 
     iq_fp8, iw_out, mqa = K.fused_q(
         pos,
@@ -791,18 +816,18 @@ def test_fused_q(num_tokens: int, num_q_heads: int, index_interleave: bool):
     mqa_nope_ref = (ql_nope.float() / s).to(FP8)
     qpe_ref = rope(
         q_pe.float(),
-        pos.unsqueeze(-1).expand(num_tokens, num_q_heads),
+        pos.unsqueeze(-1).expand(num_tokens, cfg.num_heads),
         q_cos_sin,
         interleave=True,
     )
     mqa_pe_ref = (qpe_ref / s).to(FP8)
-    assert_fp8(mqa[:, :, :KV_LORA], mqa_nope_ref, "mqa ql_nope")
-    assert_fp8(mqa[:, :, KV_LORA:], mqa_pe_ref, "mqa q_pe")
+    assert_fp8(mqa[:, :, : cfg.kv_lora], mqa_nope_ref, "mqa ql_nope")
+    assert_fp8(mqa[:, :, cfg.kv_lora :], mqa_pe_ref, "mqa q_pe")
 
     # Indexer-Q: RoPE + UE8M0 fp8 quant; index weights fold in q-scale.
     iq_ref = rope(
         index_q.float(),
-        pos.unsqueeze(-1).expand(num_tokens, INDEX_HEADS),
+        pos.unsqueeze(-1).expand(num_tokens, cfg.index_heads),
         idx_cos_sin,
         interleave=index_interleave,
     )
@@ -812,20 +837,21 @@ def test_fused_q(num_tokens: int, num_q_heads: int, index_interleave: bool):
     torch.testing.assert_close(iw_out, iw_ref, rtol=1e-3, atol=1e-3)
 
 
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
 @pytest.mark.parametrize("num_tokens", [1, 17, 512])
-def test_fused_q_no_indexer(num_tokens: int):
+def test_fused_q_no_indexer(num_tokens: int, cfg: ModelConfig):
     torch.manual_seed(3)
     dev = "cuda"
     max_pos = 8192
     pos = torch.arange(num_tokens, device=dev, dtype=torch.int64)
     q_pe = torch.randn(
-        num_tokens, NUM_HEADS, ROPE_DIM, device=dev, dtype=torch.bfloat16
+        num_tokens, cfg.num_heads, cfg.rope_dim, device=dev, dtype=torch.bfloat16
     )
     ql_nope = torch.randn(
-        num_tokens, NUM_HEADS, KV_LORA, device=dev, dtype=torch.bfloat16
+        num_tokens, cfg.num_heads, cfg.kv_lora, device=dev, dtype=torch.bfloat16
     )
     q_scale = torch.tensor([0.5], device=dev, dtype=torch.float32)
-    q_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
+    q_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
     _, _, mqa = K.fused_q(
         pos,
         q_pe,
@@ -841,19 +867,20 @@ def test_fused_q_no_indexer(num_tokens: int):
         index_rope_interleave=False,
     )
     s = q_scale.item()
-    assert_fp8(mqa[:, :, :KV_LORA], (ql_nope.float() / s).to(FP8), "mqa ql_nope")
+    assert_fp8(mqa[:, :, : cfg.kv_lora], (ql_nope.float() / s).to(FP8), "mqa ql_nope")
     qpe_ref = rope(
         q_pe.float(),
-        pos.unsqueeze(-1).expand(num_tokens, NUM_HEADS),
+        pos.unsqueeze(-1).expand(num_tokens, cfg.num_heads),
         q_cos_sin,
         interleave=True,
     )
-    assert_fp8(mqa[:, :, KV_LORA:], (qpe_ref / s).to(FP8), "mqa q_pe")
+    assert_fp8(mqa[:, :, cfg.kv_lora :], (qpe_ref / s).to(FP8), "mqa q_pe")
 
 
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
 @pytest.mark.parametrize("num_tokens", [1, 17, 512])
 @pytest.mark.parametrize("has_indexer", [True, False])
-def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool):
+def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool, cfg: ModelConfig):
     """bf16-query path (FlashMLA sparse, SM90/SM100): only the RoPE'd q_pe is
     produced (bf16, unquantized); ql_nope is consumed directly by the caller."""
     torch.manual_seed(6)
@@ -862,21 +889,27 @@ def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool):
     pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos
 
     q_pe = torch.randn(
-        num_tokens, NUM_HEADS, ROPE_DIM, device=dev, dtype=torch.bfloat16
+        num_tokens, cfg.num_heads, cfg.rope_dim, device=dev, dtype=torch.bfloat16
     )
     ql_nope = torch.randn(
-        num_tokens, NUM_HEADS, KV_LORA, device=dev, dtype=torch.bfloat16
+        num_tokens, cfg.num_heads, cfg.kv_lora, device=dev, dtype=torch.bfloat16
     )
     q_scale = torch.tensor([0.37], device=dev, dtype=torch.float32)
-    q_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
+    q_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
 
     index_q = index_w = idx_cos_sin = None
     if has_indexer:
         index_q = torch.randn(
-            num_tokens, INDEX_HEADS, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16
+            num_tokens,
+            cfg.index_heads,
+            cfg.index_head_dim,
+            device=dev,
+            dtype=torch.bfloat16,
         )
-        index_w = torch.randn(num_tokens, INDEX_HEADS, device=dev, dtype=torch.float32)
-        idx_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
+        index_w = torch.randn(
+            num_tokens, cfg.index_heads, device=dev, dtype=torch.float32
+        )
+        idx_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
 
     iq_fp8, iw_out, q_pe_out = K.fused_q(
         pos,
@@ -887,8 +920,8 @@ def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool):
         ql_nope,
         q_scale,
         index_w,
-        INDEX_HEAD_DIM**-0.5,
-        INDEX_HEADS**-0.5,
+        cfg.index_head_dim**-0.5,
+        cfg.index_heads**-0.5,
         has_indexer=has_indexer,
         index_rope_interleave=False,
         quantize_mqa=False,
@@ -896,10 +929,10 @@ def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool):
 
     # MQA query: only the RoPE'd q_pe, bf16, unquantized.
     assert q_pe_out.dtype == torch.bfloat16
-    assert q_pe_out.shape == (num_tokens, NUM_HEADS, ROPE_DIM)
+    assert q_pe_out.shape == (num_tokens, cfg.num_heads, cfg.rope_dim)
     qpe_ref = rope(
         q_pe.float(),
-        pos.unsqueeze(-1).expand(num_tokens, NUM_HEADS),
+        pos.unsqueeze(-1).expand(num_tokens, cfg.num_heads),
         q_cos_sin,
         interleave=True,
     )
@@ -910,13 +943,15 @@ def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool):
         assert index_q is not None
         iq_ref = rope(
             index_q.float(),
-            pos.unsqueeze(-1).expand(num_tokens, INDEX_HEADS),
+            pos.unsqueeze(-1).expand(num_tokens, cfg.index_heads),
             idx_cos_sin,
             interleave=False,
         )
         q_ref, scale_ref = ue8m0_quant(iq_ref)
         assert_fp8(iq_fp8, q_ref, "indexer-Q fp8 (bf16-query path)")
-        iw_ref = index_w * scale_ref * (INDEX_HEAD_DIM**-0.5) * (INDEX_HEADS**-0.5)
+        iw_ref = (
+            index_w * scale_ref * (cfg.index_head_dim**-0.5) * (cfg.index_heads**-0.5)
+        )
         torch.testing.assert_close(iw_out, iw_ref, rtol=1e-3, atol=1e-3)
 
 
@@ -956,21 +991,22 @@ def test_fused_q_triton_supports_large_token_count():
 # ── fused_eh_norm (MTP) ──────────────────────────────────────────────────────
 
 
+@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
 @pytest.mark.parametrize("num_tokens", [1, 4, 17, 512])
-def test_fused_eh_norm(num_tokens: int):
+def test_fused_eh_norm(num_tokens: int, cfg: ModelConfig):
     torch.manual_seed(4)
     dev = "cuda"
     # Mix in a position-0 token to exercise the embeds-zeroing branch.
     pos = torch.arange(num_tokens, device=dev, dtype=torch.int64)
     pos[0] = 0
-    embeds = torch.randn(num_tokens, HIDDEN, device=dev, dtype=torch.bfloat16)
-    prev = torch.randn(num_tokens, HIDDEN, device=dev, dtype=torch.bfloat16)
-    ew = torch.randn(HIDDEN, device=dev, dtype=torch.bfloat16)
-    hw = torch.randn(HIDDEN, device=dev, dtype=torch.bfloat16)
+    embeds = torch.randn(num_tokens, cfg.hidden, device=dev, dtype=torch.bfloat16)
+    prev = torch.randn(num_tokens, cfg.hidden, device=dev, dtype=torch.bfloat16)
+    ew = torch.randn(cfg.hidden, device=dev, dtype=torch.bfloat16)
+    hw = torch.randn(cfg.hidden, device=dev, dtype=torch.bfloat16)
 
     out = K.fused_eh_norm(pos, embeds, prev, ew, hw, EPS)
 
     masked = torch.where(pos.unsqueeze(-1) == 0, torch.zeros_like(embeds), embeds)
     ref = torch.cat([rms_norm(masked, ew), rms_norm(prev, hw)], dim=-1)
-    assert out.shape == (num_tokens, 2 * HIDDEN)
+    assert out.shape == (num_tokens, 2 * cfg.hidden)
     assert_bf16(out, ref, "eh_norm")
diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py
index c0b90e0be6..e5ce96cf24 100644
--- a/vllm/_aiter_ops.py
+++ b/vllm/_aiter_ops.py
@@ -3406,7 +3406,7 @@ class rocm_aiter_ops:
         X: torch.Tensor,
         W: torch.Tensor,
         w_scale: torch.Tensor,
-        Y: torch.Tensor,
+        Y: torch.Tensor | None = None,
         transpose_bm: bool | None = False,
         prequant: bool | None = False,
         y_scale: torch.Tensor | None = None,
diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py
index c37801e887..1b8fa13a51 100644
--- a/vllm/model_executor/layers/attention/mla_attention.py
+++ b/vllm/model_executor/layers/attention/mla_attention.py
@@ -1025,9 +1025,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                 mqa_q_pe = mqa_pe_padded
 
             if self.is_aiter_triton_fp4_bmm_enabled:
-                from aiter.ops.triton.batched_gemm_a16wfp4 import batched_gemm_a16wfp4
-
-                mqa_ql_nope = batched_gemm_a16wfp4(
+                mqa_ql_nope = rocm_aiter_ops.batched_gemm_a16wfp4(
                     mqa_q_nope,
                     self.W_K,
                     self.W_K_scale,
diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py
index ed702400be..d5728e0c33 100644
--- a/vllm/model_executor/models/deepseek_v2.py
+++ b/vllm/model_executor/models/deepseek_v2.py
@@ -665,6 +665,11 @@ class DeepseekV32IndexerCache(torch.nn.Module, AttentionLayerBase):
 
     def forward(self): ...
 
+    @property
+    def uses_shuffled_layout(self) -> bool:
+        """Whether this cache's reader expects the shuffled value layout."""
+        return False
+
     def get_attn_backend(self) -> type[AttentionBackend]:
         return DeepseekV32IndexerBackend
 
diff --git a/vllm/models/deepseek_v32/__init__.py b/vllm/models/deepseek_v32/__init__.py
index 0ac154505f..9456239809 100644
--- a/vllm/models/deepseek_v32/__init__.py
+++ b/vllm/models/deepseek_v32/__init__.py
@@ -7,9 +7,9 @@ MLA + a "lightning indexer" that selects the top-k tokens for a sparse MLA
 attend. The same model code serves any DSA checkpoint, including GLM-5.2
 (``glm_moe_dsa``), which reuses this architecture.
 
-The CUDA implementation selects capability-specific kernels internally and
-falls back when an optimization is unavailable. Other platforms use the
-generic implementation by default.
+The CUDA and ROCm implementations select capability-specific kernels
+internally and fall back when an optimization is unavailable. Other platforms
+use the generic implementation by default.
 """
 
 from vllm.platforms import current_platform
@@ -20,8 +20,12 @@ if current_platform.is_cuda():
     from .nvidia.model import DeepseekV32ForCausalLM
     from .nvidia.model import DeepseekV32ForCausalLM as GlmMoeDsaForCausalLM
     from .nvidia.mtp import DeepseekV32MTP
+elif current_platform.is_rocm():
+    from .amd.model import DeepseekV32ForCausalLM
+    from .amd.model import DeepseekV32ForCausalLM as GlmMoeDsaForCausalLM
+    from .amd.mtp import DeepseekV32MTP
 else:
-    # ROCm, XPU, and CPU keep the generic implementation.
+    # XPU and CPU keep the generic implementation.
     from vllm.model_executor.models.deepseek_mtp import DeepSeekMTP as DeepseekV32MTP
     from vllm.model_executor.models.deepseek_v2 import (
         DeepseekV3ForCausalLM as DeepseekV32ForCausalLM,
diff --git a/vllm/models/deepseek_v32/amd/rocm.py b/vllm/models/deepseek_v32/amd/rocm.py
index 6d8e7ea855..b9ac31f8be 100644
--- a/vllm/models/deepseek_v32/amd/rocm.py
+++ b/vllm/models/deepseek_v32/amd/rocm.py
@@ -3,12 +3,14 @@
 
 import torch
 
+from vllm._aiter_ops import rocm_aiter_ops
 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
 from vllm.forward_context import get_forward_context
 from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
 from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache
 from vllm.models.deepseek_v32.attention import DeepseekV32Attention, DeepseekV32Indexer
 from vllm.models.deepseek_v32.common.kernels import fused_norm_rope, fused_q
+from vllm.platforms import current_platform
 from vllm.utils.torch_utils import is_quantized_kv_cache
 from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerBackend
 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 (
 )
 
 
-class DeepseekV32MLASparseBackend(ROCMAiterMLASparseBackend):
-    @staticmethod
-    def get_supported_kernel_block_sizes() -> list:
-        return [16, 32]
-
-
-class DeepseekV32ROCmIndexerBackend(DeepseekV32IndexerBackend):
-    @staticmethod
-    def get_supported_kernel_block_sizes() -> list:
-        return [16, 32]
-
-
 class DeepseekV32ROCmIndexerCache(DeepseekV32IndexerCache):
     def get_attn_backend(self):
-        return DeepseekV32ROCmIndexerBackend
+        return DeepseekV32IndexerBackend
+
+    @property
+    def uses_shuffled_layout(self) -> bool:
+        # aiter's gather/insert pair shuffles the cache above block size 1:
+        # [n_blocks, blk/16, head_dim/16, 16, 16] instead of [n_blocks, blk, head_dim].
+        return self.kv_cache.ndim == 3 and self.kv_cache.shape[1] != 1
 
 
 class DeepseekV32ROCmIndexer(DeepseekV32Indexer):
@@ -46,7 +42,7 @@ class DeepseekV32MLAAttention(DeepseekV32Attention):
             config,
             prefix,
             topk_indices_buffer,
-            attn_backend=DeepseekV32MLASparseBackend,
+            attn_backend=ROCMAiterMLASparseBackend,
         )
 
         self.indexer_op: SparseAttnIndexer | None = None
@@ -100,14 +96,10 @@ class DeepseekV32MLAAttention(DeepseekV32Attention):
         q_nope = q_nope.transpose(0, 1)  # (N, tokens, P)
 
         if self.is_aiter_triton_fp4_bmm_enabled:
-            from aiter.ops.triton.batched_gemm_a16wfp4 import batched_gemm_a16wfp4
-
-            ql_nope = batched_gemm_a16wfp4(
+            ql_nope = rocm_aiter_ops.batched_gemm_a16wfp4(
                 q_nope, self.W_K, self.W_K_scale, transpose_bm=True, prequant=True
             )
         elif self.is_aiter_triton_fp8_bmm_enabled:
-            from vllm._aiter_ops import rocm_aiter_ops
-
             ql_nope = rocm_aiter_ops.triton_fp8_bmm(
                 q_nope, self.W_K, self.W_K_scale, group_size=128, transpose_bm=True
             )
@@ -152,14 +144,10 @@ class DeepseekV32MLAAttention(DeepseekV32Attention):
         )
 
         if self.is_aiter_triton_fp4_bmm_enabled:
-            from aiter.ops.triton.batched_gemm_a16wfp4 import batched_gemm_a16wfp4
-
-            batched_gemm_a16wfp4(
+            rocm_aiter_ops.batched_gemm_a16wfp4(
                 x, self.W_V, self.W_V_scale, out_view, transpose_bm=True, prequant=True
             )
         elif self.is_aiter_triton_fp8_bmm_enabled:
-            from vllm._aiter_ops import rocm_aiter_ops
-
             rocm_aiter_ops.triton_fp8_bmm(
                 x,
                 self.W_V,
@@ -202,6 +190,7 @@ class DeepseekV32MLAAttention(DeepseekV32Attention):
             indexer_k_norm_eps = self.indexer.k_norm.eps
             indexer_k_rope_cos_sin_cache = self.indexer_rope_emb.cos_sin_cache
             indexer_k_cache = self.indexer.k_cache.kv_cache
+            indexer_cache_shuffled = self.indexer.k_cache.uses_shuffled_layout
             indexer_softmax_scale = self.indexer.softmax_scale
             indexer_n_head_scale = self.indexer.n_head**-0.5
         else:
@@ -211,6 +200,7 @@ class DeepseekV32MLAAttention(DeepseekV32Attention):
             indexer_k_norm_eps = 1e-6
             indexer_k_rope_cos_sin_cache = None
             indexer_k_cache = None
+            indexer_cache_shuffled = False
             indexer_softmax_scale = 0.0
             indexer_n_head_scale = 0.0
 
@@ -241,6 +231,7 @@ class DeepseekV32MLAAttention(DeepseekV32Attention):
             self.topk_indices_buffer,
             slot_mapping=mla_slot,
             indexer_k_cache=indexer_k_cache,
+            indexer_cache_shuffled=indexer_cache_shuffled,
             mla_kv_cache=mla_kv_cache,
             mla_kv_cache_dtype=self.kv_cache_dtype,
             mla_k_scale=mla_k_scale,
@@ -282,7 +273,7 @@ class DeepseekV32MLAAttention(DeepseekV32Attention):
         num_actual = attn_metadata.num_actual_tokens  # type: ignore[attr-defined]
         kv_cache = self.kv_cache
         if self._fp8_kv_needs_view:
-            kv_cache = kv_cache.view(torch.float8_e4m3fn)
+            kv_cache = kv_cache.view(current_platform.fp8_dtype())
 
         q_for_attn = self._build_q_for_attn(ql_nope, mqa_q, num_actual)
         attn_out, _ = self.impl.forward_mqa(  # type: ignore[attr-defined]
diff --git a/vllm/models/deepseek_v32/attention.py b/vllm/models/deepseek_v32/attention.py
index b8e70e4ad3..e6fb00434a 100644
--- a/vllm/models/deepseek_v32/attention.py
+++ b/vllm/models/deepseek_v32/attention.py
@@ -327,6 +327,7 @@ class DeepseekV32Attention(MLAAttention):
             indexer_k_norm_bias = self.indexer.k_norm.bias
             indexer_k_norm_eps = self.indexer.k_norm.eps
             indexer_k_rope_cos_sin_cache = self.indexer_rope_emb.cos_sin_cache
+            indexer_cache_shuffled = self.indexer.k_cache.uses_shuffled_layout
             indexer_k_cache = None if self.use_pcp else self.indexer.k_cache.kv_cache
             index_k_out = torch.empty_like(index_k) if self.use_pcp else None
             indexer_softmax_scale = self.indexer.softmax_scale
@@ -338,6 +339,7 @@ class DeepseekV32Attention(MLAAttention):
             indexer_k_norm_eps = 1e-6
             indexer_k_rope_cos_sin_cache = None
             indexer_k_cache = None
+            indexer_cache_shuffled = False
             index_k_out = None
             indexer_softmax_scale = 0.0
             indexer_n_head_scale = 0.0
@@ -377,6 +379,7 @@ class DeepseekV32Attention(MLAAttention):
             slot_mapping=mla_slot,
             indexer_slot_mapping=indexer_slot,
             indexer_k_cache=indexer_k_cache,
+            indexer_cache_shuffled=indexer_cache_shuffled,
             mla_kv_cache=mla_kv_cache,
             mla_kv_cache_dtype=self.kv_cache_dtype,
             mla_k_scale=mla_k_scale,
@@ -544,7 +547,7 @@ class DeepseekV32Attention(MLAAttention):
             return
 
         if self._fp8_kv_needs_view:
-            kv_cache = kv_cache.view(torch.float8_e4m3fn)
+            kv_cache = kv_cache.view(current_platform.fp8_dtype())
         if self._fp8_query:
             # FlashInfer sparse: single packed fp8 query.
             mqa_q_arg: torch.Tensor | tuple[torch.Tensor, torch.Tensor] = mqa_q[
diff --git a/vllm/models/deepseek_v32/common/kernels.py b/vllm/models/deepseek_v32/common/kernels.py
index 1808d12f64..e0259c46fa 100644
--- a/vllm/models/deepseek_v32/common/kernels.py
+++ b/vllm/models/deepseek_v32/common/kernels.py
@@ -12,6 +12,18 @@ from vllm.utils.torch_utils import is_quantized_kv_cache
 # has_indexer=False path so the indexer args don't allocate every call.
 _DUMMY_CACHE: dict[tuple, torch.Tensor] = {}
 
+# Tile shape of the indexer-K cache's shuffled layout, used when the cache reports
+# uses_shuffled_layout. The block tile is a token count; the head tile is a byte
+# count, which the kernel converts to cache elements.
+_INDEXER_CACHE_BLOCK_TILE = 16
+_INDEXER_CACHE_HEAD_TILE_BYTES = 16
+
+# FNUZ (e4m3fnuz) on gfx942, OCP (e4m3fn) elsewhere. 224.0 keeps the ue8m0
+# power-of-two scale inside the FNUZ range after its ceil.
+_FP8_DTYPE = current_platform.fp8_dtype()
+_USE_FNUZ = torch.float8_e4m3fnuz == _FP8_DTYPE
+_FP8_MAX = 224.0 if _USE_FNUZ else 448.0
+
 
 def _dummy(shape: tuple, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
     key = (shape, dtype, device)
@@ -47,16 +59,17 @@ def _get_cos_sin(
 
 
 @triton.jit
-def _fp8_ue8m0_quantize(vals):
+def _fp8_ue8m0_quantize(vals, FP8_MAX: tl.constexpr, USE_FNUZ: tl.constexpr):
     """Quantize float32 values to FP8 E4M3 with a ue8m0 (power-of-2) scale.
 
     Returns (fp8_vals, scale) so the caller can store them or reuse the scale.
     """
     vals = vals.to(tl.float32)
     amax = tl.max(tl.abs(vals))
-    scale = tl.div_rn(tl.maximum(amax, 1e-4), 448.0)
+    scale = tl.div_rn(tl.maximum(amax, 1e-4), FP8_MAX)
     scale = tl.math.exp2(tl.math.ceil(tl.math.log2(scale)))
-    fp8_vals = tl.div_rn(vals, scale).to(tl.float8e4nv)
+    fp8_dtype = tl.float8e4b8 if USE_FNUZ else tl.float8e4nv
+    fp8_vals = tl.div_rn(vals, scale).to(fp8_dtype)
     return fp8_vals, scale
 
 
@@ -71,18 +84,32 @@ def _fp8_quant_and_cache_write(
     cache_block_stride,
     offsets,
     HEAD_DIM: tl.constexpr,
+    SHUFFLE: tl.constexpr,
+    BLOCK_TILE: tl.constexpr,
+    HEAD_TILE: tl.constexpr,
+    FP8_MAX: tl.constexpr,
+    USE_FNUZ: tl.constexpr,
 ):
-    k_fp8, scale = _fp8_ue8m0_quantize(vals)
+    k_fp8, scale = _fp8_ue8m0_quantize(vals, FP8_MAX, USE_FNUZ)
 
     block_idx = slot_idx // cache_block_size
     block_offset = slot_idx % cache_block_size
     block_start = block_idx * cache_block_stride
 
-    tl.store(
-        kv_cache_ptr + block_start + block_offset * HEAD_DIM + offsets,
-        k_fp8,
-        mask=mask,
-    )
+    # Shuffled layout: [blk/BLOCK_TILE, HEAD_DIM/HEAD_TILE, BLOCK_TILE, HEAD_TILE],
+    # so one contiguous run holds BLOCK_TILE tokens x HEAD_TILE bytes and the
+    # reader's coalesced loads land on the bytes we wrote.
+    if SHUFFLE:
+        value_off = (
+            block_offset // BLOCK_TILE * BLOCK_TILE * HEAD_DIM
+            + block_offset % BLOCK_TILE * HEAD_TILE
+            + offsets // HEAD_TILE * BLOCK_TILE * HEAD_TILE
+            + offsets % HEAD_TILE
+        )
+    else:
+        value_off = block_offset * HEAD_DIM + offsets
+
+    tl.store(kv_cache_ptr + block_start + value_off, k_fp8, mask=mask)
     scale_byte_off = block_start + cache_block_size * HEAD_DIM + block_offset * 4
     tl.store(kv_cache_scale_ptr + scale_byte_off // 4, scale)
 
@@ -137,6 +164,9 @@ def _fused_norm_rope_kernel(
     indexer_cache_scale_ptr,
     indexer_cache_block_size,
     indexer_cache_stride,
+    INDEXER_CACHE_SHUFFLE: tl.constexpr,
+    INDEXER_CACHE_BLOCK_TILE: tl.constexpr,
+    INDEXER_CACHE_HEAD_TILE: tl.constexpr,
     # MLA KV cache (concat kv_c_normed + k_pe_roped, uses slot_mapping_ptr)
     mla_cache_ptr,
     mla_cache_block_stride,
@@ -161,12 +191,15 @@ def _fused_norm_rope_kernel(
     HAS_INDEXER: tl.constexpr,
     INDEX_ROPE_INTERLEAVE: tl.constexpr,
     USE_PDL: tl.constexpr,
+    FP8_MAX: tl.constexpr,
+    USE_FNUZ: tl.constexpr,
 ):
     tok_idx = tl.program_id(0).to(tl.int64)
     pid = tl.program_id(1)
     if USE_PDL:
         tl.extra.cuda.gdc_wait()
         tl.extra.cuda.gdc_launch_dependents()
+    fp8_dtype = tl.float8e4b8 if USE_FNUZ else tl.float8e4nv
     if pid == 3:
         if not HAS_INDEXER:
             # Shared layer: reuse the previous indexer layer's top-k; do not
@@ -327,9 +360,11 @@ def _fused_norm_rope_kernel(
                 )
                 kv_2d = tl.reshape(kv_c, (MLA_NUM_TILES, MLA_TILE_DIM))
                 tile_amax = tl.max(tl.abs(kv_2d), axis=1, keep_dims=True)
-                tile_scale = tl.maximum(tile_amax * (1.0 / 448.0), 1e-4)
+                # scale = amax / FP8_MAX (fp8 e4m3 max), matching the reference
+                # concat_and_cache_ds_mla kernel; floored to FLT_MIN.
+                tile_scale = tl.maximum(tile_amax * (1.0 / FP8_MAX), 1e-4)
                 tile_scale = tl.math.exp2(tl.math.ceil(tl.math.log2(tile_scale)))
-                kv_c_fp8 = tl.reshape((kv_2d / tile_scale).to(tl.float8e4nv), (KV_DIM,))
+                kv_c_fp8 = tl.reshape((kv_2d / tile_scale).to(fp8_dtype), (KV_DIM,))
                 tl.store(mla_cache_ptr + byte_base + kv_block, kv_c_fp8)
                 tile_off = tl.arange(0, MLA_NUM_TILES)
                 tl.store(
@@ -349,16 +384,16 @@ def _fused_norm_rope_kernel(
             # kv_c_normed (KV_DIM elements)
             if MLA_CACHE_FP8:
                 scale = tl.load(mla_cache_scale_ptr)
-                kv_c_fp8 = (kv_c.to(tl.float32) / scale).to(tl.float8e4nv)
+                kv_c_fp8 = (kv_c.to(tl.float32) / scale).to(fp8_dtype)
                 tl.store(dst + kv_block, kv_c_fp8)
             else:
                 tl.store(dst + kv_block, kv_c)
             # k_pe_roped (from registers, interleaved layout)
             if MLA_CACHE_FP8:
-                tl.store(dst + KV_DIM + dim_off * 2, (r1 / scale).to(tl.float8e4nv))
+                tl.store(dst + KV_DIM + dim_off * 2, (r1 / scale).to(fp8_dtype))
                 tl.store(
                     dst + KV_DIM + dim_off * 2 + 1,
-                    (r2 / scale).to(tl.float8e4nv),
+                    (r2 / scale).to(fp8_dtype),
                 )
             else:
                 tl.store(dst + KV_DIM + dim_off * 2, r1)
@@ -463,6 +498,11 @@ def _fused_norm_rope_kernel(
                 indexer_cache_stride,
                 index_k_block,
                 INDEX_K_DIM,
+                INDEXER_CACHE_SHUFFLE,
+                INDEXER_CACHE_BLOCK_TILE,
+                INDEXER_CACHE_HEAD_TILE,
+                FP8_MAX,
+                USE_FNUZ,
             )
 
 
@@ -486,6 +526,7 @@ def fused_norm_rope(
     slot_mapping: torch.Tensor | None = None,
     indexer_slot_mapping: torch.Tensor | None = None,
     indexer_k_cache: torch.Tensor | None = None,
+    indexer_cache_shuffled: bool = False,
     mla_kv_cache: torch.Tensor | None = None,
     mla_kv_cache_dtype: str = "auto",
     mla_k_scale: torch.Tensor | None = None,
@@ -532,13 +573,28 @@ def fused_norm_rope(
             indexer_slot_mapping = slot_mapping
         idx_cache_scale_view = indexer_k_cache.view(torch.uint8).view(torch.float32)
         idx_cache_block_size = indexer_k_cache.shape[1]
+        idx_cache_stride = indexer_k_cache.shape[2]
+        # The caller's cache reports whether its reader expects the shuffled
+        # layout; see DeepseekV32IndexerCache.uses_shuffled_layout.
+        idx_cache_shuffle = indexer_cache_shuffled
+        if idx_cache_shuffle:
+            assert idx_cache_block_size % _INDEXER_CACHE_BLOCK_TILE == 0, (
+                f"indexer K cache block size {idx_cache_block_size} must be a "
+                f"multiple of {_INDEXER_CACHE_BLOCK_TILE} for the shuffled layout"
+            )
         idx_cache_stride = indexer_k_cache.stride(0)
         if indexer_k_cache.dtype == torch.uint8:
-            indexer_k_cache = indexer_k_cache.view(torch.float8_e4m3fn)
+            indexer_k_cache = indexer_k_cache.view(_FP8_DTYPE)
+        # The head tile is a byte count; the kernel indexes the cache in elements.
+        idx_cache_head_tile = (
+            _INDEXER_CACHE_HEAD_TILE_BYTES // indexer_k_cache.element_size()
+        )
     else:
         idx_cache_scale_view = None
         idx_cache_block_size = 1
         idx_cache_stride = 0
+        idx_cache_shuffle = False
+        idx_cache_head_tile = _INDEXER_CACHE_HEAD_TILE_BYTES
 
     # --- MLA KV cache setup ---
     mla_cache_nvfp4 = mla_kv_cache_dtype == "nvfp4_ds_mla"
@@ -561,12 +617,12 @@ def fused_norm_rope(
             mla_entry_stride = u8_cache.stride(1)
             mla_ds_scale_view = u8_cache.view(torch.float32)
             mla_ds_rope_view = u8_cache.view(torch.bfloat16)
-            mla_kv_cache = u8_cache.view(torch.float8_e4m3fn)
+            mla_kv_cache = u8_cache.view(_FP8_DTYPE)
         else:
             mla_block_stride = mla_kv_cache.stride(0)
             mla_entry_stride = mla_kv_cache.stride(1)
             if mla_cache_fp8 and mla_kv_cache.dtype == torch.uint8:
-                mla_kv_cache = mla_kv_cache.view(torch.float8_e4m3fn)
+                mla_kv_cache = mla_kv_cache.view(_FP8_DTYPE)
         if mla_k_scale is None:
             mla_k_scale = torch.ones(1, dtype=torch.float32, device=device)
     else:
@@ -592,6 +648,9 @@ def fused_norm_rope(
         assert index_k_out.shape == index_k.shape
         index_k_out_stride = index_k_out.stride(0)
     use_pdl = current_platform.is_arch_support_pdl()
+    pdl_kwargs = (
+        {"USE_PDL": use_pdl, "launch_pdl": True} if use_pdl else {"USE_PDL": use_pdl}
+    )
     _fused_norm_rope_kernel[(num_tokens, 4)](
         positions,
         # Q RMS norm
@@ -639,6 +698,9 @@ def fused_norm_rope(
         idx_cache_scale_view,
         idx_cache_block_size,
         idx_cache_stride,
+        idx_cache_shuffle,
+        _INDEXER_CACHE_BLOCK_TILE,
+        idx_cache_head_tile,
         # MLA KV cache (uses same slot_mapping)
         mla_kv_cache,
         mla_block_stride,
@@ -659,8 +721,9 @@ def fused_norm_rope(
         TOPK_BLOCK_SIZE=1024,
         HAS_INDEXER=has_indexer,
         INDEX_ROPE_INTERLEAVE=index_rope_interleave,
-        USE_PDL=use_pdl,
-        launch_pdl=use_pdl,
+        FP8_MAX=_FP8_MAX,
+        USE_FNUZ=_USE_FNUZ,
+        **pdl_kwargs,
     )
     return q_c_out
 
@@ -715,6 +778,8 @@ def _fused_q_kernel(
     INDEX_ROPE_INTERLEAVE: tl.constexpr,
     QUANTIZE_MQA: tl.constexpr,
     USE_PDL: tl.constexpr,
+    FP8_MAX: tl.constexpr,
+    USE_FNUZ: tl.constexpr,
 ):
     tok_idx = tl.program_id(0).to(tl.int64)
     pid = tl.program_id(1)
@@ -722,6 +787,7 @@ def _fused_q_kernel(
     if USE_PDL:
         tl.extra.cuda.gdc_wait()
         tl.extra.cuda.gdc_launch_dependents()
+    fp8_dtype = tl.float8e4b8 if USE_FNUZ else tl.float8e4nv
 
     if pid == 2:
         # ql_nope quantize + pack into the front of mqa_q_fp8. On the bf16
@@ -744,7 +810,7 @@ def _fused_q_kernel(
                     + ql_nope_off,
                     mask=ql_nope_mask,
                 ).to(tl.float32)
-                ql_nope_fp8 = (ql_nope / scale).to(tl.float8e4nv)
+                ql_nope_fp8 = (ql_nope / scale).to(fp8_dtype)
                 tl.store(
                     mqa_q_fp8_ptr
                     + tok_idx * mqa_q_fp8_stride0
@@ -794,7 +860,7 @@ def _fused_q_kernel(
                         + q_head_idx * mqa_q_fp8_stride1
                         + QL_NOPE_DIM
                         + rot_off * 2,
-                        (r1 / scale).to(tl.float8e4nv),
+                        (r1 / scale).to(fp8_dtype),
                     )
                     tl.store(
                         mqa_q_fp8_ptr
@@ -803,7 +869,7 @@ def _fused_q_kernel(
                         + QL_NOPE_DIM
                         + rot_off * 2
                         + 1,
-                        (r2 / scale).to(tl.float8e4nv),
+                        (r2 / scale).to(fp8_dtype),
                     )
                 else:
                     # bf16 query: write the RoPE'd q_pe unquantized.
@@ -861,7 +927,7 @@ def _fused_q_kernel(
         index_q = tl.where(in_rope, roped, index_q)
 
         # Index Q Quantize (from registers)
-        index_q_fp8, index_q_scale = _fp8_ue8m0_quantize(index_q)
+        index_q_fp8, index_q_scale = _fp8_ue8m0_quantize(index_q, FP8_MAX, USE_FNUZ)
         tl.store(
             index_q_fp8_ptr
             + tok_idx * index_q_fp8_stride0
@@ -955,7 +1021,7 @@ def fused_q(
             q_pe.shape[0],
             q_pe.shape[1],
             ql_nope.shape[2] + q_pe.shape[2],
-            dtype=torch.float8_e4m3fn,
+            dtype=_FP8_DTYPE,
             device=q_pe.device,
         )
         # Placeholder; pid 0 packs q_pe into mqa_q_fp8 instead.
@@ -967,7 +1033,7 @@ def fused_q(
         mqa_q_fp8 = q_pe_out  # unused placeholder for the fp8 pack pointer
         mqa_q = q_pe_out
 
-    index_q_fp8 = torch.empty_like(index_q, dtype=torch.float8_e4m3fn)
+    index_q_fp8 = torch.empty_like(index_q, dtype=_FP8_DTYPE)
     index_weights_out = torch.empty_like(index_weights, dtype=torch.float32)
     if cutedsl_kernel is not None:
         cutedsl_kernel(
@@ -990,6 +1056,9 @@ def fused_q(
         return index_q_fp8, index_weights_out, mqa_q
 
     use_pdl = current_platform.is_arch_support_pdl()
+    pdl_kwargs = (
+        {"USE_PDL": use_pdl, "launch_pdl": True} if use_pdl else {"USE_PDL": use_pdl}
+    )
     _fused_q_kernel[(num_tokens, 3, grid_heads)](
         positions,
         q_pe,
@@ -1031,8 +1100,9 @@ def fused_q(
         HAS_INDEXER=has_indexer,
         INDEX_ROPE_INTERLEAVE=index_rope_interleave,
         QUANTIZE_MQA=quantize_mqa,
-        USE_PDL=use_pdl,
-        launch_pdl=use_pdl,
+        FP8_MAX=_FP8_MAX,
+        USE_FNUZ=_USE_FNUZ,
+        **pdl_kwargs,
         # num_warps=1 is optimal here: each program is a single 128-element
         # rope+quant, so the kernel is program-count/occupancy bound, not
         # per-program compute bound (swept 1/2/4/8 — 1 wins or ties everywhere).