[Online quantization] Add online MXFP4 quantization support (#49347)
Signed-off-by: Felix Marty <Felix.Marty@amd.com>
This commit is contained in:
fxmarty-amd
@@ -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. <br> - 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
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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",),
|
||||
)
|
||||
@@ -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),
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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}. "
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user