| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Native CDNA4 MXFP8 linear (``RocmDotScaledMxfp8LinearKernel``). |
| 4 | |
| 5 | DeepSeek V4.1 routes every attention linear through this kernel. It used to |
| 6 | carry two BF16-dequantize workarounds on ROCm (one for all ``.attn.`` linears, |
| 7 | one for the row-parallel ``wo_b``, whose comment blamed NaNs at small token |
| 8 | counts). These tests pin the two properties those workarounds stood in for -- |
| 9 | finite output and MXFP8-level accuracy -- across the V4.1 attention shapes at |
| 10 | TP=4 and the small token counts decode actually runs. |
| 11 | """ |
| 12 | |
| 13 | import pytest |
| 14 | import torch |
| 15 | |
| 16 | from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( |
| 17 | _mxfp8_e4m3_quantize_torch, |
| 18 | _mxfp8_e4m3_quantize_triton, |
| 19 | dequant_mxfp8_to_bf16, |
| 20 | mxfp8_e4m3_quantize, |
| 21 | ) |
| 22 | from vllm.platforms import current_platform |
| 23 | |
| 24 | pytestmark = pytest.mark.skipif( |
| 25 | not current_platform.is_rocm() or not current_platform.supports_mx(), |
| 26 | reason="native MX linear requires CDNA4 (gfx95x)", |
| 27 | ) |
| 28 | |
| 29 | # (name, N, K) for the DeepSeek V4.1 attention linears at tensor_parallel=4: |
| 30 | # fused wq_a|wkv is replicated, wq_b/wo_a are column-parallel, wo_b is |
| 31 | # row-parallel (so K, not N, is sharded). |
| 32 | V41_ATTN_SHAPES = [ |
| 33 | ("fused_wqa_wkv", 1280 + 512, 5120), |
| 34 | ("wq_b", 64 * 512 // 4, 1280), |
| 35 | ("wo_a", 8 * 1024 // 4, 4096), |
| 36 | ("wo_b", 5120, 8192 // 4), |
| 37 | ] |
| 38 | |
| 39 | |
| 40 | def _make_layer(n: int, k: int, device: str) -> torch.nn.Module: |
| 41 | weight_bf16 = torch.randn(n, k, device=device, dtype=torch.bfloat16) * 0.05 |
| 42 | weight, weight_scale = _mxfp8_e4m3_quantize_torch(weight_bf16) |
| 43 | layer = torch.nn.Module() |
| 44 | layer.weight = torch.nn.Parameter(weight, requires_grad=False) |
| 45 | layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False) |
| 46 | return layer |
| 47 | |
| 48 | |
| 49 | @pytest.mark.parametrize("name,n,k", V41_ATTN_SHAPES, ids=lambda v: str(v)) |
| 50 | @pytest.mark.parametrize("num_tokens", [1, 2, 7, 16, 31, 32, 128, 1024]) |
| 51 | @torch.inference_mode() |
| 52 | def test_rocm_mxfp8_linear_matches_dequant_reference(name, n, k, num_tokens): |
| 53 | """Output is finite and matches an exact MXFP8 GEMM of the same inputs. |
| 54 | |
| 55 | The reference quantizes the activation with the same quantizer and then |
| 56 | contracts in FP32, so the tolerance covers only kernel/accumulation error, |
| 57 | not the MXFP8 quantization error both sides share. |
| 58 | """ |
| 59 | from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import ( |
| 60 | Mxfp8LinearLayerConfig, |
| 61 | ) |
| 62 | from vllm.model_executor.kernels.linear.mxfp8.rocm_native import ( |
| 63 | RocmDotScaledMxfp8LinearKernel, |
| 64 | ) |
| 65 | |
| 66 | torch.manual_seed(num_tokens) |
| 67 | device = "cuda" |
| 68 | layer = _make_layer(n, k, device) |
| 69 | kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(weight_shape=(n, k))) |
| 70 | kernel.process_weights_after_loading(layer) |
| 71 | |
| 72 | x = torch.randn(num_tokens, k, device=device, dtype=torch.bfloat16) * 0.5 |
| 73 | out = kernel.apply_weights(layer, x) |
| 74 | |
| 75 | assert out.shape == (num_tokens, n) |
| 76 | assert torch.isfinite(out).all(), f"{name}: non-finite output" |
| 77 | |
| 78 | x_q, x_scale = mxfp8_e4m3_quantize(x) |
| 79 | expected = ( |
| 80 | dequant_mxfp8_to_bf16(x_q, x_scale).float() |
| 81 | @ dequant_mxfp8_to_bf16(layer.weight, layer.weight_scale).float().T |
| 82 | ) |
| 83 | rel = (out.float() - expected).norm() / expected.norm() |
| 84 | assert rel < 5e-3, f"{name}: relative error {rel:.4f}" |
| 85 | |
| 86 | |
| 87 | @torch.inference_mode() |
| 88 | def test_rocm_mxfp8_linear_handles_extreme_magnitudes(): |
| 89 | """Rows that are all-zero or near the E4M3 range must not produce NaNs. |
| 90 | |
| 91 | A degenerate block makes the dynamic activation quantizer pick an extreme |
| 92 | E8M0 scale; a zero row makes it pick the tiny-clamped one. Both feed |
| 93 | ``tl.dot_scaled`` directly. |
| 94 | """ |
| 95 | from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import ( |
| 96 | Mxfp8LinearLayerConfig, |
| 97 | ) |
| 98 | from vllm.model_executor.kernels.linear.mxfp8.rocm_native import ( |
| 99 | RocmDotScaledMxfp8LinearKernel, |
| 100 | ) |
| 101 | |
| 102 | torch.manual_seed(0) |
| 103 | device = "cuda" |
| 104 | n, k = 5120, 2048 # wo_b at TP=4 |
| 105 | layer = _make_layer(n, k, device) |
| 106 | kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(weight_shape=(n, k))) |
| 107 | kernel.process_weights_after_loading(layer) |
| 108 | |
| 109 | x = torch.randn(8, k, device=device, dtype=torch.bfloat16) |
| 110 | x[0] = 0.0 |
| 111 | x[1] = 1e-30 |
| 112 | x[2] = 400.0 |
| 113 | x[3, ::2] = 0.0 |
| 114 | |
| 115 | out = kernel.apply_weights(layer, x) |
| 116 | assert torch.isfinite(out).all() |
| 117 | assert (out[0] == 0).all(), "an all-zero activation row must give a zero row" |
| 118 | |
| 119 | |
| 120 | # (name, N, K) of the DeepSeek-V4.1-Flash MXFP8 linears at TP=2 and TP=4. The |
| 121 | # checkpoint stores their scales in 32x32 blocks; shared_down at TP=4 has |
| 122 | # K % 128 != 0, which only the block-scaled GEMM runs natively. |
| 123 | V41_BLOCK32_SHAPES = [ |
| 124 | ("fused_wqa_wkv", 1280 + 512, 5120), |
| 125 | ("wq_b_tp2", 64 * 512 // 2, 1280), |
| 126 | ("indexer_wq_b", 32 * 128, 1280), |
| 127 | ("wo_b_tp2", 5120, 8192 // 2), |
| 128 | ("shared_gate_up_tp2", 2 * 2304 // 2, 5120), |
| 129 | ("shared_down_tp2", 5120, 2304 // 2), |
| 130 | ("shared_down_tp4", 5120, 2304 // 4), |
| 131 | ] |
| 132 | |
| 133 | |
| 134 | def _make_block32_layer(n: int, k: int, device: str) -> torch.nn.Module: |
| 135 | """An MXFP8 layer with 32x32 block scales, expanded per row as loaded.""" |
| 136 | blocks = (torch.randn(n, k, device=device) * 0.05).view(n // 32, 32, k // 32, 32) |
| 137 | amax = blocks.abs().amax(dim=(1, 3)).clamp(min=torch.finfo(torch.float32).tiny) |
| 138 | bits = (torch.ceil(torch.log2(amax / 448.0)) + 127.0).clamp(0, 254) |
| 139 | weight = (blocks * torch.exp2(127.0 - bits)[:, None, :, None]).view(n, k) |
| 140 | layer = torch.nn.Module() |
| 141 | layer.weight = torch.nn.Parameter( |
| 142 | weight.to(torch.float8_e4m3fn), requires_grad=False |
| 143 | ) |
| 144 | layer.weight_scale = torch.nn.Parameter( |
| 145 | bits.to(torch.uint8).repeat_interleave(32, dim=0), requires_grad=False |
| 146 | ) |
| 147 | return layer |
| 148 | |
| 149 | |
| 150 | @pytest.mark.parametrize("name,n,k", V41_BLOCK32_SHAPES, ids=lambda v: str(v)) |
| 151 | @pytest.mark.parametrize("num_tokens", [1, 7, 32, 33, 192, 1024]) |
| 152 | @torch.inference_mode() |
| 153 | def test_rocm_mxfp8_linear_block32_matches_dequant_reference(name, n, k, num_tokens): |
| 154 | """32x32 block scales stay compact and the block-scaled GEMM is exact. |
| 155 | |
| 156 | The token counts cover the packed, tiled and split-K configurations. |
| 157 | """ |
| 158 | from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import ( |
| 159 | Mxfp8LinearLayerConfig, |
| 160 | ) |
| 161 | from vllm.model_executor.kernels.linear.mxfp8.rocm_native import ( |
| 162 | RocmDotScaledMxfp8LinearKernel, |
| 163 | ) |
| 164 | |
| 165 | torch.manual_seed(num_tokens) |
| 166 | device = "cuda" |
| 167 | layer = _make_block32_layer(n, k, device) |
| 168 | row_scale = layer.weight_scale.data.clone() |
| 169 | kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(weight_shape=(n, k))) |
| 170 | kernel.process_weights_after_loading(layer) |
| 171 | assert layer.weight.dtype == torch.float8_e4m3fn |
| 172 | assert layer.weight_scale.shape == (n // 32, k // 32) |
| 173 | |
| 174 | x = torch.randn(num_tokens, k, device=device, dtype=torch.bfloat16) * 0.5 |
| 175 | out = kernel.apply_weights(layer, x) |
| 176 | |
| 177 | x_q, x_scale = mxfp8_e4m3_quantize(x) |
| 178 | expected = ( |
| 179 | dequant_mxfp8_to_bf16(x_q, x_scale).float() |
| 180 | @ dequant_mxfp8_to_bf16(layer.weight, row_scale).float().T |
| 181 | ) |
| 182 | assert out.shape == (num_tokens, n) |
| 183 | rel = (out.float() - expected).norm() / expected.norm() |
| 184 | assert rel < 5e-3, f"{name}: relative error {rel:.4f}" |
| 185 | |
| 186 | |
| 187 | @torch.inference_mode() |
| 188 | def test_rocm_mxfp8_linear_keeps_per_row_scales(): |
| 189 | """Scales that differ within a 32-row block are not compacted.""" |
| 190 | from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import ( |
| 191 | Mxfp8LinearLayerConfig, |
| 192 | ) |
| 193 | from vllm.model_executor.kernels.linear.mxfp8.rocm_native import ( |
| 194 | RocmDotScaledMxfp8LinearKernel, |
| 195 | ) |
| 196 | |
| 197 | n, k = 64, 256 |
| 198 | layer = _make_block32_layer(n, k, "cuda") |
| 199 | layer.weight_scale.data[1, 0] += 1 |
| 200 | kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(weight_shape=(n, k))) |
| 201 | kernel.process_weights_after_loading(layer) |
| 202 | assert layer.weight_scale.shape == (n, k // 32) |
| 203 | |
| 204 | |
| 205 | def _split_k_case(packed: bool): |
| 206 | """A 3-way in-launch split-K on shared_down TP4, with K and M tails.""" |
| 207 | from vllm.model_executor.kernels.linear.mxfp8 import rocm_block32_gemm as g |
| 208 | |
| 209 | n, k, m = 5120, 576, 19 |
| 210 | layer = _make_block32_layer(n, k, "cuda") |
| 211 | x_q, x_scale = mxfp8_e4m3_quantize( |
| 212 | torch.randn(m, k, device="cuda", dtype=torch.bfloat16) |
| 213 | ) |
| 214 | expected = ( |
| 215 | dequant_mxfp8_to_bf16(x_q, x_scale).float() |
| 216 | @ dequant_mxfp8_to_bf16(layer.weight, layer.weight_scale).float().T |
| 217 | ) |
| 218 | if packed: |
| 219 | cfg = g._packed(16, 32, 128, 1, 2, 2, 16, 0, 3) |
| 220 | else: |
| 221 | cfg = g._tiled(32, 64, 128, 1, True, 4, 2, 16, 0, 3) |
| 222 | args = (x_q, x_scale, layer.weight, layer.weight_scale.data[::32].contiguous()) |
| 223 | return g, cfg, args, expected |
| 224 | |
| 225 | |
| 226 | def _rel_err(out: torch.Tensor, expected: torch.Tensor) -> float: |
| 227 | return ((out.float() - expected).norm() / expected.norm()).item() |
| 228 | |
| 229 | |
| 230 | @pytest.mark.parametrize("packed", [True, False]) |
| 231 | @torch.inference_mode() |
| 232 | def test_rocm_mxfp8_block32_split_k_in_launch(packed): |
| 233 | """Splits summed by the tile's last program are exact and repeatable.""" |
| 234 | torch.manual_seed(0) |
| 235 | g, cfg, args, expected = _split_k_case(packed) |
| 236 | out = torch.empty_like(expected, dtype=torch.bfloat16) |
| 237 | g._launch(*args, out, cfg) |
| 238 | assert _rel_err(out, expected) < 5e-3 |
| 239 | first = out.clone() |
| 240 | for _ in range(20): |
| 241 | g._launch(*args, out, cfg) |
| 242 | assert torch.equal(out, first) |
| 243 | |
| 244 | |
| 245 | @pytest.mark.parametrize("packed", [True, False]) |
| 246 | @torch.inference_mode() |
| 247 | def test_rocm_mxfp8_block32_split_k_first_use_in_capture(packed): |
| 248 | """A stream whose first split-K launch is being captured has no counters |
| 249 | yet; it must fall back to an unfused split instead of failing.""" |
| 250 | torch.manual_seed(0) |
| 251 | g, cfg, args, expected = _split_k_case(packed) |
| 252 | out = torch.zeros_like(expected, dtype=torch.bfloat16) |
| 253 | stream = torch.cuda.Stream() |
| 254 | graph = torch.cuda.CUDAGraph() |
| 255 | with torch.cuda.graph(graph, stream=stream): |
| 256 | g._launch(*args, out, cfg) |
| 257 | graph.replay() |
| 258 | torch.accelerator.synchronize() |
| 259 | assert _rel_err(out, expected) < 5e-3 |
| 260 | |
| 261 | |
| 262 | @pytest.mark.parametrize("num_rows", [1, 64, 65]) |
| 263 | @torch.inference_mode() |
| 264 | def test_rocm_mxfp8_quantizer_matches_torch_on_degenerate_blocks(num_rows): |
| 265 | """The fused ROCm quantizer must agree with the torch one, zeros included. |
| 266 | |
| 267 | An all-zero 32-element block clamps the E8M0 scale to 0, i.e. 2**-127 -- |
| 268 | subnormal in fp32 and flushed to zero on CDNA. Dividing by it turned the |
| 269 | block into 0/0 == NaN, which then poisoned every output column of the GEMM |
| 270 | that consumed it. |
| 271 | """ |
| 272 | torch.manual_seed(num_rows) |
| 273 | device = "cuda" |
| 274 | x = torch.randn(num_rows, 128, device=device, dtype=torch.bfloat16) |
| 275 | x[0] = 0.0 # all-zero row: every block clamps to scale byte 0 |
| 276 | x[-1, :32] = 0.0 # single zero block in an otherwise normal row |
| 277 | x[-1, 32:64] = 1e-38 # denormal-ish block |
| 278 | |
| 279 | q_triton, s_triton = _mxfp8_e4m3_quantize_triton(x) |
| 280 | q_torch, s_torch = _mxfp8_e4m3_quantize_torch(x) |
| 281 | |
| 282 | assert not torch.isnan(q_triton.float()).any() |
| 283 | torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0) |
| 284 | torch.testing.assert_close( |
| 285 | q_triton.view(torch.uint8), q_torch.view(torch.uint8), rtol=0, atol=0 |
| 286 | ) |