@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 688 lines 27.9 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Reference-vs-optimized unit tests for the MiniMax-M3 AMD/ROCm fused kernels.
4
5 Each optimized kernel added for the ROCm port has a slow PyTorch reference; the
6 tests assert the two agree within tolerance:
7
8 * Gemma RMSNorm (plain + fused-add-residual) -> fp32 PyTorch normalize
9 * SwiGLU-OAI (split layout) -> fp32 PyTorch elementwise
10 * Fused MXFP8 activation quant (Triton) -> _mxfp8_e4m3_quantize_torch
11 * Native MXFP8 linear (dot_scaled) -> dequant-to-bf16 @ matmul
12 * Native MXFP8 MoE (dot_scaled grouped GEMM) -> dequant-to-bf16 MoE math
13 * Strided K/V sparse-PA cache insert -> contiguous-copy insert
14
15 The native MXFP8 GEMMs also guard the ``dot_scaled`` rhs-scale orientation: the
16 scale is loaded ``[N, K//32]`` and passed WITHOUT transpose; a stray ``.T``
17 makes the shape ``[K//32, N]`` and Triton raises before producing output, so any
18 regression there fails these tests loudly.
19
20 Hardware scope: the whole module is ROCm-only (these are the AMD path; NVIDIA
21 uses the FlashInfer kernels). The norm/activation/quant kernels run on any ROCm
22 arch; the native MXFP8 ``dot_scaled`` linear/MoE tests are additionally gated to
23 CDNA4 gfx95x (``@requires_gfx950``) since gfx942 uses the BF16 emulation path.
24
25 Run: pytest tests/kernels/test_minimax_m3_amd_ops.py -v
26 """
27
28 import pytest
29 import torch
30
31 from vllm.platforms import current_platform
32
33 if not current_platform.is_rocm():
34 pytest.skip("MiniMax-M3 AMD fused ops require ROCm.", allow_module_level=True)
35 if not torch.cuda.is_available():
36 pytest.skip("Requires a GPU.", allow_module_level=True)
37
38 from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( # noqa: E402
39 _mxfp8_e4m3_quantize_torch,
40 _mxfp8_e4m3_quantize_triton,
41 dequant_mxfp8_to_bf16,
42 )
43 from vllm.models.minimax_m3.amd.model import _kv_insert_operand # noqa: E402
44 from vllm.models.minimax_m3.amd.ops import ( # noqa: E402
45 gemma_fused_add_rmsnorm,
46 gemma_rmsnorm,
47 swiglu_oai_quantize_mxfp8,
48 swiglu_oai_split,
49 )
50 from vllm.models.minimax_m3.amd.ops.gemma_rmsnorm import _num_warps # noqa: E402
51
52 DEVICE = "cuda"
53 EPS = 1e-6
54
55
56 def _gcn_arch() -> str:
57 try:
58 return torch.cuda.get_device_properties(0).gcnArchName
59 except Exception: # pragma: no cover - no device / non-AMD
60 return ""
61
62
63 # The pure-Triton norm/activation/quant kernels run on any ROCm arch (CDNA3
64 # gfx942 and CDNA4 gfx950). The native MXFP8 ``dot_scaled`` GEMMs (linear + MoE)
65 # use CDNA4 hardware microscaling and are gated to gfx95x in the source
66 # (``RocmDotScaledMxfp8LinearKernel.is_supported``; the MoE oracle routes gfx942
67 # to the BF16 emulation path instead) — so those tests are gfx950-only.
68 requires_gfx950 = pytest.mark.skipif(
69 "gfx95" not in _gcn_arch(),
70 reason="native MXFP8 dot_scaled is a CDNA4 (gfx95x) feature; "
71 "gfx942 uses the BF16 emulation path instead.",
72 )
73
74
75 def _relerr(a: torch.Tensor, b: torch.Tensor) -> float:
76 a = a.float()
77 b = b.float()
78 return ((a - b).norm() / (b.norm() + 1e-8)).item()
79
80
81 # --------------------------------------------------------------------------- #
82 # Gemma RMSNorm
83 # --------------------------------------------------------------------------- #
84 def _ref_gemma_rmsnorm(x, w, eps, residual=None):
85 orig_dtype = x.dtype
86 xf = x.float()
87 res_out = None
88 if residual is not None:
89 xf = xf + residual.float()
90 res_out = xf.to(orig_dtype)
91 xf = xf * torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps)
92 xf = xf * (1.0 + w.float())
93 out = xf.to(orig_dtype)
94 return out if residual is None else (out, res_out)
95
96
97 @pytest.mark.parametrize("shape", [(1, 4096), (37, 6144), (128, 2048)])
98 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
99 @pytest.mark.parametrize("seed", [0, 1234])
100 @torch.inference_mode()
101 def test_gemma_rmsnorm(shape, dtype, seed):
102 torch.manual_seed(seed)
103 x = torch.randn(*shape, device=DEVICE, dtype=dtype)
104 w = torch.randn(shape[-1], device=DEVICE, dtype=dtype) * 0.1
105 got = gemma_rmsnorm(x, w, EPS)
106 ref = _ref_gemma_rmsnorm(x, w, EPS)
107 assert got.shape == x.shape
108 assert _relerr(got, ref) < 5e-3
109
110
111 @pytest.mark.parametrize("shape", [(1, 6144), (64, 4096)])
112 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
113 @torch.inference_mode()
114 def test_gemma_fused_add_rmsnorm(shape, dtype):
115 torch.manual_seed(0)
116 x = torch.randn(*shape, device=DEVICE, dtype=dtype)
117 res = torch.randn(*shape, device=DEVICE, dtype=dtype)
118 w = torch.randn(shape[-1], device=DEVICE, dtype=dtype) * 0.1
119 got_out, got_res = gemma_fused_add_rmsnorm(x, res, w, EPS)
120 ref_out, ref_res = _ref_gemma_rmsnorm(x, w, EPS, residual=res)
121 assert _relerr(got_out, ref_out) < 5e-3
122 # residual_out is the pre-norm sum (x + res): bit-for-bit identical cast.
123 assert torch.equal(got_res, ref_res)
124
125
126 @torch.inference_mode()
127 def test_gemma_rmsnorm_per_head_strided():
128 """q_norm/k_norm normalize a non-contiguous ``qkv.split`` slice over head_dim."""
129 torch.manual_seed(0)
130 T, H, D, kv = 7, 48, 128, 8
131 total = (H + 2 * kv) * D
132 qkv = torch.randn(T, total, device=DEVICE, dtype=torch.bfloat16)
133 q = qkv[..., : H * D] # non-contiguous view (row stride == total)
134 q_by_head = q.view(T, H, D)
135 assert not q_by_head.is_contiguous()
136 w = torch.randn(D, device=DEVICE, dtype=torch.bfloat16) * 0.1
137 got = gemma_rmsnorm(q_by_head, w, EPS)
138 ref = _ref_gemma_rmsnorm(q_by_head, w, EPS)
139 assert got.shape == q_by_head.shape
140 assert _relerr(got, ref) < 5e-3
141
142
143 def test_num_warps_monotonic():
144 assert _num_warps(128) <= _num_warps(2048) <= _num_warps(8192)
145
146
147 # --------------------------------------------------------------------------- #
148 # SwiGLU-OAI (split layout)
149 # --------------------------------------------------------------------------- #
150 def _ref_swiglu(gate_up, alpha, beta, limit):
151 d = gate_up.shape[-1] // 2
152 gate = gate_up[..., :d].float()
153 up = gate_up[..., d:].float()
154 if limit is not None:
155 gate = gate.clamp(max=limit)
156 up = up.clamp(min=-limit, max=limit)
157 return (gate * torch.sigmoid(alpha * gate) * (up + beta)).to(gate_up.dtype)
158
159
160 @pytest.mark.parametrize("m,inter", [(1, 768), (64, 1536), (128, 1024)])
161 @pytest.mark.parametrize("limit", [7.0, None])
162 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
163 @torch.inference_mode()
164 def test_swiglu_oai_split(m, inter, limit, dtype):
165 torch.manual_seed(0)
166 gate_up = torch.randn(m, 2 * inter, device=DEVICE, dtype=dtype)
167 got = swiglu_oai_split(gate_up, alpha=1.702, beta=1.0, limit=limit)
168 ref = _ref_swiglu(gate_up, 1.702, 1.0, limit)
169 assert got.shape == (m, inter)
170 assert _relerr(got, ref) < 5e-3
171
172
173 @torch.inference_mode()
174 def test_swiglu_oai_quantize_mxfp8_uses_e4m3_range_for_scale():
175 # Keep a small value in two MX blocks next to their maxima. The scale must
176 # use E4M3's finite range (448), otherwise that value underflows to zero.
177 # Keep the third block empty to cover the reference's finite tiny clamp.
178 m, inter = 3, 96
179 gate = torch.full((m, inter), 2.0, device=DEVICE, dtype=torch.float16)
180 up = torch.full((m, inter), 2**-10, device=DEVICE, dtype=torch.float16)
181 up[:, :64:32] = 1.0
182 up[:, 64:] = 0.0
183 gate_up = torch.cat((gate, up), dim=-1)
184
185 got_q, got_s = swiglu_oai_quantize_mxfp8(gate_up, alpha=0.0, beta=0.0, limit=None)
186 ref_q, ref_s = _mxfp8_e4m3_quantize_torch(up, is_sf_swizzled_layout=False)
187
188 assert torch.equal(got_s, ref_s)
189 assert torch.equal(got_q, ref_q)
190
191
192 # --------------------------------------------------------------------------- #
193 # Fused MXFP8 activation quant (Triton vs torch reference)
194 # --------------------------------------------------------------------------- #
195 @pytest.mark.parametrize("shape", [(64, 4096), (1, 6144), (333, 2048)])
196 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
197 @torch.inference_mode()
198 def test_mxfp8_quant_triton_matches_torch(shape, dtype):
199 torch.manual_seed(0)
200 x = torch.randn(*shape, device=DEVICE, dtype=dtype)
201 xq_t, s_t = _mxfp8_e4m3_quantize_torch(x, is_sf_swizzled_layout=False)
202 xq_k, s_k = _mxfp8_e4m3_quantize_triton(x)
203 assert s_k.shape == s_t.shape == (shape[0], shape[1] // 32)
204 # Both paths use the E4M3-aware scale calculation; allow a 1-step
205 # difference at exact powers of two due to floating-point rounding.
206 assert (s_k.int() - s_t.int()).abs().max().item() <= 1
207 # Dequantized values agree to fp8 granularity.
208 deq_t = dequant_mxfp8_to_bf16(xq_t, s_t)
209 deq_k = dequant_mxfp8_to_bf16(xq_k, s_k)
210 assert _relerr(deq_k, deq_t) < 1e-2
211
212
213 # --------------------------------------------------------------------------- #
214 # Native MXFP8 linear (dot_scaled) vs dequant-to-bf16 matmul
215 # --------------------------------------------------------------------------- #
216 @requires_gfx950
217 @pytest.mark.parametrize("m,n,k", [(64, 256, 128), (37, 512, 256), (1, 6144, 4096)])
218 @torch.inference_mode()
219 def test_mxfp8_native_linear(m, n, k):
220 from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (
221 _mxfp8_dot_scaled_linear,
222 )
223
224 torch.manual_seed(0)
225 w_bf16 = torch.randn(n, k, device=DEVICE, dtype=torch.bfloat16) * 0.1
226 w_fp8, w_scale = _mxfp8_e4m3_quantize_torch(w_bf16, is_sf_swizzled_layout=False)
227 x = torch.randn(m, k, device=DEVICE, dtype=torch.bfloat16) * 0.5
228
229 got = _mxfp8_dot_scaled_linear(x, w_fp8, w_scale)
230 # Reference: consume the SAME quantized weights (isolates activation-quant
231 # noise) -> dequant to bf16, plain matmul.
232 w_deq = dequant_mxfp8_to_bf16(w_fp8, w_scale)
233 ref = torch.nn.functional.linear(x, w_deq).to(x.dtype)
234 assert got.shape == (m, n)
235 # Only the activation is re-quantized inside the kernel -> small MX noise.
236 assert _relerr(got, ref) < 5e-2
237
238
239 # --------------------------------------------------------------------------- #
240 # Native MXFP8 MoE (dot_scaled grouped GEMM) vs dequant-to-bf16 MoE math
241 # --------------------------------------------------------------------------- #
242 def _ref_moe(x, w13, w2, topk_weights, topk_ids, alpha, beta, limit):
243 T, H = x.shape
244 inter = w2.shape[-1]
245 top_k = topk_ids.shape[1]
246 out = torch.zeros(T, H, device=x.device, dtype=torch.float32)
247 for t in range(T):
248 for j in range(top_k):
249 e = int(topk_ids[t, j].item())
250 g1 = x[t].float() @ w13[e].float().T # [2I]
251 gate = g1[:inter]
252 up = g1[inter:]
253 if limit is not None:
254 gate = gate.clamp(max=limit)
255 up = up.clamp(min=-limit, max=limit)
256 act = gate * torch.sigmoid(alpha * gate) * (up + beta)
257 g2 = act @ w2[e].float().T # [H]
258 out[t] += topk_weights[t, j].float() * g2
259 return out.to(x.dtype)
260
261
262 @requires_gfx950
263 @pytest.mark.parametrize(
264 "T,H,inter,E,top_k", [(8, 256, 512, 8, 2), (1, 512, 256, 16, 4)]
265 )
266 @torch.inference_mode()
267 def test_mxfp8_native_moe(T, H, inter, E, top_k):
268 from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import (
269 fused_moe_mxfp8_native,
270 )
271
272 torch.manual_seed(0)
273 alpha, beta, limit = 1.702, 1.0, 7.0
274 w13_bf16 = torch.randn(E, 2 * inter, H, device=DEVICE, dtype=torch.bfloat16) * 0.1
275 w2_bf16 = torch.randn(E, H, inter, device=DEVICE, dtype=torch.bfloat16) * 0.1
276 w13_fp8, w13_scale = _mxfp8_e4m3_quantize_torch(
277 w13_bf16, is_sf_swizzled_layout=False
278 )
279 w2_fp8, w2_scale = _mxfp8_e4m3_quantize_torch(w2_bf16, is_sf_swizzled_layout=False)
280
281 x = torch.randn(T, H, device=DEVICE, dtype=torch.bfloat16) * 0.5
282 logits = torch.randn(T, E, device=DEVICE, dtype=torch.float32)
283 topk_weights, topk_ids = logits.softmax(dim=-1).topk(top_k, dim=-1)
284 topk_weights = topk_weights.to(torch.float32)
285 topk_ids = topk_ids.to(torch.int32)
286
287 got = fused_moe_mxfp8_native(
288 x,
289 w13_fp8,
290 w13_scale,
291 w2_fp8,
292 w2_scale,
293 topk_weights,
294 topk_ids,
295 alpha=alpha,
296 beta=beta,
297 limit=limit,
298 global_num_experts=E,
299 expert_map=None,
300 )
301 # Reference consumes the dequantized weights (same bits the kernel reads).
302 w13_deq = dequant_mxfp8_to_bf16(w13_fp8, w13_scale)
303 w2_deq = dequant_mxfp8_to_bf16(w2_fp8, w2_scale)
304 ref = _ref_moe(x, w13_deq, w2_deq, topk_weights, topk_ids, alpha, beta, limit)
305 assert got.shape == (T, H)
306 assert _relerr(got, ref) < 5e-2
307
308
309 # --------------------------------------------------------------------------- #
310 # Native MXFP8 grouped GEMM (dot_scaled) vs pure-PyTorch grouped matmul
311 # --------------------------------------------------------------------------- #
312 def _ref_grouped_gemm(a_deq, w_deq, topk_ids, a_div, num_valid, mul_weight=None):
313 """Pure-PyTorch reference for ``_grouped_gemm_mxfp8``.
314
315 For each routed (expanded) token ``tid in [0, num_valid)`` the kernel writes
316 ``out[tid] = a[tid // a_div] @ w[expert(tid)].T`` (fp32 accumulate), optionally
317 scaled by ``mul_weight[tid]``. The expert for ``tid`` is ``topk_ids.flatten()
318 [tid]`` (row-major expansion: ``tid = token*top_k + slot``). This is computed
319 here with plain ``torch.matmul`` on the dequantized operands — independent of
320 the Triton ``dot_scaled`` path and of the (separate) aiter backend.
321 """
322 eids = topk_ids.reshape(-1)
323 n = w_deq.shape[1]
324 out = torch.empty(num_valid, n, dtype=torch.float32, device=a_deq.device)
325 for tid in range(num_valid):
326 e = int(eids[tid].item())
327 out[tid] = a_deq[tid // a_div].float() @ w_deq[e].float().T
328 if mul_weight is not None:
329 out[tid] *= float(mul_weight[tid].item())
330 return out
331
332
333 @requires_gfx950
334 @pytest.mark.parametrize("T,N,K,E,top_k", [(8, 256, 128, 8, 2), (5, 512, 256, 16, 4)])
335 @pytest.mark.parametrize("weighted", [False, True])
336 @torch.inference_mode()
337 def test_mxfp8_grouped_gemm_native(T, N, K, E, top_k, weighted):
338 """Directly exercise ``_grouped_gemm_mxfp8`` against a non-Triton reference.
339
340 Covers both call modes used by ``fused_moe_mxfp8_native``:
341 * ``weighted=False`` -> g1: ``a_div=top_k`` (a-row shared across the top_k
342 expansions of a token), no per-token weight.
343 * ``weighted=True`` -> g2: ``a_div=1`` (one a-row per expansion), output
344 scaled by ``topk_weights``.
345 """
346 from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import (
347 _grouped_gemm_mxfp8,
348 )
349 from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
350 moe_align_block_size,
351 )
352
353 torch.manual_seed(0)
354 block_m = 64
355 a_div = 1 if weighted else top_k
356 m_routed = T * top_k
357 # a-rows: g1 reads one row per token (a_div=top_k); g2 one per expansion.
358 a_rows = m_routed if weighted else T
359 a_bf16 = torch.randn(a_rows, K, device=DEVICE, dtype=torch.bfloat16) * 0.5
360 w_bf16 = torch.randn(E, N, K, device=DEVICE, dtype=torch.bfloat16) * 0.1
361 a_fp8, a_scale = _mxfp8_e4m3_quantize_torch(a_bf16, is_sf_swizzled_layout=False)
362 w_fp8, w_scale = _mxfp8_e4m3_quantize_torch(w_bf16, is_sf_swizzled_layout=False)
363
364 logits = torch.randn(T, E, device=DEVICE, dtype=torch.float32)
365 topk_weights, topk_ids = logits.softmax(dim=-1).topk(top_k, dim=-1)
366 topk_weights = topk_weights.to(torch.float32)
367 topk_ids = topk_ids.to(torch.int32)
368 mul = topk_weights.reshape(-1) if weighted else None
369
370 sorted_ids, expert_ids, num_post = moe_align_block_size(
371 topk_ids, block_m, E, None, ignore_invalid_experts=False
372 )
373 got = _grouped_gemm_mxfp8(
374 a_fp8,
375 a_scale,
376 w_fp8,
377 w_scale,
378 sorted_ids,
379 expert_ids,
380 num_post,
381 m_routed,
382 top_k,
383 block_m,
384 torch.bfloat16,
385 a_div=a_div,
386 mul_weight_by=mul,
387 )
388 # Reference: dequant the SAME bits the kernel reads, plain torch matmul.
389 a_deq = dequant_mxfp8_to_bf16(a_fp8, a_scale)
390 w_deq = dequant_mxfp8_to_bf16(w_fp8, w_scale)
391 ref = _ref_grouped_gemm(a_deq, w_deq, topk_ids, a_div, m_routed, mul)
392 assert got.shape == (m_routed, N)
393 assert _relerr(got, ref) < 5e-2
394
395
396 # --------------------------------------------------------------------------- #
397 # MXFP8 linear emulation: BF16-at-load (default) vs per-step dequant + switch
398 # --------------------------------------------------------------------------- #
399 @pytest.mark.parametrize("shape", [(512, 2048), (1, 6144)])
400 @pytest.mark.parametrize("act_dtype", [torch.bfloat16, torch.float16])
401 @pytest.mark.parametrize("dequant_at_load", [True, False])
402 @torch.inference_mode()
403 def test_mxfp8_linear_emulation_bf16_at_load(
404 shape, act_dtype, dequant_at_load, monkeypatch
405 ):
406 """EmulationMxfp8LinearKernel load-time BF16 dequant (default) and the
407 ``VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD=0`` per-step fallback must produce the
408 same result; the dtype-match (BF16/FP16 activations) must also hold."""
409 from vllm.model_executor.kernels.linear.mxfp8.emulation import (
410 EmulationMxfp8LinearKernel,
411 )
412 from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import (
413 Mxfp8LinearLayerConfig,
414 )
415
416 monkeypatch.setenv(
417 "VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD", "1" if dequant_at_load else "0"
418 )
419 N, K = shape
420 torch.manual_seed(0)
421 w_bf16 = torch.randn(N, K, device=DEVICE, dtype=torch.bfloat16)
422 w_fp8, w_scale = _mxfp8_e4m3_quantize_torch(w_bf16, is_sf_swizzled_layout=False)
423 assert w_scale.shape == (N, K // 32)
424
425 # Reference: dequant once, plain linear in the activation dtype.
426 w_ref = dequant_mxfp8_to_bf16(w_fp8, w_scale).to(act_dtype)
427 x = torch.randn(7, K, device=DEVICE, dtype=act_dtype)
428 out_ref = torch.nn.functional.linear(x, w_ref)
429
430 layer = torch.nn.Module()
431 layer.weight = torch.nn.Parameter(w_fp8.clone(), requires_grad=False)
432 layer.weight_scale = torch.nn.Parameter(w_scale.clone(), requires_grad=False)
433
434 kernel = EmulationMxfp8LinearKernel(Mxfp8LinearLayerConfig(weight_shape=(N, K)))
435 kernel.process_weights_after_loading(layer)
436
437 if dequant_at_load:
438 # weights converted to BF16 at load (>= 2-byte)
439 assert layer.weight.element_size() >= 2
440 else:
441 # opt-out: weights stay 1-byte MXFP8, dequant happens per-step
442 assert layer.weight.element_size() == 1
443
444 out = kernel.apply_weights(layer, x)
445 assert out.dtype == act_dtype # dtype-match preserved (no tl.dot/F.linear crash)
446 assert _relerr(out.float(), out_ref.float()) < 2e-2
447
448
449 # --------------------------------------------------------------------------- #
450 # Native MXFP8 linear: the K % 128 != 0 weight is dequantized once at load
451 # --------------------------------------------------------------------------- #
452 @requires_gfx950
453 # (5120, 576) and (5120, 288): DSv4.1 shared-expert down_proj at TP4 and TP8.
454 @pytest.mark.parametrize("shape", [(5120, 576), (5120, 288), (256, 2048)])
455 @torch.inference_mode()
456 def test_mxfp8_rocm_native_unaligned_k_dequantizes_at_load(shape):
457 """``dot_scaled`` tiles K by 128, so an unaligned weight must not reach it.
458
459 ``process_weights_after_loading`` converts those to BF16 once; a weight that
460 never went through it has to fall back per step rather than take the
461 invalid path, which is what keying the dispatch on the weight dtype alone
462 would do.
463 """
464 from vllm.model_executor.kernels.linear.mxfp8 import rocm_native
465 from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import (
466 Mxfp8LinearLayerConfig,
467 )
468
469 N, K = shape
470 aligned = K % 128 == 0
471 torch.manual_seed(0)
472 w_bf16 = torch.randn(N, K, device=DEVICE, dtype=torch.bfloat16) * 0.1
473 w_fp8, w_scale = _mxfp8_e4m3_quantize_torch(w_bf16, is_sf_swizzled_layout=False)
474 x = torch.randn(7, K, device=DEVICE, dtype=torch.bfloat16) * 0.5
475 out_ref = torch.nn.functional.linear(x, dequant_mxfp8_to_bf16(w_fp8, w_scale))
476 # dot_scaled re-quantizes the activation, the dequant paths do not.
477 tol = 5e-2 if aligned else 2e-3
478
479 def make_layer():
480 layer = torch.nn.Module()
481 layer.weight = torch.nn.Parameter(w_fp8.clone(), requires_grad=False)
482 layer.weight_scale = torch.nn.Parameter(w_scale.clone(), requires_grad=False)
483 return layer
484
485 kernel = rocm_native.RocmDotScaledMxfp8LinearKernel(
486 Mxfp8LinearLayerConfig(weight_shape=(N, K))
487 )
488 layer = make_layer()
489 kernel.process_weights_after_loading(layer)
490 assert (layer.weight.element_size() >= 2) is not aligned
491 assert _relerr(kernel.apply_weights(layer, x).float(), out_ref.float()) < tol
492
493 # An unaligned weight that is still FP8 must not reach dot_scaled.
494 unprocessed = make_layer()
495 unprocessed.weight_scale = torch.nn.Parameter(
496 w_scale[:N, : K // 32].contiguous(), requires_grad=False
497 )
498 calls = []
499 real = rocm_native._mxfp8_dot_scaled_linear
500
501 def _spy(*a, **k):
502 calls.append(1)
503 return real(*a, **k)
504
505 rocm_native._mxfp8_dot_scaled_linear = _spy
506 try:
507 out = kernel.apply_weights(unprocessed, x)
508 finally:
509 rocm_native._mxfp8_dot_scaled_linear = real
510 assert bool(calls) is aligned
511 assert _relerr(out.float(), out_ref.float()) < tol
512
513
514 # ── EP expert_mask handling for the FlyDSL (AITER_MXFP8) MoE ────────────────
515 # The map/mask choice lives in ``RoutedExperts.expert_map``: it hands AITER
516 # experts (``consumes_expert_mask``) the precomputed 0/1 ``expert_mask`` and
517 # everyone else the canonical -1 index map, keyed on the resolved experts kernel
518 # rather than the global VLLM_ROCM_USE_AITER switch. ``AiterMxfp8Experts.apply``
519 # forwards that mask to aiter unchanged (it no longer rebuilds it via
520 # ``(expert_map >= 0)``, which collapsed an already-0/1 mask to all-ones and
521 # produced EP garbage).
522 def _capture_expert_mask(expert_mask, *, global_num_experts):
523 """Drive the real ``AiterMxfp8Experts.apply`` mask branch and capture the
524 ``expert_mask`` it forwards to ``rocm_aiter_ops.fused_moe``."""
525 from types import SimpleNamespace
526 from unittest import mock
527
528 from vllm._aiter_ops import rocm_aiter_ops
529 from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp8_moe import (
530 AiterMxfp8Experts,
531 )
532
533 experts = object.__new__(AiterMxfp8Experts) # bypass heavy __init__
534 experts.quant_config = SimpleNamespace(gemm1_clamp_limit=None)
535 experts.w1_scale_val = None
536 experts.w2_scale_val = None
537
538 captured = {}
539
540 def _fake_fused_moe(hidden_states, w1, w2, tw, ti, *, expert_mask, **kw):
541 captured["expert_mask"] = expert_mask
542 return torch.zeros_like(hidden_states)
543
544 w1 = torch.zeros(1, device=DEVICE)
545 w2 = torch.zeros(1, device=DEVICE)
546 out = torch.zeros(4, 8, device=DEVICE, dtype=torch.bfloat16)
547 hidden = torch.zeros(4, 8, device=DEVICE, dtype=torch.bfloat16)
548 tw = torch.ones(4, 2, device=DEVICE)
549 ti = torch.zeros(4, 2, dtype=torch.int32, device=DEVICE)
550
551 with mock.patch.object(rocm_aiter_ops, "fused_moe", side_effect=_fake_fused_moe):
552 experts.apply(
553 output=out,
554 hidden_states=hidden,
555 w1=w1,
556 w2=w2,
557 topk_weights=tw,
558 topk_ids=ti,
559 activation=None,
560 global_num_experts=global_num_experts,
561 expert_map=expert_mask,
562 a1q_scale=None,
563 a2_scale=None,
564 workspace13=None,
565 workspace2=None,
566 expert_tokens_meta=None,
567 apply_router_weight_on_input=False,
568 )
569 return captured["expert_mask"]
570
571
572 @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only")
573 def test_aiter_mxfp8_apply_forwards_expert_mask_unchanged():
574 """AiterMxfp8Experts.apply forwards the precomputed 0/1 mask to aiter as-is
575 (no re-derivation that would collapse an already-0/1 mask to all-ones)."""
576 # 0/1 local-expert mask over global ids + trailing sentinel (rank owns 0..3).
577 ep_mask = torch.tensor(
578 [1, 1, 1, 1, 0, 0, 0, 0, 0], dtype=torch.int32, device=DEVICE
579 )
580 got = _capture_expert_mask(ep_mask, global_num_experts=8)
581 assert torch.equal(got.cpu().to(torch.int32), ep_mask.cpu().to(torch.int32))
582 assert got.sum().item() == 4 # the 4 local experts, not all 9
583
584
585 @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only")
586 def test_routed_experts_expert_map_delegates_to_kernel():
587 """RoutedExperts.expert_map returns the 0/1 mask only for kernels that set
588 ``consumes_expert_mask`` (AITER), and the canonical -1 map otherwise -- keyed
589 on the resolved kernel, not the global aiter switch."""
590 from types import SimpleNamespace
591
592 from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
593
594 canonical = torch.tensor([0, 1, 2, 3, -1, -1, -1, -1], dtype=torch.int32)
595 mask = torch.tensor([1, 1, 1, 1, 0, 0, 0, 0, 0], dtype=torch.int32)
596
597 def resolve(consumes_mask, *, has_moe_kernel=True):
598 moe_kernel = (
599 SimpleNamespace(
600 fused_experts=SimpleNamespace(consumes_expert_mask=consumes_mask)
601 )
602 if has_moe_kernel
603 else None
604 )
605 layer = SimpleNamespace(
606 _expert_map=canonical,
607 expert_mask=mask,
608 quant_method=SimpleNamespace(moe_kernel=moe_kernel),
609 )
610 return RoutedExperts.expert_map.fget(layer)
611
612 assert torch.equal(resolve(True), mask) # AITER kernel -> 0/1 mask
613 assert torch.equal(resolve(False), canonical) # non-AITER -> canonical map
614 assert torch.equal(resolve(False, has_moe_kernel=False), canonical) # non-modular
615
616
617 # --------------------------------------------------------------------------- #
618 # Strided K/V operand for the AITER sparse-PA cache insert
619 # --------------------------------------------------------------------------- #
620 def _fused_qkv_kv_slices(num_tokens, num_kv_heads, head_dim, num_q_heads=16):
621 """K/V as the sparse-PA insert sees them: column slices of a fused qkv."""
622 kv = num_kv_heads * head_dim
623 row = num_q_heads * head_dim + 2 * kv + 2 * head_dim
624 qkv = torch.randn((num_tokens, row), dtype=torch.bfloat16, device=DEVICE)
625 k_start = num_q_heads * head_dim
626 k = qkv[:, k_start : k_start + kv].view(num_tokens, num_kv_heads, head_dim)
627 v = qkv[:, k_start + kv : k_start + 2 * kv].view(num_tokens, num_kv_heads, head_dim)
628 return k, v
629
630
631 def _asm_kv_cache(pages, num_kv_heads, head_dim):
632 """Page-16 K/V cache views matching ``_ensure_aiter_sparse_pa_kv_cache``."""
633 x = 16 // torch.tensor([], dtype=torch.bfloat16).element_size()
634 key = torch.zeros(
635 (pages, num_kv_heads, head_dim // x, 16, x),
636 dtype=torch.bfloat16,
637 device=DEVICE,
638 )
639 value = torch.zeros(
640 (pages, num_kv_heads, 16 // x, head_dim, x),
641 dtype=torch.bfloat16,
642 device=DEVICE,
643 )
644 return key, value
645
646
647 @pytest.mark.parametrize("num_tokens", [1, 8, 129, 1024])
648 @pytest.mark.parametrize("num_kv_heads", [1, 2])
649 def test_kv_insert_operand_matches_contiguous(num_tokens, num_kv_heads):
650 """A strided K/V slice must insert bit-identically to a contiguous copy."""
651 reshape_and_cache = pytest.importorskip("aiter").reshape_and_cache
652 head_dim = 128
653 pages = num_tokens // 16 + 4
654 slots = torch.randperm(pages * 16, device=DEVICE)[:num_tokens].to(torch.int32)
655
656 torch.manual_seed(num_tokens)
657 k, v = _fused_qkv_kv_slices(num_tokens, num_kv_heads, head_dim)
658 assert num_tokens == 1 or not k.is_contiguous()
659
660 caches = []
661 for operand in (lambda t: t.contiguous(), _kv_insert_operand):
662 key_cache, value_cache = _asm_kv_cache(pages, num_kv_heads, head_dim)
663 reshape_and_cache(
664 operand(k),
665 operand(v),
666 key_cache,
667 value_cache,
668 slots,
669 kv_cache_dtype="auto",
670 asm_layout=True,
671 )
672 caches.append((key_cache, value_cache))
673
674 (ref_k, ref_v), (got_k, got_v) = caches
675 assert torch.equal(ref_k, got_k)
676 assert torch.equal(ref_v, got_v)
677
678
679 def test_kv_insert_operand_copies_unsupported_layouts():
680 """Only a row-strided slice passes through; anything else is copied."""
681 k, _ = _fused_qkv_kv_slices(32, 2, 128)
682 assert _kv_insert_operand(k) is k
683
684 inner_strided = k[:, :, ::2]
685 assert _kv_insert_operand(inner_strided).is_contiguous()
686
687 flat = k.reshape(32, -1)
688 assert _kv_insert_operand(flat).is_contiguous()