diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index 1dcee345f8..b796623b7a 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -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) diff --git a/csrc/cache_kernels.cu b/csrc/cache_kernels.cu index d1c1e9df98..24453ff1c1 100644 --- a/csrc/cache_kernels.cu +++ b/csrc/cache_kernels.cu @@ -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(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(kv_cache)[dst_scale_idx / 4] = + ((__float_as_uint(scale) >> 23) & 0xFF) * 0x01010101u; + } else { + reinterpret_cast(kv_cache)[dst_scale_idx / 4] = scale; + } } } @@ -1912,7 +1919,7 @@ void cp_gather_and_upconvert_nvfp4_kv_cache( reinterpret_cast(k.data_ptr()), \ reinterpret_cast(kv_cache.data_ptr()), \ slot_mapping.const_data_ptr(), 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(), diff --git a/tests/kernels/attention/test_deepgemm_attention.py b/tests/kernels/attention/test_deepgemm_attention.py index f332fb0fea..109cc83830 100644 --- a/tests/kernels/attention/test_deepgemm_attention.py +++ b/tests/kernels/attention/test_deepgemm_attention.py @@ -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) ] diff --git a/tests/kernels/attention/test_rocm_paged_mxfp4_indexer.py b/tests/kernels/attention/test_rocm_paged_mxfp4_indexer.py index 311ee89a8d..e270ae4a58 100644 --- a/tests/kernels/attention/test_rocm_paged_mxfp4_indexer.py +++ b/tests/kernels/attention/test_rocm_paged_mxfp4_indexer.py @@ -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, diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index cc4f4c72a9..e992422905 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -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]) diff --git a/tests/kernels/test_fused_deepseek_v32_norm_rope.py b/tests/kernels/test_fused_deepseek_v32_norm_rope.py index 63bbe0c95f..ec61b94abd 100644 --- a/tests/kernels/test_fused_deepseek_v32_norm_rope.py +++ b/tests/kernels/test_fused_deepseek_v32_norm_rope.py @@ -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() diff --git a/tests/kernels/test_fused_indexer_q_rope_quant.py b/tests/kernels/test_fused_indexer_q_rope_quant.py index 0ce3a91851..c6dfa11b0f 100644 --- a/tests/kernels/test_fused_indexer_q_rope_quant.py +++ b/tests/kernels/test_fused_indexer_q_rope_quant.py @@ -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) diff --git a/tests/kernels/test_kpool_decode_update_batched.py b/tests/kernels/test_kpool_decode_update_batched.py index 4accf527f3..5dd9eff5f1 100644 --- a/tests/kernels/test_kpool_decode_update_batched.py +++ b/tests/kernels/test_kpool_decode_update_batched.py @@ -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 diff --git a/tests/model_executor/layers/test_mla_short_prefill_indexer.py b/tests/model_executor/layers/test_mla_short_prefill_indexer.py index ba4ff895cb..7ca4cff404 100644 --- a/tests/model_executor/layers/test_mla_short_prefill_indexer.py +++ b/tests/model_executor/layers/test_mla_short_prefill_indexer.py @@ -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") diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index 464fce66f2..62c8bc32bf 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -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 diff --git a/tests/v1/attention/test_rocm_paged_mxfp4_indexer_plan.py b/tests/v1/attention/test_rocm_paged_mxfp4_indexer_plan.py index 0283ed99ee..a3dbec86dc 100644 --- a/tests/v1/attention/test_rocm_paged_mxfp4_indexer_plan.py +++ b/tests/v1/attention/test_rocm_paged_mxfp4_indexer_plan.py @@ -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], diff --git a/tools/build_deepgemm_C.py b/tools/build_deepgemm_C.py index 36824318c0..dec8b0105a 100644 --- a/tools/build_deepgemm_C.py +++ b/tools/build_deepgemm_C.py @@ -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, ) diff --git a/tools/check_wheel_deepgemm.py b/tools/check_wheel_deepgemm.py index fcc35ead58..c8da1351bb 100644 --- a/tools/check_wheel_deepgemm.py +++ b/tools/check_wheel_deepgemm.py @@ -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 diff --git a/tools/install_deepgemm.sh b/tools/install_deepgemm.sh index 2fb2351a0c..b617e6d567 100755 --- a/tools/install_deepgemm.sh +++ b/tools/install_deepgemm.sh @@ -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 diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 935e4c0d2d..3af3a161ac 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -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: diff --git a/vllm/models/deepseek_v32/common/kernels.py b/vllm/models/deepseek_v32/common/kernels.py index 255b1f5aaf..f876120a6e 100644 --- a/vllm/models/deepseek_v32/common/kernels.py +++ b/vllm/models/deepseek_v32/common/kernels.py @@ -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, diff --git a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py index 50ca86558e..703169f482 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py +++ b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py @@ -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, ) diff --git a/vllm/models/deepseek_v41/common/ops/indexer_k_store.py b/vllm/models/deepseek_v41/common/ops/indexer_k_store.py index ba8ddbd95a..fe203e9178 100644 --- a/vllm/models/deepseek_v41/common/ops/indexer_k_store.py +++ b/vllm/models/deepseek_v41/common/ops/indexer_k_store.py @@ -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)) diff --git a/vllm/models/glm5next/nvidia/ops/kpool_compress.py b/vllm/models/glm5next/nvidia/ops/kpool_compress.py index eb289c65cf..310d60dfd4 100644 --- a/vllm/models/glm5next/nvidia/ops/kpool_compress.py +++ b/vllm/models/glm5next/nvidia/ops/kpool_compress.py @@ -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), ) diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py index 154fdcda90..1a3fbedb83 100644 --- a/vllm/models/glm5next/nvidia/sparse_indexer.py +++ b/vllm/models/glm5next/nvidia/sparse_indexer.py @@ -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, diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 04c4508f0b..3dc9be5d93 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -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", diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 2d797ae662..0382a81fe0 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -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, diff --git a/vllm/v1/attention/ops/common.py b/vllm/v1/attention/ops/common.py index 5c1677cac4..6936be9fa1 100644 --- a/vllm/v1/attention/ops/common.py +++ b/vllm/v1/attention/ops/common.py @@ -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,