diff --git a/docs/features/quantization/online.md b/docs/features/quantization/online.md index 67e28a270d..c7257c8d39 100644 --- a/docs/features/quantization/online.md +++ b/docs/features/quantization/online.md @@ -20,6 +20,16 @@ llm = LLM("meta-llama/Llama-3.1-8B", quantization="fp8_per_block") # MXFP8 quantization for weights and activations llm = LLM("meta-llama/Llama-3.1-8B", quantization="mxfp8") + +# MXFP4 weight; activation quantization depends on the `linear_backend` picked +llm = LLM("meta-llama/Llama-3.1-8B", quantization="mxfp4") + +# MXFP4 MOE-only weight and activation quantization +llm = LLM( + "Qwen/Qwen3.5-35B-A3B", + quantization="mxfp4", + quantization_config={"linear": {"activation": None, "weight": None}} +) ``` Or with the CLI: @@ -28,6 +38,10 @@ Or with the CLI: vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_tensor vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_block vllm serve meta-llama/Llama-3.1-8B --quantization mxfp8 +vllm serve meta-llama/Llama-3.1-8B --quantization mxfp4 + +vllm serve Qwen/Qwen3.5-35B-A3B --quantization mxfp4 \ + --quantization-config '{"linear":{"activation":null,"weight":null}}' ``` ## Supported Schemes @@ -37,6 +51,7 @@ vllm serve meta-llama/Llama-3.1-8B --quantization mxfp8 | `fp8_per_tensor` | fp8_e4m3 data, fp32 per-tensor scale | fp8_e4m3 data, fp32 per-tensor scale | On some GPUs (Ada, Hopper) linear activations use per-token scaling for better performance | | `fp8_per_block` | fp8_e4m3 data, fp32 per-128x128-block scale | fp8_e4m3 data, fp32 per-1x128-block scale | | | `mxfp8` | fp8_e4m3 data, e8m0 per-1x32-block scale | fp8_e4m3 data, e8m0 per-1x32-block scale | Requires SM 100+ (Blackwell or newer) for w8a8, other GPUs use a w8a16 fallback | +| `mxfp4` | fp4_e2m1 data, e8m0 per-1x32-block scale ([OCP MX specs](https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf)) | - linear: fp4_e2m1 data, e8m0 per-1x32-block scale in some backends, or BF16.
- MOE: fp4_e2m1 data, e8m0 per-1x32-block scale. | Linear MXFP4 backend is auto-selected per platform, not enforcing activation dtype. Some use BF16 activation. Use `--linear-backend` to pin one (e.g. `--linear-backend flashinfer`). | ## Advanced Configuration diff --git a/tests/evals/gsm8k/configs/Qwen3-30B-A3B-MXFP4-AITER-TP2-online.yaml b/tests/evals/gsm8k/configs/Qwen3-30B-A3B-MXFP4-AITER-TP2-online.yaml new file mode 100644 index 0000000000..d954415f27 --- /dev/null +++ b/tests/evals/gsm8k/configs/Qwen3-30B-A3B-MXFP4-AITER-TP2-online.yaml @@ -0,0 +1,15 @@ +model_name: "Qwen/Qwen3-30B-A3B" +accuracy_threshold: 0.89 +tolerance: 0.03 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --max-model-len 4096 + --tensor-parallel-size 2 + --gpu-memory-utilization 0.35 + --moe-backend aiter + --quantization mxfp4 + --quantization-config '{"linear":{"activation":null,"weight":null}}' +env: + VLLM_ROCM_USE_AITER: "1" + ENABLE_CK: "0" # Avoid AITER CK-based MHA JIT compilation to save time diff --git a/tests/evals/gsm8k/configs/Qwen3-30B-A3B-MXFP4-AITER-TP2.yaml b/tests/evals/gsm8k/configs/Qwen3-30B-A3B-MXFP4-AITER-TP2.yaml new file mode 100644 index 0000000000..5bd8735453 --- /dev/null +++ b/tests/evals/gsm8k/configs/Qwen3-30B-A3B-MXFP4-AITER-TP2.yaml @@ -0,0 +1,14 @@ +# Note: this model was produced from Qwen3-30B-A3B (BF16) +model_name: "amd-quark/Qwen3-30B-A3B-MXFP4" +accuracy_threshold: 0.89 +tolerance: 0.03 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --max-model-len 4096 + --tensor-parallel-size 2 + --gpu-memory-utilization 0.35 + --moe-backend aiter +env: + VLLM_ROCM_USE_AITER: "1" + ENABLE_CK: "0" # Avoid AITER CK-based MHA JIT compilation to save time diff --git a/tests/evals/gsm8k/configs/Qwen3.5-35B-A3B-MXFP4-AITER-TP2.yaml b/tests/evals/gsm8k/configs/Qwen3.5-35B-A3B-MXFP4-AITER-TP2.yaml index ca5cc450c0..56c273db57 100644 --- a/tests/evals/gsm8k/configs/Qwen3.5-35B-A3B-MXFP4-AITER-TP2.yaml +++ b/tests/evals/gsm8k/configs/Qwen3.5-35B-A3B-MXFP4-AITER-TP2.yaml @@ -1,3 +1,4 @@ +# NOTE: this model was produced from Qwen/Qwen3.5-35B-A3B-FP8 model_name: "amd/Qwen3.5-35B-A3B-MXFP4" accuracy_threshold: 0.89 tolerance: 0.03 diff --git a/tests/evals/gsm8k/configs/Qwen3.5-35B-A3B-MXFP4-EMU-TP2.yaml b/tests/evals/gsm8k/configs/Qwen3.5-35B-A3B-MXFP4-EMU-TP2.yaml index 70217d2651..855f16fd4f 100644 --- a/tests/evals/gsm8k/configs/Qwen3.5-35B-A3B-MXFP4-EMU-TP2.yaml +++ b/tests/evals/gsm8k/configs/Qwen3.5-35B-A3B-MXFP4-EMU-TP2.yaml @@ -1,3 +1,4 @@ +# NOTE: this model was produced from Qwen/Qwen3.5-35B-A3B-FP8 model_name: "amd/Qwen3.5-35B-A3B-MXFP4" accuracy_threshold: 0.89 tolerance: 0.03 diff --git a/tests/kernels/quantization/test_mxfp4_kernel_selection.py b/tests/kernels/quantization/test_mxfp4_kernel_selection.py index c3f8bc094d..28ce615ce6 100644 --- a/tests/kernels/quantization/test_mxfp4_kernel_selection.py +++ b/tests/kernels/quantization/test_mxfp4_kernel_selection.py @@ -114,9 +114,15 @@ def test_weight_only_kernels_reject_non_mxfp4_activation(kernel_cls): "activation_quant_key", [None, kMxfp4Dynamic, kMxfp6E3M2Dynamic, kMxfp6E2M3Dynamic], ) -def test_emulation_kernel_accepts_any_config(activation_quant_key): +def test_emulation_kernel_accepts_any_config(activation_quant_key, monkeypatch): """EmulationMxfp4LinearKernel is the universal fallback: it must accept every supported activation format.""" + # `EmulationMxfp4LinearKernel.can_implement` gates on `has_quark()`, + # which we are not testing here. + monkeypatch.setattr( + "vllm.model_executor.kernels.linear.mxfp4.emulation.has_quark", + lambda: True, + ) config = MxFp4LinearLayerConfig(activation_quant_key=activation_quant_key) with patch( "vllm.model_executor.kernels.linear._get_linear_backend", @@ -126,9 +132,15 @@ def test_emulation_kernel_accepts_any_config(activation_quant_key): assert can_implement, reason -def test_emulation_kernel_derives_quant_dequant_func_from_config(): +def test_emulation_kernel_derives_quant_dequant_func_from_config(monkeypatch): """quant_dequant_func must be derived purely from the config's activation QuantKey, not set externally.""" + # `EmulationMxfp4LinearKernel.can_implement` gates on `has_quark()`, + # which we are not testing here. + monkeypatch.setattr( + "vllm.model_executor.kernels.linear.mxfp4.emulation.has_quark", + lambda: True, + ) with patch( "vllm.model_executor.kernels.linear._get_linear_backend", return_value="emulation", diff --git a/tests/quantization/test_online.py b/tests/quantization/test_online.py index 1a50c80413..3e95e92245 100644 --- a/tests/quantization/test_online.py +++ b/tests/quantization/test_online.py @@ -14,6 +14,7 @@ from tests.quantization.utils import ( is_quant_method_supported, ) from vllm import _custom_ops as ops +from vllm._aiter_ops import rocm_aiter_ops from vllm._custom_ops import scaled_fp4_quant from vllm.model_executor.layers.linear import UnquantizedLinearMethod from vllm.model_executor.layers.quantization.online.fp8 import ( @@ -27,6 +28,10 @@ from vllm.model_executor.layers.quantization.online.fp8 import ( _is_tp_sharded, ) from vllm.model_executor.layers.quantization.online.int8 import Int8OnlineMoEMethod +from vllm.model_executor.layers.quantization.online.mxfp4 import ( + Mxfp4OnlineLinearMethod, + Mxfp4OnlineMoEMethod, +) from vllm.model_executor.layers.quantization.online.nvfp4 import ( Nvfp4OnlineMoEMethod, _quantize_moe_weight_to_nvfp4, @@ -40,6 +45,17 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe +if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx942, on_gfx950 +else: + + def on_gfx950() -> bool: + return False + + def on_gfx942() -> bool: + return False + + DEVICE = current_platform.device_type @@ -84,6 +100,12 @@ DEVICE = current_platform.device_type Fp8PerTensorOnlineLinearMethod, Fp8PerTensorOnlineMoEMethod, ), + ( + "mxfp4", + None, + Mxfp4OnlineLinearMethod, + Mxfp4OnlineMoEMethod, + ), ], ) @pytest.mark.parametrize( @@ -106,8 +128,14 @@ def test_online_quantization( Does not test performance, peak memory usage, etc. """ - if use_rocm_aiter: - monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + # TODO: Relax this condition once there is a native MXFP4_MXFP4 + # linear/moe backend supported on cuda. + if quant_scheme == "mxfp4" and not (on_gfx950() or on_gfx942()): + pytest.skip("mxfp4 online quantization is only tested on AMD gfx942, gfx950.") + + if current_platform.is_rocm(): + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1" if use_rocm_aiter else "0") + rocm_aiter_ops.refresh_env_variables() if current_platform.is_xpu() and quant_scheme == "fp8_per_block": pytest.skip("Skip test for online fp8_per_block on XPU platform.") @@ -144,7 +172,10 @@ def test_online_quantization( if moe is not None: assert isinstance(moe._quant_method, expected_moe_cls) - if current_platform.is_cuda() or current_platform.is_xpu(): + if quant_scheme == "mxfp4": + # Packed e2m1 values, two per byte. + assert o_proj.weight.dtype == torch.uint8 + elif current_platform.is_cuda() or current_platform.is_xpu(): assert o_proj.weight.dtype == torch.float8_e4m3fn elif current_platform.is_rocm(): assert o_proj.weight.dtype == current_platform.fp8_dtype() diff --git a/tests/quantization/test_online_mxfp4.py b/tests/quantization/test_online_mxfp4.py new file mode 100644 index 0000000000..660d404054 --- /dev/null +++ b/tests/quantization/test_online_mxfp4.py @@ -0,0 +1,606 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Correctness tests for the online MXFP4 weight quantization backends. + +Each backend (Triton, aiter, XPU) quantizes a synthetic bf16/fp16 +tensor to MXFP4, the result is dequantized back to bf16/fp16, and compared +against the pure-torch reference in `reference_mxfp4.py`. +""" + +import pytest +import torch + +from vllm.config.model import ModelConfig +from vllm.model_executor.kernels.linear import _POSSIBLE_MXFP4_KERNELS +from vllm.model_executor.kernels.linear.mxfp4.aiter import ( + AiterMxfp4LinearKernel, +) +from vllm.model_executor.layers.fused_moe import FusedMoEFactory +from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts +from vllm.model_executor.layers.quantization.online.mxfp4 import ( + Mxfp4OnlineLinearMethod, + Mxfp4OnlineMoEMethod, +) +from vllm.model_executor.layers.quantization.quark.quark_moe import ( + QuarkOCP_MX_MoEMethod, +) +from vllm.model_executor.layers.quantization.quark.schemes.quark_ocp_mx import ( + QuarkOCP_MX, +) +from vllm.model_executor.layers.quantization.quark.utils import ( + quark_quantize_weight_to_mxfp4, +) +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( + downcast_to_mxfp, + mxfp4_quantize, + quant_dequant_mxfp4, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.utils import replace_parameter +from vllm.platforms import PlatformEnum, current_platform + +from .reference_mxfp4 import dq_mxfp4_torch, qdq_mxfp4_torch + + +def fix_negative_zeros(tensor: torch.Tensor) -> torch.Tensor: + """Normalize FP4 e2m1 negative-zero codewords (0b1000) to positive zero + (0b0000) in a packed uint8 tensor (two e2m1 values per byte, low nibble + first). + + -0.0 and +0.0 dequantize to the same float value, but different MXFP4 + quantization backends disagree on which one they emit for values that + round to zero magnitude. This normalizes both packed representations so + tensors from different backends can be compared with `torch.equal`. + """ + assert tensor.dtype == torch.uint8, ( + f"Expected a torch.uint8 tensor, got {tensor.dtype}" + ) + + low_nibble = tensor & 0x0F + high_nibble = tensor & 0xF0 + + low_nibble = torch.where( + low_nibble == 0x08, torch.zeros_like(low_nibble), low_nibble + ) + high_nibble = torch.where( + high_nibble == 0x80, torch.zeros_like(high_nibble), high_nibble + ) + + return low_nibble | high_nibble + + +def assert_quantized_weights_equal( + checkpoint_layer: torch.nn.Module, + online_layer: torch.nn.Module, + keys: tuple[str, ...], + packed_weight_keys: tuple[str, ...], +) -> None: + """Assert the checkpoint and online paths produced byte-identical weights. + + `packed_weight_keys` lists the entries holding packed FP4 codewords, which + need negative-zero normalization before comparison. + """ + for key in keys: + checkpoint_tensor = getattr(checkpoint_layer, key).view(torch.uint8) + online_tensor = getattr(online_layer, key).view(torch.uint8) + + # NOTE: AMD Quark checkpoints use exclusively **positive** zeros, + # while other mxfp4_quantize implementations from mxfp4_utils.py may not. + if key in packed_weight_keys: + checkpoint_tensor = fix_negative_zeros(checkpoint_tensor) + online_tensor = fix_negative_zeros(online_tensor) + + assert checkpoint_tensor.shape == online_tensor.shape + assert checkpoint_tensor.dtype == online_tensor.dtype + num_mismatched = (checkpoint_tensor != online_tensor).sum().item() + total = checkpoint_tensor.numel() + + assert torch.equal(checkpoint_tensor, online_tensor), ( + f"{key}: {num_mismatched}/{total} " + f"({100 * num_mismatched / total:.4f}%) mismatched bytes" + ) + + +def _skip_reason_if_unavailable(backend: str, dtype: torch.dtype) -> str | None: + """Return a skip reason if `backend` cannot run on the current host.""" + if backend == "triton": + if not (current_platform.is_cuda() or current_platform.is_rocm()): + return "Triton MXFP4 kernel requires a CUDA or ROCm GPU." + return None + if backend == "aiter": + from vllm._aiter_ops import is_aiter_found_and_supported + + if not is_aiter_found_and_supported(): + return "aiter is not available/supported on this platform." + if dtype != torch.bfloat16: + return "aiter's dynamic_mxfp4_quant only supports bfloat16 input." + return None + if backend == "xpu": + if not current_platform.is_xpu(): + return "not on XPU platform." + return None + if backend == "quark": + try: + import quark.torch.kernel.mx # noqa: F401 + except ImportError: + return "amd-quark is not installed." + return None + raise ValueError(f"Unknown backend {backend}") + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("backend", ["triton", "aiter", "xpu", "quark"]) +def test_mxfp4_quantization_correctness(backend: str, dtype: torch.dtype): + """ + Tests that the different implementations of mxfp4_quantize + in mxfp4_utils.py all match. + """ + skip_reason = _skip_reason_if_unavailable(backend, dtype) + if skip_reason is not None: + pytest.skip(skip_reason) + + torch.manual_seed(3) + + num_rows = 64 + hidden_size = 32 * 32 # multiple 32-element MXFP4 blocks + device = current_platform.device_type + + x = (torch.rand(num_rows, hidden_size, dtype=dtype, device=device) - 0.5) * 2 + # Vary the magnitude block-to-block so several scale exponents are + # exercised, rather than a single one for the whole tensor. + scalings = [2.3, 0.03, 7.3, 0.1, 0.004, 17.3, 1e4, 1e-4] + for i in range(hidden_size // 32): + x[:, i * 32 : (i + 1) * 32] *= scalings[i % len(scalings)] + + if backend == "triton": + x_fp4, x_scale, _ = downcast_to_mxfp(x, axis=-1) + elif backend == "aiter": + x_fp4, x_scale = quark_quantize_weight_to_mxfp4(x) + elif backend == "xpu": + # TODO: enable this test on XPU + pytest.skip("xpu mxfp4 quantization match to reference is untested") + # from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( + # xpu_mxfp4_quantize, + # ) + + # x_fp4, x_scale = xpu_mxfp4_quantize(x) + elif backend == "quark": + result = quant_dequant_mxfp4(x, scale_calculation_mode="even") + else: + raise ValueError(f"Unknown backend {backend}") + + if backend != "quark": + result = dq_mxfp4_torch(x_fp4, x_scale, x.dtype) + reference = qdq_mxfp4_torch(x, scale_calculation_mode="even") + + assert torch.equal(result, reference) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="Only tested on ROCm/CUDA." +) +@pytest.mark.parametrize("tp_size", [2, 4, 8]) +def test_online_mxfp4_tp_weight_quant_matches_unsharded(tp_size: int): + """TP-aligned MXFP4 shards match slices of unsharded quantization.""" + torch.manual_seed(3) + + device = current_platform.device_type + hidden_size = 1024 + intermediate_size = 512 + + dense_weight = torch.randn( + hidden_size, intermediate_size, dtype=torch.bfloat16, device=device + ) + w13_weight = torch.randn( + 2, + 2 * intermediate_size, + hidden_size, + dtype=torch.bfloat16, + device=device, + ) + w2_weight = torch.randn( + 2, + hidden_size, + intermediate_size, + dtype=torch.bfloat16, + device=device, + ) + + dense_quantized, dense_scales = mxfp4_quantize(dense_weight) + w13_quantized, w13_scales = mxfp4_quantize(w13_weight) + w2_quantized, w2_scales = mxfp4_quantize(w2_weight) + + for tp_rank in range(tp_size): + output_per_partition = hidden_size // tp_size + input_per_partition = intermediate_size // tp_size + output_start = tp_rank * output_per_partition + input_start = tp_rank * input_per_partition + + column_quantized, column_scales = mxfp4_quantize( + dense_weight.narrow(0, output_start, output_per_partition) + ) + assert torch.equal( + column_quantized, + dense_quantized.narrow(0, output_start, output_per_partition), + ) + assert torch.equal( + column_scales, + dense_scales.narrow(0, output_start, output_per_partition), + ) + + row_quantized, row_scales = mxfp4_quantize( + dense_weight.narrow(1, input_start, input_per_partition) + ) + assert torch.equal( + row_quantized, + dense_quantized.narrow(1, input_start // 2, input_per_partition // 2), + ) + assert torch.equal( + row_scales, + dense_scales.narrow(1, input_start // 32, input_per_partition // 32), + ) + + w13_quantized_shard, w13_scales_shard = mxfp4_quantize( + w13_weight.narrow(1, 2 * input_start, 2 * input_per_partition) + ) + assert torch.equal( + w13_quantized_shard, + w13_quantized.narrow(1, 2 * input_start, 2 * input_per_partition), + ) + assert torch.equal( + w13_scales_shard, + w13_scales.narrow(1, 2 * input_start, 2 * input_per_partition), + ) + + w2_quantized_shard, w2_scales_shard = mxfp4_quantize( + w2_weight.narrow(2, input_start, input_per_partition) + ) + assert torch.equal( + w2_quantized_shard, + w2_quantized.narrow(2, input_start // 2, input_per_partition // 2), + ) + assert torch.equal( + w2_scales_shard, + w2_scales.narrow(2, input_start // 32, input_per_partition // 32), + ) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="Only tested on ROCm/CUDA." +) +@pytest.mark.parametrize("moe_backend", ["aiter", "emulation"]) +@pytest.mark.parametrize( + "unpadded_hidden_size,unpadded_intermediate_size", + [ + pytest.param(256, 256, id="no_padding"), + # Not multiples of the MXFP4 block size (32), so that every backend -- + # including emulation, whose round-up granularity is 32 -- has to pad. + pytest.param(240, 208, id="padding"), + ], +) +def test_online_mxfp4_moe_matches_quark( + moe_backend: str, + unpadded_hidden_size: int, + unpadded_intermediate_size: int, + default_vllm_config, + dist_init, + monkeypatch, +): + """ + Ensures `Mxfp4OnlineMoEMethod` (online quantization) + and `QuarkOCP_MX_MoEMethod` (AMD Quark checkpoints) produce the same weights, + with same MOE backend used. + + The `padding` case additionally covers `maybe_roundup_sizes`: the online + path allocates its bf16 weights with `torch.empty_strided` and the loader + only writes the unpadded slice, so the padding is filled with NaN here to + reproduce that uninitialized state. It must be zeroed before quantization + for the two paths to agree, since Quark allocates its padding with + `torch.zeros`. + """ + if moe_backend == "aiter": + from vllm._aiter_ops import rocm_aiter_ops + + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1") + monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: True) + rocm_aiter_ops.refresh_env_variables() + + if current_platform.is_cuda(): + pytest.skip( + "mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4 requires " + "rocm_aiter_ops in weight conversion, not compatible on cuda" + ) + elif moe_backend == "emulation": + # `OCP_MXQuantizationEmulationTritonExperts.is_supported_config` gates + # on `has_quark()`, but its weight processing does not require it. + monkeypatch.setattr( + "vllm.model_executor.layers.fused_moe.experts." + "ocp_mx_emulation_moe.has_quark", + lambda: True, + ) + + default_vllm_config.model_config = ModelConfig() + + num_experts = 4 + device = current_platform.device_type + + def make_layer(prefix: str) -> RoutedExperts: + runner = FusedMoEFactory( + num_experts=num_experts, + top_k=2, + hidden_size=unpadded_hidden_size, + intermediate_size=unpadded_intermediate_size, + prefix=prefix, + ) + layer = runner.routed_experts + layer.moe_config.moe_backend = moe_backend + return layer + + # `create_weights` implementations use plain `torch.zeros`/`torch.randn` + # with no explicit device, so without this context they would default + # to CPU and diverge from the (GPU) tensors produced by + # `mxfp4_quantize`/`replace_parameter`. + with torch.device(device): + checkpoint_layer = make_layer("checkpoint_layer") + online_layer = make_layer("online_layer") + + checkpoint_method = QuarkOCP_MX_MoEMethod( + weight_config={"qscheme": "per_group", "dtype": "fp4"}, + input_config={"dtype": "fp4", "is_dynamic": True}, + moe=checkpoint_layer.moe_config, + ) + online_method = Mxfp4OnlineMoEMethod(layer=online_layer) + + # `RoutedExperts.__init__` applies this round-up in production; these + # layers are built without a quant config, so it is applied explicitly. + # Both methods must agree on the padded sizes. + roundup_arguments = dict( + hidden_size=unpadded_hidden_size, + intermediate_size_per_partition=unpadded_intermediate_size, + act_dtype=torch.bfloat16, + moe_parallel_config=online_layer.moe_config.moe_parallel_config, + ) + hidden_size, intermediate_size = online_method.maybe_roundup_sizes( + **roundup_arguments + ) + assert ( + hidden_size, + intermediate_size, + ) == checkpoint_method.maybe_roundup_sizes(**roundup_arguments) + + if unpadded_hidden_size == 240: + # Padded case. + assert hidden_size > unpadded_hidden_size + assert intermediate_size > unpadded_intermediate_size + + for layer in (checkpoint_layer, online_layer): + assert layer.moe_config.hidden_dim_unpadded == unpadded_hidden_size + assert ( + layer.moe_config.intermediate_size_per_partition_unpadded + == unpadded_intermediate_size + ) + layer.moe_config.hidden_dim = hidden_size + layer.moe_config.intermediate_size_per_partition = intermediate_size + + gate_up_weight = torch.randn( + num_experts, + 2 * unpadded_intermediate_size, + unpadded_hidden_size, + dtype=torch.bfloat16, + ) + down_weight = torch.randn( + num_experts, + unpadded_hidden_size, + unpadded_intermediate_size, + dtype=torch.bfloat16, + ) + + scalings = [2.3, 0.03, 7.3, 0.1, 0.004, 17.3, 1e4, 1e-4] + for i, start in enumerate(range(0, unpadded_hidden_size, 32)): + gate_up_weight[..., start : start + 32] *= scalings[i % len(scalings)] + for i, start in enumerate(range(0, unpadded_intermediate_size, 32)): + down_weight[..., start : start + 32] *= scalings[i % len(scalings)] + + def scatter_into_padded( + gate_up_weight: torch.Tensor, + down_weight: torch.Tensor, + padding_value: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Place the source weights into buffers of the padded size, the + way the weight loader writes only the unpadded slice of a larger + allocation.""" + w13 = torch.full( + (num_experts, 2 * intermediate_size, hidden_size), + padding_value, + dtype=torch.bfloat16, + ) + w2 = torch.full( + (num_experts, hidden_size, intermediate_size), + padding_value, + dtype=torch.bfloat16, + ) + # w13 stacks gate and up, each padded to `intermediate_size`. + w13[:, :unpadded_intermediate_size, :unpadded_hidden_size] = gate_up_weight[ + :, :unpadded_intermediate_size, : + ] + w13[ + :, + intermediate_size : intermediate_size + unpadded_intermediate_size, + :unpadded_hidden_size, + ] = gate_up_weight[:, unpadded_intermediate_size:, :] + w2[:, :unpadded_hidden_size, :unpadded_intermediate_size] = down_weight + return w13, w2 + + # Checkpoint path: pre-quantize the source weights to MXFP4 with the + # ~same RTN recipe a real Quark checkpoint, then feed the + # already-packed tensors into `QuarkOCP_MX_MoEMethod`. Quark's + # `create_weights` allocates with `torch.zeros`, so its padding is zero. + checkpoint_w13, checkpoint_w2 = scatter_into_padded( + gate_up_weight, down_weight, padding_value=0.0 + ) + checkpoint_w13, checkpoint_w13_scale = mxfp4_quantize(checkpoint_w13) + checkpoint_w2, checkpoint_w2_scale = mxfp4_quantize(checkpoint_w2) + + checkpoint_method.create_weights( + layer=checkpoint_layer, + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size_per_partition=intermediate_size, + params_dtype=torch.bfloat16, + ) + checkpoint_layer.w13_weight.data.copy_(checkpoint_w13) + checkpoint_layer.w2_weight.data.copy_(checkpoint_w2) + checkpoint_layer.w13_weight_scale.data.copy_(checkpoint_w13_scale) + checkpoint_layer.w2_weight_scale.data.copy_(checkpoint_w2_scale) + checkpoint_method.process_weights_after_loading(checkpoint_layer) + + # Online path: feed the *same* raw bf16 source weights and let + # `Mxfp4OnlineMoEMethod` quantize them during + # `process_weights_after_loading`. The padding starts out + # uninitialized, standing in for the materialized meta tensor. + online_w13, online_w2 = scatter_into_padded( + gate_up_weight, down_weight, padding_value=float("nan") + ) + + online_method.create_weights( + layer=online_layer, + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size_per_partition=intermediate_size, + params_dtype=torch.bfloat16, + ) + replace_parameter(online_layer, "w13_weight", online_w13) + replace_parameter(online_layer, "w2_weight", online_w2) + online_method.process_weights_after_loading(online_layer) + + assert checkpoint_method.mxfp4_backend == online_method.mxfp4_backend + + assert_quantized_weights_equal( + checkpoint_layer, + online_layer, + keys=("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale"), + packed_weight_keys=("w13_weight", "w2_weight"), + ) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="Only tested on ROCm/CUDA." +) +@pytest.mark.parametrize("linear_backend", ["emulation", "aiter", "marlin"]) +def test_online_mxfp4_dense_matches_quark( + linear_backend: str, default_vllm_config, dist_init, monkeypatch +): + """ + Ensures `Mxfp4OnlineLinearMethod` (online quantization) + and `QuarkOCP_MX` (AMD Quark checkpoints) produce the same weights, + with same linear backend used. + """ + if linear_backend == "marlin" and not current_platform.is_cuda(): + # `MarlinMxFp4LinearKernel.process_weights_after_loading` is CUDA-only + pytest.skip( + "MarlinMxFp4LinearKernel.process_weights_after_loading is CUDA-only." + ) + + if linear_backend == "aiter": + # `AiterMxfp4LinearKernel` is registered only for ROCm in + # `_POSSIBLE_MXFP4_KERNELS`, and `is_supported` gates on + # `current_platform.supports_mx()`. Force it onto the CUDA + # kernel list and bypass the platform gate so this backend can + # be exercised on CUDA hosts too. + + monkeypatch.setattr( + AiterMxfp4LinearKernel, + "is_supported", + classmethod(lambda cls, compute_capability=None: (True, None)), + ) + monkeypatch.setitem( + _POSSIBLE_MXFP4_KERNELS, + PlatformEnum.CUDA, + [ + AiterMxfp4LinearKernel, + *_POSSIBLE_MXFP4_KERNELS.get(PlatformEnum.CUDA, []), + ], + ) + elif linear_backend == "emulation": + # `EmulationMxfp4LinearKernel.can_implement` gates on `has_quark()`, + # EmulationMxfp4LinearKernel.process_weights_after_loading does not require it. + monkeypatch.setattr( + "vllm.model_executor.kernels.linear.mxfp4.emulation.has_quark", + lambda: True, + ) + + default_vllm_config.model_config = ModelConfig() + default_vllm_config.kernel_config.linear_backend = linear_backend + + input_size = 256 + output_size = 128 + device = current_platform.device_type + + # `create_weights` implementations use plain `torch.empty` with no explicit + # device, so without this context they would default to CPU and diverge + # from the (GPU) tensors produced by `mxfp4_quantize`/`replace_parameter`. + with torch.device(device): + checkpoint_layer = torch.nn.Module() + online_layer = torch.nn.Module() + for layer in (checkpoint_layer, online_layer): + # `LinearBase` subclasses normally set these attributes on the + # layer before `create_weights` is called; the marlin kernel's + # `process_weights_after_loading` reads them directly off the + # layer, so a bare `torch.nn.Module` needs them set explicitly. + layer.input_size = input_size + layer.output_size = output_size + layer.input_size_per_partition = input_size + layer.output_size_per_partition = output_size + layer.params_dtype = torch.bfloat16 + + weight = torch.randn(output_size, input_size, dtype=torch.bfloat16) + + scalings = [2.3, 0.03, 7.3, 0.1, 0.004, 17.3, 1e4, 1e-4] + for i in range(input_size // 32): + weight[:, i * 32 : (i + 1) * 32] *= scalings[i % len(scalings)] + + # Checkpoint path: pre-quantize the source weight to MXFP4 with the + # ~same RTN recipe as a real Quark checkpoint, then feed the + # already-packed tensors into `QuarkOCP_MX`. + checkpoint_weight, checkpoint_weight_scale = mxfp4_quantize(weight) + + checkpoint_scheme = QuarkOCP_MX( + weight_quant_spec={"qscheme": "per_group", "dtype": "fp4"}, + input_quant_spec={"dtype": "fp4", "is_dynamic": True}, + ) + checkpoint_scheme.create_weights( + layer=checkpoint_layer, + output_partition_sizes=[output_size], + input_size_per_partition=input_size, + params_dtype=torch.bfloat16, + weight_loader=default_weight_loader, + ) + checkpoint_layer.weight.data.copy_(checkpoint_weight) + checkpoint_layer.weight_scale.data.copy_(checkpoint_weight_scale) + checkpoint_scheme.process_weights_after_loading(checkpoint_layer) + + # Online path: feed the *same* raw bf16 source weight and let + # `Mxfp4OnlineLinearMethod` quantize it during + # `process_weights_after_loading`. + online_method = Mxfp4OnlineLinearMethod() + online_method.create_weights( + layer=online_layer, + input_size_per_partition=input_size, + output_partition_sizes=[output_size], + input_size=input_size, + output_size=output_size, + params_dtype=torch.bfloat16, + weight_loader=default_weight_loader, + ) + replace_parameter(online_layer, "weight", weight.clone()) + online_method.process_weights_after_loading(online_layer) + + assert type(checkpoint_scheme.ocp_mx_linear) is type(online_method.kernel) + + assert_quantized_weights_equal( + checkpoint_layer, + online_layer, + keys=("weight", "weight_scale"), + packed_weight_keys=("weight",), + ) diff --git a/vllm/config/quantization.py b/vllm/config/quantization.py index 4f2a40426c..93ea6736b7 100644 --- a/vllm/config/quantization.py +++ b/vllm/config/quantization.py @@ -17,6 +17,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8StaticTensorSym, kInt8StaticChannelSym, kMxfp4Dynamic, + kMxfp4Static, kMxfp8Dynamic, kNvfp4Static, ) @@ -131,6 +132,10 @@ _ONLINE_SHORTHANDS: dict[str, QuantizationConfigArgs] = { linear=QuantSpec(weight=kMxfp8Dynamic), moe=QuantSpec(weight=kMxfp8Dynamic), ), + "mxfp4": QuantizationConfigArgs( + linear=QuantSpec(weight=kMxfp4Static), + moe=QuantSpec(weight=kMxfp4Static), + ), # INT8 weight-only on MoE; linear stays unquantized (no `linear` field). "int8_per_channel_weight_only": QuantizationConfigArgs( moe=QuantSpec(weight=kInt8StaticChannelSym), diff --git a/vllm/model_executor/kernels/linear/mxfp4/emulation.py b/vllm/model_executor/kernels/linear/mxfp4/emulation.py index 66436521b8..ddd0143c22 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/emulation.py +++ b/vllm/model_executor/kernels/linear/mxfp4/emulation.py @@ -23,6 +23,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kMxfp6E3M2Dynamic, ) from vllm.platforms import current_platform +from vllm.utils.import_utils import has_quark from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig @@ -68,6 +69,9 @@ class EmulationMxfp4LinearKernel(MxFp4LinearKernel): ): return False, "only supports MXFP4 or MXFP6 or unquantized activations" + if not has_quark(): + return False, "amd-quark package not available" + if ( current_platform.is_rocm() and current_platform.supports_mx() diff --git a/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py b/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py index 09fc03028c..9107a52ed8 100644 --- a/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py @@ -26,7 +26,11 @@ from vllm.model_executor.layers.quantization.utils.mxfp6_utils import dequant_mx from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import ( OCP_MX_Scheme, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, +) from vllm.platforms import current_platform +from vllm.utils.import_utils import has_quark logger = init_logger(__name__) @@ -92,6 +96,25 @@ class OCP_MXQuantizationEmulationTritonExperts(TritonExperts): ]: self._quant_dtype = current_platform.fp8_dtype() + @staticmethod + def is_supported_config( + cls: type["mk.FusedMoEExperts"], + moe_config: FusedMoEConfig, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + activation_format: mk.FusedMoEActivationFormat, + ) -> tuple[bool, str | None]: + if not has_quark(): + return False, "kernel requires amd-quark package" + + return TritonExperts.is_supported_config( + cls=cls, + moe_config=moe_config, + weight_key=weight_key, + activation_key=activation_key, + activation_format=activation_format, + ) + @property def quant_dtype(self) -> torch.dtype | str | None: return self._quant_dtype diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index cdb9c335d0..3722528c21 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -1650,6 +1650,21 @@ def convert_weight_to_mxfp4_moe_kernel_format( w13_bias, w2_bias, ) + elif mxfp4_backend in ( + Mxfp4MoeBackend.AITER_MXFP4_MXFP4, + Mxfp4MoeBackend.AITER_MXFP4_FP8, + ): + return convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( + mxfp4_backend=mxfp4_backend, + layer=layer, + w13_weight=w13_weight, + w2_weight=w2_weight, + w13_weight_scale=w13_weight_scale, + w2_weight_scale=w2_weight_scale, + w13_bias=w13_bias, + w2_bias=w2_bias, + _cache_permute_indices=_cache_permute_indices, + ) else: raise ValueError( f"Unsupported mxfp4_backend for Mxfp4MoEMethod: {mxfp4_backend}. " diff --git a/vllm/model_executor/layers/quantization/online/base.py b/vllm/model_executor/layers/quantization/online/base.py index 3f8a741d35..c9e7a43e3e 100644 --- a/vllm/model_executor/layers/quantization/online/base.py +++ b/vllm/model_executor/layers/quantization/online/base.py @@ -36,6 +36,10 @@ from vllm.model_executor.layers.quantization.online.fp8 import ( from vllm.model_executor.layers.quantization.online.int8 import ( Int8OnlineMoEMethod, ) +from vllm.model_executor.layers.quantization.online.mxfp4 import ( + Mxfp4OnlineLinearMethod, + Mxfp4OnlineMoEMethod, +) from vllm.model_executor.layers.quantization.online.mxfp8 import ( Mxfp8OnlineLinearMethod, Mxfp8OnlineMoEMethod, @@ -49,6 +53,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8StaticChannelSym, kFp8StaticTensorSym, kInt8StaticChannelSym, + kMxfp4Static, kMxfp8Dynamic, kNvfp4Static, ) @@ -64,6 +69,7 @@ _ONLINE_LINEAR_METHODS: dict[QuantKey, type] = { kFp8Static128BlockSym: Fp8PerBlockOnlineLinearMethod, kFp8StaticChannelSym: Fp8PtpcOnlineLinearMethod, kMxfp8Dynamic: Mxfp8OnlineLinearMethod, + kMxfp4Static: Mxfp4OnlineLinearMethod, } _ONLINE_MOE_METHODS: dict[QuantKey, type] = { @@ -71,6 +77,7 @@ _ONLINE_MOE_METHODS: dict[QuantKey, type] = { kFp8Static128BlockSym: Fp8PerBlockOnlineMoEMethod, kFp8StaticChannelSym: Fp8PtpcOnlineMoEMethod, kMxfp8Dynamic: Mxfp8OnlineMoEMethod, + kMxfp4Static: Mxfp4OnlineMoEMethod, kInt8StaticChannelSym: Int8OnlineMoEMethod, kNvfp4Static: Nvfp4OnlineMoEMethod, } diff --git a/vllm/model_executor/layers/quantization/online/fp8.py b/vllm/model_executor/layers/quantization/online/fp8.py index 38b6267968..a8720b2d23 100644 --- a/vllm/model_executor/layers/quantization/online/fp8.py +++ b/vllm/model_executor/layers/quantization/online/fp8.py @@ -639,36 +639,6 @@ class Fp8PerBlockOnlineMoEMethod(_Fp8OnlineMoEBase): round_up(intermediate_size_per_partition, block_size), ) - def _zero_padding(self, layer: Module) -> None: - hidden_size = layer.moe_config.hidden_dim_unpadded - intermediate_size = layer.moe_config.intermediate_size_per_partition_unpadded - - w13_shard = layer.w13_weight.shape[1] // self.moe.w13_num_shards - if w13_shard > intermediate_size: - for shard in range(self.moe.w13_num_shards): - start = shard * w13_shard + intermediate_size - layer.w13_weight[:, start : (shard + 1) * w13_shard, :] = 0 - if layer.w13_weight.shape[2] > hidden_size: - layer.w13_weight[:, :, hidden_size:] = 0 - - if layer.w2_weight.shape[1] > hidden_size: - layer.w2_weight[:, hidden_size:, :] = 0 - if layer.w2_weight.shape[2] > intermediate_size: - layer.w2_weight[:, :, intermediate_size:] = 0 - - if getattr(layer, "w13_bias", None) is not None: - w13_bias_shard = layer.w13_bias.shape[1] // self.moe.w13_num_shards - if w13_bias_shard > intermediate_size: - for shard in range(self.moe.w13_num_shards): - start = shard * w13_bias_shard + intermediate_size - layer.w13_bias[:, start : (shard + 1) * w13_bias_shard] = 0 - - if ( - getattr(layer, "w2_bias", None) is not None - and layer.w2_bias.shape[1] > hidden_size - ): - layer.w2_bias[:, hidden_size:] = 0 - def process_weights_after_loading(self, layer: Module) -> None: if getattr(layer, "_already_called_process_weights_after_loading", False): return diff --git a/vllm/model_executor/layers/quantization/online/moe_base.py b/vllm/model_executor/layers/quantization/online/moe_base.py index 6aced453a1..8f47768a64 100644 --- a/vllm/model_executor/layers/quantization/online/moe_base.py +++ b/vllm/model_executor/layers/quantization/online/moe_base.py @@ -95,6 +95,36 @@ class OnlineMoEMethodBase(FusedMoEMethodBase): initialize_online_processing(layer) + def _zero_padding(self, layer: torch.nn.Module) -> None: + hidden_size = layer.moe_config.hidden_dim_unpadded + intermediate_size = layer.moe_config.intermediate_size_per_partition_unpadded + + w13_shard = layer.w13_weight.shape[1] // self.moe.w13_num_shards + if w13_shard > intermediate_size: + for shard in range(self.moe.w13_num_shards): + start = shard * w13_shard + intermediate_size + layer.w13_weight[:, start : (shard + 1) * w13_shard, :] = 0 + if layer.w13_weight.shape[2] > hidden_size: + layer.w13_weight[:, :, hidden_size:] = 0 + + if layer.w2_weight.shape[1] > hidden_size: + layer.w2_weight[:, hidden_size:, :] = 0 + if layer.w2_weight.shape[2] > intermediate_size: + layer.w2_weight[:, :, intermediate_size:] = 0 + + if getattr(layer, "w13_bias", None) is not None: + w13_bias_shard = layer.w13_bias.shape[1] // self.moe.w13_num_shards + if w13_bias_shard > intermediate_size: + for shard in range(self.moe.w13_num_shards): + start = shard * w13_bias_shard + intermediate_size + layer.w13_bias[:, start : (shard + 1) * w13_bias_shard] = 0 + + if ( + getattr(layer, "w2_bias", None) is not None + and layer.w2_bias.shape[1] > hidden_size + ): + layer.w2_bias[:, hidden_size:] = 0 + @abstractmethod def process_weights_after_loading(self, layer: torch.nn.Module) -> None: pass diff --git a/vllm/model_executor/layers/quantization/online/mxfp4.py b/vllm/model_executor/layers/quantization/online/mxfp4.py new file mode 100644 index 0000000000..877819cfa1 --- /dev/null +++ b/vllm/model_executor/layers/quantization/online/mxfp4.py @@ -0,0 +1,317 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Online MXFP4 (microscaling FP4, block-32) quantization methods.""" + +from typing import TYPE_CHECKING + +import torch +from torch.nn import Module + +if TYPE_CHECKING: + import vllm.model_executor.layers.fused_moe.modular_kernel as mk + from vllm.model_executor.layers.fused_moe import ( + FusedMoEParallelConfig, + FusedMoEQuantConfig, + RoutedExperts, + ) + +from vllm.model_executor.kernels.linear import init_mxfp4_linear_kernel +from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + TRITON_BACKENDS, + Mxfp4MoeBackend, + convert_weight_to_mxfp4_moe_kernel_format, + make_mxfp4_moe_kernel, + make_mxfp4_moe_quant_config, + mxfp4_round_up_hidden_size_and_intermediate_size, + select_mxfp4_moe_backend, +) +from vllm.model_executor.layers.quantization.online.fp8 import ( + _Fp8OnlineLinearBase, +) +from vllm.model_executor.layers.quantization.online.moe_base import ( + OnlineMoEMethodBase, +) +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( + mxfp4_quantize, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Dynamic +from vllm.model_executor.utils import replace_parameter + +# MXFP4 block size (elements per shared e8m0 scale), same as the checkpoint +# and MoE oracle paths. +MXFP4_BLOCK_SIZE = 32 + + +def _quantize_mxfp4_moe_weight( + weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Batch quantization: bf16/fp16 stacked expert weights -> MXFP4. + + Returns packed FP4 weights (uint8, two values per byte) and per-block + (group-32) e8m0 scales (uint8), one pair per expert. + """ + num_experts = weight.size(0) + first_quant, first_scale = mxfp4_quantize(weight[0]) + # Pre-allocate the output tensors rather than stacking. + # This is important for consistent memory layout. + w_quant = torch.empty( + (num_experts, *first_quant.shape), + dtype=first_quant.dtype, + device=weight.device, + ) + w_scales = torch.empty( + (num_experts, *first_scale.shape), + dtype=first_scale.dtype, + device=weight.device, + ) + w_quant[0] = first_quant + w_scales[0] = first_scale + for i in range(1, num_experts): + w_quant[i], w_scales[i] = mxfp4_quantize(weight[i]) + + return w_quant, w_scales + + +class Mxfp4OnlineLinearMethod(_Fp8OnlineLinearBase): + """Online MXFP4 linear method. + Loads bf16/fp16 checkpoints and quantizes weights to MXFP4 (microscaling + FP4 with block-32 scales) during weight loading. + """ + + activation_quant_key = kMxfp4Dynamic + + def __init__(self): + super().__init__() + self.kernel = init_mxfp4_linear_kernel( + activation_quant_key=self.activation_quant_key + ) + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + if input_size_per_partition % MXFP4_BLOCK_SIZE != 0: + raise ValueError( + f"MXFP4 requires input_size_per_partition " + f"({input_size_per_partition}) to be divisible by " + f"{MXFP4_BLOCK_SIZE}." + ) + + super().create_weights( + layer, + input_size_per_partition, + output_partition_sizes, + input_size, + output_size, + params_dtype, + **extra_weight_attrs, + ) + + def process_weights_after_loading(self, layer: Module) -> None: + if getattr(layer, "_already_called_process_weights_after_loading", False): + return + + weight_fp4, weight_scale = mxfp4_quantize(layer.weight.contiguous()) + + layer.input_scale = None + replace_parameter(layer, "weight", weight_fp4.data) + replace_parameter(layer, "weight_scale", weight_scale.data) + + self.kernel.process_weights_after_loading(layer) + + layer._already_called_process_weights_after_loading = True + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.kernel.apply_weights(layer, x, bias) + + +class Mxfp4OnlineMoEMethod(OnlineMoEMethodBase): + """MoE method for online MXFP4 (block) quantization.""" + + mxfp4_backend: Mxfp4MoeBackend + experts_cls: "type[mk.FusedMoEExperts] | None" + activation_quant_key = kMxfp4Dynamic + + def __init__(self, *, layer: torch.nn.Module): + super().__init__(layer.moe_config) + self.weight_block_size: list[int] = [1, MXFP4_BLOCK_SIZE] + self.weight_scale_name = "weight_scale" + + self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend( + config=self.moe, activation_key=self.activation_quant_key + ) + + def maybe_roundup_sizes( + self, + hidden_size: int, + intermediate_size_per_partition: int, + act_dtype: torch.dtype, + moe_parallel_config: "FusedMoEParallelConfig", + ) -> tuple[int, int]: + hidden_size, intermediate_size_per_partition = super().maybe_roundup_sizes( + hidden_size=hidden_size, + intermediate_size_per_partition=intermediate_size_per_partition, + act_dtype=act_dtype, + moe_parallel_config=moe_parallel_config, + ) + return mxfp4_round_up_hidden_size_and_intermediate_size( + self.mxfp4_backend, hidden_size, intermediate_size_per_partition + ) + + def create_weights( + self, + layer: Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + if ( + hidden_size % MXFP4_BLOCK_SIZE != 0 + or intermediate_size_per_partition % MXFP4_BLOCK_SIZE != 0 + ): + raise ValueError( + "Online MXFP4 MoE requires hidden/intermediate sizes divisible " + f"by {MXFP4_BLOCK_SIZE}." + ) + + super().create_weights( + layer=layer, + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size_per_partition=intermediate_size_per_partition, + params_dtype=params_dtype, + **extra_weight_attrs, + ) + + layer.weight_block_size = [1, MXFP4_BLOCK_SIZE] + + def _setup_kernel( + self, + layer: "RoutedExperts", + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + w13_bias: torch.Tensor | None, + w2_bias: torch.Tensor | None, + ) -> None: + w13, w2, w13_scale, w2_scale, w13_bias, w2_bias = ( + convert_weight_to_mxfp4_moe_kernel_format( + mxfp4_backend=self.mxfp4_backend, + layer=layer, + w13_weight=w13, + w2_weight=w2, + w13_weight_scale=w13_scale, + w2_weight_scale=w2_scale, + w13_bias=w13_bias, + w2_bias=w2_bias, + ) + ) + + # Handle weight/scale assignment based on backend type + # Copied from `quark_moe.py`. + # TODO: This should not be here, replace_parameter should handle + # triton_kernels.tensor.Tensor. + if self.mxfp4_backend in TRITON_BACKENDS or self.mxfp4_backend in ( + Mxfp4MoeBackend.AITER_MXFP4_FP8, + ): + # Triton-based backends: w13/w2 are triton_kernels.tensor.Tensor + # Store on layer for apply(), scales are PrecisionConfig + layer.w13_weight = w13 + layer.w2_weight = w2 + self.w13_precision_config = w13_scale + self.w2_precision_config = w2_scale + else: + # Standard backends: replace parameters + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w2_weight", w2) + replace_parameter(layer, "w13_weight_scale", w13_scale) + replace_parameter(layer, "w2_weight_scale", w2_scale) + + if w13_bias is not None and w2_bias is not None: + replace_parameter(layer, "w13_bias", w13_bias) + replace_parameter(layer, "w2_bias", w2_bias) + + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + if self.moe_quant_config is not None and self.experts_cls is not None: + self.moe_kernel = make_mxfp4_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + mxfp4_backend=self.mxfp4_backend, + experts_cls=self.experts_cls, + routing_tables=layer._expert_routing_tables(), + layer=layer, + ) + + def get_fused_moe_quant_config( + self, layer: torch.nn.Module + ) -> "FusedMoEQuantConfig | None": + # NOTE: unlike `QuarkOCP_MX_MoEMethod.get_fused_moe_quant_config`, this + # does not branch on `self.mxfp4_backend in TRITON_BACKENDS or + # AITER_MXFP4_FP8` to read `self.w13_precision_config`/ + # `self.w2_precision_config` instead of `layer.w13_weight_scale`/ + # `layer.w2_weight_scale`. For those backends, + # `convert_weight_to_mxfp4_moe_kernel_format` deletes the latter, so + # `getattr` below would raise `AttributeError`. This is currently + # unreachable because `activation_quant_key` is hardcoded to + # `kMxfp4Dynamic` + # TODO: When supporting online activation quant key override, + # fix this convert_weight_to_mxfp4_moe_kernel_format issue. + w1_scale = getattr(layer, f"w13_{self.weight_scale_name}") + w2_scale = getattr(layer, f"w2_{self.weight_scale_name}") + + return make_mxfp4_moe_quant_config( + mxfp4_backend=self.mxfp4_backend, + w1_scale=w1_scale, + w2_scale=w2_scale, + w1_bias=getattr(layer, "w13_bias", None), + w2_bias=getattr(layer, "w2_bias", None), + a1_scale=layer.w13_input_scale, + a2_scale=layer.w2_input_scale, + swiglu_limit=getattr(layer, "swiglu_limit", None), + gemm1_alpha=getattr(layer, "swiglu_alpha", None), + gemm1_beta=getattr(layer, "swiglu_beta", None), + layer=layer, + ) + + def process_weights_after_loading(self, layer: Module) -> None: + if getattr(layer, "_already_called_process_weights_after_loading", False): + return + + self._zero_padding(layer) + + if self.mxfp4_backend == Mxfp4MoeBackend.NONE: + layer._already_called_process_weights_after_loading = True + return + + layer.w13_input_scale = None + layer.w2_input_scale = None + + w13, w13_scale = _quantize_mxfp4_moe_weight(layer.w13_weight) + w2, w2_scale = _quantize_mxfp4_moe_weight(layer.w2_weight) + + self._setup_kernel( + layer, + w13, + w2, + w13_scale, + w2_scale, + getattr(layer, "w13_bias", None), + getattr(layer, "w2_bias", None), + ) + + layer._already_called_process_weights_after_loading = True diff --git a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py index 0d3bf88cbf..65348a8222 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py @@ -6,6 +6,7 @@ import torch from vllm.logger import init_logger from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton from vllm.utils.import_utils import has_triton_kernels from vllm.utils.torch_utils import direct_register_custom_op, is_torch_equal_or_newer @@ -184,3 +185,319 @@ except AttributeError as error: def xpu_mxfp4_quantize(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: return torch.ops.vllm.xpu_mxfp4_quantize(x) + + +def permute_shape(shape: tuple[int, ...], permute_order: list[int]) -> tuple[int, ...]: + return tuple(shape[i] for i in permute_order) + + +@triton.jit +def _compute_quant_and_scale(src_tensor, valid_src_mask): + BLOCK_SIZE_OUT_DIM: tl.constexpr = src_tensor.shape[0] + BLOCK_SIZE_QUANT_DIM: tl.constexpr = src_tensor.shape[1] + BLOCK_SIZE_QUANT_MX_SCALE: tl.constexpr = src_tensor.shape[1] // 32 + + # Explicit cast to fp32 since most ops are not supported on bfloat16. We + # avoid needless conversions to and from bf16 + f32_tensor = src_tensor.to(tl.float32) + abs_tensor = tl.abs(f32_tensor) + abs_tensor = tl.where( + valid_src_mask, abs_tensor, -1.0 + ) # Don't consider padding tensors in scale computation + abs_tensor = tl.reshape( + abs_tensor, [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 32] + ) + max_val = tl.max(abs_tensor, axis=2, keep_dims=True) + # Round-to-nearest-even block scale (OCP MX spec), matching hardware/aiter + # MXFP4 quantizers: compute 2 ** (floor(log2(rounding(max_abs(v)))) - max_exp). + max_val = max_val.to(tl.int32, bitcast=True) + max_val = (max_val + 0x200000).to(tl.uint32, bitcast=True) & 0x7F800000 + max_val = max_val.to(tl.float32, bitcast=True) + # Add epsilon to prevent log2(0) = -inf when max_val is zero (all-zero input block) + eps = tl.where(max_val == 0.0, 2 ** (-126), 0.0) + scale_e8m0_unbiased = tl.log2(max_val + eps).floor() - 2 + scale_e8m0_unbiased = tl.clamp(scale_e8m0_unbiased, min=-127, max=127) + dequant_scale_rounded = tl.exp2(scale_e8m0_unbiased) + dequant_scale_exponent = dequant_scale_rounded.to(tl.uint32, bitcast=True) + + dequant_scale_rounded = dequant_scale_exponent.to(tl.float32, bitcast=True) + quant_scale = tl.where(dequant_scale_rounded == 0, 0, 1.0 / dequant_scale_rounded) + + f32_tensor = tl.reshape( + f32_tensor, [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 32] + ) + quant_tensor = f32_tensor * quant_scale + + # Reshape the tensors after scaling + quant_tensor = quant_tensor.reshape([BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_DIM]) + # Set the invalid portions of the tensor to 0. This will ensure that any + # padding tensors are 0 in the mx format. + quant_tensor = tl.where(valid_src_mask, quant_tensor, 0) + dequant_scale_exponent = dequant_scale_exponent.reshape( + [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE] + ) + + # First, we simply extract the exponent part of the scales and store the result + dequant_scale_exponent = (dequant_scale_exponent >> 23).to(tl.uint8) + # Now we must convert the tensor to packed e2m1 (mxfp4). + quant_tensor = quant_tensor.to(tl.uint32, bitcast=True) + signs = quant_tensor & 0x80000000 + exponents = (quant_tensor >> 23) & 0xFF + mantissas = quant_tensor & 0x7FFFFF + + # 0.25 <= x < 0.75 maps to 0.5, a denormal number + E8_BIAS = 127 + E2_BIAS = 1 + # Move implicit bit 1 at the beginning to mantissa for denormals. Bits + # shifted out are OR-reduced into the LSB as a sticky bit, so the + # round-to-even logic below can tell an exact tie from a value that is + # merely close to one. + adjusted_exponents = tl.core.sub(E8_BIAS, exponents + 1, sanitize_overflow=False) + folded_mantissas = 0x400000 | (mantissas >> 1) + denormal_mantissas = folded_mantissas >> adjusted_exponents + denormal_mantissas |= ( + (denormal_mantissas << adjusted_exponents) != folded_mantissas + ).to(tl.uint32) + mantissas = tl.where(exponents < E8_BIAS, denormal_mantissas, mantissas) + + # For normal numbers, we change the bias from 127 to 1, and for subnormals, + # we keep exponent as 0. + exponents = tl.maximum(exponents, E8_BIAS - E2_BIAS) - (E8_BIAS - E2_BIAS) + + # Combine sign, exponent, and mantissa, rounding to nearest with ties + # broken to even (OCP MX spec, matching hardware/aiter MXFP4 + # quantizers) via the standard "add a rounding bias, then truncate" + # trick: the bias's LSB is the kept mantissa bit, so an exact tie + # rounds to whichever candidate is even, while any lower (sticky) bit + # being set always rounds up. + rounding_bias = ((mantissas >> 22) & 1) + 0x1FFFFF + rounded_mantissa = (mantissas + rounding_bias) >> 22 # in {0, 1, 2} + e2m1_tmp = tl.minimum((exponents << 1) + rounded_mantissa, 0x7) + e2m1_value = ((signs >> 28) | e2m1_tmp).to(tl.uint8) + + e2m1_value = tl.reshape( + e2m1_value, [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_DIM // 2, 2] + ) + evens, odds = tl.split(e2m1_value) + out_tensor = evens | (odds << 4) + + return out_tensor, dequant_scale_exponent + + +@triton.jit +def _downcast_to_mxfp( + mx_tensor_ptr, + stride_mxt_outer, + stride_mxt_quant: tl.constexpr, + mx_scale_ptr, + stride_mx_scale_outer, + stride_mx_scale_quant, + src_ptr, + stride_src_outer, + stride_src_quant, + outer_dim, + quant_dim, + BLOCK_SIZE_OUT_DIM: tl.constexpr, + BLOCK_SIZE_QUANT_DIM: tl.constexpr, +): + tl.static_assert( + stride_mxt_quant == 1, f"Output stride, {stride_mxt_quant=} must be 1." + ) + tl.static_assert( + BLOCK_SIZE_QUANT_DIM % 32 == 0, + f"{BLOCK_SIZE_QUANT_DIM=} must be a multiple of 32", + ) + + # uint8 signifies two fp4 e2m1 values packed into a single byte + tl.static_assert( + mx_tensor_ptr.dtype.element_ty == tl.uint8, + f"Invalid {mx_tensor_ptr.dtype.element_ty=}. Must be uint8.", + ) + + src_dtype: tl.constexpr = src_ptr.dtype.element_ty + tl.static_assert( + mx_scale_ptr.dtype.element_ty == tl.uint8, + f"{mx_scale_ptr.dtype.element_ty=} must be uint8", + ) + tl.static_assert( + (src_dtype == tl.bfloat16) or (src_dtype == tl.float16), + f"{src_dtype=} must be bfloat16 or float16", + ) + + outer_block = tl.program_id(0).to(tl.int64) + quant_block = tl.program_id(1).to(tl.int64) + + BLOCK_SIZE_QUANT_MX_SCALE: tl.constexpr = BLOCK_SIZE_QUANT_DIM // 32 + BLOCK_SIZE_QUANT_MX_TENSOR: tl.constexpr = BLOCK_SIZE_QUANT_DIM // 2 + + start_src_quant = quant_block * BLOCK_SIZE_QUANT_DIM + start_mx_scale_quant = quant_block * BLOCK_SIZE_QUANT_MX_SCALE + start_mx_quant = quant_block * BLOCK_SIZE_QUANT_MX_TENSOR + start_out = outer_block * BLOCK_SIZE_OUT_DIM + + src_ptr += start_src_quant * stride_src_quant + start_out * stride_src_outer + mx_scale_ptr += ( + start_mx_scale_quant * stride_mx_scale_quant + start_out * stride_mx_scale_outer + ) + mx_tensor_ptr += start_mx_quant * stride_mxt_quant + start_out * stride_mxt_outer + + offs_src_quant = tl.arange(0, BLOCK_SIZE_QUANT_DIM)[None, :].to(tl.int64) + offs_mxt_quant = tl.arange(0, BLOCK_SIZE_QUANT_MX_TENSOR)[None, :].to(tl.int64) + offs_scale_quant = tl.arange(0, BLOCK_SIZE_QUANT_MX_SCALE)[None, :].to(tl.int64) + offs_outer = tl.arange(0, BLOCK_SIZE_OUT_DIM)[:, None].to(tl.int64) + + mask_src_quant = start_src_quant + offs_src_quant < quant_dim + mask_n = start_out + offs_outer < outer_dim + full_mask_src = mask_src_quant and mask_n + + mask_mxt_quant = start_mx_quant + offs_mxt_quant < tl.cdiv(quant_dim, 2) + full_mask_mxt = mask_mxt_quant and mask_n + + scale_mask_k = start_mx_scale_quant + offs_scale_quant < tl.cdiv(quant_dim, 32) + full_scale_mask = scale_mask_k and mask_n + + src_tensor_offsets = ( + offs_src_quant * stride_src_quant + offs_outer * stride_src_outer + ) + mx_scale_offsets = ( + offs_scale_quant * stride_mx_scale_quant + offs_outer * stride_mx_scale_outer + ) + mx_tensor_offsets = ( + offs_mxt_quant * stride_mxt_quant + offs_outer * stride_mxt_outer + ) + src_tensor = tl.load(src_ptr + src_tensor_offsets, mask=full_mask_src) + + out_tensor, scale_tensor = _compute_quant_and_scale(src_tensor, full_mask_src) + + tl.store(mx_scale_ptr + mx_scale_offsets, scale_tensor, mask=full_scale_mask) + tl.store(mx_tensor_ptr + mx_tensor_offsets, out_tensor, mask=full_mask_mxt) + + +def downcast_to_mxfp( + src_tensor: torch.Tensor, + axis: int, + out_quant_tensor: torch.Tensor | None = None, + out_scale: torch.Tensor | None = None, + BLOCK_OUT_DIM: int = 128, + BLOCK_QUANT_DIM: int = 32, +): + """ + Convert the src weights to MXFP4. The src weight is quantized along the + axis dimension into packed e2m1 values (torch.uint8, two values per + byte), so the size of that dimension in the output is half of the + logical (unpacked) size. + """ + out_quant_type = torch.uint8 + ndim = src_tensor.ndim + assert -ndim <= axis < ndim, f"Invalid axis {axis=}" + axis = axis if axis >= 0 else axis + ndim + + L = src_tensor.shape[axis] + # We make this assertion since we can't track if the "real" shape was odd, + # and we padded it to be even. + # We want to maintain the property dequant(quant(x)).shape == x.shape + assert L % 2 == 0, f"axis dim must be divisible by 2 for e2m1. Got {L}" + + device = src_tensor.device + + packed_quant_dim = triton.cdiv(L, 2) + out_scale_dim = triton.cdiv(L, 32) + + # Move the quantization axis to the end for the kernel, then permute back. + permute_order = list(range(ndim)) + permute_order[axis], permute_order[-1] = permute_order[-1], permute_order[axis] + + prmted_quant_tensor_shape = permute_shape(src_tensor.shape, permute_order)[:-1] + ( + packed_quant_dim, + ) + prmted_scale_shape = permute_shape(src_tensor.shape, permute_order)[:-1] + ( + out_scale_dim, + ) + prmted_src_tensor = src_tensor.permute(permute_order) + + if out_quant_tensor is None: + out_quant_tensor = torch.empty( + prmted_quant_tensor_shape, dtype=out_quant_type, device=device + ) + else: + expected_shape = ( + src_tensor.shape[:axis] + (packed_quant_dim,) + src_tensor.shape[axis + 1 :] + ) + assert out_quant_tensor.shape == expected_shape, ( + f"{out_quant_tensor.shape=} != {expected_shape=}" + ) + assert out_quant_tensor.dtype == out_quant_type, ( + f"{out_quant_tensor.dtype=} != {out_quant_type=}" + ) + assert out_quant_tensor.stride(axis) == 1, ( + f"{out_quant_tensor.stride(axis)=} != 1" + ) + # We expect the axis dimension to be last, so permute the tensor + out_quant_tensor = out_quant_tensor.permute(permute_order) + + if out_scale is None: + out_scale = torch.empty(prmted_scale_shape, dtype=torch.uint8, device=device) + else: + expected_scale_shape = permute_shape(prmted_scale_shape, permute_order) + assert out_scale.shape == expected_scale_shape, ( + f"{out_scale.shape=} {expected_scale_shape=}" + ) + assert out_scale.dtype == torch.uint8, f"{out_scale.dtype=} != torch.uint8" + out_scale = out_scale.permute(permute_order) + + # Flatten input tensor for kernel. This will typically make a copy + reshaped_src_tensor = prmted_src_tensor.reshape(-1, L) + blocks_quant_dim = triton.cdiv(reshaped_src_tensor.shape[-1], BLOCK_QUANT_DIM) + blocks_out_dim = triton.cdiv(reshaped_src_tensor.shape[0], BLOCK_OUT_DIM) + + # Flatten the output tensors for the kernel, this should be a view always + kernel_quant_tensor = out_quant_tensor.reshape(-1, packed_quant_dim) + kernel_scale = out_scale.reshape(-1, out_scale_dim) + assert kernel_quant_tensor.data_ptr() == out_quant_tensor.data_ptr() + assert kernel_scale.data_ptr() == out_scale.data_ptr() + + _downcast_to_mxfp[(blocks_out_dim, blocks_quant_dim)]( + kernel_quant_tensor, + *kernel_quant_tensor.stride(), + kernel_scale, + *kernel_scale.stride(), + reshaped_src_tensor, + *reshaped_src_tensor.stride(), + *reshaped_src_tensor.shape, + BLOCK_OUT_DIM, + BLOCK_QUANT_DIM, + num_warps=8, + ) + + out_quant_tensor = out_quant_tensor.permute(permute_order) + out_scale = out_scale.permute(permute_order).contiguous() + return out_quant_tensor, out_scale, permute_shape(prmted_scale_shape, permute_order) + + +def mxfp4_quantize(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize a bf16/fp16 tensor to MXFP4 along its last dimension. + + Dispatches to the fastest backend available on the current platform: + + - the native XPU custom op on XPU, + - aiter on ROCm when `aiter` is installed, + - and the portable Triton kernel (``downcast_to_mxfp``) otherwise. + + Returns packed FP4 values (uint8, two values per byte) and per-block + (group-32) e8m0 scales (uint8). + """ + if current_platform.is_xpu(): + return xpu_mxfp4_quantize(x) + + from vllm._aiter_ops import is_aiter_found_and_supported + + if is_aiter_found_and_supported() and x.dtype == torch.bfloat16: + from vllm.model_executor.layers.quantization.quark.utils import ( + quark_quantize_weight_to_mxfp4, + ) + + return quark_quantize_weight_to_mxfp4(x) + + quant_tensor, scale, _ = downcast_to_mxfp(x, axis=-1) + + return quant_tensor, scale diff --git a/vllm/utils/import_utils.py b/vllm/utils/import_utils.py index 3df83148ce..9f6ea0fec4 100644 --- a/vllm/utils/import_utils.py +++ b/vllm/utils/import_utils.py @@ -565,6 +565,11 @@ def has_humming() -> bool: return _has_module("humming") +def has_quark(): + """Whether the optional `quark` package is available.""" + return _has_module("quark") + + def check_torchcodec_available(): """Whether the optional `torchcodec` package is available.""" try: