iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm][Model][Bugfix] Enable GLM-5.2-MXFP4 on the deepseek_v32 path and fix sparse attention correctness (#51915)

Browse Source 
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>
This commit is contained in:
Jack Hu
2026-09-22 18:08:51 -05:00
committed by GitHub
co-authored by James E T Smith Douglas Lehr Douglas Lehr
parent 9646f53064
commit 1c0eee919d
8 changed files with 318 additions and 211 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/test_fused_deepseek_v32_norm_rope.py
+186 -150
View File
@@ -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")
vllm/_aiter_ops.py
+1 -1
View File
@@ -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,
vllm/model_executor/layers/attention/mla_attention.py
+1 -3
View File
@@ -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,
vllm/model_executor/models/deepseek_v2.py
+5
View File
@@ -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
vllm/models/deepseek_v32/__init__.py
+8 -4
View File
@@ -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,
vllm/models/deepseek_v32/amd/rocm.py
+16 -25
View File
@@ -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]
vllm/models/deepseek_v32/attention.py
+4 -1
View File
@@ -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[
vllm/models/deepseek_v32/common/kernels.py
+97 -27
View File
@@ -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).