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

[Kernel][Build] Bump DeepGEMM to the 26/09/30 release and adapt the sparse indexer (#61149)

Browse Source 
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yongye Zhu
2026-10-12 23:24:43 -07:00
committed by GitHub
co-authored by Claude Opus 5.5
parent 86b1c2ef15
commit a6009ea85b
23 changed files with 516 additions and 208 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmake/external_projects/deepgemm.cmake
+6 -6
View File
@@ -29,12 +29,11 @@ if(DEEPGEMM_SRC_DIR)
else()
# Keep in sync with tools/install_deepgemm.sh
set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/vllm-project/DeepGEMM.git")
# Pinned to the fork's dev branch: upstream 2.8.0 plus the SM120
# port, the SM90 paged-MQA kv_block=32/next_n=4 port, configurable SwiGLU
# alpha/beta, and SiTU for FP8/FP4 Mega MoE. Also includes the CUDA 12.x
# layout header fix from vllm-project/DeepGEMM#12.
# TORCH_LIBRARY build with a CPython 3.10 stable-ABI floor.
set(_DEEPGEMM_UPSTREAM_TAG "1e1842a833699298f7afc02eefb2ed168fff6938")
# Pinned to the fork's dev branch: upstream's 26/09/30 release (2.8.1)
# plus the SM120 port, the SM90 paged-MQA kv_block=32/next_n=4 port,
# configurable SwiGLU alpha/beta, SiTU for FP8/FP4 Mega MoE, and NVFP4
# Mega MoE. TORCH_LIBRARY build with a CPython 3.10 stable-ABI floor.
set(_DEEPGEMM_UPSTREAM_TAG "6ce1c5ec6627a6bf11dbcc698b0b1e19c9248432")
set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}")
if(NOT _deepgemm_fc_root)
@@ -146,6 +145,7 @@ if(DEEPGEMM_ARCHS)
#
install(FILES
"${deepgemm_SOURCE_DIR}/deep_gemm/__init__.py"
"${deepgemm_SOURCE_DIR}/deep_gemm/locality_domain.py"
DESTINATION vllm/third_party/deep_gemm
COMPONENT _deep_gemm_C)
csrc/cache_kernels.cu
+16 -4
View File
@@ -630,7 +630,8 @@ __global__ void indexer_k_quant_and_cache_kernel(
const int cache_block_size, // cache block size
const int64_t cache_block_stride, // stride for each block in kv_cache
const bool use_ue8m0 // use ue8m0 scale format
const bool use_ue8m0, // use ue8m0 scale format
const bool use_mx_scale // store the ue8m0 scale as packed MX bytes
) {
constexpr int VEC_SIZE = 4;
const int64_t token_idx = blockIdx.x;
@@ -679,7 +680,13 @@ __global__ void indexer_k_quant_and_cache_kernel(
const int64_t dst_scale_idx =
block_idx * cache_block_stride + cache_block_size * head_dim +
(block_offset * head_dim + head_dim_idx) * 4 / quant_block_size;
reinterpret_cast<float*>(kv_cache)[dst_scale_idx / 4] = scale;
if (use_mx_scale) {
// One UE8M0 byte per 32 elements, all equal to the shared scale's.
reinterpret_cast<uint32_t*>(kv_cache)[dst_scale_idx / 4] =
((__float_as_uint(scale) >> 23) & 0xFF) * 0x01010101u;
} else {
reinterpret_cast<float*>(kv_cache)[dst_scale_idx / 4] = scale;
}
}
}
@@ -1912,7 +1919,7 @@ void cp_gather_and_upconvert_nvfp4_kv_cache(
reinterpret_cast<KV_T*>(k.data_ptr()), \
reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
slot_mapping.const_data_ptr<int64_t>(), head_dim, quant_block_size, \
cache_block_size, cache_block_stride, use_ue8m0);
cache_block_size, cache_block_stride, use_ue8m0, use_mx_scale);
void indexer_k_quant_and_cache(
torch::stable::Tensor& k, // [num_tokens, head_dim]
@@ -1924,8 +1931,13 @@ void indexer_k_quant_and_cache(
int head_dim = k.size(1);
int cache_block_size = kv_cache.size(1);
int64_t cache_block_stride = kv_cache.stride(0);
bool use_ue8m0 = scale_fmt == "ue8m0";
// "mxfp8" stores the ue8m0 scale as four packed UE8M0 bytes, the MX scales
// DeepGEMM's SM100 MQA logits kernels read.
bool use_mx_scale = scale_fmt == "mxfp8";
bool use_ue8m0 = use_mx_scale || scale_fmt == "ue8m0";
STD_TORCH_CHECK(!use_mx_scale || (head_dim == 128 && quant_block_size == 128),
"mxfp8 scales require head_dim == quant_block_size == 128");
STD_TORCH_CHECK(k.device() == kv_cache.device(),
"k and kv_cache must be on the same device");
STD_TORCH_CHECK(k.device() == slot_mapping.device(),
tests/kernels/attention/test_deepgemm_attention.py
+116 -67
View File
@@ -5,6 +5,7 @@ import random
import pytest
import torch
from vllm.model_executor.layers.sparse_attn_indexer import as_fp32_logits
from vllm.platforms import current_platform
from vllm.utils.deep_gemm import (
_ceil_to_ue8m0,
@@ -13,18 +14,29 @@ from vllm.utils.deep_gemm import (
fp8_fp4_paged_mqa_logits,
get_num_sms,
get_paged_mqa_logits_metadata,
mqa_logits_mx_inputs,
mqa_logits_use_mx_scales,
native_next_n_supported,
)
from vllm.utils.import_utils import has_deep_gemm
from vllm.utils.math_utils import cdiv
def encode_indexer_scales(sf: torch.Tensor) -> torch.Tensor:
"""Per-token power-of-two fp32 scales as the indexer K cache stores them:
four equal packed UE8M0 bytes under MX scales, else the fp32 itself."""
if not mqa_logits_use_mx_scales():
return sf
ue8m0 = (sf.view(torch.int32) >> 23).to(torch.uint8)
return ue8m0[..., None].repeat_interleave(4, dim=-1).view(torch.int32)[..., 0]
def kv_cache_cast_to_fp8(x: torch.Tensor) -> torch.Tensor:
# x: (num_blocks, block_size, 1, head_dim)
num_blocks, block_size, num_heads, head_dim = x.shape
assert num_heads == 1
x_amax = x.abs().float().amax(dim=3, keepdim=True).clamp(1e-4)
sf = x_amax / 448.0
sf = _ceil_to_ue8m0(x_amax / 448.0)
x_scaled = (x * (1.0 / sf)).to(torch.float8_e4m3fn)
x_fp8 = torch.empty(
(num_blocks, block_size * (head_dim + 4)),
@@ -34,9 +46,9 @@ def kv_cache_cast_to_fp8(x: torch.Tensor) -> torch.Tensor:
x_fp8[:, : block_size * head_dim] = x_scaled.view(
num_blocks, block_size * head_dim
).view(dtype=torch.uint8)
x_fp8[:, block_size * head_dim :] = sf.view(num_blocks, block_size).view(
dtype=torch.uint8
)
x_fp8[:, block_size * head_dim :] = encode_indexer_scales(
sf.view(num_blocks, block_size)
).view(dtype=torch.uint8)
return x_fp8.view(num_blocks, block_size, num_heads, head_dim + 4)
@@ -91,64 +103,72 @@ def _ref_fp8_mqa_logits(
return logits
def _packed_requests_test_data(seq_len: int, seq_len_kv: int):
# Two causal prefill requests packed into one KV buffer, so the second
# request's rows start at a nonzero K position.
half, half_kv = seq_len // 2, seq_len_kv // 2
rows = torch.arange(seq_len, dtype=torch.int, device="cuda")
ks = torch.where(rows < half, 0, half_kv).to(torch.int)
ke = ks + half_kv - half + 1 + rows % half
return ks, ke
@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
@pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available")
@pytest.mark.skipif(
not current_platform.has_device_capability(90), reason="SM90 and SM100 only"
)
@pytest.mark.parametrize("clean_logits", [True, False])
def test_deepgemm_fp8_mqa_logits(clean_logits: bool):
@pytest.mark.parametrize("layout", ["cp", "no_cp", "packed"])
def test_deepgemm_fp8_mqa_logits(layout: str):
"""FP8 indexer prefill logits, compressed to each row's valid K span."""
torch.manual_seed(0)
random.seed(0)
num_heads, head_dim = 32, 128
for seq_len in (512,):
for seq_len_kv in (1024,):
for disable_cp in (False, True):
q = torch.randn(
seq_len,
num_heads,
head_dim,
device="cuda",
dtype=torch.bfloat16,
)
kv = torch.randn(
seq_len_kv, head_dim, device="cuda", dtype=torch.bfloat16
)
weights = torch.randn(
seq_len, num_heads, device="cuda", dtype=torch.float32
)
seq_len, seq_len_kv = 512, 1024
q = torch.randn(seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16)
kv = torch.randn(seq_len_kv, head_dim, device="cuda", dtype=torch.bfloat16)
weights = torch.randn(seq_len, num_heads, device="cuda", dtype=torch.float32)
if disable_cp:
ks = torch.zeros(seq_len, dtype=torch.int, device="cuda")
ke = torch.arange(seq_len, dtype=torch.int, device="cuda") + (
seq_len_kv - seq_len
)
else:
ks, ke = _generate_cp_test_data(seq_len, seq_len_kv)
if layout == "cp":
ks, ke = _generate_cp_test_data(seq_len, seq_len_kv)
elif layout == "no_cp":
ks = torch.zeros(seq_len, dtype=torch.int, device="cuda")
ke = torch.arange(seq_len, dtype=torch.int, device="cuda") + (
seq_len_kv - seq_len
)
else:
ks, ke = _packed_requests_test_data(seq_len, seq_len_kv)
q_fp8 = q.to(torch.float8_e4m3fn)
kv_fp8 = per_custom_dims_cast_to_fp8(kv, (0,), False)
logits = fp8_fp4_mqa_logits(
(q_fp8, None), kv_fp8, weights, ks, ke, clean_logits=clean_logits
)
q_fp8 = q.to(torch.float8_e4m3fn)
kv_fp8, kv_sf = per_custom_dims_cast_to_fp8(kv, (0,), True)
q_scale, kernel_weights = mqa_logits_mx_inputs(q_fp8, None, weights)
max_seqlen_k = int((ke - ks).max())
logits = fp8_fp4_mqa_logits(
(q_fp8, q_scale),
(kv_fp8, encode_indexer_scales(kv_sf)),
kernel_weights,
ks,
ke,
max_seqlen_k,
)
assert logits.shape == (seq_len, max_seqlen_k)
logits = as_fp32_logits(logits, ke - ks)
ref_logits = _ref_fp8_mqa_logits(
q=q,
kv=kv,
weights=weights,
cu_seqlen_ks=ks,
cu_seqlen_ke=ke,
)
ref_neginf_mask = ref_logits == float("-inf")
if clean_logits:
neginf_mask = logits == float("-inf")
assert torch.equal(neginf_mask, ref_neginf_mask)
ref_logits = ref_logits.masked_fill(ref_neginf_mask, 0)
logits = logits.masked_fill(ref_neginf_mask, 0)
diff = calc_diff(logits, ref_logits)
assert diff < 1e-3, f"{diff=}"
ref_logits = _ref_fp8_mqa_logits(
q=q,
kv=kv,
weights=kernel_weights.float(),
cu_seqlen_ks=ks,
cu_seqlen_ke=ke,
)
# Compressed row i holds K positions [ks[i], ke[i]) from column zero.
cols = ks[:, None] + torch.arange(max_seqlen_k, device="cuda")[None, :]
valid = cols < ke[:, None]
ref_logits = ref_logits.gather(1, cols.clamp(max=seq_len_kv - 1).long())
ref_logits = ref_logits.masked_fill(~valid, 0)
logits = logits.masked_fill(~valid, 0)
diff = calc_diff(logits, ref_logits)
assert diff < 1e-3, f"{diff=}"
def _ref_fp8_fp4_paged_mqa_logits(
@@ -212,9 +232,6 @@ def test_deepgemm_fp8_fp4_paged_mqa_logits(batch_size: int, next_n: int):
if not native_next_n_supported(next_n):
pytest.skip(f"next_n={next_n} has no native kernel on this architecture")
# NOTE: clean_logits=True is incompatible with the 2D context_lens
# required by csrc/apis/attention.hpp; only the False path is exercised.
clean_logits = False
torch.manual_seed(0)
random.seed(0)
@@ -273,26 +290,42 @@ def test_deepgemm_fp8_fp4_paged_mqa_logits(batch_size: int, next_n: int):
context_lens_2d = (
context_lens.unsqueeze(-1) - next_n + 1 + next_n_arange
).contiguous()
q_scale, kernel_weights = mqa_logits_mx_inputs(q_fp8, None, weights)
kernel_q = (q_fp8, q_scale)
kernel_context_lens, kernel_block_tables = context_lens_2d, block_tables
indices = None
if mqa_logits_use_mx_scales():
# SM100 only takes varlen rows: one Q token each, tied to its
# request by `indices`.
assert q_scale is not None
kernel_q = (
q_fp8.flatten(0, 1)[:, None],
q_scale.flatten(0, 1)[:, None],
)
kernel_context_lens = context_lens_2d.view(-1, 1)
kernel_block_tables = block_tables.repeat_interleave(next_n, dim=0)
indices = torch.arange(batch_size, device="cuda", dtype=torch.int32)
indices = indices.repeat_interleave(next_n)
schedule_metadata = get_paged_mqa_logits_metadata(
context_lens_2d,
blocksize,
get_num_sms(),
kernel_context_lens, blocksize, get_num_sms(), indices=indices
)
logits = fp8_fp4_paged_mqa_logits(
(q_fp8, None),
kernel_q,
kv_cache_fp8,
weights,
context_lens_2d,
block_tables,
kernel_weights,
kernel_context_lens,
kernel_block_tables,
schedule_metadata,
max_model_len,
clean_logits=clean_logits,
indices=indices,
)
if indices is not None:
logits = as_fp32_logits(logits, kernel_context_lens.view(-1))
ref_logits = _ref_fp8_fp4_paged_mqa_logits(
q,
kv_cache,
weights,
kernel_weights.float(),
context_lens,
block_tables,
max_model_len,
@@ -309,7 +342,7 @@ def test_deepgemm_fp8_fp4_paged_mqa_logits(batch_size: int, next_n: int):
context_lens[row_indices] - next_n + next_n_offset
).unsqueeze(1)
logits = logits.masked_fill(~mask, 0)
logits = logits.float().masked_fill(~mask, 0)
ref_logits = ref_logits.masked_fill(~mask, 0)
diff = calc_diff(logits, ref_logits)
assert diff < 1e-3, f"{diff=}"
@@ -349,12 +382,28 @@ def test_deepgemm_paged_mqa_packed_manager_block_stride():
q = torch.randn(1, 1, num_heads, head_dim, device="cuda", dtype=torch.bfloat16).to(
torch.float8_e4m3fn
)
weights = torch.randn(1, num_heads, device="cuda", dtype=torch.float32)
q_scale, weights = mqa_logits_mx_inputs(
q, None, torch.randn(1, num_heads, device="cuda", dtype=torch.float32)
)
context_lens = torch.full((1, 1), block_size, device="cuda", dtype=torch.int32)
schedule = get_paged_mqa_logits_metadata(context_lens, page_size, get_num_sms())
indices = (
torch.zeros(1, device="cuda", dtype=torch.int32)
if mqa_logits_use_mx_scales()
else None
)
schedule = get_paged_mqa_logits_metadata(
context_lens, page_size, get_num_sms(), indices=indices
)
logits = [
fp8_fp4_paged_mqa_logits(
(q, None), cache, weights, context_lens, table, schedule, block_size, False
(q, q_scale),
cache,
weights,
context_lens,
table,
schedule,
block_size,
indices=indices,
)
for cache, table in zip((compact_pages, packed_pages), tables)
]
tests/kernels/attention/test_rocm_paged_mxfp4_indexer.py
+1
View File
@@ -253,6 +253,7 @@ def _prefill_metadata(
cu_seq_lens=torch.zeros(1, dtype=torch.int32, device=DEVICE),
token_to_seq=torch.zeros(1, dtype=torch.int32, device=DEVICE),
total_seq_lens=1,
max_seqlen_k=int(ends.max()),
token_start=t0,
token_end=t1,
num_reqs=req_hi - req_lo,
tests/kernels/test_compressor_kv_cache.py
+32 -10
View File
@@ -1422,7 +1422,10 @@ def test_dequantize_and_gather_k_cache(
@pytest.mark.parametrize("num_tokens", [1, 4, 8, 17])
@pytest.mark.parametrize("block_size", [16, 64])
def test_indexer_quant_cache_roundtrip(num_tokens: int, block_size: int):
@pytest.mark.parametrize("scale_fmt", ["ue8m0", "mxfp8"])
def test_indexer_quant_cache_roundtrip(
num_tokens: int, block_size: int, scale_fmt: str
):
"""K → indexer_k_quant_and_cache → cp_gather_indexer_k_quant_cache
→ manual dequant → compare against original."""
HEAD_DIM = 128
@@ -1442,7 +1445,9 @@ def test_indexer_quant_cache_roundtrip(num_tokens: int, block_size: int):
)
slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
ops.indexer_k_quant_and_cache(k, kv_cache, slot_mapping, QUANT_BLOCK_SIZE, "ue8m0")
ops.indexer_k_quant_and_cache(
k, kv_cache, slot_mapping, QUANT_BLOCK_SIZE, scale_fmt
)
# ── Gather ──────────────────────────────────────────────────────────
max_blocks_per_seq = num_blocks
@@ -1465,7 +1470,12 @@ def test_indexer_quant_cache_roundtrip(num_tokens: int, block_size: int):
# ── Manual dequant ──────────────────────────────────────────────────
k_fp8 = dst_k.view(torch.float8_e4m3fn).float() # [num_tokens, 128]
scale = dst_scale.view(torch.float32) # [num_tokens, 1]
if scale_fmt == "mxfp8":
# One UE8M0 byte per 32 elements, all equal to the token's scale.
assert (dst_scale == dst_scale[:, :1]).all()
scale = (dst_scale[:, :1].to(torch.int32) << 23).view(torch.float32)
else:
scale = dst_scale.view(torch.float32) # [num_tokens, 1]
k_recovered = k_fp8 * scale # [num_tokens, 128]
# ── Compare ─────────────────────────────────────────────────────────
@@ -1749,8 +1759,13 @@ def _reference_kv_compress_norm_rope(
),
],
)
def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: bool):
@pytest.mark.parametrize("mx_scale", [False, True])
def test_fused_kv_insert_indexer(
num_tokens: int, kv_block_size: int, use_fp4: bool, mx_scale: bool
):
"""Fused K compress+norm+rope+quant+insert for the indexer KV cache."""
if use_fp4 and mx_scale:
pytest.skip("MX scales only change the FP8 cache")
HEAD_DIM = 128
ROPE_DIM = 64
BLOCK_SIZE = 16
@@ -1833,6 +1848,7 @@ def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: b
TOKEN_STRIDE=TOKEN_STRIDE,
SCALE_DIM=SCALE_DIM,
KV_BLOCK_STRIDE=kv_cache.stride(0),
MX_SCALE=mx_scale,
num_warps=1,
)
@@ -1876,12 +1892,18 @@ def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: b
), f"token {i}: FP8 bytes differ"
scale_off = kv_block_size * TOKEN_STRIDE + pos * SCALE_DIM
actual_scale = kv_cache[blk, scale_off : scale_off + SCALE_DIM].view(
torch.float32
)
assert torch.equal(actual_scale, scale[i : i + 1]), (
f"token {i}: scale {actual_scale.item()} != {scale[i].item()}"
)
actual_scale = kv_cache[blk, scale_off : scale_off + SCALE_DIM]
if mx_scale:
# The power-of-two scale's exponent as 4 equal UE8M0 bytes.
ue8m0 = (scale[i : i + 1].view(torch.int32) >> 23).to(torch.uint8)
assert torch.equal(actual_scale, ue8m0.repeat(SCALE_DIM)), (
f"token {i}: mx scale {actual_scale.tolist()} != {ue8m0.item()}"
)
else:
actual_scale = actual_scale.view(torch.float32)
assert torch.equal(actual_scale, scale[i : i + 1]), (
f"token {i}: scale {actual_scale.item()} != {scale[i].item()}"
)
@pytest.mark.parametrize("compress_ratio", [4, 128])
tests/kernels/test_fused_deepseek_v32_norm_rope.py
+13 -2
View File
@@ -32,6 +32,7 @@ import torch
from vllm.models.deepseek_v32.common import kernels as K
from vllm.platforms import current_platform
from vllm.utils.deep_gemm import mqa_logits_use_mx_scales
FP8 = torch.float8_e4m3fn
FP8_MAX = 448.0
@@ -116,6 +117,16 @@ def rope(
return out
def indexer_scales(scale_bytes: torch.Tensor) -> torch.Tensor:
"""Per-token fp32 scales from the indexer cache's 4 scale bytes per token."""
if not mqa_logits_use_mx_scales():
return scale_bytes.view(torch.float32)
# MX scales: the power-of-two scale's exponent as 4 equal UE8M0 bytes.
ue8m0 = scale_bytes.view(-1, 4)
assert (ue8m0 == ue8m0[:, :1]).all(), "MX scale bytes differ within a token"
return (ue8m0[:, 0].to(torch.int32) << 23).view(torch.float32)
def ue8m0_quant(vals: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""Per-row (last dim) UE8M0 fp8 quant matching kernels._fp8_ue8m0_quantize."""
amax = vals.float().abs().amax(dim=-1, keepdim=True)
@@ -243,7 +254,7 @@ def test_fused_norm_rope(
q_ref, s_ref = ue8m0_quant(ik_ref)
flat = idx_cache[0].reshape(-1)
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)
scales = indexer_scales(flat[bs * cfg.index_head_dim :])
assert_fp8(vals[:num_tokens], q_ref, "indexer-K fp8")
torch.testing.assert_close(scales[:num_tokens], s_ref, rtol=0, atol=0)
@@ -318,7 +329,7 @@ def test_fused_norm_rope_packed_indexer_block_stride(cfg: ModelConfig):
.view(FP8)
.view(block_size, cfg.index_head_dim)
)
scales = packed[block_size * cfg.index_head_dim :].view(torch.float32)
scales = indexer_scales(packed[block_size * cfg.index_head_dim :])
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()
tests/kernels/test_fused_indexer_q_rope_quant.py
+5
View File
@@ -26,6 +26,7 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
)
from vllm.models.deepseek_v4.common.ops import fused_indexer_q_rope_quant
from vllm.platforms import current_platform
from vllm.utils.deep_gemm import mqa_logits_use_mx_scales
from vllm.utils.import_utils import has_cutedsl
HEAD_DIM = 128
@@ -347,6 +348,10 @@ def _indexer_k_reference(
exponent = k.float().abs().amax(dim=-1, keepdim=True).clamp(min=1e-4) / 448
scale = exponent.log2().ceil().exp2()
values = (k.float() / scale).clamp(-448, 448).to(torch.float8_e4m3fn)
if mqa_logits_use_mx_scales():
# The power-of-two scale's exponent as 4 equal UE8M0 bytes.
ue8m0 = (scale.view(torch.int32) >> 23).to(torch.uint8)
return values.view(torch.uint8), ue8m0.repeat(1, 4)
return values.view(torch.uint8), scale.view(torch.uint8)
tests/kernels/test_kpool_decode_update_batched.py
+7 -1
View File
@@ -29,6 +29,7 @@ import pytest
import torch
from vllm.platforms import current_platform
from vllm.utils.deep_gemm import mqa_logits_use_mx_scales
if current_platform.is_rocm():
from vllm.models.glm5next.amd.ops.kpool_compress import (
@@ -206,7 +207,12 @@ def _torch_reference(
k_off = page_base + loc_tok * HEAD_DIM + torch.arange(HEAD_DIM)
s_off = page_base + k_region + loc_tok * 4
kv_flat[k_off] = quantized.view(torch.uint8)
kv_flat[s_off : s_off + 4] = scale.detach().reshape(1).view(torch.uint8)
scale_bytes = scale.detach().reshape(1).view(torch.uint8)
if mqa_logits_use_mx_scales():
# The power-of-two scale's exponent as 4 equal UE8M0 bytes.
ue8m0 = scale.reshape(1).view(torch.int32) >> 23
scale_bytes = ue8m0.to(torch.uint8).repeat(4)
kv_flat[s_off : s_off + 4] = scale_bytes
# stash -- gated on the TOKEN-granular tail slot, not on pos_valid.
# pos_valid keys off the pool-granular cache_loc, which is -1 for
tests/model_executor/layers/test_mla_short_prefill_indexer.py
+18 -21
View File
@@ -771,6 +771,8 @@ def test_sparse_mqa_logits_prefill_matches_dense_masked(capture: bool):
candidates = _make_candidates(lens, num_candidates, cbk).cuda()
# Reference: dense bf16 logits masked to the candidates, then torch.topk.
# The logits are compressed: row r starts at K position ks[r].
max_len = int(lens.max())
dense = deep_gemm.fp8_fp4_mqa_logits(
(
q_fp.view(torch.int8).view(rows, num_heads, head_dim // 2),
@@ -780,19 +782,13 @@ def test_sparse_mqa_logits_prefill_matches_dense_masked(capture: bool):
weights,
ks,
ke,
clean_logits=False,
logits_dtype=torch.bfloat16,
max_len,
)
cols = torch.arange(total, device="cuda")
in_range = (cols - ks.unsqueeze(1) >= 0) & (
cols - ks.unsqueeze(1) < lens.unsqueeze(1)
)
masked = dense.masked_fill(
~(in_range & _candidate_token_mask(candidates, cbk, total, ks)),
float("-inf"),
).float()
ref_abs = _reference_topk(masked, topk) # packed-workspace positions
ref = torch.where(ref_abs >= 0, ref_abs - ks.unsqueeze(1), -1)
in_range = torch.arange(max_len, device="cuda") < lens.unsqueeze(1)
# A row's last candidate block may extend past max_len.
in_cand = _candidate_token_mask(candidates, cbk, max_len + cbk)[:, :max_len]
masked = dense.masked_fill(~(in_range & in_cand), float("-inf")).float()
ref = _reference_topk(masked, topk)
out = torch.full((rows, topk), -1, dtype=torch.int32, device="cuda")
scratch = _sparse_scratch(rows, num_candidates, topk)
@@ -827,11 +823,11 @@ def test_sparse_mqa_logits_prefill_matches_dense_masked(capture: bool):
graph.replay()
else:
kernel_metadata = run_sparse()
_assert_same_selection(dense, out, ref, row_starts=ks)
_assert_same_selection(dense, out, ref)
out.fill_(-1)
run_sparse(kernel_metadata)
_assert_same_selection(dense, out, ref, row_starts=ks)
_assert_same_selection(dense, out, ref)
def _build_fp4_paged_cache(req_lens: torch.Tensor, page_kv: int, head_dim: int):
@@ -902,21 +898,22 @@ def test_sparse_mqa_logits_paged_matches_dense_masked(next_n: int, capture: bool
candidates = _make_candidates(row_lens.flatten(), num_candidates, cbk).cuda()
def reference_topk():
context_lens_2d = row_lens.view(batch, next_n)
# SM100 paged MQA logits take varlen rows: one Q token per row.
context_lens_2d = row_lens.view(rows, 1)
indices = torch.arange(batch, dtype=torch.int32, device="cuda")
indices = indices.repeat_interleave(next_n)
sched = deep_gemm.get_paged_mqa_logits_metadata(
context_lens_2d, page_kv, deep_gemm.get_num_sms(), None
context_lens_2d, page_kv, deep_gemm.get_num_sms(), indices
)
logits = deep_gemm.fp8_fp4_paged_mqa_logits(
(q, q_scale),
(q.view(rows, 1, num_heads, -1), q_scale.view(rows, 1, num_heads)),
kv_cache,
weights,
context_lens_2d,
block_table,
block_table.repeat_interleave(next_n, dim=0),
sched,
int(req_lens.max()),
False,
torch.bfloat16,
None,
indices,
)
lens_flat = row_lens.flatten()
cols = torch.arange(logits.shape[1], device="cuda")
tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py
+10 -1
View File
@@ -14,6 +14,7 @@ from vllm.models.deepseek_v4.sparse_mla import DeepseekV4SparseMLABackend
from vllm.models.deepseek_v41.sparse_mla import (
DeepseekV4SparseMLABackend as DeepseekV41SparseMLABackend,
)
from vllm.utils.deep_gemm import mqa_logits_use_mx_scales
from vllm.v1.attention.backend import CommonAttentionMetadata
from vllm.v1.attention.backends.mla import indexer as indexer_module
from vllm.v1.attention.backends.mla.compressor_utils import (
@@ -113,10 +114,17 @@ def test_indexer_draft_decode_metadata_update_is_capture_safe():
num_reqs, 4
)
num_sms = torch.cuda.get_device_properties(device).multi_processor_count
# SM100's paged MQA logits only take varlen rows.
indices = (
torch.arange(num_reqs, dtype=torch.int32, device=device)
if mqa_logits_use_mx_scales()
else None
)
schedule_metadata = indexer_module.get_paged_mqa_logits_metadata(
torch.div(seq_lens, 16, rounding_mode="floor").unsqueeze(1),
64,
num_sms,
indices=indices,
).clone()
builder = object.__new__(DeepseekV32IndexerMetadataBuilder)
@@ -139,6 +147,7 @@ def test_indexer_draft_decode_metadata_update_is_capture_safe():
decode_lens=torch.zeros(num_reqs, dtype=torch.int32, device=device),
requires_padding=False,
schedule_metadata=schedule_metadata,
indices=indices,
),
block_table=block_table,
)
@@ -157,7 +166,7 @@ def test_indexer_draft_decode_metadata_update_is_capture_safe():
torch.testing.assert_close(
schedule_metadata,
indexer_module.get_paged_mqa_logits_metadata(
expected_decode_seq_lens.unsqueeze(1), 64, num_sms
expected_decode_seq_lens.unsqueeze(1), 64, num_sms, indices=indices
),
)
assert metadata.decode is not None
tests/v1/attention/test_rocm_paged_mxfp4_indexer_plan.py
+1
View File
@@ -36,6 +36,7 @@ def _chunk(block_table, token_start, token_end, row_ends):
cu_seq_lens=torch.zeros(1, dtype=torch.int32),
token_to_seq=torch.zeros(1, dtype=torch.int32),
total_seq_lens=1,
max_seqlen_k=int(ends.max()),
token_start=token_start,
token_end=token_end,
num_reqs=block_table.shape[0],
tools/build_deepgemm_C.py
+2 -1
View File
@@ -35,8 +35,9 @@ env = os.environ.copy()
env.pop("DG_SKIP_CUDA_BUILD", None)
print(f"[build_deepgemm_C] building in {src} with {sys.executable}", flush=True)
# CMake decides when to rebuild; setuptools would only check python_api.cpp.
subprocess.check_call(
[sys.executable, "setup.py", "build_ext", "--inplace"],
[sys.executable, "setup.py", "build_ext", "--inplace", "--force"],
cwd=src,
env=env,
)
tools/check_wheel_deepgemm.py
+9 -2
View File
@@ -4,10 +4,12 @@
"""Assert the vendored DeepGEMM package has the TORCH_LIBRARY binding layout.
Expects ``deep_gemm/_C.py`` plus a single ``_C_extension*.so`` (abi3) under
``vllm.third_party.deep_gemm``. Run after vLLM is installed, e.g. the H100
deepgemm kernel tests in .buildkite/test_areas/kernels.yaml.
``vllm.third_party.deep_gemm``, and that the package imports: vLLM treats an
import failure as DeepGEMM being absent. Run after vLLM is installed, e.g. the
H100 deepgemm kernel tests in .buildkite/test_areas/kernels.yaml.
"""
import importlib
import importlib.util
import sys
from pathlib import Path
@@ -38,6 +40,11 @@ def main() -> int:
if missing:
print(f"missing: {missing}", file=sys.stderr)
return 1
try:
importlib.import_module("vllm.third_party.deep_gemm")
except Exception as e:
print(f"vllm.third_party.deep_gemm failed to import: {e!r}", file=sys.stderr)
return 1
return 0
tools/install_deepgemm.sh
+4 -4
View File
@@ -7,10 +7,10 @@ set -e
# Default values
# Keep DEEPGEMM_GIT_REF in sync with cmake/external_projects/deepgemm.cmake
DEEPGEMM_GIT_REPO="https://github.com/vllm-project/DeepGEMM.git"
# NOTE: This targets the vLLM fork's dev branch, which carries the sm120
# and sm90 paged-MQA ports plus the SwiGLU alpha/beta and SiTU Mega MoE
# activations, plus the CUDA 12.x layout header fix from vllm-project/DeepGEMM#12.
DEEPGEMM_GIT_REF="1e1842a833699298f7afc02eefb2ed168fff6938"
# NOTE: This targets the vLLM fork's dev branch: upstream's 26/09/30 release
# plus the sm120 and sm90 paged-MQA ports, the SwiGLU alpha/beta and SiTU
# Mega MoE activations, and NVFP4 Mega MoE.
DEEPGEMM_GIT_REF="6ce1c5ec6627a6bf11dbcc698b0b1e19c9248432"
WHEEL_DIR=""
# Parse command line arguments
vllm/model_executor/layers/sparse_attn_indexer.py
+84 -26
View File
@@ -42,6 +42,8 @@ from vllm.utils.deep_gemm import (
fp8_fp4_mqa_logits,
fp8_fp4_paged_mqa_logits,
has_deep_gemm,
mqa_logits_mx_inputs,
mqa_logits_use_mx_scales,
)
from vllm.utils.import_utils import has_cutedsl
from vllm.utils.torch_utils import (
@@ -286,6 +288,42 @@ def fused_indexer_q_rope_quant(
return q_fp8, weights_out
@triton.jit
def _upcast_logits_kernel(src_ptr, dst_ptr, row_ends_ptr, stride, BLOCK: tl.constexpr):
row = tl.program_id(0)
end = tl.load(row_ends_ptr + row)
src_ptr += row.to(tl.int64) * stride
dst_ptr += row.to(tl.int64) * stride
for start in range(tl.program_id(1) * BLOCK, end, tl.num_programs(1) * BLOCK):
offs = start + tl.arange(0, BLOCK)
mask = offs < end
vals = tl.load(src_ptr + offs, mask=mask)
tl.store(dst_ptr + offs, vals.to(tl.float32), mask=mask)
def as_fp32_logits(logits: torch.Tensor, row_ends: torch.Tensor) -> torch.Tensor:
"""Upcast SM100's BF16 MQA logits for the FP32 top-k kernels.
Only columns below ``row_ends`` (int32, one per row) are converted, the
span the consumers read; DeepGEMM's aligned row stride is kept.
"""
if logits.dtype == torch.float32:
return logits
out = torch.empty_strided(
logits.shape, logits.stride(), dtype=torch.float32, device=logits.device
)
num_rows, width = logits.shape
if num_rows == 0 or width == 0:
return out
assert row_ends.numel() == num_rows and row_ends.dtype == torch.int32
block = 1024
splits = min(triton.cdiv(width, block), max(1, 1024 // num_rows))
_upcast_logits_kernel[(num_rows, splits)](
logits, out, row_ends.contiguous(), logits.stride(0), BLOCK=block
)
return out
def _gather_workspace_shapes(
total_seq_lens: int,
head_dim: int,
@@ -537,7 +575,7 @@ def sparse_attn_indexer(
kv_cache,
slot_mapping_for_cache,
quant_block_size,
scale_fmt,
"mxfp8" if mqa_logits_use_mx_scales() else scale_fmt,
)
# The indexer and main MLA may classify the same short extend differently
# because they use independent decode thresholds. Only the main MLA route
@@ -649,6 +687,9 @@ def sparse_attn_indexer(
q_scale_slice = q_scale[row_start:row_end] if q_scale is not None else None
topk_indices = topk_indices_buffer[row_start:row_end, :topk_tokens]
# DeepGEMM returns compressed logits: row i starts at K position
# cu_seqlen_ks[i]. None means every row starts at column zero.
row_starts: torch.Tensor | None = None
if chunk.local_total_seq_lens == 0 or q_slice.shape[0] == 0:
logits = q_slice.new_empty((q_slice.shape[0], 0), dtype=torch.float32)
topk_indices.fill_(-1)
@@ -662,7 +703,9 @@ def sparse_attn_indexer(
else:
q_slice_cast = q_slice
k_quant_cast = k_quant
k_scale_cast = k_scale.view(torch.float32).squeeze(-1)
k_scale_cast = k_scale.view(
torch.int32 if mqa_logits_use_mx_scales() else torch.float32
).squeeze(-1)
if current_platform.is_xpu():
if q_scale_slice is not None:
raise RuntimeError("XPU fp8_mqa_logits does not support FP4 Q")
@@ -674,14 +717,23 @@ def sparse_attn_indexer(
cu_seqlen_ks,
cu_seqlen_ke,
)
row_starts = cu_seqlen_ks
row_ends = cu_seqlen_ke
else:
logits = fp8_fp4_mqa_logits(
(q_slice_cast, q_scale_slice),
(k_quant_cast, k_scale_cast),
weights[row_start:row_end],
cu_seqlen_ks,
cu_seqlen_ke,
clean_logits=False,
q_scale_slice, chunk_weights = mqa_logits_mx_inputs(
q_slice_cast, q_scale_slice, weights[row_start:row_end]
)
row_ends = cu_seqlen_ke - cu_seqlen_ks
logits = as_fp32_logits(
fp8_fp4_mqa_logits(
(q_slice_cast, q_scale_slice),
(k_quant_cast, k_scale_cast),
chunk_weights,
cu_seqlen_ks,
cu_seqlen_ke,
chunk.max_seqlen_k,
),
row_ends,
)
num_rows = logits.shape[0]
if candidate_blocks is not None:
@@ -694,8 +746,8 @@ def sparse_attn_indexer(
if candidate_write:
_select_candidate_blocks(
logits,
cu_seqlen_ks,
cu_seqlen_ke,
row_starts,
row_ends,
chunk_candidates.shape[1],
candidate_block_size,
chunk_candidates,
@@ -703,15 +755,15 @@ def sparse_attn_indexer(
else:
_apply_candidate_mask(
logits,
cu_seqlen_ks,
cu_seqlen_ke,
row_starts,
row_ends,
chunk_candidates,
candidate_block_size,
)
ops.top_k_per_row_prefill(
logits,
cu_seqlen_ks,
cu_seqlen_ke,
torch.zeros_like(row_ends) if row_starts is None else row_starts,
row_ends,
topk_indices,
num_rows,
logits.stride(0),
@@ -729,7 +781,7 @@ def sparse_attn_indexer(
dcp_rank,
dcp_world_size,
cp_kv_cache_interleave_size,
row_starts=cu_seqlen_ks,
row_starts=row_starts,
)
if row_shard is not None:
@@ -822,16 +874,22 @@ def sparse_attn_indexer(
max_model_len,
)
else:
logits = fp8_fp4_paged_mqa_logits(
(padded_q_quant_cast, padded_q_scale),
kv_cache,
weights[:num_padded_tokens],
seq_lens,
decode_metadata.block_table,
decode_metadata.schedule_metadata,
max_model_len=max_model_len,
clean_logits=False,
indices=decode_metadata.indices,
padded_q_scale, decode_weights = mqa_logits_mx_inputs(
padded_q_quant_cast, padded_q_scale, weights[:num_padded_tokens]
)
# SM100 decode is varlen: one row and context length per token.
logits = as_fp32_logits(
fp8_fp4_paged_mqa_logits(
(padded_q_quant_cast, padded_q_scale),
kv_cache,
decode_weights,
seq_lens,
decode_metadata.block_table,
decode_metadata.schedule_metadata,
max_model_len=max_model_len,
indices=decode_metadata.indices,
),
seq_lens.view(-1),
)
num_rows = logits.shape[0]
if candidate_blocks is not None:
vllm/models/deepseek_v32/common/kernels.py
+13 -1
View File
@@ -6,7 +6,9 @@ import torch
from vllm.platforms import current_platform
from vllm.triton_utils import tl, triton
from vllm.utils.deep_gemm import mqa_logits_use_mx_scales
from vllm.utils.torch_utils import is_quantized_kv_cache
from vllm.v1.attention.ops.common import pack_ue8m0_x4
# Cache of tiny 1-element dummy tensors (per device, dtype) reused by the
# has_indexer=False path so the indexer args don't allocate every call.
@@ -87,6 +89,7 @@ def _fp8_quant_and_cache_write(
SHUFFLE: tl.constexpr,
BLOCK_TILE: tl.constexpr,
HEAD_TILE: tl.constexpr,
MX_SCALE: tl.constexpr,
FP8_MAX: tl.constexpr,
USE_FNUZ: tl.constexpr,
):
@@ -111,6 +114,8 @@ def _fp8_quant_and_cache_write(
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
if MX_SCALE:
scale = pack_ue8m0_x4(scale)
tl.store(kv_cache_scale_ptr + scale_byte_off // 4, scale)
@@ -167,6 +172,7 @@ def _fused_norm_rope_kernel(
INDEXER_CACHE_SHUFFLE: tl.constexpr,
INDEXER_CACHE_BLOCK_TILE: tl.constexpr,
INDEXER_CACHE_HEAD_TILE: tl.constexpr,
INDEXER_CACHE_MX_SCALE: tl.constexpr,
# MLA KV cache (concat kv_c_normed + k_pe_roped, uses slot_mapping_ptr)
mla_cache_ptr,
mla_cache_block_stride,
@@ -501,6 +507,7 @@ def _fused_norm_rope_kernel(
INDEXER_CACHE_SHUFFLE,
INDEXER_CACHE_BLOCK_TILE,
INDEXER_CACHE_HEAD_TILE,
INDEXER_CACHE_MX_SCALE,
FP8_MAX,
USE_FNUZ,
)
@@ -571,7 +578,10 @@ def fused_norm_rope(
assert slot_mapping is not None
if indexer_slot_mapping is None:
indexer_slot_mapping = slot_mapping
idx_cache_scale_view = indexer_k_cache.view(torch.uint8).view(torch.float32)
idx_cache_mx_scale = mqa_logits_use_mx_scales()
idx_cache_scale_view = indexer_k_cache.view(torch.uint8).view(
torch.int32 if idx_cache_mx_scale else 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
@@ -591,6 +601,7 @@ def fused_norm_rope(
)
else:
idx_cache_scale_view = None
idx_cache_mx_scale = False
idx_cache_block_size = 1
idx_cache_stride = 0
idx_cache_shuffle = False
@@ -701,6 +712,7 @@ def fused_norm_rope(
idx_cache_shuffle,
_INDEXER_CACHE_BLOCK_TILE,
idx_cache_head_tile,
idx_cache_mx_scale,
# MLA KV cache (uses same slot_mapping)
mla_kv_cache,
mla_block_stride,
vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py
+11 -1
View File
@@ -33,7 +33,9 @@ from vllm.model_executor.warmup.jit_warmup_triton_helper import (
)
from vllm.platforms import current_platform
from vllm.triton_utils import tl, triton
from vllm.utils.deep_gemm import mqa_logits_use_mx_scales
from vllm.utils.math_utils import round_up
from vllm.v1.attention.ops.common import pack_ue8m0_x4
if current_platform.is_rocm():
from vllm.platforms.rocm import _ON_GFX950
@@ -742,6 +744,7 @@ def _fused_kv_compress_norm_rope_insert_indexer_attn(
TOKEN_STRIDE: tl.constexpr, # 128 for indexer
SCALE_DIM: tl.constexpr, # 4 for indexer (1 float32)
KV_BLOCK_STRIDE: tl.constexpr,
MX_SCALE: tl.constexpr, # store the scale as 4 packed UE8M0 bytes
):
"""Fused compress → RMSNorm → RoPE → FP8 quant → store.
@@ -878,7 +881,10 @@ def _fused_kv_compress_norm_rope_insert_indexer_attn(
# Single float32 scale
scale_val = tl.exp2(exponent)
tl.store(scale_ptr.to(tl.pointer_type(tl.float32)), scale_val)
if MX_SCALE:
tl.store(scale_ptr.to(tl.pointer_type(tl.int32)), pack_ue8m0_x4(scale_val))
else:
tl.store(scale_ptr.to(tl.pointer_type(tl.float32)), scale_val)
# =============================================================================
@@ -927,6 +933,7 @@ def _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn(
TOKEN_STRIDE: tl.constexpr, # HEAD_SIZE // 2 = 64 packed bytes/token
SCALE_DIM: tl.constexpr, # HEAD_SIZE // QUANT_BLOCK = 4 ue8m0 bytes/token
KV_BLOCK_STRIDE: tl.constexpr,
MX_SCALE: tl.constexpr, # unused for MXFP4 (kept for signature parity)
):
"""Fused compress → RMSNorm → RoPE → MXFP4 quant → store.
@@ -1102,6 +1109,7 @@ class FusedKVCompressNormRopeInsertIndexerTritonKernel(
block_size: int
kv_cache_block_size: int
kv_block_stride: int
mx_scale: bool
# Canonical kernel for the launcher's arg_names mapping; the launched
# variant (fp8/mxfp4, same arg_names) is picked per call in __call__.
@@ -1181,6 +1189,7 @@ class FusedKVCompressNormRopeInsertIndexerTritonKernel(
block_size=state_block_size,
kv_cache_block_size=kv_cache_block_size,
kv_block_stride=kv_block_stride,
mx_scale=not use_fp4_cache and mqa_logits_use_mx_scales(),
)
def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
@@ -1301,6 +1310,7 @@ class FusedKVCompressNormRopeInsertIndexerTritonKernel(
TOKEN_STRIDE=token_stride,
SCALE_DIM=scale_dim,
KV_BLOCK_STRIDE=kv_cache.stride(0),
MX_SCALE=not use_fp4_cache and mqa_logits_use_mx_scales(),
num_warps=1,
**pdl_kwargs,
)
vllm/models/deepseek_v41/common/ops/indexer_k_store.py
+13 -2
View File
@@ -17,6 +17,8 @@ import torch
from vllm.platforms import current_platform
from vllm.triton_utils import tl, triton
from vllm.utils.deep_gemm import mqa_logits_use_mx_scales
from vllm.v1.attention.ops.common import pack_ue8m0_x4
from . import MXFP4_BLOCK_SIZE, _fp32x2_to_fp4x2
@@ -51,7 +53,8 @@ def indexer_k_norm_rope_store(
kv_slot_mapping: [num_tokens] slots in the indexer cache (-1 = skip).
compress_ratio: group size; keys are emitted at group boundaries.
use_fp4_cache: MXFP4 (2 nibbles/byte + ue8m0 per 32) when True, else
per-token FP8 with a single fp32 scale.
per-token FP8 with a single fp32 scale (four packed ue8m0 bytes
under `mqa_logits_use_mx_scales`).
"""
num_tokens = kv_slot_mapping.numel()
@@ -111,6 +114,7 @@ def indexer_k_norm_rope_store(
KV_BLOCK_STRIDE=k_cache.stride(0),
FP8_MAX=448.0,
USE_FP4=use_fp4_cache,
MX_SCALE=not use_fp4_cache and mqa_logits_use_mx_scales(),
SHUFFLE=shuffle,
BLOCK_TILE_SIZE=_BLOCK_TILE_SIZE,
HEAD_TILE_SIZE=_HEAD_TILE_SIZE,
@@ -139,6 +143,7 @@ def _indexer_k_norm_rope_quant_store_kernel(
KV_BLOCK_STRIDE: tl.constexpr,
FP8_MAX: tl.constexpr,
USE_FP4: tl.constexpr,
MX_SCALE: tl.constexpr,
SHUFFLE: tl.constexpr,
BLOCK_TILE_SIZE: tl.constexpr,
HEAD_TILE_SIZE: tl.constexpr,
@@ -252,4 +257,10 @@ def _indexer_k_norm_rope_quant_store_kernel(
tl.store(val_ptr + tiled, x_uint8)
else:
tl.store(val_ptr + block, x_uint8)
tl.store(scale_ptr.to(tl.pointer_type(tl.float32)), tl.exp2(exponent))
if MX_SCALE:
tl.store(
scale_ptr.to(tl.pointer_type(tl.int32)),
pack_ue8m0_x4(tl.exp2(exponent)),
)
else:
tl.store(scale_ptr.to(tl.pointer_type(tl.float32)), tl.exp2(exponent))
vllm/models/glm5next/nvidia/ops/kpool_compress.py
+26 -4
View File
@@ -15,6 +15,8 @@ from __future__ import annotations
import torch
from vllm.triton_utils import tl, triton
from vllm.utils.deep_gemm import mqa_logits_use_mx_scales
from vllm.v1.attention.ops.common import pack_ue8m0_x4
# The GLM-5.3-Flash indexer head dimension is fixed at 128.
INDEX_HEAD_DIM = 128
@@ -157,6 +159,7 @@ def _kpool_softmax_rotate_write_cache_kernel(
HEAD_DIM: tl.constexpr,
S_OFFSET_NBYTES_IN_PAGE: tl.constexpr,
ROUND_SCALE: tl.constexpr,
MX_SCALE: tl.constexpr,
HAS_WRITE_MASK: tl.constexpr,
RETURN_COMPRESSED: tl.constexpr,
WRITE_CACHE: tl.constexpr,
@@ -245,7 +248,14 @@ def _kpool_softmax_rotate_write_cache_kernel(
+ loc_token_offset_in_page
)
tl.store(buf_fp8_ptr + out_k_offsets, quantized, mask=mask)
tl.store(buf_fp32_ptr + out_s_offset, scale, mask=do_write)
if MX_SCALE:
tl.store(
buf_fp32_ptr.to(tl.pointer_type(tl.int32)) + out_s_offset,
pack_ue8m0_x4(scale),
mask=do_write,
)
else:
tl.store(buf_fp32_ptr + out_s_offset, scale, mask=do_write)
if RETURN_COMPRESSED:
tl.store(
@@ -324,6 +334,7 @@ def kpool_compress_and_write_cache(
buf_fp8 = buf.view(torch.float8_e4m3fn)
buf_fp32 = buf.view(torch.float32)
mx_scale = write_cache and mqa_logits_use_mx_scales()
# bytes per page (last dim of kv_cache) viewed as uint8
buf_numel_per_page = buf.stride(0)
s_offset_nbytes_in_page = page_size * head_dim
@@ -361,7 +372,9 @@ def kpool_compress_and_write_cache(
POOL_SIZE=slot_k.shape[1],
HEAD_DIM=head_dim,
S_OFFSET_NBYTES_IN_PAGE=s_offset_nbytes_in_page,
ROUND_SCALE=round_scale,
# MX scales need power-of-two scales.
ROUND_SCALE=round_scale or mx_scale,
MX_SCALE=mx_scale,
HAS_WRITE_MASK=has_write_mask,
RETURN_COMPRESSED=return_compressed,
WRITE_CACHE=write_cache,
@@ -479,6 +492,7 @@ def _kpool_decode_update_batched_kernel(
HEAD_DIM: tl.constexpr,
S_OFFSET_NBYTES_IN_PAGE: tl.constexpr,
ROUND_SCALE: tl.constexpr,
MX_SCALE: tl.constexpr,
BLOCK_D: tl.constexpr,
):
"""One program per request; iterates its NEXT_N verify tokens in order.
@@ -604,7 +618,13 @@ def _kpool_decode_update_batched_kernel(
+ loc_token_offset_in_page
)
tl.store(buf_fp8_ptr + out_k_offsets, quantized, mask=dim_mask)
tl.store(buf_fp32_ptr + out_s_offset, scale)
if MX_SCALE:
tl.store(
buf_fp32_ptr.to(tl.pointer_type(tl.int32)) + out_s_offset,
pack_ue8m0_x4(scale),
)
else:
tl.store(buf_fp32_ptr + out_s_offset, scale)
# Stash the current token AFTER any completion read so the completion
# uses prior stashes (and the current token's own key/score via
@@ -718,7 +738,9 @@ def kpool_decode_update_and_maybe_write_cache_batched(
KPOOL_HEAD=tail_kv_cache.stride(1),
HEAD_DIM=head_dim,
S_OFFSET_NBYTES_IN_PAGE=page_size * head_dim,
ROUND_SCALE=round_scale,
# MX scales need power-of-two scales.
ROUND_SCALE=round_scale or mqa_logits_use_mx_scales(),
MX_SCALE=mqa_logits_use_mx_scales(),
BLOCK_D=triton.next_power_of_2(head_dim),
)
vllm/models/glm5next/nvidia/sparse_indexer.py
+46 -26
View File
@@ -16,6 +16,7 @@ from vllm.model_executor.layers.sparse_attn_indexer import (
PrefillRowShard,
_merge_dcp_topk_global,
_reserve_row_shard_exchange_transient,
as_fp32_logits,
)
from vllm.models.glm5next.common.sparse_indexer import (
RADIX_TOPK_WORKSPACE_SIZE,
@@ -29,7 +30,13 @@ from vllm.models.glm5next.common.sparse_indexer import (
)
from vllm.models.glm5next.nvidia.ops import kpool_compress as kpool_ops
from vllm.platforms import current_platform
from vllm.utils.deep_gemm import has_deep_gemm
from vllm.utils.deep_gemm import (
fp8_fp4_mqa_logits,
fp8_fp4_paged_mqa_logits,
has_deep_gemm,
mqa_logits_mx_inputs,
mqa_logits_use_mx_scales,
)
from vllm.utils.import_utils import has_cutedsl
from vllm.utils.torch_utils import (
LayerNameType,
@@ -253,7 +260,7 @@ def sparse_attn_indexer_kpool(
kv_cache,
slot_mapping,
quant_block_size,
scale_fmt,
"mxfp8" if mqa_logits_use_mx_scales() else scale_fmt,
)
topk_indices_buffer[: hidden_states.shape[0]] = -1
@@ -343,20 +350,29 @@ def sparse_attn_indexer_kpool(
else:
q_slice_cast = q_slice
k_quant_cast = k_quant
k_scale_cast = k_scale.view(torch.float32).squeeze(-1)
from vllm.utils.deep_gemm import fp8_fp4_mqa_logits
k_scale_cast = k_scale.view(
torch.int32 if mqa_logits_use_mx_scales() else torch.float32
).squeeze(-1)
num_rows = q_slice.shape[0]
if chunk.local_total_seq_lens == 0:
logits = q_slice.new_empty((num_rows, 0), dtype=torch.float32)
else:
logits = fp8_fp4_mqa_logits(
(q_slice_cast, q_scale_slice),
(k_quant_cast, k_scale_cast),
weights[row_start:row_end],
cu_seqlen_ks,
cu_seqlen_ke,
clean_logits=False,
q_scale_slice, chunk_weights = mqa_logits_mx_inputs(
q_slice_cast, q_scale_slice, weights[row_start:row_end]
)
# Compressed: row i starts at K position cu_seqlen_ks[i].
row_ends = cu_seqlen_ke - cu_seqlen_ks
logits = as_fp32_logits(
fp8_fp4_mqa_logits(
(q_slice_cast, q_scale_slice),
(k_quant_cast, k_scale_cast),
chunk_weights,
cu_seqlen_ks,
cu_seqlen_ke,
chunk.max_seqlen_k,
),
row_ends,
)
# kpool: logits are pool-granular (compress_ratio == index_kpool),
@@ -376,8 +392,8 @@ def sparse_attn_indexer_kpool(
else:
torch.ops._C.top_k_per_row_prefill(
logits,
cu_seqlen_ks,
cu_seqlen_ke,
torch.zeros_like(row_ends),
row_ends,
topk_dst,
num_rows,
logits.stride(0),
@@ -392,7 +408,7 @@ def sparse_attn_indexer_kpool(
dcp_rank,
dcp_world_size,
pool_cp_interleave,
row_starts=cu_seqlen_ks,
row_starts=None,
)
if index_kpool > 1:
@@ -594,18 +610,22 @@ def sparse_attn_indexer_kpool(
if use_fp4_cache
else padded_q_quant_decode_tokens
)
from vllm.utils.deep_gemm import fp8_fp4_paged_mqa_logits
logits = fp8_fp4_paged_mqa_logits(
(padded_q_quant_cast, padded_q_scale),
kv_cache,
padded_weights[:num_padded_tokens],
seq_lens,
decode_metadata.block_table,
decode_metadata.schedule_metadata,
max_model_len=max_pool_len,
clean_logits=False,
indices=decode_metadata.indices,
padded_q_scale, decode_weights = mqa_logits_mx_inputs(
padded_q_quant_cast, padded_q_scale, padded_weights[:num_padded_tokens]
)
# SM100 decode is varlen: one row and context length per token.
logits = as_fp32_logits(
fp8_fp4_paged_mqa_logits(
(padded_q_quant_cast, padded_q_scale),
kv_cache,
decode_weights,
seq_lens,
decode_metadata.block_table,
decode_metadata.schedule_metadata,
max_model_len=max_pool_len,
indices=decode_metadata.indices,
),
seq_lens.view(-1),
)
num_rows = logits.shape[0]
# kpool: logits are pool-granular -> select topk_tokens//kpool pools,
vllm/utils/deep_gemm.py
+72 -29
View File
@@ -591,37 +591,79 @@ def transform_sf_into_required_layout(*args, **kwargs):
)
@functools.cache
def mqa_logits_use_mx_scales() -> bool:
"""Whether the MQA logits kernels take only MX-scaled inputs.
On SM100 they need MXFP8/MXFP4 Q and K with packed UE8M0 scales and BF16
weights, and return BF16 logits. FP8 indexer K caches then store each
token's power-of-two scale as four equal UE8M0 bytes instead of an fp32.
"""
return current_platform.is_cuda() and current_platform.is_device_capability_family(
100
)
# Four packed UE8M0 bytes encoding 1.0.
_UNIT_UE8M0_X4 = 0x7F7F7F7F
# Grown buffers are kept alive: CUDA graphs captured earlier still read them.
_unit_mx_scales: dict[torch.device, list[torch.Tensor]] = {}
def mqa_logits_mx_inputs(
q: torch.Tensor, q_scale: torch.Tensor | None, weights: torch.Tensor
) -> tuple[torch.Tensor | None, torch.Tensor]:
"""Adapt indexer Q scales and weights to the SM100 MQA logits kernels.
FP8 Q folds its per-(token, head) power-of-two scale into ``weights``, so
it is passed with unit MX scales, which is exact since relu commutes with
a positive scale. Returns ``(q_scale, weights)`` unchanged off SM100.
"""
if not mqa_logits_use_mx_scales():
return q_scale, weights
if q_scale is None:
numel = q.numel() // q.shape[-1]
bufs = _unit_mx_scales.setdefault(q.device, [])
if not bufs or bufs[-1].numel() < numel:
size = max(numel, 2 * bufs[-1].numel() if bufs else 0)
bufs.append(
torch.full((size,), _UNIT_UE8M0_X4, dtype=torch.int32, device=q.device)
)
q_scale = bufs[-1][:numel].view(q.shape[:-1])
return q_scale, weights.to(torch.bfloat16)
def fp8_fp4_mqa_logits(
q: tuple[torch.Tensor, torch.Tensor | None],
kv: tuple[torch.Tensor, torch.Tensor],
weights: torch.Tensor,
cu_seqlen_ks: torch.Tensor,
cu_seqlen_ke: torch.Tensor,
clean_logits: bool,
max_seqlen_k: int,
) -> torch.Tensor:
"""Compute MQA logits for a single sequence without KV paging.
Unified FP8/FP4 dispatch — the underlying DeepGEMM kernel takes
``q = (values, scales_or_None)`` where ``scales`` is None for FP8 Q
(per-token scale is folded into ``weights``) and a packed block-scale
tensor for MXFP4 Q.
``q = (values, scales_or_None)``. SM100 takes MXFP8 or MXFP4 with packed
UE8M0 block scales; SM90/SM120 also take FP8 Q with ``q_scale=None``
(per-token scale folded into ``weights``).
Args:
q: Tuple ``(q_values, q_scale)``. FP8 path: q_values is [M, H, D]
float8_e4m3fn and q_scale is None (per-token scale is folded
into ``weights``). FP4 path: q_values is packed uint8 and
q_scale is the companion block-scale tensor.
kv: Tuple `(k_packed, k_scales)` — FP8 layout is [N, D]
float8_e4m3fn plus fp32 scales [N]; FP4 layout is packed uint8.
weights: weights of shape [M, H], dtype `torch.float32`.
q: Tuple ``(q_values, q_scale)``, q_values of shape [M, H, D] (packed
to D/2 bytes for MXFP4) and q_scale int32 [M, H] or None.
kv: Tuple ``(k_values, k_scales)`` of shape [N, D] and [N]; k_scales
is int32 packed UE8M0 with MX scales, fp32 otherwise.
weights: Tensor of shape [M, H], bf16 on SM100 and fp32 elsewhere.
cu_seqlen_ks: Start indices (inclusive) for valid K per query
position, shape [M], dtype int32.
cu_seqlen_ke: End indices (exclusive) for valid K per query
position, shape [M], dtype int32.
clean_logits: Whether to clean the unfilled logits into `-inf`.
max_seqlen_k: Upper bound on ``cu_seqlen_ke - cu_seqlen_ks``.
Returns:
Logits tensor of shape [M, N], dtype `torch.float32`.
Compressed logits of shape [M, max_seqlen_k]: row ``i`` holds K
positions ``[cu_seqlen_ks[i], cu_seqlen_ke[i])`` starting at column
zero, and the tail is not cleaned. bf16 on SM100, fp32 elsewhere.
"""
_lazy_init()
@@ -633,7 +675,7 @@ def fp8_fp4_mqa_logits(
weights,
cu_seqlen_ks,
cu_seqlen_ke,
clean_logits=clean_logits,
max_seqlen_k,
)
@@ -699,24 +741,25 @@ def fp8_fp4_paged_mqa_logits(
block_tables: torch.Tensor,
schedule_metadata: torch.Tensor,
max_model_len: int,
clean_logits: bool,
indices: torch.Tensor | None = None,
) -> torch.Tensor:
"""Compute MQA logits using a paged KV-cache.
Unified FP8/FP4 dispatch — the underlying DeepGEMM kernel takes
``q = (values, scales_or_None)``; pass ``(q_tensor, None)`` for the FP8
path and ``(q_values, q_scale)`` for MXFP4.
``q = (values, scales_or_None)``; see `fp8_fp4_mqa_logits` for the
per-architecture formats. SM100 requires the varlen form (``indices``
with next_n == 1).
Args:
q: Tuple ``(q_values, q_scale)``. FP8 path: q_values is
[B, next_n, H, D] float8_e4m3fn and q_scale is None. FP4 path:
q_values is packed uint8 and q_scale is the companion
block-scale tensor.
kv_cache: Paged KV-cache. FP8 layout is [num_blocks, block_size, 1,
D+4], dtype `torch.uint8`, with the last 4 bytes per (block, pos)
storing the float dequant scale.
weights: Tensor of shape [B * next_n, H], dtype `torch.float32`.
q: Tuple ``(q_values, q_scale)``, q_values of shape
[B, next_n, H, D] (packed to D/2 bytes for MXFP4) and q_scale
int32 [B, next_n, H] or None.
kv_cache: Paged KV-cache of shape [num_blocks, block_size, 1, D+4]
(D/2+4 for MXFP4), dtype `torch.uint8`; the last 4 bytes per
(block, pos) hold the int32 packed UE8M0 scales with MX scales,
or the fp32 dequant scale otherwise.
weights: Tensor of shape [B * next_n, H], bf16 on SM100 and fp32
elsewhere.
context_lens: Tensor of shape [B], dtype int32; effective context length
for each batch element.
block_tables: Tensor of shape [B, max_blocks], dtype int32; maps logical
@@ -724,12 +767,11 @@ def fp8_fp4_paged_mqa_logits(
schedule_metadata: Returned by `get_paged_mqa_logits_metadata`;
used to distribute work across SMs.
max_model_len: Maximum sequence length used to size the logits output.
clean_logits: Whether to clean the unfilled logits into `-inf`.
indices: Optional request index for each varlen row.
Returns:
Logits tensor of shape [B * next_n, max_model_len], dtype
`torch.float32`.
Logits tensor of shape [B * next_n, max_model_len]; columns past each
row's context length are not cleaned. bf16 on SM100, fp32 elsewhere.
"""
_lazy_init()
@@ -751,7 +793,6 @@ def fp8_fp4_paged_mqa_logits(
block_tables,
schedule_metadata,
max_model_len,
clean_logits=clean_logits,
**kwargs,
)
@@ -1052,6 +1093,8 @@ __all__ = [
"fp8_fp4_mqa_logits",
"fp8_fp4_paged_mqa_logits",
"get_paged_mqa_logits_metadata",
"mqa_logits_mx_inputs",
"mqa_logits_use_mx_scales",
"native_next_n_supported",
"per_block_cast_to_fp8",
"is_deep_gemm_e8m0_used",
vllm/v1/attention/backends/mla/indexer.py
+4
View File
@@ -467,6 +467,9 @@ class DeepseekV32IndexerPrefillChunkMetadata:
cu_seq_lens: torch.Tensor
token_to_seq: torch.Tensor
total_seq_lens: int
# Host upper bound on cu_seqlen_ke - cu_seqlen_ks, sizing the compressed
# DeepGEMM prefill logits.
max_seqlen_k: int
token_start: int
token_end: int
num_reqs: int
@@ -1957,6 +1960,7 @@ def build_prefill_chunk_metadata(
cu_seq_lens=cu_seq_lens,
token_to_seq=token_to_seq,
total_seq_lens=total_seq_lens,
max_seqlen_k=int(compressed_seq_lens_cpu[start_idx:end_idx].max()),
block_table=block_table[start_idx:end_idx],
token_start=token_start,
token_end=token_end,
vllm/v1/attention/ops/common.py
+7
View File
@@ -154,6 +154,13 @@ class PackSeqTritonKernel(VllmTritonJitKernel["PackSeqTritonKernel.CompileKey"])
)
@triton.jit
def pack_ue8m0_x4(scale):
"""Encode a power-of-two fp32 scale as four equal packed UE8M0 bytes, the
MX scales of a 128-wide row (see `mqa_logits_use_mx_scales`)."""
return ((scale.to(tl.int32, bitcast=True) >> 23) & 0xFF) * 0x01010101
def pack_seq_triton(
x: torch.Tensor,
lengths: torch.Tensor,