@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 225 lines 6.65 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """ROCm kernel correctness tests for AITER MLA decode.
4
5 Compares ``rocm_aiter_ops.mla_decode_fwd`` against a pure PyTorch reference
6 under the absorbed MLA formulation (DeepSeek-V3/V4) with varied batch sizes,
7 head counts, and sequence lengths.
8 """
9
10 import pytest
11 import torch
12
13 from vllm.platforms import current_platform
14 from vllm.utils.torch_utils import set_random_seed
15
16 _SKIP_NON_MI3XX = True
17 if current_platform.is_rocm():
18 from vllm.platforms.rocm import on_mi3xx
19
20 _SKIP_NON_MI3XX = not on_mi3xx()
21
22 pytestmark = [
23 pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific tests"),
24 pytest.mark.skipif(_SKIP_NON_MI3XX, reason="MI300/MI350 ROCm only"),
25 ]
26
27 # DeepSeek-V3/V4 MLA config.
28 Q_HEAD_DIM = 576 # kv_lora_rank + qk_rope_head_dim
29 V_HEAD_DIM = 512 # kv_lora_rank
30 SM_SCALE = Q_HEAD_DIM**-0.5
31
32 NUM_HEADS = [16, 128]
33 BATCH_SIZES = [1, 4, 16]
34 KV_SEQ_LENS = [16, 256]
35
36 ATOL, RTOL = 1.5e-2, 1e-2
37
38
39 def _require_aiter():
40 from vllm._aiter_ops import is_aiter_found_and_supported
41
42 if not is_aiter_found_and_supported():
43 pytest.skip("aiter is required on supported ROCm hardware for this test")
44
45
46 def _ref_mla_decode(
47 q: torch.Tensor,
48 kv_buffer: torch.Tensor,
49 kv_indptr: torch.Tensor,
50 kv_indices: torch.Tensor,
51 ) -> torch.Tensor:
52 """Pure PyTorch reference for absorbed MLA decode.
53
54 K uses full q_head_dim (576) for scoring, V uses first v_head_dim (512)
55 for output. This mirrors the production kernel's absorbed formulation.
56 """
57 batch_size, num_heads, _ = q.shape
58 output = torch.zeros(
59 batch_size, num_heads, V_HEAD_DIM, dtype=q.dtype, device=q.device
60 )
61
62 for b in range(batch_size):
63 start = kv_indptr[b].item()
64 end = kv_indptr[b + 1].item()
65 token_indices = kv_indices[start:end]
66
67 k = kv_buffer[token_indices].float()
68 v = kv_buffer[token_indices, :V_HEAD_DIM].float()
69
70 for h in range(num_heads):
71 q_h = q[b, h, :].float()
72 scores = torch.mv(k, q_h) * SM_SCALE
73 weights = torch.softmax(scores, dim=0)
74 output[b, h, :] = torch.mv(v.t(), weights).to(q.dtype)
75
76 return output
77
78
79 def _make_inputs(
80 batch_size: int,
81 nhead: int,
82 kv_seq_len: int,
83 *,
84 contiguous_indices: bool = True,
85 ):
86 """Build valid MLA decode inputs on the current CUDA device.
87
88 When contiguous_indices=False, tokens are randomly scattered in a 2x pool
89 to simulate real paged allocation.
90 """
91 from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAHelper
92
93 actual_nhead = AiterMLAHelper.get_actual_mla_num_heads(nhead)
94 total_kv_tokens = batch_size * kv_seq_len
95
96 pool_size = total_kv_tokens if contiguous_indices else total_kv_tokens * 2
97
98 q = torch.randn(
99 batch_size, actual_nhead, Q_HEAD_DIM, dtype=torch.bfloat16, device="cuda"
100 )
101 kv_buffer = torch.randn(
102 pool_size, 1, Q_HEAD_DIM, dtype=torch.bfloat16, device="cuda"
103 )
104 o = torch.zeros(
105 batch_size, actual_nhead, V_HEAD_DIM, dtype=torch.bfloat16, device="cuda"
106 )
107
108 qo_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda")
109 kv_indptr = (
110 torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda") * kv_seq_len
111 )
112 kv_last_page_lens = torch.ones(batch_size, dtype=torch.int32, device="cuda")
113
114 if contiguous_indices:
115 kv_indices = torch.arange(0, total_kv_tokens, dtype=torch.int32, device="cuda")
116 else:
117 perm = torch.randperm(pool_size, device="cuda")[:total_kv_tokens]
118 kv_indices = perm.to(torch.int32)
119
120 return {
121 "q": q,
122 "kv_buffer": kv_buffer,
123 "o": o,
124 "qo_indptr": qo_indptr,
125 "kv_indptr": kv_indptr,
126 "kv_indices": kv_indices,
127 "kv_last_page_lens": kv_last_page_lens,
128 "actual_nhead": actual_nhead,
129 }
130
131
132 def _run_kernel(inputs: dict) -> torch.Tensor:
133 """Run decode through the production rocm_aiter_ops path."""
134 from vllm._aiter_ops import rocm_aiter_ops
135
136 rocm_aiter_ops.mla_decode_fwd(
137 inputs["q"],
138 inputs["kv_buffer"],
139 inputs["o"],
140 SM_SCALE,
141 inputs["qo_indptr"],
142 1,
143 inputs["kv_indptr"],
144 inputs["kv_indices"],
145 inputs["kv_last_page_lens"],
146 )
147 return inputs["o"]
148
149
150 def _ref_output(inputs: dict) -> torch.Tensor:
151 """Compute reference output from the same inputs."""
152 kv_flat = inputs["kv_buffer"].squeeze(1)
153 return _ref_mla_decode(
154 inputs["q"], kv_flat, inputs["kv_indptr"], inputs["kv_indices"]
155 )
156
157
158 @pytest.mark.parametrize("nhead", NUM_HEADS)
159 @pytest.mark.parametrize("batch_size", BATCH_SIZES)
160 @pytest.mark.parametrize("kv_seq_len", KV_SEQ_LENS)
161 @torch.inference_mode()
162 def test_mla_decode_accuracy(
163 nhead: int,
164 batch_size: int,
165 kv_seq_len: int,
166 ) -> None:
167 """BF16 decode accuracy vs PyTorch reference (contiguous indices)."""
168 _require_aiter()
169 set_random_seed(0)
170
171 inputs = _make_inputs(batch_size, nhead, kv_seq_len)
172 output = _run_kernel(inputs)
173 output_ref = _ref_output(inputs)
174
175 torch.testing.assert_close(output, output_ref, atol=ATOL, rtol=RTOL)
176
177
178 @pytest.mark.parametrize("nhead", NUM_HEADS)
179 @pytest.mark.parametrize("batch_size", [1, 4, 8])
180 @pytest.mark.parametrize("kv_seq_len", [16, 128])
181 @torch.inference_mode()
182 def test_mla_decode_noncontiguous_indices(
183 nhead: int,
184 batch_size: int,
185 kv_seq_len: int,
186 ) -> None:
187 """BF16 decode accuracy with shuffled kv_indices (paged allocation)."""
188 _require_aiter()
189 set_random_seed(0)
190
191 inputs = _make_inputs(batch_size, nhead, kv_seq_len, contiguous_indices=False)
192 output = _run_kernel(inputs)
193 output_ref = _ref_output(inputs)
194
195 torch.testing.assert_close(output, output_ref, atol=ATOL, rtol=RTOL)
196
197
198 @torch.inference_mode()
199 def test_mla_decode_determinism() -> None:
200 """Repeated decode calls produce bitwise-identical results."""
201 _require_aiter()
202 set_random_seed(0)
203
204 inputs = _make_inputs(batch_size=4, nhead=128, kv_seq_len=16)
205 out_first = _run_kernel(inputs).clone()
206
207 for _ in range(3):
208 inputs["o"] = torch.zeros_like(inputs["o"])
209 out = _run_kernel(inputs)
210 torch.testing.assert_close(out, out_first, atol=0, rtol=0)
211
212
213 @torch.inference_mode()
214 def test_mla_decode_smoke() -> None:
215 """Basic sanity: correct shape, dtype, finite, non-zero."""
216 _require_aiter()
217 set_random_seed(0)
218
219 inputs = _make_inputs(batch_size=4, nhead=128, kv_seq_len=64)
220 output = _run_kernel(inputs)
221
222 assert output.shape == (4, 128, V_HEAD_DIM)
223 assert output.dtype == torch.bfloat16
224 assert torch.isfinite(output).all()
225 assert not torch.all(output == 0)