[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>
Samuel Kim committed
Sep 25, 2026 at 04:54 UTC
f55419c43ca6a1fca13746c84fde79ad617e1296
16 files changed
+221
-57
tests/kernels/core/test_fused_q_kv_rmsnorm.py
+1
-1
@@ -223,7 +223,6 @@ def test_shared_query_quant_preserves_projection(
223
with set_current_vllm_config(
224
VllmConfig(compilation_config=CompilationConfig(mode=CompilationMode.NONE))
225
):
226
- kernel = cls(Mxfp8LinearLayerConfig())
226
# Match the checkpoint's TP4 attention and replicated indexer widths.
227
if projection == "attention":
228
linear = ColumnParallelLinear.__new__(ColumnParallelLinear)
@@ -247,6 +246,7 @@ def test_shared_query_quant_preserves_projection(
246
)
247
linear.weight = torch.nn.Parameter(w, requires_grad=False)
248
linear.weight_scale = torch.nn.Parameter(ws, requires_grad=False)
249
+ kernel = cls(Mxfp8LinearLayerConfig(weight_shape=w.shape))
250
quant_config = DeepseekV4FP8Config.from_config(
251
{
252
"quant_method": "fp8",
tests/kernels/quantization/test_flashinfer_mxfp8_trtllm.py
+9
-3
@@ -50,7 +50,9 @@ def test_flashinfer_trtllm_mxfp8_linear_numerics(
50
x = torch.randn((m, k), dtype=torch.bfloat16, device="cuda")
51
weight = torch.randn((n, k), dtype=torch.bfloat16, device="cuda")
52
layer = _make_layer(weight)
53
- kernel = FlashInferTrtllmMxfp8LinearKernel(Mxfp8LinearLayerConfig())
53
+ kernel = FlashInferTrtllmMxfp8LinearKernel(
54
+ Mxfp8LinearLayerConfig(weight_shape=(n, k))
55
+ )
56
kernel.process_weights_after_loading(layer)
57
58
output = kernel.apply_weights(layer, x)
@@ -69,7 +71,9 @@ def test_flashinfer_trtllm_mxfp8_custom_ops() -> None:
71
x = torch.randn((7, 512), dtype=torch.bfloat16, device="cuda")
72
weight = torch.randn((256, 512), dtype=torch.bfloat16, device="cuda")
73
layer = _make_layer(weight)
72
- kernel = FlashInferTrtllmMxfp8LinearKernel(Mxfp8LinearLayerConfig())
74
+ kernel = FlashInferTrtllmMxfp8LinearKernel(
75
+ Mxfp8LinearLayerConfig(weight_shape=(256, 512))
76
+ )
77
kernel.process_weights_after_loading(layer)
78
79
torch.library.opcheck(
@@ -104,7 +108,9 @@ def test_flashinfer_trtllm_mxfp8_linear_cuda_graph() -> None:
108
m, n, k = 7, 130, 512
109
weight = torch.randn((n, k), dtype=torch.bfloat16, device="cuda")
110
layer = _make_layer(weight)
107
- kernel = FlashInferTrtllmMxfp8LinearKernel(Mxfp8LinearLayerConfig())
111
+ kernel = FlashInferTrtllmMxfp8LinearKernel(
112
+ Mxfp8LinearLayerConfig(weight_shape=(n, k))
113
+ )
114
kernel.process_weights_after_loading(layer)
115
116
static_x = torch.randn((m, k), dtype=torch.bfloat16, device="cuda")
tests/kernels/quantization/test_marlin_tile_padding.py
+3
-1
@@ -408,7 +408,9 @@ def test_fp8_block_marlin_padded_round_trip(shape):
408
_gpu_marlin_unsupported() or not is_fp8_marlin_supported(),
409
reason="FP8 Marlin is not supported on this GPU type.",
410
)
411
-@pytest.mark.parametrize("shape", [(200, 288), (4640, 512)])
411
+@pytest.mark.parametrize(
412
+ "shape", [(32, 256), (96, 256), (130, 256), (200, 288), (4640, 512)]
413
+)
414
def test_mxfp8_marlin_padded_round_trip(shape):
415
"""MXFP8 exercises the e8m0 scale path, where padded 0.0 scales clamp to
416
2^-127 instead of zero and must still contribute nothing."""
tests/kernels/quantization/test_mxfp8_kernel_selection.py
new
+135
@@ -0,0 +1,135 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for MXFP8 linear kernel selection logic (CPU-only)
4
+
5
+Run `pytest tests/kernels/quantization/test_mxfp8_kernel_selection.py`.
6
+"""
7
+
8
+from contextlib import contextmanager
9
+from unittest.mock import patch
10
+
11
+import pytest
12
+
13
+from vllm.model_executor.kernels.linear import (
14
+ FlashInferCutedslMxfp8LinearKernel,
15
+ FlashInferCutlassMxfp8LinearKernel,
16
+ MarlinMxfp8LinearKernel,
17
+ Mxfp8LinearLayerConfig,
18
+ init_mxfp8_linear_kernel,
19
+)
20
+from vllm.platforms import PlatformEnum
21
+
22
+pytestmark = pytest.mark.cpu_test
23
+
24
+# Kernels backed by FlashInfer mm_mxfp8, which requires N, K >= 128 and
25
+# K % 32 == 0.
26
+_MM_MXFP8_KERNELS = [
27
+ FlashInferCutedslMxfp8LinearKernel,
28
+ FlashInferCutlassMxfp8LinearKernel,
29
+]
30
+
31
+_SUPPORTED_SHAPE = (4096, 4096)
32
+# (N, K): N < 128, K < 128, K % 32 != 0.
33
+_UNSUPPORTED_SHAPES = [(64, 4096), (4096, 64), (4096, 4112)]
34
+
35
+# (N, K) with N < 128 or K < 128, which only Marlin runs.
36
+_MARLIN_ONLY_SHAPES = [(32, 256), (64, 256), (96, 256), (256, 64)]
37
+# (N, K): N <= 0, K <= 0, K % 32 != 0.
38
+_MARLIN_UNSUPPORTED_SHAPES = [(0, 256), (256, 0), (256, 100)]
39
+
40
+_SELECTION_CASES = [
41
+ (100, (128, 128), FlashInferCutedslMxfp8LinearKernel),
42
+ (100, (130, 256), FlashInferCutedslMxfp8LinearKernel),
43
+ (103, (128, 128), FlashInferCutedslMxfp8LinearKernel),
44
+ (103, (130, 256), FlashInferCutedslMxfp8LinearKernel),
45
+ (120, (128, 128), FlashInferCutlassMxfp8LinearKernel),
46
+ (120, (130, 256), MarlinMxfp8LinearKernel),
47
+ (121, (128, 128), FlashInferCutlassMxfp8LinearKernel),
48
+ (121, (130, 256), MarlinMxfp8LinearKernel),
49
+ *[
50
+ (capability, shape, MarlinMxfp8LinearKernel)
51
+ for capability in (100, 103, 120, 121)
52
+ for shape in _MARLIN_ONLY_SHAPES
53
+ ],
54
+]
55
+
56
+
57
+@pytest.mark.parametrize("kernel_cls", _MM_MXFP8_KERNELS)
58
+def test_mm_mxfp8_kernels_accept_supported_shape(kernel_cls):
59
+ config = Mxfp8LinearLayerConfig(weight_shape=_SUPPORTED_SHAPE)
60
+ can_implement, reason = kernel_cls.can_implement(config)
61
+ assert can_implement, reason
62
+
63
+
64
+@pytest.mark.parametrize("kernel_cls", _MM_MXFP8_KERNELS)
65
+@pytest.mark.parametrize("weight_shape", _UNSUPPORTED_SHAPES)
66
+def test_mm_mxfp8_kernels_reject_unsupported_shape(kernel_cls, weight_shape):
67
+ config = Mxfp8LinearLayerConfig(weight_shape=weight_shape)
68
+ can_implement, reason = kernel_cls.can_implement(config)
69
+ assert not can_implement
70
+ assert reason
71
+
72
+
73
+@pytest.mark.parametrize("weight_shape", _MARLIN_UNSUPPORTED_SHAPES)
74
+def test_marlin_mxfp8_rejects_unsupported_shape(weight_shape):
75
+ config = Mxfp8LinearLayerConfig(weight_shape=weight_shape)
76
+ can_implement, reason = MarlinMxfp8LinearKernel.can_implement(config)
77
+ assert not can_implement
78
+ assert reason
79
+
80
+
81
+@contextmanager
82
+def _patch_cuda_platform(capability: int):
83
+ """Patch the platform so kernel selection runs as on a CUDA GPU with the
84
+ given compute capability. The FlashInfer and Marlin kernels report as
85
+ supported, CuTe-DSL only on SM100/103."""
86
+ with (
87
+ patch("vllm.model_executor.kernels.linear.current_platform") as platform,
88
+ patch(
89
+ "vllm.model_executor.kernels.linear.mxfp8.flashinfer.current_platform."
90
+ "is_device_capability_family",
91
+ side_effect=lambda cap, device_id=0: cap // 10 == capability // 10,
92
+ ),
93
+ patch.object(
94
+ FlashInferCutedslMxfp8LinearKernel,
95
+ "is_supported",
96
+ return_value=(capability in (100, 103), None),
97
+ ),
98
+ patch.object(
99
+ FlashInferCutlassMxfp8LinearKernel,
100
+ "is_supported",
101
+ return_value=(True, None),
102
+ ),
103
+ patch.object(
104
+ MarlinMxfp8LinearKernel, "is_supported", return_value=(True, None)
105
+ ),
106
+ ):
107
+ platform._enum = PlatformEnum.CUDA
108
+ yield
109
+
110
+
111
+@pytest.mark.parametrize(
112
+ ("capability", "expected"),
113
+ [(100, True), (103, True), (120, False), (121, False)],
114
+)
115
+def test_cutlass_mxfp8_requires_aligned_n_only_on_sm12x(capability, expected):
116
+ """(130, 256) passes the generic mm_mxfp8 limits. Only the SM12x CUTLASS
117
+ backend needs N % 32 == 0, so SM100/103 must keep accepting it."""
118
+ config = Mxfp8LinearLayerConfig(weight_shape=(130, 256))
119
+ with _patch_cuda_platform(capability):
120
+ can_implement, _ = FlashInferCutlassMxfp8LinearKernel.can_implement(config)
121
+ assert can_implement == expected
122
+
123
+
124
+@pytest.mark.parametrize(
125
+ ("capability", "weight_shape", "expected_kernel_cls"), _SELECTION_CASES
126
+)
127
+def test_init_mxfp8_linear_kernel(capability, weight_shape, expected_kernel_cls):
128
+ """A layer that mm_mxfp8 cannot handle must fall through to the next
129
+ kernel in the CUDA priority list instead of being selected and failing
130
+ in apply_weights."""
131
+ with _patch_cuda_platform(capability):
132
+ kernel = init_mxfp8_linear_kernel(weight_shape=weight_shape)
133
+
134
+ assert isinstance(kernel, expected_kernel_cls)
135
+ assert kernel.config == Mxfp8LinearLayerConfig(weight_shape=weight_shape)
tests/kernels/quantization/test_rocm_mxfp8_linear.py
+2
-2
@@ -66,7 +66,7 @@ def test_rocm_mxfp8_linear_matches_dequant_reference(name, n, k, num_tokens):
66
torch.manual_seed(num_tokens)
67
device = "cuda"
68
layer = _make_layer(n, k, device)
69
- kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(bmm_batch_size=None))
69
+ kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(weight_shape=(n, k)))
70
kernel.process_weights_after_loading(layer)
71
72
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():
103
device = "cuda"
104
n, k = 5120, 2048 # wo_b at TP=4
105
layer = _make_layer(n, k, device)
106
- kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(bmm_batch_size=None))
106
+ kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(weight_shape=(n, k)))
107
kernel.process_weights_after_loading(layer)
108
109
x = torch.randn(8, k, device=device, dtype=torch.bfloat16)
tests/kernels/test_minimax_m3_amd_ops.py
+4
-2
@@ -431,7 +431,7 @@ def test_mxfp8_linear_emulation_bf16_at_load(
431
layer.weight = torch.nn.Parameter(w_fp8.clone(), requires_grad=False)
432
layer.weight_scale = torch.nn.Parameter(w_scale.clone(), requires_grad=False)
433
434
- kernel = EmulationMxfp8LinearKernel(Mxfp8LinearLayerConfig())
434
+ kernel = EmulationMxfp8LinearKernel(Mxfp8LinearLayerConfig(weight_shape=(N, K)))
435
kernel.process_weights_after_loading(layer)
436
437
if dequant_at_load:
@@ -482,7 +482,9 @@ def test_mxfp8_rocm_native_unaligned_k_dequantizes_at_load(shape):
482
layer.weight_scale = torch.nn.Parameter(w_scale.clone(), requires_grad=False)
483
return layer
484
485
- kernel = rocm_native.RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig())
485
+ kernel = rocm_native.RocmDotScaledMxfp8LinearKernel(
486
+ Mxfp8LinearLayerConfig(weight_shape=(N, K))
487
+ )
488
layer = make_layer()
489
kernel.process_weights_after_loading(layer)
490
assert (layer.weight.element_size() >= 2) is not aligned
tests/model_executor/kernels/test_b12x_linear.py
+2
-2
@@ -66,7 +66,7 @@ from vllm.platforms import PlatformEnum
66
"MarlinMxfp8LinearKernel",
67
"EmulationMxfp8LinearKernel",
68
init_mxfp8_linear_kernel,
69
- {},
69
+ {"weight_shape": (2048, 2048)},
70
),
71
(
72
B12xTensorFP8ScaledMMLinearKernel,
@@ -365,7 +365,7 @@ def test_b12x_tensor_fp8_apply_quantizes_and_uses_packed_weight(
365
366
def test_b12x_mxfp8_can_implement_supported_config() -> None:
367
can_implement, reason = B12xMxfp8LinearKernel.can_implement(
368
- Mxfp8LinearLayerConfig()
368
+ Mxfp8LinearLayerConfig(weight_shape=(256, 512))
369
)
370
371
assert can_implement
tests/quantization/test_auto_round.py
+2
-2
@@ -1249,7 +1249,7 @@ def test_inc_mxfp8_linear_scheme_delegates_to_kernel(monkeypatch) -> None:
1249
kernel = DummyKernel()
1250
monkeypatch.setattr(
1251
"vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp8_linear.init_mxfp8_linear_kernel",
1252
- lambda: kernel,
1252
+ lambda weight_shape: kernel,
1253
)
1254
monkeypatch.setattr(
1255
"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:
1288
def test_inc_mxfp8_linear_scheme_requires_block_32_input(monkeypatch) -> None:
1289
monkeypatch.setattr(
1290
"vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp8_linear.init_mxfp8_linear_kernel",
1291
- lambda: object(),
1291
+ lambda weight_shape: object(),
1292
)
1293
scheme = INCMxfp8LinearScheme()
1294
vllm/model_executor/kernels/linear/__init__.py
+8
-3
@@ -880,10 +880,15 @@ def choose_mp_linear_kernel(
880
)
881
882
883
-def init_mxfp8_linear_kernel(*, bmm_batch_size: int | None = None) -> Mxfp8LinearKernel:
883
+def init_mxfp8_linear_kernel(
884
+ *, weight_shape: tuple[int, int], bmm_batch_size: int | None = None
885
+) -> Mxfp8LinearKernel:
886
"""Select and instantiate the best MXFP8 linear kernel for the
885
- current platform."""
886
- config = Mxfp8LinearLayerConfig(bmm_batch_size=bmm_batch_size)
887
+ current platform and `(N, K)` weight shape."""
888
+ config = Mxfp8LinearLayerConfig(
889
+ weight_shape=weight_shape,
890
+ bmm_batch_size=bmm_batch_size,
891
+ )
892
893
platform = current_platform._enum
894
possible: list[type[Mxfp8LinearKernel]]
vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py
+5
@@ -15,8 +15,13 @@ class Mxfp8LinearLayerConfig:
15
16
All MXFP8 layers share the same structure: FP8-E4M3 weights with
17
uint8 (E8M0) per-block scales at block size 32.
18
+
19
+ Attributes:
20
+ weight_shape: The layer's `(out_features, in_features)`, i.e. `(N, K)`.
21
+
22
"""
23
24
+ weight_shape: tuple[int, int]
25
bmm_batch_size: int | None = None
26
27
vllm/model_executor/kernels/linear/mxfp8/flashinfer.py
+24
-29
@@ -24,6 +24,20 @@ from vllm.utils.flashinfer import has_flashinfer, has_flashinfer_cutedsl
24
from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
25
26
27
+def _check_mm_mxfp8_shape(weight_shape: tuple[int, int]) -> tuple[bool, str | None]:
28
+ N, K = weight_shape
29
+ if K < 128:
30
+ return False, f"mm_mxfp8 requires K >= 128, got K={K}."
31
+ if K % MXFP8_BLOCK_SIZE != 0:
32
+ return (
33
+ False,
34
+ f"mm_mxfp8 requires K to be divisible by {MXFP8_BLOCK_SIZE}, got K={K}.",
35
+ )
36
+ if N < 128:
37
+ return False, f"mm_mxfp8 requires N >= 128, got N={N}."
38
+ return True, None
39
+
40
+
41
class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel):
42
"""MXFP8 W8A8 GEMM via FlashInfer CUTLASS (SM100+)."""
43
@@ -43,6 +57,15 @@ class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel):
57
58
@classmethod
59
def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
60
+ can_implement, reason = _check_mm_mxfp8_shape(c.weight_shape)
61
+ if not can_implement:
62
+ return False, reason
63
+ N, _ = c.weight_shape
64
+ if current_platform.is_device_capability_family(120) and N % 32 != 0:
65
+ return (
66
+ False,
67
+ f"mm_mxfp8 on SM12x requires N to be divisible by 32, got N={N}.",
68
+ )
69
return True, None
70
71
def input_quant_key(self) -> QuantKey:
@@ -72,20 +95,6 @@ class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel):
95
weight_scale = layer.weight_scale
96
N, K = weight.shape
97
75
- min_dim = 128
76
-
77
- assert min_dim <= K, (
78
- f"mm_mxfp8 requires K >= {min_dim}, got K={K}. "
79
- f"in_features is too small for mm_mxfp8."
80
- )
81
- assert K % MXFP8_BLOCK_SIZE == 0, (
82
- f"mm_mxfp8 requires K to be divisible by {MXFP8_BLOCK_SIZE}, got K={K}."
83
- )
84
- assert min_dim <= N, (
85
- f"mm_mxfp8 requires N >= {min_dim}, got N={N}. "
86
- f"out_features is too small for mm_mxfp8."
87
- )
88
-
98
qa = as_quantized_activation(x, self.input_quant_key())
99
if qa is not None:
100
input_mxfp8, input_scale = qa.data, qa.scale
@@ -136,7 +145,7 @@ class FlashInferCutedslMxfp8LinearKernel(Mxfp8LinearKernel):
145
146
@classmethod
147
def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
139
- return True, None
148
+ return _check_mm_mxfp8_shape(c.weight_shape)
149
150
def input_quant_key(self) -> QuantKey:
151
# Activations use FlashInfer's F8_128x4 swizzled scale layout.
@@ -166,20 +175,6 @@ class FlashInferCutedslMxfp8LinearKernel(Mxfp8LinearKernel):
175
weight_scale = layer.weight_scale
176
K, N = weight.shape
177
169
- min_dim = 128
170
-
171
- assert min_dim <= K, (
172
- f"mm_mxfp8 requires K >= {min_dim}, got K={K}. "
173
- f"in_features is too small for mm_mxfp8."
174
- )
175
- assert K % MXFP8_BLOCK_SIZE == 0, (
176
- f"mm_mxfp8 requires K to be divisible by {MXFP8_BLOCK_SIZE}, got K={K}."
177
- )
178
- assert min_dim <= N, (
179
- f"mm_mxfp8 requires N >= {min_dim}, got N={N}. "
180
- f"out_features is too small for mm_mxfp8."
181
- )
182
-
178
qa = as_quantized_activation(x, self.input_quant_key())
179
if qa is not None:
180
input_mxfp8, input_scale = qa.data, qa.scale
vllm/model_executor/kernels/linear/mxfp8/marlin.py
+14
@@ -3,6 +3,10 @@
3
4
import torch
5
6
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
7
+ MXFP8_BLOCK_SIZE,
8
+)
9
+
10
from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
11
12
@@ -23,6 +27,16 @@ class MarlinMxfp8LinearKernel(Mxfp8LinearKernel):
27
28
@classmethod
29
def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
30
+ N, K = c.weight_shape
31
+ if N <= 0 or K <= 0:
32
+ return False, f"Marlin MXFP8 requires positive N and K, got N={N}, K={K}."
33
+ if K % MXFP8_BLOCK_SIZE != 0:
34
+ return (
35
+ False,
36
+ f"Marlin MXFP8 requires K to be divisible by {MXFP8_BLOCK_SIZE}, "
37
+ f"got K={K}.",
38
+ )
39
+ # Tile misalignment is fixed by zero-padding at weight prep.
40
return True, None
41
42
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_mxfp8.py
+2
-3
@@ -33,9 +33,6 @@ class CompressedTensorsW8A8Mxfp8(CompressedTensorsScheme):
33
- Activations dynamically quantized to MXFP8 during inference
34
"""
35
36
- def __init__(self):
37
- self.kernel = init_mxfp8_linear_kernel()
38
-
36
@classmethod
37
def get_min_capability(cls) -> int:
38
return 75
@@ -79,6 +76,8 @@ class CompressedTensorsW8A8Mxfp8(CompressedTensorsScheme):
76
)
77
layer.register_parameter("weight_scale", weight_scale)
78
79
+ self.kernel = init_mxfp8_linear_kernel(weight_shape=layer.weight.shape)
80
+
81
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
82
self.kernel.process_weights_after_loading(layer)
83
vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp8_linear.py
+2
-3
@@ -17,9 +17,6 @@ from .inc_scheme import INCLinearScheme
17
18
19
class INCMxfp8LinearScheme(INCLinearScheme):
20
- def __init__(self) -> None:
21
- self.kernel = init_mxfp8_linear_kernel()
22
-
20
@classmethod
21
def get_min_capability(cls) -> int:
22
return 75
@@ -72,6 +69,8 @@ class INCMxfp8LinearScheme(INCLinearScheme):
69
)
70
layer.register_parameter("weight_scale", weight_scale)
71
72
+ self.kernel = init_mxfp8_linear_kernel(weight_shape=layer.weight.shape)
73
+
74
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
75
self.kernel.process_weights_after_loading(layer)
76
vllm/model_executor/layers/quantization/modelopt.py
+6
-2
@@ -2320,7 +2320,8 @@ def select_linear_kernel(
2320
return init_nvfp4_linear_kernel(use_a16=spec.activation is None)
2321
if w.scale.dtype == MXFP8_SCALE_DTYPE:
2322
return init_mxfp8_linear_kernel(
2323
- bmm_batch_size=getattr(layer, "bmm_batch_size", None)
2323
+ weight_shape=weight_shape or layer.weight.shape,
2324
+ bmm_batch_size=getattr(layer, "bmm_batch_size", None),
2325
)
2326
# fp8 family: init_fp8 routes block-vs-plain itself off the activation key,
2327
# and needs a real key -- weight-only fp8 is not a ModelOpt format.
@@ -2526,7 +2527,10 @@ class ModelOptLinearMethod(LinearMethodBase):
2527
2528
def process_weights_after_loading(self, layer) -> None:
2529
if self.spec.weight == kMxfp8Static and getattr(layer, "is_bmm", False):
2529
- self.kernel = init_mxfp8_linear_kernel(bmm_batch_size=layer.bmm_batch_size)
2530
+ self.kernel = init_mxfp8_linear_kernel(
2531
+ weight_shape=(layer.weight.shape[-2], layer.weight.shape[-1]),
2532
+ bmm_batch_size=layer.bmm_batch_size,
2533
+ )
2534
if is_weights_pre_processed():
2535
if not self.supports_pre_processed_weights:
2536
raise RuntimeError(
vllm/model_executor/layers/quantization/online/mxfp8.py
+2
-4
@@ -42,10 +42,6 @@ class Mxfp8OnlineLinearMethod(OnlineLinearBase):
42
FP8 with block-32 scales) during weight loading.
43
"""
44
45
- def __init__(self):
46
- super().__init__()
47
- self.kernel = init_mxfp8_linear_kernel()
48
-
45
def create_weights(
46
self,
47
layer: torch.nn.Module,
@@ -73,6 +69,8 @@ class Mxfp8OnlineLinearMethod(OnlineLinearBase):
69
**extra_weight_attrs,
70
)
71
72
+ self.kernel = init_mxfp8_linear_kernel(weight_shape=layer.weight.shape)
73
+
74
def process_weights_after_loading(self, layer: Module) -> None:
75
if getattr(layer, "_already_called_process_weights_after_loading", False):
76
return