diff --git a/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py b/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py index db2c1becc4..6523588769 100644 --- a/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py +++ b/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py @@ -7,9 +7,9 @@ v_head_dim=256, kv_lora_rank=512 -> head_size 512. These tests pin the selection contract only (no GPU required): * ``get_supported_head_sizes`` advertises 512 so the backend is a candidate; -* ``supports_combination`` admits exactly two shapes for 512 -- a quantized - DS-MLA cache format (zero-padded 576/656B envelope) and a rope-free bf16 - cache on SM90 -- and rejects everything else; +* ``supports_combination`` admits 512 only for rope-free models on SM90, with + a bf16 cache or an fp8_ds_mla cache (zero RoPE slot), and rejects + everything else; * adding 512 must NOT reorder the SM100 priority list for bf16, which is why the bf16 arm is restricted to SM90. """ @@ -19,6 +19,7 @@ from types import SimpleNamespace import pytest import torch +from vllm.platforms import current_platform from vllm.platforms.interface import DeviceCapability from vllm.v1.attention.backends.mla.flashmla_sparse import ( QUANTIZED_DS_MLA_CACHE_FORMATS, @@ -29,7 +30,7 @@ SM90 = DeviceCapability(major=9, minor=0) SM100 = DeviceCapability(major=10, minor=0) BF16_CACHE_DTYPES = (None, "auto", "bfloat16", "float16") -QUANTIZED_CACHE_DTYPES = ("fp8_ds_mla", "nvfp4_ds_mla") +NOPE_SM90_CACHE_DTYPES = (*BF16_CACHE_DTYPES, "fp8_ds_mla") def _supports_combination(head_size, kv_cache_dtype, capability): @@ -86,34 +87,26 @@ def test_quantized_ds_mla_formats_are_the_envelope_set(): assert frozenset({"fp8_ds_mla", "nvfp4_ds_mla"}) == QUANTIZED_DS_MLA_CACHE_FORMATS -@pytest.mark.parametrize("kv_cache_dtype", QUANTIZED_CACHE_DTYPES) -def test_nope_512_quantized_rejected_here(kv_cache_dtype): - # Quantized DS-MLA NoPE-512 is served by the zero-padded 576/656B - # envelope, wired up in a separate change; this backend must reject it. - reason = _supports_combination(512, kv_cache_dtype, SM90) - assert reason is not None - - -@pytest.mark.parametrize("kv_cache_dtype", BF16_CACHE_DTYPES) -def test_nope_512_bf16_sm90_rope_free_accepted(kv_cache_dtype, rope_free_model): +@pytest.mark.parametrize("kv_cache_dtype", NOPE_SM90_CACHE_DTYPES) +def test_nope_512_sm90_rope_free_accepted(kv_cache_dtype, rope_free_model): assert _supports_combination(512, kv_cache_dtype, SM90) is None -@pytest.mark.parametrize("kv_cache_dtype", BF16_CACHE_DTYPES) -def test_nope_512_bf16_sm90_rope_carrying_rejected(kv_cache_dtype, rope_carrying_model): +@pytest.mark.parametrize("kv_cache_dtype", NOPE_SM90_CACHE_DTYPES) +def test_nope_512_sm90_rope_carrying_rejected(kv_cache_dtype, rope_carrying_model): reason = _supports_combination(512, kv_cache_dtype, SM90) assert reason is not None and "rope-free" in reason -@pytest.mark.parametrize("kv_cache_dtype", BF16_CACHE_DTYPES) -def test_nope_512_bf16_sm100_rejected(kv_cache_dtype, rope_free_model): +@pytest.mark.parametrize("kv_cache_dtype", NOPE_SM90_CACHE_DTYPES) +def test_nope_512_sm100_rejected(kv_cache_dtype, rope_free_model): # The SM90 restriction is what keeps the SM100 priority order unchanged. reason = _supports_combination(512, kv_cache_dtype, SM100) assert reason is not None -@pytest.mark.parametrize("kv_cache_dtype", ("fp8", "fp8_e4m3")) -def test_nope_512_plain_fp8_rejected(kv_cache_dtype, rope_free_model): +@pytest.mark.parametrize("kv_cache_dtype", ("fp8", "fp8_e4m3", "nvfp4_ds_mla")) +def test_nope_512_other_quantized_rejected(kv_cache_dtype, rope_free_model): reason = _supports_combination(512, kv_cache_dtype, SM90) assert reason is not None @@ -174,3 +167,15 @@ def test_sm100_bf16_512_priority_unchanged(num_heads): def test_sm100_576_priority_unchanged(): order = _sparse_order("auto", 576, num_heads=32) assert order[0] == "FLASHMLA_SPARSE", order + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA backends") +def test_sm90_nope_fp8_ds_mla_resolves_to_flashmla(): + """The sparse backends ranked ahead of FLASHMLA_SPARSE for SM90 NoPE-512 + reject fp8_ds_mla, so the auto-selector lands on FlashMLA.""" + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + order = _sparse_order("fp8_ds_mla", 512, num_heads=64, capability=SM90) + for name in order[: order.index("FLASHMLA_SPARSE")]: + backend_cls = AttentionBackendEnum[name].get_class() + assert not backend_cls.supports_kv_cache_dtype("fp8_ds_mla"), name diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index 1c192f3f1c..9c504f31ef 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -538,10 +538,9 @@ def create_and_prepopulate_kv_cache( if fp8_attention: if use_fp8_ds_mla: kv_lora_rank = kv_c_contexts[0].shape[-1] - rope_dim = k_pe_contexts[0].shape[-1] # 4 * 4: 4 float32 scale values for 128-element tiles - # 2 * rope_dim: 16-bit RoPE values - kv_entry_size = kv_lora_rank + 4 * 4 + 2 * rope_dim + # 2 * 64: 16-bit RoPE values (zero-filled for NoPE models) + kv_entry_size = kv_lora_rank + 4 * 4 + 2 * 64 elif use_nvfp4_ds_mla: kv_lora_rank = kv_c_contexts[0].shape[-1] rope_dim = k_pe_contexts[0].shape[-1] diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 468e36a787..3d32aed5d6 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -295,7 +295,8 @@ def _quantize_dequantize_fp8_ds_mla( rope_dim = k_pe.shape[-1] num_tokens = kv_c.shape[0] num_blocks = max(1, math.ceil(num_tokens / block_size)) - entry_size = kv_lora_rank + 4 * 4 + 2 * rope_dim + # The 64-dim RoPE slot is always there; NoPE rows zero-fill it. + entry_size = kv_lora_rank + 4 * 4 + 2 * 64 tmp_cache = torch.zeros( num_blocks, block_size, entry_size, dtype=torch.uint8, device=kv_c.device @@ -413,6 +414,13 @@ def _quantize_dequantize_nvfp4_ds_mla( @pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4]) @pytest.mark.parametrize("block_size", [32, 64]) @pytest.mark.parametrize(("q_scale", "k_scale"), [(1.0, 1.0), (2.0, 3.0)]) +@pytest.mark.parametrize( + ("total_num_heads", "qk_nope_head_dim", "qk_rope_head_dim", "v_head_dim"), + [ + pytest.param(128, 128, 64, 128, id="deepseek"), + pytest.param(64, 256, 0, 256, id="glm53_flash_nope"), + ], +) def test_sparse_backend_decode_correctness( default_vllm_config, dist_init, @@ -424,8 +432,18 @@ def test_sparse_backend_decode_correctness( workspace_init, q_scale: float, k_scale: float, + total_num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, monkeypatch, ): + if qk_rope_head_dim == 0 and ( + backend_cls != FlashMLASparseBackend + or kv_cache_dtype not in ("auto", "fp8_ds_mla") + or torch.cuda.get_device_capability()[0] != 9 + ): + pytest.skip("NoPE sparse MLA is served by FlashMLA on SM90 only") if ( batch_name == "large_q_pure_prefill" and backend_cls == FlashMLASparseBackend @@ -484,15 +502,10 @@ def test_sparse_backend_decode_correctness( device = torch.device(DEVICE_TYPE) dtype = torch.bfloat16 - # Model hyper-parameters (kept intentionally small for the unit test) - total_num_heads = 128 # Compute per-rank heads for simulated TP num_heads = max(1, total_num_heads // tensor_parallel_size) kv_lora_rank = 512 - qk_nope_head_dim = 128 - qk_rope_head_dim = 64 - v_head_dim = 128 head_size = kv_lora_rank + qk_rope_head_dim topk_tokens = 128 @@ -3809,6 +3822,7 @@ def test_flashmla_fp8_paths_accept_decode_subset(monkeypatch, use_mixed_batch: b topk_indices_buffer=topk_indices, num_heads=2, kv_lora_rank=1, + q_head_size=3, index_group=None, index_group_index=0, dcp_world_size=1, diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 319b9977f4..ffb68582ea 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -816,8 +816,9 @@ class AttentionImplBase(ABC, Generic[T]): is_sparse: ClassVar[bool] = False # Whether this impl provides a dense-MHA prefill path (forward_mha). Sparse - # impls without one run the top-k MQA path for all requests. - supports_dense_mha_prefill: ClassVar[bool] = True + # impls without one run the top-k MQA path for all requests. Impls may + # override it per instance, depending on the kv-cache layout. + supports_dense_mha_prefill: bool = True # Required attributes that all impls should have num_heads: int diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index 1c3756e6c4..dea9b81116 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, ClassVar import numpy as np import torch import torch.distributed as dist +import torch.nn.functional as F from vllm import _custom_ops as ops from vllm.config import VllmConfig, get_current_vllm_config @@ -127,6 +128,9 @@ QUANTIZED_DS_MLA_CACHE_FORMATS: frozenset[str] = frozenset( {"fp8_ds_mla", "nvfp4_ds_mla"} ) +# RoPE dims in an fp8_ds_mla row (the bf16 tail of the 656B layout). +FP8_DS_MLA_ROPE_DIM = 64 + class FlashMLASparseBackend(AttentionBackend): supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] @@ -186,16 +190,16 @@ class FlashMLASparseBackend(AttentionBackend): ) -> str | None: if head_size == 512: # GLM5Next NoPE (qk_rope_head_dim == 0, kv_lora_rank == 512) has - # head_size 512 and is served here by the direct 512-wide bf16 - # cache on SM90. Quantized DS-MLA caches need the zero-padded - # 576/656B envelope, which is wired up separately; plain fp8, - # SM100 bf16, and rope-carrying 512 models must fall through to + # head_size 512 and is served here on SM90, either by the direct + # 512-wide bf16 cache or by the 576 fp8_ds_mla layout with a zero + # RoPE slot (see FlashMLASparseImpl.q_head_size). Plain fp8, + # SM100, and rope-carrying 512 models must fall through to # FlashInfer/TRITON. if ( - kv_cache_dtype in (None, "auto", "bfloat16", "float16") + kv_cache_dtype in (None, "auto", "bfloat16", "float16", "fp8_ds_mla") and device_capability.major == 9 ): - # Direct bf16 NoPE-512 is only correct for rope-free models. + # NoPE-512 is only correct for rope-free models. # Precedent for reading hf_text_config in supports_combination: # flashinfer_mla_sparse.py. from vllm.config import get_current_vllm_config @@ -210,8 +214,8 @@ class FlashMLASparseBackend(AttentionBackend): ) else: return ( - "FLASHMLA_SPARSE supports head_size 512 only with bf16 " - "kv-cache on SM90 (NoPE), got " + "FLASHMLA_SPARSE supports head_size 512 only with bf16 or " + "fp8_ds_mla kv-cache on SM90 (NoPE), got " f"kv_cache_dtype={kv_cache_dtype}, " f"capability={device_capability}" ) @@ -732,6 +736,16 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): ) self.fp8_decode_padded_heads = self._compute_fp8_decode_padded_heads(num_heads) + # The kernels read queries as wide as the cache rows. A NoPE model on + # fp8_ds_mla keeps the DeepSeek 576 layout: the cache writer zero-fills + # the RoPE slot and forward_mqa zero-pads q_pe to match (q_pe . 0 = 0). + self.q_head_size = head_size + if kv_cache_dtype == "fp8_ds_mla" and self.qk_rope_head_dim == 0: + self.q_head_size = self.kv_lora_rank + FP8_DS_MLA_ROPE_DIM + # Dense-MHA prefill would gather the cache into a + # kv_lora_rank-wide workspace, too narrow for the 576 rows. + self.supports_dense_mha_prefill = False + vllm_config = get_current_vllm_config() max_tokens = vllm_config.scheduler_config.max_num_batched_tokens q_concat_heads = num_heads @@ -741,7 +755,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): // self.prefill_padding * self.prefill_padding ) - q_concat_shape = (max_tokens, q_concat_heads, head_size) + q_concat_shape = (max_tokens, q_concat_heads, self.q_head_size) if is_quantized_kv_cache(kv_cache_dtype): assert kv_cache_dtype in QUANTIZED_DS_MLA_CACHE_FORMATS, ( "FlashMLA Sparse Attention backend only supports the " @@ -777,11 +791,11 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): # PCP+DCP upconverts this rank's KV shard, then all-gathers the # shards into a workspace of the full prefill size. shard_rows //= parallel_config.decode_context_parallel_size - self.prefill_workspace_shape = (shard_rows, head_size) + self.prefill_workspace_shape = (shard_rows, self.q_head_size) self.workspace_specs.append((self.prefill_workspace_shape, torch.bfloat16)) if self.pcp_dcp_kv_gather: self.workspace_specs.append( - ((prefill_workspace_size, head_size), torch.bfloat16) + ((prefill_workspace_size, self.q_head_size), torch.bfloat16) ) prefill_query_heads = num_heads if self.pcp_dcp_kv_gather and self.dcp_world_size > self.pcp_world_size: @@ -792,7 +806,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): self.workspace_specs.extend( (shape, torch.bfloat16) for shape in ( - (max_tokens, padded_prefill_query_heads, head_size), + (max_tokens, padded_prefill_query_heads, self.q_head_size), (max_tokens, padded_prefill_query_heads, self.kv_lora_rank), (max_tokens, prefill_query_heads, self.kv_lora_rank), ) @@ -1370,11 +1384,16 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): if q_pe.size(-1) == 0: # NoPE (GLM5Next): concat_mla_q requires rope_dim == 64, # copy directly into the head-padded buffer instead. - q[:, : ql_nope.shape[1]].copy_(ql_nope) + q[:, : ql_nope.shape[1], : self.kv_lora_rank].copy_(ql_nope) + if self.q_head_size > self.kv_lora_rank: + q[..., self.kv_lora_rank :].zero_() else: ops.concat_mla_q(ql_nope, q_pe, q) else: actual_num_heads = q.shape[1] + if q.shape[-1] < self.q_head_size: + # DCP gathers the unpadded NoPE query. + q = F.pad(q, (0, self.q_head_size - q.shape[-1])) num_actual_toks = q.shape[0]