From 91d7324cb19d301c72d849e457221ee8dd645024 Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Tue, 22 Sep 2026 17:03:59 +0800 Subject: [PATCH] [perf] wire FA and FlashMLA for sm90 GLM5Next NoPE SparseMLA (#55385) Signed-off-by: JaredforReal Co-authored-by: Kimi Code Co-authored-by: Leoyzen Co-authored-by: Lucas Wilkinson --- cmake/external_projects/vllm_flash_attn.cmake | 2 +- ...st_flashmla_nope_sm90_backend_selection.py | 176 ++++++++++++++++++ .../backends/mla/flashattn_mla_sparse.py | 10 +- .../attention/backends/mla/flashmla_sparse.py | 42 ++++- 4 files changed, 225 insertions(+), 5 deletions(-) create mode 100644 tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index 1b3be244de..2a26023f74 100644 --- a/cmake/external_projects/vllm_flash_attn.cmake +++ b/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 506341a143fcabd4bb79052a7605ada727d6b3f5 + 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 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 0000000000..db2c1becc4 --- /dev/null +++ b/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 diff --git a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py index 43720fb3ed..2d1e218640 100644 --- a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py @@ -309,8 +309,16 @@ class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]): cu_seqlens_q = torch.arange( 0, q_rope.shape[0] + 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, diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index aa022066f0..2e71f0a35c 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -156,8 +156,8 @@ class FlashMLASparseBackend(AttentionBackend): @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: @@ -184,6 +184,37 @@ class FlashMLASparseBackend(AttentionBackend): 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 " @@ -1332,7 +1363,12 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): *self.workspace_specs ) q = 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]