[Bugfix][Quantization] Fix MXFP8 startup crash on layers below mm_mxfp8 shape limits (#54223)

Signed-off-by: samuelkim7 <samuelmwkim@gmail.com>
Signed-off-by: mgoin <mgoin64@gmail.com>
Co-authored-by: mgoin <mgoin64@gmail.com>

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",

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")

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."""

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)

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)

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 test_mxfp8_rocm_native_unaligned_k_dequantizes_at_load(shape):
         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

tests/model_executor/kernels/test_b12x_linear.py

@@ -66,7 +66,7 @@ from vllm.platforms import PlatformEnum
             "MarlinMxfp8LinearKernel",
             "EmulationMxfp8LinearKernel",
             init_mxfp8_linear_kernel,
-            {},
+            {"weight_shape": (2048, 2048)},
         ),
         (
             B12xTensorFP8ScaledMMLinearKernel,
@@ -365,7 +365,7 @@ def test_b12x_tensor_fp8_apply_quantizes_and_uses_packed_weight(
 
 def test_b12x_mxfp8_can_implement_supported_config() -> None:
     can_implement, reason = B12xMxfp8LinearKernel.can_implement(
-        Mxfp8LinearLayerConfig()
+        Mxfp8LinearLayerConfig(weight_shape=(256, 512))
     )
 
     assert can_implement

tests/quantization/test_auto_round.py

@@ -1249,7 +1249,7 @@ def test_inc_mxfp8_linear_scheme_delegates_to_kernel(monkeypatch) -> 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 test_inc_mxfp8_linear_scheme_delegates_to_kernel(monkeypatch) -> 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()
 

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]]

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
 
 

vllm/model_executor/kernels/linear/mxfp8/flashinfer.py

@@ -24,6 +24,20 @@ from vllm.utils.flashinfer import has_flashinfer, has_flashinfer_cutedsl
 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 @@ class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel):
 
     @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 @@ class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel):
         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 @@ class FlashInferCutedslMxfp8LinearKernel(Mxfp8LinearKernel):
 
     @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 @@ class FlashInferCutedslMxfp8LinearKernel(Mxfp8LinearKernel):
         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

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 @@ class MarlinMxfp8LinearKernel(Mxfp8LinearKernel):
 
     @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:

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 @@ class CompressedTensorsW8A8Mxfp8(CompressedTensorsScheme):
         )
         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)
 

vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp8_linear.py

@@ -17,9 +17,6 @@ from .inc_scheme import INCLinearScheme
 
 
 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 @@ class INCMxfp8LinearScheme(INCLinearScheme):
         )
         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)
 

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 @@ class ModelOptLinearMethod(LinearMethodBase):
 
     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(

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 @@ class Mxfp8OnlineLinearMethod(OnlineLinearBase):
             **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