diff --git a/docs/user_guide/quantization/mxfp4.md b/docs/user_guide/quantization/mxfp4.md index 401a55ad4..b18b3e7b3 100644 --- a/docs/user_guide/quantization/mxfp4.md +++ b/docs/user_guide/quantization/mxfp4.md @@ -39,7 +39,7 @@ vLLM-Omni provides two quantization methods with different scale structures: | NVIDIA Blackwell GPU (SM 100+) | ⭕ | | NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ | | NVIDIA Ampere GPU (SM 80+) | ⭕ | -| AMD ROCm | ⭕ | +| AMD ROCm (gfx950 / MI355X) | ✅ | | Intel XPU | ⭕ | | Ascend NPU (Atlas 950 A5) | ✅ | diff --git a/docs/user_guide/quantization/overview.md b/docs/user_guide/quantization/overview.md index 38d027f12..3a4a45573 100644 --- a/docs/user_guide/quantization/overview.md +++ b/docs/user_guide/quantization/overview.md @@ -20,7 +20,7 @@ type has a different quantization scope. | NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ | | NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ | | NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ | ⭕ | ✅ | ✅ | ❌ | -| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ❌ | +| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ⭕ | ⭕ | ❌ | | Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ❌ | | Ascend NPU | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | diff --git a/tests/diffusion/quantization/test_mxfp4_config.py b/tests/diffusion/quantization/test_mxfp4_config.py index af3a9c9c4..240ee2509 100644 --- a/tests/diffusion/quantization/test_mxfp4_config.py +++ b/tests/diffusion/quantization/test_mxfp4_config.py @@ -537,3 +537,118 @@ def test_single_scale_row_parallel_tp2_loader_simulation(): for rank in range(_TP2): assert _shard(ckpt_weight_scale, layer.weight_scale, rank, _TP2, "input_dim").shape == layer.weight_scale.shape + + +# --------------------------------------------------------------------------- +# ROCm MXFP4 (gfx950) — get_quant_method dispatch +# +# These run on CPU: the platform check, gcnArchName probe and aiter op +# registration are all mocked (stdlib unittest.mock + monkeypatch, so no +# pytest-mock dependency). They cover only the dispatch + weight-allocation +# logic added by the ROCm PR — the AITER GEMM / quant kernels require real +# gfx950 hardware and are intentionally NOT exercised here. +# --------------------------------------------------------------------------- + + +@pytest.fixture +def _rocm_platform(monkeypatch: pytest.MonkeyPatch): + """Make current_omni_platform report ROCm, and stub the aiter + custom-op registration so ROCmMxfp4*Method can be constructed without aiter.""" + from vllm_omni.platforms import current_omni_platform + from vllm_omni.quantization import mxfp4_config + + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: False) + monkeypatch.setattr(current_omni_platform, "is_rocm", lambda: True) + monkeypatch.setattr(mxfp4_config, "_register_rocm_mxfp4_op", lambda: None) + + +def _patch_gcn_arch(monkeypatch: pytest.MonkeyPatch, arch: str) -> None: + """Patch torch.cuda.get_device_properties(...).gcnArchName to return `arch`.""" + from types import SimpleNamespace + + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0) + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda *a, **k: SimpleNamespace(gcnArchName=arch), + ) + + +def _fake_linear_layer(): + """A stand-in that passes isinstance(layer, LinearBase) without a real layer.""" + from unittest.mock import MagicMock + + from vllm.model_executor.layers.linear import LinearBase + + return MagicMock(spec=LinearBase) + + +def test_rocm_online_dispatch_returns_rocm_method(_rocm_platform, monkeypatch): + """ROCm + gfx950 + online checkpoint must return ROCmMxfp4OnlineLinearMethod.""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, ROCmMxfp4OnlineLinearMethod + + _patch_gcn_arch(monkeypatch, "gfx950:sramecc+:xnack-") + cfg = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=False) + + method = cfg.get_quant_method(_fake_linear_layer(), "blocks.0.attn1.to_q") + assert isinstance(method, ROCmMxfp4OnlineLinearMethod) + + +def test_rocm_ignored_layer_returns_unquantized(_rocm_platform, monkeypatch): + """A prefix in ignored_layers must return UnquantizedLinearMethod before the gfx950 probe.""" + from vllm.model_executor.layers.linear import UnquantizedLinearMethod + + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + _patch_gcn_arch(monkeypatch, "gfx950:sramecc+:xnack-") + cfg = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=False, ignored_layers=["proj_out"]) + + assert isinstance(cfg.get_quant_method(_fake_linear_layer(), "proj_out"), UnquantizedLinearMethod) + + +def test_rocm_non_gfx950_raises(_rocm_platform, monkeypatch): + """MXFP4 on ROCm requires gfx950; any other arch must raise.""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + _patch_gcn_arch(monkeypatch, "gfx942:sramecc+:xnack-") + cfg = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=False) + + with pytest.raises(NotImplementedError, match="gfx950"): + cfg.get_quant_method(_fake_linear_layer(), "blocks.0.attn1.to_q") + + +# --------------------------------------------------------------------------- +# ROCm MXFP4 — create_weights (online lazy meta-device placeholder) +# +# The base ROCmMxfp4LinearMethod is abstract (no create_weights); the online +# subclass gets create_weights from _LazyWeightMixin, which registers a BF16 +# weight on the meta device to be materialised at load time. +# --------------------------------------------------------------------------- + + +def test_rocm_create_weights_column_parallel_tp2(_rocm_platform): + """Column-parallel TP=2: meta BF16 weight has the output halved, input full.""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, ROCmMxfp4OnlineLinearMethod + + method = ROCmMxfp4OnlineLinearMethod(DiffusionMXFP4Config()) + layer = _create_weights(method, input_size_per_partition=_TP2_K, output_partition_sizes=[_TP2_N // _TP2]) + + assert layer.weight.shape == (_TP2_N // _TP2, _TP2_K) + assert layer.weight.dtype == torch.bfloat16 + assert layer.weight.device.type == "meta" + assert layer.weight.input_dim == 1 + assert layer.weight.output_dim == 0 + assert layer.logical_widths == [_TP2_N // _TP2] + assert layer.weight_block_size is None + + +def test_rocm_create_weights_row_parallel_tp2(_rocm_platform): + """Row-parallel TP=2: meta BF16 weight has the input halved, output full.""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, ROCmMxfp4OnlineLinearMethod + + method = ROCmMxfp4OnlineLinearMethod(DiffusionMXFP4Config()) + layer = _create_weights(method, input_size_per_partition=_TP2_K // _TP2, output_partition_sizes=[_TP2_N]) + + assert layer.weight.shape == (_TP2_N, _TP2_K // _TP2) + assert layer.input_size_per_partition == _TP2_K // _TP2 + assert layer.output_size_per_partition == _TP2_N diff --git a/vllm_omni/quantization/mxfp4_config.py b/vllm_omni/quantization/mxfp4_config.py index 6a02d5edb..a6b29ea79 100644 --- a/vllm_omni/quantization/mxfp4_config.py +++ b/vllm_omni/quantization/mxfp4_config.py @@ -7,6 +7,8 @@ Architecture mirrors mxfp8_config.py: MXFPLinearMethodBase – platform-agnostic skeleton (imported from mxfp8_config) NPUMxfp4LinearMethod – NPU single-scale offline (W4A4 MXFP4) NPUMxfp4OnlineLinearMethod – NPU single-scale online (BF16 → FP4) + ROCmMxfp4LinearMethod – ROCm base class (AITER quant + shuffle; online-only) + ROCmMxfp4OnlineLinearMethod – ROCm online (BF16 → FP4 via AITER) NPUMxfp4DualScaleLinearMethod – NPU dual-scale offline (W4A4 MXFP4 DualScale) NPUMxfp4DualScaleOnlineLinearMethod – NPU dual-scale online (BF16 → FP4) @@ -153,9 +155,16 @@ class DiffusionMXFP4Config(QuantizationConfig): if self.is_checkpoint_mxfp4_serialized: return NPUMxfp4LinearMethod(self) return NPUMxfp4OnlineLinearMethod(self) + if current_omni_platform.is_rocm(): + gcn_arch = torch.cuda.get_device_properties(torch.accelerator.current_device_index()).gcnArchName + if "gfx950" not in gcn_arch: + raise NotImplementedError(f"MXFP4 on ROCm requires gfx950 (MI355X). Detected: {gcn_arch}") + if self.is_checkpoint_mxfp4_serialized: + raise NotImplementedError("Pre-quantized MXFP4 checkpoints are not yet supported on ROCm.") + return ROCmMxfp4OnlineLinearMethod(self) raise NotImplementedError( "DiffusionMXFP4Config (W4A4 MXFP4) is currently only supported " - "on NPU (Ascend) platforms. CUDA support is not yet implemented." + "on NPU (Ascend) and ROCm (AMD, gfx950) platforms." ) return None @@ -332,6 +341,151 @@ class NPUMxfp4OnlineLinearMethod(_LazyWeightMixin, NPUMxfp4LinearMethod): layer._already_called_process_weights_after_loading = True +# --------------------------------------------------------------------------- +# ROCm MXFP4 base method (AITER) — used as base class for the online subclass. +# Online-only for now; offline (pre-quantized checkpoint) not yet supported. +# --------------------------------------------------------------------------- + + +def _register_rocm_mxfp4_op() -> None: + """Register the vllm_omni::rocm_mxfp4_gemm custom op for torch.compile. + + Wraps activation quantization + GEMM in a single opaque op so that + torch.compile/inductor doesn't try to trace through AITER internals. + """ + import aiter + + @torch.library.custom_op("vllm_omni::rocm_mxfp4_gemm", mutates_args=()) + def _rocm_mxfp4_gemm( + a: torch.Tensor, + w_quant: torch.Tensor, + w_scale: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + quant_func = aiter.get_hip_quant(aiter.QuantType.per_1x32) + a_quant, a_scale = quant_func(a, shuffle=True) + return aiter.gemm_a4w4( + a_quant, + w_quant, + a_scale, + w_scale, + bpreshuffle=True, + bias=bias, + ) + + @_rocm_mxfp4_gemm.register_fake + def _rocm_mxfp4_gemm_fake( + a: torch.Tensor, + w_quant: torch.Tensor, + w_scale: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + M, _ = a.shape + N, _ = w_quant.shape + return torch.empty(M, N, dtype=a.dtype, device=a.device) + + +class ROCmMxfp4LinearMethod(MXFPLinearMethodBase): + """ROCm W4A4 MXFP4 linear method using AITER (gemm_a4w4). + + Weight buffers after process_weights_after_loading: + weight_shuffle : FP4 quantized + shuffled via shuffle_weight(layout=(16,16)) + weight_scale : per-group-of-32 scales from AITER per_1x32 + + Forward path: + _quantize_activation: pass-through (activation quant is inside the custom op) + _quant_matmul: torch.ops.vllm_omni.rocm_mxfp4_gemm wraps activation quant + then gemm_a4w4 in one custom op + """ + + def __init__(self, quant_config: DiffusionMXFP4Config) -> None: + self.quant_config = quant_config + self.out_dtype = torch.get_default_dtype() + if not hasattr(torch.ops.vllm_omni, "rocm_mxfp4_gemm"): + _register_rocm_mxfp4_op() + + def process_weights_after_loading(self, layer: Module) -> None: + if getattr(layer, "_already_called_process_weights_after_loading", False): + return + + import aiter + from aiter.ops.shuffle import shuffle_weight + + quant_func = aiter.get_hip_quant(aiter.QuantType.per_1x32) + weight_quant, weight_scale = quant_func(layer.weight.data, shuffle=True) + weight_shuffled = shuffle_weight(weight_quant, layout=(16, 16)) + + # Store quantized tensors as non-parameter buffers; delete original weight. + layer.register_buffer("weight_shuffle", weight_shuffled, persistent=True) + layer.register_buffer("weight_scale", weight_scale, persistent=True) + + # Remove the original weight parameter to free memory. + if hasattr(layer, "weight") and isinstance(layer.weight, torch.nn.Parameter): + delattr(layer, "weight") + layer.register_parameter("weight", None) + + layer._already_called_process_weights_after_loading = True + + def _quantize_activation(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + # Activation quantization is inside the custom op called by _quant_matmul, + # so pass through the raw activation here. + return x, None + + def _quant_matmul( + self, + x_q: torch.Tensor, + x_scale: torch.Tensor, + layer: torch.nn.Module, + bias: torch.Tensor | None, + ori_dtype: torch.dtype, + ) -> torch.Tensor: + output = torch.ops.vllm_omni.rocm_mxfp4_gemm( + x_q, + layer.weight_shuffle, + layer.weight_scale, + bias, + ) + if output.dtype != ori_dtype: + output = output.to(ori_dtype) + return output + + +# --------------------------------------------------------------------------- +# ROCm MXFP4 online method (BF16 checkpoint → quantize at load time) +# --------------------------------------------------------------------------- + + +class ROCmMxfp4OnlineLinearMethod(_LazyWeightMixin, ROCmMxfp4LinearMethod): + """ROCm W4A4 MXFP4 online linear method using AITER. + + MRO: ROCmMxfp4OnlineLinearMethod → _LazyWeightMixin → ROCmMxfp4LinearMethod + → MXFPLinearMethodBase → LinearMethodBase + + create_weights : _LazyWeightMixin (meta device + patched loader) + process_weights : ROCmMxfp4OnlineLinearMethod (meta → materialize, then AITER quant + shuffle) + apply / ops : ROCmMxfp4LinearMethod / MXFPLinearMethodBase (shared) + """ + + def process_weights_after_loading(self, layer: Module) -> None: + if getattr(layer, "_already_called_process_weights_after_loading", False): + return + + # Materialise from meta device if needed (same pattern as NPU online). + if layer.weight is not None and layer.weight.device == torch.device("meta"): + weight = ModelWeightParameter( + data=torch.empty_like(layer.weight, device=layer._load_device), + input_dim=1, + output_dim=0, + weight_loader=layer.weight.weight_loader, + ) + _copy_missing_attrs(layer.weight, weight) + layer.register_parameter("weight", weight) + initialize_single_dummy_weight(layer.weight) + + # Delegate to the base class which does AITER quant + shuffle. + ROCmMxfp4LinearMethod.process_weights_after_loading(self, layer) + + # --------------------------------------------------------------------------- # NPU MXFP4 dual-scale offline method (W4A4_MXFP4_DUALSCALE checkpoint) # ---------------------------------------------------------------------------