diff --git a/tests/kernels/core/test_fused_q_kv_rmsnorm.py b/tests/kernels/core/test_fused_q_kv_rmsnorm.py index 3d0e2ac11e35..485a35348eaf 100644 --- a/tests/kernels/core/test_fused_q_kv_rmsnorm.py +++ b/tests/kernels/core/test_fused_q_kv_rmsnorm.py @@ -223,7 +223,6 @@ def test_shared_query_quant_preserves_projection( with set_current_vllm_config( VllmConfig(compilation_config=CompilationConfig(mode=CompilationMode.NONE)) ): - kernel = cls(Mxfp8LinearLayerConfig()) # Match the checkpoint's TP4 attention and replicated indexer widths. if projection == "attention": linear = ColumnParallelLinear.__new__(ColumnParallelLinear) @@ -247,6 +246,7 @@ def test_shared_query_quant_preserves_projection( ) linear.weight = torch.nn.Parameter(w, requires_grad=False) linear.weight_scale = torch.nn.Parameter(ws, requires_grad=False) + kernel = cls(Mxfp8LinearLayerConfig(weight_shape=w.shape)) quant_config = DeepseekV4FP8Config.from_config( { "quant_method": "fp8", diff --git a/tests/kernels/quantization/test_flashinfer_mxfp8_trtllm.py b/tests/kernels/quantization/test_flashinfer_mxfp8_trtllm.py index 315e35b782a0..1b6a3730715f 100644 --- a/tests/kernels/quantization/test_flashinfer_mxfp8_trtllm.py +++ b/tests/kernels/quantization/test_flashinfer_mxfp8_trtllm.py @@ -50,7 +50,9 @@ def test_flashinfer_trtllm_mxfp8_linear_numerics( x = torch.randn((m, k), dtype=torch.bfloat16, device="cuda") weight = torch.randn((n, k), dtype=torch.bfloat16, device="cuda") layer = _make_layer(weight) - kernel = FlashInferTrtllmMxfp8LinearKernel(Mxfp8LinearLayerConfig()) + kernel = FlashInferTrtllmMxfp8LinearKernel( + Mxfp8LinearLayerConfig(weight_shape=(n, k)) + ) kernel.process_weights_after_loading(layer) output = kernel.apply_weights(layer, x) @@ -69,7 +71,9 @@ def test_flashinfer_trtllm_mxfp8_custom_ops() -> None: x = torch.randn((7, 512), dtype=torch.bfloat16, device="cuda") weight = torch.randn((256, 512), dtype=torch.bfloat16, device="cuda") layer = _make_layer(weight) - kernel = FlashInferTrtllmMxfp8LinearKernel(Mxfp8LinearLayerConfig()) + kernel = FlashInferTrtllmMxfp8LinearKernel( + Mxfp8LinearLayerConfig(weight_shape=(256, 512)) + ) kernel.process_weights_after_loading(layer) torch.library.opcheck( @@ -104,7 +108,9 @@ def test_flashinfer_trtllm_mxfp8_linear_cuda_graph() -> None: m, n, k = 7, 130, 512 weight = torch.randn((n, k), dtype=torch.bfloat16, device="cuda") layer = _make_layer(weight) - kernel = FlashInferTrtllmMxfp8LinearKernel(Mxfp8LinearLayerConfig()) + kernel = FlashInferTrtllmMxfp8LinearKernel( + Mxfp8LinearLayerConfig(weight_shape=(n, k)) + ) kernel.process_weights_after_loading(layer) static_x = torch.randn((m, k), dtype=torch.bfloat16, device="cuda") diff --git a/tests/kernels/quantization/test_marlin_tile_padding.py b/tests/kernels/quantization/test_marlin_tile_padding.py index d224b0a2b257..8b55910a786e 100644 --- a/tests/kernels/quantization/test_marlin_tile_padding.py +++ b/tests/kernels/quantization/test_marlin_tile_padding.py @@ -408,7 +408,9 @@ def test_fp8_block_marlin_padded_round_trip(shape): _gpu_marlin_unsupported() or not is_fp8_marlin_supported(), reason="FP8 Marlin is not supported on this GPU type.", ) -@pytest.mark.parametrize("shape", [(200, 288), (4640, 512)]) +@pytest.mark.parametrize( + "shape", [(32, 256), (96, 256), (130, 256), (200, 288), (4640, 512)] +) def test_mxfp8_marlin_padded_round_trip(shape): """MXFP8 exercises the e8m0 scale path, where padded 0.0 scales clamp to 2^-127 instead of zero and must still contribute nothing.""" diff --git a/tests/kernels/quantization/test_mxfp8_kernel_selection.py b/tests/kernels/quantization/test_mxfp8_kernel_selection.py new file mode 100644 index 000000000000..f559eca4ddd1 --- /dev/null +++ b/tests/kernels/quantization/test_mxfp8_kernel_selection.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for MXFP8 linear kernel selection logic (CPU-only) + +Run `pytest tests/kernels/quantization/test_mxfp8_kernel_selection.py`. +""" + +from contextlib import contextmanager +from unittest.mock import patch + +import pytest + +from vllm.model_executor.kernels.linear import ( + FlashInferCutedslMxfp8LinearKernel, + FlashInferCutlassMxfp8LinearKernel, + MarlinMxfp8LinearKernel, + Mxfp8LinearLayerConfig, + init_mxfp8_linear_kernel, +) +from vllm.platforms import PlatformEnum + +pytestmark = pytest.mark.cpu_test + +# Kernels backed by FlashInfer mm_mxfp8, which requires N, K >= 128 and +# K % 32 == 0. +_MM_MXFP8_KERNELS = [ + FlashInferCutedslMxfp8LinearKernel, + FlashInferCutlassMxfp8LinearKernel, +] + +_SUPPORTED_SHAPE = (4096, 4096) +# (N, K): N < 128, K < 128, K % 32 != 0. +_UNSUPPORTED_SHAPES = [(64, 4096), (4096, 64), (4096, 4112)] + +# (N, K) with N < 128 or K < 128, which only Marlin runs. +_MARLIN_ONLY_SHAPES = [(32, 256), (64, 256), (96, 256), (256, 64)] +# (N, K): N <= 0, K <= 0, K % 32 != 0. +_MARLIN_UNSUPPORTED_SHAPES = [(0, 256), (256, 0), (256, 100)] + +_SELECTION_CASES = [ + (100, (128, 128), FlashInferCutedslMxfp8LinearKernel), + (100, (130, 256), FlashInferCutedslMxfp8LinearKernel), + (103, (128, 128), FlashInferCutedslMxfp8LinearKernel), + (103, (130, 256), FlashInferCutedslMxfp8LinearKernel), + (120, (128, 128), FlashInferCutlassMxfp8LinearKernel), + (120, (130, 256), MarlinMxfp8LinearKernel), + (121, (128, 128), FlashInferCutlassMxfp8LinearKernel), + (121, (130, 256), MarlinMxfp8LinearKernel), + *[ + (capability, shape, MarlinMxfp8LinearKernel) + for capability in (100, 103, 120, 121) + for shape in _MARLIN_ONLY_SHAPES + ], +] + + +@pytest.mark.parametrize("kernel_cls", _MM_MXFP8_KERNELS) +def test_mm_mxfp8_kernels_accept_supported_shape(kernel_cls): + config = Mxfp8LinearLayerConfig(weight_shape=_SUPPORTED_SHAPE) + can_implement, reason = kernel_cls.can_implement(config) + assert can_implement, reason + + +@pytest.mark.parametrize("kernel_cls", _MM_MXFP8_KERNELS) +@pytest.mark.parametrize("weight_shape", _UNSUPPORTED_SHAPES) +def test_mm_mxfp8_kernels_reject_unsupported_shape(kernel_cls, weight_shape): + config = Mxfp8LinearLayerConfig(weight_shape=weight_shape) + can_implement, reason = kernel_cls.can_implement(config) + assert not can_implement + assert reason + + +@pytest.mark.parametrize("weight_shape", _MARLIN_UNSUPPORTED_SHAPES) +def test_marlin_mxfp8_rejects_unsupported_shape(weight_shape): + config = Mxfp8LinearLayerConfig(weight_shape=weight_shape) + can_implement, reason = MarlinMxfp8LinearKernel.can_implement(config) + assert not can_implement + assert reason + + +@contextmanager +def _patch_cuda_platform(capability: int): + """Patch the platform so kernel selection runs as on a CUDA GPU with the + given compute capability. The FlashInfer and Marlin kernels report as + supported, CuTe-DSL only on SM100/103.""" + with ( + patch("vllm.model_executor.kernels.linear.current_platform") as platform, + patch( + "vllm.model_executor.kernels.linear.mxfp8.flashinfer.current_platform." + "is_device_capability_family", + side_effect=lambda cap, device_id=0: cap // 10 == capability // 10, + ), + patch.object( + FlashInferCutedslMxfp8LinearKernel, + "is_supported", + return_value=(capability in (100, 103), None), + ), + patch.object( + FlashInferCutlassMxfp8LinearKernel, + "is_supported", + return_value=(True, None), + ), + patch.object( + MarlinMxfp8LinearKernel, "is_supported", return_value=(True, None) + ), + ): + platform._enum = PlatformEnum.CUDA + yield + + +@pytest.mark.parametrize( + ("capability", "expected"), + [(100, True), (103, True), (120, False), (121, False)], +) +def test_cutlass_mxfp8_requires_aligned_n_only_on_sm12x(capability, expected): + """(130, 256) passes the generic mm_mxfp8 limits. Only the SM12x CUTLASS + backend needs N % 32 == 0, so SM100/103 must keep accepting it.""" + config = Mxfp8LinearLayerConfig(weight_shape=(130, 256)) + with _patch_cuda_platform(capability): + can_implement, _ = FlashInferCutlassMxfp8LinearKernel.can_implement(config) + assert can_implement == expected + + +@pytest.mark.parametrize( + ("capability", "weight_shape", "expected_kernel_cls"), _SELECTION_CASES +) +def test_init_mxfp8_linear_kernel(capability, weight_shape, expected_kernel_cls): + """A layer that mm_mxfp8 cannot handle must fall through to the next + kernel in the CUDA priority list instead of being selected and failing + in apply_weights.""" + with _patch_cuda_platform(capability): + kernel = init_mxfp8_linear_kernel(weight_shape=weight_shape) + + assert isinstance(kernel, expected_kernel_cls) + assert kernel.config == Mxfp8LinearLayerConfig(weight_shape=weight_shape) diff --git a/tests/kernels/quantization/test_rocm_mxfp8_linear.py b/tests/kernels/quantization/test_rocm_mxfp8_linear.py index 1d5e62c6c4b1..ccf0179ec4e5 100644 --- a/tests/kernels/quantization/test_rocm_mxfp8_linear.py +++ b/tests/kernels/quantization/test_rocm_mxfp8_linear.py @@ -66,7 +66,7 @@ def test_rocm_mxfp8_linear_matches_dequant_reference(name, n, k, num_tokens): torch.manual_seed(num_tokens) device = "cuda" layer = _make_layer(n, k, device) - kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(bmm_batch_size=None)) + kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(weight_shape=(n, k))) kernel.process_weights_after_loading(layer) x = torch.randn(num_tokens, k, device=device, dtype=torch.bfloat16) * 0.5 @@ -103,7 +103,7 @@ def test_rocm_mxfp8_linear_handles_extreme_magnitudes(): device = "cuda" n, k = 5120, 2048 # wo_b at TP=4 layer = _make_layer(n, k, device) - kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(bmm_batch_size=None)) + kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(weight_shape=(n, k))) kernel.process_weights_after_loading(layer) x = torch.randn(8, k, device=device, dtype=torch.bfloat16) diff --git a/tests/kernels/test_minimax_m3_amd_ops.py b/tests/kernels/test_minimax_m3_amd_ops.py index 27bd74ce41dc..ffb2a5d77896 100644 --- a/tests/kernels/test_minimax_m3_amd_ops.py +++ b/tests/kernels/test_minimax_m3_amd_ops.py @@ -431,7 +431,7 @@ def test_mxfp8_linear_emulation_bf16_at_load( layer.weight = torch.nn.Parameter(w_fp8.clone(), requires_grad=False) layer.weight_scale = torch.nn.Parameter(w_scale.clone(), requires_grad=False) - kernel = EmulationMxfp8LinearKernel(Mxfp8LinearLayerConfig()) + kernel = EmulationMxfp8LinearKernel(Mxfp8LinearLayerConfig(weight_shape=(N, K))) kernel.process_weights_after_loading(layer) if dequant_at_load: @@ -482,7 +482,9 @@ def make_layer(): layer.weight_scale = torch.nn.Parameter(w_scale.clone(), requires_grad=False) return layer - kernel = rocm_native.RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig()) + kernel = rocm_native.RocmDotScaledMxfp8LinearKernel( + Mxfp8LinearLayerConfig(weight_shape=(N, K)) + ) layer = make_layer() kernel.process_weights_after_loading(layer) assert (layer.weight.element_size() >= 2) is not aligned diff --git a/tests/model_executor/kernels/test_b12x_linear.py b/tests/model_executor/kernels/test_b12x_linear.py index 9d4f9beeb9a3..ac72cfb8452f 100644 --- a/tests/model_executor/kernels/test_b12x_linear.py +++ b/tests/model_executor/kernels/test_b12x_linear.py @@ -66,7 +66,7 @@ "MarlinMxfp8LinearKernel", "EmulationMxfp8LinearKernel", init_mxfp8_linear_kernel, - {}, + {"weight_shape": (2048, 2048)}, ), ( B12xTensorFP8ScaledMMLinearKernel, @@ -365,7 +365,7 @@ def mm( def test_b12x_mxfp8_can_implement_supported_config() -> None: can_implement, reason = B12xMxfp8LinearKernel.can_implement( - Mxfp8LinearLayerConfig() + Mxfp8LinearLayerConfig(weight_shape=(256, 512)) ) assert can_implement diff --git a/tests/quantization/test_auto_round.py b/tests/quantization/test_auto_round.py index ef5f29602fde..f2eeb9e9d9ce 100644 --- a/tests/quantization/test_auto_round.py +++ b/tests/quantization/test_auto_round.py @@ -1249,7 +1249,7 @@ def apply_weights(self, layer, x, bias=None): kernel = DummyKernel() monkeypatch.setattr( "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp8_linear.init_mxfp8_linear_kernel", - lambda: kernel, + lambda weight_shape: kernel, ) monkeypatch.setattr( "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp8_linear.ModelWeightParameter", @@ -1288,7 +1288,7 @@ def apply_weights(self, layer, x, bias=None): def test_inc_mxfp8_linear_scheme_requires_block_32_input(monkeypatch) -> None: monkeypatch.setattr( "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp8_linear.init_mxfp8_linear_kernel", - lambda: object(), + lambda weight_shape: object(), ) scheme = INCMxfp8LinearScheme() diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index fc4696f7054b..e5fed78c40a7 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -880,10 +880,15 @@ def choose_mp_linear_kernel( ) -def init_mxfp8_linear_kernel(*, bmm_batch_size: int | None = None) -> Mxfp8LinearKernel: +def init_mxfp8_linear_kernel( + *, weight_shape: tuple[int, int], bmm_batch_size: int | None = None +) -> Mxfp8LinearKernel: """Select and instantiate the best MXFP8 linear kernel for the - current platform.""" - config = Mxfp8LinearLayerConfig(bmm_batch_size=bmm_batch_size) + current platform and `(N, K)` weight shape.""" + config = Mxfp8LinearLayerConfig( + weight_shape=weight_shape, + bmm_batch_size=bmm_batch_size, + ) platform = current_platform._enum possible: list[type[Mxfp8LinearKernel]] diff --git a/vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py b/vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py index fdcfeb264b29..3a8b886b4c8d 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py +++ b/vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py @@ -15,8 +15,13 @@ class Mxfp8LinearLayerConfig: All MXFP8 layers share the same structure: FP8-E4M3 weights with uint8 (E8M0) per-block scales at block size 32. + + Attributes: + weight_shape: The layer's `(out_features, in_features)`, i.e. `(N, K)`. + """ + weight_shape: tuple[int, int] bmm_batch_size: int | None = None diff --git a/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py b/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py index cc9b35d34bdd..889b260435a5 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py +++ b/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py @@ -24,6 +24,20 @@ from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig +def _check_mm_mxfp8_shape(weight_shape: tuple[int, int]) -> tuple[bool, str | None]: + N, K = weight_shape + if K < 128: + return False, f"mm_mxfp8 requires K >= 128, got K={K}." + if K % MXFP8_BLOCK_SIZE != 0: + return ( + False, + f"mm_mxfp8 requires K to be divisible by {MXFP8_BLOCK_SIZE}, got K={K}.", + ) + if N < 128: + return False, f"mm_mxfp8 requires N >= 128, got N={N}." + return True, None + + class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel): """MXFP8 W8A8 GEMM via FlashInfer CUTLASS (SM100+).""" @@ -43,6 +57,15 @@ def is_supported( @classmethod def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]: + can_implement, reason = _check_mm_mxfp8_shape(c.weight_shape) + if not can_implement: + return False, reason + N, _ = c.weight_shape + if current_platform.is_device_capability_family(120) and N % 32 != 0: + return ( + False, + f"mm_mxfp8 on SM12x requires N to be divisible by 32, got N={N}.", + ) return True, None def input_quant_key(self) -> QuantKey: @@ -72,20 +95,6 @@ def apply_weights( weight_scale = layer.weight_scale N, K = weight.shape - min_dim = 128 - - assert min_dim <= K, ( - f"mm_mxfp8 requires K >= {min_dim}, got K={K}. " - f"in_features is too small for mm_mxfp8." - ) - assert K % MXFP8_BLOCK_SIZE == 0, ( - f"mm_mxfp8 requires K to be divisible by {MXFP8_BLOCK_SIZE}, got K={K}." - ) - assert min_dim <= N, ( - f"mm_mxfp8 requires N >= {min_dim}, got N={N}. " - f"out_features is too small for mm_mxfp8." - ) - qa = as_quantized_activation(x, self.input_quant_key()) if qa is not None: input_mxfp8, input_scale = qa.data, qa.scale @@ -136,7 +145,7 @@ def is_supported( @classmethod def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]: - return True, None + return _check_mm_mxfp8_shape(c.weight_shape) def input_quant_key(self) -> QuantKey: # Activations use FlashInfer's F8_128x4 swizzled scale layout. @@ -166,20 +175,6 @@ def apply_weights( weight_scale = layer.weight_scale K, N = weight.shape - min_dim = 128 - - assert min_dim <= K, ( - f"mm_mxfp8 requires K >= {min_dim}, got K={K}. " - f"in_features is too small for mm_mxfp8." - ) - assert K % MXFP8_BLOCK_SIZE == 0, ( - f"mm_mxfp8 requires K to be divisible by {MXFP8_BLOCK_SIZE}, got K={K}." - ) - assert min_dim <= N, ( - f"mm_mxfp8 requires N >= {min_dim}, got N={N}. " - f"out_features is too small for mm_mxfp8." - ) - qa = as_quantized_activation(x, self.input_quant_key()) if qa is not None: input_mxfp8, input_scale = qa.data, qa.scale diff --git a/vllm/model_executor/kernels/linear/mxfp8/marlin.py b/vllm/model_executor/kernels/linear/mxfp8/marlin.py index 6ebdbd647405..45824e025877 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/marlin.py +++ b/vllm/model_executor/kernels/linear/mxfp8/marlin.py @@ -3,6 +3,10 @@ import torch +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + MXFP8_BLOCK_SIZE, +) + from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig @@ -23,6 +27,16 @@ def is_supported( @classmethod def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]: + N, K = c.weight_shape + if N <= 0 or K <= 0: + return False, f"Marlin MXFP8 requires positive N and K, got N={N}, K={K}." + if K % MXFP8_BLOCK_SIZE != 0: + return ( + False, + f"Marlin MXFP8 requires K to be divisible by {MXFP8_BLOCK_SIZE}, " + f"got K={K}.", + ) + # Tile misalignment is fixed by zero-padding at weight prep. return True, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_mxfp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_mxfp8.py index 0f72be13a554..83bf4bb9cb35 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_mxfp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_mxfp8.py @@ -33,9 +33,6 @@ class CompressedTensorsW8A8Mxfp8(CompressedTensorsScheme): - Activations dynamically quantized to MXFP8 during inference """ - def __init__(self): - self.kernel = init_mxfp8_linear_kernel() - @classmethod def get_min_capability(cls) -> int: return 75 @@ -79,6 +76,8 @@ def create_weights( ) layer.register_parameter("weight_scale", weight_scale) + self.kernel = init_mxfp8_linear_kernel(weight_shape=layer.weight.shape) + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: self.kernel.process_weights_after_loading(layer) diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp8_linear.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp8_linear.py index 775b5f5bba65..b084c697ea3b 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp8_linear.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp8_linear.py @@ -17,9 +17,6 @@ class INCMxfp8LinearScheme(INCLinearScheme): - def __init__(self) -> None: - self.kernel = init_mxfp8_linear_kernel() - @classmethod def get_min_capability(cls) -> int: return 75 @@ -72,6 +69,8 @@ def create_weights( ) layer.register_parameter("weight_scale", weight_scale) + self.kernel = init_mxfp8_linear_kernel(weight_shape=layer.weight.shape) + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: self.kernel.process_weights_after_loading(layer) diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 7c06a2332520..88c63d97b8f2 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -2320,7 +2320,8 @@ def select_linear_kernel( return init_nvfp4_linear_kernel(use_a16=spec.activation is None) if w.scale.dtype == MXFP8_SCALE_DTYPE: return init_mxfp8_linear_kernel( - bmm_batch_size=getattr(layer, "bmm_batch_size", None) + weight_shape=weight_shape or layer.weight.shape, + bmm_batch_size=getattr(layer, "bmm_batch_size", None), ) # fp8 family: init_fp8 routes block-vs-plain itself off the activation key, # and needs a real key -- weight-only fp8 is not a ModelOpt format. @@ -2526,7 +2527,10 @@ def create_weights( def process_weights_after_loading(self, layer) -> None: if self.spec.weight == kMxfp8Static and getattr(layer, "is_bmm", False): - self.kernel = init_mxfp8_linear_kernel(bmm_batch_size=layer.bmm_batch_size) + self.kernel = init_mxfp8_linear_kernel( + weight_shape=(layer.weight.shape[-2], layer.weight.shape[-1]), + bmm_batch_size=layer.bmm_batch_size, + ) if is_weights_pre_processed(): if not self.supports_pre_processed_weights: raise RuntimeError( diff --git a/vllm/model_executor/layers/quantization/online/mxfp8.py b/vllm/model_executor/layers/quantization/online/mxfp8.py index 5165762d5b88..9d5814206c1d 100644 --- a/vllm/model_executor/layers/quantization/online/mxfp8.py +++ b/vllm/model_executor/layers/quantization/online/mxfp8.py @@ -42,10 +42,6 @@ class Mxfp8OnlineLinearMethod(OnlineLinearBase): FP8 with block-32 scales) during weight loading. """ - def __init__(self): - super().__init__() - self.kernel = init_mxfp8_linear_kernel() - def create_weights( self, layer: torch.nn.Module, @@ -73,6 +69,8 @@ def create_weights( **extra_weight_attrs, ) + self.kernel = init_mxfp8_linear_kernel(weight_shape=layer.weight.shape) + def process_weights_after_loading(self, layer: Module) -> None: if getattr(layer, "_already_called_process_weights_after_loading", False): return