| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Correctness tests for the online MXFP4 weight quantization backends. |
| 4 | |
| 5 | Each backend (Triton, aiter, XPU) quantizes a synthetic bf16/fp16 |
| 6 | tensor to MXFP4, the result is dequantized back to bf16/fp16, and compared |
| 7 | against the pure-torch reference in `reference_mxfp4.py`. |
| 8 | """ |
| 9 | |
| 10 | import pytest |
| 11 | import torch |
| 12 | |
| 13 | from vllm.config.kernel import MoEBackend |
| 14 | from vllm.config.model import ModelConfig |
| 15 | from vllm.model_executor.kernels.linear import _POSSIBLE_MXFP4_KERNELS |
| 16 | from vllm.model_executor.kernels.linear.mxfp4.aiter import ( |
| 17 | AiterMxfp4LinearKernel, |
| 18 | ) |
| 19 | from vllm.model_executor.layers.fused_moe import FusedMoEFactory |
| 20 | from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts |
| 21 | from vllm.model_executor.layers.quantization.online.mxfp4 import ( |
| 22 | Mxfp4OnlineLinearMethod, |
| 23 | Mxfp4OnlineMoEMethod, |
| 24 | ) |
| 25 | from vllm.model_executor.layers.quantization.quark.quark_moe import ( |
| 26 | QuarkOCP_MX_MoEMethod, |
| 27 | ) |
| 28 | from vllm.model_executor.layers.quantization.quark.schemes.quark_ocp_mx import ( |
| 29 | QuarkOCP_MX, |
| 30 | ) |
| 31 | from vllm.model_executor.layers.quantization.quark.utils import ( |
| 32 | quark_quantize_weight_to_mxfp4, |
| 33 | ) |
| 34 | from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( |
| 35 | downcast_to_mxfp, |
| 36 | mxfp4_quantize, |
| 37 | quant_dequant_mxfp4, |
| 38 | ) |
| 39 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 40 | kMxfp4Dynamic, |
| 41 | kMxfp4Static, |
| 42 | ) |
| 43 | from vllm.model_executor.model_loader.weight_utils import default_weight_loader |
| 44 | from vllm.model_executor.utils import replace_parameter |
| 45 | from vllm.platforms import PlatformEnum, current_platform |
| 46 | |
| 47 | from .reference_mxfp4 import dq_mxfp4_torch, qdq_mxfp4_torch |
| 48 | |
| 49 | |
| 50 | def fix_negative_zeros(tensor: torch.Tensor) -> torch.Tensor: |
| 51 | """Normalize FP4 e2m1 negative-zero codewords (0b1000) to positive zero |
| 52 | (0b0000) in a packed uint8 tensor (two e2m1 values per byte, low nibble |
| 53 | first). |
| 54 | |
| 55 | -0.0 and +0.0 dequantize to the same float value, but different MXFP4 |
| 56 | quantization backends disagree on which one they emit for values that |
| 57 | round to zero magnitude. This normalizes both packed representations so |
| 58 | tensors from different backends can be compared with `torch.equal`. |
| 59 | """ |
| 60 | assert tensor.dtype == torch.uint8, ( |
| 61 | f"Expected a torch.uint8 tensor, got {tensor.dtype}" |
| 62 | ) |
| 63 | |
| 64 | low_nibble = tensor & 0x0F |
| 65 | high_nibble = tensor & 0xF0 |
| 66 | |
| 67 | low_nibble = torch.where( |
| 68 | low_nibble == 0x08, torch.zeros_like(low_nibble), low_nibble |
| 69 | ) |
| 70 | high_nibble = torch.where( |
| 71 | high_nibble == 0x80, torch.zeros_like(high_nibble), high_nibble |
| 72 | ) |
| 73 | |
| 74 | return low_nibble | high_nibble |
| 75 | |
| 76 | |
| 77 | def assert_quantized_weights_equal( |
| 78 | checkpoint_layer: torch.nn.Module, |
| 79 | online_layer: torch.nn.Module, |
| 80 | keys: tuple[str, ...], |
| 81 | packed_weight_keys: tuple[str, ...], |
| 82 | ) -> None: |
| 83 | """Assert the checkpoint and online paths produced byte-identical weights. |
| 84 | |
| 85 | `packed_weight_keys` lists the entries holding packed FP4 codewords, which |
| 86 | need negative-zero normalization before comparison. |
| 87 | """ |
| 88 | for key in keys: |
| 89 | checkpoint_tensor = getattr(checkpoint_layer, key).view(torch.uint8) |
| 90 | online_tensor = getattr(online_layer, key).view(torch.uint8) |
| 91 | |
| 92 | # NOTE: AMD Quark checkpoints use exclusively **positive** zeros, |
| 93 | # while other mxfp4_quantize implementations from mxfp4_utils.py may not. |
| 94 | if key in packed_weight_keys: |
| 95 | checkpoint_tensor = fix_negative_zeros(checkpoint_tensor) |
| 96 | online_tensor = fix_negative_zeros(online_tensor) |
| 97 | |
| 98 | assert checkpoint_tensor.shape == online_tensor.shape |
| 99 | assert checkpoint_tensor.dtype == online_tensor.dtype |
| 100 | num_mismatched = (checkpoint_tensor != online_tensor).sum().item() |
| 101 | total = checkpoint_tensor.numel() |
| 102 | |
| 103 | assert torch.equal(checkpoint_tensor, online_tensor), ( |
| 104 | f"{key}: {num_mismatched}/{total} " |
| 105 | f"({100 * num_mismatched / total:.4f}%) mismatched bytes" |
| 106 | ) |
| 107 | |
| 108 | |
| 109 | def _skip_reason_if_unavailable(backend: str, dtype: torch.dtype) -> str | None: |
| 110 | """Return a skip reason if `backend` cannot run on the current host.""" |
| 111 | if backend == "triton": |
| 112 | if not (current_platform.is_cuda() or current_platform.is_rocm()): |
| 113 | return "Triton MXFP4 kernel requires a CUDA or ROCm GPU." |
| 114 | return None |
| 115 | if backend == "aiter": |
| 116 | from vllm._aiter_ops import is_aiter_found_and_supported |
| 117 | |
| 118 | if not is_aiter_found_and_supported(): |
| 119 | return "aiter is not available/supported on this platform." |
| 120 | if dtype != torch.bfloat16: |
| 121 | return "aiter's dynamic_mxfp4_quant only supports bfloat16 input." |
| 122 | return None |
| 123 | if backend == "xpu": |
| 124 | if not current_platform.is_xpu(): |
| 125 | return "not on XPU platform." |
| 126 | return None |
| 127 | if backend == "quark": |
| 128 | try: |
| 129 | import quark.torch.kernel.mx # noqa: F401 |
| 130 | except ImportError: |
| 131 | return "amd-quark is not installed." |
| 132 | return None |
| 133 | raise ValueError(f"Unknown backend {backend}") |
| 134 | |
| 135 | |
| 136 | @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) |
| 137 | @pytest.mark.parametrize("backend", ["triton", "aiter", "xpu", "quark"]) |
| 138 | def test_mxfp4_quantization_correctness(backend: str, dtype: torch.dtype): |
| 139 | """Tests that the different implementations of mxfp4_quantize |
| 140 | in mxfp4_utils.py all match. |
| 141 | """ |
| 142 | skip_reason = _skip_reason_if_unavailable(backend, dtype) |
| 143 | if skip_reason is not None: |
| 144 | pytest.skip(skip_reason) |
| 145 | |
| 146 | torch.manual_seed(3) |
| 147 | |
| 148 | num_rows = 64 |
| 149 | hidden_size = 32 * 32 # multiple 32-element MXFP4 blocks |
| 150 | device = current_platform.device_type |
| 151 | |
| 152 | x = (torch.rand(num_rows, hidden_size, dtype=dtype, device=device) - 0.5) * 2 |
| 153 | # Vary the magnitude block-to-block so several scale exponents are |
| 154 | # exercised, rather than a single one for the whole tensor. |
| 155 | scalings = [2.3, 0.03, 7.3, 0.1, 0.004, 17.3, 1e4, 1e-4] |
| 156 | for i in range(hidden_size // 32): |
| 157 | x[:, i * 32 : (i + 1) * 32] *= scalings[i % len(scalings)] |
| 158 | |
| 159 | if backend == "triton": |
| 160 | x_fp4, x_scale, _ = downcast_to_mxfp(x, axis=-1) |
| 161 | elif backend == "aiter": |
| 162 | x_fp4, x_scale = quark_quantize_weight_to_mxfp4(x) |
| 163 | elif backend == "xpu": |
| 164 | # TODO: enable this test on XPU |
| 165 | pytest.skip("xpu mxfp4 quantization match to reference is untested") |
| 166 | # from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( |
| 167 | # xpu_mxfp4_quantize, |
| 168 | # ) |
| 169 | |
| 170 | # x_fp4, x_scale = xpu_mxfp4_quantize(x) |
| 171 | elif backend == "quark": |
| 172 | result = quant_dequant_mxfp4(x, scale_calculation_mode="even") |
| 173 | else: |
| 174 | raise ValueError(f"Unknown backend {backend}") |
| 175 | |
| 176 | if backend != "quark": |
| 177 | result = dq_mxfp4_torch(x_fp4, x_scale, x.dtype) |
| 178 | reference = qdq_mxfp4_torch(x, scale_calculation_mode="even") |
| 179 | |
| 180 | assert torch.equal(result, reference) |
| 181 | |
| 182 | |
| 183 | @pytest.mark.skipif( |
| 184 | not current_platform.is_cuda_alike(), reason="Only tested on ROCm/CUDA." |
| 185 | ) |
| 186 | @pytest.mark.parametrize("tp_size", [2, 4, 8]) |
| 187 | def test_online_mxfp4_tp_weight_quant_matches_unsharded(tp_size: int): |
| 188 | """TP-aligned MXFP4 shards match slices of unsharded quantization.""" |
| 189 | torch.manual_seed(3) |
| 190 | |
| 191 | device = current_platform.device_type |
| 192 | hidden_size = 1024 |
| 193 | intermediate_size = 512 |
| 194 | |
| 195 | dense_weight = torch.randn( |
| 196 | hidden_size, intermediate_size, dtype=torch.bfloat16, device=device |
| 197 | ) |
| 198 | w13_weight = torch.randn( |
| 199 | 2, |
| 200 | 2 * intermediate_size, |
| 201 | hidden_size, |
| 202 | dtype=torch.bfloat16, |
| 203 | device=device, |
| 204 | ) |
| 205 | w2_weight = torch.randn( |
| 206 | 2, |
| 207 | hidden_size, |
| 208 | intermediate_size, |
| 209 | dtype=torch.bfloat16, |
| 210 | device=device, |
| 211 | ) |
| 212 | |
| 213 | dense_quantized, dense_scales = mxfp4_quantize(dense_weight) |
| 214 | w13_quantized, w13_scales = mxfp4_quantize(w13_weight) |
| 215 | w2_quantized, w2_scales = mxfp4_quantize(w2_weight) |
| 216 | |
| 217 | for tp_rank in range(tp_size): |
| 218 | output_per_partition = hidden_size // tp_size |
| 219 | input_per_partition = intermediate_size // tp_size |
| 220 | output_start = tp_rank * output_per_partition |
| 221 | input_start = tp_rank * input_per_partition |
| 222 | |
| 223 | column_quantized, column_scales = mxfp4_quantize( |
| 224 | dense_weight.narrow(0, output_start, output_per_partition) |
| 225 | ) |
| 226 | assert torch.equal( |
| 227 | column_quantized, |
| 228 | dense_quantized.narrow(0, output_start, output_per_partition), |
| 229 | ) |
| 230 | assert torch.equal( |
| 231 | column_scales, |
| 232 | dense_scales.narrow(0, output_start, output_per_partition), |
| 233 | ) |
| 234 | |
| 235 | row_quantized, row_scales = mxfp4_quantize( |
| 236 | dense_weight.narrow(1, input_start, input_per_partition) |
| 237 | ) |
| 238 | assert torch.equal( |
| 239 | row_quantized, |
| 240 | dense_quantized.narrow(1, input_start // 2, input_per_partition // 2), |
| 241 | ) |
| 242 | assert torch.equal( |
| 243 | row_scales, |
| 244 | dense_scales.narrow(1, input_start // 32, input_per_partition // 32), |
| 245 | ) |
| 246 | |
| 247 | w13_quantized_shard, w13_scales_shard = mxfp4_quantize( |
| 248 | w13_weight.narrow(1, 2 * input_start, 2 * input_per_partition) |
| 249 | ) |
| 250 | assert torch.equal( |
| 251 | w13_quantized_shard, |
| 252 | w13_quantized.narrow(1, 2 * input_start, 2 * input_per_partition), |
| 253 | ) |
| 254 | assert torch.equal( |
| 255 | w13_scales_shard, |
| 256 | w13_scales.narrow(1, 2 * input_start, 2 * input_per_partition), |
| 257 | ) |
| 258 | |
| 259 | w2_quantized_shard, w2_scales_shard = mxfp4_quantize( |
| 260 | w2_weight.narrow(2, input_start, input_per_partition) |
| 261 | ) |
| 262 | assert torch.equal( |
| 263 | w2_quantized_shard, |
| 264 | w2_quantized.narrow(2, input_start // 2, input_per_partition // 2), |
| 265 | ) |
| 266 | assert torch.equal( |
| 267 | w2_scales_shard, |
| 268 | w2_scales.narrow(2, input_start // 32, input_per_partition // 32), |
| 269 | ) |
| 270 | |
| 271 | |
| 272 | @pytest.mark.skipif( |
| 273 | not current_platform.is_cuda_alike(), reason="Only tested on ROCm/CUDA." |
| 274 | ) |
| 275 | @pytest.mark.parametrize("moe_backend", ["aiter", "emulation"]) |
| 276 | @pytest.mark.parametrize( |
| 277 | "unpadded_hidden_size,unpadded_intermediate_size", |
| 278 | [ |
| 279 | pytest.param(256, 256, id="no_padding"), |
| 280 | # Not multiples of the MXFP4 block size (32), so that every backend -- |
| 281 | # including emulation, whose round-up granularity is 32 -- has to pad. |
| 282 | pytest.param(240, 208, id="padding"), |
| 283 | ], |
| 284 | ) |
| 285 | def test_online_mxfp4_moe_matches_quark( |
| 286 | moe_backend: MoEBackend, |
| 287 | unpadded_hidden_size: int, |
| 288 | unpadded_intermediate_size: int, |
| 289 | default_vllm_config, |
| 290 | dist_init, |
| 291 | monkeypatch, |
| 292 | ): |
| 293 | """Ensures `Mxfp4OnlineMoEMethod` (online quantization) |
| 294 | and `QuarkOCP_MX_MoEMethod` (AMD Quark checkpoints) produce the same weights, |
| 295 | with same MOE backend used. |
| 296 | |
| 297 | The `padding` case additionally covers `maybe_roundup_sizes`: the online |
| 298 | path allocates its bf16 weights with `torch.empty_strided` and the loader |
| 299 | only writes the unpadded slice, so the padding is filled with NaN here to |
| 300 | reproduce that uninitialized state. It must be zeroed before quantization |
| 301 | for the two paths to agree, since Quark allocates its padding with |
| 302 | `torch.zeros`. |
| 303 | """ |
| 304 | if moe_backend == "aiter": |
| 305 | from vllm._aiter_ops import rocm_aiter_ops |
| 306 | |
| 307 | monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") |
| 308 | monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1") |
| 309 | monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: True) |
| 310 | rocm_aiter_ops.refresh_env_variables() |
| 311 | |
| 312 | if current_platform.is_cuda(): |
| 313 | pytest.skip( |
| 314 | "mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4 requires " |
| 315 | "rocm_aiter_ops in weight conversion, not compatible on cuda" |
| 316 | ) |
| 317 | elif moe_backend == "emulation": |
| 318 | # `OCP_MXQuantizationEmulationTritonExperts.is_supported_config` gates |
| 319 | # on `has_quark()`, but its weight processing does not require it. |
| 320 | monkeypatch.setattr( |
| 321 | "vllm.model_executor.layers.fused_moe.experts." |
| 322 | "ocp_mx_emulation_moe.has_quark", |
| 323 | lambda: True, |
| 324 | ) |
| 325 | |
| 326 | default_vllm_config.model_config = ModelConfig() |
| 327 | |
| 328 | num_experts = 4 |
| 329 | device = current_platform.device_type |
| 330 | |
| 331 | def make_layer(prefix: str) -> RoutedExperts: |
| 332 | runner = FusedMoEFactory( |
| 333 | num_experts=num_experts, |
| 334 | top_k=2, |
| 335 | hidden_size=unpadded_hidden_size, |
| 336 | intermediate_size=unpadded_intermediate_size, |
| 337 | prefix=prefix, |
| 338 | ) |
| 339 | layer = runner.routed_experts |
| 340 | layer.moe_config.moe_backend = moe_backend |
| 341 | return layer |
| 342 | |
| 343 | # `create_weights` implementations use plain `torch.zeros`/`torch.randn` |
| 344 | # with no explicit device, so without this context they would default |
| 345 | # to CPU and diverge from the (GPU) tensors produced by |
| 346 | # `mxfp4_quantize`/`replace_parameter`. |
| 347 | with torch.device(device): |
| 348 | checkpoint_layer = make_layer("checkpoint_layer") |
| 349 | online_layer = make_layer("online_layer") |
| 350 | |
| 351 | checkpoint_method = QuarkOCP_MX_MoEMethod( |
| 352 | moe=checkpoint_layer.moe_config, |
| 353 | weight_quant_key=kMxfp4Static, |
| 354 | activation_quant_key=kMxfp4Dynamic, |
| 355 | ) |
| 356 | online_method = Mxfp4OnlineMoEMethod(moe=online_layer.moe_config) |
| 357 | |
| 358 | # `RoutedExperts.__init__` applies this round-up in production; these |
| 359 | # layers are built without a quant config, so it is applied explicitly. |
| 360 | # Both methods must agree on the padded sizes. |
| 361 | roundup_arguments = dict( |
| 362 | hidden_size=unpadded_hidden_size, |
| 363 | intermediate_size_per_partition=unpadded_intermediate_size, |
| 364 | act_dtype=torch.bfloat16, |
| 365 | moe_parallel_config=online_layer.moe_config.moe_parallel_config, |
| 366 | ) |
| 367 | hidden_size, intermediate_size = online_method.maybe_roundup_sizes( |
| 368 | **roundup_arguments |
| 369 | ) |
| 370 | assert ( |
| 371 | hidden_size, |
| 372 | intermediate_size, |
| 373 | ) == checkpoint_method.maybe_roundup_sizes(**roundup_arguments) |
| 374 | |
| 375 | if unpadded_hidden_size == 240: |
| 376 | # Padded case. |
| 377 | assert hidden_size > unpadded_hidden_size |
| 378 | assert intermediate_size > unpadded_intermediate_size |
| 379 | |
| 380 | for layer in (checkpoint_layer, online_layer): |
| 381 | assert layer.moe_config.hidden_dim_unpadded == unpadded_hidden_size |
| 382 | assert ( |
| 383 | layer.moe_config.intermediate_size_per_partition_unpadded |
| 384 | == unpadded_intermediate_size |
| 385 | ) |
| 386 | layer.moe_config.hidden_dim = hidden_size |
| 387 | layer.moe_config.intermediate_size_per_partition = intermediate_size |
| 388 | |
| 389 | gate_up_weight = torch.randn( |
| 390 | num_experts, |
| 391 | 2 * unpadded_intermediate_size, |
| 392 | unpadded_hidden_size, |
| 393 | dtype=torch.bfloat16, |
| 394 | ) |
| 395 | down_weight = torch.randn( |
| 396 | num_experts, |
| 397 | unpadded_hidden_size, |
| 398 | unpadded_intermediate_size, |
| 399 | dtype=torch.bfloat16, |
| 400 | ) |
| 401 | |
| 402 | scalings = [2.3, 0.03, 7.3, 0.1, 0.004, 17.3, 1e4, 1e-4] |
| 403 | for i, start in enumerate(range(0, unpadded_hidden_size, 32)): |
| 404 | gate_up_weight[..., start : start + 32] *= scalings[i % len(scalings)] |
| 405 | for i, start in enumerate(range(0, unpadded_intermediate_size, 32)): |
| 406 | down_weight[..., start : start + 32] *= scalings[i % len(scalings)] |
| 407 | |
| 408 | def scatter_into_padded( |
| 409 | gate_up_weight: torch.Tensor, |
| 410 | down_weight: torch.Tensor, |
| 411 | padding_value: float, |
| 412 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 413 | """Place the source weights into buffers of the padded size, the |
| 414 | way the weight loader writes only the unpadded slice of a larger |
| 415 | allocation.""" |
| 416 | w13 = torch.full( |
| 417 | (num_experts, 2 * intermediate_size, hidden_size), |
| 418 | padding_value, |
| 419 | dtype=torch.bfloat16, |
| 420 | ) |
| 421 | w2 = torch.full( |
| 422 | (num_experts, hidden_size, intermediate_size), |
| 423 | padding_value, |
| 424 | dtype=torch.bfloat16, |
| 425 | ) |
| 426 | # w13 stacks gate and up, each padded to `intermediate_size`. |
| 427 | w13[:, :unpadded_intermediate_size, :unpadded_hidden_size] = gate_up_weight[ |
| 428 | :, :unpadded_intermediate_size, : |
| 429 | ] |
| 430 | w13[ |
| 431 | :, |
| 432 | intermediate_size : intermediate_size + unpadded_intermediate_size, |
| 433 | :unpadded_hidden_size, |
| 434 | ] = gate_up_weight[:, unpadded_intermediate_size:, :] |
| 435 | w2[:, :unpadded_hidden_size, :unpadded_intermediate_size] = down_weight |
| 436 | return w13, w2 |
| 437 | |
| 438 | # Checkpoint path: pre-quantize the source weights to MXFP4 with the |
| 439 | # ~same RTN recipe a real Quark checkpoint, then feed the |
| 440 | # already-packed tensors into `QuarkOCP_MX_MoEMethod`. Quark's |
| 441 | # `create_weights` allocates with `torch.zeros`, so its padding is zero. |
| 442 | checkpoint_w13, checkpoint_w2 = scatter_into_padded( |
| 443 | gate_up_weight, down_weight, padding_value=0.0 |
| 444 | ) |
| 445 | checkpoint_w13, checkpoint_w13_scale = mxfp4_quantize(checkpoint_w13) |
| 446 | checkpoint_w2, checkpoint_w2_scale = mxfp4_quantize(checkpoint_w2) |
| 447 | |
| 448 | checkpoint_method.create_weights( |
| 449 | layer=checkpoint_layer, |
| 450 | num_experts=num_experts, |
| 451 | hidden_size=hidden_size, |
| 452 | intermediate_size_per_partition=intermediate_size, |
| 453 | params_dtype=torch.bfloat16, |
| 454 | ) |
| 455 | checkpoint_layer.w13_weight.data.copy_(checkpoint_w13) |
| 456 | checkpoint_layer.w2_weight.data.copy_(checkpoint_w2) |
| 457 | checkpoint_layer.w13_weight_scale.data.copy_(checkpoint_w13_scale) |
| 458 | checkpoint_layer.w2_weight_scale.data.copy_(checkpoint_w2_scale) |
| 459 | checkpoint_method.process_weights_after_loading(checkpoint_layer) |
| 460 | |
| 461 | # Online path: feed the *same* raw bf16 source weights and let |
| 462 | # `Mxfp4OnlineMoEMethod` quantize them during |
| 463 | # `process_weights_after_loading`. The padding starts out |
| 464 | # uninitialized, standing in for the materialized meta tensor. |
| 465 | online_w13, online_w2 = scatter_into_padded( |
| 466 | gate_up_weight, down_weight, padding_value=float("nan") |
| 467 | ) |
| 468 | |
| 469 | online_method.create_weights( |
| 470 | layer=online_layer, |
| 471 | num_experts=num_experts, |
| 472 | hidden_size=hidden_size, |
| 473 | intermediate_size_per_partition=intermediate_size, |
| 474 | params_dtype=torch.bfloat16, |
| 475 | ) |
| 476 | replace_parameter(online_layer, "w13_weight", online_w13) |
| 477 | replace_parameter(online_layer, "w2_weight", online_w2) |
| 478 | online_method.process_weights_after_loading(online_layer) |
| 479 | |
| 480 | assert checkpoint_method.mxfp4_backend == online_method.mxfp4_backend |
| 481 | |
| 482 | assert_quantized_weights_equal( |
| 483 | checkpoint_layer, |
| 484 | online_layer, |
| 485 | keys=("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale"), |
| 486 | packed_weight_keys=("w13_weight", "w2_weight"), |
| 487 | ) |
| 488 | |
| 489 | |
| 490 | @pytest.mark.skipif( |
| 491 | not current_platform.is_cuda_alike(), reason="Only tested on ROCm/CUDA." |
| 492 | ) |
| 493 | @pytest.mark.parametrize("linear_backend", ["emulation", "aiter", "marlin"]) |
| 494 | def test_online_mxfp4_dense_matches_quark( |
| 495 | linear_backend: str, default_vllm_config, dist_init, monkeypatch |
| 496 | ): |
| 497 | """Ensures `Mxfp4OnlineLinearMethod` (online quantization) |
| 498 | and `QuarkOCP_MX` (AMD Quark checkpoints) produce the same weights, |
| 499 | with same linear backend used. |
| 500 | """ |
| 501 | if linear_backend == "marlin" and not current_platform.is_cuda(): |
| 502 | # `MarlinMxFp4LinearKernel.process_weights_after_loading` is CUDA-only |
| 503 | pytest.skip( |
| 504 | "MarlinMxFp4LinearKernel.process_weights_after_loading is CUDA-only." |
| 505 | ) |
| 506 | |
| 507 | if linear_backend == "aiter": |
| 508 | # `AiterMxfp4LinearKernel` is registered only for ROCm in |
| 509 | # `_POSSIBLE_MXFP4_KERNELS`, and `is_supported` gates on |
| 510 | # `current_platform.supports_mx()`. Force it onto the CUDA |
| 511 | # kernel list and bypass the platform gate so this backend can |
| 512 | # be exercised on CUDA hosts too. |
| 513 | |
| 514 | monkeypatch.setattr( |
| 515 | AiterMxfp4LinearKernel, |
| 516 | "is_supported", |
| 517 | classmethod(lambda cls, compute_capability=None: (True, None)), |
| 518 | ) |
| 519 | monkeypatch.setitem( |
| 520 | _POSSIBLE_MXFP4_KERNELS, |
| 521 | PlatformEnum.CUDA, |
| 522 | [ |
| 523 | AiterMxfp4LinearKernel, |
| 524 | *_POSSIBLE_MXFP4_KERNELS.get(PlatformEnum.CUDA, []), |
| 525 | ], |
| 526 | ) |
| 527 | elif linear_backend == "emulation": |
| 528 | # `EmulationMxfp4LinearKernel.can_implement` gates on `has_quark()`, |
| 529 | # EmulationMxfp4LinearKernel.process_weights_after_loading does not require it. |
| 530 | monkeypatch.setattr( |
| 531 | "vllm.model_executor.kernels.linear.mxfp4.emulation.has_quark", |
| 532 | lambda: True, |
| 533 | ) |
| 534 | |
| 535 | default_vllm_config.model_config = ModelConfig() |
| 536 | default_vllm_config.kernel_config.linear_backend = linear_backend |
| 537 | |
| 538 | input_size = 256 |
| 539 | output_size = 128 |
| 540 | device = current_platform.device_type |
| 541 | |
| 542 | # `create_weights` implementations use plain `torch.empty` with no explicit |
| 543 | # device, so without this context they would default to CPU and diverge |
| 544 | # from the (GPU) tensors produced by `mxfp4_quantize`/`replace_parameter`. |
| 545 | with torch.device(device): |
| 546 | checkpoint_layer = torch.nn.Module() |
| 547 | online_layer = torch.nn.Module() |
| 548 | for layer in (checkpoint_layer, online_layer): |
| 549 | # `LinearBase` subclasses normally set these attributes on the |
| 550 | # layer before `create_weights` is called; the marlin kernel's |
| 551 | # `process_weights_after_loading` reads them directly off the |
| 552 | # layer, so a bare `torch.nn.Module` needs them set explicitly. |
| 553 | layer.input_size = input_size |
| 554 | layer.output_size = output_size |
| 555 | layer.input_size_per_partition = input_size |
| 556 | layer.output_size_per_partition = output_size |
| 557 | layer.params_dtype = torch.bfloat16 |
| 558 | |
| 559 | weight = torch.randn(output_size, input_size, dtype=torch.bfloat16) |
| 560 | |
| 561 | scalings = [2.3, 0.03, 7.3, 0.1, 0.004, 17.3, 1e4, 1e-4] |
| 562 | for i in range(input_size // 32): |
| 563 | weight[:, i * 32 : (i + 1) * 32] *= scalings[i % len(scalings)] |
| 564 | |
| 565 | # Checkpoint path: pre-quantize the source weight to MXFP4 with the |
| 566 | # ~same RTN recipe as a real Quark checkpoint, then feed the |
| 567 | # already-packed tensors into `QuarkOCP_MX`. |
| 568 | checkpoint_weight, checkpoint_weight_scale = mxfp4_quantize(weight) |
| 569 | |
| 570 | checkpoint_scheme = QuarkOCP_MX( |
| 571 | weight_quant_key=kMxfp4Static, |
| 572 | activation_quant_key=kMxfp4Dynamic, |
| 573 | ) |
| 574 | checkpoint_scheme.create_weights( |
| 575 | layer=checkpoint_layer, |
| 576 | output_partition_sizes=[output_size], |
| 577 | input_size_per_partition=input_size, |
| 578 | params_dtype=torch.bfloat16, |
| 579 | weight_loader=default_weight_loader, |
| 580 | ) |
| 581 | checkpoint_layer.weight.data.copy_(checkpoint_weight) |
| 582 | checkpoint_layer.weight_scale.data.copy_(checkpoint_weight_scale) |
| 583 | checkpoint_scheme.process_weights_after_loading(checkpoint_layer) |
| 584 | |
| 585 | # Online path: feed the *same* raw bf16 source weight and let |
| 586 | # `Mxfp4OnlineLinearMethod` quantize it during |
| 587 | # `process_weights_after_loading`. |
| 588 | online_method = Mxfp4OnlineLinearMethod() |
| 589 | online_method.create_weights( |
| 590 | layer=online_layer, |
| 591 | input_size_per_partition=input_size, |
| 592 | output_partition_sizes=[output_size], |
| 593 | input_size=input_size, |
| 594 | output_size=output_size, |
| 595 | params_dtype=torch.bfloat16, |
| 596 | weight_loader=default_weight_loader, |
| 597 | ) |
| 598 | replace_parameter(online_layer, "weight", weight.clone()) |
| 599 | online_method.process_weights_after_loading(online_layer) |
| 600 | |
| 601 | assert type(checkpoint_scheme.ocp_mx_linear) is type(online_method.kernel) |
| 602 | |
| 603 | assert_quantized_weights_equal( |
| 604 | checkpoint_layer, |
| 605 | online_layer, |
| 606 | keys=("weight", "weight_scale"), |
| 607 | packed_weight_keys=("weight",), |
| 608 | ) |
| 609 | |
| 610 | |
| 611 | def test_quark_ocp_mx_dense_rejects_unaligned_input_partition(): |
| 612 | scheme = object.__new__(QuarkOCP_MX) |
| 613 | layer = torch.nn.Module() |
| 614 | layer.prefix = "model.layers.0.mlp.shared_expert.down_proj" |
| 615 | |
| 616 | with pytest.raises( |
| 617 | ValueError, |
| 618 | match=( |
| 619 | "input size per partition of 80, which must be divisible by " |
| 620 | "the OCP MX group size 32" |
| 621 | ), |
| 622 | ): |
| 623 | scheme.create_weights( |
| 624 | layer=layer, |
| 625 | output_partition_sizes=[2560], |
| 626 | input_size_per_partition=80, |
| 627 | params_dtype=torch.bfloat16, |
| 628 | weight_loader=lambda *_: None, |
| 629 | ) |