| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Behavior checks for FlashInfer SM120 sparse MLA backend selection.""" |
| 4 | |
| 5 | from types import SimpleNamespace |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | from vllm.config import set_current_vllm_config |
| 10 | from vllm.models.deepseek_v4.nvidia.flashinfer_sparse import ( |
| 11 | _required_sm120_sparse_topk, |
| 12 | ) |
| 13 | from vllm.platforms.interface import DeviceCapability |
| 14 | from vllm.utils import flashinfer as fi_utils |
| 15 | from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( |
| 16 | FlashInferMLASparseSM120Backend, |
| 17 | ) |
| 18 | from vllm.v1.attention.backends.registry import AttentionBackendEnum |
| 19 | |
| 20 | |
| 21 | def _fake_vllm_config(model_type: str) -> SimpleNamespace: |
| 22 | return SimpleNamespace( |
| 23 | model_config=SimpleNamespace( |
| 24 | hf_text_config=SimpleNamespace(model_type=model_type, index_topk=2048), |
| 25 | ), |
| 26 | ) |
| 27 | |
| 28 | |
| 29 | def test_sm120_backend_uses_dedicated_backend_name() -> None: |
| 30 | assert FlashInferMLASparseSM120Backend.get_name() == "FLASHINFER_MLA_SPARSE_SM120" |
| 31 | assert ( |
| 32 | AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM120.get_class() |
| 33 | is FlashInferMLASparseSM120Backend |
| 34 | ) |
| 35 | |
| 36 | |
| 37 | def test_sm120_backend_uses_sparse_mqa_for_prefill() -> None: |
| 38 | impl_cls = FlashInferMLASparseSM120Backend.get_impl_cls() |
| 39 | |
| 40 | assert impl_cls.is_sparse |
| 41 | assert not impl_cls.supports_dense_mha_prefill |
| 42 | |
| 43 | |
| 44 | def test_v32_glm_sm120_backend_accepts_glm_block_size( |
| 45 | monkeypatch, |
| 46 | ) -> None: |
| 47 | monkeypatch.setattr(fi_utils, "has_flashinfer_sparse_mla_sm120", lambda: True) |
| 48 | |
| 49 | with set_current_vllm_config(_fake_vllm_config("glm4_moe")): |
| 50 | invalid_reasons = FlashInferMLASparseSM120Backend.validate_configuration( |
| 51 | head_size=576, |
| 52 | dtype=torch.bfloat16, |
| 53 | kv_cache_dtype="fp8", |
| 54 | block_size=256, |
| 55 | use_mla=True, |
| 56 | has_sink=False, |
| 57 | use_sparse=True, |
| 58 | use_mm_prefix=False, |
| 59 | use_per_head_quant_scales=False, |
| 60 | device_capability=DeviceCapability(12, 0), |
| 61 | attn_type="decoder", |
| 62 | ) |
| 63 | |
| 64 | assert invalid_reasons == [] |
| 65 | |
| 66 | |
| 67 | def test_sm120_dsv4_capability_checks_exact_dispatch_shape(monkeypatch) -> None: |
| 68 | fake_module = SimpleNamespace( |
| 69 | _DECODE_DSV4_DISPATCH=frozenset({(32, 128), (32, 192)}) |
| 70 | ) |
| 71 | monkeypatch.setattr(fi_utils, "has_flashinfer_sparse_mla_sm120", lambda: True) |
| 72 | monkeypatch.setattr(fi_utils, "_get_submodule", lambda _name: fake_module) |
| 73 | fi_utils.has_flashinfer_sparse_mla_sm120_config.cache_clear() |
| 74 | |
| 75 | assert fi_utils.has_flashinfer_sparse_mla_sm120_config(32, 128) |
| 76 | assert fi_utils.has_flashinfer_sparse_mla_sm120_config(32, 192) |
| 77 | assert not fi_utils.has_flashinfer_sparse_mla_sm120_config(32, 256) |
| 78 | assert not fi_utils.has_flashinfer_sparse_mla_sm120_config(16, 192) |
| 79 | |
| 80 | fi_utils.has_flashinfer_sparse_mla_sm120_config.cache_clear() |
| 81 | |
| 82 | |
| 83 | def test_sm120_dsv4_required_topk_tracks_dspark_width() -> None: |
| 84 | causal = SimpleNamespace( |
| 85 | attention_config=SimpleNamespace(use_non_causal=False), |
| 86 | speculative_config=SimpleNamespace(num_speculative_tokens=5), |
| 87 | ) |
| 88 | dspark = SimpleNamespace( |
| 89 | attention_config=SimpleNamespace(use_non_causal=True), |
| 90 | speculative_config=SimpleNamespace(num_speculative_tokens=5), |
| 91 | ) |
| 92 | |
| 93 | assert _required_sm120_sparse_topk(causal, 128) == 128 |
| 94 | assert _required_sm120_sparse_topk(dspark, 128) == 192 |