@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 176 lines 6.11 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """CPU tests for the FlashMLA sparse backend NoPE-512 (GLM5Next) gates.
4
5 GLM5Next is rope-free MLA: qk_nope_head_dim=256, qk_rope_head_dim=0,
6 v_head_dim=256, kv_lora_rank=512 -> head_size 512. These tests pin the
7 selection contract only (no GPU required):
8
9 * ``get_supported_head_sizes`` advertises 512 so the backend is a candidate;
10 * ``supports_combination`` admits exactly two shapes for 512 -- a quantized
11 DS-MLA cache format (zero-padded 576/656B envelope) and a rope-free bf16
12 cache on SM90 -- and rejects everything else;
13 * adding 512 must NOT reorder the SM100 priority list for bf16, which is why
14 the bf16 arm is restricted to SM90.
15 """
16
17 from types import SimpleNamespace
18
19 import pytest
20 import torch
21
22 from vllm.platforms.interface import DeviceCapability
23 from vllm.v1.attention.backends.mla.flashmla_sparse import (
24 QUANTIZED_DS_MLA_CACHE_FORMATS,
25 FlashMLASparseBackend,
26 )
27
28 SM90 = DeviceCapability(major=9, minor=0)
29 SM100 = DeviceCapability(major=10, minor=0)
30
31 BF16_CACHE_DTYPES = (None, "auto", "bfloat16", "float16")
32 QUANTIZED_CACHE_DTYPES = ("fp8_ds_mla", "nvfp4_ds_mla")
33
34
35 def _supports_combination(head_size, kv_cache_dtype, capability):
36 return FlashMLASparseBackend.supports_combination(
37 head_size,
38 torch.bfloat16,
39 kv_cache_dtype,
40 64,
41 use_mla=True,
42 has_sink=False,
43 use_sparse=True,
44 use_mm_prefix=False,
45 device_capability=capability,
46 )
47
48
49 @pytest.fixture
50 def rope_free_model(monkeypatch):
51 """Patch get_current_vllm_config so the model looks like GLM5Next NoPE."""
52 import vllm.config as cfg
53
54 monkeypatch.setattr(
55 cfg,
56 "get_current_vllm_config",
57 lambda: SimpleNamespace(
58 model_config=SimpleNamespace(
59 hf_text_config=SimpleNamespace(qk_rope_head_dim=0)
60 )
61 ),
62 )
63
64
65 @pytest.fixture
66 def rope_carrying_model(monkeypatch):
67 """Patch get_current_vllm_config with a rope-carrying (576-style) model."""
68 import vllm.config as cfg
69
70 monkeypatch.setattr(
71 cfg,
72 "get_current_vllm_config",
73 lambda: SimpleNamespace(
74 model_config=SimpleNamespace(
75 hf_text_config=SimpleNamespace(qk_rope_head_dim=64)
76 )
77 ),
78 )
79
80
81 def test_supported_head_sizes_include_512():
82 assert FlashMLASparseBackend.get_supported_head_sizes() == [576, 512]
83
84
85 def test_quantized_ds_mla_formats_are_the_envelope_set():
86 assert frozenset({"fp8_ds_mla", "nvfp4_ds_mla"}) == QUANTIZED_DS_MLA_CACHE_FORMATS
87
88
89 @pytest.mark.parametrize("kv_cache_dtype", QUANTIZED_CACHE_DTYPES)
90 def test_nope_512_quantized_rejected_here(kv_cache_dtype):
91 # Quantized DS-MLA NoPE-512 is served by the zero-padded 576/656B
92 # envelope, wired up in a separate change; this backend must reject it.
93 reason = _supports_combination(512, kv_cache_dtype, SM90)
94 assert reason is not None
95
96
97 @pytest.mark.parametrize("kv_cache_dtype", BF16_CACHE_DTYPES)
98 def test_nope_512_bf16_sm90_rope_free_accepted(kv_cache_dtype, rope_free_model):
99 assert _supports_combination(512, kv_cache_dtype, SM90) is None
100
101
102 @pytest.mark.parametrize("kv_cache_dtype", BF16_CACHE_DTYPES)
103 def test_nope_512_bf16_sm90_rope_carrying_rejected(kv_cache_dtype, rope_carrying_model):
104 reason = _supports_combination(512, kv_cache_dtype, SM90)
105 assert reason is not None and "rope-free" in reason
106
107
108 @pytest.mark.parametrize("kv_cache_dtype", BF16_CACHE_DTYPES)
109 def test_nope_512_bf16_sm100_rejected(kv_cache_dtype, rope_free_model):
110 # The SM90 restriction is what keeps the SM100 priority order unchanged.
111 reason = _supports_combination(512, kv_cache_dtype, SM100)
112 assert reason is not None
113
114
115 @pytest.mark.parametrize("kv_cache_dtype", ("fp8", "fp8_e4m3"))
116 def test_nope_512_plain_fp8_rejected(kv_cache_dtype, rope_free_model):
117 reason = _supports_combination(512, kv_cache_dtype, SM90)
118 assert reason is not None
119
120
121 @pytest.mark.parametrize("capability", [SM90, SM100], ids=["sm90", "sm100"])
122 @pytest.mark.parametrize(
123 "kv_cache_dtype",
124 ("auto", "bfloat16", "fp8", "fp8_ds_mla", "nvfp4_ds_mla"),
125 )
126 def test_576_unchanged_for_all_dtypes(kv_cache_dtype, capability):
127 # DeepSeek 576 must be untouched by the NoPE-512 gates, except for the
128 # pre-existing SM100-only nvfp4 rule.
129 reason = _supports_combination(576, kv_cache_dtype, capability)
130 if kv_cache_dtype == "nvfp4_ds_mla" and capability.major != 10:
131 assert reason is not None and "SM100" in reason
132 else:
133 assert reason is None, reason
134
135
136 def _sparse_order(kv_cache_dtype, head_size, num_heads=32, capability=SM100):
137 try:
138 from vllm.platforms.cuda import _get_backend_priorities
139 except ImportError as e: # pragma: no cover - depends on host
140 pytest.skip(f"CUDA platform unavailable on this host: {e}")
141
142 priorities = _get_backend_priorities(
143 use_mla=True,
144 device_capability=capability,
145 num_heads=num_heads,
146 kv_cache_dtype=kv_cache_dtype,
147 head_size=head_size,
148 )
149 sparse = {
150 "FLASH_ATTN_MLA_SPARSE",
151 "FLASHMLA_SPARSE",
152 "FLASHINFER_MLA_SPARSE",
153 "FLASHINFER_MLA_SPARSE_SM90",
154 }
155 return [b.name for b in priorities if b.name in sparse]
156
157
158 @pytest.mark.parametrize("num_heads", [32, 64], ids=["tp2", "tp1"])
159 def test_sm100_bf16_512_priority_unchanged(num_heads):
160 """Adding head_size 512 must not reorder the SM100 bf16 sparse list.
161
162 On SM100 with >16 heads the hardcoded bf16 order is
163 [FLASHMLA_SPARSE, FLASHINFER_MLA_SPARSE] (FlashMLA first). head_size 512
164 must leave this untouched (identical to the 576 order); the guard that
165 keeps bf16-512 from *serving* through FlashMLA on SM100 lives in
166 supports_combination (gated to SM90), tested separately.
167 """
168 order512 = _sparse_order("auto", 512, num_heads=num_heads)
169 order576 = _sparse_order("auto", 576, num_heads=num_heads)
170 assert order512 == order576, order512
171 assert order512[0] == "FLASHMLA_SPARSE", order512
172
173
174 def test_sm100_576_priority_unchanged():
175 order = _sparse_order("auto", 576, num_heads=32)
176 assert order[0] == "FLASHMLA_SPARSE", order