| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Tests for SM100 CUTLASS MXFP4 x MXFP4 grouped MoE kernels.""" |
| 5 | |
| 6 | import random |
| 7 | |
| 8 | import pytest |
| 9 | import torch |
| 10 | |
| 11 | from tests.kernels.utils import torch_moe_single |
| 12 | from vllm import _custom_ops as ops |
| 13 | from vllm.platforms import current_platform |
| 14 | from vllm.utils.torch_utils import set_random_seed |
| 15 | |
| 16 | random.seed(42) |
| 17 | set_random_seed(42) |
| 18 | |
| 19 | MXFP4_BLOCK_SIZE = 32 |
| 20 | |
| 21 | |
| 22 | def align(val: int, alignment: int = 128) -> int: |
| 23 | return int((val + alignment - 1) // alignment * alignment) |
| 24 | |
| 25 | |
| 26 | def calc_diff(x, y): |
| 27 | x, y = x.double(), y.double() |
| 28 | denominator = (x * x + y * y).sum() |
| 29 | sim = 2 * (x * y).sum() / denominator |
| 30 | return 1 - sim |
| 31 | |
| 32 | |
| 33 | def is_sm100_supported() -> bool: |
| 34 | return current_platform.is_cuda() and current_platform.is_device_capability_family( |
| 35 | 100 |
| 36 | ) |
| 37 | |
| 38 | |
| 39 | def compute_ref_output( |
| 40 | input_tensor: torch.Tensor, |
| 41 | weight_list: list[torch.Tensor], |
| 42 | expert_offsets: list[int], |
| 43 | expert_offset: int, |
| 44 | num_experts: int, |
| 45 | ) -> torch.Tensor: |
| 46 | """Reference output using torch_moe_single with top-1 routing.""" |
| 47 | score = torch.full( |
| 48 | (expert_offset, num_experts), |
| 49 | -1e9, |
| 50 | device=input_tensor.device, |
| 51 | dtype=torch.float32, |
| 52 | ) |
| 53 | for g in range(num_experts): |
| 54 | start = expert_offsets[g] |
| 55 | end = expert_offsets[g + 1] if g + 1 < num_experts else expert_offset |
| 56 | score[start:end, g] = 0.0 |
| 57 | |
| 58 | return torch_moe_single( |
| 59 | input_tensor, torch.stack(weight_list, dim=0), score, topk=1 |
| 60 | ) |
| 61 | |
| 62 | |
| 63 | @pytest.mark.skipif( |
| 64 | not is_sm100_supported(), |
| 65 | reason="cutlass_mxfp4_group_mm requires CUDA SM100", |
| 66 | ) |
| 67 | @pytest.mark.parametrize("num_experts", [8, 16, 32]) |
| 68 | @pytest.mark.parametrize("out_dtype", [torch.bfloat16]) |
| 69 | def test_cutlass_mxfp4_grouped_mm(num_experts, out_dtype): |
| 70 | """Test the MXFP4 grouped GEMM kernel by: |
| 71 | 1. Creating random per-expert inputs and weights |
| 72 | 2. Quantizing both to MXFP4 using the CUDA kernel |
| 73 | 3. Running the CUTLASS grouped GEMM |
| 74 | 4. Comparing against BF16 reference |
| 75 | """ |
| 76 | device = "cuda" |
| 77 | alignment = 128 |
| 78 | # N and K must be multiples of 128 for clean swizzle layout |
| 79 | n_g = random.randint(1, 16) * alignment |
| 80 | k_g = random.randint(1, 16) * alignment |
| 81 | |
| 82 | expert_offset = 0 |
| 83 | expert_offsets_input = [] |
| 84 | problem_sizes = [] |
| 85 | input_list = [] |
| 86 | weight_list = [] |
| 87 | |
| 88 | for g in range(num_experts): |
| 89 | m_g = random.randint(1, 256) |
| 90 | expert_offsets_input.append(expert_offset) |
| 91 | expert_offset += m_g |
| 92 | problem_sizes.append([m_g, n_g, k_g]) |
| 93 | |
| 94 | input_list.append( |
| 95 | torch.normal(0.0, std=0.5, size=(m_g, k_g), device=device, dtype=out_dtype) |
| 96 | ) |
| 97 | weight_list.append( |
| 98 | torch.normal(0.0, std=0.5, size=(n_g, k_g), device=device, dtype=out_dtype) |
| 99 | ) |
| 100 | |
| 101 | input_tensor = torch.concat(input_list, dim=0) # [M_total, K] |
| 102 | |
| 103 | # --- Quantize INPUTS via mxfp4_experts_quant --- |
| 104 | input_bs_offsets = [] |
| 105 | tot = 0 |
| 106 | for g in range(num_experts): |
| 107 | input_bs_offsets.append(tot) |
| 108 | tot += align(problem_sizes[g][0], 128) |
| 109 | input_bs_offsets.append(tot) |
| 110 | |
| 111 | _inp_expert_offsets = torch.tensor( |
| 112 | expert_offsets_input + [expert_offset], device=device, dtype=torch.int32 |
| 113 | ) |
| 114 | _inp_bs_offsets = torch.tensor(input_bs_offsets, device=device, dtype=torch.int32) |
| 115 | |
| 116 | input_quant, input_sf = ops.mxfp4_experts_quant( |
| 117 | input_tensor, |
| 118 | _inp_expert_offsets, |
| 119 | _inp_bs_offsets, |
| 120 | num_experts, |
| 121 | topk=1, |
| 122 | ) |
| 123 | |
| 124 | # --- Quantize WEIGHTS via mxfp4_experts_quant --- |
| 125 | # Treat each expert's N weight rows as an "expert" with N tokens |
| 126 | weight_tensor = torch.concat(weight_list, dim=0) # [E*N, K] |
| 127 | weight_expert_offsets = [g * n_g for g in range(num_experts)] + [num_experts * n_g] |
| 128 | # N is always multiple of 128, so blockscale offsets are clean |
| 129 | weight_bs_offsets = [g * n_g for g in range(num_experts)] + [num_experts * n_g] |
| 130 | |
| 131 | _wt_expert_offsets = torch.tensor( |
| 132 | weight_expert_offsets, device=device, dtype=torch.int32 |
| 133 | ) |
| 134 | _wt_bs_offsets = torch.tensor(weight_bs_offsets, device=device, dtype=torch.int32) |
| 135 | |
| 136 | weight_quant, weight_sf = ops.mxfp4_experts_quant( |
| 137 | weight_tensor, |
| 138 | _wt_expert_offsets, |
| 139 | _wt_bs_offsets, |
| 140 | num_experts, |
| 141 | topk=1, |
| 142 | ) |
| 143 | |
| 144 | # Reshape weight quantized data to [E, N, K//2] |
| 145 | weight_quant = weight_quant[: num_experts * n_g].view(num_experts, n_g, k_g // 2) |
| 146 | |
| 147 | # Reshape weight scale factors to [E, N, K//32] |
| 148 | # The quant kernel produces uint8 SF buffer. Each row has K//32 SFs. |
| 149 | scales_per_row = k_g // MXFP4_BLOCK_SIZE |
| 150 | weight_sf_flat = weight_sf.view(-1)[: num_experts * n_g * scales_per_row] |
| 151 | weight_sf_3d = weight_sf_flat.view(num_experts, n_g, scales_per_row) |
| 152 | |
| 153 | # Output |
| 154 | output = torch.empty((expert_offset, n_g), device=device, dtype=out_dtype) |
| 155 | |
| 156 | _problem_sizes = torch.tensor(problem_sizes, device=device, dtype=torch.int32) |
| 157 | _expert_offsets = torch.tensor( |
| 158 | expert_offsets_input, device=device, dtype=torch.int32 |
| 159 | ) |
| 160 | _input_bs = torch.tensor(input_bs_offsets[:-1], device=device, dtype=torch.int32) |
| 161 | |
| 162 | # Run the MXFP4 grouped GEMM |
| 163 | ops.cutlass_mxfp4_moe_mm( |
| 164 | output, |
| 165 | input_quant, |
| 166 | weight_quant, |
| 167 | input_sf, |
| 168 | weight_sf_3d, |
| 169 | _problem_sizes, |
| 170 | _expert_offsets, |
| 171 | _input_bs, |
| 172 | ) |
| 173 | |
| 174 | # Reference: BF16 matmul |
| 175 | ref_output = compute_ref_output( |
| 176 | input_tensor=input_tensor, |
| 177 | weight_list=weight_list, |
| 178 | expert_offsets=expert_offsets_input, |
| 179 | expert_offset=expert_offset, |
| 180 | num_experts=num_experts, |
| 181 | ) |
| 182 | |
| 183 | # Compare per-expert |
| 184 | for g in range(num_experts): |
| 185 | start = expert_offsets_input[g] |
| 186 | end = expert_offsets_input[g + 1] if g + 1 < num_experts else expert_offset |
| 187 | if start == end: |
| 188 | continue |
| 189 | baseline = ref_output[start:end] |
| 190 | actual = output[start:end] |
| 191 | diff = calc_diff(actual, baseline) |
| 192 | print( |
| 193 | f"m_g={end - start} n_g={n_g} k_g={k_g} " |
| 194 | f"num_experts={num_experts}, " |
| 195 | f"out_dtype={out_dtype}, diff={diff:.5f}" |
| 196 | ) |
| 197 | # FP4 quantization is very lossy (~4 bits precision) |
| 198 | # Comparing quantized vs full-precision gives cosine diff of 0.05-0.15 |
| 199 | assert diff < 0.15, f"Expert {g}: diff={diff:.5f} exceeds threshold" |
| 200 | |
| 201 | |
| 202 | @pytest.mark.skipif( |
| 203 | not is_sm100_supported(), |
| 204 | reason="mxfp4_experts_quant requires CUDA SM100", |
| 205 | ) |
| 206 | def test_mxfp4_experts_quant_basic(): |
| 207 | """Basic smoke test for the MXFP4 experts quantization kernel.""" |
| 208 | device = "cuda" |
| 209 | num_experts = 4 |
| 210 | k = 256 |
| 211 | tokens_per_expert = 16 |
| 212 | |
| 213 | total_tokens = tokens_per_expert * num_experts |
| 214 | input_tensor = torch.randn(total_tokens, k, device=device, dtype=torch.bfloat16) / 5 |
| 215 | |
| 216 | expert_offsets = [i * tokens_per_expert for i in range(num_experts + 1)] |
| 217 | blockscale_offsets = [ |
| 218 | align(i * tokens_per_expert, 128) for i in range(num_experts + 1) |
| 219 | ] |
| 220 | |
| 221 | _expert_offsets = torch.tensor(expert_offsets, device=device, dtype=torch.int32) |
| 222 | _blockscale_offsets = torch.tensor( |
| 223 | blockscale_offsets, device=device, dtype=torch.int32 |
| 224 | ) |
| 225 | |
| 226 | output, output_sf = ops.mxfp4_experts_quant( |
| 227 | input_tensor, |
| 228 | _expert_offsets, |
| 229 | _blockscale_offsets, |
| 230 | num_experts, |
| 231 | topk=1, |
| 232 | ) |
| 233 | |
| 234 | assert output.shape == (total_tokens, k // 2) |
| 235 | assert output.dtype == torch.uint8 |
| 236 | assert output_sf.dtype == torch.uint8 |
| 237 | assert output.any(), "Quantized output is all zeros" |
| 238 | print( |
| 239 | f"MXFP4 experts quant: output shape={output.shape}, sf shape={output_sf.shape}" |
| 240 | ) |
| 241 | print("PASSED") |
| 242 | |
| 243 | |
| 244 | def untile_cutlass_scale(scale_raw: torch.Tensor, rows: int, K: int) -> torch.Tensor: |
| 245 | """Convert CUTLASS tiled scale back to flat [M, K//32] layout. |
| 246 | |
| 247 | CUTLASS tiled layout: [numMTiles, numKTiles, 32(outerM), 4(innerM), 4(innerK)] |
| 248 | Produced by: padded.reshape(numMTiles, 4, 32, numKTiles, 4).permute(0,3,2,1,4) |
| 249 | To undo: tiled.permute(0, 3, 2, 1, 4).reshape(padded_M, padded_sK) |
| 250 | """ |
| 251 | num_scale_cols = K // MXFP4_BLOCK_SIZE |
| 252 | num_m_tiles = (rows + 127) // 128 |
| 253 | num_k_tiles = (num_scale_cols + 3) // 4 |
| 254 | padded_M = num_m_tiles * 128 |
| 255 | padded_sK = num_k_tiles * 4 |
| 256 | |
| 257 | scale_bytes = scale_raw.view(torch.uint8).flatten() |
| 258 | total_bytes = padded_M * padded_sK |
| 259 | tiled = scale_bytes[:total_bytes].reshape(num_m_tiles, num_k_tiles, 32, 4, 4) |
| 260 | undone = tiled.permute(0, 3, 2, 1, 4).contiguous() |
| 261 | return undone.reshape(padded_M, padded_sK)[:rows, :num_scale_cols] |
| 262 | |
| 263 | |
| 264 | def compute_reference_e8m0_scale(block_max: float) -> int: |
| 265 | """Compute the expected OCP MX spec E8M0 scale for a given block max. |
| 266 | |
| 267 | The CUTLASS kernel uses round-to-nearest on the mantissa: |
| 268 | rounded_bits = (float_bits + (1 << 21)) & 0xFF800000 |
| 269 | biased_exp = (rounded_bits >> 23) & 0xFF |
| 270 | scale_exp = max(biased_exp - 2, 0) |
| 271 | |
| 272 | This ensures max_val / scale <= 6.0 for most inputs. |
| 273 | """ |
| 274 | import struct |
| 275 | |
| 276 | if block_max <= 0: |
| 277 | return 0 |
| 278 | # Replicate the kernel's rounding logic in Python |
| 279 | float_bytes = struct.pack("f", block_max) |
| 280 | max_bits = struct.unpack("I", float_bytes)[0] |
| 281 | rounded_bits = (max_bits + (1 << 21)) & 0xFF800000 |
| 282 | biased_exp = (rounded_bits >> 23) & 0xFF |
| 283 | scale_exp = max(int(biased_exp) - 2, 0) |
| 284 | scale_exp = min(scale_exp, 254) |
| 285 | return scale_exp |
| 286 | |
| 287 | |
| 288 | @pytest.mark.skipif( |
| 289 | not is_sm100_supported(), |
| 290 | reason="mxfp4_experts_quant requires CUDA SM100", |
| 291 | ) |
| 292 | @pytest.mark.parametrize("k", [256, 7168]) |
| 293 | @pytest.mark.parametrize("m", [16, 64]) |
| 294 | def test_mxfp4_experts_quant_e8m0_scale_correctness(m, k): |
| 295 | """Test that mxfp4_experts_quant computes E8M0 block scales correctly |
| 296 | per OCP MX spec (not the NVFP4 formula). |
| 297 | |
| 298 | The old buggy kernel used: floor(log2(max/6)) + 127 |
| 299 | The fixed kernel uses: round_nearest_exp(max) - 2 |
| 300 | |
| 301 | This test verifies: |
| 302 | 1. Scales match the expected OCP MX formula for all blocks |
| 303 | 2. No block max exceeds the representable range (no unexpected saturation) |
| 304 | 3. Reconstruction error is within expected bounds for MXFP4 |
| 305 | """ |
| 306 | device = "cuda" |
| 307 | |
| 308 | # Generate input with controlled range |
| 309 | input_tensor = torch.randn(m, k, device=device, dtype=torch.bfloat16) * 0.5 |
| 310 | |
| 311 | # Quantize |
| 312 | num_experts = 1 |
| 313 | expert_offsets = torch.tensor([0, m], device=device, dtype=torch.int32) |
| 314 | num_k_tiles = (k // MXFP4_BLOCK_SIZE + 3) // 4 |
| 315 | blockscale_offsets = torch.tensor( |
| 316 | [0, align(m, 128) * num_k_tiles], device=device, dtype=torch.int32 |
| 317 | ) |
| 318 | |
| 319 | output_fp4, output_sf = ops.mxfp4_experts_quant( |
| 320 | input_tensor, expert_offsets, blockscale_offsets, num_experts, topk=1 |
| 321 | ) |
| 322 | |
| 323 | # Untile scale to flat layout for verification |
| 324 | scale_flat = untile_cutlass_scale(output_sf, m, k) |
| 325 | assert scale_flat.shape == (m, k // MXFP4_BLOCK_SIZE) |
| 326 | |
| 327 | # Verify each block's scale matches the OCP MX spec formula |
| 328 | num_blocks = k // MXFP4_BLOCK_SIZE |
| 329 | mismatches = 0 |
| 330 | buggy_pattern = 0 # count blocks where scale is 1-2 lower than expected |
| 331 | |
| 332 | for row in range(m): |
| 333 | for blk in range(num_blocks): |
| 334 | block_start = blk * MXFP4_BLOCK_SIZE |
| 335 | block_end = block_start + MXFP4_BLOCK_SIZE |
| 336 | block_max = ( |
| 337 | input_tensor[row, block_start:block_end].float().abs().max().item() |
| 338 | ) |
| 339 | |
| 340 | actual_scale = scale_flat[row, blk].item() |
| 341 | expected_scale = compute_reference_e8m0_scale(block_max) |
| 342 | |
| 343 | if actual_scale != expected_scale: |
| 344 | mismatches += 1 |
| 345 | if actual_scale < expected_scale: |
| 346 | buggy_pattern += 1 |
| 347 | |
| 348 | total_blocks = m * num_blocks |
| 349 | match_rate = (total_blocks - mismatches) / total_blocks |
| 350 | |
| 351 | print( |
| 352 | f" m={m}, k={k}: scale match rate = {match_rate * 100:.2f}% " |
| 353 | f"({mismatches}/{total_blocks} mismatches)" |
| 354 | ) |
| 355 | |
| 356 | # The fixed kernel should match the reference formula exactly |
| 357 | assert match_rate > 0.99, ( |
| 358 | f"E8M0 scale match rate too low: {match_rate * 100:.2f}%. " |
| 359 | f"Buggy pattern (scale too low): {buggy_pattern}/{mismatches}. " |
| 360 | f"This suggests the NVFP4 formula bug is present." |
| 361 | ) |
| 362 | |
| 363 | # Extra check: if most mismatches show scale < expected, it's the old bug |
| 364 | if mismatches > 0: |
| 365 | assert buggy_pattern / mismatches < 0.5, ( |
| 366 | f"Most scale mismatches show scale too LOW ({buggy_pattern}/{mismatches}). " |
| 367 | "This is the signature of the NVFP4 formula bug in nvfp4_utils.cuh." |
| 368 | ) |
| 369 | |
| 370 | # Verify reconstruction error is within MXFP4 expected bounds |
| 371 | # Dequantize and check cosine similarity |
| 372 | fp4_lut = torch.tensor( |
| 373 | [0, 0.5, 1, 1.5, 2, 3, 4, 6, 0, -0.5, -1, -1.5, -2, -3, -4, -6], |
| 374 | device=device, |
| 375 | dtype=torch.float32, |
| 376 | ) |
| 377 | lo = (output_fp4 & 0x0F).long() |
| 378 | hi = ((output_fp4 >> 4) & 0x0F).long() |
| 379 | unpacked = torch.stack([lo, hi], dim=-1).reshape(m, k) |
| 380 | fp4_vals = fp4_lut[unpacked] |
| 381 | |
| 382 | scales_expanded = 2.0 ** (scale_flat.float() - 127.0) |
| 383 | scales_expanded = scales_expanded.unsqueeze(-1).expand(-1, -1, MXFP4_BLOCK_SIZE) |
| 384 | scales_expanded = scales_expanded.reshape(m, k) |
| 385 | recon = (fp4_vals * scales_expanded).bfloat16() |
| 386 | |
| 387 | # Cosine similarity should be > 0.99 for well-behaved MXFP4 quantization |
| 388 | cos_sim = torch.nn.functional.cosine_similarity( |
| 389 | recon.float().flatten().unsqueeze(0), |
| 390 | input_tensor.float().flatten().unsqueeze(0), |
| 391 | ).item() |
| 392 | max_abs_diff = (recon.float() - input_tensor.float()).abs().max().item() |
| 393 | |
| 394 | print( |
| 395 | f" Reconstruction: cosine_sim={cos_sim:.6f}, max_abs_diff={max_abs_diff:.4f}" |
| 396 | ) |
| 397 | |
| 398 | assert cos_sim > 0.99, ( |
| 399 | f"Reconstruction cosine similarity too low: {cos_sim:.6f}. " |
| 400 | f"Expected > 0.99 for correct MXFP4 quantization." |
| 401 | ) |
| 402 | # With correct E8M0, max abs diff should be bounded by scale * 6 |
| 403 | # (worst case: value just below threshold rounds to wrong FP4 code) |
| 404 | assert max_abs_diff < 1.0, ( |
| 405 | f"Max reconstruction error too large: {max_abs_diff:.4f}. " |
| 406 | "Likely caused by incorrect E8M0 scale (values saturating to ±6)." |
| 407 | ) |
| 408 | |
| 409 | |
| 410 | @pytest.mark.skipif( |
| 411 | not is_sm100_supported(), |
| 412 | reason="mxfp4_experts_quant requires CUDA SM100", |
| 413 | ) |
| 414 | def test_mxfp4_experts_quant_no_saturation(): |
| 415 | """Test that the E8M0 scale is large enough to avoid unexpected saturation. |
| 416 | |
| 417 | With the buggy NVFP4 formula, the scale was too small causing most values |
| 418 | to saturate to ±6 in FP4. The fixed OCP MX formula should ensure that |
| 419 | block_max / scale <= 6.0 (the max E2M1 value) in almost all cases. |
| 420 | """ |
| 421 | device = "cuda" |
| 422 | |
| 423 | m, k = 128, 1024 |
| 424 | # Use inputs with known range to make saturation detectable |
| 425 | input_tensor = torch.randn(m, k, device=device, dtype=torch.bfloat16) * 0.5 |
| 426 | |
| 427 | num_experts = 1 |
| 428 | expert_offsets = torch.tensor([0, m], device=device, dtype=torch.int32) |
| 429 | num_k_tiles = (k // MXFP4_BLOCK_SIZE + 3) // 4 |
| 430 | blockscale_offsets = torch.tensor( |
| 431 | [0, align(m, 128) * num_k_tiles], device=device, dtype=torch.int32 |
| 432 | ) |
| 433 | |
| 434 | output_fp4, output_sf = ops.mxfp4_experts_quant( |
| 435 | input_tensor, expert_offsets, blockscale_offsets, num_experts, topk=1 |
| 436 | ) |
| 437 | |
| 438 | # Check saturation rate: count FP4 values that are ±6 (codes 7 and 15) |
| 439 | lo = output_fp4 & 0x0F |
| 440 | hi = (output_fp4 >> 4) & 0x0F |
| 441 | # Code 7 = +6.0, code 15 = -6.0 |
| 442 | saturated = ((lo == 7) | (lo == 15) | (hi == 7) | (hi == 15)).sum().item() |
| 443 | total_values = m * k |
| 444 | saturation_rate = saturated / total_values |
| 445 | |
| 446 | print( |
| 447 | f" Saturation rate: {saturation_rate * 100:.2f}% " |
| 448 | f"({saturated}/{total_values} values at ±6)" |
| 449 | ) |
| 450 | |
| 451 | # For Gaussian input with std=0.5, saturation should be very rare |
| 452 | # (±6 * scale is far from the typical range). |
| 453 | # The buggy kernel had ~30-50% saturation; fixed should be < 5%. |
| 454 | assert saturation_rate < 0.05, ( |
| 455 | f"FP4 saturation rate too high: {saturation_rate * 100:.2f}%. " |
| 456 | "This suggests the E8M0 scale is too small (NVFP4 formula bug). " |
| 457 | "Expected < 5% for Gaussian(0, 0.5) input with correct OCP MX scale." |
| 458 | ) |
| 459 | |
| 460 | |
| 461 | if __name__ == "__main__": |
| 462 | pytest.main([__file__, "-v", "-s"]) |