diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu b/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu index 8d4ba1accc7c..25be762b9f03 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu +++ b/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu @@ -173,3 +173,37 @@ void silu_and_mul_scaled_fp4_experts_quant( STD_TORCH_CHECK_NOT_IMPLEMENTED( false, "No compiled silu_and_mul nvfp4 experts quantization kernel"); } + +// MXFP4 link stubs for archs without a compiled MXFP4 kernel. On SM10.x the +// real impls in `mxfp4_experts_quant.cu` are compiled (same CMake block that +// defines `ENABLE_NVFP4_SM100`); on SM 8.x the real file is skipped, leaving +// `torch_bindings.cpp`'s unconditional `&mxfp4_experts_quant` reference with +// no symbol to bind to, so the `.so` fails to load. These stubs satisfy the +// linker on unsupported archs and raise NOT_IMPLEMENTED if anyone actually +// calls them. +#if !(defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) +void mxfp4_experts_quant( + torch::stable::Tensor& output, torch::stable::Tensor& output_scale, + torch::stable::Tensor const& input, + torch::stable::Tensor const& input_offset_by_experts, + torch::stable::Tensor const& output_scale_offset_by_experts, + int64_t n_experts) { + STD_TORCH_CHECK_NOT_IMPLEMENTED( + false, "No compiled mxfp4 experts quantization kernel for SM ", + get_sm_version_num(), + ". Recompile with SM 10.x (Blackwell) for MXFP4 support."); +} + +void silu_and_mul_mxfp4_experts_quant( + torch::stable::Tensor& output, torch::stable::Tensor& output_scale, + torch::stable::Tensor const& input, + torch::stable::Tensor const& input_offset_by_experts, + torch::stable::Tensor const& output_scale_offset_by_experts, + int64_t n_experts) { + STD_TORCH_CHECK_NOT_IMPLEMENTED( + false, + "No compiled silu_and_mul mxfp4 experts quantization kernel for SM ", + get_sm_version_num(), + ". Recompile with SM 10.x (Blackwell) for MXFP4 support."); +} +#endif // !ENABLE_NVFP4_SM100 diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index 0b0df945a885..00329e112e43 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -141,7 +141,8 @@ Priority is **1 = highest** (tried first). | 2 | `FLASHMLA` | | 3 | `FLASHINFER_MLA` | | 4 | `TRITON_MLA` | -| 5 | `FLASHMLA_SPARSE` | +| 5 | `TRITON_MLA_SPARSE` | +| 6 | `FLASHMLA_SPARSE` | > **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` is preferred for low query-head counts (<= 16), while `FLASHMLA_SPARSE` is preferred otherwise. > @@ -226,4 +227,5 @@ MLA decode backends are selected using the standard | `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | | `TOKENSPEED_MLA` | fp16, bf16 | `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x | | `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | Any | +| `TRITON_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | Any | | `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | Any | diff --git a/tests/kernels/attention/test_mqa_logits_triton.py b/tests/kernels/attention/test_mqa_logits_triton.py new file mode 100644 index 000000000000..57849bcac88e --- /dev/null +++ b/tests/kernels/attention/test_mqa_logits_triton.py @@ -0,0 +1,317 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Correctness tests for the Triton MQA logits kernels.""" + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.utils.math_utils import cdiv +from vllm.v1.attention.ops.mqa_logits_triton import ( + fp8_mqa_logits_triton, + fp8_paged_mqa_logits_triton, +) + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="Triton MQA logits kernels require CUDA/ROCm", +) + + +def _quantize_k_per_row( + k_bf16: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + amax = k_bf16.abs().float().amax(dim=-1, keepdim=True).clamp_min(1e-4) + sf = amax / 448.0 + k_fp8 = (k_bf16.float() / sf).to(torch.float8_e4m3fn) + return k_fp8, sf.squeeze(-1) + + +def _pack_paged_kv(kv_bf16: torch.Tensor) -> torch.Tensor: + """Pack BF16 KV into the layout produced by `indexer_k_quant_and_cache`. + + Physical bytes per block are segregated: all `block_size * head_dim` fp8 K + bytes first, then `block_size * 4` fp32 scale bytes. The outer + `[NB, BS, 1, D+4]` shape matches the production cache allocation. + """ + num_blocks, block_size, head_dim = kv_bf16.shape + amax = kv_bf16.abs().float().amax(dim=-1, keepdim=True).clamp_min(1e-4) + sf = (amax / 448.0).to(torch.float32) + k_fp8 = (kv_bf16.float() / sf).to(torch.float8_e4m3fn) + packed = torch.empty( + (num_blocks, block_size, 1, head_dim + 4), + dtype=torch.uint8, + device=kv_bf16.device, + ) + flat = packed.view(num_blocks, -1) + k_end = block_size * head_dim + flat[:, :k_end] = k_fp8.reshape(num_blocks, -1).view(torch.uint8) + flat[:, k_end:] = sf.reshape(num_blocks, -1).view(torch.uint8) + return packed + + +# References adapted from DeepGEMM (test_attention.py) — used as the spec +# the Triton kernels must agree with. +def _fp8_mqa_logits_ref( + q: torch.Tensor, + kv: tuple[torch.Tensor, torch.Tensor], + weights: torch.Tensor, + cu_seqlen_ks: torch.Tensor, + cu_seqlen_ke: torch.Tensor, +) -> torch.Tensor: + k_fp8, scale = kv + seq_len_kv = k_fp8.shape[0] + k = k_fp8.to(torch.bfloat16) + q = q.to(torch.bfloat16) + arange = torch.arange(0, seq_len_kv, device=q.device)[None, :] + mask = (arange >= cu_seqlen_ks[:, None]) & (arange < cu_seqlen_ke[:, None]) + score = torch.einsum("mhd,nd->hmn", q, k).float() * scale + logits = (score.relu() * weights.unsqueeze(-1).transpose(0, 1)).sum(dim=0) + return logits.masked_fill(~mask, float("-inf")) + + +def _fp8_paged_mqa_logits_ref( + q: torch.Tensor, + kv_cache: torch.Tensor, + weights: torch.Tensor, + context_lens: torch.Tensor, + block_tables: torch.Tensor, + max_model_len: int, +) -> torch.Tensor: + fp8_dtype = torch.float8_e4m3fn + batch_size, next_n, _, dim = q.size() + num_blocks, block_size = kv_cache.shape[0], kv_cache.shape[1] + flat = kv_cache.view(num_blocks, -1) + k_end = block_size * dim + kv_data = ( + flat[:, :k_end].reshape(num_blocks, block_size, 1, dim).view(fp8_dtype).float() + ) + scale = flat[:, k_end:].view(torch.float32).reshape(num_blocks, block_size, 1, 1) + q = q.float() + kv_data = kv_data * scale + logits = torch.full( + [batch_size * next_n, max_model_len], + float("-inf"), + device=q.device, + dtype=torch.float32, + ) + context_lens_list = context_lens.tolist() + for i in range(batch_size): + context_len = context_lens_list[i] + q_offsets = torch.arange(context_len - next_n, context_len, device=q.device) + weight_slice = ( + weights[i * next_n : (i + 1) * next_n, :].transpose(0, 1).contiguous() + ) + for block_rk in range(cdiv(context_len, block_size)): + block_idx = block_tables[i][block_rk] + qx, kx = q[i], kv_data[block_idx] + k_offsets = torch.arange( + block_rk * block_size, + (block_rk + 1) * block_size, + device=q.device, + ) + mask = (k_offsets[None, :] < context_len) & ( + k_offsets[None, :] <= q_offsets[:, None] + ) + s = torch.where( + mask[None, :, :], + (qx.transpose(0, 1) @ kx.transpose(0, 1).transpose(1, 2)).to( + logits.dtype + ), + float("-inf"), + ) + s = (torch.relu(s) * weight_slice[..., None]).sum(dim=0) + logits[ + i * next_n : (i + 1) * next_n, + block_rk * block_size : (block_rk + 1) * block_size, + ] = torch.where(k_offsets[None, :] <= q_offsets[:, None], s, float("-inf")) + return logits + + +# Looser tolerance to accommodate FP8 rounding and the paged torch +# reference using fp32 matmul while the triton kernel uses bf16 matmul +# (with an fp32 accumulator, matching the DeepGEMM path). +_ATOL = 1.0 +_RTOL = 0.2 + + +@pytest.mark.parametrize("M,N", [(64, 64), (128, 256), (256, 512)]) +@pytest.mark.parametrize("num_heads", [16, 32]) +@pytest.mark.parametrize("partial_mask", [False, True]) +@pytest.mark.parametrize("clean_logits", [True, False]) +def test_fp8_mqa_logits_triton_matches_torch( + M, N, num_heads, partial_mask, clean_logits +): + """`clean_logits=True` requires full-tensor agreement with the reference; + `clean_logits=False` is only contractually correct on the in-range + `[ks, ke)` slots, which is what `finite` already masks here (the reference + writes -inf outside that range).""" + torch.manual_seed(0) + head_dim = 128 + device = "cuda" + + q_bf16 = torch.randn(M, num_heads, head_dim, dtype=torch.bfloat16, device=device) + k_bf16 = torch.randn(N, head_dim, dtype=torch.bfloat16, device=device) + weights = torch.randn(M, num_heads, dtype=torch.float32, device=device) + q_fp8 = q_bf16.to(torch.float8_e4m3fn) + k_fp8, k_scales = _quantize_k_per_row(k_bf16) + + if partial_mask: + # Exercise the non-trivial ks/ke masking branch (chunked prefill). + ks = torch.arange(M, dtype=torch.int32, device=device) % (N // 4) + ke = ks + torch.randint(1, N // 2, (M,), dtype=torch.int32, device=device) + ke = torch.minimum(ke, torch.tensor(N, dtype=torch.int32, device=device)) + else: + ks = torch.zeros(M, dtype=torch.int32, device=device) + ke = torch.full((M,), N, dtype=torch.int32, device=device) + + out_torch = _fp8_mqa_logits_ref(q_fp8, (k_fp8, k_scales), weights, ks, ke) + out_triton = fp8_mqa_logits_triton( + q_fp8, (k_fp8, k_scales), weights, ks, ke, clean_logits=clean_logits + ) + + if clean_logits: + assert torch.equal( + torch.isinf(out_torch) & (out_torch < 0), + torch.isinf(out_triton) & (out_triton < 0), + ) + finite = ~torch.isinf(out_torch) + if finite.any(): + torch.testing.assert_close( + out_triton[finite], out_torch[finite], atol=_ATOL, rtol=_RTOL + ) + + +def test_fp8_mqa_logits_triton_clean_logits_false_overwrites_masked(): + """Regression test for PR #38476 comment 4398225404: when `clean_logits=False` + skips the `-inf` pre-fill, the kernel itself must still write `-inf` to + masked positions. Otherwise the K-tile early-exit branch leaves + uninitialized memory in the output, downstream top-k reads garbage, and the + model collapses to repeating a single token. + + Trigger: a narrow per-row `[ks, ke)` so that BLOCK_N tiles past every row's + `ke` hit the early-exit path. We assert that `clean_logits=False` produces + bit-identical output to `clean_logits=True` on every position (the flag is + a perf opt, not a behavior change).""" + torch.manual_seed(0) + M, N, num_heads, head_dim = 256, 1024, 32, 128 + device = "cuda" + + q_fp8 = torch.randn(M, num_heads, head_dim, dtype=torch.bfloat16, device=device).to( + torch.float8_e4m3fn + ) + k_bf16 = torch.randn(N, head_dim, dtype=torch.bfloat16, device=device) + weights = torch.randn(M, num_heads, dtype=torch.float32, device=device) + k_fp8, k_scales = _quantize_k_per_row(k_bf16) + + # Narrow per-row range: most rows' `[ks, ke)` cover ~1/16 of N, so most + # K-tiles past `ke` hit the kernel's early-exit branch. + ks = torch.arange(M, dtype=torch.int32, device=device) % (N // 8) + ke = ks + N // 16 + ke = torch.minimum(ke, torch.full_like(ke, N)) + + out_clean = fp8_mqa_logits_triton( + q_fp8, (k_fp8, k_scales), weights, ks, ke, clean_logits=True + ) + out_dirty = fp8_mqa_logits_triton( + q_fp8, (k_fp8, k_scales), weights, ks, ke, clean_logits=False + ) + + # Every position the clean run wrote `-inf` to (i.e. outside `[ks, ke)`) + # must also be `-inf` in the dirty run — otherwise the early-exit left + # uninitialized memory and downstream top-k will pick from garbage. + inf_mask = torch.isinf(out_clean) & (out_clean < 0) + bad = inf_mask & ~(torch.isinf(out_dirty) & (out_dirty < 0)) + if bad.any(): + # Surface a useful error: how many positions, and a sample of the + # leaked values so it's obvious they're not `-inf`. + leaked_count = int(bad.sum().item()) + sample = out_dirty[bad].flatten()[:8].tolist() + raise AssertionError( + f"clean_logits=False left {leaked_count} positions un-initialised " + f"(should be -inf). Sample leaked values: {sample}. " + "Likely cause: the kernel's K-tile early-exit branch returns " + "without writing -inf, so downstream top-k reads garbage." + ) + + +@pytest.mark.parametrize( + "batch_size,next_n,context_len", + [ + (1, 1, 128), + (1, 1, 512), + (2, 1, 256), + (1, 4, 512), # speculative decoding with next_n=4 + ], +) +@pytest.mark.parametrize("num_heads", [16, 32]) +@pytest.mark.parametrize("clean_logits", [True, False]) +def test_fp8_paged_mqa_logits_triton_matches_torch( + batch_size, next_n, context_len, num_heads, clean_logits +): + """`clean_logits=True` requires whole-tensor agreement; `clean_logits=False` + is only correct on `[:, :context_len]` (downstream topk reads only that + range), so all comparisons below are sliced accordingly.""" + torch.manual_seed(0) + head_dim = 128 + block_size = 64 + device = "cuda" + + total_blocks = 64 + max_blocks = (context_len + block_size - 1) // block_size + 4 + + kv_bf16 = torch.randn( + total_blocks, block_size, head_dim, dtype=torch.bfloat16, device=device + ) + kv_packed = _pack_paged_kv(kv_bf16) + + q_fp8 = torch.randn( + batch_size, + next_n, + num_heads, + head_dim, + dtype=torch.bfloat16, + device=device, + ).to(torch.float8_e4m3fn) + + weights = torch.randn( + batch_size * next_n, num_heads, dtype=torch.float32, device=device + ) + + context_lens = torch.full( + (batch_size,), context_len, dtype=torch.int32, device=device + ) + block_tables = torch.randint( + 0, + total_blocks, + (batch_size, max_blocks), + dtype=torch.int32, + device=device, + ) + + max_model_len = max_blocks * block_size + + out_torch = _fp8_paged_mqa_logits_ref( + q_fp8, kv_packed, weights, context_lens, block_tables, max_model_len + ) + out_triton = fp8_paged_mqa_logits_triton( + q_fp8, + kv_packed, + weights, + context_lens, + block_tables, + max_model_len, + clean_logits=clean_logits, + ) + + inf_torch = (torch.isinf(out_torch) & (out_torch < 0))[:, :context_len] + inf_triton = (torch.isinf(out_triton) & (out_triton < 0))[:, :context_len] + assert torch.equal(inf_torch, inf_triton) + finite = ~inf_torch + if finite.any(): + torch.testing.assert_close( + out_triton[:, :context_len][finite], + out_torch[:, :context_len][finite], + atol=_ATOL, + rtol=_RTOL, + ) diff --git a/tests/kernels/attention/test_triton_mla_sparse_kernel.py b/tests/kernels/attention/test_triton_mla_sparse_kernel.py new file mode 100644 index 000000000000..0522341d9e5d --- /dev/null +++ b/tests/kernels/attention/test_triton_mla_sparse_kernel.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Correctness tests for the Triton sparse MLA kernel. + +Compares split-KV against the single-pass (`num_kv_splits=1`) path +produced by the same kernel — both paths must agree to within bf16 ULPs. +""" + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.v1.attention.ops.triton_mla_sparse_kernel import ( + _DIM_QK, + triton_mla_sparse_attention, +) + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="Triton sparse MLA kernel requires CUDA/ROCm", +) + + +@pytest.fixture(scope="module") +def kv_cache(): + torch.manual_seed(0) + return torch.randn(32768, 1, _DIM_QK, dtype=torch.bfloat16, device="cuda") + + +def _assert_split_matches_single_pass( + num_tokens: int, + num_heads: int, + topk: int, + num_kv_splits: int | None, + kv_cache: torch.Tensor, +) -> None: + torch.manual_seed(0) + q = torch.randn(num_tokens, num_heads, _DIM_QK, dtype=torch.bfloat16, device="cuda") + indices = torch.randint( + 0, kv_cache.shape[0], (num_tokens, 1, topk), dtype=torch.int32, device="cuda" + ) + out_ref = triton_mla_sparse_attention( + q, + kv_cache, + indices, + sm_scale=0.1, + num_kv_splits=1, + ) + out = triton_mla_sparse_attention( + q, + kv_cache, + indices, + sm_scale=0.1, + num_kv_splits=num_kv_splits, + ) + torch.testing.assert_close( + out.float(), + out_ref.float(), + atol=5e-2, + rtol=5e-3, + ) + + +@pytest.mark.parametrize( + "num_tokens,num_heads", + [(1, 16), (1, 128), (8, 32), (32, 128), (128, 16)], +) +@pytest.mark.parametrize("topk", [1024, 2048, 4096]) +@pytest.mark.parametrize("num_kv_splits", [2, 4, 8]) +def test_split_kv_matches_single_pass( + num_tokens, num_heads, topk, num_kv_splits, kv_cache +): + _assert_split_matches_single_pass( + num_tokens, + num_heads, + topk, + num_kv_splits, + kv_cache, + ) + + +@pytest.mark.parametrize("num_tokens", [1, 8, 32, 128]) +def test_auto_split_matches_single_pass(num_tokens, kv_cache): + _assert_split_matches_single_pass( + num_tokens, + num_heads=128, + topk=2048, + num_kv_splits=None, + kv_cache=kv_cache, + ) + + +@pytest.mark.parametrize("num_kv_splits", [1, 2, 4, 8]) +def test_short_prefill_no_nan(num_kv_splits, kv_cache): + """Regression: short prefill where most topk slots are -1 sentinels. + + The indexer fills 2048 topk positions with only a handful of valid + indices; the rest are -1. Before the NEG_LARGE sentinel fix, the online + softmax produced NaN via `max(-inf, -inf) = -inf` and + `exp2(-inf − -inf) = NaN`, poisoning every split. + """ + torch.manual_seed(0) + num_tokens, num_heads, topk = 5, 16, 2048 + q = torch.randn(num_tokens, num_heads, _DIM_QK, dtype=torch.bfloat16, device="cuda") + indices = torch.full((num_tokens, 1, topk), -1, dtype=torch.int32, device="cuda") + # Only the first `t+1` slots of each query hold valid indices; the + # remaining ~2045 slots are -1, producing many all-invalid BLOCK_N tiles. + for t in range(num_tokens): + indices[t, 0, : t + 1] = torch.arange( + 64, 64 + t + 1, dtype=torch.int32, device="cuda" + ) + out = triton_mla_sparse_attention( + q, kv_cache, indices, sm_scale=0.0417, num_kv_splits=num_kv_splits + ) + assert not torch.isnan(out).any() + assert not torch.isinf(out).any() diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 4bf52a49c43f..a42765e98258 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -13,7 +13,7 @@ from vllm.utils.deep_gemm import ( fp8_fp4_mqa_logits, fp8_fp4_paged_mqa_logits, - has_deep_gemm, + is_deep_gemm_supported, ) from vllm.utils.torch_utils import ( LayerNameType, @@ -25,6 +25,10 @@ DeepseekV32IndexerMetadata, ) from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton +from vllm.v1.attention.ops.mqa_logits_triton import ( + fp8_mqa_logits_triton, + fp8_paged_mqa_logits_triton, +) from vllm.v1.worker.workspace import current_workspace_manager if current_platform.is_cuda_alike(): @@ -175,6 +179,12 @@ def sparse_attn_indexer( ) topk_indices_buffer[: hidden_states.shape[0]] = -1 + # DeepGEMM availability is constant per process; check once for both branches. + use_deep_gemm = is_deep_gemm_supported() + if not use_deep_gemm: + assert not use_fp4_cache, ( + "Triton sparse-MLA fallback does not support FP4 KV cache" + ) if has_prefill: prefill_metadata = attn_metadata_narrowed.prefill assert prefill_metadata is not None @@ -220,14 +230,24 @@ def sparse_attn_indexer( q_slice_cast = q_slice k_quant_cast = k_quant k_scale_cast = k_scale.view(torch.float32).squeeze(-1) - logits = fp8_fp4_mqa_logits( - (q_slice_cast, q_scale_slice), - (k_quant_cast, k_scale_cast), - weights[chunk.token_start : chunk.token_end], - chunk.cu_seqlen_ks, - chunk.cu_seqlen_ke, - clean_logits=False, - ) + if use_deep_gemm: + logits = fp8_fp4_mqa_logits( + (q_slice_cast, q_scale_slice), + (k_quant_cast, k_scale_cast), + weights[chunk.token_start : chunk.token_end], + chunk.cu_seqlen_ks, + chunk.cu_seqlen_ke, + clean_logits=False, + ) + else: + logits = fp8_mqa_logits_triton( + q_slice_cast, + (k_quant_cast, k_scale_cast), + weights[chunk.token_start : chunk.token_end], + chunk.cu_seqlen_ks, + chunk.cu_seqlen_ke, + clean_logits=False, + ) num_rows = logits.shape[0] topk_indices = topk_indices_buffer[ @@ -307,16 +327,31 @@ def sparse_attn_indexer( if use_fp4_cache else padded_q_quant_decode_tokens ) - 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, - ) + if use_deep_gemm: + 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, + ) + else: + # SM80/SM121 Triton fallback. Downstream topk reads only up to + # `seq_lens`, so size the buffer to the active batch max rather + # than the configured model max. + active_max_model_len = attn_metadata_narrowed.max_seq_len + logits = fp8_paged_mqa_logits_triton( + padded_q_quant_cast, + kv_cache, + weights[:num_padded_tokens], + seq_lens, + decode_metadata.block_table, + max_model_len=active_max_model_len, + clean_logits=False, + ) num_rows = logits.shape[0] topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens] @@ -438,9 +473,15 @@ def __init__( self.topk_indices_buffer = topk_indices_buffer self.skip_k_cache_insert = skip_k_cache_insert self.use_fp4_cache = use_fp4_cache - if current_platform.is_cuda() and not has_deep_gemm(): - raise RuntimeError( - "Sparse Attention Indexer CUDA op requires DeepGEMM to be installed." + # On SM80/SM121 (A100, GB10) DeepGEMM is unavailable — fall back to + # the Triton sparse-MLA path. is_deep_gemm_supported() encodes the + # SM-arch + has_deep_gemm() gate; if not supported, downgrade the + # hard error from upstream to a one-time warning so the indexer + # routes through the Triton kernels in `mqa_logits_triton.py`. + if current_platform.is_cuda() and not is_deep_gemm_supported(): + logger.warning_once( + "DeepGEMM not supported on this platform; " + "using Triton fallback for sparse attention indexer." ) def forward_native( diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index e990b46dd92d..c592be70a34b 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -126,6 +126,7 @@ def _get_backend_priorities( AttentionBackendEnum.FLASHMLA, AttentionBackendEnum.FLASHINFER_MLA, AttentionBackendEnum.TRITON_MLA, + AttentionBackendEnum.TRITON_MLA_SPARSE, AttentionBackendEnum.FLASHMLA_SPARSE, ] else: diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 7c0715a9e8b6..6e2509696e97 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -11,7 +11,7 @@ from vllm.triton_utils import tl, triton from vllm.utils.deep_gemm import ( get_paged_mqa_logits_metadata, - has_deep_gemm, + is_deep_gemm_supported, ) from vllm.utils.math_utils import cdiv from vllm.utils.platform_utils import num_compute_units @@ -611,7 +611,7 @@ def build( seq_lens = seq_lens.unsqueeze(-1) # DeepGEMM is required for the paged MQA logits on CUDA devices - if current_platform.is_cuda() and has_deep_gemm(): + if current_platform.is_cuda() and is_deep_gemm_supported(): self.scheduler_metadata_buffer[:] = get_paged_mqa_logits_metadata( seq_lens, self.kv_cache_spec.storage_block_size, diff --git a/vllm/v1/attention/backends/mla/triton_mla_sparse.py b/vllm/v1/attention/backends/mla/triton_mla_sparse.py new file mode 100644 index 000000000000..f9a3b0975488 --- /dev/null +++ b/vllm/v1/attention/backends/mla/triton_mla_sparse.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Pure-Triton sparse MLA backend for SM80 (A100) / SM121 (GB10).""" + +from typing import ClassVar + +import torch + +from vllm.config import get_current_vllm_config_or_none +from vllm.config.cache import CacheDType +from vllm.platforms.interface import DeviceCapability +from vllm.utils.platform_utils import num_compute_units +from vllm.v1.attention.backend import AttentionBackend, AttentionCGSupport +from vllm.v1.attention.backends.mla.xpu_mla_sparse import ( + XPUMLASparseImpl, + XPUMLASparseMetadata, + XPUMLASparseMetadataBuilder, +) +from vllm.v1.attention.ops.mqa_logits_triton import ( + warmup_fp8_mqa_logits_triton, + warmup_fp8_paged_mqa_logits_triton, +) +from vllm.v1.attention.ops.triton_mla_sparse_kernel import ( + _DIM_QK, + KV_SPLITS_CANDIDATES, + triton_mla_sparse_attention, +) + +# V3.2 indexers don't expose `n_head`; GLM-5.1-NVFP4 sets index_n_heads=32. +# Autotune key includes (num_heads, head_dim), so a wrong warmup shape forces +# a re-tune on first real request. +_INDEXER_NUM_HEADS = 64 +_INDEXER_HEAD_DIM = 128 + + +class TritonMLASparseMetadataBuilder(XPUMLASparseMetadataBuilder): + # XPU base keeps NEVER (not validated under cudagraph); this subclass + # claims UNIFORM_BATCH for the CUDA/Triton path. + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH + + +class TritonMLASparseImpl(XPUMLASparseImpl): + """Triton sparse-MLA impl with split-KV decode (3-7× faster than the + single-pass XPU base for single-query decode on SM80 / SM121).""" + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self._sm_count: int | None = None + if self.topk_indices_buffer is not None: + self._sm_count = num_compute_units(self.topk_indices_buffer.device.index) + self._warmup_autotune(kwargs["indexer"]) + + def _warmup_autotune(self, indexer) -> None: + """Prime `@triton.autotune` caches at init so the first request + doesn't pay the inline config-sweep cost.""" + if self.topk_indices_buffer is None: + return + device = self.topk_indices_buffer.device + topk = self.topk_indices_buffer.shape[-1] + q = torch.empty(1, self.num_heads, _DIM_QK, dtype=torch.bfloat16, device=device) + kv = torch.empty(64, 1, _DIM_QK, dtype=torch.bfloat16, device=device) + indices = torch.zeros(1, 1, topk, dtype=torch.int32, device=device) + for splits in KV_SPLITS_CANDIDATES: + triton_mla_sparse_attention( + q, + kv, + indices, + sm_scale=self.softmax_scale, + num_kv_splits=splits, + sm_count=self._sm_count, + ) + indexer_num_heads = getattr(indexer, "n_head", _INDEXER_NUM_HEADS) + indexer_head_dim = getattr(indexer, "head_dim", _INDEXER_HEAD_DIM) + warmup_fp8_mqa_logits_triton( + num_heads=indexer_num_heads, head_dim=indexer_head_dim, device=device + ) + cfg = get_current_vllm_config_or_none() + if cfg is not None: + warmup_fp8_paged_mqa_logits_triton( + num_heads=indexer_num_heads, + head_dim=indexer_head_dim, + block_size=cfg.cache_config.block_size, + device=device, + ) + + def _forward_bf16_kv( + self, + q: torch.Tensor, # [sq, heads, d_qk] + kv_c_and_k_pe_cache: torch.Tensor, # [blocks, heads, d_qk] + topk_indices: torch.Tensor, # [sq, topk] + attn_metadata: XPUMLASparseMetadata, + ) -> torch.Tensor: + num_tokens = q.shape[0] + kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view( + -1, 1, kv_c_and_k_pe_cache.shape[-1] + ) + topk_indices = topk_indices.view(num_tokens, 1, -1) + output = triton_mla_sparse_attention( + q, + kv_c_and_k_pe_cache, + topk_indices, + sm_scale=self.softmax_scale, + sm_count=self._sm_count, + ) + return output[:, : self.num_heads, :] + + +class TritonMLASparseBackend(AttentionBackend): + supported_dtypes: ClassVar[list[torch.dtype]] = [ + torch.float16, + torch.bfloat16, + ] + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ + "auto", + "float16", + "bfloat16", + ] + + @staticmethod + def get_name() -> str: + return "TRITON_MLA_SPARSE" + + @staticmethod + def get_metadata_cls() -> type[XPUMLASparseMetadata]: + return XPUMLASparseMetadata + + @staticmethod + def get_builder_cls() -> type["TritonMLASparseMetadataBuilder"]: + return TritonMLASparseMetadataBuilder + + @staticmethod + def get_impl_cls() -> type["TritonMLASparseImpl"]: + return TritonMLASparseImpl + + @classmethod + def is_mla(cls) -> bool: + return True + + @classmethod + def is_sparse(cls) -> bool: + return True + + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: + return (num_blocks, block_size, head_size) + + @classmethod + def get_supported_head_sizes(cls) -> list[int]: + return [_DIM_QK] + + @classmethod + def supports_compute_capability(cls, capability: DeviceCapability) -> bool: + return True diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index 87abb6884313..147f16f9b2cf 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -71,6 +71,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): "FlashInferMLASparseBackend" ) TRITON_MLA = "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend" + TRITON_MLA_SPARSE = ( + "vllm.v1.attention.backends.mla.triton_mla_sparse.TritonMLASparseBackend" + ) CUTLASS_MLA = "vllm.v1.attention.backends.mla.cutlass_mla.CutlassMLABackend" FLASHMLA = "vllm.v1.attention.backends.mla.flashmla.FlashMLABackend" FLASHMLA_SPARSE = ( diff --git a/vllm/v1/attention/ops/mqa_logits_triton.py b/vllm/v1/attention/ops/mqa_logits_triton.py new file mode 100644 index 000000000000..d27b2411ba92 --- /dev/null +++ b/vllm/v1/attention/ops/mqa_logits_triton.py @@ -0,0 +1,455 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Triton fallback for DeepGEMM's fp8_mqa_logits / fp8_paged_mqa_logits.""" + +import torch + +from vllm.triton_utils import tl, triton + +# Paged decode: num_warps=4 dominated on A100/SM80 across {2,4,8}; the others +# were 1.5–1.7× slower at (num_heads=32, head_dim=128, block_size=64), so +# narrow the sweep to keep autotune from latching onto a bad pick under noise. +_PAGED_AUTOTUNE_CONFIGS = [ + triton.Config({}, num_warps=4, num_stages=ns) for ns in (2, 4) +] + +# Prefill kernel adds BLOCK_N as a free tile axis. num_warps=8 was 1.5–3× +# worse than {2,4} across the sweep; keep BLOCK_N ∈ {32, 64, 128} so autotune +# can pick per shape (BN=128 wins for GLM-5.1 long chunks). +_PREFILL_AUTOTUNE_CONFIGS = [ + triton.Config({"BLOCK_N": bn}, num_warps=nw, num_stages=ns) + for bn in (32, 64, 128) + for nw in (2, 4) + for ns in (2, 4) +] + +# Warmup shape mirrors the chunked-prefill regime (small M, long N) so +# autotune picks a tile sized for real serving rather than a launch-overhead- +# dominated dummy grid. +_PREFILL_WARMUP_M = 8 +_PREFILL_WARMUP_N = 8192 + + +_E4M3FN_BF16_LUT_CACHE: dict[torch.device, torch.Tensor] = {} + + +def _get_e4m3fn_bf16_lut(device: torch.device) -> torch.Tensor: + lut = _E4M3FN_BF16_LUT_CACHE.get(device) + if lut is not None: + return lut + lut = ( + torch.arange(256, dtype=torch.uint8, device=device) + .view(torch.float8_e4m3fn) + .to(torch.bfloat16) + ) + lut[0x7F] = 480.0 + lut[0xFF] = -480.0 + _E4M3FN_BF16_LUT_CACHE[device] = lut + return lut + + +@triton.jit +def _decode_e4m3fn_bf16_lut(u, lut_ptr): + return tl.load(lut_ptr + u.to(tl.uint32)) + + +@triton.autotune( + configs=_PAGED_AUTOTUNE_CONFIGS, + key=["num_heads", "head_dim", "block_size"], +) +@triton.jit +def _fp8_paged_mqa_logits_kernel( + q_ptr, + kv_fp8_ptr, + kv_scale_ptr, + weights_ptr, + fp8_lut_ptr, + context_lens_ptr, + block_tables_ptr, + logits_ptr, + stride_q_b, + stride_q_n, + stride_q_h, + stride_q_d, + stride_kvf_block, + stride_kvf_s, + stride_kvf_d, + stride_kvs_block, + stride_kvs_s, + stride_w_t, + stride_w_h, + stride_bt_b, + stride_bt_k, + stride_l_t, + stride_l_n, + next_n: tl.constexpr, + num_heads: tl.constexpr, + head_dim: tl.constexpr, + block_size: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_N: tl.constexpr, +): + token_id = tl.program_id(0) + block_rk = tl.program_id(1) + + batch_id = token_id // next_n + next_n_id = token_id % next_n + + context_len = tl.load(context_lens_ptr + batch_id) + if block_rk * block_size >= context_len: + return + + q_offset = context_len - next_n + next_n_id + + block_idx = tl.load( + block_tables_ptr + batch_id * stride_bt_b + block_rk * stride_bt_k + ) + + offs_h = tl.arange(0, BLOCK_H) + offs_d = tl.arange(0, BLOCK_D) + offs_n = tl.arange(0, BLOCK_N) + mask_h = offs_h < num_heads + mask_d = offs_d < head_dim + mask_n = offs_n < block_size + + q_base = q_ptr + batch_id * stride_q_b + next_n_id * stride_q_n + q_byte = tl.load( + q_base + offs_h[:, None] * stride_q_h + offs_d[None, :] * stride_q_d, + mask=mask_h[:, None] & mask_d[None, :], + other=0, + ) + q = _decode_e4m3fn_bf16_lut(q_byte, fp8_lut_ptr) + + kvf_base = kv_fp8_ptr + block_idx * stride_kvf_block + k_byte = tl.load( + kvf_base + offs_n[:, None] * stride_kvf_s + offs_d[None, :] * stride_kvf_d, + mask=mask_n[:, None] & mask_d[None, :], + other=0, + ) + kvs_base = kv_scale_ptr + block_idx * stride_kvs_block + k_scale = tl.load( + kvs_base + offs_n * stride_kvs_s, + mask=mask_n, + other=0.0, + ) + k = _decode_e4m3fn_bf16_lut(k_byte, fp8_lut_ptr) + # Scale in fp32 after the dot to avoid an extra bf16 round-trip on K. + s = tl.dot(q, tl.trans(k)) * k_scale[None, :] + + w = tl.load( + weights_ptr + token_id * stride_w_t + offs_h * stride_w_h, + mask=mask_h, + other=0.0, + ) + s = tl.where(s > 0, s, 0.0) * w[:, None] + out = tl.sum(s, axis=0) + + k_offset = block_rk * block_size + offs_n + valid = mask_n & (k_offset < context_len) & (k_offset <= q_offset) + out = tl.where(valid, out, float("-inf")) + + tl.store( + logits_ptr + token_id * stride_l_t + k_offset * stride_l_n, + out, + mask=mask_n, + ) + + +def fp8_paged_mqa_logits_triton( + q: torch.Tensor, + kv_cache: torch.Tensor, + weights: torch.Tensor, + context_lens: torch.Tensor, + block_tables: torch.Tensor, + max_model_len: int, + clean_logits: bool = True, +) -> torch.Tensor: + """Triton implementation of DeepGEMM's fp8_paged_mqa_logits. + + Args: + q: [B, next_n, H, D] fp8_e4m3fn + kv_cache: [num_blocks, block_size, 1, D+4] uint8 (FP8 + fp32 scale) + weights: [B*next_n, H] float32 + context_lens: [B] int32 + block_tables: [B, max_blocks] int32 + max_model_len: output width. Caller passes the active batch max so + the logits buffer and grid stay tight. + clean_logits: when False, skip the -inf pre-fill of the output + (indexer top-k reads only `[:context_len]` per row). + Returns: + logits: [B*next_n, max_model_len] float32 + """ + B, next_n, num_heads, head_dim = q.shape + _, block_size, one, d_plus_4 = kv_cache.shape + assert one == 1 + assert d_plus_4 == head_dim + 4 + + # Cache layout from `indexer_k_quant_and_cache`: per block, FP8 K bytes + # (block_size * head_dim) followed by fp32 scales (block_size * 4). The + # `[NB, block_size, 1, head_dim+4]` shape is a stride trick; re-slice flat. + # Kernel decodes FP8 from uint8 via LUT (SM80 Triton can't load fp8e4nv). + num_blocks = kv_cache.shape[0] + kv_flat = kv_cache.view(num_blocks, -1) + k_end = block_size * head_dim + kv_byte = kv_flat[:, :k_end].as_strided( + (num_blocks, block_size, head_dim), + (kv_flat.stride(0), head_dim, 1), + ) + kv_scale = kv_flat[:, k_end:].view(torch.float32) + q_byte = q.view(torch.uint8) + + if clean_logits: + logits = torch.full( + (B * next_n, max_model_len), + float("-inf"), + dtype=torch.float32, + device=q.device, + ) + else: + logits = torch.empty( + (B * next_n, max_model_len), dtype=torch.float32, device=q.device + ) + + BLOCK_H = max(16, triton.next_power_of_2(num_heads)) + BLOCK_D = triton.next_power_of_2(head_dim) + BLOCK_N = triton.next_power_of_2(block_size) + + fp8_lut = _get_e4m3fn_bf16_lut(q.device) + grid = (B * next_n, block_tables.shape[1]) + _fp8_paged_mqa_logits_kernel[grid]( + q_byte, + kv_byte, + kv_scale, + weights, + fp8_lut, + context_lens, + block_tables, + logits, + q_byte.stride(0), + q_byte.stride(1), + q_byte.stride(2), + q_byte.stride(3), + kv_byte.stride(0), + kv_byte.stride(1), + kv_byte.stride(2), + kv_scale.stride(0), + kv_scale.stride(1), + weights.stride(0), + weights.stride(1), + block_tables.stride(0), + block_tables.stride(1), + logits.stride(0), + logits.stride(1), + next_n=next_n, + num_heads=num_heads, + head_dim=head_dim, + block_size=block_size, + BLOCK_H=BLOCK_H, + BLOCK_D=BLOCK_D, + BLOCK_N=BLOCK_N, + ) + return logits + + +@triton.autotune( + configs=_PREFILL_AUTOTUNE_CONFIGS, + # Per-program work is N-independent; key on (heads, dim) only so chunked + # prefill with varying N doesn't re-tune on every new chunk size. + key=["num_heads", "head_dim"], +) +@triton.jit +def _fp8_mqa_logits_kernel( + q_ptr, + k_ptr, + k_scale_ptr, + weights_ptr, + ks_ptr, + ke_ptr, + logits_ptr, + stride_q_m, + stride_q_h, + stride_q_d, + stride_k_n, + stride_k_d, + stride_w_m, + stride_w_h, + stride_l_m, + stride_l_n, + num_heads: tl.constexpr, + head_dim: tl.constexpr, + N, + BLOCK_H: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_N: tl.constexpr, +): + # bf16 q/k inputs: the wrapper pre-decodes FP8 → bf16. At compute-bound + # prefill this is ~2× the in-kernel LUT (LUT lookups contend with the + # matmul for ALU/regs). Paged-decode keeps the LUT path. + m = tl.program_id(0) + n_block = tl.program_id(1) + + n_start = n_block * BLOCK_N + offs_n = n_start + tl.arange(0, BLOCK_N) + mask_n = offs_n < N + # Early-exit when this row's `[ks, ke)` range doesn't overlap the tile. + # Chunked prefill produces many such all-masked tiles per row. + ks = tl.load(ks_ptr + m) + ke = tl.load(ke_ptr + m) + if (n_start >= ke) | (n_start + BLOCK_N <= ks): + # When `clean_logits=False` the caller skipped the -inf pre-fill, so + # write -inf here for the early-exit tile. + tl.store( + logits_ptr + m * stride_l_m + offs_n * stride_l_n, + tl.full([BLOCK_N], float("-inf"), dtype=tl.float32), + mask=mask_n, + ) + return + + offs_h = tl.arange(0, BLOCK_H) + offs_d = tl.arange(0, BLOCK_D) + mask_h = offs_h < num_heads + mask_d = offs_d < head_dim + + q = tl.load( + q_ptr + + m * stride_q_m + + offs_h[:, None] * stride_q_h + + offs_d[None, :] * stride_q_d, + mask=mask_h[:, None] & mask_d[None, :], + other=0.0, + ) + + k = tl.load( + k_ptr + offs_n[:, None] * stride_k_n + offs_d[None, :] * stride_k_d, + mask=mask_n[:, None] & mask_d[None, :], + other=0.0, + ) + k_scale = tl.load(k_scale_ptr + offs_n, mask=mask_n, other=0.0) + s = tl.dot(q, tl.trans(k)) * k_scale[None, :] + + w = tl.load( + weights_ptr + m * stride_w_m + offs_h * stride_w_h, + mask=mask_h, + other=0.0, + ) + s = tl.where(s > 0, s, 0.0) * w[:, None] + out = tl.sum(s, axis=0) + + valid = mask_n & (offs_n >= ks) & (offs_n < ke) + out = tl.where(valid, out, float("-inf")) + + tl.store( + logits_ptr + m * stride_l_m + offs_n * stride_l_n, + out, + mask=mask_n, + ) + + +def fp8_mqa_logits_triton( + q: torch.Tensor, + kv: tuple[torch.Tensor, torch.Tensor], + weights: torch.Tensor, + cu_seqlen_ks: torch.Tensor, + cu_seqlen_ke: torch.Tensor, + clean_logits: bool = True, +) -> torch.Tensor: + """Triton implementation of DeepGEMM's fp8_mqa_logits. + + Args: + q: [M, H, D] fp8_e4m3fn + kv: (k_fp8 [N, D], k_scales [N]) — fp8_e4m3fn, float32 + weights: [M, H] float32 + cu_seqlen_ks: [M] int32 + cu_seqlen_ke: [M] int32 + clean_logits: when False, skip the -inf pre-fill of the output + (indexer top-k reads only `[ks, ke)` per row). Matches DeepGEMM. + Returns: + logits: [M, N] float32 + """ + k_fp8, k_scales = kv + k_scales = k_scales.reshape(-1) + + M, num_heads, head_dim = q.shape + N = k_fp8.shape[0] + + if clean_logits: + logits = torch.full((M, N), float("-inf"), dtype=torch.float32, device=q.device) + else: + logits = torch.empty((M, N), dtype=torch.float32, device=q.device) + + BLOCK_H = max(16, triton.next_power_of_2(num_heads)) + BLOCK_D = triton.next_power_of_2(head_dim) + + # Pre-decode FP8 → bf16; the kernel runs a straight `tl.dot`. + q_bf16 = q.to(torch.bfloat16) + k_bf16 = k_fp8.to(torch.bfloat16) + + # Grid depends on the autotuned BLOCK_N. + grid = lambda meta: (M, triton.cdiv(N, meta["BLOCK_N"])) # noqa: E731 + _fp8_mqa_logits_kernel[grid]( + q_bf16, + k_bf16, + k_scales, + weights, + cu_seqlen_ks, + cu_seqlen_ke, + logits, + q_bf16.stride(0), + q_bf16.stride(1), + q_bf16.stride(2), + k_bf16.stride(0), + k_bf16.stride(1), + weights.stride(0), + weights.stride(1), + logits.stride(0), + logits.stride(1), + num_heads=num_heads, + head_dim=head_dim, + N=N, + BLOCK_H=BLOCK_H, + BLOCK_D=BLOCK_D, + ) + return logits + + +def warmup_fp8_mqa_logits_triton( + num_heads: int, + head_dim: int, + device: torch.device, +) -> None: + """Prime the prefill `@triton.autotune` cache so first-call doesn't pay + the inline sweep (~5–8 s on A100 SM80). N is a runtime scalar, so one + small-M / long-N shape covers all chunk lengths.""" + max_block_n = max(c.kwargs["BLOCK_N"] for c in _PREFILL_AUTOTUNE_CONFIGS) + m = _PREFILL_WARMUP_M + n = max(_PREFILL_WARMUP_N, max_block_n) + q = torch.empty(m, num_heads, head_dim, dtype=torch.float8_e4m3fn, device=device) + k = torch.empty(n, head_dim, dtype=torch.float8_e4m3fn, device=device) + scales = torch.zeros(n, dtype=torch.float32, device=device) + weights = torch.zeros(m, num_heads, dtype=torch.float32, device=device) + ks = torch.zeros(m, dtype=torch.int32, device=device) + ke = torch.full((m,), n, dtype=torch.int32, device=device) + fp8_mqa_logits_triton(q, (k, scales), weights, ks, ke) + + +def warmup_fp8_paged_mqa_logits_triton( + num_heads: int, + head_dim: int, + block_size: int, + device: torch.device, +) -> None: + """Prime the paged-decode `@triton.autotune` cache for the indexer's + logits kernel (see `warmup_fp8_mqa_logits_triton` for rationale). + """ + num_blocks = 2 + q = torch.empty(1, 1, num_heads, head_dim, dtype=torch.float8_e4m3fn, device=device) + kv_cache = torch.zeros( + num_blocks, block_size, 1, head_dim + 4, dtype=torch.uint8, device=device + ) + weights = torch.zeros(1, num_heads, dtype=torch.float32, device=device) + context_lens = torch.tensor([block_size], dtype=torch.int32, device=device) + block_tables = torch.zeros(1, 1, dtype=torch.int32, device=device) + fp8_paged_mqa_logits_triton( + q, kv_cache, weights, context_lens, block_tables, max_model_len=block_size + ) diff --git a/vllm/v1/attention/ops/triton_mla_sparse_kernel.py b/vllm/v1/attention/ops/triton_mla_sparse_kernel.py new file mode 100644 index 000000000000..0654c019a7c4 --- /dev/null +++ b/vllm/v1/attention/ops/triton_mla_sparse_kernel.py @@ -0,0 +1,550 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Triton sparse MLA attention with split-KV for low-batch decode.""" + +import functools + +import torch + +from vllm.triton_utils import LOG2E, LOGE2, tl, triton +from vllm.utils.platform_utils import num_compute_units + +# DeepSeek-V3.2 / GLM-5 sparse MLA shape constants. +_BLOCK_DMODEL = 512 +_BLOCK_DPE = 64 +_BLOCK_DV = 512 +_DIM_QK = _BLOCK_DMODEL + _BLOCK_DPE # 576 + +_BLOCK_H = 16 +# Smallest BLOCK_N the autotune sweep offers; only used for the topk-divisibility +# check at dispatch time. +_MIN_BLOCK_N = 16 + +# Merge kernel grid is spread across heads and DV tiles to avoid a (1,1) +# launch starving the SMs (pattern from FlashMLA's combine kernel). +_MERGE_BLOCK_H = 1 +_MERGE_BLOCK_DV_TILE = 128 +assert _BLOCK_DV % _MERGE_BLOCK_DV_TILE == 0 +_NUM_MERGE_DV_TILES = _BLOCK_DV // _MERGE_BLOCK_DV_TILE + +# Final (prefill) and split (decode) kernels each tune to their own regime. +_FINAL_AUTOTUNE_CONFIGS = [ + triton.Config({"BLOCK_N": 16}, num_warps=nw, num_stages=ns) + for nw in (2, 4) + for ns in (2, 4) +] +_SPLIT_AUTOTUNE_CONFIGS = [ + triton.Config({"BLOCK_N": 32}, num_warps=4, num_stages=ns) for ns in (2, 4) +] + +# Split-count candidates for `_choose_num_kv_splits`; also the set pre-compiled +# by `_warmup_autotune`. +KV_SPLITS_CANDIDATES = (1, 2, 4, 8, 16) + +_MIN_TOPK_PER_SPLIT = 128 # below this, per-split work is too small to amortize +_SPLIT_MAX_OCCUPANCY = 4 # skip split when baseline grid fills >=1/4 of SMs + + +@triton.jit +def _sparse_mla_compute_tile( + q_buffer, + k_buffer, # V is the first BLOCK_DV lanes of each row of k_buffer. + indices_ptr, + cur_q, + cur_head, + cur_kv_head_id, + mask_h, + split_start, + split_end, + seq_kv, + stride_q_token, + stride_q_head, + stride_kv_token, + stride_kv_head, + stride_indices_token, + stride_indices_head, + sm_scale, + BLOCK_H: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_DV: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_DPE: tl.constexpr, +): + """Shared stage-1 body: load Q, run the sparse online-softmax loop over + `[split_start, split_end)` of the topk axis, return accumulators.""" + offs_d = tl.arange(0, BLOCK_DMODEL) + offs_dpe = BLOCK_DMODEL + tl.arange(0, BLOCK_DPE) + offs_dv = tl.arange(0, BLOCK_DV) + mask_dpe = offs_dpe < BLOCK_DMODEL + BLOCK_DPE + + q = tl.load( + q_buffer + + cur_q * stride_q_token + + cur_head[:, None] * stride_q_head + + offs_d[None, :], + mask=mask_h[:, None], + other=0.0, + ) + qpe = tl.load( + q_buffer + + cur_q * stride_q_token + + cur_head[:, None] * stride_q_head + + offs_dpe[None, :], + mask=(mask_h[:, None]) & (mask_dpe[None, :]), + other=0.0, + ) + + # Finite sentinel (not -inf) — when an entire BLOCK_N tile is masked, + # `-inf - -inf = NaN` poisons the softmax; `sentinel - sentinel = 0` + # gives `exp2(0) = 1` and the matching V rows are already 0. + NEG_LARGE = -1.0e30 + e_max = tl.zeros([BLOCK_H], dtype=tl.float32) + NEG_LARGE + e_sum = tl.zeros([BLOCK_H], dtype=tl.float32) + acc = tl.zeros([BLOCK_H, BLOCK_DV], dtype=tl.float32) + + for start_indice in range(split_start, split_end, BLOCK_N): + offs_indice = start_indice + tl.arange(0, BLOCK_N) + mask_indice = offs_indice < split_end + indices = tl.load( + indices_ptr + + cur_q * stride_indices_token + + cur_kv_head_id * stride_indices_head + + offs_indice, + mask=mask_indice, + other=-1, + ) + mask_kv = (indices >= 0) & (indices < seq_kv) + + offs_k = ( + indices[None, :] * stride_kv_token + + cur_kv_head_id * stride_kv_head + + offs_d[:, None] + ) + k = tl.load(k_buffer + offs_k, mask=mask_kv[None, :], other=0.0) + qk = tl.dot(q, k.to(q.dtype)) + + offs_kpe = ( + indices[None, :] * stride_kv_token + + cur_kv_head_id * stride_kv_head + + offs_dpe[:, None] + ) + kpe = tl.load( + k_buffer + offs_kpe, + mask=(mask_kv[None, :]) & (mask_dpe[:, None]), + other=0.0, + ) + qk += tl.dot(qpe, kpe.to(q.dtype)) + + qk *= sm_scale + qk = tl.where((mask_h[:, None]) & (mask_kv[None, :]), qk, NEG_LARGE) + + offs_v = ( + indices[:, None] * stride_kv_token + + cur_kv_head_id * stride_kv_head + + offs_dv[None, :] + ) + v = tl.load(k_buffer + offs_v, mask=mask_kv[:, None], other=0.0) + + n_e_max = tl.maximum(tl.max(qk, 1), e_max) + re_scale = tl.exp2(e_max - n_e_max) + p = tl.exp2(qk - n_e_max[:, None]) + acc *= re_scale[:, None] + acc += tl.dot(p.to(v.dtype), v) + e_sum = e_sum * re_scale + tl.sum(p, 1) + e_max = n_e_max + + return acc, e_max, e_sum + + +@triton.autotune(configs=_FINAL_AUTOTUNE_CONFIGS, key=["index_topk", "kv_group_num"]) +@triton.jit +def _sparse_mla_kernel_final( + q_buffer, + k_buffer, + indices_ptr, + out_ptr, + seq_kv, + h_q, + stride_q_token, + stride_q_head, + stride_kv_token, + stride_kv_head, + stride_out_token, + stride_out_head, + stride_indices_token, + stride_indices_head, + sm_scale, + index_topk: tl.constexpr, + kv_group_num: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_DV: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_DPE: tl.constexpr, +): + """Single-pass fast path: full topk, write final bf16 output directly.""" + cur_q = tl.program_id(0) + cur_head_id = tl.program_id(1) + cur_kv_head_id = cur_head_id // tl.cdiv(kv_group_num, BLOCK_H) + + VALID_BLOCK_H: tl.constexpr = BLOCK_H if kv_group_num > BLOCK_H else kv_group_num + cur_head = cur_head_id * VALID_BLOCK_H + tl.arange(0, BLOCK_H) + mask_h = (cur_head < (cur_head_id + 1) * VALID_BLOCK_H) & (cur_head < h_q) + + acc, e_max, e_sum = _sparse_mla_compute_tile( + q_buffer, + k_buffer, + indices_ptr, + cur_q, + cur_head, + cur_kv_head_id, + mask_h, + 0, + index_topk, + seq_kv, + stride_q_token, + stride_q_head, + stride_kv_token, + stride_kv_head, + stride_indices_token, + stride_indices_head, + sm_scale, + BLOCK_H, + BLOCK_N, + BLOCK_DV, + BLOCK_DMODEL, + BLOCK_DPE, + ) + + # Guard against queries with zero valid KV (e_sum == 0 → NaN from 0/0). + e_sum_safe = tl.where(e_sum > 0, e_sum, 1.0) + offs_dv = tl.arange(0, BLOCK_DV) + tl.store( + out_ptr + + cur_q * stride_out_token + + cur_head[:, None] * stride_out_head + + offs_dv[None, :], + (acc / e_sum_safe[:, None]).to(tl.bfloat16), + mask=mask_h[:, None], + ) + + +@triton.autotune( + configs=_SPLIT_AUTOTUNE_CONFIGS, + key=["index_topk", "NUM_KV_SPLITS", "kv_group_num"], +) +@triton.jit +def _sparse_mla_kernel_split( + q_buffer, + k_buffer, + indices_ptr, + mid_out_ptr, + seq_kv, + h_q, + stride_q_token, + stride_q_head, + stride_kv_token, + stride_kv_head, + stride_mid_token, + stride_mid_head, + stride_mid_split, + stride_indices_token, + stride_indices_head, + sm_scale, + index_topk: tl.constexpr, + NUM_KV_SPLITS: tl.constexpr, + kv_group_num: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_DV: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_DPE: tl.constexpr, + LOGE2: tl.constexpr, +): + """Stage 1 of split-KV: process one slice of the topk axis and write + its `(out_partial, lse_partial)` into the mid buffer.""" + cur_q = tl.program_id(0) + cur_head_id = tl.program_id(1) + split_kv_id = tl.program_id(2) + cur_kv_head_id = cur_head_id // tl.cdiv(kv_group_num, BLOCK_H) + + VALID_BLOCK_H: tl.constexpr = BLOCK_H if kv_group_num > BLOCK_H else kv_group_num + cur_head = cur_head_id * VALID_BLOCK_H + tl.arange(0, BLOCK_H) + mask_h = (cur_head < (cur_head_id + 1) * VALID_BLOCK_H) & (cur_head < h_q) + + split_topk: tl.constexpr = tl.cdiv(index_topk, NUM_KV_SPLITS) + split_start = split_kv_id * split_topk + split_end = tl.minimum(split_start + split_topk, index_topk) + + acc, e_max, e_sum = _sparse_mla_compute_tile( + q_buffer, + k_buffer, + indices_ptr, + cur_q, + cur_head, + cur_kv_head_id, + mask_h, + split_start, + split_end, + seq_kv, + stride_q_token, + stride_q_head, + stride_kv_token, + stride_kv_head, + stride_indices_token, + stride_indices_head, + sm_scale, + BLOCK_H, + BLOCK_N, + BLOCK_DV, + BLOCK_DMODEL, + BLOCK_DPE, + ) + + # Partial output and natural-log LSE for stage-2 merge. + # When a split has no valid KV (`e_sum == 0`), guard the divide so the + # mid buffer holds 0 instead of NaN; otherwise the `0 * NaN = NaN` term + # in stage 2 would poison every other split. + e_sum_safe = tl.where(e_sum > 0, e_sum, 1.0) + offs_dv = tl.arange(0, BLOCK_DV) + mid_base_2d = ( + mid_out_ptr + + cur_q * stride_mid_token + + cur_head[:, None] * stride_mid_head + + split_kv_id * stride_mid_split + ) + tl.store( + mid_base_2d + offs_dv[None, :], + acc / e_sum_safe[:, None], + mask=mask_h[:, None], + ) + mid_lse_ptr = ( + mid_out_ptr + + cur_q * stride_mid_token + + cur_head * stride_mid_head + + split_kv_id * stride_mid_split + + BLOCK_DV + ) + tl.store(mid_lse_ptr, (e_max + tl.log2(e_sum)) * LOGE2, mask=mask_h) + + +@triton.jit +def _sparse_mla_merge_kernel( + mid_out_ptr, + out_ptr, + h_q, + stride_mid_token, + stride_mid_head, + stride_mid_split, + stride_out_token, + stride_out_head, + NUM_KV_SPLITS: tl.constexpr, + kv_group_num: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_DV: tl.constexpr, + BLOCK_DV_TILE: tl.constexpr, +): + """Stage 2: N-way online-softmax merge of per-split `(out, lse)` tiles. + + Grid is `(num_tokens, num_head_groups, num_dv_tiles)`. Each program handles + `BLOCK_H` heads × `BLOCK_DV_TILE` output-dim lanes. The LSE reduction is + identical across DV tiles for the same (token, head) — each program + recomputes it locally, which is cheap (O(NUM_KV_SPLITS) scalars) and + avoids inter-CTA synchronization. + """ + cur_q = tl.program_id(0) + cur_head_id = tl.program_id(1) + cur_dv_tile = tl.program_id(2) + + VALID_BLOCK_H: tl.constexpr = BLOCK_H if kv_group_num > BLOCK_H else kv_group_num + cur_head = cur_head_id * VALID_BLOCK_H + tl.arange(0, BLOCK_H) + mask_h = (cur_head < (cur_head_id + 1) * VALID_BLOCK_H) & (cur_head < h_q) + + offs_dv = cur_dv_tile * BLOCK_DV_TILE + tl.arange(0, BLOCK_DV_TILE) + mask_dv = offs_dv < BLOCK_DV + # Finite sentinel — same NaN guard as the split kernel for empty splits. + e_max = tl.zeros([BLOCK_H], dtype=tl.float32) - 1.0e30 + e_sum = tl.zeros([BLOCK_H], dtype=tl.float32) + acc = tl.zeros([BLOCK_H, BLOCK_DV_TILE], dtype=tl.float32) + + mid_base_2d = ( + mid_out_ptr + cur_q * stride_mid_token + cur_head[:, None] * stride_mid_head + ) + mid_lse_1d = ( + mid_out_ptr + cur_q * stride_mid_token + cur_head * stride_mid_head + BLOCK_DV + ) + + for split_kv_id in range(NUM_KV_SPLITS): + tv = tl.load( + mid_base_2d + split_kv_id * stride_mid_split + offs_dv[None, :], + mask=mask_h[:, None] & mask_dv[None, :], + other=0.0, + ) + tlogic = tl.load( + mid_lse_1d + split_kv_id * stride_mid_split, + mask=mask_h, + other=-float("inf"), + ) + n_e_max = tl.maximum(tlogic, e_max) + old_scale = tl.exp(e_max - n_e_max) + exp_logic = tl.exp(tlogic - n_e_max) + acc = acc * old_scale[:, None] + exp_logic[:, None] * tv + e_sum = e_sum * old_scale + exp_logic + e_max = n_e_max + + e_sum_safe = tl.where(e_sum > 0, e_sum, 1.0) + tl.store( + out_ptr + + cur_q * stride_out_token + + cur_head[:, None] * stride_out_head + + offs_dv[None, :], + (acc / e_sum_safe[:, None]).to(tl.bfloat16), + mask=mask_h[:, None] & mask_dv[None, :], + ) + + +@functools.lru_cache(maxsize=256) +def _choose_num_kv_splits( + num_tokens: int, num_head_groups: int, index_topk: int, sm_count: int +) -> int: + """Pick a power-of-2 split count that fills the device without dropping + per-split work below _MIN_TOPK_PER_SPLIT. Returns 1 when the single-pass + grid already reaches ~1/_SPLIT_MAX_OCCUPANCY utilization. + """ + baseline = num_tokens * num_head_groups + if baseline == 0 or baseline * _SPLIT_MAX_OCCUPANCY >= sm_count: + return 1 + ideal = triton.next_power_of_2(max(1, index_topk // _MIN_TOPK_PER_SPLIT)) + max_splits = max(1, sm_count // baseline) + max_splits = 1 << (max_splits.bit_length() - 1) # floor to power of 2 + num_kv_splits = min(ideal, max_splits) + while num_kv_splits > 1 and index_topk % num_kv_splits != 0: + num_kv_splits //= 2 + return max(1, num_kv_splits) + + +def triton_mla_sparse_attention( + q: torch.Tensor, + kv: torch.Tensor, + indices: torch.Tensor, + sm_scale: float, + num_kv_splits: int | None = None, + sm_count: int | None = None, +) -> torch.Tensor: + """Sparse MLA attention over topk indices. + + Args: + q: [num_tokens, num_heads_q, dim_qk] bf16 + kv: [seq_kv, num_heads_kv=1, dim_qk] bf16 + indices: [num_tokens, num_heads_kv=1, topk] int32 + sm_scale: softmax scale + num_kv_splits: override auto-heuristic; None/0 = auto, 1 = force single-pass. + sm_count: cached device SM count for the split heuristic. + + Returns: + out: [num_tokens, num_heads_q, _BLOCK_DV] bf16 + """ + num_tokens, num_heads_q, dim_qk = q.shape + assert dim_qk == _DIM_QK, ( + f"sparse MLA kernel requires dim_qk={_DIM_QK} (DeepSeek-V3.2 / GLM-5), " + f"got {dim_qk}" + ) + assert kv.shape[1] == 1 and kv.shape[2] == _DIM_QK + index_topk = indices.shape[2] + assert index_topk % _MIN_BLOCK_N == 0, ( + f"topk ({index_topk}) must be a multiple of the smallest autotune " + f"BLOCK_N ({_MIN_BLOCK_N})" + ) + + kv_group_num = num_heads_q + num_head_groups = triton.cdiv(num_heads_q, min(_BLOCK_H, kv_group_num)) + + if num_kv_splits is None or num_kv_splits == 0: + if sm_count is None: + sm_count = num_compute_units(q.device.index) + num_kv_splits = _choose_num_kv_splits( + num_tokens, num_head_groups, index_topk, sm_count + ) + + out = torch.empty( + (num_tokens, num_heads_q, _BLOCK_DV), + dtype=torch.bfloat16, + device=q.device, + ) + + if num_kv_splits == 1: + _sparse_mla_kernel_final[(num_tokens, num_head_groups)]( + q_buffer=q, + k_buffer=kv, + indices_ptr=indices, + out_ptr=out, + seq_kv=kv.shape[0], + h_q=num_heads_q, + stride_q_token=q.stride(0), + stride_q_head=q.stride(1), + stride_kv_token=kv.stride(0), + stride_kv_head=kv.stride(1), + stride_out_token=out.stride(0), + stride_out_head=out.stride(1), + stride_indices_token=indices.stride(0), + stride_indices_head=indices.stride(1), + sm_scale=sm_scale * LOG2E, + index_topk=index_topk, + kv_group_num=kv_group_num, + BLOCK_H=_BLOCK_H, + BLOCK_DV=_BLOCK_DV, + BLOCK_DMODEL=_BLOCK_DMODEL, + BLOCK_DPE=_BLOCK_DPE, + ) + return out + + # Split-KV: partial fp32 output + LSE per (token, head, split). + mid_out = torch.empty( + (num_tokens, num_heads_q, num_kv_splits, _BLOCK_DV + 1), + dtype=torch.float32, + device=q.device, + ) + _sparse_mla_kernel_split[(num_tokens, num_head_groups, num_kv_splits)]( + q_buffer=q, + k_buffer=kv, + indices_ptr=indices, + mid_out_ptr=mid_out, + seq_kv=kv.shape[0], + h_q=num_heads_q, + stride_q_token=q.stride(0), + stride_q_head=q.stride(1), + stride_kv_token=kv.stride(0), + stride_kv_head=kv.stride(1), + stride_mid_token=mid_out.stride(0), + stride_mid_head=mid_out.stride(1), + stride_mid_split=mid_out.stride(2), + stride_indices_token=indices.stride(0), + stride_indices_head=indices.stride(1), + sm_scale=sm_scale * LOG2E, + index_topk=index_topk, + NUM_KV_SPLITS=num_kv_splits, + kv_group_num=kv_group_num, + BLOCK_H=_BLOCK_H, + BLOCK_DV=_BLOCK_DV, + BLOCK_DMODEL=_BLOCK_DMODEL, + BLOCK_DPE=_BLOCK_DPE, + LOGE2=LOGE2, + ) + + _sparse_mla_merge_kernel[(num_tokens, num_heads_q, _NUM_MERGE_DV_TILES)]( + mid_out_ptr=mid_out, + out_ptr=out, + h_q=num_heads_q, + stride_mid_token=mid_out.stride(0), + stride_mid_head=mid_out.stride(1), + stride_mid_split=mid_out.stride(2), + stride_out_token=out.stride(0), + stride_out_head=out.stride(1), + NUM_KV_SPLITS=num_kv_splits, + kv_group_num=kv_group_num, + BLOCK_H=_MERGE_BLOCK_H, + BLOCK_DV=_BLOCK_DV, + BLOCK_DV_TILE=_MERGE_BLOCK_DV_TILE, + num_warps=2, + ) + return out