From 891390fc223e990832714083bffdf092e268f533 Mon Sep 17 00:00:00 2001 From: Leoyzen Date: Sun, 6 Sep 2026 02:53:29 +0800 Subject: [PATCH 1/5] [Attention][MLA] Accept NoPE-512 head size on FLASHMLA_SPARSE for quantized DS-MLA (fp8-ds-mla-nope-sm90 2.1+2.2) head_size=512 (rope-free NoPE models) is only accepted when the kv-cache dtype is a quantized DS-MLA packed format (fp8_ds_mla/nvfp4_ds_mla); the model attention layer serves it via a zero-padded 576/656B envelope with bf16-zero rope bytes [640:768]. bf16/auto NoPE-512 keeps flowing to the FlashInfer/TRITON backends unchanged. Refs openspec change fp8-ds-mla-nope-sm90 Signed-off-by: Leoyzen --- .../attention/backends/mla/flashmla_sparse.py | 23 ++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index d8d8246423ab..3bd0cbb48fa2 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -97,6 +97,15 @@ FlashMLA dequant thread needs are contiguous. See the layout comment in `csrc/libtorch_stable/cache_kernels.cu`. +NOTE: rope-free NoPE models (kv_lora_rank=512, qk_rope_head_dim=0, e.g. +GLM-5.3-Flash) are served through the same fixed 576/656B geometry via a +zero-padded envelope: the model attention layer emits zero `q_pe [T, H, 64]` +and zero `k_pe [T, 64]` (bf16) so the RoPE bytes [640:768] of every 656B +row are bf16 zeros. RoPE is baked at cache-write time, so `q_pe . 0 = 0` +exactly contributes nothing to the attention scores. This path is only +accepted for quantized DS-MLA cache formats (fp8_ds_mla/nvfp4_ds_mla); +bf16 NoPE-512 continues to be served by other backends. + """ # Quantized DS-MLA cache formats served by the FP8/NVFP4 sparse decode kernel @@ -136,7 +145,10 @@ def get_impl_cls() -> type["FlashMLASparseImpl"]: @classmethod def get_supported_head_sizes(cls) -> list[int]: # DeepSeek V3.2 layout: 512 NoPE + 64 RoPE = 576. - return [576] + # NoPE-512 (rope-free models, e.g. GLM-5.3-Flash) is served via the + # zero-padded 576/656B envelope -- see supports_combination and the + # module docstring. + return [576, 512] @classmethod def is_mla(cls) -> bool: @@ -168,6 +180,15 @@ def supports_combination( f"FLASHMLA_SPARSE only supports the {kv_cache_dtype} kv-cache " "dtype on SM100 (Blackwell)" ) + if head_size == 512 and kv_cache_dtype not in QUANTIZED_DS_MLA_CACHE_FORMATS: + # NoPE-512 rides the zero-padded 576/656B envelope, which only + # exists for the quantized DS-MLA packed formats. bf16/auto + # NoPE-512 must keep flowing to the FlashInfer/TRITON backends. + return ( + "FLASHMLA_SPARSE only supports head_size=512 with a quantized " + f"DS-MLA kv-cache dtype ({sorted(QUANTIZED_DS_MLA_CACHE_FORMATS)}), " + f"got kv_cache_dtype={kv_cache_dtype}" + ) return None From 0e4611eb43e9647f508805c06fa870dc0e2871b0 Mon Sep 17 00:00:00 2001 From: Leoyzen Date: Sun, 6 Sep 2026 02:56:39 +0800 Subject: [PATCH 2/5] [Attention][MLA] Zero-padded-rope shim for NoPE MLA on quantized DS-MLA caches (fp8-ds-mla-nope-sm90 3.1+3.2) Shared helper nope_zero_rope_pad in mla_attention.py: when qk_rope_head_dim == 0 and the effective KV dtype is a quantized DS-MLA packed format served by FLASHMLA_SPARSE, MLAAttention promotes its rope dims to the fixed 576/656B geometry and MultiHeadLatentAttentionWrapper's forward injects zero q_pe [T, H, 64] (cat, mirrors the rope>0 path) and a persistent bf16-zero k_pe [T, 1, 64] workspace buffer (re-zeroed per step, CUDA-graph capture stable) ahead of concat_and_cache_mla, so the csrc asserts (kv_lora_rank == 512, pe_dim == 64, 656B row) pass unchanged. rope_dim > 0 model paths are untouched. Refs openspec change fp8-ds-mla-nope-sm90 Signed-off-by: Leoyzen --- .../layers/attention/mla_attention.py | 90 +++++++++++++++++++ vllm/model_executor/layers/mla.py | 11 +++ vllm/models/glm5next/nvidia/attention.py | 19 ++++ 3 files changed, 120 insertions(+) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 00e10e914234..ab7d56ab2800 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -370,6 +370,81 @@ def _canonicalize_sparse_mla_kv_cache_dtype( return kv_cache_dtype +# Rope bytes carried by the zero-padded envelope NoPE models use to ride the +# fixed 576/656B DS-MLA cache geometry (see flashmla_sparse.py). +_NOPE_ZERO_ROPE_PAD_DIM = 64 + +# Persistent bf16-zero k_pe buffer for the zero-padded envelope. Allocated +# workspace-style (see FlashMLASparseImpl's q_concat_buffer) so CUDA-graph +# capture is allocation-stable; the slice is re-zeroed each step, which the +# capture records as a graph node, so replays keep the rope bytes zero. +_NOPE_K_PE_ZERO_BUFFER_MAX_TOKENS = 16384 +_nope_k_pe_zero_buffer: torch.Tensor | None = None + + +def needs_nope_zero_rope_pad(qk_rope_head_dim: int, kv_cache_dtype: str | None) -> bool: + """Whether the zero-padded-rope envelope shim should be active. + + Active only for rope-free MLA (qk_rope_head_dim == 0) whose effective KV + cache is a quantized DS-MLA packed format; the padded 576/656B envelope + only exists for those. bf16/auto NoPE models are unaffected. + """ + return qk_rope_head_dim == 0 and kv_cache_dtype in ( + "fp8_ds_mla", + "nvfp4_ds_mla", + ) + + +def _get_nope_k_pe_zero(num_tokens: int, device: torch.device) -> torch.Tensor: + """Zero k_pe [num_tokens, 1, 64] view of the persistent bf16 zero buffer.""" + global _nope_k_pe_zero_buffer + if ( + _nope_k_pe_zero_buffer is None + or _nope_k_pe_zero_buffer.device != device + or _nope_k_pe_zero_buffer.shape[0] < num_tokens + ): + max_tokens = max(num_tokens, _NOPE_K_PE_ZERO_BUFFER_MAX_TOKENS) + _nope_k_pe_zero_buffer = torch.zeros( + (max_tokens, 1, _NOPE_ZERO_ROPE_PAD_DIM), + dtype=torch.bfloat16, + device=device, + ) + assert _nope_k_pe_zero_buffer is not None + buf = _nope_k_pe_zero_buffer[:num_tokens] + # The buffer is only ever written with zeros, but re-affirm in case a + # foreign writer touched it (defense in depth for the csrc 656B-row + # contract: rope bytes [640:768] must be bf16 zero). Recorded as a graph + # node under capture, so replays keep the rope bytes zero. + buf.zero_() + return buf + + +def nope_zero_rope_pad( + q: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Zero-pad rope-free NoPE MLA onto the fixed 576/656B DS-MLA envelope. + + RoPE is baked at cache-write time, so a zero q_pe/k_pe contributes + ``q_pe . 0 = 0`` exactly: the padded query is ``cat(q, zeros [T, H, 64])`` + and k_pe becomes a persistent bf16-zero ``[T, 1, 64]`` buffer. The + downstream ``concat_and_cache_mla`` asserts (kv_lora_rank == 512, + pe_dim == 64, 656B row) then pass unchanged. + + Returns: + (q_padded, kv_c_normed, k_pe_zero) with q padded to + ``qk_nope_head_dim + 64`` per head and k_pe shaped [T, 1, 64]. + """ + num_tokens = q.shape[0] + zero_pe = _get_nope_k_pe_zero(num_tokens, q.device) + q_pe = zero_pe.expand(num_tokens, q.shape[1], _NOPE_ZERO_ROPE_PAD_DIM).to(q.dtype) + q_padded = torch.cat([q, q_pe], dim=-1) + # k_pe: reuse the same [T, 1, 64] bf16 zeros (MQA single KV head). + k_pe_zero = zero_pe if k_pe.shape[-1] == 0 else k_pe + return q_padded, kv_c_normed, k_pe_zero + + def _get_kv_b_proj_input_dtype( kv_b_proj: ColumnParallelLinear, use_fp8_prefill: bool ) -> torch.dtype | None: @@ -519,6 +594,21 @@ def __init__( self.kv_cache_dtype = kv_cache_dtype _init_kv_cache_quant(self, quant_config, prefix) + # Zero-padded-rope envelope for rope-free NoPE on a quantized DS-MLA + # cache: promote the rope dims so the fixed 576/656B geometry holds. + # The wrapper layer pads q with zero q_pe and swaps in zero k_pe + # (nope_zero_rope_pad) before calling forward, so the csrc + # concat_and_cache_mla asserts (kv_lora_rank == 512, pe_dim == 64, + # 656B row) pass unchanged. bf16/auto NoPE and rope>0 models are + # untouched (needs_nope_zero_rope_pad gates on the quantized formats). + if ( + needs_nope_zero_rope_pad(qk_rope_head_dim, self.kv_cache_dtype) + and self.attn_backend.get_name() == "FLASHMLA_SPARSE" + ): + self.qk_rope_head_dim = _NOPE_ZERO_ROPE_PAD_DIM + self.qk_head_dim = self.qk_nope_head_dim + _NOPE_ZERO_ROPE_PAD_DIM + self.head_size = self.kv_lora_rank + _NOPE_ZERO_ROPE_PAD_DIM + if ( cache_config is not None and cache_config.enable_prefix_caching diff --git a/vllm/model_executor/layers/mla.py b/vllm/model_executor/layers/mla.py index 421181306236..74364ce87f4b 100644 --- a/vllm/model_executor/layers/mla.py +++ b/vllm/model_executor/layers/mla.py @@ -7,6 +7,7 @@ from vllm.config import CacheConfig from vllm.model_executor.custom_op import PluggableLayer from vllm.model_executor.layers.attention import MLAAttention +from vllm.model_executor.layers.attention.mla_attention import nope_zero_rope_pad from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.models.common.ops import fused_q_kv_rmsnorm from vllm.platforms import current_platform @@ -216,6 +217,16 @@ def forward( heads *= q_proj_layer.group_size q = q.view(-1, heads, self.qk_head_dim) + # Rope-free NoPE riding the zero-padded 576/656B DS-MLA envelope: the + # MLA attention layer below has promoted its rope dim (see + # MLAAttention.__init__), so pad q with zero q_pe [T, H, 64] and swap + # in the persistent zero k_pe [T, 1, 64] here, ahead of the + # concat_and_cache_mla path. rope_dim > 0 models never enter this. + if self.qk_rope_head_dim == 0 and ( + self.mla_attn.qk_rope_head_dim != self.qk_rope_head_dim + ): + q, kv_c_normed, k_pe = nope_zero_rope_pad(q, kv_c_normed, k_pe) + if self.rotary_emb is not None: q[..., self.qk_nope_head_dim :], k_pe = self.rotary_emb( positions, q[..., self.qk_nope_head_dim :], k_pe diff --git a/vllm/models/glm5next/nvidia/attention.py b/vllm/models/glm5next/nvidia/attention.py index 0be749348494..ed533a4da059 100644 --- a/vllm/models/glm5next/nvidia/attention.py +++ b/vllm/models/glm5next/nvidia/attention.py @@ -13,6 +13,9 @@ get_tensor_model_parallel_world_size, ) from vllm.logger import init_logger +from vllm.model_executor.layers.attention.mla_attention import ( + needs_nope_zero_rope_pad, +) from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -583,6 +586,22 @@ def __init__( fuse_qkv_rmsnorm=True, ) + # NoPE (rope_dim == 0) on a quantized DS-MLA cache rides the + # zero-padded 576/656B envelope: MLAAttention promotes its rope dims + # and the wrapper's forward injects zero q_pe/k_pe + # (nope_zero_rope_pad in mla_attention.py) ahead of + # concat_and_cache_mla, so the csrc asserts (kv_lora_rank == 512, + # pe_dim == 64, 656B row) pass unchanged. rope_dim > 0 checkpoints + # never activate this. + if needs_nope_zero_rope_pad( + self.qk_rope_head_dim, + cache_config.cache_dtype if cache_config is not None else None, + ): + logger.info_once( + "GLM-5.3-Flash NoPE MLA using zero-padded 576/656B " + "DS-MLA envelope (rope bytes [640:768] are bf16 zeros)." + ) + def forward( self, hidden_states: torch.Tensor, positions: torch.Tensor ) -> torch.Tensor: From f49706ab0f587460ec997e0a4f37e6227b14f2df Mon Sep 17 00:00:00 2001 From: Leoyzen Date: Sun, 6 Sep 2026 02:56:58 +0800 Subject: [PATCH 3/5] [Attention] Prefer FLASHMLA_SPARSE over FlashInfer SM90 for fp8_ds_mla NoPE-512 (fp8-ds-mla-nope-sm90 4.1) In the SM90 sparse-MLA priority tail, when head_size == 512 and the kv-cache dtype is fp8_ds_mla, order FLASHMLA_SPARSE (packed 656B path with decode LSE for DCP/MTP) ahead of FLASHINFER_MLA_SPARSE_SM90. Plain fp8 keeps selecting FlashInfer (stopgap preserved, no silent dtype reinterpretation); bf16/auto and head_size == 576 priorities are unchanged. Refs openspec change fp8-ds-mla-nope-sm90 Signed-off-by: Leoyzen --- vllm/platforms/cuda.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index d1df152c63de..41aca4c29abe 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -140,7 +140,17 @@ def _get_backend_priorities( ] flashinfer_sparse = AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM90 if head_size == 512: - sparse_tail.insert(0, flashinfer_sparse) + if kv_cache_dtype == "fp8_ds_mla": + # NoPE-512 on the packed fp8_ds_mla format rides the + # FlashMLA zero-padded 576/656B envelope (decode LSE for + # DCP/MTP); order it ahead of the FlashInfer SM90 + # stopgap. Plain fp8 keeps FlashInfer -- no silent dtype + # reinterpretation (see _canonicalize_sparse_mla_ + # kv_cache_dtype in mla_attention.py, which only promotes + # fp8 -> fp8_ds_mla for FLASHMLA_SPARSE). + sparse_tail.insert(2, flashinfer_sparse) + else: + sparse_tail.insert(0, flashinfer_sparse) else: sparse_tail.append(flashinfer_sparse) return [ From c3e08b9b77d51a2172568f65046afaf82227528c Mon Sep 17 00:00:00 2001 From: Leoyzen Date: Sun, 6 Sep 2026 04:00:06 +0800 Subject: [PATCH 4/5] [Bugfix][Attention] Size sparse-MLA chunked prefill workspace from the runtime layers (fp8-ds-mla-nope-sm90) The NoPE zero-padded-rope shim promotes the rope dim on the MLAAttention layer (kv_lora_rank 512, rope 0 -> padded 576/656B envelope), but SparseMLACommonMetadataBuilder still derived its MLA dims from the raw HF config, allocating the chunked-prefill workspace 512 rows wide. On the fp8_ds_mla context-gather path (cp_gather_and_upconvert_fp8_kv_cache, called from MLACommonImpl._compute_prefill_context) the csrc reads head_dim from the workspace row width and asserts 576, so GLM-5.3-Flash NoPE + fp8_ds_mla on FLASHMLA_SPARSE crashed at cache_kernels.cu:1669 on the first chunked prefill. Derive the builder's mla_dims from the instantiated layers in the static forward context (mirroring MLACommonMetadataBuilder), so the workspace follows the promoted 576 geometry. Raw-config dims are unchanged for every other model (rope>0 and bf16 NoPE paths read the same values from both sources). Zero csrc; zero cache-geometry changes (656B row width was always state_content_bytes-driven). Updates the dcp_direct_a2a_lse_reduce builder test to mock the layer dims on the forward-context layer instead of the (now unused) get_mla_dims hook. Refs openspec change fp8-ds-mla-nope-sm90 (bugfix on 4fcc182a0c) Signed-off-by: Leoyzen --- .../test_dcp_direct_a2a_lse_reduce.py | 14 +++++----- .../layers/attention/sparse_mla_attention.py | 27 +++++++++++++++++-- 2 files changed, 33 insertions(+), 8 deletions(-) diff --git a/tests/distributed/test_dcp_direct_a2a_lse_reduce.py b/tests/distributed/test_dcp_direct_a2a_lse_reduce.py index e5764954bd53..a27492cb7c7f 100644 --- a/tests/distributed/test_dcp_direct_a2a_lse_reduce.py +++ b/tests/distributed/test_dcp_direct_a2a_lse_reduce.py @@ -516,15 +516,17 @@ def test_sparse_mla_builder_initializes_dcp_manager(monkeypatch): "get_dcp_group", lambda: MagicMock(world_size=2), ) - monkeypatch.setattr( - sparse_mla, - "get_mla_dims", - lambda _: MagicMock(kv_lora_rank=8, qk_rope_head_dim=4), - ) manager = object.__new__(dcp.MLADCPManager) manager.init_kv_gather = MagicMock() - layer = MagicMock(dcp_manager=manager) + layer = MagicMock( + dcp_manager=manager, + q_lora_rank=None, + kv_lora_rank=8, + qk_nope_head_dim=4, + qk_rope_head_dim=4, + v_head_dim=4, + ) config = MagicMock() config.model_config.dtype = torch.bfloat16 config.model_config.max_model_len = 64 diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index ab6946132f4b..844b84183785 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -19,10 +19,10 @@ MLACommonBaseImpl, MLACommonMetadata, MLACommonPrefillMetadata, + MLADims, accumulate_mla_context_chunk, align_mla_chunked_context_workspace_size, build_mla_chunked_context_metadata, - get_mla_dims, init_mla_context_partial, ) from vllm.platforms import current_platform @@ -116,6 +116,29 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]): metadata_cls: type[T] require_uniform_decodes: ClassVar[bool] = False + @staticmethod + def _get_mla_dims_from_layers( + vllm_config: "VllmConfig", layer_names: list[str] + ) -> MLADims: + """MLA dims as the runtime attention layers see them. + + Reads the instantiated layers from the static forward context rather + than the raw HF config: the NoPE zero-padded-rope shim promotes the + rope dim on the MLAAttention layer (see MLAAttention.__init__ / + nope_zero_rope_pad), so consumers like the chunked-prefill workspace + must size from the promoted (576) geometry, not the config's raw + rope-free 512 -- cp_gather_and_upconvert_fp8_kv_cache asserts the + workspace is 576 wide. Mirrors MLACommonMetadataBuilder. + """ + layer = vllm_config.compilation_config.static_forward_context[layer_names[0]] + return MLADims( + q_lora_rank=layer.q_lora_rank, + kv_lora_rank=layer.kv_lora_rank, + qk_nope_head_dim=layer.qk_nope_head_dim, + qk_rope_head_dim=layer.qk_rope_head_dim, + v_head_dim=layer.v_head_dim, + ) + def __init__( self, kv_cache_spec: "AttentionSpec", @@ -127,7 +150,7 @@ def __init__( self.vllm_config = vllm_config self.device = device self.model_config = vllm_config.model_config - self.mla_dims = get_mla_dims(self.model_config) + self.mla_dims = self._get_mla_dims_from_layers(vllm_config, layer_names) self.topk_tokens: int = vllm_config.model_config.hf_config.index_topk self.req_id_per_token_buffer = torch.empty( (vllm_config.scheduler_config.max_num_batched_tokens,), From 351c919f3e766368a361666b945e242ab0eca69a Mon Sep 17 00:00:00 2001 From: Leoyzen Date: Sun, 6 Sep 2026 03:11:44 +0800 Subject: [PATCH 5/5] [Test] NoPE-512 fp8_ds_mla SM90 envelope: selection matrix, cache-write, kernel-equivalence, ragged batch (fp8-ds-mla-nope-sm90 5.1-5.4) - Backend-selection matrix (CPU-eligible): NoPE-512 accepted only for quantized DS-MLA formats; bf16/auto/plain-fp8 keep flowing to FlashInfer/TRITON; 576 unchanged for every dtype. SM90 priority tests pin (512, fp8_ds_mla) -> FlashMLA sparse ahead of FlashInfer SM90, plain fp8/bf16 routings unchanged (skips where the CUDA platform module is unavailable, e.g. non-CUDA hosts). - Cache-write: shim's zero k_pe -> concat_and_cache_mla (fp8_ds_mla) passes the csrc asserts unchanged; rope bytes [640:768] of every 656B row are bf16 zero; NoPE bytes match per-128-block fp8 quantization reference (GPU, auto-skipped without CUDA). - Kernel-equivalence: padded-envelope fp8 vs bf16 NoPE-512 reference within fp8 block-quant noise, with a garbage-rope control proving the test has teeth (GPU, auto-skipped without CUDA/FlashMLA). - Ragged mixed-batch: kpool>1 valid counts incl. 0 and 1 with -1 tails; outputs finite, all-invalid rows neutralize to (0, -inf) per the DCP merge contract (GPU, auto-skipped without CUDA). Refs openspec change fp8-ds-mla-nope-sm90 Signed-off-by: Leoyzen --- ...st_flashmla_nope_sm90_backend_selection.py | 132 +++++++ .../test_flashmla_nope_sm90_fp8_ds_mla.py | 353 ++++++++++++++++++ 2 files changed, 485 insertions(+) create mode 100644 tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py create mode 100644 tests/v1/attention/test_flashmla_nope_sm90_fp8_ds_mla.py diff --git a/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py b/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py new file mode 100644 index 000000000000..b8c1c99c0afb --- /dev/null +++ b/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Backend-selection matrix for the rope-free NoPE zero-padded envelope on +SM90 (fp8_ds_mla NoPE-512 -> FLASHMLA_SPARSE). + +These tests are CPU-eligible: they exercise the static backend acceptance +gates and the SM90 priority ordering without touching GPU kernels. On hosts +without the compiled CUDA extension, the priority test skips. +""" + +import pytest +import torch + +from vllm.platforms.interface import DeviceCapability +from vllm.v1.attention.backends.mla.flashmla_sparse import ( + QUANTIZED_DS_MLA_CACHE_FORMATS, + FlashMLASparseBackend, +) + +SM90 = DeviceCapability(major=9, minor=0) + + +def test_supported_head_sizes_include_nope_512(): + assert FlashMLASparseBackend.get_supported_head_sizes() == [576, 512] + + +@pytest.mark.parametrize( + "head_size,kv_cache_dtype,accepted", + [ + # NoPE-512 rides the zero-padded envelope only for quantized DS-MLA. + (512, "fp8_ds_mla", True), + # nvfp4_ds_mla is SM100-only, so it is rejected on SM90 regardless. + (512, "nvfp4_ds_mla", False), + # bf16/auto and plain fp8 NoPE-512 must keep flowing to + # FlashInfer/TRITON backends (existing traffic unchanged). + (512, "auto", False), + (512, "bfloat16", False), + (512, "fp8", False), + (512, "fp8_e4m3", False), + (512, None, False), + # 576 traffic (DeepSeek family) is unchanged for every dtype. + (576, "auto", True), + (576, "bfloat16", True), + (576, "fp8", True), + (576, "fp8_ds_mla", True), + ], +) +def test_supports_combination_nope_gate(head_size, kv_cache_dtype, accepted): + reason = FlashMLASparseBackend.supports_combination( + head_size, + torch.bfloat16, + kv_cache_dtype, + 64, + use_mla=True, + has_sink=False, + use_sparse=True, + use_mm_prefix=False, + device_capability=SM90, + ) + if accepted: + assert reason is None, reason + else: + assert reason is not None + + +def test_nope_gate_only_quantized_ds_mla_formats(): + assert frozenset({"fp8_ds_mla", "nvfp4_ds_mla"}) == QUANTIZED_DS_MLA_CACHE_FORMATS + + +def test_supports_combination_nope_gate_sm100_nvfp4(): + """On SM100 the nvfp4_ds_mla NoPE-512 envelope is accepted (the SM90 + rejection above comes from the SM100-only check, not the NoPE gate).""" + reason = FlashMLASparseBackend.supports_combination( + 512, + torch.bfloat16, + "nvfp4_ds_mla", + 64, + use_mla=True, + has_sink=False, + use_sparse=True, + use_mm_prefix=False, + device_capability=DeviceCapability(major=10, minor=0), + ) + assert reason is None, reason + + +def _sparse_order(kv_cache_dtype, head_size): + try: + from vllm.platforms.cuda import _get_backend_priorities + except ImportError as e: + pytest.skip(f"CUDA platform unavailable on this host: {e}") + + priorities = _get_backend_priorities( + use_mla=True, + device_capability=SM90, + num_heads=32, + kv_cache_dtype=kv_cache_dtype, + head_size=head_size, + ) + sparse = { + "FLASH_ATTN_MLA_SPARSE", + "FLASHMLA_SPARSE", + "FLASHINFER_MLA_SPARSE_SM90", + } + return [b.name for b in priorities if b.name in sparse] + + +def test_sm90_nope_512_fp8_ds_mla_prefers_flashmla(): + order = _sparse_order("fp8_ds_mla", 512) + assert order[0] == "FLASHMLA_SPARSE" + assert order[1] == "FLASHINFER_MLA_SPARSE_SM90" + + +def test_sm90_nope_512_plain_fp8_keeps_flashinfer(): + order = _sparse_order("fp8", 512) + assert order[0] == "FLASHINFER_MLA_SPARSE_SM90" + assert "FLASHMLA_SPARSE" in order + + +def test_sm90_nope_512_bf16_unchanged(): + order = _sparse_order("auto", 512) + assert order[0] == "FLASHINFER_MLA_SPARSE_SM90" + + +def test_sm90_576_unchanged_for_all_dtypes(): + for kv in ("auto", "fp8", "fp8_ds_mla"): + order = _sparse_order(kv, 576) + assert order == [ + "FLASH_ATTN_MLA_SPARSE", + "FLASHMLA_SPARSE", + "FLASHINFER_MLA_SPARSE_SM90", + ], kv diff --git a/tests/v1/attention/test_flashmla_nope_sm90_fp8_ds_mla.py b/tests/v1/attention/test_flashmla_nope_sm90_fp8_ds_mla.py new file mode 100644 index 000000000000..69e3e698dbcd --- /dev/null +++ b/tests/v1/attention/test_flashmla_nope_sm90_fp8_ds_mla.py @@ -0,0 +1,353 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GPU tests for the rope-free NoPE zero-padded 576/656B DS-MLA envelope on +SM90 (FLASHMLA_SPARSE serving head_size=512 via fp8_ds_mla). + +See openspec change fp8-ds-mla-nope-sm90: +- Cache-write contract: rope bytes [640:768] of every 656B row are bf16 zero + and the NoPE bytes carry per-128-block fp8 quantization. +- Kernel equivalence: padded-envelope fp8 matches a bf16 NoPE-512 reference + within fp8 block-quant noise; a garbage-rope control diverges. +- Ragged mixed-batch rows (kpool>1 valid counts incl. 0 and 1, -1 tails) + stay finite and all-invalid rows neutralize to (0, -inf). +""" + +from types import MethodType, SimpleNamespace + +import pytest +import torch + +from vllm.model_executor.layers.attention.mla_attention import ( + _NOPE_ZERO_ROPE_PAD_DIM, + needs_nope_zero_rope_pad, + nope_zero_rope_pad, +) +from vllm.platforms import current_platform +from vllm.v1.attention.backends.mla.flashmla_sparse import ( + FlashMLASparseImpl, + FlashMLASparseMetadata, +) + +if not current_platform.is_cuda(): + pytest.skip( + "fp8_ds_mla NoPE SM90 GPU tests need CUDA (packed-cache ops and the " + "FlashMLA kernels); CPU-eligible selection tests live in " + "test_flashmla_nope_sm90_backend_selection.py", + allow_module_level=True, + ) + + +def test_nope_zero_rope_pad_helper_shapes_and_zeros(): + """The shim provides the exact padded-envelope contract: zero q_pe + [T, H, 64] appended to q and a persistent bf16-zero k_pe [T, 1, 64].""" + device = torch.device("cuda") + torch.manual_seed(0) + q = torch.randn(7, 4, 256, dtype=torch.bfloat16, device=device) + kv_c = torch.randn(7, 512, dtype=torch.bfloat16, device=device) + k_pe_empty = torch.empty(7, 1, 0, dtype=torch.bfloat16, device=device) + + q_padded, _, k_pe = nope_zero_rope_pad(q, kv_c, k_pe_empty) + assert q_padded.shape == (7, 4, 256 + _NOPE_ZERO_ROPE_PAD_DIM) + assert torch.equal(q_padded[..., :256], q) + assert q_padded[..., 256:].abs().max().item() == 0.0 + assert k_pe.shape == (7, 1, _NOPE_ZERO_ROPE_PAD_DIM) + assert k_pe.dtype == torch.bfloat16 + assert k_pe.abs().max().item() == 0.0 + + # The k_pe buffer is persistent: a second call reuses the same storage + # (CUDA-graph capture stability) and keeps it zero. + k_pe2 = nope_zero_rope_pad(q, kv_c, k_pe_empty)[2] + assert k_pe2.data_ptr() == k_pe.data_ptr() + assert k_pe2.abs().max().item() == 0.0 + + +def test_needs_nope_zero_rope_pad_gate(): + assert needs_nope_zero_rope_pad(0, "fp8_ds_mla") + assert needs_nope_zero_rope_pad(0, "nvfp4_ds_mla") + # bf16/auto NoPE and rope>0 models never activate the shim. + assert not needs_nope_zero_rope_pad(0, "auto") + assert not needs_nope_zero_rope_pad(0, "bfloat16") + assert not needs_nope_zero_rope_pad(0, "fp8") + assert not needs_nope_zero_rope_pad(0, None) + assert not needs_nope_zero_rope_pad(64, "fp8_ds_mla") + + +def _block_quant_fp8(x: torch.Tensor, group_size: int = 128): + """Per-128-element block fp8 quantization reference (fp8_ds_mla layout).""" + rows, cols = x.shape + assert cols % group_size == 0 + groups = cols // group_size + xg = x.float().view(rows, groups, group_size) + amax = xg.abs().amax(dim=-1, keepdim=True).clamp(min=1e-12) + finfo = torch.finfo(torch.float8_e4m3fn) + scale = (finfo.max / amax).float() + q = (xg * scale).clamp(min=finfo.min, max=finfo.max).to(torch.float8_e4m3fn) + return q.view(rows, cols), scale.view(rows, groups) + + +def test_nope_fp8_ds_mla_cache_write_zero_rope_bytes(): + """The shim writes bf16 zeros into rope bytes [640:768] of every 656B row + and the packed NoPE bytes carry per-128-block fp8 quantization.""" + from vllm import _custom_ops as ops + + device = torch.device("cuda") + torch.manual_seed(0) + num_tokens, block_size = 12, 4 + kv_lora_rank = 512 + kv_c = torch.randn(num_tokens, kv_lora_rank, dtype=torch.bfloat16, device=device) + + # Shim: zero q_pe/k_pe ahead of the concat path. + _, _, k_pe = nope_zero_rope_pad( + torch.empty(num_tokens, 1, 0, dtype=torch.bfloat16, device=device), + kv_c, + torch.empty(num_tokens, 1, 0, dtype=torch.bfloat16, device=device), + ) + assert k_pe.shape == (num_tokens, 1, _NOPE_ZERO_ROPE_PAD_DIM) + assert k_pe.dtype == torch.bfloat16 + assert k_pe.abs().max().item() == 0.0 + + kv_cache = torch.zeros( + (num_tokens + block_size - 1) // block_size, + block_size, + 656, + dtype=torch.uint8, + device=device, + ) + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) + scale = torch.ones(1, dtype=torch.float32, device=device) + + # concat_and_cache_mla's fp8_ds_mla asserts (kv_lora_rank == 512, + # pe_dim == 64, 656B row, 2-byte elements) must pass unchanged. + ops.concat_and_cache_mla( + kv_c, + k_pe.squeeze(1), + kv_cache, + slot_mapping, + kv_cache_dtype="fp8_ds_mla", + scale=scale, + ) + + rows = kv_cache.view(-1, 656)[:num_tokens] + # RoPE bytes [640:768] are 64 bf16 zeros. + rope = rows[:, 640:768].view(torch.bfloat16) + assert rope.abs().max().item() == 0.0 + # NoPE bytes [0:512] decode to the block-quantized reference. + q_ref, scale_ref = _block_quant_fp8(kv_c.cpu()) + packed = rows[:, :512].view(torch.float8_e4m3fn).cpu() + torch.testing.assert_close(packed.float(), q_ref.float(), rtol=0, atol=0) + # Scale bytes [512:528]: 4 fp32 scales per token. + scales = rows[:, 512:528].view(torch.float32).cpu() + torch.testing.assert_close(scales, scale_ref, rtol=1e-6, atol=1e-8) + + +def test_nope_padded_envelope_matches_bf16_reference(): + """Zero-padded-rope fp8 sparse attention matches a bf16 NoPE-512 + reference within fp8 block-quant noise (rel-err ~2.5e-2 observed on the + H20 probe); a garbage-rope control must diverge, proving the test has + teeth.""" + import vllm.v1.attention.ops.flashmla as fm + + ok, reason = fm.is_flashmla_sparse_supported() + if not ok: + pytest.skip(reason) + + device = torch.device("cuda") + torch.manual_seed(0) + num_tokens = 8 + num_heads = 64 + kv_lora_rank = 512 + topk = 128 + page_block_size = 64 + + kv_c = torch.randn(num_tokens, kv_lora_rank, dtype=torch.bfloat16, device=device) + q_latent = torch.randn( + num_tokens, num_heads, kv_lora_rank, dtype=torch.bfloat16, device=device + ) + softmax_scale = (kv_lora_rank + _NOPE_ZERO_ROPE_PAD_DIM) ** -0.5 + + # Build the padded 656B envelope cache via the shim's zero k_pe. + k_pe = nope_zero_rope_pad( + torch.empty(num_tokens, 1, 0, dtype=torch.bfloat16, device=device), + kv_c, + torch.empty(num_tokens, 1, 0, dtype=torch.bfloat16, device=device), + )[2] + kv_cache = torch.zeros(1, page_block_size, 1, 656, dtype=torch.uint8, device=device) + assert num_tokens <= page_block_size + from vllm import _custom_ops as ops + + ops.concat_and_cache_mla( + kv_c, + k_pe.squeeze(1), + kv_cache, + torch.arange(num_tokens, dtype=torch.int64, device=device), + kv_cache_dtype="fp8_ds_mla", + scale=torch.ones(1, dtype=torch.float32, device=device), + ) + + # Padded query: q_latent + zero q_pe -> [T, H, 576]. + q_padded = torch.cat( + [ + q_latent, + torch.zeros( + num_tokens, + num_heads, + _NOPE_ZERO_ROPE_PAD_DIM, + dtype=torch.bfloat16, + device=device, + ), + ], + dim=-1, + ) + + cache_seqlens = torch.full( + (num_tokens,), num_tokens, dtype=torch.int32, device=device + ) + tile_md, num_splits = fm.get_mla_metadata( + cache_seqlens, + num_heads, + 1, + num_heads_q=num_heads, + topk=topk, + is_fp8_kvcache=True, + ) + indices = torch.arange(topk, dtype=torch.int32, device=device) + indices = indices.clamp(max=num_tokens - 1).expand(num_tokens, topk).contiguous() + + def run(cache): + out, _ = fm.flash_mla_with_kvcache( + q=q_padded.unsqueeze(0), + k_cache=cache.squeeze(2).unsqueeze(-2), + block_table=torch.zeros((num_tokens, 1), dtype=torch.int32, device=device), + head_dim_v=512, + cache_seqlens=cache_seqlens, + tile_scheduler_metadata=tile_md, + num_splits=num_splits, + is_fp8_kvcache=True, + indices=indices.unsqueeze(0), + softmax_scale=softmax_scale, + ) + return out.squeeze(0) # (T, H, 512) + + out = run(kv_cache) + + # bf16 NoPE-512 reference: softmax(q . k) @ k over the latent dim only -- + # the zero rope part contributes exactly nothing. + ref_scores = ( + torch.einsum("xhd,wd->xhw", q_latent.float(), kv_c.float()) * softmax_scale + ) + ref_probs = torch.softmax(ref_scores, dim=1) + ref_out = torch.einsum("xhw,wd->xhd", ref_probs, kv_c.float()) + + rel_err = ((out.float() - ref_out).norm() / ref_out.norm().clamp(min=1e-6)).item() + # fp8 block-quant noise floor (~2.5e-2 observed on the probe). + assert rel_err < 0.08, f"padded-envelope fp8 rel-err {rel_err:.3e}" + + # Garbage-rope control: nonzero rope bytes must diverge from the NoPE + # reference -- the padding contract is what keeps the fp8 path + # NoPE-correct, not luck. + garbage_pe = torch.randn( + num_tokens, _NOPE_ZERO_ROPE_PAD_DIM, dtype=torch.bfloat16, device=device + ) + kv_cache_garbage = kv_cache.clone() + ops.concat_and_cache_mla( + kv_c, + garbage_pe, + kv_cache_garbage, + torch.arange(num_tokens, dtype=torch.int64, device=device), + kv_cache_dtype="fp8_ds_mla", + scale=torch.ones(1, dtype=torch.float32, device=device), + ) + rel_err_garbage = ( + (run(kv_cache_garbage).float() - ref_out).norm() + / ref_out.norm().clamp(min=1e-6) + ).item() + assert rel_err_garbage > 2 * rel_err + 1e-3, ( + f"garbage-rope control did not diverge ({rel_err_garbage:.3e} vs {rel_err:.3e})" + ) + + +def test_fp8_mixed_batch_ragged_valid_counts(): + """Ragged top-k rows (valid counts 0, 1, partial with -1 tails, as the + kpool>1 indexer produces) must give finite outputs and neutralize + all-invalid rows to (0, -inf) per the DCP merge contract.""" + import vllm.v1.attention.backends.mla.flashmla_sparse as fms + + num_tokens, num_heads, head_dim = 5, 2, 3 + device = torch.device("cuda") + q = torch.randn( + num_tokens, num_heads, head_dim, dtype=torch.bfloat16, device=device + ) + # Ragged rows: full, 0-valid, 1-valid, partial with -1 tail, all -1. + local_indices = torch.tensor( + [ + [0, 1, 2, 3], + [-1, -1, -1, -1], + [5, -1, -1, -1], + [7, 8, -1, -1], + [-1, -1, -1, -1], + ], + dtype=torch.int32, + device=device, + ) + + def convert_indices(*args, **kwargs): # noqa: ARG001 + return kwargs["output"].copy_(local_indices) + + orig = fms.triton_filter_and_convert_dcp_index + fms.triton_filter_and_convert_dcp_index = convert_indices + try: + + def run_kernel(**kwargs): + out = torch.full((1, num_tokens, num_heads, 1), float("nan"), device=device) + lse = torch.full((1, num_heads, num_tokens), float("nan"), device=device) + for token_id in (0, 2, 3): # rows with local candidates + out[0, token_id] = float(token_id + 1) + lse[0, :, token_id] = float(token_id + 1) + return out, lse + + metadata = SimpleNamespace( + fp8_extra_metadata=FlashMLASparseMetadata.FP8KernelMetadata( + scheduler_metadata=object(), # type: ignore[arg-type] + dummy_block_table=torch.empty(1, 1, dtype=torch.int32, device=device), + cache_lens=torch.empty(1, dtype=torch.int32, device=device), + ), + req_id_per_token=torch.empty(num_tokens, dtype=torch.int32, device=device), + block_table=torch.empty(1, 1, dtype=torch.int32, device=device), + block_size=64, + cp_kv_cache_interleave_size=1, + fp8_use_mixed_batch=True, + physical_topk_indices=torch.empty_like(local_indices), + physical_topk_valid_counts=torch.empty( + num_tokens, dtype=torch.int32, device=device + ), + physical_topk_is_valid=False, + ) + impl = SimpleNamespace( + kv_cache_dtype="fp8_ds_mla", + topk_indices_buffer=local_indices, + dcp_world_size=2, + dcp_rank=0, + need_to_return_lse_for_decode=True, + _fp8_flash_mla_kernel=run_kernel, + ) + impl._forward_fp8_kv_mixed_batch = MethodType( + FlashMLASparseImpl._forward_fp8_kv_mixed_batch, impl + ) + + out, lse = FlashMLASparseImpl.forward_mqa( + impl, q, torch.empty(0, device=device), metadata, None + ) + finally: + fms.triton_filter_and_convert_dcp_index = orig + + assert out is not None and lse is not None + assert not out.isnan().any() + assert not lse.isnan().any() + # All-invalid rows neutralize to the DCP merge identity (0, -inf). + for token_id in (1, 4): + assert torch.equal(out[token_id], torch.zeros_like(out[token_id])) + assert torch.isneginf(lse[token_id]).all() + # Rows with local candidates keep their kernel values. + for token_id in (0, 2, 3): + assert torch.equal(out[token_id], torch.full_like(out[token_id], token_id + 1)) + assert torch.equal(lse[token_id], torch.full_like(lse[token_id], token_id + 1))