[Attention][MLA] Guard NoPE-512 SM90 selection; fix FA3 sparse zero-rope MQA

cmake/external_projects/vllm_flash_attn.cmake

@@ -39,7 +39,7 @@ else()
   FetchContent_Declare(
           vllm-flash-attn
           GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG 06bdd47c0d0383daf6a2ff0c418faff9c6da16e5
+          GIT_TAG 9cd61de38763d712bb6ce56e2a02cc2bf718c89f
           GIT_PROGRESS TRUE
           # Don't share the vllm-flash-attn build between build types
           BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn

tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py

@@ -0,0 +1,176 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU tests for the FlashMLA sparse backend NoPE-512 (GLM5Next) gates.
+
+GLM5Next is rope-free MLA: qk_nope_head_dim=256, qk_rope_head_dim=0,
+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;
+* adding 512 must NOT reorder the SM100 priority list for bf16, which is why
+  the bf16 arm is restricted to SM90.
+"""
+
+from types import SimpleNamespace
+
+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)
+SM100 = DeviceCapability(major=10, minor=0)
+
+BF16_CACHE_DTYPES = (None, "auto", "bfloat16", "float16")
+QUANTIZED_CACHE_DTYPES = ("fp8_ds_mla", "nvfp4_ds_mla")
+
+
+def _supports_combination(head_size, kv_cache_dtype, capability):
+    return 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=capability,
+    )
+
+
+@pytest.fixture
+def rope_free_model(monkeypatch):
+    """Patch get_current_vllm_config so the model looks like GLM5Next NoPE."""
+    import vllm.config as cfg
+
+    monkeypatch.setattr(
+        cfg,
+        "get_current_vllm_config",
+        lambda: SimpleNamespace(
+            model_config=SimpleNamespace(
+                hf_text_config=SimpleNamespace(qk_rope_head_dim=0)
+            )
+        ),
+    )
+
+
+@pytest.fixture
+def rope_carrying_model(monkeypatch):
+    """Patch get_current_vllm_config with a rope-carrying (576-style) model."""
+    import vllm.config as cfg
+
+    monkeypatch.setattr(
+        cfg,
+        "get_current_vllm_config",
+        lambda: SimpleNamespace(
+            model_config=SimpleNamespace(
+                hf_text_config=SimpleNamespace(qk_rope_head_dim=64)
+            )
+        ),
+    )
+
+
+def test_supported_head_sizes_include_512():
+    assert FlashMLASparseBackend.get_supported_head_sizes() == [576, 512]
+
+
+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):
+    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):
+    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):
+    # 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):
+    reason = _supports_combination(512, kv_cache_dtype, SM90)
+    assert reason is not None
+
+
+@pytest.mark.parametrize("capability", [SM90, SM100], ids=["sm90", "sm100"])
+@pytest.mark.parametrize(
+    "kv_cache_dtype",
+    ("auto", "bfloat16", "fp8", "fp8_ds_mla", "nvfp4_ds_mla"),
+)
+def test_576_unchanged_for_all_dtypes(kv_cache_dtype, capability):
+    # DeepSeek 576 must be untouched by the NoPE-512 gates, except for the
+    # pre-existing SM100-only nvfp4 rule.
+    reason = _supports_combination(576, kv_cache_dtype, capability)
+    if kv_cache_dtype == "nvfp4_ds_mla" and capability.major != 10:
+        assert reason is not None and "SM100" in reason
+    else:
+        assert reason is None, reason
+
+
+def _sparse_order(kv_cache_dtype, head_size, num_heads=32, capability=SM100):
+    try:
+        from vllm.platforms.cuda import _get_backend_priorities
+    except ImportError as e:  # pragma: no cover - depends on host
+        pytest.skip(f"CUDA platform unavailable on this host: {e}")
+
+    priorities = _get_backend_priorities(
+        use_mla=True,
+        device_capability=capability,
+        num_heads=num_heads,
+        kv_cache_dtype=kv_cache_dtype,
+        head_size=head_size,
+    )
+    sparse = {
+        "FLASH_ATTN_MLA_SPARSE",
+        "FLASHMLA_SPARSE",
+        "FLASHINFER_MLA_SPARSE",
+        "FLASHINFER_MLA_SPARSE_SM90",
+    }
+    return [b.name for b in priorities if b.name in sparse]
+
+
+@pytest.mark.parametrize("num_heads", [32, 64], ids=["tp2", "tp1"])
+def test_sm100_bf16_512_priority_unchanged(num_heads):
+    """Adding head_size 512 must not reorder the SM100 bf16 sparse list.
+
+    On SM100 with >16 heads the hardcoded bf16 order is
+    [FLASHMLA_SPARSE, FLASHINFER_MLA_SPARSE] (FlashMLA first). head_size 512
+    must leave this untouched (identical to the 576 order); the guard that
+    keeps bf16-512 from *serving* through FlashMLA on SM100 lives in
+    supports_combination (gated to SM90), tested separately.
+    """
+    order512 = _sparse_order("auto", 512, num_heads=num_heads)
+    order576 = _sparse_order("auto", 576, num_heads=num_heads)
+    assert order512 == order576, order512
+    assert order512[0] == "FLASHMLA_SPARSE", order512
+
+
+def test_sm100_576_priority_unchanged():
+    order = _sparse_order("auto", 576, num_heads=32)
+    assert order[0] == "FLASHMLA_SPARSE", order

vllm/v1/attention/backends/mla/flashattn_mla_sparse.py

@@ -230,8 +230,16 @@ def forward_mqa(
         cu_seqlens_q = torch.arange(
             0, num_actual_toks + 1, dtype=torch.int32, device=q_rope.device
         )
-        k_cache = kv_rows[:, self.kv_lora_rank :].unsqueeze(1).unsqueeze(1)
         v_cache = kv_rows[:, : self.kv_lora_rank].unsqueeze(1).unsqueeze(1)
+        if self.qk_rope_head_dim == 0:
+            # FA3's QV path requires the 64-wide Q/K specialization. For NoPE
+            # MLA, zero Q preserves QV-only attention scores while providing a
+            # valid TMA shape.
+            _FA3_QV_HEAD_DIM = 64
+            q_rope = q_nope.new_zeros(*q_nope.shape[:-1], _FA3_QV_HEAD_DIM)
+            k_cache = v_cache[..., :_FA3_QV_HEAD_DIM]
+        else:
+            k_cache = kv_rows[:, self.kv_lora_rank :].unsqueeze(1).unsqueeze(1)
 
         out = flash_attn_varlen_func(
             q=q_rope,

vllm/v1/attention/backends/mla/flashmla_sparse.py

@@ -135,8 +135,8 @@ 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]
+        # DeepSeek V3.2: 512 NoPE + 64 RoPE = 576; GLM5Next NoPE: 512.
+        return [576, 512]
 
     @classmethod
     def is_mla(cls) -> bool:
@@ -163,6 +163,37 @@ def supports_combination(
         use_mm_prefix: bool,
         device_capability: DeviceCapability,
     ) -> 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
+            # FlashInfer/TRITON.
+            if (
+                kv_cache_dtype in (None, "auto", "bfloat16", "float16")
+                and device_capability.major == 9
+            ):
+                # Direct bf16 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
+
+                vllm_config = get_current_vllm_config()
+                if vllm_config.model_config is not None:
+                    hf_text_config = vllm_config.model_config.hf_text_config
+                    if getattr(hf_text_config, "qk_rope_head_dim", 64) != 0:
+                        return (
+                            "FLASHMLA_SPARSE supports head_size 512 only for "
+                            "rope-free (NoPE) models"
+                        )
+            else:
+                return (
+                    "FLASHMLA_SPARSE supports head_size 512 only with bf16 "
+                    "kv-cache on SM90 (NoPE), got "
+                    f"kv_cache_dtype={kv_cache_dtype}, "
+                    f"capability={device_capability}"
+                )
         if kv_cache_dtype == "nvfp4_ds_mla" and device_capability.major != 10:
             return (
                 f"FLASHMLA_SPARSE only supports the {kv_cache_dtype} kv-cache "
@@ -996,7 +1027,12 @@ def forward_mqa(
         if isinstance(q, tuple):
             ql_nope, q_pe = q
             q = self.q_concat_buffer[: ql_nope.shape[0]]
-            ops.concat_mla_q(ql_nope, q_pe, q)
+            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)
+            else:
+                ops.concat_mla_q(ql_nope, q_pe, q)
         else:
             actual_num_heads = q.shape[1]
 

vllm/v1/attention/backends/mla/prefill/flash_attn.py

@@ -327,6 +327,13 @@ def supports_mla_dimensions(cls, mla_dimensions: MLADimensions) -> bool:
                 qk_rope_head_dim=64,
                 v_head_dim=128,
             ),
+            # GLM5Next NoPE: qk_head_dim = 256 + 0, same kernel dims as the
+            # (192, 64, 256) layout above.
+            MLADimensions(
+                qk_nope_head_dim=256,
+                qk_rope_head_dim=0,
+                v_head_dim=256,
+            ),
         ]
 
     def __init__(