@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 381 lines 13.8 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import pytest
4 import torch
5 import torch.nn.functional as F
6 from torch import Tensor
7
8 from vllm.platforms import current_platform
9
10 FLASHINFER_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
11
12 if not current_platform.is_cuda() or not current_platform.has_device_capability(90):
13 pytest.skip(
14 reason="FlashInfer MLA requires CUDA compute capability 9.0 or above.",
15 allow_module_level=True,
16 )
17 else:
18 from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
19 from flashinfer.mla import BatchMLAPagedAttentionWrapper
20
21 requires_sm90 = pytest.mark.skipif(
22 not current_platform.is_device_capability_family(90),
23 reason="This test requires an SM90 GPU.",
24 )
25 requires_sm10x = pytest.mark.skipif(
26 not current_platform.is_device_capability_family(100),
27 reason="This test requires an SM10x GPU.",
28 )
29
30 # Deepseek R1 MLA config.
31 NUM_HEADS = 128
32 KV_LORA_RANK = 512
33 QK_NOPE_HEAD_DIM = 128
34 QK_ROPE_HEAD_DIM = 64
35 QK_HEAD_DIM = KV_LORA_RANK + QK_ROPE_HEAD_DIM
36 SCALE = (QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM) ** -0.5
37
38
39 def _make_decode_inputs(bs: int, block_size: int, dtype: torch.dtype):
40 """Build valid trtllm MLA decode inputs on the current CUDA device."""
41 max_seq_len_cap = 1024
42 seq_lens = [torch.randint(2, max_seq_len_cap, (1,)).item() for _ in range(bs)]
43 seq_lens[-1] = max_seq_len_cap
44 max_seq_len = max(seq_lens)
45 seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int32)
46
47 # Generate block tables with random but unique block IDs
48 # From https://github.com/flashinfer-ai/flashinfer/pull/1222
49 blocks_per_seq = (seq_lens_tensor + block_size - 1) // block_size
50 max_num_blocks_per_seq = max(blocks_per_seq.max().item(), 4)
51 total_blocks_needed = int(sum(blocks_per_seq))
52 all_block_ids = torch.randperm(total_blocks_needed)
53
54 block_tables = torch.zeros((bs, max_num_blocks_per_seq), dtype=torch.int32)
55 block_id = 0
56 for i in range(bs):
57 num_blocks_needed = blocks_per_seq[i]
58 block_tables[i, :num_blocks_needed] = all_block_ids[
59 block_id : block_id + num_blocks_needed
60 ]
61 block_id += num_blocks_needed
62
63 kv_cache = torch.randn(block_tables.numel(), block_size, QK_HEAD_DIM).to(dtype)
64 q = torch.randn(bs, NUM_HEADS, QK_HEAD_DIM).to(dtype)
65 return q, kv_cache, block_tables, seq_lens_tensor, max_seq_len
66
67
68 def ref_mla(
69 out: Tensor, # (bs, num_heads, v_head_dim)
70 query: Tensor, # (bs, num_heads, head_dim)
71 kv_cache: Tensor, # (num_blocks, block_size, head_dim)
72 scale: float,
73 block_tables: Tensor, # (bs, max_num_blocks)
74 seq_lens: Tensor, # (bs,)
75 ):
76 bs, num_heads, v_head_dim = out.shape
77 head_dim = query.shape[2]
78
79 for i in range(bs):
80 # gather and flatten KV-cache
81 kv = kv_cache[block_tables[i]] # (max_num_blocks, block_size, head_dim)
82 kv = kv.view(1, -1, head_dim)[:, : seq_lens[i]] # (1, seq_len, head_dim)
83 v = kv[:, :, :v_head_dim]
84
85 q = query[i].view(num_heads, 1, head_dim)
86 o = F.scaled_dot_product_attention(q, kv, v, scale=scale, enable_gqa=True)
87 out[i] = o.view(num_heads, v_head_dim)
88
89 return out
90
91
92 @pytest.mark.parametrize("dtype", [torch.bfloat16])
93 @pytest.mark.parametrize("bs", [1, 2, 4, 16])
94 @pytest.mark.parametrize("block_size", [32, 64])
95 @requires_sm10x
96 def test_flashinfer_mla_decode(dtype: torch.dtype, bs: int, block_size: int):
97 torch.set_default_device("cuda")
98 torch.manual_seed(42)
99
100 q, kv_cache, block_tables, seq_lens_tensor, max_seq_len = _make_decode_inputs(
101 bs, block_size, dtype
102 )
103
104 out_ref = q.new_zeros(bs, NUM_HEADS, KV_LORA_RANK)
105 ref_mla(out_ref, q, kv_cache, SCALE, block_tables, seq_lens_tensor)
106
107 workspace_buffer = torch.zeros(
108 FLASHINFER_WORKSPACE_BUFFER_SIZE,
109 dtype=torch.uint8,
110 device=q.device,
111 )
112 # Flashinfer MLA expects the query to be of shape
113 # (bs, q_len_per_request, num_heads, qk_head_dim),
114 # where q_len_per_request is the MTP query length (=1 without MTP)
115 q = q.unsqueeze(1)
116
117 out_ans = trtllm_batch_decode_with_kv_cache_mla(
118 query=q,
119 kv_cache=kv_cache.unsqueeze(1),
120 workspace_buffer=workspace_buffer,
121 qk_nope_head_dim=QK_NOPE_HEAD_DIM,
122 kv_lora_rank=KV_LORA_RANK,
123 qk_rope_head_dim=QK_ROPE_HEAD_DIM,
124 block_tables=block_tables,
125 seq_lens=seq_lens_tensor,
126 max_seq_len=max_seq_len,
127 bmm1_scale=SCALE,
128 )
129 out_ans = out_ans.squeeze(1)
130 torch.testing.assert_close(out_ans, out_ref, atol=1e-2, rtol=1e-2)
131
132
133 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn])
134 @pytest.mark.parametrize("mode", ["tp", "dcp", "hisparse"])
135 @requires_sm10x
136 def test_flashinfer_trtllm_sparse_mla_decode_without_rope(dtype, mode, monkeypatch):
137 """TP skips preparation; empty DCP contributions and HiSparse outputs vanish."""
138 import vllm.v1.attention.backends.mla.flashinfer_mla_sparse as backend
139 from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
140
141 torch.set_default_device("cuda")
142 torch.manual_seed(42)
143 batch_size, num_heads, block_size, num_blocks, sparse_topk = 3, 8, 64, 4, 2176
144 valid_lens = [1 if mode == "tp" else 0, 17, 73]
145 counts = torch.tensor(valid_lens, dtype=torch.int32)
146 query = torch.randn(batch_size, 1, num_heads, KV_LORA_RANK).to(dtype)
147 kv_cache = torch.randn(num_blocks, block_size, KV_LORA_RANK).to(dtype)
148 num_slots = num_blocks * block_size
149 indices = torch.full((batch_size, sparse_topk), -1, dtype=torch.int32)
150 for token, length in enumerate(valid_lens):
151 indices[token, :length] = torch.randperm(num_slots)[:length].int()
152 empty = counts == 0
153 reference_indices = indices.clone()
154 reference_indices[:, 0].masked_fill_(empty, 0)
155 workspace = torch.empty(FLASHINFER_WORKSPACE_BUFFER_SIZE, dtype=torch.int8)
156
157 impl = object.__new__(backend.FlashInferMLASparseImpl)
158 impl.qk_nope_head_dim = QK_NOPE_HEAD_DIM
159 impl.kv_lora_rank = KV_LORA_RANK
160 impl.qk_rope_head_dim = 0
161 impl.index_group = (
162 object.__new__(HiSparseMLAIndexGroup) if mode == "hisparse" else None
163 )
164 impl._workspace_buffer = workspace
165 impl.bmm1_scale = QK_NOPE_HEAD_DIM**-0.5
166 impl.bmm2_scale = 1.0
167 impl.is_nope_mla = True
168 impl.need_to_return_lse_for_decode = mode == "dcp"
169 if mode == "tp":
170
171 def unexpected_helper(*args, **kwargs):
172 pytest.fail("TP must not launch empty-query preparation")
173
174 monkeypatch.setattr(
175 backend, "prepare_sparse_mla_safe_lengths", unexpected_helper
176 )
177
178 out, lse = impl._run_mqa_kernel(query.squeeze(1), kv_cache, indices, counts)
179 reference = trtllm_batch_decode_with_kv_cache_mla(
180 query=query,
181 kv_cache=kv_cache.unsqueeze(1),
182 workspace_buffer=workspace,
183 qk_nope_head_dim=QK_NOPE_HEAD_DIM,
184 kv_lora_rank=KV_LORA_RANK,
185 qk_rope_head_dim=0,
186 block_tables=reference_indices.unsqueeze(1),
187 seq_lens=counts,
188 max_seq_len=sparse_topk,
189 sparse_mla_top_k=sparse_topk,
190 sparse_mla_top_k_lens=counts.clamp(min=1),
191 bmm1_scale=QK_NOPE_HEAD_DIM**-0.5,
192 bmm2_scale=1.0,
193 return_lse=mode == "dcp",
194 )
195 if mode == "dcp":
196 reference, reference_lse = reference
197 reference_lse = reference_lse.reshape(batch_size, num_heads)
198 reference_lse.masked_fill_(empty[:, None], float("-inf"))
199 torch.testing.assert_close(lse, reference_lse, atol=0, rtol=0)
200
201 from vllm.v1.attention.ops.dcp import CPTritonContext, correct_attn_out
202
203 # The peer owns only our empty selections. Check our contribution after
204 # the production merge, without zeroing the local attention output.
205 peer_lse = torch.full_like(lse, float("-inf")).masked_fill(empty[:, None], 0)
206 out, _ = correct_attn_out(
207 out, torch.stack((lse, peer_lse)), cp_rank=0, ctx=CPTritonContext()
208 )
209 else:
210 assert lse is None
211 reference = reference.squeeze(1).masked_fill(empty[:, None, None], 0)
212 torch.testing.assert_close(out, reference, atol=0, rtol=0)
213
214 if dtype == torch.bfloat16:
215 flat_cache = kv_cache.view(num_slots, KV_LORA_RANK).float()
216 refs = []
217 for token, length in enumerate(valid_lens):
218 selected_kv = flat_cache[reference_indices[token, :length].long()]
219 scores = torch.einsum("hd,kd->hk", query[token, 0].float(), selected_kv)
220 probs = torch.softmax(scores * QK_NOPE_HEAD_DIM**-0.5, dim=-1)
221 refs.append(torch.einsum("hk,kd->hd", probs, selected_kv))
222 ref = torch.stack(refs).to(torch.bfloat16)
223 torch.testing.assert_close(out, ref, atol=2e-2, rtol=2e-2)
224
225
226 @requires_sm10x
227 def test_sparse_mla_safe_lengths():
228 from vllm.v1.attention.backends.mla.sparse_utils import (
229 prepare_sparse_mla_safe_lengths,
230 )
231
232 device = "cuda"
233 num_tokens = 259
234 # Include a partial tile and expose accidental writes into adjacent storage.
235 counts_storage = torch.full((num_tokens, 2), 17, dtype=torch.int32, device=device)
236 counts = counts_storage[:, 0]
237 counts[::3] = 0
238 original_counts = counts_storage.clone()
239 empty = counts == 0
240 indices_storage = torch.full(
241 (num_tokens, 256), 42, dtype=torch.int32, device=device
242 )
243 indices = indices_storage[:, ::2]
244 expected_indices = indices_storage.clone()
245 expected_indices[:, 0].masked_fill_(empty, 0)
246 safe = prepare_sparse_mla_safe_lengths(indices, counts)
247 torch.testing.assert_close(safe, counts.clamp(min=1), rtol=0, atol=0)
248 torch.testing.assert_close(indices_storage, expected_indices, rtol=0, atol=0)
249 torch.testing.assert_close(counts_storage, original_counts, rtol=0, atol=0)
250
251
252 @requires_sm90
253 def test_flashinfer_sm90_fp8_mla_decode_without_rope():
254 """Hopper FA3 supports BF16 queries over an FP8 cache without KPE."""
255 torch.manual_seed(42)
256 device = torch.device("cuda")
257 batch_size = 2
258 num_heads = 16
259 page_size = 16
260 num_pages = 6
261
262 q_nope = torch.randn(
263 batch_size,
264 num_heads,
265 KV_LORA_RANK,
266 dtype=torch.bfloat16,
267 device=device,
268 )
269 q_pe = torch.empty(
270 batch_size,
271 num_heads,
272 0,
273 dtype=torch.bfloat16,
274 device=device,
275 )
276
277 ckv = torch.randn(
278 num_pages,
279 page_size,
280 KV_LORA_RANK,
281 device=device,
282 )
283 fp8_max = torch.finfo(torch.float8_e4m3fn).max
284 ckv_scale = ckv.abs().max().item() / fp8_max
285 ckv_fp8 = (ckv / ckv_scale).clamp(-fp8_max, fp8_max).to(torch.float8_e4m3fn)
286 scale_bf16 = torch.tensor(ckv_scale, dtype=torch.bfloat16, device=device)
287 ckv_ref = ckv_fp8.to(torch.bfloat16) * scale_bf16
288 kpe_fp8 = torch.empty(
289 num_pages,
290 page_size,
291 0,
292 dtype=torch.float8_e4m3fn,
293 device=device,
294 )
295 kpe_ref = torch.empty(
296 num_pages,
297 page_size,
298 0,
299 dtype=torch.bfloat16,
300 device=device,
301 )
302
303 qo_indptr = torch.tensor([0, 1, 2], dtype=torch.int32, device=device)
304 kv_indptr = torch.tensor([0, 3, 5], dtype=torch.int32, device=device)
305 kv_indices = torch.tensor([4, 1, 3, 0, 5], dtype=torch.int32, device=device)
306 kv_lens = torch.tensor([45, 29], dtype=torch.int32, device=device)
307 sm_scale = QK_NOPE_HEAD_DIM**-0.5
308
309 def run(
310 ckv_cache: torch.Tensor,
311 kpe_cache: torch.Tensor,
312 **kwargs,
313 ) -> torch.Tensor:
314 workspace = torch.empty(
315 FLASHINFER_WORKSPACE_BUFFER_SIZE,
316 dtype=torch.uint8,
317 device=device,
318 )
319 wrapper = BatchMLAPagedAttentionWrapper(workspace, backend="fa3")
320 wrapper.plan(
321 qo_indptr,
322 kv_indptr,
323 kv_indices,
324 kv_lens,
325 num_heads,
326 KV_LORA_RANK,
327 0,
328 page_size,
329 False,
330 sm_scale,
331 q_data_type=torch.bfloat16,
332 kv_data_type=ckv_cache.dtype,
333 )
334 return wrapper.run(q_nope, q_pe, ckv_cache, kpe_cache, **kwargs)
335
336 out_ref = run(ckv_ref, kpe_ref)
337 out = run(ckv_fp8, kpe_fp8, ckv_scale=ckv_scale, kpe_scale=1.0)
338 torch.testing.assert_close(out, out_ref, atol=2e-2, rtol=2e-2)
339
340
341 @requires_sm10x
342 def test_flashinfer_mla_decode_workspace_supports_autotune():
343 """VLLM's FlashInfer MLA decode workspace must be int8 for autotuning.
344
345 Model Runner V2's warmup autotunes ``trtllm_batch_decode_mla``, which makes
346 the FlashInfer autotuner enumerate the CuteDSL tactic. That tactic asserts
347 ``workspace_buffer.dtype == torch.int8``; the trtllm-gen path (used for
348 normal, non-autotuned inference) instead views the buffer as uint8, so a
349 uint8 workspace only fails once the autotuner tries CuteDSL. That regressed
350 every DeepSeek MLA test on Blackwell under V2 with
351 ``workspace_buffer must be torch.int8`` (vllm-project/vllm#46646).
352 """
353 from flashinfer.autotuner import autotune
354
355 from vllm.v1.attention.backends.mla.flashinfer_mla import _get_workspace_buffer
356
357 torch.set_default_device("cuda")
358 torch.manual_seed(0)
359
360 workspace_buffer = _get_workspace_buffer(return_lse=False)
361 assert workspace_buffer.dtype == torch.int8
362
363 q, kv_cache, block_tables, seq_lens_tensor, max_seq_len = _make_decode_inputs(
364 bs=1, block_size=64, dtype=torch.bfloat16
365 )
366
367 # Under the autotuner the CuteDSL tactic is instantiated with our workspace;
368 # a uint8 buffer raises AssertionError here, an int8 buffer succeeds.
369 with torch.inference_mode(), autotune(True):
370 trtllm_batch_decode_with_kv_cache_mla(
371 query=q.unsqueeze(1),
372 kv_cache=kv_cache.unsqueeze(1),
373 workspace_buffer=workspace_buffer,
374 qk_nope_head_dim=QK_NOPE_HEAD_DIM,
375 kv_lora_rank=KV_LORA_RANK,
376 qk_rope_head_dim=QK_ROPE_HEAD_DIM,
377 block_tables=block_tables,
378 seq_lens=seq_lens_tensor,
379 max_seq_len=max_seq_len,
380 bmm1_scale=SCALE,
381 )