From fbd9b57c9b30137092809bb12028e8b873b89de9 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Fri, 30 Jan 2026 13:08:00 -0500 Subject: [PATCH 01/32] Initial implementation Signed-off-by: Matthew Bonanni --- docs/design/attention_backends.md | 2 ++ vllm/platforms/cuda.py | 27 +++++++++++++++++++++++++- vllm/v1/attention/backends/registry.py | 4 ++++ 3 files changed, 32 insertions(+), 1 deletion(-) diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index ef96d002d7fc..23ca28b5c84a 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -128,6 +128,7 @@ Priority is **1 = highest** (tried first). | 4 | `FLASHMLA` | | 5 | `TRITON_MLA` | | 6 | `FLASHMLA_SPARSE` | +| 7 | `FLASHINFER_MLA_SPARSE` | **Ampere/Hopper (SM 8.x-9.x):** @@ -203,6 +204,7 @@ configuration. |---------|--------|-----------|-------------|------------|------|--------|-----------|-----------------|--------------| | `CUTLASS_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | Decoder | 10.x | | `FLASHINFER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | Decoder | 10.x | +| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | 576 | ❌ | ✅ | ❌ | Decoder | 10.x | | `FLASHMLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | Decoder | 9.x-10.x | | `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | Decoder | 9.x-10.x | | `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | Decoder | 9.x | diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index dbbd0876e0f0..12f21ba8fae4 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -56,6 +56,7 @@ def _get_backend_priorities( AttentionBackendEnum.FLASHMLA, AttentionBackendEnum.TRITON_MLA, AttentionBackendEnum.FLASHMLA_SPARSE, + AttentionBackendEnum.FLASHINFER_MLA_SPARSE, ] else: return [ @@ -179,6 +180,7 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None: use_flashmla = False use_cutlass_mla = False use_flashinfer_mla = False + use_flashinfer_mla_sparse = False from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported @@ -187,6 +189,16 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None: hf_text_config = model_config.hf_text_config qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1) if ( + cls.is_device_capability_family(100) + and use_sparse + and qk_nope_head_dim == 128 + ): + # Blackwell + sparse => Use FlashInfer MLA Sparse + use_flashinfer_mla_sparse = True + vllm_config.attention_config.backend = ( + AttentionBackendEnum.FLASHINFER_MLA_SPARSE + ) + elif ( cls.is_device_capability_family(100) and not use_sparse and qk_nope_head_dim == 128 @@ -214,6 +226,9 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None: use_flashmla = backend == AttentionBackendEnum.FLASHMLA use_cutlass_mla = backend == AttentionBackendEnum.CUTLASS_MLA use_flashinfer_mla = backend == AttentionBackendEnum.FLASHINFER_MLA + use_flashinfer_mla_sparse = ( + backend == AttentionBackendEnum.FLASHINFER_MLA_SPARSE + ) if ( use_flashmla @@ -239,8 +254,18 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None: "Forcing kv cache block size to 64 for FlashInferMLA backend." ) + if use_flashinfer_mla_sparse and cache_config.block_size != 64: + cache_config.block_size = 64 + logger.info( + "Forcing kv cache block size to 64 for FlashInferMLASparse backend." + ) + # TODO(Chen): remove this hacky code - if use_sparse and cache_config.block_size != 64: + if ( + use_sparse + and not use_flashinfer_mla_sparse + and cache_config.block_size != 64 + ): cache_config.block_size = 64 logger.info( "Forcing kv cache block size to 64 for FlashMLASparse backend." diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index bd45702fa587..155f5aaef5fa 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -62,6 +62,10 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): FLASHINFER_MLA = ( "vllm.v1.attention.backends.mla.flashinfer_mla.FlashInferMLABackend" ) + FLASHINFER_MLA_SPARSE = ( + "vllm.v1.attention.backends.mla.flashinfer_mla_sparse." + "FlashInferMLASparseBackend" + ) TRITON_MLA = "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend" CUTLASS_MLA = "vllm.v1.attention.backends.mla.cutlass_mla.CutlassMLABackend" FLASHMLA = "vllm.v1.attention.backends.mla.flashmla.FlashMLABackend" From dee20921c7cc68c431bd7f26fc5496584a512c79 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Fri, 30 Jan 2026 13:09:05 -0500 Subject: [PATCH 02/32] Add implementation Signed-off-by: Matthew Bonanni --- .../backends/mla/flashinfer_mla_sparse.py | 543 ++++++++++++++++++ 1 file changed, 543 insertions(+) create mode 100644 vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py new file mode 100644 index 000000000000..5e6866e3618d --- /dev/null +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -0,0 +1,543 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""FlashInfer MLA Sparse Attention Backend. + +This backend uses the FlashInfer TRT-LLM MLA kernel with sparse_mla_top_k +for models like DeepSeek-V3.2 that use index-based sparse attention. + +For sparse MLA: +- block_tables shape changes from [batch_size, max_num_blocks] (dense) + to [batch_size, q_len_per_request, sparse_mla_top_k] (sparse) +- The sparse indices represent physical cache slot positions to attend to +- sparse_mla_top_k parameter must be set to the topk value +""" + +from dataclasses import dataclass +from typing import TYPE_CHECKING, ClassVar + +import numpy as np +import torch +from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla + +from vllm import _custom_ops as ops +from vllm.config import VllmConfig +from vllm.config.cache import CacheDType +from vllm.logger import init_logger +from vllm.model_executor.layers.attention.mla_attention import ( + MLACommonBaseImpl, + get_mla_dims, +) +from vllm.platforms.interface import DeviceCapability +from vllm.v1.attention.backend import ( + AttentionBackend, + AttentionCGSupport, + AttentionLayer, + AttentionMetadata, + AttentionMetadataBuilder, + AttentionType, + CommonAttentionMetadata, + MultipleOf, +) +from vllm.v1.attention.backends.mla.flashmla_sparse import ( + triton_convert_req_index_to_global_index, +) +from vllm.v1.attention.backends.utils import ( + KVCacheLayoutType, + split_decodes_and_prefills, +) +from vllm.v1.kv_cache_interface import AttentionSpec + +if TYPE_CHECKING: + from vllm.model_executor.models.deepseek_v2 import Indexer + +logger = init_logger(__name__) + +FLASHINFER_MLA_SPARSE_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024 + + +class FlashInferMLASparseBackend(AttentionBackend): + """FlashInfer MLA backend with sparse attention support. + + This backend uses the FlashInfer TRT-LLM MLA kernel with sparse_mla_top_k + for models like DeepSeek-V3.2 that use index-based sparse attention. + """ + + accept_output_buffer: bool = True + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ + "auto", + "bfloat16", + "fp8", + "fp8_e4m3", + ] + + @staticmethod + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: + # FlashInfer MLA sparse requires block size 64 + return [64] + + @staticmethod + def get_name() -> str: + return "FLASHINFER_MLA_SPARSE" + + @staticmethod + def get_impl_cls() -> type["FlashInferMLASparseImpl"]: + return FlashInferMLASparseImpl + + @staticmethod + def get_builder_cls() -> type["FlashInferMLASparseMetadataBuilder"]: + return FlashInferMLASparseMetadataBuilder + + @classmethod + def get_supported_head_sizes(cls) -> list[int]: + return [576] + + @classmethod + def is_mla(cls) -> bool: + return True + + @classmethod + def is_sparse(cls) -> bool: + return True + + @classmethod + def supports_compute_capability(cls, capability: DeviceCapability) -> bool: + # FlashInfer sparse MLA targets Blackwell (SM 10.x) + return capability.major == 10 + + @classmethod + def supports_combination( + cls, + head_size: int, + dtype: torch.dtype, + kv_cache_dtype: CacheDType | None, + block_size: int, + use_mla: bool, + has_sink: bool, + use_sparse: bool, + device_capability: DeviceCapability, + ) -> str | None: + # FlashInfer MLA sparse kernel requires qk_nope_head_dim == 128 + from vllm.config import get_current_vllm_config + + vllm_config = get_current_vllm_config() + if vllm_config.model_config is not None: + hf_text_config = vllm_config.model_config.hf_text_config + qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1) + if qk_nope_head_dim != 128: + return ( + f"FlashInfer MLA Sparse kernel requires qk_nope_head_dim == 128, " + f"but got {qk_nope_head_dim}" + ) + # Check for index_topk which indicates sparse model + if not hasattr(hf_text_config, "index_topk"): + return "FlashInfer MLA Sparse requires model with index_topk config" + return None + + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, # assumed to be 1 for MLA + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: + return (num_blocks, block_size, head_size) + + @classmethod + def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None": + return "HND" + + +@dataclass +class FlashInferMLASparseDecodeMetadata: + """Decode-specific metadata for FlashInfer MLA Sparse.""" + + block_table: torch.Tensor + seq_lens: torch.Tensor + + +@dataclass +class FlashInferMLASparseMetadata(AttentionMetadata): + """Attention metadata for FlashInfer MLA Sparse backend.""" + + num_reqs: int + max_query_len: int + max_seq_len: int + num_actual_tokens: int + + # Query start locations + query_start_loc: torch.Tensor + slot_mapping: torch.Tensor + block_table: torch.Tensor + req_id_per_token: torch.Tensor + + # Sparse-specific + block_size: int = 64 + topk_tokens: int = 2048 + + # Decode metadata (None during prefill-only) + decode: FlashInferMLASparseDecodeMetadata | None = None + + # For split batches + num_decodes: int = 0 + num_prefills: int = 0 + num_decode_tokens: int = 0 + num_prefill_tokens: int = 0 + + +class FlashInferMLASparseMetadataBuilder( + AttentionMetadataBuilder[FlashInferMLASparseMetadata] +): + """Builder for FlashInfer MLA Sparse attention metadata.""" + + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH + + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ) -> None: + self.vllm_config = vllm_config + self.layer_names = layer_names + self.kv_cache_spec = kv_cache_spec + self.model_config = vllm_config.model_config + self.device = device + + # Treat requests with query length <= 1 as decodes + self._init_reorder_batch_threshold(1, supports_spec_as_decode=True) + + self.mla_dims = get_mla_dims(self.model_config) + self.topk_tokens = vllm_config.model_config.hf_config.index_topk + + self.req_id_per_token_buffer = torch.empty( + (vllm_config.scheduler_config.max_num_batched_tokens,), + dtype=torch.int32, + device=device, + ) + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> FlashInferMLASparseMetadata: + cm = common_attn_metadata + num_tokens = cm.num_actual_tokens + + # Build req_id_per_token mapping + starts = np.asarray(cm.query_start_loc_cpu, dtype=np.int32) + seg_lengths = np.diff(starts) + req_id_per_token = np.repeat( + np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths + ) + + # Zero-fill for cudagraphs + self.req_id_per_token_buffer.fill_(0) + self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_( + torch.from_numpy(req_id_per_token), non_blocking=True + ) + req_id_per_token_tensor = self.req_id_per_token_buffer[:num_tokens] + + # Split into decode and prefill + (num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens) = ( + split_decodes_and_prefills( + cm, + decode_threshold=self.reorder_batch_threshold or 1, + require_uniform=True, + ) + ) + + # Build decode metadata if we have decode tokens + decode_metadata = None + if num_decodes > 0: + decode_block_table = cm.block_table_tensor[:num_decodes] + decode_seq_lens = cm.seq_lens[:num_decodes] + decode_metadata = FlashInferMLASparseDecodeMetadata( + block_table=decode_block_table, + seq_lens=decode_seq_lens, + ) + + return FlashInferMLASparseMetadata( + num_reqs=cm.num_reqs, + max_query_len=cm.max_query_len, + max_seq_len=cm.max_seq_len, + num_actual_tokens=cm.num_actual_tokens, + query_start_loc=cm.query_start_loc, + slot_mapping=cm.slot_mapping, + block_table=cm.block_table_tensor, + req_id_per_token=req_id_per_token_tensor, + block_size=self.kv_cache_spec.block_size, + topk_tokens=self.topk_tokens, + decode=decode_metadata, + num_decodes=num_decodes, + num_prefills=num_prefills, + num_decode_tokens=num_decode_tokens, + num_prefill_tokens=num_prefill_tokens, + ) + + +# Global workspace buffer (lazily initialized) +_fi_sparse_workspace: torch.Tensor | None = None + + +def _get_workspace_buffer(device: torch.device) -> torch.Tensor: + global _fi_sparse_workspace + if _fi_sparse_workspace is None: + _fi_sparse_workspace = torch.zeros( + FLASHINFER_MLA_SPARSE_WORKSPACE_BUFFER_SIZE, + dtype=torch.uint8, + device=device, + ) + return _fi_sparse_workspace + + +class FlashInferMLASparseImpl(MLACommonBaseImpl[FlashInferMLASparseMetadata]): + """FlashInfer MLA Sparse implementation. + + Uses the TRT-LLM MLA kernel with sparse_mla_top_k parameter for + sparse attention computation. + """ + + def __init__( + self, + num_heads: int, + head_size: int, + scale: float, + num_kv_heads: int, + alibi_slopes: list[float] | None, + sliding_window: int | None, + kv_cache_dtype: str, + logits_soft_cap: float | None, + attn_type: str, + kv_sharing_target_layer_name: str | None, + # MLA Specific Arguments + topk_indice_buffer: torch.Tensor | None = None, + indexer: "Indexer | None" = None, + **mla_args, + ) -> None: + super().__init__( + num_heads, + head_size, + scale, + num_kv_heads, + alibi_slopes, + sliding_window, + kv_cache_dtype, + logits_soft_cap, + attn_type, + kv_sharing_target_layer_name, + indexer=indexer, + **mla_args, + ) + + unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap] + if any(unsupported_features): + raise NotImplementedError( + "FlashInferMLASparseImpl does not support one of the following: " + "alibi_slopes, sliding_window, logits_soft_cap" + ) + + if attn_type != AttentionType.DECODER: + raise NotImplementedError( + "Encoder self-attention and " + "encoder/decoder cross-attention " + "are not implemented for " + "FlashInferMLASparseImpl" + ) + + assert indexer is not None, "Indexer required for sparse MLA" + self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer + + self._workspace_buffer: torch.Tensor | None = None + self.bmm1_scale: float | None = None + self.bmm2_scale: float | None = None + + def _forward_decode( + self, + q: torch.Tensor, + kv_c_and_k_pe_cache: torch.Tensor, + topk_indices: torch.Tensor, + attn_metadata: FlashInferMLASparseMetadata, + layer: AttentionLayer, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Forward pass for decode with sparse attention. + + For sparse mode: + - topk_indices contains logical indices per token [num_tokens, topk] + - These are converted to physical cache slots using the block table + - The resulting indices are passed as block_tables with shape + [batch_size, q_len_per_request, sparse_mla_top_k] + """ + assert kv_c_and_k_pe_cache.numel() > 0 + assert attn_metadata.decode is not None + + num_decodes = attn_metadata.num_decodes + num_decode_tokens = attn_metadata.num_decode_tokens + + # Convert per-request topk indices to global cache slots + # topk_indices: (num_decode_tokens, topk) -> physical cache slots + topk_indices_physical = triton_convert_req_index_to_global_index( + attn_metadata.req_id_per_token[:num_decode_tokens], + attn_metadata.block_table, + topk_indices, + BLOCK_SIZE=attn_metadata.block_size, + NUM_TOPK_TOKENS=topk_indices.shape[1], + ) + + # Reshape q for batch decode: (num_decode_tokens, num_heads, head_dim) + # -> (num_decodes, q_len, num_heads, head_dim) + q_len = num_decode_tokens // num_decodes + q = q.view(num_decodes, q_len, q.shape[-2], q.shape[-1]) + + # Reshape topk_indices: (num_decode_tokens, topk) + # -> (num_decodes, q_len, topk) + topk_indices_physical = topk_indices_physical.view(num_decodes, q_len, -1) + + if self._workspace_buffer is None: + self._workspace_buffer = _get_workspace_buffer(q.device) + + if self.bmm1_scale is None: + self.bmm1_scale = layer._q_scale_float * layer._k_scale_float * self.scale + if self.bmm2_scale is None: + self.bmm2_scale = layer._v_scale_float + + # Call FlashInfer with sparse_mla_top_k + # For sparse mode: + # - block_tables shape is [batch_size, q_len_per_request, sparse_mla_top_k] + # - This contains the physical cache slot indices to attend to + o = trtllm_batch_decode_with_kv_cache_mla( + query=q, + kv_cache=kv_c_and_k_pe_cache.unsqueeze(1), + workspace_buffer=self._workspace_buffer, + qk_nope_head_dim=self.qk_nope_head_dim, + kv_lora_rank=self.kv_lora_rank, + qk_rope_head_dim=self.qk_rope_head_dim, + block_tables=topk_indices_physical, # Sparse indices as block table + seq_lens=attn_metadata.decode.seq_lens, + max_seq_len=attn_metadata.max_seq_len, + bmm1_scale=self.bmm1_scale, + bmm2_scale=self.bmm2_scale, + sparse_mla_top_k=attn_metadata.topk_tokens, + ) + + # Reshape output: (num_decodes, q_len, num_heads, head_dim_v) + # -> (num_decode_tokens, num_heads, head_dim_v) + o = o.view(-1, o.shape[-2], o.shape[-1]) + + return o, None + + def forward( + self, + layer: AttentionLayer, + q: torch.Tensor, + k_c_normed: torch.Tensor, + k_pe: torch.Tensor, + kv_cache: torch.Tensor, + attn_metadata: FlashInferMLASparseMetadata | None, + output: torch.Tensor | None = None, + output_scale: torch.Tensor | None = None, + output_block_scale: torch.Tensor | None = None, + ) -> torch.Tensor: + """Forward pass for FlashInfer MLA Sparse attention.""" + assert output is not None, "Output tensor must be provided." + + if output_scale is not None or output_block_scale is not None: + raise NotImplementedError( + "Fused output quantization is not yet supported for " + "FlashInferMLASparseImpl" + ) + + if attn_metadata is None: + # Dummy run + return output.fill_(0) + + num_actual_toks = attn_metadata.num_actual_tokens + + # Slice inputs to actual token count (may be padded for CUDA graphs) + q = q[:num_actual_toks, ...] + k_c_normed = k_c_normed[:num_actual_toks, ...] + k_pe = k_pe[:num_actual_toks, ...] + + assert self.topk_indices_buffer is not None + topk_indices = self.topk_indices_buffer[:num_actual_toks] + + # Compute ql_nope = q_nope @ W_UK^T for decode-style MLA + q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + # (B, N, P) -> (N, B, P) + q_nope = q_nope.transpose(0, 1) + # (N, B, P) x (N, P, L) -> (N, B, L) + ql_nope = torch.bmm(q_nope, self.W_UK_T) + # (N, B, L) -> (B, N, L) + ql_nope = ql_nope.transpose(0, 1) + + q = torch.cat([ql_nope, q_pe], dim=-1) + + # Write to KV cache + if kv_cache.numel() > 0: + ops.concat_and_cache_mla( + k_c_normed, + k_pe.squeeze(1), + kv_cache, + attn_metadata.slot_mapping.flatten(), + kv_cache_dtype=self.kv_cache_dtype, + scale=layer._k_scale, + ) + + num_decode_tokens = attn_metadata.num_decode_tokens + num_prefill_tokens = attn_metadata.num_prefill_tokens + + # FlashInfer sparse MLA only supports decode + # Prefill would need to fall back to dense attention + if num_decode_tokens > 0 and num_prefill_tokens == 0: + # Pure decode batch + attn_out, _ = self._forward_decode( + q, + kv_cache, + topk_indices, + attn_metadata, + layer, + ) + elif num_decode_tokens > 0: + # Mixed batch: handle decode portion + attn_out = q.new_empty( + (num_actual_toks, self.num_heads, self.kv_lora_rank), + dtype=q.dtype, + device=q.device, + ) + + decode_out, _ = self._forward_decode( + q[:num_decode_tokens], + kv_cache, + topk_indices[:num_decode_tokens], + attn_metadata, + layer, + ) + attn_out[:num_decode_tokens] = decode_out + + # For prefill tokens, sparse attention is not supported + # This should not happen in normal operation since the scheduler + # should use a dense backend for prefill + if num_prefill_tokens > 0: + logger.warning_once( + "FlashInfer MLA Sparse does not support prefill. " + "Prefill tokens will use zero attention output." + ) + attn_out[num_decode_tokens:] = 0 + else: + # Pure prefill - not supported for sparse + logger.warning_once( + "FlashInfer MLA Sparse does not support prefill. " + "Using zero attention output." + ) + attn_out = q.new_zeros( + (num_actual_toks, self.num_heads, self.kv_lora_rank), + dtype=q.dtype, + device=q.device, + ) + + # V up-projection + self._v_up_proj(attn_out, out=output[:num_actual_toks]) + return output From aaca43b54826dbf643365cedb36796d0e1758c60 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Fri, 30 Jan 2026 20:21:19 +0000 Subject: [PATCH 03/32] Update test Signed-off-by: Matthew Bonanni --- .../v1/attention/test_sparse_mla_backends.py | 72 +++++++++++-------- 1 file changed, 41 insertions(+), 31 deletions(-) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 8d1f5cc46ba9..92d5e4450305 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for the FlashMLA sparse backend utilities.""" +"""Unit tests for the sparse MLA backends and utilities.""" import math from types import MethodType, SimpleNamespace @@ -25,6 +25,9 @@ from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.platforms import current_platform from vllm.utils.math_utils import cdiv +from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( + FlashInferMLASparseBackend, +) from vllm.v1.attention.backends.mla.flashmla_sparse import ( FlashMLASparseBackend, triton_convert_req_index_to_global_index, @@ -156,31 +159,42 @@ def _quantize_dequantize_fp8_ds_mla( return dequant_kv_c, dequant_k_pe +@pytest.mark.parametrize( + "backend_cls", + [FlashMLASparseBackend, FlashInferMLASparseBackend], + ids=["FlashMLA", "FlashInfer"], +) @pytest.mark.parametrize("batch_name", list(SPARSE_BACKEND_BATCH_SPECS.keys())) -@pytest.mark.parametrize("kv_cache_dtype", ["fp8_ds_mla", "auto"]) +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_ds_mla"]) @pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4]) -@pytest.mark.skipif( - torch.cuda.get_device_capability() < (9, 0), - reason="FlashMLASparseBackend requires CUDA 9.0 or higher", -) def test_sparse_backend_decode_correctness( default_vllm_config, dist_init, + backend_cls, batch_name, kv_cache_dtype, tensor_parallel_size, workspace_init, ): if current_platform.is_rocm(): - pytest.skip("ROCm does not support fp8_ds_mla data type for kv cache.") + pytest.skip("ROCm does not support sparse MLA backends.") - if not torch.cuda.is_available(): - pytest.skip("CUDA is required for sparse MLA decode test") + if kv_cache_dtype not in backend_cls.supported_kv_cache_dtypes: + pytest.skip(f"{backend_cls.get_name()} does not support {kv_cache_dtype}") - device = torch.device("cuda") - dtype = torch.bfloat16 + if backend_cls == FlashMLASparseBackend: + ok, reason = flashmla.is_flashmla_sparse_supported() + if not ok: + pytest.skip(reason) + elif backend_cls == FlashInferMLASparseBackend: + if not current_platform.has_device_capability(100): + pytest.skip("FlashInferMLASparseBackend requires SM 10.0 or higher") batch_spec = SPARSE_BACKEND_BATCH_SPECS[batch_name] + use_fp8_ds_mla_quantization = kv_cache_dtype == "fp8_ds_mla" + + device = torch.device("cuda") + dtype = torch.bfloat16 # Model hyper-parameters (kept intentionally small for the unit test) total_num_heads = 128 @@ -268,22 +282,22 @@ def test_sparse_backend_decode_correctness( kv_c_full = torch.rand(s_len, kv_lora_rank, dtype=dtype, device=device) k_pe_full = torch.rand(s_len, 1, qk_rope_head_dim, dtype=dtype, device=device) - # SM100 (Blackwell) uses float -> e8m0 -> bf16 scale conversion - # which truncates scales to powers of 2. Simulate this in reference. - is_sm100 = torch.cuda.get_device_capability()[0] >= 10 - kv_c_full, k_pe_full = _quantize_dequantize_fp8_ds_mla( - kv_c_full, - k_pe_full.squeeze(1), - block_size=vllm_config.cache_config.block_size, - scale=kv_cache_scale, - simulate_sm100_e8m0_scales=is_sm100, - ) + if use_fp8_ds_mla_quantization: + is_sm100 = torch.cuda.get_device_capability()[0] >= 10 + kv_c_full, k_pe_squeezed = _quantize_dequantize_fp8_ds_mla( + kv_c_full, + k_pe_full.squeeze(1), + block_size=block_size, + scale=kv_cache_scale, + simulate_sm100_e8m0_scales=is_sm100, + ) + k_pe_full = k_pe_squeezed.unsqueeze(1) q_nope, q_pe = q_c.split([qk_nope_head_dim, qk_rope_head_dim], dim=-1) ql_nope = torch.einsum("qnh,lnh->qnl", q_nope, W_UK) q_mqa = torch.cat([ql_nope, q_pe], dim=-1) - k_mqa = torch.cat([kv_c_full, k_pe_full], dim=-1) + k_mqa = torch.cat([kv_c_full, k_pe_full.squeeze(1)], dim=-1) k_mqa = k_mqa.unsqueeze(1).expand(-1, num_heads, -1) v_mqa = kv_c_full.unsqueeze(1).expand(-1, num_heads, -1) @@ -334,11 +348,11 @@ def test_sparse_backend_decode_correctness( num_blocks=vllm_config.cache_config.num_gpu_blocks, common_attn_metadata=common_attn_metadata, randomize_blocks=False, - kv_cache_dtype=vllm_config.cache_config.cache_dtype, + kv_cache_dtype=kv_cache_dtype if use_fp8_ds_mla_quantization else "auto", scale=kv_cache_scale, ) - builder_cls = FlashMLASparseBackend.get_builder_cls() + builder_cls = backend_cls.get_builder_cls() builder = builder_cls(kv_cache_spec, ["placeholder"], vllm_config, device) metadata = builder.build( common_prefix_len=0, common_attn_metadata=common_attn_metadata @@ -362,15 +376,11 @@ def test_sparse_backend_decode_correctness( causal_mask, debug_indices, torch.full_like(debug_indices, -1) ) - # FlashMLASparseImpl now reads top-k indices from the indexer-provided + # Sparse backends read top-k indices from the indexer-provided # buffer, so emulate that contract with a simple namespace mock. debug_indices = debug_indices.expand(metadata.num_actual_tokens, -1).clone() mock_indexer = SimpleNamespace(topk_indices_buffer=debug_indices) - ok, reason = flashmla.is_flashmla_sparse_supported() - if not ok: - pytest.skip(reason) - kv_b_proj_weight = torch.cat([W_UK, W_UV], dim=-1) kv_b_proj_weight = kv_b_proj_weight.view( kv_lora_rank, num_heads * (qk_nope_head_dim + v_head_dim) @@ -383,7 +393,7 @@ def test_sparse_backend_decode_correctness( ).to(device=device, dtype=dtype) mock_kv_b_proj.weight = torch.nn.Parameter(kv_b_proj_weight.T.contiguous()) - impl_cls = FlashMLASparseBackend.get_impl_cls() + impl_cls = backend_cls.get_impl_cls() with set_current_vllm_config(vllm_config): impl = impl_cls( num_heads=num_heads, @@ -430,7 +440,7 @@ def test_sparse_backend_decode_correctness( # FP8 quantization introduces some error, but should be within reasonable bounds # BF16 (auto) should be very accurate, FP8 allows slightly more tolerance - if kv_cache_dtype == "fp8_ds_mla": + if kv_cache_dtype.startswith("fp8"): torch.testing.assert_close(backend_output, sdpa_reference, rtol=0.05, atol=0.05) else: torch.testing.assert_close(backend_output, sdpa_reference, rtol=0.01, atol=0.01) From 6f3bbc9aae794145d354638c0c3f8d9873a57719 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Fri, 30 Jan 2026 20:23:13 +0000 Subject: [PATCH 04/32] Remove unnecessary skip Signed-off-by: Matthew Bonanni --- tests/v1/attention/test_sparse_mla_backends.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 92d5e4450305..99da0dbec43e 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -176,9 +176,6 @@ def test_sparse_backend_decode_correctness( tensor_parallel_size, workspace_init, ): - if current_platform.is_rocm(): - pytest.skip("ROCm does not support sparse MLA backends.") - if kv_cache_dtype not in backend_cls.supported_kv_cache_dtypes: pytest.skip(f"{backend_cls.get_name()} does not support {kv_cache_dtype}") From 77d5461b7285fc279eaf0458a8cc091a4131bdf7 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Fri, 30 Jan 2026 23:55:52 +0000 Subject: [PATCH 05/32] Cleanup Signed-off-by: Matthew Bonanni --- .../backends/mla/flashinfer_mla_sparse.py | 59 +++---------------- 1 file changed, 7 insertions(+), 52 deletions(-) diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 5e6866e3618d..30b45350492b 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -486,58 +486,13 @@ def forward( scale=layer._k_scale, ) - num_decode_tokens = attn_metadata.num_decode_tokens - num_prefill_tokens = attn_metadata.num_prefill_tokens - - # FlashInfer sparse MLA only supports decode - # Prefill would need to fall back to dense attention - if num_decode_tokens > 0 and num_prefill_tokens == 0: - # Pure decode batch - attn_out, _ = self._forward_decode( - q, - kv_cache, - topk_indices, - attn_metadata, - layer, - ) - elif num_decode_tokens > 0: - # Mixed batch: handle decode portion - attn_out = q.new_empty( - (num_actual_toks, self.num_heads, self.kv_lora_rank), - dtype=q.dtype, - device=q.device, - ) - - decode_out, _ = self._forward_decode( - q[:num_decode_tokens], - kv_cache, - topk_indices[:num_decode_tokens], - attn_metadata, - layer, - ) - attn_out[:num_decode_tokens] = decode_out - - # For prefill tokens, sparse attention is not supported - # This should not happen in normal operation since the scheduler - # should use a dense backend for prefill - if num_prefill_tokens > 0: - logger.warning_once( - "FlashInfer MLA Sparse does not support prefill. " - "Prefill tokens will use zero attention output." - ) - attn_out[num_decode_tokens:] = 0 - else: - # Pure prefill - not supported for sparse - logger.warning_once( - "FlashInfer MLA Sparse does not support prefill. " - "Using zero attention output." - ) - attn_out = q.new_zeros( - (num_actual_toks, self.num_heads, self.kv_lora_rank), - dtype=q.dtype, - device=q.device, - ) + attn_out, _ = self._forward_decode( + q, + kv_cache, + topk_indices, + attn_metadata, + layer, + ) - # V up-projection self._v_up_proj(attn_out, out=output[:num_actual_toks]) return output From 38e6fc996ff89c6c838151b36c0b50c1783be6e5 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Mon, 2 Feb 2026 21:39:50 +0000 Subject: [PATCH 06/32] Address refactor Signed-off-by: Matthew Bonanni --- .../backends/mla/flashinfer_mla_sparse.py | 71 +------------------ 1 file changed, 1 insertion(+), 70 deletions(-) diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 30b45350492b..5bd9c5138f94 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -19,7 +19,6 @@ import torch from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla -from vllm import _custom_ops as ops from vllm.config import VllmConfig from vllm.config.cache import CacheDType from vllm.logger import init_logger @@ -355,7 +354,7 @@ def __init__( self.bmm1_scale: float | None = None self.bmm2_scale: float | None = None - def _forward_decode( + def forward_mqa( self, q: torch.Tensor, kv_c_and_k_pe_cache: torch.Tensor, @@ -428,71 +427,3 @@ def _forward_decode( o = o.view(-1, o.shape[-2], o.shape[-1]) return o, None - - def forward( - self, - layer: AttentionLayer, - q: torch.Tensor, - k_c_normed: torch.Tensor, - k_pe: torch.Tensor, - kv_cache: torch.Tensor, - attn_metadata: FlashInferMLASparseMetadata | None, - output: torch.Tensor | None = None, - output_scale: torch.Tensor | None = None, - output_block_scale: torch.Tensor | None = None, - ) -> torch.Tensor: - """Forward pass for FlashInfer MLA Sparse attention.""" - assert output is not None, "Output tensor must be provided." - - if output_scale is not None or output_block_scale is not None: - raise NotImplementedError( - "Fused output quantization is not yet supported for " - "FlashInferMLASparseImpl" - ) - - if attn_metadata is None: - # Dummy run - return output.fill_(0) - - num_actual_toks = attn_metadata.num_actual_tokens - - # Slice inputs to actual token count (may be padded for CUDA graphs) - q = q[:num_actual_toks, ...] - k_c_normed = k_c_normed[:num_actual_toks, ...] - k_pe = k_pe[:num_actual_toks, ...] - - assert self.topk_indices_buffer is not None - topk_indices = self.topk_indices_buffer[:num_actual_toks] - - # Compute ql_nope = q_nope @ W_UK^T for decode-style MLA - q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) - # (B, N, P) -> (N, B, P) - q_nope = q_nope.transpose(0, 1) - # (N, B, P) x (N, P, L) -> (N, B, L) - ql_nope = torch.bmm(q_nope, self.W_UK_T) - # (N, B, L) -> (B, N, L) - ql_nope = ql_nope.transpose(0, 1) - - q = torch.cat([ql_nope, q_pe], dim=-1) - - # Write to KV cache - if kv_cache.numel() > 0: - ops.concat_and_cache_mla( - k_c_normed, - k_pe.squeeze(1), - kv_cache, - attn_metadata.slot_mapping.flatten(), - kv_cache_dtype=self.kv_cache_dtype, - scale=layer._k_scale, - ) - - attn_out, _ = self._forward_decode( - q, - kv_cache, - topk_indices, - attn_metadata, - layer, - ) - - self._v_up_proj(attn_out, out=output[:num_actual_toks]) - return output From ead72ab648a5a18bfe748d5a6ee1ae6e8b95d7e9 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Mon, 2 Feb 2026 21:45:38 +0000 Subject: [PATCH 07/32] Move kernel to sparse_utils Signed-off-by: Matthew Bonanni --- .../backends/mla/flashinfer_mla_sparse.py | 2 +- .../attention/backends/mla/flashmla_sparse.py | 164 +---------------- .../v1/attention/backends/mla/sparse_utils.py | 167 ++++++++++++++++++ 3 files changed, 171 insertions(+), 162 deletions(-) create mode 100644 vllm/v1/attention/backends/mla/sparse_utils.py diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 5bd9c5138f94..3eac1be136b7 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -37,7 +37,7 @@ CommonAttentionMetadata, MultipleOf, ) -from vllm.v1.attention.backends.mla.flashmla_sparse import ( +from vllm.v1.attention.backends.mla.sparse_utils import ( triton_convert_req_index_to_global_index, ) from vllm.v1.attention.backends.utils import ( diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index 80e402a4d4eb..799c77d73ad2 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -15,7 +15,6 @@ ) from vllm.platforms import current_platform from vllm.platforms.interface import DeviceCapability -from vllm.triton_utils import tl, triton from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -26,6 +25,9 @@ MultipleOf, SparseMLAAttentionImpl, ) +from vllm.v1.attention.backends.mla.sparse_utils import ( + triton_convert_req_index_to_global_index, +) from vllm.v1.attention.backends.utils import ( reshape_attn_output_for_spec_decode, reshape_query_for_spec_decode, @@ -203,166 +205,6 @@ class Chunk: fp8_use_mixed_batch: bool = False -# Kernel with prefill workspace support -@triton.jit -def _convert_req_index_to_global_index_kernel( - req_id_ptr, # int32 [num_tokens] - block_table_ptr, # int32 [num_requests, max_num_blocks_per_req] - token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS] - out_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS] - prefill_request_id_ptr, # int32 [num_tokens], -1 for decode, >=0 for prefill - workspace_starts_ptr, # int32 [num_prefill_reqs+1] or nullptr - # shapes (compile-time where possible) - max_num_blocks_per_req: tl.constexpr, - BLOCK_SIZE: tl.constexpr, - BLOCK_N: tl.constexpr, # tile width along columns - HAS_PREFILL: tl.constexpr, - # strides (in elements) - bt_stride0, - bt_stride1, - ti_stride0, - ti_stride1, - out_stride0, - out_stride1, -): - # program_id(0) -> token_id (row) - # program_id(1) -> tile index along columns - token_id = tl.program_id(0) - tile_id = tl.program_id(1) - - # Each program covers BLOCK_N consecutive columns - indice_id = tile_id * BLOCK_N + tl.arange(0, BLOCK_N) - - # Load request id for this token (no mask: grid is exact) - req = tl.load(req_id_ptr + token_id) - - # Load token indices for this tile - ti_ptr = token_indices_ptr + token_id * ti_stride0 + indice_id * ti_stride1 - tok = tl.load(ti_ptr) # int32 - - # Only token == -1 should propagate as -1 - is_invalid_tok = tok < 0 - is_prefill = False - if HAS_PREFILL: - prefill_req_id = tl.load(prefill_request_id_ptr + token_id) - is_prefill = prefill_req_id >= 0 - # Compute block id and in-block offset - block_id = tok // BLOCK_SIZE - inblock_off = tok % BLOCK_SIZE - - # Guard block_table access - valid_block = (block_id < max_num_blocks_per_req) & (block_id >= 0) - bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1 - is_invalid_tok |= ~valid_block - base = tl.load(bt_ptr, mask=valid_block & ~is_prefill, other=0) - out_val = base * BLOCK_SIZE + inblock_off - - # Override with prefill output if prefill is enabled - if HAS_PREFILL: - workspace_start = tl.load( - workspace_starts_ptr + prefill_req_id, mask=is_prefill, other=0 - ) - prefill_out = workspace_start + tok - out_val = tl.where(is_prefill, prefill_out, out_val) - out_val = tl.where(is_invalid_tok, -1, out_val) - - # Store results - out_ptr_ij = out_ptr + token_id * out_stride0 + indice_id * out_stride1 - tl.store(out_ptr_ij, out_val) - - -def triton_convert_req_index_to_global_index( - req_id: torch.Tensor, # int32 [num_tokens] - block_table: torch.Tensor, # int32 [num_requests, max_num_blocks_per_req] - token_indices: torch.Tensor, # int32 [num_tokens, NUM_TOPK_TOKENS] - BLOCK_SIZE: int = 64, - NUM_TOPK_TOKENS: int = 2048, - BLOCK_N: int = 128, # tile width along columns - HAS_PREFILL_WORKSPACE: bool = False, - prefill_workspace_request_ids: torch.Tensor | None = None, - prefill_workspace_starts: torch.Tensor | None = None, -): - """ - out[token_id, indice_id] = - block_table[req_id[token_id], - token_indices[token_id, indice_id] // BLOCK_SIZE] * BLOCK_SIZE - + token_indices[token_id, indice_id] % BLOCK_SIZE - - Only when token_indices[token_id, indice_id] == -1 do we output -1. - For safety, we also output -1 if the derived block_id would be - out-of-bounds. - - When HAS_PREFILL_WORKSPACE is True, prefill tokens are mapped to workspace offsets - instead of global cache slots. prefill_workspace_request_ids and - prefill_workspace_starts must be provided. - - prefill_workspace_request_ids: int32 [num_tokens], -1 for decode else - prefill request index (maps to prefill_workspace_starts) - prefill_workspace_starts: int32 [num_prefills], 0-indexed workspace - starts for each prefill request - """ - assert req_id.dtype == torch.int32 - assert block_table.dtype == torch.int32 - assert token_indices.dtype == torch.int32 - assert token_indices.shape[1] == NUM_TOPK_TOKENS - assert NUM_TOPK_TOKENS % BLOCK_N == 0, ( - f"NUM_TOPK_TOKENS ({NUM_TOPK_TOKENS}) must be divisible by BLOCK_N ({BLOCK_N})" - ) - - if HAS_PREFILL_WORKSPACE: - assert prefill_workspace_request_ids is not None - assert prefill_workspace_starts is not None - assert prefill_workspace_request_ids.dtype == torch.int32 - assert prefill_workspace_starts.dtype == torch.int32 - - num_tokens = req_id.shape[0] - max_num_blocks_per_req = block_table.shape[1] - tiles_per_row = NUM_TOPK_TOKENS // BLOCK_N - - # Ensure contiguous tensors on the same device - req_id_c = req_id.contiguous() - block_table_c = block_table.contiguous() - token_indices_c = token_indices.contiguous() - out = torch.empty_like(token_indices_c) - - # Strides in elements - bt_stride0, bt_stride1 = block_table_c.stride() - ti_stride0, ti_stride1 = token_indices_c.stride() - out_stride0, out_stride1 = out.stride() - - # Prepare prefill pointers - if HAS_PREFILL_WORKSPACE: - assert prefill_workspace_request_ids is not None # for mypy - assert prefill_workspace_starts is not None # for mypy - assert prefill_workspace_request_ids.is_contiguous() - assert prefill_workspace_starts.is_contiguous() - - # Exact 2D grid: tokens × column tiles - grid = (num_tokens, tiles_per_row) - - _convert_req_index_to_global_index_kernel[grid]( - req_id_c, - block_table_c, - token_indices_c, - out, - prefill_workspace_request_ids, - prefill_workspace_starts, - # shapes / constexprs - max_num_blocks_per_req, - BLOCK_SIZE, - BLOCK_N, - HAS_PREFILL_WORKSPACE, - # strides - bt_stride0, - bt_stride1, - ti_stride0, - ti_stride1, - out_stride0, - out_stride1, - ) - return out - - def get_prefill_workspace_size(max_model_len: int): # NOTE(Lucas): 5 is a magic number for controlling the prefill buffer size. # May be tuned later. diff --git a/vllm/v1/attention/backends/mla/sparse_utils.py b/vllm/v1/attention/backends/mla/sparse_utils.py new file mode 100644 index 000000000000..cccc8f7506ac --- /dev/null +++ b/vllm/v1/attention/backends/mla/sparse_utils.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Utility functions for sparse MLA backends.""" + +import torch + +from vllm.triton_utils import tl, triton + + +# Kernel with prefill workspace support +@triton.jit +def _convert_req_index_to_global_index_kernel( + req_id_ptr, # int32 [num_tokens] + block_table_ptr, # int32 [num_requests, max_num_blocks_per_req] + token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS] + out_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS] + prefill_request_id_ptr, # int32 [num_tokens], -1 for decode, >=0 for prefill + workspace_starts_ptr, # int32 [num_prefill_reqs+1] or nullptr + # shapes (compile-time where possible) + max_num_blocks_per_req: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + BLOCK_N: tl.constexpr, # tile width along columns + HAS_PREFILL: tl.constexpr, + # strides (in elements) + bt_stride0, + bt_stride1, + ti_stride0, + ti_stride1, + out_stride0, + out_stride1, +): + # program_id(0) -> token_id (row) + # program_id(1) -> tile index along columns + token_id = tl.program_id(0) + tile_id = tl.program_id(1) + + # Each program covers BLOCK_N consecutive columns + indice_id = tile_id * BLOCK_N + tl.arange(0, BLOCK_N) + + # Load request id for this token (no mask: grid is exact) + req = tl.load(req_id_ptr + token_id) + + # Load token indices for this tile + ti_ptr = token_indices_ptr + token_id * ti_stride0 + indice_id * ti_stride1 + tok = tl.load(ti_ptr) # int32 + + # Only token == -1 should propagate as -1 + is_invalid_tok = tok < 0 + is_prefill = False + if HAS_PREFILL: + prefill_req_id = tl.load(prefill_request_id_ptr + token_id) + is_prefill = prefill_req_id >= 0 + # Compute block id and in-block offset + block_id = tok // BLOCK_SIZE + inblock_off = tok % BLOCK_SIZE + + # Guard block_table access + valid_block = (block_id < max_num_blocks_per_req) & (block_id >= 0) + bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1 + is_invalid_tok |= ~valid_block + base = tl.load(bt_ptr, mask=valid_block & ~is_prefill, other=0) + out_val = base * BLOCK_SIZE + inblock_off + + # Override with prefill output if prefill is enabled + if HAS_PREFILL: + workspace_start = tl.load( + workspace_starts_ptr + prefill_req_id, mask=is_prefill, other=0 + ) + prefill_out = workspace_start + tok + out_val = tl.where(is_prefill, prefill_out, out_val) + out_val = tl.where(is_invalid_tok, -1, out_val) + + # Store results + out_ptr_ij = out_ptr + token_id * out_stride0 + indice_id * out_stride1 + tl.store(out_ptr_ij, out_val) + + +def triton_convert_req_index_to_global_index( + req_id: torch.Tensor, # int32 [num_tokens] + block_table: torch.Tensor, # int32 [num_requests, max_num_blocks_per_req] + token_indices: torch.Tensor, # int32 [num_tokens, NUM_TOPK_TOKENS] + BLOCK_SIZE: int = 64, + NUM_TOPK_TOKENS: int = 2048, + BLOCK_N: int = 128, # tile width along columns + HAS_PREFILL_WORKSPACE: bool = False, + prefill_workspace_request_ids: torch.Tensor | None = None, + prefill_workspace_starts: torch.Tensor | None = None, +): + """ + out[token_id, indice_id] = + block_table[req_id[token_id], + token_indices[token_id, indice_id] // BLOCK_SIZE] * BLOCK_SIZE + + token_indices[token_id, indice_id] % BLOCK_SIZE + + Only when token_indices[token_id, indice_id] == -1 do we output -1. + For safety, we also output -1 if the derived block_id would be + out-of-bounds. + + When HAS_PREFILL_WORKSPACE is True, prefill tokens are mapped to workspace offsets + instead of global cache slots. prefill_workspace_request_ids and + prefill_workspace_starts must be provided. + + prefill_workspace_request_ids: int32 [num_tokens], -1 for decode else + prefill request index (maps to prefill_workspace_starts) + prefill_workspace_starts: int32 [num_prefills], 0-indexed workspace + starts for each prefill request + """ + assert req_id.dtype == torch.int32 + assert block_table.dtype == torch.int32 + assert token_indices.dtype == torch.int32 + assert token_indices.shape[1] == NUM_TOPK_TOKENS + assert NUM_TOPK_TOKENS % BLOCK_N == 0, ( + f"NUM_TOPK_TOKENS ({NUM_TOPK_TOKENS}) must be divisible by BLOCK_N ({BLOCK_N})" + ) + + if HAS_PREFILL_WORKSPACE: + assert prefill_workspace_request_ids is not None + assert prefill_workspace_starts is not None + assert prefill_workspace_request_ids.dtype == torch.int32 + assert prefill_workspace_starts.dtype == torch.int32 + + num_tokens = req_id.shape[0] + max_num_blocks_per_req = block_table.shape[1] + tiles_per_row = NUM_TOPK_TOKENS // BLOCK_N + + # Ensure contiguous tensors on the same device + req_id_c = req_id.contiguous() + block_table_c = block_table.contiguous() + token_indices_c = token_indices.contiguous() + out = torch.empty_like(token_indices_c) + + # Strides in elements + bt_stride0, bt_stride1 = block_table_c.stride() + ti_stride0, ti_stride1 = token_indices_c.stride() + out_stride0, out_stride1 = out.stride() + + # Prepare prefill pointers + if HAS_PREFILL_WORKSPACE: + assert prefill_workspace_request_ids is not None # for mypy + assert prefill_workspace_starts is not None # for mypy + assert prefill_workspace_request_ids.is_contiguous() + assert prefill_workspace_starts.is_contiguous() + + # Exact 2D grid: tokens × column tiles + grid = (num_tokens, tiles_per_row) + + _convert_req_index_to_global_index_kernel[grid]( + req_id_c, + block_table_c, + token_indices_c, + out, + prefill_workspace_request_ids, + prefill_workspace_starts, + # shapes / constexprs + max_num_blocks_per_req, + BLOCK_SIZE, + BLOCK_N, + HAS_PREFILL_WORKSPACE, + # strides + bt_stride0, + bt_stride1, + ti_stride0, + ti_stride1, + out_stride0, + out_stride1, + ) + return out From 49cd2094f7f9579a5a3e994333027764bb682b4b Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Mon, 2 Feb 2026 21:52:20 +0000 Subject: [PATCH 08/32] Fix from refactor Signed-off-by: Matthew Bonanni --- vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 3eac1be136b7..29f477c75e4b 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -23,7 +23,6 @@ from vllm.config.cache import CacheDType from vllm.logger import init_logger from vllm.model_executor.layers.attention.mla_attention import ( - MLACommonBaseImpl, get_mla_dims, ) from vllm.platforms.interface import DeviceCapability @@ -36,6 +35,7 @@ AttentionType, CommonAttentionMetadata, MultipleOf, + SparseMLAAttentionImpl, ) from vllm.v1.attention.backends.mla.sparse_utils import ( triton_convert_req_index_to_global_index, @@ -293,7 +293,7 @@ def _get_workspace_buffer(device: torch.device) -> torch.Tensor: return _fi_sparse_workspace -class FlashInferMLASparseImpl(MLACommonBaseImpl[FlashInferMLASparseMetadata]): +class FlashInferMLASparseImpl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata]): """FlashInfer MLA Sparse implementation. Uses the TRT-LLM MLA kernel with sparse_mla_top_k parameter for From 14044783390475acd6ae5039a25a25d7c6cb085e Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Mon, 2 Feb 2026 21:55:22 +0000 Subject: [PATCH 09/32] Super init call not necessary Signed-off-by: Matthew Bonanni --- .../backends/mla/flashinfer_mla_sparse.py | 15 --------------- 1 file changed, 15 deletions(-) diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 29f477c75e4b..132beeeb22a4 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -317,21 +317,6 @@ def __init__( indexer: "Indexer | None" = None, **mla_args, ) -> None: - super().__init__( - num_heads, - head_size, - scale, - num_kv_heads, - alibi_slopes, - sliding_window, - kv_cache_dtype, - logits_soft_cap, - attn_type, - kv_sharing_target_layer_name, - indexer=indexer, - **mla_args, - ) - unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap] if any(unsupported_features): raise NotImplementedError( From 349da8788d1f0d54afbaa474f5c87b70cc06495b Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Mon, 2 Feb 2026 22:04:03 +0000 Subject: [PATCH 10/32] More fixes Signed-off-by: Matthew Bonanni --- .../backends/mla/flashinfer_mla_sparse.py | 22 +++++++++---------- 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 132beeeb22a4..a5df2a0a8d77 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -341,9 +341,8 @@ def __init__( def forward_mqa( self, - q: torch.Tensor, + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], kv_c_and_k_pe_cache: torch.Tensor, - topk_indices: torch.Tensor, attn_metadata: FlashInferMLASparseMetadata, layer: AttentionLayer, ) -> tuple[torch.Tensor, torch.Tensor | None]: @@ -355,27 +354,26 @@ def forward_mqa( - The resulting indices are passed as block_tables with shape [batch_size, q_len_per_request, sparse_mla_top_k] """ - assert kv_c_and_k_pe_cache.numel() > 0 - assert attn_metadata.decode is not None + # Concatenate q if it's a tuple (ql_nope, q_pe) + if isinstance(q, tuple): + q = torch.cat(q, dim=-1) - num_decodes = attn_metadata.num_decodes - num_decode_tokens = attn_metadata.num_decode_tokens + num_actual_toks = q.shape[0] + + # Get topk indices + assert self.topk_indices_buffer is not None + topk_indices = self.topk_indices_buffer[:num_actual_toks] # Convert per-request topk indices to global cache slots # topk_indices: (num_decode_tokens, topk) -> physical cache slots topk_indices_physical = triton_convert_req_index_to_global_index( - attn_metadata.req_id_per_token[:num_decode_tokens], + attn_metadata.req_id_per_token[:num_actual_toks], attn_metadata.block_table, topk_indices, BLOCK_SIZE=attn_metadata.block_size, NUM_TOPK_TOKENS=topk_indices.shape[1], ) - # Reshape q for batch decode: (num_decode_tokens, num_heads, head_dim) - # -> (num_decodes, q_len, num_heads, head_dim) - q_len = num_decode_tokens // num_decodes - q = q.view(num_decodes, q_len, q.shape[-2], q.shape[-1]) - # Reshape topk_indices: (num_decode_tokens, topk) # -> (num_decodes, q_len, topk) topk_indices_physical = topk_indices_physical.view(num_decodes, q_len, -1) From d4536e27ed1691767fe1011d5039cf8f089bcc6c Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Tue, 3 Feb 2026 17:58:40 +0000 Subject: [PATCH 11/32] Fix uniform query lengths Signed-off-by: Matthew Bonanni --- docs/design/attention_backends.md | 2 +- .../backends/mla/flashinfer_mla_sparse.py | 177 +++++++++++++++--- 2 files changed, 149 insertions(+), 30 deletions(-) diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index a9e1565d6184..839889f37a28 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -204,7 +204,7 @@ configuration. |---------|--------|-----------|-------------|------------|------|--------|-----------|-----------------|--------------| | `CUTLASS_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | Decoder | 10.x | | `FLASHINFER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | Decoder | 10.x | -| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | 576 | ❌ | ✅ | ❌ | Decoder | 10.x | +| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | 64 | 576 | ❌ | ✅ | ❌ | Decoder | 10.x | | `FLASHMLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | Decoder | 9.x-10.x | | `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | Decoder | 9.x-10.x | | `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | Decoder | 9.x | diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index a5df2a0a8d77..58963bbdb855 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -66,8 +66,6 @@ class FlashInferMLASparseBackend(AttentionBackend): supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ "auto", "bfloat16", - "fp8", - "fp8_e4m3", ] @staticmethod @@ -171,6 +169,9 @@ class FlashInferMLASparseMetadata(AttentionMetadata): block_table: torch.Tensor req_id_per_token: torch.Tensor + # Sequence lengths for all requests (context + query) + seq_lens: torch.Tensor + # Sparse-specific block_size: int = 64 topk_tokens: int = 2048 @@ -268,6 +269,7 @@ def build( slot_mapping=cm.slot_mapping, block_table=cm.block_table_tensor, req_id_per_token=req_id_per_token_tensor, + seq_lens=cm.seq_lens, block_size=self.kv_cache_spec.block_size, topk_tokens=self.topk_tokens, decode=decode_metadata, @@ -332,6 +334,17 @@ def __init__( "FlashInferMLASparseImpl" ) + self.num_heads = num_heads + self.head_size = head_size + self.scale = float(scale) + self.num_kv_heads = num_kv_heads + self.kv_cache_dtype = kv_cache_dtype + + # MLA-specific dimensions + self.kv_lora_rank: int = mla_args["kv_lora_rank"] + self.qk_nope_head_dim: int = mla_args["qk_nope_head_dim"] + self.qk_rope_head_dim: int = mla_args["qk_rope_head_dim"] + assert indexer is not None, "Indexer required for sparse MLA" self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer @@ -339,6 +352,41 @@ def __init__( self.bmm1_scale: float | None = None self.bmm2_scale: float | None = None + def _call_kernel( + self, + q: torch.Tensor, + kv_c_and_k_pe_cache: torch.Tensor, + topk_indices: torch.Tensor, + seq_lens: torch.Tensor, + max_seq_len: int, + topk_tokens: int, + ) -> torch.Tensor: + """Call the FlashInfer sparse MLA kernel for a batch with uniform q_len. + + Args: + q: Query tensor of shape (batch_size, q_len, num_heads, head_dim) + kv_c_and_k_pe_cache: KV cache tensor + topk_indices: Physical indices of shape (batch_size, q_len, topk) + seq_lens: Sequence lengths tensor of shape (batch_size,) + max_seq_len: Maximum sequence length + topk_tokens: Number of top-k tokens for sparse attention + """ + o = trtllm_batch_decode_with_kv_cache_mla( + query=q, + kv_cache=kv_c_and_k_pe_cache.unsqueeze(1), + workspace_buffer=self._workspace_buffer, + qk_nope_head_dim=self.qk_nope_head_dim, + kv_lora_rank=self.kv_lora_rank, + qk_rope_head_dim=self.qk_rope_head_dim, + block_tables=topk_indices, + seq_lens=seq_lens, + max_seq_len=max_seq_len, + bmm1_scale=self.bmm1_scale, + bmm2_scale=self.bmm2_scale, + sparse_mla_top_k=topk_tokens, + ) + return o + def forward_mqa( self, q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], @@ -353,6 +401,9 @@ def forward_mqa( - These are converted to physical cache slots using the block table - The resulting indices are passed as block_tables with shape [batch_size, q_len_per_request, sparse_mla_top_k] + + Mixed batches (decode + prefill) are handled by processing each group + separately, as the kernel requires uniform query lengths within a batch. """ # Concatenate q if it's a tuple (ql_nope, q_pe) if isinstance(q, tuple): @@ -365,7 +416,6 @@ def forward_mqa( topk_indices = self.topk_indices_buffer[:num_actual_toks] # Convert per-request topk indices to global cache slots - # topk_indices: (num_decode_tokens, topk) -> physical cache slots topk_indices_physical = triton_convert_req_index_to_global_index( attn_metadata.req_id_per_token[:num_actual_toks], attn_metadata.block_table, @@ -374,10 +424,6 @@ def forward_mqa( NUM_TOPK_TOKENS=topk_indices.shape[1], ) - # Reshape topk_indices: (num_decode_tokens, topk) - # -> (num_decodes, q_len, topk) - topk_indices_physical = topk_indices_physical.view(num_decodes, q_len, -1) - if self._workspace_buffer is None: self._workspace_buffer = _get_workspace_buffer(q.device) @@ -386,27 +432,100 @@ def forward_mqa( if self.bmm2_scale is None: self.bmm2_scale = layer._v_scale_float - # Call FlashInfer with sparse_mla_top_k - # For sparse mode: - # - block_tables shape is [batch_size, q_len_per_request, sparse_mla_top_k] - # - This contains the physical cache slot indices to attend to - o = trtllm_batch_decode_with_kv_cache_mla( - query=q, - kv_cache=kv_c_and_k_pe_cache.unsqueeze(1), - workspace_buffer=self._workspace_buffer, - qk_nope_head_dim=self.qk_nope_head_dim, - kv_lora_rank=self.kv_lora_rank, - qk_rope_head_dim=self.qk_rope_head_dim, - block_tables=topk_indices_physical, # Sparse indices as block table - seq_lens=attn_metadata.decode.seq_lens, - max_seq_len=attn_metadata.max_seq_len, - bmm1_scale=self.bmm1_scale, - bmm2_scale=self.bmm2_scale, - sparse_mla_top_k=attn_metadata.topk_tokens, - ) + num_decodes = attn_metadata.num_decodes + num_prefills = attn_metadata.num_prefills + num_decode_tokens = attn_metadata.num_decode_tokens + num_prefill_tokens = attn_metadata.num_prefill_tokens + + # Handle pure decode case (most common, optimized path) + if num_prefills == 0 and num_decodes > 0: + assert attn_metadata.decode is not None + decode_q_len = num_decode_tokens // num_decodes + + q_reshaped = q.view(num_decodes, decode_q_len, q.shape[-2], q.shape[-1]) + indices_reshaped = topk_indices_physical.view(num_decodes, decode_q_len, -1) + + o = self._call_kernel( + q_reshaped, + kv_c_and_k_pe_cache, + indices_reshaped, + attn_metadata.decode.seq_lens, + attn_metadata.max_seq_len, + attn_metadata.topk_tokens, + ) + return o.view(-1, o.shape[-2], o.shape[-1]), None + + # Handle pure prefill case + if num_decodes == 0 and num_prefills > 0: + prefill_q_len = num_prefill_tokens // num_prefills - # Reshape output: (num_decodes, q_len, num_heads, head_dim_v) - # -> (num_decode_tokens, num_heads, head_dim_v) - o = o.view(-1, o.shape[-2], o.shape[-1]) + q_reshaped = q.view(num_prefills, prefill_q_len, q.shape[-2], q.shape[-1]) + indices_reshaped = topk_indices_physical.view( + num_prefills, prefill_q_len, -1 + ) + + o = self._call_kernel( + q_reshaped, + kv_c_and_k_pe_cache, + indices_reshaped, + attn_metadata.seq_lens, + attn_metadata.max_seq_len, + attn_metadata.topk_tokens, + ) + return o.view(-1, o.shape[-2], o.shape[-1]), None + + # Handle mixed batch: process decode and prefill separately + # Output tensor to hold results + output = q.new_empty(num_actual_toks, self.num_heads, self.kv_lora_rank) + + # Process decode tokens + if num_decode_tokens > 0: + assert attn_metadata.decode is not None + decode_q_len = num_decode_tokens // num_decodes + + decode_q = q[:num_decode_tokens].view( + num_decodes, decode_q_len, q.shape[-2], q.shape[-1] + ) + decode_indices = topk_indices_physical[:num_decode_tokens].view( + num_decodes, decode_q_len, -1 + ) + + decode_out = self._call_kernel( + decode_q, + kv_c_and_k_pe_cache, + decode_indices, + attn_metadata.decode.seq_lens, + attn_metadata.max_seq_len, + attn_metadata.topk_tokens, + ) + output[:num_decode_tokens] = decode_out.view( + num_decode_tokens, decode_out.shape[-2], decode_out.shape[-1] + ) + + # Process prefill tokens + if num_prefill_tokens > 0: + prefill_q_len = num_prefill_tokens // num_prefills + + prefill_q = q[num_decode_tokens:].view( + num_prefills, prefill_q_len, q.shape[-2], q.shape[-1] + ) + prefill_indices = topk_indices_physical[num_decode_tokens:].view( + num_prefills, prefill_q_len, -1 + ) + + # Get seq_lens for prefill requests (they come after decode requests) + prefill_seq_lens = attn_metadata.seq_lens[num_decodes:] + + prefill_out = self._call_kernel( + prefill_q, + kv_c_and_k_pe_cache, + prefill_indices, + prefill_seq_lens, + attn_metadata.max_seq_len, + attn_metadata.topk_tokens, + ) + output[num_decode_tokens:] = prefill_out.view( + num_prefill_tokens, prefill_out.shape[-2], prefill_out.shape[-1] + ) - return o, None + return output, None From 98530a27a726b5bb1cf6097a42894221c119d152 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Tue, 3 Feb 2026 19:06:54 +0000 Subject: [PATCH 12/32] Update check_and_update_config Signed-off-by: Matthew Bonanni --- vllm/platforms/cuda.py | 41 ++++++++++++++++------------------------- 1 file changed, 16 insertions(+), 25 deletions(-) diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 549bc73248b3..077760c0cd1b 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -180,6 +180,7 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None: use_flashmla = False use_cutlass_mla = False use_flashinfer_mla = False + use_flashmla_sparse = False use_flashinfer_mla_sparse = False from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported @@ -189,16 +190,6 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None: hf_text_config = model_config.hf_text_config qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1) if ( - cls.is_device_capability_family(100) - and use_sparse - and qk_nope_head_dim == 128 - ): - # Blackwell + sparse => Use FlashInfer MLA Sparse - use_flashinfer_mla_sparse = True - vllm_config.attention_config.backend = ( - AttentionBackendEnum.FLASHINFER_MLA_SPARSE - ) - elif ( cls.is_device_capability_family(100) and not use_sparse and qk_nope_head_dim == 128 @@ -226,6 +217,7 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None: use_flashmla = backend == AttentionBackendEnum.FLASHMLA use_cutlass_mla = backend == AttentionBackendEnum.CUTLASS_MLA use_flashinfer_mla = backend == AttentionBackendEnum.FLASHINFER_MLA + use_flashmla_sparse = backend == AttentionBackendEnum.FLASHMLA_SPARSE use_flashinfer_mla_sparse = ( backend == AttentionBackendEnum.FLASHINFER_MLA_SPARSE ) @@ -254,22 +246,21 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None: "Forcing kv cache block size to 64 for FlashInferMLA backend." ) - if use_flashinfer_mla_sparse and cache_config.block_size != 64: - cache_config.block_size = 64 - logger.info( - "Forcing kv cache block size to 64 for FlashInferMLASparse backend." - ) + if use_sparse: + if not (use_flashmla_sparse or use_flashinfer_mla_sparse): + use_flashmla_sparse = True - # TODO(Chen): remove this hacky code - if ( - use_sparse - and not use_flashinfer_mla_sparse - and cache_config.block_size != 64 - ): - cache_config.block_size = 64 - logger.info( - "Forcing kv cache block size to 64 for FlashMLASparse backend." - ) + if use_flashmla_sparse and cache_config.block_size != 64: + cache_config.block_size = 64 + logger.info( + "Forcing kv cache block size to 64 for FlashMLASparse backend." + ) + elif use_flashinfer_mla_sparse and cache_config.block_size != 64: + cache_config.block_size = 64 + logger.info( + "Forcing kv cache block size to 64 for FlashInferMLASparse " + "backend." + ) scheduler_config = vllm_config.scheduler_config # Note: model_config may be None during testing From bb4e94e5c76fb913c3689e838df177e6e40aa83b Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Tue, 3 Feb 2026 19:08:33 +0000 Subject: [PATCH 13/32] Update block size support Signed-off-by: Matthew Bonanni --- docs/design/attention_backends.md | 2 +- vllm/platforms/cuda.py | 5 ++++- vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py | 3 +-- 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index 839889f37a28..8e523c90fb8e 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -204,7 +204,7 @@ configuration. |---------|--------|-----------|-------------|------------|------|--------|-----------|-----------------|--------------| | `CUTLASS_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | Decoder | 10.x | | `FLASHINFER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | Decoder | 10.x | -| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | 64 | 576 | ❌ | ✅ | ❌ | Decoder | 10.x | +| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | 32, 64 | 576 | ❌ | ✅ | ❌ | Decoder | 10.x | | `FLASHMLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | Decoder | 9.x-10.x | | `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | Decoder | 9.x-10.x | | `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | Decoder | 9.x | diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 077760c0cd1b..f884257e428e 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -255,7 +255,10 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None: logger.info( "Forcing kv cache block size to 64 for FlashMLASparse backend." ) - elif use_flashinfer_mla_sparse and cache_config.block_size != 64: + elif use_flashinfer_mla_sparse and cache_config.block_size not in ( + 32, + 64, + ): cache_config.block_size = 64 logger.info( "Forcing kv cache block size to 64 for FlashInferMLASparse " diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 58963bbdb855..887905a0734e 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -70,8 +70,7 @@ class FlashInferMLASparseBackend(AttentionBackend): @staticmethod def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: - # FlashInfer MLA sparse requires block size 64 - return [64] + return [32, 64] @staticmethod def get_name() -> str: From ad8d06ed6a910f4afad31c153cd5e8a452a6b908 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Tue, 3 Feb 2026 19:11:07 +0000 Subject: [PATCH 14/32] Parameterize block sizes Signed-off-by: Matthew Bonanni --- tests/v1/attention/test_sparse_mla_backends.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index fd3fcf0fe97c..defb249318fb 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -167,6 +167,7 @@ def _quantize_dequantize_fp8_ds_mla( @pytest.mark.parametrize("batch_name", list(SPARSE_BACKEND_BATCH_SPECS.keys())) @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_ds_mla"]) @pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4]) +@pytest.mark.parametrize("block_size", [32, 64]) def test_sparse_backend_decode_correctness( default_vllm_config, dist_init, @@ -174,11 +175,18 @@ def test_sparse_backend_decode_correctness( batch_name, kv_cache_dtype, tensor_parallel_size, + block_size, workspace_init, ): if kv_cache_dtype not in backend_cls.supported_kv_cache_dtypes: pytest.skip(f"{backend_cls.get_name()} does not support {kv_cache_dtype}") + supported_block_sizes = backend_cls.get_supported_kernel_block_sizes() + if block_size not in supported_block_sizes: + pytest.skip( + f"{backend_cls.get_name()} does not support block_size={block_size}" + ) + if backend_cls == FlashMLASparseBackend: ok, reason = flashmla.is_flashmla_sparse_supported() if not ok: @@ -207,7 +215,6 @@ def test_sparse_backend_decode_correctness( max_seqlen = max(batch_spec.seq_lens) total_cache_tokens = sum(batch_spec.seq_lens) - block_size = 64 # Note: We use TP=1 to avoid multi-GPU requirements in CI. # The test simulates head partitioning via mocked methods below. From f604adb03a5fa5b3a640ca0724a6b4734fbb5b59 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Tue, 3 Feb 2026 19:42:56 +0000 Subject: [PATCH 15/32] Update benchmark Signed-off-by: Matthew Bonanni --- benchmarks/attention_benchmarks/common.py | 36 ++++- benchmarks/attention_benchmarks/mla_runner.py | 132 +++++++++++++----- .../layers/attention/mla_attention.py | 53 +++---- 3 files changed, 157 insertions(+), 64 deletions(-) diff --git a/benchmarks/attention_benchmarks/common.py b/benchmarks/attention_benchmarks/common.py index 7155bdc3fc5b..6516733eb7d6 100644 --- a/benchmarks/attention_benchmarks/common.py +++ b/benchmarks/attention_benchmarks/common.py @@ -21,7 +21,7 @@ class MockHfConfig: """Mock HuggingFace config that satisfies vLLM's requirements.""" - def __init__(self, mla_dims: dict): + def __init__(self, mla_dims: dict, index_topk: int | None = None): self.num_attention_heads = mla_dims["num_q_heads"] self.num_key_value_heads = mla_dims["num_kv_heads"] self.hidden_size = mla_dims["head_dim"] * mla_dims["num_q_heads"] @@ -32,6 +32,8 @@ def __init__(self, mla_dims: dict): self.qk_rope_head_dim = mla_dims["qk_rope_head_dim"] self.v_head_dim = mla_dims["v_head_dim"] self.qk_head_dim = mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"] + if index_topk is not None: + self.index_topk = index_topk def get_text_config(self): return self @@ -82,6 +84,38 @@ def __call__(self, x: torch.Tensor) -> tuple[torch.Tensor]: return (result,) # Return as tuple to match ColumnParallelLinear API +class MockIndexer: + """Mock Indexer for sparse MLA backends. + + Provides topk_indices_buffer that sparse MLA backends use to determine + which KV cache slots to attend to for each token. + """ + + def __init__( + self, + max_num_tokens: int, + topk_tokens: int, + device: torch.device, + ): + self.topk_tokens = topk_tokens + self.topk_indices_buffer = torch.zeros( + (max_num_tokens, topk_tokens), + dtype=torch.int32, + device=device, + ) + + def fill_random_indices(self, num_tokens: int, max_kv_len: int): + """Fill topk_indices_buffer with random valid indices for benchmarking.""" + indices = torch.randint( + 0, + max_kv_len, + (num_tokens, self.topk_tokens), + dtype=torch.int32, + device=self.topk_indices_buffer.device, + ) + self.topk_indices_buffer[:num_tokens] = indices + + class MockLayer(AttentionLayerBase): """Mock attention layer with scale parameters and impl. diff --git a/benchmarks/attention_benchmarks/mla_runner.py b/benchmarks/attention_benchmarks/mla_runner.py index 2c6c3aaac360..46a856429dfd 100644 --- a/benchmarks/attention_benchmarks/mla_runner.py +++ b/benchmarks/attention_benchmarks/mla_runner.py @@ -16,6 +16,7 @@ from common import ( BenchmarkResult, MockHfConfig, + MockIndexer, MockKVBProj, MockLayer, setup_mla_dims, @@ -62,6 +63,7 @@ def create_minimal_vllm_config( block_size: int = 128, max_num_seqs: int = 256, mla_dims: dict | None = None, + index_topk: int | None = None, ) -> VllmConfig: """ Create minimal VllmConfig for MLA benchmarks. @@ -73,6 +75,8 @@ def create_minimal_vllm_config( max_num_seqs: Maximum number of sequences mla_dims: Optional custom MLA dimensions dict. If not provided, uses setup_mla_dims(model_name) + index_topk: Optional topk value for sparse MLA backends. If provided, + the config will include index_topk for sparse attention. Returns: VllmConfig for benchmarking @@ -82,7 +86,7 @@ def create_minimal_vllm_config( mla_dims = setup_mla_dims(model_name) # Create mock HF config first (avoids downloading from HuggingFace) - mock_hf_config = MockHfConfig(mla_dims) + mock_hf_config = MockHfConfig(mla_dims, index_topk=index_topk) # Create a temporary minimal config.json to avoid HF downloads # This ensures consistent ModelConfig construction without network access @@ -186,6 +190,8 @@ def create_minimal_vllm_config( "flashmla": "FlashMLA", "flashinfer_mla": "FlashInferMLA", "cutlass_mla": "CutlassMLA", + "flashinfer_mla_sparse": "FlashInferMLASparse", + "flashmla_sparse": "FlashMLASparse", } # Special properties that differ from defaults @@ -197,6 +203,15 @@ def create_minimal_vllm_config( "flashinfer_mla": { "block_size": 64, # FlashInfer MLA only supports 32 or 64 }, + "flashinfer_mla_sparse": { + "block_size": 64, # FlashInfer MLA sparse only supports 32 or 64 + "is_sparse": True, + }, + "flashmla_sparse": { + "query_format": "concat", # Single concatenated tensor (vs tuple) + "block_size": 64, # FlashMLA sparse uses fixed block size + "is_sparse": True, + }, } @@ -230,6 +245,7 @@ def _get_backend_config(backend: str) -> dict: "builder_class": f"{name}MetadataBuilder", "query_format": props.get("query_format", "tuple"), "block_size": props.get("block_size", None), + "is_sparse": props.get("is_sparse", False), } @@ -447,6 +463,8 @@ def _create_backend_impl( mla_dims: dict, vllm_config: VllmConfig, device: torch.device, + max_num_tokens: int = 8192, + index_topk: int | None = None, ): """ Create backend implementation instance. @@ -456,9 +474,11 @@ def _create_backend_impl( mla_dims: MLA dimension configuration vllm_config: VllmConfig instance device: Target device + max_num_tokens: Maximum number of tokens for sparse indexer buffer + index_topk: Topk value for sparse MLA backends Returns: - Tuple of (impl, layer, builder_instance) + Tuple of (impl, layer, builder_instance, indexer) """ # Import backend classes backend_module = importlib.import_module(backend_cfg["module"]) @@ -474,26 +494,44 @@ def _create_backend_impl( v_head_dim=mla_dims["v_head_dim"], ) + # Create indexer for sparse backends + indexer = None + if backend_cfg.get("is_sparse", False): + if index_topk is None: + index_topk = 2048 # Default topk for sparse MLA + indexer = MockIndexer( + max_num_tokens=max_num_tokens, + topk_tokens=index_topk, + device=device, + ) + + # Build impl kwargs + impl_kwargs = { + "num_heads": mla_dims["num_q_heads"], + "head_size": mla_dims["head_dim"], + "scale": scale, + "num_kv_heads": mla_dims["num_kv_heads"], + "alibi_slopes": None, + "sliding_window": None, + "kv_cache_dtype": "auto", + "logits_soft_cap": None, + "attn_type": "decoder", + "kv_sharing_target_layer_name": None, + "q_lora_rank": None, + "kv_lora_rank": mla_dims["kv_lora_rank"], + "qk_nope_head_dim": mla_dims["qk_nope_head_dim"], + "qk_rope_head_dim": mla_dims["qk_rope_head_dim"], + "qk_head_dim": mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"], + "v_head_dim": mla_dims["v_head_dim"], + "kv_b_proj": mock_kv_b_proj, + } + + # Add indexer for sparse backends + if indexer is not None: + impl_kwargs["indexer"] = indexer + # Create impl - impl = impl_class( - num_heads=mla_dims["num_q_heads"], - head_size=mla_dims["head_dim"], - scale=scale, - num_kv_heads=mla_dims["num_kv_heads"], - alibi_slopes=None, - sliding_window=None, - kv_cache_dtype="auto", - logits_soft_cap=None, - attn_type="decoder", - kv_sharing_target_layer_name=None, - q_lora_rank=None, - kv_lora_rank=mla_dims["kv_lora_rank"], - qk_nope_head_dim=mla_dims["qk_nope_head_dim"], - qk_rope_head_dim=mla_dims["qk_rope_head_dim"], - qk_head_dim=mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"], - v_head_dim=mla_dims["v_head_dim"], - kv_b_proj=mock_kv_b_proj, - ) + impl = impl_class(**impl_kwargs) # Initialize DCP attributes if not hasattr(impl, "dcp_world_size") or impl.dcp_world_size in (None, -1): @@ -529,7 +567,7 @@ def _create_backend_impl( device=device, ) - return impl, layer, builder_instance + return impl, layer, builder_instance, indexer # ============================================================================ @@ -594,6 +632,7 @@ def _run_single_benchmark( backend_cfg: dict, mla_dims: dict, device: torch.device, + indexer=None, ) -> BenchmarkResult: """ Run a single benchmark iteration. @@ -606,6 +645,7 @@ def _run_single_benchmark( backend_cfg: Backend configuration dict mla_dims: MLA dimension configuration device: Target device + indexer: Optional MockIndexer for sparse backends Returns: BenchmarkResult with timing statistics @@ -613,7 +653,9 @@ def _run_single_benchmark( # Parse batch spec requests = parse_batch_spec(config.batch_spec) q_lens = [r.q_len for r in requests] + kv_lens = [r.kv_len for r in requests] total_q = sum(q_lens) + max_kv_len = max(kv_lens) # Determine block size block_size = backend_cfg["block_size"] or config.block_size @@ -641,8 +683,16 @@ def _run_single_benchmark( torch.bfloat16, ) - # Determine which forward method to use based on metadata - if metadata.decode is not None: + # Fill indexer with random indices for sparse backends + is_sparse = backend_cfg.get("is_sparse", False) + if is_sparse and indexer is not None: + indexer.fill_random_indices(total_q, max_kv_len) + + # Determine which forward method to use + if is_sparse: + # Sparse backends use forward_mqa + forward_fn = lambda: impl.forward_mqa(decode_inputs, kv_cache, metadata, layer) + elif metadata.decode is not None: forward_fn = lambda: impl._forward_decode( decode_inputs, kv_cache, metadata, layer ) @@ -693,11 +743,13 @@ def _run_single_benchmark( def _run_mla_benchmark_batched( backend: str, configs_with_params: list[tuple], # [(config, threshold, num_splits), ...] + index_topk: int = 2048, ) -> list[BenchmarkResult]: """ Unified batched MLA benchmark runner for all backends. - Works for: flashattn_mla, flashmla, flashinfer_mla, cutlass_mla + Works for: flashattn_mla, flashmla, flashinfer_mla, cutlass_mla, + flashinfer_mla_sparse, flashmla_sparse This function reuses backend initialization across multiple benchmarks to avoid setup/teardown overhead. @@ -707,6 +759,7 @@ def _run_mla_benchmark_batched( configs_with_params: List of (config, threshold, num_splits) tuples - threshold: reorder_batch_threshold (FlashAttn/FlashMLA only) - num_splits: num_kv_splits (CUTLASS only) + index_topk: Topk value for sparse MLA backends (default 2048) Returns: List of BenchmarkResult objects @@ -730,19 +783,27 @@ def _run_mla_benchmark_batched( if mla_dims is None: mla_dims = setup_mla_dims("deepseek-v3") + # Determine if this is a sparse backend + is_sparse = backend_cfg.get("is_sparse", False) + # Create and set vLLM config for MLA (reused across all benchmarks) vllm_config = create_minimal_vllm_config( model_name="deepseek-v3", # Used only for model path block_size=block_size, mla_dims=mla_dims, # Use custom dims from config or default + index_topk=index_topk if is_sparse else None, ) results = [] with set_current_vllm_config(vllm_config): - # Create backend impl, layer, and builder (reused across benchmarks) - impl, layer, builder_instance = _create_backend_impl( - backend_cfg, mla_dims, vllm_config, device + # Create backend impl, layer, builder, and indexer (reused across benchmarks) + impl, layer, builder_instance, indexer = _create_backend_impl( + backend_cfg, + mla_dims, + vllm_config, + device, + index_topk=index_topk if is_sparse else None, ) # Run each benchmark with the shared impl @@ -768,6 +829,7 @@ def _run_mla_benchmark_batched( backend_cfg, mla_dims, device, + indexer=indexer, ) results.append(result) @@ -793,20 +855,24 @@ def run_mla_benchmark( config, reorder_batch_threshold: int | None = None, num_kv_splits: int | None = None, + index_topk: int = 2048, ) -> BenchmarkResult | list[BenchmarkResult]: """ Unified MLA benchmark runner for all backends. - Works for: flashattn_mla, flashmla, flashinfer_mla, cutlass_mla + Works for: flashattn_mla, flashmla, flashinfer_mla, cutlass_mla, + flashinfer_mla_sparse, flashmla_sparse Always uses batched execution internally for optimal performance. Args: - backend: Backend name (flashattn_mla, flashmla, flashinfer_mla, cutlass_mla) + backend: Backend name (flashattn_mla, flashmla, flashinfer_mla, cutlass_mla, + flashinfer_mla_sparse, flashmla_sparse) config: BenchmarkConfig or list of (BenchmarkConfig, param) tuples reorder_batch_threshold: Threshold override for FlashAttn/FlashMLA (single config mode only) num_kv_splits: Number of KV splits for CUTLASS (single config mode only) + index_topk: Topk value for sparse MLA backends (default 2048) Returns: BenchmarkResult (single mode) or list of BenchmarkResult (batched mode) @@ -816,9 +882,9 @@ def run_mla_benchmark( # Already in batched format if len(config) > 0 and isinstance(config[0], tuple): # Format: [(cfg, param), ...] where param is threshold or num_splits - if backend in ("flashattn_mla", "flashmla"): + if backend in ("flashattn_mla", "flashmla", "flashmla_sparse"): configs_with_params = [(cfg, param, None) for cfg, param in config] - else: # cutlass_mla or flashinfer_mla + else: # cutlass_mla, flashinfer_mla, or sparse backends configs_with_params = [(cfg, None, param) for cfg, param in config] else: # Format: [cfg, ...] - just configs @@ -830,7 +896,7 @@ def run_mla_benchmark( return_single = True # Use unified batched execution - results = _run_mla_benchmark_batched(backend, configs_with_params) + results = _run_mla_benchmark_batched(backend, configs_with_params, index_topk) # Return single result or list based on input return results[0] if return_single else results diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 501b939c11b1..1e9ce46f0663 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -522,22 +522,6 @@ def forward_impl( k_c_normed = k_c_normed[:num_actual_toks, ...] k_pe = k_pe[:num_actual_toks, ...] - assert ( - attn_metadata.num_decodes is not None - and attn_metadata.num_prefills is not None - and attn_metadata.num_decode_tokens is not None - ) - - has_decode = attn_metadata.num_decodes > 0 - has_prefill = attn_metadata.num_prefills > 0 - num_decode_tokens = attn_metadata.num_decode_tokens - - decode_q = q[:num_decode_tokens] - - prefill_q = q[num_decode_tokens:] - prefill_k_pe = k_pe[num_decode_tokens:] - prefill_k_c_normed = k_c_normed[num_decode_tokens:] - # write the latent and rope to kv cache if kv_cache.numel() > 0: ops.concat_and_cache_mla( @@ -555,27 +539,34 @@ def forward_impl( # Sparse MLA impls only support forward_mqa (decode-style attention) is_sparse_impl = isinstance(self.impl, SparseMLAAttentionImpl) - if has_prefill and not is_sparse_impl: + if is_sparse_impl: + has_decode = True + has_prefill = False + num_decode_tokens = q.size(0) + else: + assert ( + attn_metadata.num_decodes is not None + and attn_metadata.num_prefills is not None + and attn_metadata.num_decode_tokens is not None + ) + has_decode = attn_metadata.num_decodes > 0 + has_prefill = attn_metadata.num_prefills > 0 + num_decode_tokens = attn_metadata.num_decode_tokens + + if has_prefill: self.impl.forward_mha( - prefill_q, - prefill_k_c_normed, - prefill_k_pe, + q[num_decode_tokens:], + k_c_normed[num_decode_tokens:], + k_pe[num_decode_tokens:], kv_cache, attn_metadata, self._k_scale, output=output[num_decode_tokens:], ) - if has_decode or (has_prefill and is_sparse_impl): - # For sparse impl, we always use forward_mqa for all tokens - # For non-sparse impl, we only use forward_mqa for decode tokens - if is_sparse_impl: - mqa_q = q - mqa_output_slice = output - else: - assert attn_metadata.decode is not None - mqa_q = decode_q - mqa_output_slice = output[:num_decode_tokens] + if has_decode: + mqa_q = q[:num_decode_tokens] + mqa_output_slice = output[:num_decode_tokens] mqa_q_nope, mqa_q_pe = mqa_q.split( [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 @@ -644,6 +635,8 @@ def forward_impl( mqa_q = get_dcp_group().all_gather(mqa_q, dim=1) # call decode attn + if not is_sparse_impl: + assert attn_metadata.decode is not None attn_out, lse = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self) # correct dcp attn_out with lse. From f281b288e2863ed65d03c1d4f9de3d62cbfdc9d8 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Tue, 3 Feb 2026 21:01:58 +0000 Subject: [PATCH 16/32] Clean up Signed-off-by: Matthew Bonanni --- .../backends/mla/flashinfer_mla_sparse.py | 169 ++---------------- 1 file changed, 18 insertions(+), 151 deletions(-) diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 887905a0734e..62c3327f33d8 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -40,10 +40,7 @@ from vllm.v1.attention.backends.mla.sparse_utils import ( triton_convert_req_index_to_global_index, ) -from vllm.v1.attention.backends.utils import ( - KVCacheLayoutType, - split_decodes_and_prefills, -) +from vllm.v1.attention.backends.utils import KVCacheLayoutType from vllm.v1.kv_cache_interface import AttentionSpec if TYPE_CHECKING: @@ -145,14 +142,6 @@ def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None": return "HND" -@dataclass -class FlashInferMLASparseDecodeMetadata: - """Decode-specific metadata for FlashInfer MLA Sparse.""" - - block_table: torch.Tensor - seq_lens: torch.Tensor - - @dataclass class FlashInferMLASparseMetadata(AttentionMetadata): """Attention metadata for FlashInfer MLA Sparse backend.""" @@ -175,15 +164,6 @@ class FlashInferMLASparseMetadata(AttentionMetadata): block_size: int = 64 topk_tokens: int = 2048 - # Decode metadata (None during prefill-only) - decode: FlashInferMLASparseDecodeMetadata | None = None - - # For split batches - num_decodes: int = 0 - num_prefills: int = 0 - num_decode_tokens: int = 0 - num_prefill_tokens: int = 0 - class FlashInferMLASparseMetadataBuilder( AttentionMetadataBuilder[FlashInferMLASparseMetadata] @@ -205,9 +185,6 @@ def __init__( self.model_config = vllm_config.model_config self.device = device - # Treat requests with query length <= 1 as decodes - self._init_reorder_batch_threshold(1, supports_spec_as_decode=True) - self.mla_dims = get_mla_dims(self.model_config) self.topk_tokens = vllm_config.model_config.hf_config.index_topk @@ -240,25 +217,6 @@ def build( ) req_id_per_token_tensor = self.req_id_per_token_buffer[:num_tokens] - # Split into decode and prefill - (num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens) = ( - split_decodes_and_prefills( - cm, - decode_threshold=self.reorder_batch_threshold or 1, - require_uniform=True, - ) - ) - - # Build decode metadata if we have decode tokens - decode_metadata = None - if num_decodes > 0: - decode_block_table = cm.block_table_tensor[:num_decodes] - decode_seq_lens = cm.seq_lens[:num_decodes] - decode_metadata = FlashInferMLASparseDecodeMetadata( - block_table=decode_block_table, - seq_lens=decode_seq_lens, - ) - return FlashInferMLASparseMetadata( num_reqs=cm.num_reqs, max_query_len=cm.max_query_len, @@ -271,11 +229,6 @@ def build( seq_lens=cm.seq_lens, block_size=self.kv_cache_spec.block_size, topk_tokens=self.topk_tokens, - decode=decode_metadata, - num_decodes=num_decodes, - num_prefills=num_prefills, - num_decode_tokens=num_decode_tokens, - num_prefill_tokens=num_prefill_tokens, ) @@ -393,28 +346,27 @@ def forward_mqa( attn_metadata: FlashInferMLASparseMetadata, layer: AttentionLayer, ) -> tuple[torch.Tensor, torch.Tensor | None]: - """Forward pass for decode with sparse attention. + """Forward pass with sparse attention. For sparse mode: - topk_indices contains logical indices per token [num_tokens, topk] - These are converted to physical cache slots using the block table - The resulting indices are passed as block_tables with shape - [batch_size, q_len_per_request, sparse_mla_top_k] + [num_reqs, q_len_per_request, sparse_mla_top_k] - Mixed batches (decode + prefill) are handled by processing each group - separately, as the kernel requires uniform query lengths within a batch. + Each token is treated independently, so there is no distinction between + decode and prefill phases. """ - # Concatenate q if it's a tuple (ql_nope, q_pe) if isinstance(q, tuple): q = torch.cat(q, dim=-1) num_actual_toks = q.shape[0] + num_reqs = attn_metadata.num_reqs + q_len = num_actual_toks // num_reqs - # Get topk indices assert self.topk_indices_buffer is not None topk_indices = self.topk_indices_buffer[:num_actual_toks] - # Convert per-request topk indices to global cache slots topk_indices_physical = triton_convert_req_index_to_global_index( attn_metadata.req_id_per_token[:num_actual_toks], attn_metadata.block_table, @@ -431,100 +383,15 @@ def forward_mqa( if self.bmm2_scale is None: self.bmm2_scale = layer._v_scale_float - num_decodes = attn_metadata.num_decodes - num_prefills = attn_metadata.num_prefills - num_decode_tokens = attn_metadata.num_decode_tokens - num_prefill_tokens = attn_metadata.num_prefill_tokens - - # Handle pure decode case (most common, optimized path) - if num_prefills == 0 and num_decodes > 0: - assert attn_metadata.decode is not None - decode_q_len = num_decode_tokens // num_decodes - - q_reshaped = q.view(num_decodes, decode_q_len, q.shape[-2], q.shape[-1]) - indices_reshaped = topk_indices_physical.view(num_decodes, decode_q_len, -1) - - o = self._call_kernel( - q_reshaped, - kv_c_and_k_pe_cache, - indices_reshaped, - attn_metadata.decode.seq_lens, - attn_metadata.max_seq_len, - attn_metadata.topk_tokens, - ) - return o.view(-1, o.shape[-2], o.shape[-1]), None - - # Handle pure prefill case - if num_decodes == 0 and num_prefills > 0: - prefill_q_len = num_prefill_tokens // num_prefills - - q_reshaped = q.view(num_prefills, prefill_q_len, q.shape[-2], q.shape[-1]) - indices_reshaped = topk_indices_physical.view( - num_prefills, prefill_q_len, -1 - ) - - o = self._call_kernel( - q_reshaped, - kv_c_and_k_pe_cache, - indices_reshaped, - attn_metadata.seq_lens, - attn_metadata.max_seq_len, - attn_metadata.topk_tokens, - ) - return o.view(-1, o.shape[-2], o.shape[-1]), None - - # Handle mixed batch: process decode and prefill separately - # Output tensor to hold results - output = q.new_empty(num_actual_toks, self.num_heads, self.kv_lora_rank) - - # Process decode tokens - if num_decode_tokens > 0: - assert attn_metadata.decode is not None - decode_q_len = num_decode_tokens // num_decodes - - decode_q = q[:num_decode_tokens].view( - num_decodes, decode_q_len, q.shape[-2], q.shape[-1] - ) - decode_indices = topk_indices_physical[:num_decode_tokens].view( - num_decodes, decode_q_len, -1 - ) - - decode_out = self._call_kernel( - decode_q, - kv_c_and_k_pe_cache, - decode_indices, - attn_metadata.decode.seq_lens, - attn_metadata.max_seq_len, - attn_metadata.topk_tokens, - ) - output[:num_decode_tokens] = decode_out.view( - num_decode_tokens, decode_out.shape[-2], decode_out.shape[-1] - ) + q_reshaped = q.view(num_reqs, q_len, q.shape[-2], q.shape[-1]) + indices_reshaped = topk_indices_physical.view(num_reqs, q_len, -1) - # Process prefill tokens - if num_prefill_tokens > 0: - prefill_q_len = num_prefill_tokens // num_prefills - - prefill_q = q[num_decode_tokens:].view( - num_prefills, prefill_q_len, q.shape[-2], q.shape[-1] - ) - prefill_indices = topk_indices_physical[num_decode_tokens:].view( - num_prefills, prefill_q_len, -1 - ) - - # Get seq_lens for prefill requests (they come after decode requests) - prefill_seq_lens = attn_metadata.seq_lens[num_decodes:] - - prefill_out = self._call_kernel( - prefill_q, - kv_c_and_k_pe_cache, - prefill_indices, - prefill_seq_lens, - attn_metadata.max_seq_len, - attn_metadata.topk_tokens, - ) - output[num_decode_tokens:] = prefill_out.view( - num_prefill_tokens, prefill_out.shape[-2], prefill_out.shape[-1] - ) - - return output, None + o = self._call_kernel( + q_reshaped, + kv_c_and_k_pe_cache, + indices_reshaped, + attn_metadata.seq_lens, + attn_metadata.max_seq_len, + attn_metadata.topk_tokens, + ) + return o.view(-1, o.shape[-2], o.shape[-1]), None From 9a6fb87c274b3fa9d27bd91517e1cd05b7e2c2dd Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Tue, 3 Feb 2026 21:09:31 +0000 Subject: [PATCH 17/32] Fix Signed-off-by: Matthew Bonanni --- .../backends/mla/flashinfer_mla_sparse.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 62c3327f33d8..6f3f7bcd1b27 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -351,18 +351,11 @@ def forward_mqa( For sparse mode: - topk_indices contains logical indices per token [num_tokens, topk] - These are converted to physical cache slots using the block table - - The resulting indices are passed as block_tables with shape - [num_reqs, q_len_per_request, sparse_mla_top_k] - - Each token is treated independently, so there is no distinction between - decode and prefill phases. """ if isinstance(q, tuple): q = torch.cat(q, dim=-1) num_actual_toks = q.shape[0] - num_reqs = attn_metadata.num_reqs - q_len = num_actual_toks // num_reqs assert self.topk_indices_buffer is not None topk_indices = self.topk_indices_buffer[:num_actual_toks] @@ -383,14 +376,17 @@ def forward_mqa( if self.bmm2_scale is None: self.bmm2_scale = layer._v_scale_float - q_reshaped = q.view(num_reqs, q_len, q.shape[-2], q.shape[-1]) - indices_reshaped = topk_indices_physical.view(num_reqs, q_len, -1) + q_reshaped = q.view(num_actual_toks, 1, q.shape[-2], q.shape[-1]) + indices_reshaped = topk_indices_physical.view(num_actual_toks, 1, -1) + seq_lens_per_token = attn_metadata.seq_lens[ + attn_metadata.req_id_per_token[:num_actual_toks] + ] o = self._call_kernel( q_reshaped, kv_c_and_k_pe_cache, indices_reshaped, - attn_metadata.seq_lens, + seq_lens_per_token, attn_metadata.max_seq_len, attn_metadata.topk_tokens, ) From 5dcf05874b3747aad19a1c4f433b2fbb9d8922e2 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Tue, 3 Feb 2026 21:35:32 +0000 Subject: [PATCH 18/32] Use single batch layout Signed-off-by: Matthew Bonanni --- .../backends/mla/flashinfer_mla_sparse.py | 22 ++++++++----------- 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 6f3f7bcd1b27..ba304be2641a 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -163,6 +163,7 @@ class FlashInferMLASparseMetadata(AttentionMetadata): # Sparse-specific block_size: int = 64 topk_tokens: int = 2048 + topk_tokens_tensor: torch.Tensor | None = None class FlashInferMLASparseMetadataBuilder( @@ -193,6 +194,9 @@ def __init__( dtype=torch.int32, device=device, ) + self.topk_tokens_tensor = torch.tensor( + [self.topk_tokens], dtype=torch.int32, device=device + ) def build( self, @@ -229,6 +233,7 @@ def build( seq_lens=cm.seq_lens, block_size=self.kv_cache_spec.block_size, topk_tokens=self.topk_tokens, + topk_tokens_tensor=self.topk_tokens_tensor, ) @@ -346,12 +351,6 @@ def forward_mqa( attn_metadata: FlashInferMLASparseMetadata, layer: AttentionLayer, ) -> tuple[torch.Tensor, torch.Tensor | None]: - """Forward pass with sparse attention. - - For sparse mode: - - topk_indices contains logical indices per token [num_tokens, topk] - - These are converted to physical cache slots using the block table - """ if isinstance(q, tuple): q = torch.cat(q, dim=-1) @@ -376,18 +375,15 @@ def forward_mqa( if self.bmm2_scale is None: self.bmm2_scale = layer._v_scale_float - q_reshaped = q.view(num_actual_toks, 1, q.shape[-2], q.shape[-1]) - indices_reshaped = topk_indices_physical.view(num_actual_toks, 1, -1) - seq_lens_per_token = attn_metadata.seq_lens[ - attn_metadata.req_id_per_token[:num_actual_toks] - ] + q_reshaped = q.view(1, num_actual_toks, q.shape[-2], q.shape[-1]) + indices_reshaped = topk_indices_physical.view(1, num_actual_toks, -1) o = self._call_kernel( q_reshaped, kv_c_and_k_pe_cache, indices_reshaped, - seq_lens_per_token, - attn_metadata.max_seq_len, + attn_metadata.topk_tokens_tensor, + attn_metadata.topk_tokens, attn_metadata.topk_tokens, ) return o.view(-1, o.shape[-2], o.shape[-1]), None From 9bdfa661f2ec1f3a21ccb9235365ed0069a72177 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Tue, 3 Feb 2026 21:55:00 +0000 Subject: [PATCH 19/32] Remove unnecessary abstraction Signed-off-by: Matthew Bonanni --- .../backends/mla/flashinfer_mla_sparse.py | 55 +++++-------------- 1 file changed, 13 insertions(+), 42 deletions(-) diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index ba304be2641a..3abfa38898db 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -309,41 +309,6 @@ def __init__( self.bmm1_scale: float | None = None self.bmm2_scale: float | None = None - def _call_kernel( - self, - q: torch.Tensor, - kv_c_and_k_pe_cache: torch.Tensor, - topk_indices: torch.Tensor, - seq_lens: torch.Tensor, - max_seq_len: int, - topk_tokens: int, - ) -> torch.Tensor: - """Call the FlashInfer sparse MLA kernel for a batch with uniform q_len. - - Args: - q: Query tensor of shape (batch_size, q_len, num_heads, head_dim) - kv_c_and_k_pe_cache: KV cache tensor - topk_indices: Physical indices of shape (batch_size, q_len, topk) - seq_lens: Sequence lengths tensor of shape (batch_size,) - max_seq_len: Maximum sequence length - topk_tokens: Number of top-k tokens for sparse attention - """ - o = trtllm_batch_decode_with_kv_cache_mla( - query=q, - kv_cache=kv_c_and_k_pe_cache.unsqueeze(1), - workspace_buffer=self._workspace_buffer, - qk_nope_head_dim=self.qk_nope_head_dim, - kv_lora_rank=self.kv_lora_rank, - qk_rope_head_dim=self.qk_rope_head_dim, - block_tables=topk_indices, - seq_lens=seq_lens, - max_seq_len=max_seq_len, - bmm1_scale=self.bmm1_scale, - bmm2_scale=self.bmm2_scale, - sparse_mla_top_k=topk_tokens, - ) - return o - def forward_mqa( self, q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], @@ -378,12 +343,18 @@ def forward_mqa( q_reshaped = q.view(1, num_actual_toks, q.shape[-2], q.shape[-1]) indices_reshaped = topk_indices_physical.view(1, num_actual_toks, -1) - o = self._call_kernel( - q_reshaped, - kv_c_and_k_pe_cache, - indices_reshaped, - attn_metadata.topk_tokens_tensor, - attn_metadata.topk_tokens, - attn_metadata.topk_tokens, + o = trtllm_batch_decode_with_kv_cache_mla( + query=q_reshaped, + kv_cache=kv_c_and_k_pe_cache.unsqueeze(1), + workspace_buffer=self._workspace_buffer, + qk_nope_head_dim=self.qk_nope_head_dim, + kv_lora_rank=self.kv_lora_rank, + qk_rope_head_dim=self.qk_rope_head_dim, + block_tables=indices_reshaped, + seq_lens=attn_metadata.topk_tokens_tensor, + max_seq_len=attn_metadata.topk_tokens, + bmm1_scale=self.bmm1_scale, + bmm2_scale=self.bmm2_scale, + sparse_mla_top_k=attn_metadata.topk_tokens, ) return o.view(-1, o.shape[-2], o.shape[-1]), None From 9043242e2e4e673694bfad8d933ad20f7f82790f Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Tue, 3 Feb 2026 22:49:20 +0000 Subject: [PATCH 20/32] Fix indexing Signed-off-by: Matthew Bonanni --- .../backends/mla/flashinfer_mla_sparse.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 3abfa38898db..6d7638681a92 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -161,9 +161,9 @@ class FlashInferMLASparseMetadata(AttentionMetadata): seq_lens: torch.Tensor # Sparse-specific + topk_tokens_tensor: torch.Tensor block_size: int = 64 topk_tokens: int = 2048 - topk_tokens_tensor: torch.Tensor | None = None class FlashInferMLASparseMetadataBuilder( @@ -194,8 +194,11 @@ def __init__( dtype=torch.int32, device=device, ) - self.topk_tokens_tensor = torch.tensor( - [self.topk_tokens], dtype=torch.int32, device=device + self.topk_tokens_tensor = torch.full( + (vllm_config.scheduler_config.max_num_batched_tokens,), + self.topk_tokens, + dtype=torch.int32, + device=device, ) def build( @@ -340,18 +343,17 @@ def forward_mqa( if self.bmm2_scale is None: self.bmm2_scale = layer._v_scale_float - q_reshaped = q.view(1, num_actual_toks, q.shape[-2], q.shape[-1]) - indices_reshaped = topk_indices_physical.view(1, num_actual_toks, -1) + seq_lens = attn_metadata.topk_tokens_tensor[:num_actual_toks] o = trtllm_batch_decode_with_kv_cache_mla( - query=q_reshaped, + query=q.unsqueeze(1), kv_cache=kv_c_and_k_pe_cache.unsqueeze(1), workspace_buffer=self._workspace_buffer, qk_nope_head_dim=self.qk_nope_head_dim, kv_lora_rank=self.kv_lora_rank, qk_rope_head_dim=self.qk_rope_head_dim, - block_tables=indices_reshaped, - seq_lens=attn_metadata.topk_tokens_tensor, + block_tables=topk_indices_physical.unsqueeze(1), + seq_lens=seq_lens, max_seq_len=attn_metadata.topk_tokens, bmm1_scale=self.bmm1_scale, bmm2_scale=self.bmm2_scale, From 1fb6fb0191c902e7deee55c5882a2093e0c6c1d9 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Tue, 3 Feb 2026 23:19:41 +0000 Subject: [PATCH 21/32] Improve test Signed-off-by: Matthew Bonanni --- .../v1/attention/test_sparse_mla_backends.py | 102 +++++++++++------- 1 file changed, 64 insertions(+), 38 deletions(-) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index defb249318fb..374e655d1dd5 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -5,7 +5,6 @@ import math from types import MethodType, SimpleNamespace -import numpy as np import pytest import torch @@ -211,7 +210,7 @@ def test_sparse_backend_decode_correctness( qk_rope_head_dim = 64 v_head_dim = 128 head_size = kv_lora_rank + qk_rope_head_dim - topk_tokens = 2048 + topk_tokens = 128 max_seqlen = max(batch_spec.seq_lens) total_cache_tokens = sum(batch_spec.seq_lens) @@ -265,11 +264,45 @@ def test_sparse_backend_decode_correctness( seq_lens = batch_spec.seq_lens query_lens = batch_spec.query_lens + # Pre-compute positions and sparse indices for all tokens. + # We need these BEFORE computing the reference to use sparse attention masks. + total_query_tokens = sum(query_lens) + positions = [] + for i in range(batch_spec.batch_size): + s_len = seq_lens[i] + q_len = query_lens[i] + ctx_len = s_len - q_len + for q_idx in range(q_len): + positions.append(ctx_len + q_idx) + + # Create sparse indices with UNIQUE per-token offsets to catch bugs where + # the kernel uses wrong indices for some tokens (e.g., due to incorrect + # tensor shapes like [1, num_tokens, ...] instead of [num_tokens, 1, ...]). + base_indices = torch.arange(topk_tokens, device=device, dtype=torch.int32) + sparse_indices = torch.empty( + total_query_tokens, topk_tokens, dtype=torch.int32, device=device + ) + for tok_idx in range(total_query_tokens): + max_valid_idx = positions[tok_idx] + offset = tok_idx * 7 # Prime number for varied offsets + if max_valid_idx > 0: + tok_indices = (base_indices + offset) % (max_valid_idx + 1) + else: + tok_indices = torch.zeros_like(base_indices) + # Mask out indices beyond topk (though with small topk this won't apply) + tok_indices = torch.where( + base_indices < topk_tokens, + tok_indices, + torch.full_like(tok_indices, -1), + ) + sparse_indices[tok_idx] = tok_indices + all_q_vllm, all_kv_c_vllm, all_k_pe_vllm = [], [], [] kv_c_contexts, k_pe_contexts = [], [] reference_outputs = [] kv_cache_scale = torch.tensor(1.0, dtype=torch.float32, device=device) + global_token_idx = 0 for i in range(batch_spec.batch_size): s_len = seq_lens[i] @@ -302,24 +335,37 @@ def test_sparse_backend_decode_correctness( q_mqa = torch.cat([ql_nope, q_pe], dim=-1) k_mqa = torch.cat([kv_c_full, k_pe_full.squeeze(1)], dim=-1) - k_mqa = k_mqa.unsqueeze(1).expand(-1, num_heads, -1) - v_mqa = kv_c_full.unsqueeze(1).expand(-1, num_heads, -1) + v_mqa = kv_c_full - attn_mask = torch.ones(q_len, s_len, dtype=torch.bool, device=device) - causal_mask = torch.tril(torch.ones(q_len, q_len, device=device)) - attn_mask[:, ctx_len:] = causal_mask + # Compute sparse SDPA reference per query token using its sparse indices + for q_idx in range(q_len): + tok_sparse_idx = sparse_indices[global_token_idx] + valid_mask = tok_sparse_idx >= 0 + valid_indices = tok_sparse_idx[valid_mask].long() - q_sdpa_in = q_mqa.unsqueeze(0).transpose(1, 2) - k_sdpa_in = k_mqa.unsqueeze(0).transpose(1, 2) - v_sdpa_in = v_mqa.unsqueeze(0).transpose(1, 2) + q_tok = q_mqa[q_idx : q_idx + 1] # [1, num_heads, head_dim] + k_sparse = k_mqa[valid_indices] # [num_valid, head_dim] + v_sparse = v_mqa[valid_indices] # [num_valid, kv_lora_rank] - sdpa_out = torch.nn.functional.scaled_dot_product_attention( - q_sdpa_in, k_sdpa_in, v_sdpa_in, attn_mask=attn_mask, scale=scale - ) - sdpa_out = sdpa_out.transpose(1, 2).squeeze(0) + k_sparse = k_sparse.unsqueeze(1).expand(-1, num_heads, -1) + v_sparse = v_sparse.unsqueeze(1).expand(-1, num_heads, -1) - sdpa_out = torch.einsum("qnl,lnv->qnv", sdpa_out, W_UV) - reference_outputs.append(sdpa_out.flatten(start_dim=-2)) + # SDPA: [1, num_heads, 1, head_dim] x [1, num_heads, num_valid, head_dim] + q_sdpa_in = q_tok.unsqueeze(0).transpose(1, 2) + k_sdpa_in = k_sparse.unsqueeze(0).transpose(1, 2) + v_sdpa_in = v_sparse.unsqueeze(0).transpose(1, 2) + + sdpa_out = torch.nn.functional.scaled_dot_product_attention( + q_sdpa_in, k_sdpa_in, v_sdpa_in, scale=scale + ) + sdpa_out = sdpa_out.transpose(1, 2).squeeze( + 0 + ) # [1, num_heads, kv_lora_rank] + + sdpa_out = torch.einsum("qnl,lnv->qnv", sdpa_out, W_UV) + reference_outputs.append(sdpa_out.flatten(start_dim=-2)) + + global_token_idx += 1 all_q_vllm.append(q_c) all_kv_c_vllm.append(kv_c_full[ctx_len:]) @@ -362,28 +408,8 @@ def test_sparse_backend_decode_correctness( common_prefix_len=0, common_attn_metadata=common_attn_metadata ) - starts = np.asarray(common_attn_metadata.query_start_loc_cpu, dtype=np.int32) - seg_lengths = np.diff(starts) - positions = np.arange(starts[-1], dtype=np.int32) - np.repeat( - starts[:-1], seg_lengths - ) - seq_lengths = np.asarray(common_attn_metadata.seq_lens.cpu(), dtype=np.int32) - prefix_lengths = seq_lengths - seg_lengths - positions += np.repeat(prefix_lengths, seg_lengths) - - pos_gpu = torch.as_tensor(positions, device=device, dtype=torch.int32) - topk = metadata.topk_tokens - debug_indices = torch.arange(topk, device=device, dtype=torch.int32).unsqueeze(0) - token_positions = pos_gpu.unsqueeze(1) - causal_mask = debug_indices <= token_positions - debug_indices = torch.where( - causal_mask, debug_indices, torch.full_like(debug_indices, -1) - ) - - # Sparse backends read top-k indices from the indexer-provided - # buffer, so emulate that contract with a simple namespace mock. - debug_indices = debug_indices.expand(metadata.num_actual_tokens, -1).clone() - mock_indexer = SimpleNamespace(topk_indices_buffer=debug_indices) + # Use the pre-computed sparse_indices for the mock indexer + mock_indexer = SimpleNamespace(topk_indices_buffer=sparse_indices) kv_b_proj_weight = torch.cat([W_UK, W_UV], dim=-1) kv_b_proj_weight = kv_b_proj_weight.view( From c35a9c6cd2ee890a8667b4416414e7f80c6120b4 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Thu, 29 Jan 2026 13:40:00 -0500 Subject: [PATCH 22/32] fix up attention benchmarks Signed-off-by: Lucas Wilkinson --- benchmarks/attention_benchmarks/batch_spec.py | 37 +++ benchmarks/attention_benchmarks/common.py | 13 +- .../configs/standard_attention.yaml | 18 +- benchmarks/attention_benchmarks/runner.py | 241 +++++++++++------- 4 files changed, 214 insertions(+), 95 deletions(-) diff --git a/benchmarks/attention_benchmarks/batch_spec.py b/benchmarks/attention_benchmarks/batch_spec.py index 41681796e2e6..9f15f1d8096e 100644 --- a/benchmarks/attention_benchmarks/batch_spec.py +++ b/benchmarks/attention_benchmarks/batch_spec.py @@ -229,3 +229,40 @@ def get_batch_stats(requests: list[BatchRequest]) -> dict: sum(r.kv_len for r in requests) / len(requests) if requests else 0 ), } + + +def get_batch_type(batch_spec: str, spec_decode_threshold: int = 8) -> str: + """ + Classify a batch spec into a type string. + + Args: + batch_spec: Batch specification string (e.g., "q2k", "8q1s1k", "2q2k_8q1s1k") + spec_decode_threshold: Max q_len to be considered spec-decode vs extend + + Returns: + Type string: "prefill", "decode", "spec-decode", "extend", or "mixed (types...)" + """ + requests = parse_batch_spec(batch_spec) + + # Classify each request + types_present = set() + for req in requests: + if req.is_decode: + types_present.add("decode") + elif req.is_prefill: + types_present.add("prefill") + elif req.is_extend: + # Distinguish spec-decode (small q_len) from extend (chunked prefill) + if req.q_len <= spec_decode_threshold: + types_present.add("spec-decode") + else: + types_present.add("extend") + + if len(types_present) == 1: + return types_present.pop() + elif len(types_present) > 1: + # Sort for consistent output + sorted_types = sorted(types_present) + return f"mixed ({'+'.join(sorted_types)})" + else: + return "unknown" diff --git a/benchmarks/attention_benchmarks/common.py b/benchmarks/attention_benchmarks/common.py index 7155bdc3fc5b..190b2f97717e 100644 --- a/benchmarks/attention_benchmarks/common.py +++ b/benchmarks/attention_benchmarks/common.py @@ -12,6 +12,7 @@ import numpy as np import torch +from batch_spec import get_batch_type, parse_batch_spec from rich.console import Console from rich.table import Table @@ -316,12 +317,14 @@ def print_table( backends: List of backend names being compared compare_to_fastest: Show percentage comparison to fastest """ - # Group by batch spec + # Group by batch spec, preserving first-occurrence order by_spec = {} + specs_order = [] for r in results: spec = r.config.batch_spec if spec not in by_spec: by_spec[spec] = {} + specs_order.append(spec) by_spec[spec][r.config.backend] = r # Create shortened backend names for display @@ -337,6 +340,8 @@ def shorten_backend_name(name: str) -> str: table = Table(title="Attention Benchmark Results") table.add_column("Batch\nSpec", no_wrap=True) + table.add_column("Type", no_wrap=True) + table.add_column("Batch\nSize", justify="right", no_wrap=True) multi = len(backends) > 1 for backend in backends: @@ -350,12 +355,14 @@ def shorten_backend_name(name: str) -> str: table.add_column(col_rel, justify="right", no_wrap=False) # Add rows - for spec in sorted(by_spec.keys()): + for spec in specs_order: spec_results = by_spec[spec] times = {b: r.mean_time for b, r in spec_results.items() if r.success} best_time = min(times.values()) if times else 0.0 - row = [spec] + batch_type = get_batch_type(spec) + batch_size = len(parse_batch_spec(spec)) + row = [spec, batch_type, str(batch_size)] for backend in backends: if backend in spec_results: r = spec_results[backend] diff --git a/benchmarks/attention_benchmarks/configs/standard_attention.yaml b/benchmarks/attention_benchmarks/configs/standard_attention.yaml index c0bdb98fbf62..440f0286c662 100644 --- a/benchmarks/attention_benchmarks/configs/standard_attention.yaml +++ b/benchmarks/attention_benchmarks/configs/standard_attention.yaml @@ -1,4 +1,10 @@ # Standard attention backend benchmark configuration +# +# Tests Flash Attention (FA4 on Blackwell, FA3 on Hopper, FA2 on older GPUs) +# and Triton attention backends with various workloads. +# +# Note: FlashInfer requires additional model layer setup and is not supported +# in standalone benchmarks. Use 'flash' and 'triton' backends only. model: num_layers: 32 @@ -25,10 +31,18 @@ batch_specs: - "4q1k_16q1s2k" # 4 prefill + 16 decode - "2q4k_32q1s1k" # 2 large prefill + 32 decode - # Context extension - - "q1ks2k" # 1k query, 2k sequence (chunked prefill) + # Speculative decode (q <= 8) + - "16q2s1k" # 16 requests, 2 spec tokens, 1k KV cache + - "16q4s1k" # 16 requests, 4 spec tokens, 1k KV cache + - "16q8s1k" # 16 requests, 8 spec tokens, 1k KV cache + - "32q4s2k" # 32 requests, 4 spec tokens, 2k KV cache + - "8q8s4k" # 8 requests, 8 spec tokens, 4k KV cache + + # Context extension (chunked prefill) + - "q1ks2k" # 1k query, 2k sequence - "2q1ks4k" # 2 requests: 1k query, 4k sequence +# Available backends: flash, triton, flashinfer backends: - flash - triton diff --git a/benchmarks/attention_benchmarks/runner.py b/benchmarks/attention_benchmarks/runner.py index bf08a1550c0c..79bfca681803 100644 --- a/benchmarks/attention_benchmarks/runner.py +++ b/benchmarks/attention_benchmarks/runner.py @@ -8,7 +8,9 @@ (FlashAttention, Triton, FlashInfer) with real vLLM integration. """ +import logging import types +from contextlib import contextmanager import numpy as np import torch @@ -24,8 +26,13 @@ ParallelConfig, SchedulerConfig, VllmConfig, + set_current_vllm_config, +) +from vllm.v1.attention.backends.utils import ( + CommonAttentionMetadata, + get_kv_cache_layout, + set_kv_cache_layout, ) -from vllm.v1.attention.backends.utils import CommonAttentionMetadata from vllm.v1.kv_cache_interface import FullAttentionSpec # ============================================================================ @@ -37,22 +44,14 @@ "flash": { "module": "vllm.v1.attention.backends.flash_attn", "backend_class": "FlashAttentionBackend", - "dtype": torch.float16, - "cache_layout": "standard", - # ^ [2, num_blocks, block_size, num_kv_heads, head_dim] }, "triton": { "module": "vllm.v1.attention.backends.triton_attn", "backend_class": "TritonAttentionBackend", - "dtype": torch.float32, - "cache_layout": "standard", }, "flashinfer": { "module": "vllm.v1.attention.backends.flashinfer", "backend_class": "FlashInferBackend", - "dtype": torch.float16, - "cache_layout": "flashinfer", - # ^ [num_blocks, 2, block_size, num_kv_heads, head_dim] }, } @@ -66,6 +65,18 @@ def _get_backend_config(backend: str) -> dict: return _BACKEND_CONFIG[backend] +@contextmanager +def log_warnings_and_errors_only(): + """Temporarily set vLLM logger to WARNING level.""" + logger = logging.getLogger("vllm") + old_level = logger.level + logger.setLevel(logging.WARNING) + try: + yield + finally: + logger.setLevel(old_level) + + # ============================================================================ # Metadata Building Helpers # ============================================================================ @@ -88,11 +99,7 @@ def _build_common_attn_metadata( query_start_loc_cpu = query_start_loc.cpu() seq_lens = torch.tensor(kv_lens, dtype=torch.int32, device=device) - seq_lens_cpu = seq_lens.cpu() - max_seq_len = int(seq_lens_cpu.max()) - - context_lens = [kv - q for kv, q in zip(kv_lens, q_lens)] - num_computed_tokens_cpu = torch.tensor(context_lens, dtype=torch.int32) + max_seq_len = int(seq_lens.max().item()) max_blocks = (max(kv_lens) + block_size - 1) // block_size num_blocks = batch_size * max_blocks @@ -107,8 +114,6 @@ def _build_common_attn_metadata( query_start_loc=query_start_loc, query_start_loc_cpu=query_start_loc_cpu, seq_lens=seq_lens, - seq_lens_cpu=seq_lens_cpu, - num_computed_tokens_cpu=num_computed_tokens_cpu, num_reqs=batch_size, num_actual_tokens=total_tokens, max_query_len=max_query_len, @@ -121,7 +126,6 @@ def _build_common_attn_metadata( def _create_vllm_config( config: BenchmarkConfig, - dtype: torch.dtype, max_num_blocks: int, ) -> VllmConfig: """Create a VllmConfig for benchmarking with mock model methods.""" @@ -129,7 +133,7 @@ def _create_vllm_config( model="meta-llama/Meta-Llama-3-8B", tokenizer="meta-llama/Meta-Llama-3-8B", trust_remote_code=False, - dtype=dtype, + dtype="auto", # Use model's native dtype seed=0, max_model_len=1024, ) @@ -198,6 +202,7 @@ def _create_backend_impl( backend_cfg: dict, config: BenchmarkConfig, device: torch.device, + dtype: torch.dtype, ): """Create backend implementation instance.""" import importlib @@ -206,7 +211,6 @@ def _create_backend_impl( backend_class = getattr(backend_module, backend_cfg["backend_class"]) scale = get_attention_scale(config.head_dim) - dtype = backend_cfg["dtype"] impl = backend_class.get_impl_cls()( num_heads=config.num_q_heads, @@ -227,7 +231,7 @@ def _create_backend_impl( layer = MockLayer(device, kv_cache_spec=kv_cache_spec) - return backend_class, impl, layer, dtype + return backend_class, impl, layer def _create_metadata_builder( @@ -235,11 +239,44 @@ def _create_metadata_builder( kv_cache_spec: FullAttentionSpec, vllm_config: VllmConfig, device: torch.device, + backend_name: str = "", ): """Create metadata builder instance.""" - return backend_class.get_builder_cls()( + layer_names = ["layer_0"] + builder_cls = backend_class.get_builder_cls() + + # Flashinfer needs get_per_layer_parameters mocked since we don't have + # real model layers registered + if backend_name == "flashinfer": + import unittest.mock + + from vllm.v1.attention.backends.utils import PerLayerParameters + + def mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls): + head_size = vllm_config.model_config.get_head_size() + return { + layer_name: PerLayerParameters( + window_left=-1, # No sliding window + logits_soft_cap=0.0, # No soft cap + sm_scale=1.0 / (head_size**0.5), # Standard scale + ) + for layer_name in layer_names + } + + with unittest.mock.patch( + "vllm.v1.attention.backends.flashinfer.get_per_layer_parameters", + mock_get_per_layer_parameters, + ): + return builder_cls( + kv_cache_spec=kv_cache_spec, + layer_names=layer_names, + vllm_config=vllm_config, + device=device, + ) + + return builder_cls( kv_cache_spec=kv_cache_spec, - layer_names=["layer_0"], + layer_names=layer_names, vllm_config=vllm_config, device=device, ) @@ -281,39 +318,44 @@ def _create_input_tensors( def _create_kv_cache( config: BenchmarkConfig, max_num_blocks: int, - cache_layout: str, + backend_class, device: torch.device, dtype: torch.dtype, ) -> list: - """Create KV cache tensors for all layers.""" - if cache_layout == "flashinfer": - # FlashInfer layout: [num_blocks, 2, block_size, num_kv_heads, head_dim] - cache_list = [ - torch.zeros( - max_num_blocks, - 2, - config.block_size, - config.num_kv_heads, - config.head_dim, - device=device, - dtype=dtype, - ) - for _ in range(config.num_layers) - ] - else: - # Standard layout: [2, num_blocks, block_size, num_kv_heads, head_dim] - cache_list = [ - torch.zeros( - 2, - max_num_blocks, - config.block_size, - config.num_kv_heads, - config.head_dim, - device=device, - dtype=dtype, - ) - for _ in range(config.num_layers) - ] + """Create KV cache tensors for all layers using the backend's methods. + + Uses the backend's get_kv_cache_shape() and get_kv_cache_stride_order() + to create the cache with the correct shape and memory layout. + """ + # Get the logical shape from the backend + cache_shape = backend_class.get_kv_cache_shape( + num_blocks=max_num_blocks, + block_size=config.block_size, + num_kv_heads=config.num_kv_heads, + head_size=config.head_dim, + ) + + # Get the stride order for custom memory layout + try: + stride_order = backend_class.get_kv_cache_stride_order() + assert len(stride_order) == len(cache_shape) + except (AttributeError, NotImplementedError): + stride_order = tuple(range(len(cache_shape))) + + # Permute shape to physical layout order + physical_shape = tuple(cache_shape[i] for i in stride_order) + + # Compute inverse permutation to get back to logical view + inv_order = [stride_order.index(i) for i in range(len(stride_order))] + + cache_list = [] + for _ in range(config.num_layers): + # Allocate in physical layout order (contiguous in memory) + cache = torch.zeros(*physical_shape, device=device, dtype=dtype) + # Permute to logical view + cache = cache.permute(*inv_order) + cache_list.append(cache) + return cache_list @@ -418,53 +460,72 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult: kv_lens = [r.kv_len for r in requests] total_q = sum(q_lens) max_kv = max(kv_lens) + batch_size = len(q_lens) - max_num_blocks = (max_kv + config.block_size - 1) // config.block_size - - backend_class, impl, layer, dtype = _create_backend_impl( - backend_cfg, config, device - ) + # Calculate total blocks needed: batch_size * max_blocks_per_request + max_blocks_per_request = (max_kv + config.block_size - 1) // config.block_size + max_num_blocks = batch_size * max_blocks_per_request + + # Suppress vLLM logs during setup to reduce spam + with log_warnings_and_errors_only(): + # Create vllm_config first - uses model's native dtype via "auto" + vllm_config = _create_vllm_config(config, max_num_blocks) + dtype = vllm_config.model_config.dtype + + # Wrap everything in set_current_vllm_config context + # This is required for backends like flashinfer that need global config + with set_current_vllm_config(vllm_config): + backend_class, impl, layer = _create_backend_impl( + backend_cfg, config, device, dtype + ) - common_metadata = _build_common_attn_metadata( - q_lens, kv_lens, config.block_size, device - ) + # Set KV cache layout if the backend requires a specific one + # (e.g., FlashInfer requires HND on SM100/Blackwell for TRTLLM attention) + required_layout = backend_class.get_required_kv_cache_layout() + if required_layout is not None: + set_kv_cache_layout(required_layout) + get_kv_cache_layout.cache_clear() - kv_cache_spec = FullAttentionSpec( - block_size=config.block_size, - num_kv_heads=config.num_kv_heads, - head_size=config.head_dim, - dtype=dtype, - ) + common_metadata = _build_common_attn_metadata( + q_lens, kv_lens, config.block_size, device + ) - vllm_config = _create_vllm_config(config, dtype, max_num_blocks) + kv_cache_spec = FullAttentionSpec( + block_size=config.block_size, + num_kv_heads=config.num_kv_heads, + head_size=config.head_dim, + dtype=dtype, + ) - builder = _create_metadata_builder( - backend_class, kv_cache_spec, vllm_config, device - ) + builder = _create_metadata_builder( + backend_class, kv_cache_spec, vllm_config, device, config.backend + ) - attn_metadata = builder.build( - common_prefix_len=0, - common_attn_metadata=common_metadata, - ) + attn_metadata = builder.build( + common_prefix_len=0, + common_attn_metadata=common_metadata, + ) - q_list, k_list, v_list = _create_input_tensors(config, total_q, device, dtype) + q_list, k_list, v_list = _create_input_tensors( + config, total_q, device, dtype + ) - cache_list = _create_kv_cache( - config, max_num_blocks, backend_cfg["cache_layout"], device, dtype - ) + cache_list = _create_kv_cache( + config, max_num_blocks, backend_class, device, dtype + ) - times, mem_stats = _run_single_benchmark( - config, - impl, - layer, - q_list, - k_list, - v_list, - cache_list, - attn_metadata, - device, - dtype, - ) + times, mem_stats = _run_single_benchmark( + config, + impl, + layer, + q_list, + k_list, + v_list, + cache_list, + attn_metadata, + device, + dtype, + ) mean_time = np.mean(times) throughput = total_q / mean_time if mean_time > 0 else 0 From 24a54002ecaebcf6e7ee112635fe0e5d9e1a60dc Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Wed, 4 Feb 2026 12:43:32 -0500 Subject: [PATCH 23/32] clean Signed-off-by: Matthew Bonanni --- .../attention_benchmarks/configs/standard_attention.yaml | 6 ------ 1 file changed, 6 deletions(-) diff --git a/benchmarks/attention_benchmarks/configs/standard_attention.yaml b/benchmarks/attention_benchmarks/configs/standard_attention.yaml index 440f0286c662..591db6837871 100644 --- a/benchmarks/attention_benchmarks/configs/standard_attention.yaml +++ b/benchmarks/attention_benchmarks/configs/standard_attention.yaml @@ -1,10 +1,4 @@ # Standard attention backend benchmark configuration -# -# Tests Flash Attention (FA4 on Blackwell, FA3 on Hopper, FA2 on older GPUs) -# and Triton attention backends with various workloads. -# -# Note: FlashInfer requires additional model layer setup and is not supported -# in standalone benchmarks. Use 'flash' and 'triton' backends only. model: num_layers: 32 From e8e0f4e8fdc90f82d9e7138f25ecaf58f30d58d8 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Wed, 4 Feb 2026 13:05:15 -0500 Subject: [PATCH 24/32] add smoke Signed-off-by: Matthew Bonanni --- .buildkite/test_areas/benchmarks.yaml | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/.buildkite/test_areas/benchmarks.yaml b/.buildkite/test_areas/benchmarks.yaml index 574b642d407b..57080c46f04e 100644 --- a/.buildkite/test_areas/benchmarks.yaml +++ b/.buildkite/test_areas/benchmarks.yaml @@ -17,3 +17,14 @@ steps: - tests/benchmarks/ commands: - pytest -v -s benchmarks/ + +- label: Attention Benchmarks Smoke Test (B200) + device: b200 + num_gpus: 2 + optional: true + timeout_in_minutes: 10 + source_file_dependencies: + - benchmarks/attention_benchmarks/ + - vllm/v1/attention/ + commands: + - python benchmarks/attention_benchmarks/benchmark.py --backends flash flashinfer --batch-specs "8q1s1k" --repeats 1 --warmup-iters 1 From 1daa3d0a7c9594fa41b439e32deca93a098db5b5 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Wed, 4 Feb 2026 19:25:40 +0000 Subject: [PATCH 25/32] Fix and update test Signed-off-by: Matthew Bonanni --- .../v1/attention/test_sparse_mla_backends.py | 90 ++++++++++++++++--- .../backends/mla/flashinfer_mla_sparse.py | 13 +-- .../v1/attention/backends/mla/sparse_utils.py | 28 +++++- 3 files changed, 108 insertions(+), 23 deletions(-) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 374e655d1dd5..fe9ca828963d 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -278,23 +278,33 @@ def test_sparse_backend_decode_correctness( # Create sparse indices with UNIQUE per-token offsets to catch bugs where # the kernel uses wrong indices for some tokens (e.g., due to incorrect # tensor shapes like [1, num_tokens, ...] instead of [num_tokens, 1, ...]). - base_indices = torch.arange(topk_tokens, device=device, dtype=torch.int32) + # Also include -1 masked indices to verify the kernel handles them correctly. sparse_indices = torch.empty( total_query_tokens, topk_tokens, dtype=torch.int32, device=device ) for tok_idx in range(total_query_tokens): max_valid_idx = positions[tok_idx] offset = tok_idx * 7 # Prime number for varied offsets - if max_valid_idx > 0: - tok_indices = (base_indices + offset) % (max_valid_idx + 1) + # Use only half the topk indices as valid, mask the rest with -1 + # This tests that the kernel correctly ignores -1 indices + num_valid = min(topk_tokens // 2, max_valid_idx + 1) + if num_valid > 0: + valid_range = torch.arange(num_valid, device=device, dtype=torch.int32) + tok_indices = (valid_range + offset) % (max_valid_idx + 1) + # Pad with -1 for the remaining positions + tok_indices = torch.cat( + [ + tok_indices, + torch.full( + (topk_tokens - num_valid,), -1, device=device, dtype=torch.int32 + ), + ] + ) else: - tok_indices = torch.zeros_like(base_indices) - # Mask out indices beyond topk (though with small topk this won't apply) - tok_indices = torch.where( - base_indices < topk_tokens, - tok_indices, - torch.full_like(tok_indices, -1), - ) + tok_indices = torch.full( + (topk_tokens,), -1, device=device, dtype=torch.int32 + ) + tok_indices[0] = 0 # At least one valid index sparse_indices[tok_idx] = tok_indices all_q_vllm, all_kv_c_vllm, all_k_pe_vllm = [], [], [] @@ -676,3 +686,63 @@ def test_triton_convert_req_index_to_global_index_with_prefill_workspace(block_s def test_split_prefill_chunks(seq_lens, max_buf, expected): out = split_prefill_chunks(seq_lens, max_buf) assert out == expected + + +def test_triton_convert_returns_valid_counts(): + """Test that return_valid_counts correctly counts non-negative indices.""" + device = torch.device("cuda") + num_tokens = 8 + num_requests = 2 + max_blocks_per_req = 10 + block_size = 64 + num_topk_tokens = 128 + + req_id = torch.tensor([0, 0, 0, 0, 1, 1, 1, 1], dtype=torch.int32, device=device) + block_table = torch.arange( + num_requests * max_blocks_per_req, dtype=torch.int32, device=device + ).view(num_requests, max_blocks_per_req) + + # Create token indices with varying numbers of valid entries + # Token 0: 64 valid, 64 invalid (-1) + # Token 1: 32 valid, 96 invalid + # Token 2: 128 valid (all) + # Token 3: 1 valid, 127 invalid + # etc. + token_indices = torch.full( + (num_tokens, num_topk_tokens), -1, dtype=torch.int32, device=device + ) + expected_valid = [] + for i in range(num_tokens): + num_valid = [64, 32, 128, 1, 64, 32, 128, 1][i] + token_indices[i, :num_valid] = torch.arange( + num_valid, dtype=torch.int32, device=device + ) % (block_size * max_blocks_per_req) + expected_valid.append(num_valid) + + expected_valid_tensor = torch.tensor( + expected_valid, dtype=torch.int32, device=device + ) + + # Test with return_valid_counts=True + result, valid_counts = triton_convert_req_index_to_global_index( + req_id, + block_table, + token_indices, + BLOCK_SIZE=block_size, + NUM_TOPK_TOKENS=num_topk_tokens, + return_valid_counts=True, + ) + + torch.testing.assert_close(valid_counts, expected_valid_tensor, rtol=0, atol=0) + + # Test that return_valid_counts=False returns only the indices + result_only = triton_convert_req_index_to_global_index( + req_id, + block_table, + token_indices, + BLOCK_SIZE=block_size, + NUM_TOPK_TOKENS=num_topk_tokens, + return_valid_counts=False, + ) + assert isinstance(result_only, torch.Tensor) + torch.testing.assert_close(result_only, result, rtol=0, atol=0) diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 6d7638681a92..21a0d99c20c5 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -161,7 +161,6 @@ class FlashInferMLASparseMetadata(AttentionMetadata): seq_lens: torch.Tensor # Sparse-specific - topk_tokens_tensor: torch.Tensor block_size: int = 64 topk_tokens: int = 2048 @@ -194,12 +193,6 @@ def __init__( dtype=torch.int32, device=device, ) - self.topk_tokens_tensor = torch.full( - (vllm_config.scheduler_config.max_num_batched_tokens,), - self.topk_tokens, - dtype=torch.int32, - device=device, - ) def build( self, @@ -236,7 +229,6 @@ def build( seq_lens=cm.seq_lens, block_size=self.kv_cache_spec.block_size, topk_tokens=self.topk_tokens, - topk_tokens_tensor=self.topk_tokens_tensor, ) @@ -327,12 +319,13 @@ def forward_mqa( assert self.topk_indices_buffer is not None topk_indices = self.topk_indices_buffer[:num_actual_toks] - topk_indices_physical = triton_convert_req_index_to_global_index( + topk_indices_physical, seq_lens = triton_convert_req_index_to_global_index( attn_metadata.req_id_per_token[:num_actual_toks], attn_metadata.block_table, topk_indices, BLOCK_SIZE=attn_metadata.block_size, NUM_TOPK_TOKENS=topk_indices.shape[1], + return_valid_counts=True, ) if self._workspace_buffer is None: @@ -343,8 +336,6 @@ def forward_mqa( if self.bmm2_scale is None: self.bmm2_scale = layer._v_scale_float - seq_lens = attn_metadata.topk_tokens_tensor[:num_actual_toks] - o = trtllm_batch_decode_with_kv_cache_mla( query=q.unsqueeze(1), kv_cache=kv_c_and_k_pe_cache.unsqueeze(1), diff --git a/vllm/v1/attention/backends/mla/sparse_utils.py b/vllm/v1/attention/backends/mla/sparse_utils.py index cccc8f7506ac..e4bd0cf425e1 100644 --- a/vllm/v1/attention/backends/mla/sparse_utils.py +++ b/vllm/v1/attention/backends/mla/sparse_utils.py @@ -7,13 +7,14 @@ from vllm.triton_utils import tl, triton -# Kernel with prefill workspace support +# Kernel with prefill workspace support and valid count tracking @triton.jit def _convert_req_index_to_global_index_kernel( req_id_ptr, # int32 [num_tokens] block_table_ptr, # int32 [num_requests, max_num_blocks_per_req] token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS] out_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS] + valid_count_ptr, # int32 [num_tokens] - output valid count per row prefill_request_id_ptr, # int32 [num_tokens], -1 for decode, >=0 for prefill workspace_starts_ptr, # int32 [num_prefill_reqs+1] or nullptr # shapes (compile-time where possible) @@ -21,6 +22,7 @@ def _convert_req_index_to_global_index_kernel( BLOCK_SIZE: tl.constexpr, BLOCK_N: tl.constexpr, # tile width along columns HAS_PREFILL: tl.constexpr, + COUNT_VALID: tl.constexpr, # whether to count valid indices # strides (in elements) bt_stride0, bt_stride1, @@ -74,6 +76,11 @@ def _convert_req_index_to_global_index_kernel( out_ptr_ij = out_ptr + token_id * out_stride0 + indice_id * out_stride1 tl.store(out_ptr_ij, out_val) + # Count valid indices in this tile and atomically add to row total + if COUNT_VALID: + tile_valid_count = tl.sum((~is_invalid_tok).to(tl.int32)) + tl.atomic_add(valid_count_ptr + token_id, tile_valid_count) + def triton_convert_req_index_to_global_index( req_id: torch.Tensor, # int32 [num_tokens] @@ -85,7 +92,8 @@ def triton_convert_req_index_to_global_index( HAS_PREFILL_WORKSPACE: bool = False, prefill_workspace_request_ids: torch.Tensor | None = None, prefill_workspace_starts: torch.Tensor | None = None, -): + return_valid_counts: bool = False, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """ out[token_id, indice_id] = block_table[req_id[token_id], @@ -104,6 +112,9 @@ def triton_convert_req_index_to_global_index( prefill request index (maps to prefill_workspace_starts) prefill_workspace_starts: int32 [num_prefills], 0-indexed workspace starts for each prefill request + + When return_valid_counts is True, also returns the count of valid (non -1) + indices per row, computed during the same kernel pass (no extra overhead). """ assert req_id.dtype == torch.int32 assert block_table.dtype == torch.int32 @@ -129,6 +140,13 @@ def triton_convert_req_index_to_global_index( token_indices_c = token_indices.contiguous() out = torch.empty_like(token_indices_c) + # Allocate valid count buffer if needed (must be zero-initialized for atomics) + valid_counts: torch.Tensor | None = None + if return_valid_counts: + valid_counts = torch.zeros( + num_tokens, dtype=torch.int32, device=token_indices.device + ) + # Strides in elements bt_stride0, bt_stride1 = block_table_c.stride() ti_stride0, ti_stride1 = token_indices_c.stride() @@ -149,6 +167,7 @@ def triton_convert_req_index_to_global_index( block_table_c, token_indices_c, out, + valid_counts, prefill_workspace_request_ids, prefill_workspace_starts, # shapes / constexprs @@ -156,6 +175,7 @@ def triton_convert_req_index_to_global_index( BLOCK_SIZE, BLOCK_N, HAS_PREFILL_WORKSPACE, + return_valid_counts, # strides bt_stride0, bt_stride1, @@ -164,4 +184,8 @@ def triton_convert_req_index_to_global_index( out_stride0, out_stride1, ) + + if return_valid_counts: + assert valid_counts is not None + return out, valid_counts return out From f326062c94949b969bda4a7be395938c3afe28e7 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Thu, 5 Feb 2026 15:12:00 +0000 Subject: [PATCH 26/32] Fix benchmarks Signed-off-by: Matthew Bonanni --- benchmarks/attention_benchmarks/benchmark.py | 37 +++++++++++-------- benchmarks/attention_benchmarks/common.py | 29 ++++++++++----- .../configs/mla_decode.yaml | 4 +- .../configs/reorder_threshold.yaml | 9 ++--- .../configs/speculative_decode.yaml | 9 ++--- benchmarks/attention_benchmarks/mla_runner.py | 4 -- 6 files changed, 51 insertions(+), 41 deletions(-) diff --git a/benchmarks/attention_benchmarks/benchmark.py b/benchmarks/attention_benchmarks/benchmark.py index ba11fca7452f..fff255f174d1 100644 --- a/benchmarks/attention_benchmarks/benchmark.py +++ b/benchmarks/attention_benchmarks/benchmark.py @@ -496,15 +496,18 @@ def main(): if "description" in yaml_config: console.print(f"[dim]{yaml_config['description']}[/]") - # Override args with YAML values - # (YAML takes precedence unless CLI arg was explicitly set) - # Backend(s) - if "backend" in yaml_config: - args.backend = yaml_config["backend"] - args.backends = None - elif "backends" in yaml_config: - args.backends = yaml_config["backends"] - args.backend = None + # Override args with YAML values, but CLI args take precedence + # Check if CLI provided backends (they would be non-None and not default) + cli_backends_provided = args.backends is not None or args.backend is not None + + # Backend(s) - only use YAML if CLI didn't specify + if not cli_backends_provided: + if "backend" in yaml_config: + args.backend = yaml_config["backend"] + args.backends = None + elif "backends" in yaml_config: + args.backends = yaml_config["backends"] + args.backend = None # Check for special modes if "mode" in yaml_config: @@ -544,13 +547,15 @@ def main(): args.num_kv_heads = model.get("num_kv_heads", args.num_kv_heads) args.block_size = model.get("block_size", args.block_size) - # Benchmark settings - if "benchmark" in yaml_config: - bench = yaml_config["benchmark"] - args.device = bench.get("device", args.device) - args.repeats = bench.get("repeats", args.repeats) - args.warmup_iters = bench.get("warmup_iters", args.warmup_iters) - args.profile_memory = bench.get("profile_memory", args.profile_memory) + # Benchmark settings (top-level keys) + if "device" in yaml_config: + args.device = yaml_config["device"] + if "repeats" in yaml_config: + args.repeats = yaml_config["repeats"] + if "warmup_iters" in yaml_config: + args.warmup_iters = yaml_config["warmup_iters"] + if "profile_memory" in yaml_config: + args.profile_memory = yaml_config["profile_memory"] # Parameter sweep configuration if "parameter_sweep" in yaml_config: diff --git a/benchmarks/attention_benchmarks/common.py b/benchmarks/attention_benchmarks/common.py index c321c7048247..87ec434455bd 100644 --- a/benchmarks/attention_benchmarks/common.py +++ b/benchmarks/attention_benchmarks/common.py @@ -527,18 +527,29 @@ def get_attention_scale(head_dim: int) -> float: def is_mla_backend(backend: str) -> bool: """ - Check if backend is an MLA backend using the backend's is_mla() property. + Check if backend is an MLA backend. + + Uses a simple name-based heuristic to avoid import issues when CUDA ops + are not built. This is more robust than trying to import the backend class. Args: - backend: Backend name (e.g., "CUTLASS_MLA", "FLASHINFER_MLA") + backend: Backend name (e.g., "cutlass_mla", "flashinfer_mla_sparse") Returns: True if the backend is an MLA backend, False otherwise """ - from vllm.v1.attention.backends.registry import AttentionBackendEnum - - try: - backend_class = AttentionBackendEnum[backend.upper()].get_class() - return backend_class.is_mla() - except (KeyError, ValueError, ImportError): - return False + # Known MLA backends (lowercase for comparison) + mla_backends = { + "cutlass_mla", + "flashinfer_mla", + "flashinfer_mla_sparse", + "flashmla", + "flashmla_sparse", + "flashattn_mla", + "flash_attn_mla", + "triton_mla", + "rocm_aiter_mla", + "rocm_aiter_mla_sparse", + "rocm_aiter_triton_mla", + } + return backend.lower() in mla_backends diff --git a/benchmarks/attention_benchmarks/configs/mla_decode.yaml b/benchmarks/attention_benchmarks/configs/mla_decode.yaml index aaf4eec9b1c8..7fea501b7ebb 100644 --- a/benchmarks/attention_benchmarks/configs/mla_decode.yaml +++ b/benchmarks/attention_benchmarks/configs/mla_decode.yaml @@ -46,8 +46,8 @@ backends: - flashmla # Hopper only device: "cuda:0" -repeats: 5 -warmup_iters: 3 +repeats: 100 +warmup_iters: 10 profile_memory: true # Backend-specific tuning diff --git a/benchmarks/attention_benchmarks/configs/reorder_threshold.yaml b/benchmarks/attention_benchmarks/configs/reorder_threshold.yaml index 1ea0a12b5338..8ce7944838be 100644 --- a/benchmarks/attention_benchmarks/configs/reorder_threshold.yaml +++ b/benchmarks/attention_benchmarks/configs/reorder_threshold.yaml @@ -62,11 +62,10 @@ model: block_size: 128 # Benchmark settings -benchmark: - device: "cuda:0" - repeats: 15 # More repeats for spec decode variance - warmup_iters: 5 - profile_memory: false +device: "cuda:0" +repeats: 15 # More repeats for spec decode variance +warmup_iters: 5 +profile_memory: false # Output output: diff --git a/benchmarks/attention_benchmarks/configs/speculative_decode.yaml b/benchmarks/attention_benchmarks/configs/speculative_decode.yaml index 56d2428fe74f..e28f7db83aea 100644 --- a/benchmarks/attention_benchmarks/configs/speculative_decode.yaml +++ b/benchmarks/attention_benchmarks/configs/speculative_decode.yaml @@ -48,11 +48,10 @@ backends: # - flashinfer_mla # Benchmark settings -benchmark: - device: "cuda:0" - repeats: 10 # More repeats for statistical significance - warmup_iters: 5 - profile_memory: false +device: "cuda:0" +repeats: 10 # More repeats for statistical significance +warmup_iters: 5 +profile_memory: false # Test these threshold values for optimization parameter_sweep: diff --git a/benchmarks/attention_benchmarks/mla_runner.py b/benchmarks/attention_benchmarks/mla_runner.py index 46a856429dfd..4816ad17e630 100644 --- a/benchmarks/attention_benchmarks/mla_runner.py +++ b/benchmarks/attention_benchmarks/mla_runner.py @@ -124,16 +124,12 @@ def create_minimal_vllm_config( seed=0, max_model_len=32768, quantization=None, - quantization_param_path=None, enforce_eager=False, - max_context_len_to_capture=None, - max_seq_len_to_capture=8192, max_logprobs=20, disable_sliding_window=False, skip_tokenizer_init=True, served_model_name=None, limit_mm_per_prompt=None, - use_async_output_proc=True, config_format="auto", ) finally: From 553b1401395d6f96a77b9ed5603ed25f0ce6c7b3 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Thu, 5 Feb 2026 15:27:29 +0000 Subject: [PATCH 27/32] More benchmark ux improvements Signed-off-by: Matthew Bonanni --- benchmarks/attention_benchmarks/common.py | 31 ++---- .../configs/mla_decode.yaml | 14 +-- .../configs/mla_mixed_batch.yaml | 8 +- .../configs/reorder_threshold.yaml | 2 +- .../configs/speculative_decode.yaml | 6 +- .../configs/standard_attention.yaml | 8 +- benchmarks/attention_benchmarks/mla_runner.py | 103 ++++++++---------- benchmarks/attention_benchmarks/runner.py | 51 ++++----- 8 files changed, 102 insertions(+), 121 deletions(-) diff --git a/benchmarks/attention_benchmarks/common.py b/benchmarks/attention_benchmarks/common.py index 87ec434455bd..c617925a4ed7 100644 --- a/benchmarks/attention_benchmarks/common.py +++ b/benchmarks/attention_benchmarks/common.py @@ -527,29 +527,20 @@ def get_attention_scale(head_dim: int) -> float: def is_mla_backend(backend: str) -> bool: """ - Check if backend is an MLA backend. - - Uses a simple name-based heuristic to avoid import issues when CUDA ops - are not built. This is more robust than trying to import the backend class. + Check if backend is an MLA backend using the AttentionBackendEnum. Args: - backend: Backend name (e.g., "cutlass_mla", "flashinfer_mla_sparse") + backend: Backend name matching AttentionBackendEnum exactly + (e.g., "FLASHMLA_SPARSE") Returns: True if the backend is an MLA backend, False otherwise """ - # Known MLA backends (lowercase for comparison) - mla_backends = { - "cutlass_mla", - "flashinfer_mla", - "flashinfer_mla_sparse", - "flashmla", - "flashmla_sparse", - "flashattn_mla", - "flash_attn_mla", - "triton_mla", - "rocm_aiter_mla", - "rocm_aiter_mla_sparse", - "rocm_aiter_triton_mla", - } - return backend.lower() in mla_backends + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + try: + backend_enum = AttentionBackendEnum[backend] + backend_class = backend_enum.get_class() + return backend_class.is_mla() + except (KeyError, ValueError, ImportError, AttributeError): + return False diff --git a/benchmarks/attention_benchmarks/configs/mla_decode.yaml b/benchmarks/attention_benchmarks/configs/mla_decode.yaml index 7fea501b7ebb..9c3c440a8f82 100644 --- a/benchmarks/attention_benchmarks/configs/mla_decode.yaml +++ b/benchmarks/attention_benchmarks/configs/mla_decode.yaml @@ -40,10 +40,10 @@ batch_specs: - "32q1s32k" # 32 requests, 32k KV cache backends: - - cutlass_mla - - flashinfer_mla - - flashattn_mla # Hopper only - - flashmla # Hopper only + - CUTLASS_MLA + - FLASHINFER_MLA + - FLASH_ATTN_MLA # Hopper only + - FLASHMLA # Hopper only device: "cuda:0" repeats: 100 @@ -51,11 +51,11 @@ warmup_iters: 10 profile_memory: true # Backend-specific tuning -cutlass_mla: +CUTLASS_MLA: num_kv_splits: auto # or specific value like 4, 8, 16 -flashattn_mla: +FLASH_ATTN_MLA: reorder_batch_threshold: 512 -flashmla: +FLASHMLA: reorder_batch_threshold: 1 diff --git a/benchmarks/attention_benchmarks/configs/mla_mixed_batch.yaml b/benchmarks/attention_benchmarks/configs/mla_mixed_batch.yaml index ad3c0dced6ec..b555d90cbf62 100644 --- a/benchmarks/attention_benchmarks/configs/mla_mixed_batch.yaml +++ b/benchmarks/attention_benchmarks/configs/mla_mixed_batch.yaml @@ -45,10 +45,10 @@ batch_specs: - "4q4k_60q1s4k" # 4 prefill + 60 decode backends: - - cutlass_mla - - flashinfer_mla - - flashattn_mla # Hopper only - - flashmla # Hopper only + - CUTLASS_MLA + - FLASHINFER_MLA + - FLASH_ATTN_MLA # Hopper only + - FLASHMLA # Hopper only device: "cuda:0" repeats: 5 diff --git a/benchmarks/attention_benchmarks/configs/reorder_threshold.yaml b/benchmarks/attention_benchmarks/configs/reorder_threshold.yaml index 8ce7944838be..0d76ef0a358c 100644 --- a/benchmarks/attention_benchmarks/configs/reorder_threshold.yaml +++ b/benchmarks/attention_benchmarks/configs/reorder_threshold.yaml @@ -6,7 +6,7 @@ description: "Decode vs Prefill pipeline crossover analysis" # Test FlashAttn MLA -backend: flashattn_mla +backend: FLASH_ATTN_MLA # Mode: decode_vs_prefill comparison (special sweep mode) # For each batch spec, we'll test both decode and prefill pipelines diff --git a/benchmarks/attention_benchmarks/configs/speculative_decode.yaml b/benchmarks/attention_benchmarks/configs/speculative_decode.yaml index e28f7db83aea..47b6d3604d1d 100644 --- a/benchmarks/attention_benchmarks/configs/speculative_decode.yaml +++ b/benchmarks/attention_benchmarks/configs/speculative_decode.yaml @@ -41,11 +41,11 @@ batch_specs: # Backends that support query length > 1 backends: - - flashattn_mla # reorder_batch_threshold = 512 - - flashmla # reorder_batch_threshold = 1 (tunable) + - FLASH_ATTN_MLA # reorder_batch_threshold = 512 + - FLASHMLA # reorder_batch_threshold = 1 (tunable) # FlashInfer-MLA also supports uniform spec-as-decode but with different mechanism -# - flashinfer_mla +# - FLASHINFER_MLA # Benchmark settings device: "cuda:0" diff --git a/benchmarks/attention_benchmarks/configs/standard_attention.yaml b/benchmarks/attention_benchmarks/configs/standard_attention.yaml index 591db6837871..deb5a4b27ff3 100644 --- a/benchmarks/attention_benchmarks/configs/standard_attention.yaml +++ b/benchmarks/attention_benchmarks/configs/standard_attention.yaml @@ -36,11 +36,11 @@ batch_specs: - "q1ks2k" # 1k query, 2k sequence - "2q1ks4k" # 2 requests: 1k query, 4k sequence -# Available backends: flash, triton, flashinfer +# Available backends: FLASH_ATTN, TRITON_ATTN, FLASHINFER backends: - - flash - - triton - - flashinfer + - FLASH_ATTN + - TRITON_ATTN + - FLASHINFER device: "cuda:0" repeats: 5 diff --git a/benchmarks/attention_benchmarks/mla_runner.py b/benchmarks/attention_benchmarks/mla_runner.py index 4816ad17e630..ffcfa457217a 100644 --- a/benchmarks/attention_benchmarks/mla_runner.py +++ b/benchmarks/attention_benchmarks/mla_runner.py @@ -8,8 +8,6 @@ needing full VllmConfig integration. """ -import importlib - import numpy as np import torch from batch_spec import parse_batch_spec @@ -180,68 +178,65 @@ def create_minimal_vllm_config( # ============================================================================ -# Backend name to class name prefix mapping -_BACKEND_NAME_MAP = { - "flashattn_mla": "FlashAttnMLA", - "flashmla": "FlashMLA", - "flashinfer_mla": "FlashInferMLA", - "cutlass_mla": "CutlassMLA", - "flashinfer_mla_sparse": "FlashInferMLASparse", - "flashmla_sparse": "FlashMLASparse", -} - -# Special properties that differ from defaults +# Backend-specific properties that can't be inferred from the backend class +# Keys are AttentionBackendEnum names (uppercase) _BACKEND_PROPERTIES = { - "flashmla": { + "FLASHMLA": { "query_format": "concat", # Single concatenated tensor (vs tuple) - "block_size": 64, # FlashMLA uses fixed block size - }, - "flashinfer_mla": { - "block_size": 64, # FlashInfer MLA only supports 32 or 64 }, - "flashinfer_mla_sparse": { - "block_size": 64, # FlashInfer MLA sparse only supports 32 or 64 - "is_sparse": True, - }, - "flashmla_sparse": { + "FLASHMLA_SPARSE": { "query_format": "concat", # Single concatenated tensor (vs tuple) - "block_size": 64, # FlashMLA sparse uses fixed block size - "is_sparse": True, }, } def _get_backend_config(backend: str) -> dict: """ - Get backend configuration using naming conventions. - - All MLA backends follow the pattern: - - Module: vllm.v1.attention.backends.mla.{backend} - - Impl: {Name}Impl - - Metadata: {Name}Metadata (or MLACommonMetadata) - - DecodeMetadata: {Name}DecodeMetadata (or MLACommonDecodeMetadata) - - MetadataBuilder: {Name}MetadataBuilder + Get backend configuration from AttentionBackendEnum. + + Uses the registry to get the backend class and extract configuration + from its methods (get_impl_cls, get_builder_cls, is_sparse, etc.). + + Args: + backend: Backend name matching AttentionBackendEnum exactly + (e.g., "FLASHMLA_SPARSE") + + Returns: + Dict with backend configuration """ - if backend not in _BACKEND_NAME_MAP: - raise ValueError(f"Unknown backend: {backend}") + from vllm.v1.attention.backends.registry import AttentionBackendEnum - name = _BACKEND_NAME_MAP[backend] + try: + backend_enum = AttentionBackendEnum[backend] + backend_class = backend_enum.get_class() + except (KeyError, ValueError) as e: + valid_backends = [e.name for e in AttentionBackendEnum if e.name != "CUSTOM"] + raise ValueError( + f"Unknown backend: {backend}. " + f"Valid MLA backends: {[b for b in valid_backends if 'MLA' in b]}" + ) from e + + # Get block size from backend class + block_sizes = backend_class.get_supported_kernel_block_sizes() + # Use first supported block size (backends typically support one for MLA) + block_size = block_sizes[0] if block_sizes else None + if hasattr(block_size, "value"): + # Handle MultipleOf enum + block_size = None + + # Check if sparse via class method if available + is_sparse = getattr(backend_class, "is_sparse", lambda: False)() + + # Get properties that can't be inferred props = _BACKEND_PROPERTIES.get(backend, {}) - # Check if backend uses common metadata (FlashInfer, CUTLASS) - uses_common = backend in ("flashinfer_mla", "cutlass_mla") - return { - "module": f"vllm.v1.attention.backends.mla.{backend}", - "impl_class": f"{name}Impl", - "metadata_class": "MLACommonMetadata" if uses_common else f"{name}Metadata", - "decode_metadata_class": "MLACommonDecodeMetadata" - if uses_common - else f"{name}DecodeMetadata", - "builder_class": f"{name}MetadataBuilder", + "backend_class": backend_class, + "impl_class": backend_class.get_impl_cls(), + "builder_class": backend_class.get_builder_cls(), "query_format": props.get("query_format", "tuple"), - "block_size": props.get("block_size", None), - "is_sparse": props.get("is_sparse", False), + "block_size": block_size, + "is_sparse": is_sparse, } @@ -466,7 +461,7 @@ def _create_backend_impl( Create backend implementation instance. Args: - backend_cfg: Backend configuration dict + backend_cfg: Backend configuration dict from _get_backend_config() mla_dims: MLA dimension configuration vllm_config: VllmConfig instance device: Target device @@ -476,9 +471,9 @@ def _create_backend_impl( Returns: Tuple of (impl, layer, builder_instance, indexer) """ - # Import backend classes - backend_module = importlib.import_module(backend_cfg["module"]) - impl_class = getattr(backend_module, backend_cfg["impl_class"]) + # Get classes from backend config (already resolved by _get_backend_config) + impl_class = backend_cfg["impl_class"] + builder_class = backend_cfg["builder_class"] # Calculate scale scale = 1.0 / np.sqrt(mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"]) @@ -549,9 +544,7 @@ def _create_backend_impl( # Create builder instance if needed builder_instance = None - if backend_cfg["builder_class"]: - builder_class = getattr(backend_module, backend_cfg["builder_class"]) - + if builder_class: # Populate static_forward_context so builder can find the layer # MockLayer inherits from AttentionLayerBase, so isinstance checks pass vllm_config.compilation_config.static_forward_context = {"placeholder": layer} diff --git a/benchmarks/attention_benchmarks/runner.py b/benchmarks/attention_benchmarks/runner.py index 79bfca681803..6457a599ab91 100644 --- a/benchmarks/attention_benchmarks/runner.py +++ b/benchmarks/attention_benchmarks/runner.py @@ -40,29 +40,29 @@ # ============================================================================ -_BACKEND_CONFIG = { - "flash": { - "module": "vllm.v1.attention.backends.flash_attn", - "backend_class": "FlashAttentionBackend", - }, - "triton": { - "module": "vllm.v1.attention.backends.triton_attn", - "backend_class": "TritonAttentionBackend", - }, - "flashinfer": { - "module": "vllm.v1.attention.backends.flashinfer", - "backend_class": "FlashInferBackend", - }, -} +def _get_backend_config(backend: str) -> dict: + """ + Get backend configuration from AttentionBackendEnum. + Args: + backend: Backend name matching AttentionBackendEnum exactly + (e.g., "FLASH_ATTN", "TRITON_ATTN", "FLASHINFER") -def _get_backend_config(backend: str) -> dict: - if backend not in _BACKEND_CONFIG: + Returns: + Dict with backend_class + """ + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + try: + backend_enum = AttentionBackendEnum[backend] + backend_class = backend_enum.get_class() + except (KeyError, ValueError) as e: + valid_backends = [b.name for b in AttentionBackendEnum if b.name != "CUSTOM"] raise ValueError( - f"Unknown backend: {backend}. " - f"Available: {', '.join(_BACKEND_CONFIG.keys())}" - ) - return _BACKEND_CONFIG[backend] + f"Unknown backend: {backend}. Valid backends: {valid_backends}" + ) from e + + return {"backend_class": backend_class} @contextmanager @@ -205,10 +205,7 @@ def _create_backend_impl( dtype: torch.dtype, ): """Create backend implementation instance.""" - import importlib - - backend_module = importlib.import_module(backend_cfg["module"]) - backend_class = getattr(backend_module, backend_cfg["backend_class"]) + backend_class = backend_cfg["backend_class"] scale = get_attention_scale(config.head_dim) @@ -247,7 +244,7 @@ def _create_metadata_builder( # Flashinfer needs get_per_layer_parameters mocked since we don't have # real model layers registered - if backend_name == "flashinfer": + if backend_name == "FLASHINFER": import unittest.mock from vllm.v1.attention.backends.utils import PerLayerParameters @@ -438,7 +435,7 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult: """ Run standard attention benchmark with real kernels. - Supports: flash, triton, flashinfer + Supports: FLASH_ATTN, TRITON_ATTN, FLASHINFER Args: config: Benchmark configuration @@ -453,7 +450,7 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult: requests = parse_batch_spec(config.batch_spec) - if config.backend == "flashinfer": + if config.backend == "FLASHINFER": requests = reorder_for_flashinfer(requests) q_lens = [r.q_len for r in requests] From 1e29943f0f9d2d1b8ea0a59b9f9a548498dd416a Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Thu, 5 Feb 2026 15:51:26 +0000 Subject: [PATCH 28/32] Update mla_decode Signed-off-by: Matthew Bonanni --- .../attention_benchmarks/configs/mla_decode.yaml | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/benchmarks/attention_benchmarks/configs/mla_decode.yaml b/benchmarks/attention_benchmarks/configs/mla_decode.yaml index 9c3c440a8f82..d758654dbe80 100644 --- a/benchmarks/attention_benchmarks/configs/mla_decode.yaml +++ b/benchmarks/attention_benchmarks/configs/mla_decode.yaml @@ -3,7 +3,7 @@ model: name: "deepseek-v3" num_layers: 60 - num_q_heads: 128 + num_q_heads: 128 # Base value, can be swept for TP simulation num_kv_heads: 1 # MLA uses single latent KV head_dim: 576 kv_lora_rank: 512 @@ -12,6 +12,13 @@ model: v_head_dim: 128 block_size: 128 # CUTLASS MLA and FlashAttn MLA use 128 +# Model parameter sweep: simulate tensor parallelism by varying num_q_heads +# TP=1: 128 heads, TP=2: 64 heads, TP=4: 32 heads, TP=8: 16 heads +model_parameter_sweep: + param_name: "num_q_heads" + values: [128, 64, 32, 16] + label_format: "{backend}_{value}h" + batch_specs: # Small batches, varying sequence lengths - "16q1s512" # 16 requests, 512 KV cache @@ -34,6 +41,8 @@ batch_specs: # Very large batches - "128q1s1k" # 128 requests, 1k KV cache - "128q1s2k" # 128 requests, 2k KV cache + - "128q1s4k" # 128 requests, 4k KV cache + - "128q1s8k" # 128 requests, 8k KV cache # Long context - "32q1s16k" # 32 requests, 16k KV cache From c19e1131ade50668cd67a81b593cd86a34210b64 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Thu, 5 Feb 2026 15:59:04 +0000 Subject: [PATCH 29/32] Sort benchmark output Signed-off-by: Matthew Bonanni --- benchmarks/attention_benchmarks/benchmark.py | 10 +++++--- benchmarks/attention_benchmarks/common.py | 26 +++++++++++++++++--- 2 files changed, 30 insertions(+), 6 deletions(-) diff --git a/benchmarks/attention_benchmarks/benchmark.py b/benchmarks/attention_benchmarks/benchmark.py index fff255f174d1..de56cbac8474 100644 --- a/benchmarks/attention_benchmarks/benchmark.py +++ b/benchmarks/attention_benchmarks/benchmark.py @@ -43,6 +43,7 @@ ModelParameterSweep, ParameterSweep, ResultsFormatter, + batch_spec_sort_key, is_mla_backend, ) @@ -218,10 +219,13 @@ def run_model_parameter_sweep( by_param_and_spec[key].append(r) break - # Sort by param value then spec + # Sort by param value then spec (batch_size, q_len, kv_len) sorted_keys = sorted( by_param_and_spec.keys(), - key=lambda x: (int(x[0]) if x[0].isdigit() else x[0], x[1]), + key=lambda x: ( + int(x[0]) if x[0].isdigit() else x[0], + batch_spec_sort_key(x[1]), + ), ) current_param_value = None @@ -330,7 +334,7 @@ def run_parameter_sweep( by_spec[spec] = [] by_spec[spec].append(r) - for spec in sorted(by_spec.keys()): + for spec in sorted(by_spec.keys(), key=batch_spec_sort_key): results = by_spec[spec] best = min(results, key=lambda r: r.mean_time) console.print( diff --git a/benchmarks/attention_benchmarks/common.py b/benchmarks/attention_benchmarks/common.py index c617925a4ed7..63993e40812f 100644 --- a/benchmarks/attention_benchmarks/common.py +++ b/benchmarks/attention_benchmarks/common.py @@ -16,6 +16,25 @@ from rich.console import Console from rich.table import Table + +def batch_spec_sort_key(spec: str) -> tuple[int, int, int]: + """ + Extract sorting key from batch spec: (batch_size, max_q_len, max_kv_len). + + This ensures results are sorted by batch size first, then query length, + then sequence length, rather than alphabetically. + """ + try: + requests = parse_batch_spec(spec) + batch_size = len(requests) + max_q_len = max(r.q_len for r in requests) if requests else 0 + max_kv_len = max(r.kv_len for r in requests) if requests else 0 + return (batch_size, max_q_len, max_kv_len) + except Exception: + # Fallback for unparseable specs + return (0, 0, 0) + + # Mock classes for vLLM attention infrastructure @@ -351,16 +370,17 @@ def print_table( backends: List of backend names being compared compare_to_fastest: Show percentage comparison to fastest """ - # Group by batch spec, preserving first-occurrence order + # Group by batch spec by_spec = {} - specs_order = [] for r in results: spec = r.config.batch_spec if spec not in by_spec: by_spec[spec] = {} - specs_order.append(spec) by_spec[spec][r.config.backend] = r + # Sort specs by (batch_size, q_len, kv_len) instead of alphabetically + specs_order = sorted(by_spec.keys(), key=batch_spec_sort_key) + # Create shortened backend names for display def shorten_backend_name(name: str) -> str: """Shorten long backend names for table display.""" From 5cb64d62b104cb33688dabb557d5928afab1e1cf Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Fri, 6 Feb 2026 17:48:30 +0000 Subject: [PATCH 30/32] Add mla prefill case Signed-off-by: Matthew Bonanni --- .../configs/mla_prefill.yaml | 62 +++++++++++++++++++ 1 file changed, 62 insertions(+) create mode 100644 benchmarks/attention_benchmarks/configs/mla_prefill.yaml diff --git a/benchmarks/attention_benchmarks/configs/mla_prefill.yaml b/benchmarks/attention_benchmarks/configs/mla_prefill.yaml new file mode 100644 index 000000000000..ef6b2cb07dc7 --- /dev/null +++ b/benchmarks/attention_benchmarks/configs/mla_prefill.yaml @@ -0,0 +1,62 @@ +# MLA prefill-only benchmark configuration for sparse backends + +model: + name: "deepseek-v3" + num_layers: 60 + num_q_heads: 128 + num_kv_heads: 1 + head_dim: 576 + kv_lora_rank: 512 + qk_nope_head_dim: 128 + qk_rope_head_dim: 64 + v_head_dim: 128 + block_size: 128 + +# Model parameter sweep: simulate tensor parallelism by varying num_q_heads +# TP=1: 128 heads, TP=2: 64 heads, TP=4: 32 heads, TP=8: 16 heads +model_parameter_sweep: + param_name: "num_q_heads" + values: [128, 64, 32, 16] + label_format: "{backend}_{value}h" + +batch_specs: + # Pure prefill + - "1q512" + - "1q1k" + - "1q2k" + - "1q4k" + - "1q8k" + + # Batched pure prefill + - "2q512" + - "2q1k" + - "2q2k" + - "2q4k" + - "2q8k" + - "4q512" + - "4q1k" + - "4q2k" + - "4q4k" + - "4q8k" + - "8q512" + - "8q1k" + - "8q2k" + - "8q4k" + - "8q8k" + + # Extend + - "1q512s4k" + - "1q512s8k" + - "1q1ks8k" + - "1q2ks8k" + - "1q2ks16k" + - "1q4ks16k" + +backends: + - FLASHMLA_SPARSE + - FLASHINFER_MLA_SPARSE + +device: "cuda:0" +repeats: 10 +warmup_iters: 3 +profile_memory: true From 9a31d73f48ff6cdf46b2fe5aed7789522e9e1e27 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Mon, 9 Feb 2026 10:58:25 -0500 Subject: [PATCH 31/32] Prefer FlashInfer at low head counts Signed-off-by: Matthew Bonanni --- .../generate_attention_backend_docs.py | 44 ++++++++++++++++++- .../layers/attention/mla_attention.py | 1 + vllm/platforms/cuda.py | 22 ++++++++-- vllm/platforms/interface.py | 1 + vllm/v1/attention/selector.py | 7 ++- 5 files changed, 68 insertions(+), 7 deletions(-) diff --git a/tools/pre_commit/generate_attention_backend_docs.py b/tools/pre_commit/generate_attention_backend_docs.py index 3cca4959da79..b39a4d70323a 100644 --- a/tools/pre_commit/generate_attention_backend_docs.py +++ b/tools/pre_commit/generate_attention_backend_docs.py @@ -788,10 +788,50 @@ def parse_cuda_priority_lists() -> dict[str, list[str]]: def _get_backends_from_return(stmts: list) -> list[str]: - """Extract backend names from return statements in a list of statements.""" + """Extract backend names from return statements in a list of statements. + + Handles starred unpacking (e.g. ``*sparse_backends``) by resolving the + variable from assignments found in the same statement list. When the + variable is conditionally assigned (inside an ``if/else``), the ``else`` + branch value is used as the representative default. + """ + # Collect variable assignments so we can resolve starred expressions. + # For conditional assignments, last-written (else branch) wins. + var_assigns: dict[str, list[str]] = {} + for stmt in stmts: + if isinstance(stmt, ast.Assign) and isinstance(stmt.value, ast.List): + for target in stmt.targets: + if isinstance(target, ast.Name): + var_assigns[target.id] = [ + e.attr for e in stmt.value.elts if isinstance(e, ast.Attribute) + ] + elif isinstance(stmt, ast.If): + for branch in (stmt.body, stmt.orelse): + for branch_stmt in branch: + if isinstance(branch_stmt, ast.Assign) and isinstance( + branch_stmt.value, ast.List + ): + for target in branch_stmt.targets: + if isinstance(target, ast.Name): + var_assigns[target.id] = [ + e.attr + for e in branch_stmt.value.elts + if isinstance(e, ast.Attribute) + ] + for stmt in stmts: if isinstance(stmt, ast.Return) and isinstance(stmt.value, ast.List): - return [e.attr for e in stmt.value.elts if isinstance(e, ast.Attribute)] + backends: list[str] = [] + for e in stmt.value.elts: + if isinstance(e, ast.Attribute): + backends.append(e.attr) + elif ( + isinstance(e, ast.Starred) + and isinstance(e.value, ast.Name) + and e.value.id in var_assigns + ): + backends.extend(var_assigns[e.value.id]) + return backends return [] diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 1b719330e2a6..3b428d26f917 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -336,6 +336,7 @@ def __init__( block_size, use_mla=True, use_sparse=use_sparse, + num_heads=self.num_heads, ) if ( diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 6ddef6f462c4..b7efe24dc999 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -45,18 +45,29 @@ def _get_backend_priorities( use_mla: bool, device_capability: DeviceCapability, + num_heads: int | None = None, ) -> list[AttentionBackendEnum]: """Get backend priorities with lazy import to avoid circular dependency.""" if use_mla: if device_capability.major == 10: + # Prefer FlashInfer at low head counts (FlashMLA uses padding) + if num_heads is not None and num_heads <= 16: + sparse_backends = [ + AttentionBackendEnum.FLASHINFER_MLA_SPARSE, + AttentionBackendEnum.FLASHMLA_SPARSE, + ] + else: + sparse_backends = [ + AttentionBackendEnum.FLASHMLA_SPARSE, + AttentionBackendEnum.FLASHINFER_MLA_SPARSE, + ] return [ AttentionBackendEnum.FLASHINFER_MLA, AttentionBackendEnum.CUTLASS_MLA, AttentionBackendEnum.FLASH_ATTN_MLA, AttentionBackendEnum.FLASHMLA, AttentionBackendEnum.TRITON_MLA, - AttentionBackendEnum.FLASHMLA_SPARSE, - AttentionBackendEnum.FLASHINFER_MLA_SPARSE, + *sparse_backends, ] else: return [ @@ -295,6 +306,7 @@ def get_valid_backends( cls, device_capability: DeviceCapability, attn_selector_config: "AttentionSelectorConfig", + num_heads: int | None = None, ) -> tuple[ list[tuple["AttentionBackendEnum", int]], dict["AttentionBackendEnum", list[str]], @@ -303,7 +315,9 @@ def get_valid_backends( invalid_reasons = {} backend_priorities = _get_backend_priorities( - attn_selector_config.use_mla, device_capability + attn_selector_config.use_mla, + device_capability, + num_heads, ) for priority, backend in enumerate(backend_priorities): try: @@ -326,6 +340,7 @@ def get_attn_backend_cls( cls, selected_backend: "AttentionBackendEnum", attn_selector_config: "AttentionSelectorConfig", + num_heads: int | None = None, ) -> str: device_capability = cls.get_device_capability() assert device_capability is not None @@ -355,6 +370,7 @@ def get_attn_backend_cls( valid_backends_priorities, invalid_reasons = cls.get_valid_backends( device_capability=device_capability, attn_selector_config=attn_selector_config, + num_heads=num_heads, ) reasons_str = ( "{" diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index c3b189e013e1..748071418682 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -230,6 +230,7 @@ def get_attn_backend_cls( cls, selected_backend: "AttentionBackendEnum", attn_selector_config: "AttentionSelectorConfig", + num_heads: int | None = None, ) -> str: """Get the attention backend class of a device.""" return "" diff --git a/vllm/v1/attention/selector.py b/vllm/v1/attention/selector.py index e364c3235cfe..9580c1d5f355 100644 --- a/vllm/v1/attention/selector.py +++ b/vllm/v1/attention/selector.py @@ -53,6 +53,7 @@ def get_attn_backend( use_sparse: bool = False, use_mm_prefix: bool = False, attn_type: str | None = None, + num_heads: int | None = None, ) -> type[AttentionBackend]: """Selects which attention backend to use and lazily imports it.""" @@ -66,7 +67,6 @@ def get_attn_backend( from vllm.config import get_current_vllm_config vllm_config = get_current_vllm_config() - backend_enum = vllm_config.attention_config.backend attn_selector_config = AttentionSelectorConfig( head_size=head_size, @@ -81,8 +81,9 @@ def get_attn_backend( ) return _cached_get_attn_backend( - backend=backend_enum, + backend=vllm_config.attention_config.backend, attn_selector_config=attn_selector_config, + num_heads=num_heads, ) @@ -90,12 +91,14 @@ def get_attn_backend( def _cached_get_attn_backend( backend, attn_selector_config: AttentionSelectorConfig, + num_heads: int | None = None, ) -> type[AttentionBackend]: from vllm.platforms import current_platform attention_cls = current_platform.get_attn_backend_cls( backend, attn_selector_config=attn_selector_config, + num_heads=num_heads, ) if not attention_cls: raise ValueError( From 6d7b2c327f32bcbc3373bc8db56f0cba676d7d3f Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Mon, 9 Feb 2026 11:39:30 -0500 Subject: [PATCH 32/32] Update other platforms Signed-off-by: Matthew Bonanni --- vllm/platforms/cpu.py | 1 + vllm/platforms/rocm.py | 1 + vllm/platforms/xpu.py | 1 + 3 files changed, 3 insertions(+) diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index 60180b2725cf..7b578dfa122d 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -129,6 +129,7 @@ def get_attn_backend_cls( cls, selected_backend: "AttentionBackendEnum", attn_selector_config: "AttentionSelectorConfig", + num_heads: int | None = None, ) -> str: if selected_backend and selected_backend != AttentionBackendEnum.CPU_ATTN: logger.info("Cannot use %s backend on CPU.", selected_backend) diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 2545e46209a4..1dab8d89e061 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -269,6 +269,7 @@ def get_attn_backend_cls( cls, selected_backend: "AttentionBackendEnum", attn_selector_config: "AttentionSelectorConfig", + num_heads: int | None = None, ) -> str: from vllm._aiter_ops import rocm_aiter_ops diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index 3a0ea8b122c3..8daa2d47f17e 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -48,6 +48,7 @@ def get_attn_backend_cls( cls, selected_backend: "AttentionBackendEnum", attn_selector_config: "AttentionSelectorConfig", + num_heads: int | None = None, ) -> str: from vllm.v1.attention.backends.utils import set_kv_cache_layout