@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 224 lines 7.87 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Correctness tests for the persistent cp.async-pipelined SITU + block-FP8
4 quant kernel (``situ_and_mul_quant``).
5
6 Mirrors ``test_silu_mul_fp8_quant_deep_gemm.py`` in structure: build a reference
7 activation + block-FP8 quant, drive random per-row scale magnitudes, exercise
8 the masked contiguous layout, and compare per valid row.
9
10 The kernel computes SITU in fp32 and converts to fp8 with hardware rounding,
11 which torch cannot reproduce bit-exactly. It is therefore validated against the
12 *ideal* FP8 quantizer of the exact fp64 activation: the difference must stay
13 within the unavoidable FP8 block-quant error (within 1 fp8 code for ~all
14 elements, never off by more than 2).
15 """
16
17 import pytest
18 import torch
19
20 from vllm.model_executor.layers.fused_moe.activation import situ_and_mul_quant
21 from vllm.platforms import current_platform
22 from vllm.utils.torch_utils import set_random_seed
23
24 fp8_dtype = current_platform.fp8_dtype()
25 DEVICE = current_platform.device_type
26 FP8_MAX = torch.finfo(fp8_dtype).max
27
28 # The kernel uses fp8 e4m3 conversion, so it is CUDA-only and needs SM89+
29 # (Ada/Hopper). There is no Triton fallback.
30 pytestmark = pytest.mark.skipif(
31 not current_platform.is_cuda() or not current_platform.has_device_capability(89),
32 reason="situ_and_mul_quant persistent kernel requires CUDA SM89+ (fp8 e4m3).",
33 )
34
35 # Kimi-K3 baked SITU params. Only (d=3072, beta=4.0, linear_beta=25.0,
36 # group_size=128, fp16/bf16) routes to the persistent cp.async-pipelined kernel;
37 # anything else falls back to the scalar group kernel.
38 SITU_BETA = 4.0
39 SITU_LINEAR_BETA = 25.0
40 SITU_D = 3072
41 GROUP_SIZE = 128
42
43 # (num_tokens, num_valid_tokens). None means all rows are valid. 2048/4096 rows
44 # span multiple grid-stride waves of the persistent grid (GRID_DIM = 132*8).
45 CASES = [
46 (1, None),
47 (17, None),
48 (128, None),
49 (2048, None),
50 (4096, None),
51 (2048, 1500), # DeepEP v2 contiguous padding
52 (2048, 1),
53 (4096, 4095),
54 (512, 0), # whole tensor is padding
55 ]
56
57
58 def token_random(num_tokens: int, twod: int, dtype: torch.dtype) -> torch.Tensor:
59 """Per-row random exponent so groups span many scale magnitudes."""
60 base = torch.randn(num_tokens, twod, dtype=torch.float32, device=DEVICE)
61 exps = torch.randint(1, 13, (num_tokens, 1), device=DEVICE).float()
62 return (base * (2.0**exps)).to(dtype)
63
64
65 def ref_situ_block_fp8_quant(
66 inp: torch.Tensor, beta: float, linear_beta: float, group_size: int
67 ) -> tuple[torch.Tensor, torch.Tensor]:
68 """Ideal FP8 block-quant of the exact fp64 SITU activation.
69
70 The kernel rounds the activation to ``scalar_t`` before quantizing (it stores
71 ``(float)(scalar_t)situ(...)``), so the reference does the same before taking
72 the per-group absmax.
73 """
74 d = inp.shape[-1] // 2
75 gate, up = inp.double().chunk(2, dim=-1)
76 gate_out = beta * torch.tanh(gate / beta) * torch.sigmoid(gate)
77 up_out = linear_beta * torch.tanh(up / linear_beta) if linear_beta > 0 else up
78 act = (gate_out * up_out).to(inp.dtype).float()
79
80 m = act.shape[0]
81 ng = d // group_size
82 a = act.view(m, ng, group_size)
83 scale = a.abs().amax(dim=2).clamp_min(1e-30) / FP8_MAX
84 q = (a / scale.unsqueeze(2)).clamp(-FP8_MAX, FP8_MAX).to(fp8_dtype)
85 return q.view(m, d), scale
86
87
88 def fp8_code_dist(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
89 """Signed-magnitude fp8 code distance (valid for e4m3fn and e4m3fnuz)."""
90
91 def ordered(x: torch.Tensor) -> torch.Tensor:
92 x = x.view(torch.uint8).to(torch.int32)
93 return torch.where((x >> 7) == 1, -(x & 0x7F), x & 0x7F)
94
95 return (ordered(a) - ordered(b)).abs()
96
97
98 @pytest.mark.parametrize("num_tokens,valid", CASES)
99 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half])
100 @torch.inference_mode()
101 def test_situ_and_mul_quant_pipelined(
102 num_tokens: int, valid: int | None, dtype: torch.dtype
103 ) -> None:
104 set_random_seed(42)
105 d = SITU_D
106 inp = token_random(num_tokens, 2 * d, dtype)
107
108 # Sentinels: padding rows must keep their bytes (out) and get scale 1.0.
109 out = torch.full((num_tokens, d), 17.0, device=DEVICE).to(fp8_dtype)
110 scale = torch.full(
111 (num_tokens, d // GROUP_SIZE), -1.0, dtype=torch.float32, device=DEVICE
112 )
113 out_sentinel = out.clone()
114 num_valid_tokens = (
115 None if valid is None else torch.tensor(valid, dtype=torch.int32, device=DEVICE)
116 )
117
118 situ_and_mul_quant(
119 out,
120 scale,
121 inp,
122 beta=SITU_BETA,
123 linear_beta=SITU_LINEAR_BETA,
124 group_size=GROUP_SIZE,
125 num_valid_tokens=num_valid_tokens,
126 )
127
128 # Padding contract: skipped rows keep their bytes; their scale becomes 1.0.
129 if valid is not None and valid < num_tokens:
130 assert torch.equal(scale[valid:], torch.ones_like(scale[valid:]))
131 assert torch.equal(
132 out[valid:].view(torch.uint8), out_sentinel[valid:].view(torch.uint8)
133 )
134
135 rows = num_tokens if valid is None else valid
136 if rows == 0:
137 return
138
139 ref_q, ref_s = ref_situ_block_fp8_quant(
140 inp[:rows], SITU_BETA, SITU_LINEAR_BETA, GROUP_SIZE
141 )
142
143 assert torch.isfinite(scale[:rows]).all()
144 assert not torch.isnan(out[:rows].float()).any()
145
146 # Scales: the kernel's fp32 SITU shifts the per-group absmax only
147 # marginally vs the fp64 reference.
148 torch.testing.assert_close(scale[:rows], ref_s, rtol=3e-2, atol=1e-5)
149
150 # FP8 codes: must land within 1 code of the ideal quantizer for ~all
151 # elements, never off by more than 2.
152 dist = fp8_code_dist(out[:rows], ref_q)
153 assert dist.max().item() <= 2
154 assert (dist <= 1).float().mean().item() >= 0.99
155
156
157 @pytest.mark.parametrize("tokens,topk", [(300, 8), (1, 4), (256, 6)])
158 @torch.inference_mode()
159 def test_situ_and_mul_quant_topk_row_bound(tokens: int, topk: int) -> None:
160 """The kernel expands num_valid_tokens by topk to bound rows."""
161 set_random_seed(7)
162 d = SITU_D
163 num_tokens = tokens * topk + 37 # padding tail past the valid rows
164 inp = token_random(num_tokens, 2 * d, torch.bfloat16)
165
166 out = torch.full((num_tokens, d), 17.0, device=DEVICE).to(fp8_dtype)
167 scale = torch.full(
168 (num_tokens, d // GROUP_SIZE), -1.0, dtype=torch.float32, device=DEVICE
169 )
170 out_sentinel = out.clone()
171 num_valid_tokens = torch.tensor(tokens, dtype=torch.int32, device=DEVICE)
172
173 situ_and_mul_quant(
174 out,
175 scale,
176 inp,
177 beta=SITU_BETA,
178 linear_beta=SITU_LINEAR_BETA,
179 group_size=GROUP_SIZE,
180 num_valid_tokens=num_valid_tokens,
181 topk=topk,
182 )
183
184 rows = tokens * topk
185 assert torch.equal(scale[rows:], torch.ones_like(scale[rows:]))
186 assert torch.equal(
187 out[rows:].view(torch.uint8), out_sentinel[rows:].view(torch.uint8)
188 )
189
190 ref_q, ref_s = ref_situ_block_fp8_quant(
191 inp[:rows], SITU_BETA, SITU_LINEAR_BETA, GROUP_SIZE
192 )
193 torch.testing.assert_close(scale[:rows], ref_s, rtol=3e-2, atol=1e-5)
194 assert fp8_code_dist(out[:rows], ref_q).max().item() <= 2
195
196
197 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half])
198 @torch.inference_mode()
199 def test_situ_and_mul_quant_pipelined_deterministic(dtype: torch.dtype) -> None:
200 """Two runs on the same input must be byte-identical."""
201 set_random_seed(0)
202 d = SITU_D
203 num_tokens = 2048
204 inp = token_random(num_tokens, 2 * d, dtype)
205
206 def run() -> tuple[torch.Tensor, torch.Tensor]:
207 out = torch.empty(num_tokens, d, dtype=fp8_dtype, device=DEVICE)
208 scale = torch.empty(
209 num_tokens, d // GROUP_SIZE, dtype=torch.float32, device=DEVICE
210 )
211 situ_and_mul_quant(
212 out,
213 scale,
214 inp,
215 beta=SITU_BETA,
216 linear_beta=SITU_LINEAR_BETA,
217 group_size=GROUP_SIZE,
218 )
219 return out, scale
220
221 out1, s1 = run()
222 out2, s2 = run()
223 assert torch.equal(out1.view(torch.uint8), out2.view(torch.uint8))
224 assert torch.equal(s1, s2)