[Perf][GLM-5.3-Flash] Dense/masked-MHA sparse prefill for the NoPE (256, 0, 256) layout + skip the NoPE K concat

tests/v1/attention/test_mla_backends.py

@@ -25,6 +25,7 @@
 from vllm.model_executor.layers.attention import mla_attention as mla_attention_module
 from vllm.model_executor.layers.attention.mla_attention import (
     MLAAttention,
+    MLACommonBaseImpl,
     QueryLenSupport,
     _DecodeConcatQuantFP8,
     _use_masked_mha,
@@ -166,6 +167,25 @@ def test_glm5_flashinfer_masked_mha_routing(
     )
 
 
+@pytest.mark.parametrize("qk_rope_head_dim", [64, 0], ids=["rope", "nope"])
+def test_concat_k_nope_k_pe_matches_torch_cat(qk_rope_head_dim):
+    """The K concat used by the MLA prefill context loop must equal torch.cat of
+    k_nope with the broadcast k_pe; with no RoPE part it returns k_nope itself
+    instead of allocating and copying."""
+    torch.manual_seed(0)
+    num_tokens, num_heads, qk_nope_head_dim = 5, 4, 256
+    k_nope = torch.randn(num_tokens, num_heads, qk_nope_head_dim, dtype=torch.bfloat16)
+    k_pe = torch.randn(num_tokens, 1, qk_rope_head_dim, dtype=torch.bfloat16)
+    impl = SimpleNamespace(_use_flashinfer_concat_mla_k=False)
+
+    k = MLACommonBaseImpl._concat_k_nope_k_pe(impl, k_nope, k_pe)
+
+    expected = torch.cat([k_nope, k_pe.expand(-1, num_heads, -1)], dim=-1)
+    assert k.shape == (num_tokens, num_heads, qk_nope_head_dim + qk_rope_head_dim)
+    torch.testing.assert_close(k, expected, rtol=0, atol=0)
+    assert (k.data_ptr() == k_nope.data_ptr()) == (qk_rope_head_dim == 0)
+
+
 def test_masked_mha_routing_is_dimension_specific():
     assert _use_masked_mha(
         backend_name="FLASHMLA_SPARSE",

tests/v1/attention/test_mla_prefill_selector.py

@@ -313,6 +313,40 @@ def test_backend_supported_dimension_validation(self):
             )
             assert invalid_reasons == []
 
+    def test_flash_attn_accepts_glm53_flash_nope_dimensions(self):
+        """(256, 0, 256) runs the same FA kernels as GLM-5's (192, 64, 256); a
+        RoPE-carrying 320-wide query does not."""
+        try:
+            from vllm.v1.attention.backends.mla.prefill.flash_attn import (
+                FlashAttnPrefillBackend,
+            )
+        except ImportError:
+            pytest.skip("MLA prefill backend not available")
+            return
+
+        capability = DeviceCapability(major=10, minor=0)
+
+        def validate(qk_rope_head_dim: int) -> list[str]:
+            selector_config = MLAPrefillSelectorConfig(
+                dtype=torch.bfloat16,
+                mla_dimensions=MLADimensions(
+                    qk_nope_head_dim=256,
+                    qk_rope_head_dim=qk_rope_head_dim,
+                    v_head_dim=256,
+                ),
+            )
+            with patch.object(
+                FlashAttnPrefillBackend, "is_available", return_value=True
+            ):
+                return FlashAttnPrefillBackend.validate_configuration(
+                    capability, selector_config
+                )
+
+        assert validate(qk_rope_head_dim=0) == []
+        invalid_reasons = validate(qk_rope_head_dim=64)
+        assert len(invalid_reasons) == 1
+        assert "supported MLA dimensions" in invalid_reasons[0]
+
 
 @pytest.mark.skipif(
     not current_platform.is_cuda_alike(),

tests/v1/attention/test_sparse_mla_backends.py

@@ -32,6 +32,7 @@
     GLOBAL_TOPK_MASK_MAX_BYTES,
     SparseMLACommonImpl,
     SparseMLAPrefillMetadata,
+    _is_masked_mha_available,
     _masked_mha_workspace_fits,
     _topk_mask_shape,
     _use_dense_mha_prefill,
@@ -1165,6 +1166,33 @@ def test_masked_mha_workspace_guards_long_routing_policy(
     )
 
 
+@pytest.mark.parametrize(
+    ("model_dims", "kv_cache_dtype", "fa_version", "expected"),
+    [
+        pytest.param((128, 512, 128, 64, 128), "auto", 4, True, id="deepseek_v32"),
+        pytest.param((64, 512, 192, 64, 256), "auto", 4, True, id="glm5"),
+        pytest.param((64, 512, 256, 0, 256), "auto", 4, True, id="glm53_flash_nope"),
+        pytest.param((64, 512, 256, 0, 256), "fp8", 4, False, id="quantized_kv"),
+        pytest.param((64, 512, 256, 0, 256), "auto", 3, False, id="no_fa4"),
+        pytest.param((64, 512, 256, 64, 256), "auto", 4, False, id="rope_320"),
+        pytest.param((128, 512, 256, 0, 256), "auto", 4, False, id="wrong_heads"),
+    ],
+)
+def test_is_masked_mha_available_model_dims(
+    monkeypatch, model_dims, kv_cache_dtype, fa_version, expected
+):
+    """The allow-list gates masked MHA per exact model geometry: the DeepSeek-V3.2,
+    GLM-5 and NoPE GLM-5.3-Flash layouts on an SM100-family GPU with FA4 and an
+    unquantized KV cache, nothing else."""
+    import vllm.model_executor.layers.attention.sparse_mla_attention as mod
+
+    monkeypatch.setattr(
+        mod.current_platform, "is_device_capability_family", lambda family: True
+    )
+    monkeypatch.setattr(mod, "get_flash_attn_version", lambda **kwargs: fa_version)
+    assert _is_masked_mha_available(*model_dims, kv_cache_dtype) is expected
+
+
 def test_masked_mha_workspace_fits_accounts_for_batch_and_context():
     """Request count and context chunk length are independent multipliers."""
     base = dict(batch_size=2, max_query_len=2048, max_context_chunk_seq_len=2048)
@@ -1198,6 +1226,7 @@ def test_masked_mha_workspace_fits_accounts_for_batch_and_context():
     [
         pytest.param(128, 128, 64, 128, id="deepseek_hd192_v128"),
         pytest.param(64, 192, 64, 256, id="glm5_hd256_v256"),
+        pytest.param(64, 256, 0, 256, id="glm53_flash_nope_hd256_v256"),
     ],
 )
 def test_sparse_backend_prefill_correctness(

vllm/model_executor/layers/attention/mla_attention.py

@@ -2871,6 +2871,10 @@ def _concat_k_nope_k_pe(
         Returns:
             Tensor of shape [..., nope_dim + pe_dim]
         """
+        if k_pe.shape[-1] == 0:
+            # NoPE MLA: nothing to append, so no copy either.
+            return k_nope
+
         k = torch.empty(
             (*k_nope.shape[:-1], k_nope.shape[-1] + k_pe.shape[-1]),
             dtype=k_nope.dtype,

vllm/model_executor/layers/attention/sparse_mla_attention.py

@@ -111,6 +111,8 @@ def _is_masked_mha_available(
     if model_dims not in (
         (128, 512, 128, 64, 128),
         (64, 512, 192, 64, 256),
+        # GLM-5.3-Flash: NoPE, qk_head_dim 256 == the (192, 64, 256) kernel.
+        (64, 512, 256, 0, 256),
     ):
         return False
     qk_head_dim = qk_nope_head_dim + qk_rope_head_dim

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 layout: qk_head_dim 256 + 0 and v_head_dim 256 run
+            # the same kernels as the (192, 64, 256) DeepSeek-V3.2 layout.
+            MLADimensions(
+                qk_nope_head_dim=256,
+                qk_rope_head_dim=0,
+                v_head_dim=256,
+            ),
         ]
 
     def __init__(