A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 07:15 UTC
[Refactor] Move MXFP8 GEMM management into MxFp8LinearKernel (#39205)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Apr 10, 2026 at 23:02 UTC
11e2375fe28d2a39fecd80125bc65eab5aa30d75
10 files changed
+365
-258
vllm/model_executor/kernels/linear/__init__.py
+69
@@ -58,6 +58,19 @@ from vllm.model_executor.kernels.linear.mixed_precision.xpu import (
58
XPUW4A8IntLinearKernel,
59
XPUwNa16LinearKernel,
60
)
61
+from vllm.model_executor.kernels.linear.mxfp8 import (
62
+ Mxfp8LinearKernel,
63
+ Mxfp8LinearLayerConfig,
64
+)
65
+from vllm.model_executor.kernels.linear.mxfp8.emulation import (
66
+ EmulationMxfp8LinearKernel,
67
+)
68
+from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
69
+ FlashInferCutlassMxfp8LinearKernel,
70
+)
71
+from vllm.model_executor.kernels.linear.mxfp8.marlin import (
72
+ MarlinMxfp8LinearKernel,
73
+)
74
from vllm.model_executor.kernels.linear.nvfp4 import (
75
NvFp4LinearKernel,
76
NvFp4LinearLayerConfig,
@@ -221,6 +234,17 @@ _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = {
234
}
235
236
# in priority/performance order (when available)
237
+_POSSIBLE_MXFP8_KERNELS: dict[PlatformEnum, list[type[Mxfp8LinearKernel]]] = {
238
+ PlatformEnum.CUDA: [
239
+ FlashInferCutlassMxfp8LinearKernel,
240
+ MarlinMxfp8LinearKernel,
241
+ EmulationMxfp8LinearKernel,
242
+ ],
243
+ PlatformEnum.ROCM: [
244
+ EmulationMxfp8LinearKernel,
245
+ ],
246
+}
247
+
248
_POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = {
249
PlatformEnum.CUDA: [
250
FlashInferCutlassNvFp4LinearKernel,
@@ -482,6 +506,41 @@ def choose_mp_linear_kernel(
506
)
507
508
509
+def init_mxfp8_linear_kernel() -> Mxfp8LinearKernel:
510
+ """Select and instantiate the best MXFP8 linear kernel for the
511
+ current platform."""
512
+ config = Mxfp8LinearLayerConfig()
513
+
514
+ platform = current_platform._enum
515
+ possible = _POSSIBLE_MXFP8_KERNELS.get(platform, [])
516
+
517
+ failure_reasons = []
518
+ for kernel_cls in possible:
519
+ if kernel_cls.__name__ in envs.VLLM_DISABLED_KERNELS:
520
+ failure_reasons.append(
521
+ f" {kernel_cls.__name__} disabled by environment variable"
522
+ )
523
+ continue
524
+
525
+ is_supported, reason = kernel_cls.is_supported()
526
+ if not is_supported:
527
+ failure_reasons.append(f"{kernel_cls.__name__}: {reason}")
528
+ continue
529
+
530
+ can_implement, reason = kernel_cls.can_implement(config)
531
+ if not can_implement:
532
+ failure_reasons.append(f"{kernel_cls.__name__}: {reason}")
533
+ continue
534
+
535
+ logger.info_once("Using %s for MXFP8 GEMM", kernel_cls.__name__)
536
+ return kernel_cls(config)
537
+
538
+ raise ValueError(
539
+ "Failed to find a kernel that can implement the "
540
+ "MXFP8 linear layer. Reasons: \n" + "\n".join(failure_reasons)
541
+ )
542
+
543
+
544
def init_wfp8_a16_linear_kernel(
545
weight_quant_key: QuantKey,
546
activation_quant_key: QuantKey,
@@ -628,6 +687,10 @@ def register_linear_kernel(
687
if platform not in _POSSIBLE_FP8_KERNELS:
688
_POSSIBLE_FP8_KERNELS[platform] = []
689
_POSSIBLE_FP8_KERNELS[platform].append(kernel_class)
690
+ elif kernel_type == "mxfp8":
691
+ if platform not in _POSSIBLE_MXFP8_KERNELS:
692
+ _POSSIBLE_MXFP8_KERNELS[platform] = []
693
+ _POSSIBLE_MXFP8_KERNELS[platform].append(kernel_class)
694
elif kernel_type == "nvfp4":
695
if platform not in _POSSIBLE_NVFP4_KERNELS:
696
_POSSIBLE_NVFP4_KERNELS[platform] = []
@@ -674,6 +737,12 @@ __all__ = [
737
"TritonW4A16LinearKernel",
738
"XPUW4A8IntLinearKernel",
739
"XPUwNa16LinearKernel",
740
+ "init_mxfp8_linear_kernel",
741
+ "Mxfp8LinearKernel",
742
+ "Mxfp8LinearLayerConfig",
743
+ "FlashInferCutlassMxfp8LinearKernel",
744
+ "MarlinMxfp8LinearKernel",
745
+ "EmulationMxfp8LinearKernel",
746
"CutlassNvFp4LinearKernel",
747
"EmulationNvFp4LinearKernel",
748
"FbgemmNvFp4LinearKernel",
vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py
new
+56
@@ -0,0 +1,56 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from abc import ABC, abstractmethod
5
+from dataclasses import dataclass
6
+
7
+import torch
8
+
9
+
10
+@dataclass
11
+class Mxfp8LinearLayerConfig:
12
+ """Configuration for an MXFP8 linear layer.
13
+
14
+ All MXFP8 layers share the same structure: FP8-E4M3 weights with
15
+ uint8 (E8M0) per-block scales at block size 32.
16
+ """
17
+
18
+ pass
19
+
20
+
21
+class Mxfp8LinearKernel(ABC):
22
+ """Base class for MXFP8 quantized linear kernels.
23
+
24
+ Each subclass implements a specific GEMM backend (FlashInfer CUTLASS,
25
+ Marlin, emulation).
26
+ """
27
+
28
+ def __init__(self, c: Mxfp8LinearLayerConfig) -> None:
29
+ assert self.can_implement(c)[0]
30
+ assert self.is_supported()[0]
31
+ self.config = c
32
+
33
+ @classmethod
34
+ @abstractmethod
35
+ def is_supported(
36
+ cls, compute_capability: int | None = None
37
+ ) -> tuple[bool, str | None]:
38
+ raise NotImplementedError
39
+
40
+ @classmethod
41
+ @abstractmethod
42
+ def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
43
+ raise NotImplementedError
44
+
45
+ @abstractmethod
46
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
47
+ raise NotImplementedError
48
+
49
+ @abstractmethod
50
+ def apply_weights(
51
+ self,
52
+ layer: torch.nn.Module,
53
+ x: torch.Tensor,
54
+ bias: torch.Tensor | None = None,
55
+ ) -> torch.Tensor:
56
+ raise NotImplementedError
vllm/model_executor/kernels/linear/mxfp8/__init__.py
new
+12
@@ -0,0 +1,12 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import (
5
+ Mxfp8LinearKernel,
6
+ Mxfp8LinearLayerConfig,
7
+)
8
+
9
+__all__ = [
10
+ "Mxfp8LinearKernel",
11
+ "Mxfp8LinearLayerConfig",
12
+]
vllm/model_executor/kernels/linear/mxfp8/emulation.py
new
+60
@@ -0,0 +1,60 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+from torch.nn.parameter import Parameter
6
+
7
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
8
+ MXFP8_BLOCK_SIZE,
9
+ MXFP8_SCALE_DTYPE,
10
+ dequant_mxfp8_to_bf16,
11
+)
12
+
13
+from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
14
+
15
+
16
+class EmulationMxfp8LinearKernel(Mxfp8LinearKernel):
17
+ """Software emulation fallback for MXFP8 (dequant to BF16)."""
18
+
19
+ @classmethod
20
+ def is_supported(
21
+ cls, compute_capability: int | None = None
22
+ ) -> tuple[bool, str | None]:
23
+ return True, None
24
+
25
+ @classmethod
26
+ def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
27
+ return True, None
28
+
29
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
30
+ weight = layer.weight.data # [N, K]
31
+ N, K = weight.shape
32
+ scale_k = K // MXFP8_BLOCK_SIZE
33
+
34
+ weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous()
35
+
36
+ layer.weight = Parameter(weight.contiguous(), requires_grad=False)
37
+ layer.weight_scale = Parameter(weight_scale, requires_grad=False)
38
+
39
+ def apply_weights(
40
+ self,
41
+ layer: torch.nn.Module,
42
+ x: torch.Tensor,
43
+ bias: torch.Tensor | None = None,
44
+ ) -> torch.Tensor:
45
+ weight_scale = layer.weight_scale
46
+ if weight_scale.dtype != MXFP8_SCALE_DTYPE:
47
+ raise ValueError(
48
+ f"Emulation backend requires {MXFP8_SCALE_DTYPE} "
49
+ f"weight_scale dtype, got {weight_scale.dtype}."
50
+ )
51
+ if weight_scale.ndim != 2:
52
+ raise ValueError(
53
+ f"Emulation backend requires 2D weight_scale, "
54
+ f"got {weight_scale.ndim}D. "
55
+ f"Ensure process_weights_after_loading was called."
56
+ )
57
+
58
+ weight_bf16 = dequant_mxfp8_to_bf16(layer.weight, weight_scale)
59
+ output = torch.nn.functional.linear(x, weight_bf16, bias)
60
+ return output.to(x.dtype)
vllm/model_executor/kernels/linear/mxfp8/flashinfer.py
new
+103
@@ -0,0 +1,103 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+from torch.nn.parameter import Parameter
6
+
7
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
8
+ MXFP8_BLOCK_SIZE,
9
+ mxfp8_e4m3_quantize,
10
+ swizzle_mxfp8_scale,
11
+)
12
+from vllm.platforms import current_platform
13
+from vllm.utils import flashinfer as vllm_flashinfer
14
+
15
+from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
16
+
17
+
18
+class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel):
19
+ """MXFP8 W8A8 GEMM via FlashInfer CUTLASS (SM100+)."""
20
+
21
+ @classmethod
22
+ def is_supported(
23
+ cls, compute_capability: int | None = None
24
+ ) -> tuple[bool, str | None]:
25
+ if current_platform.has_device_capability(100):
26
+ return True, None
27
+ return False, "requires >=sm_100 (Blackwell)"
28
+
29
+ @classmethod
30
+ def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
31
+ return True, None
32
+
33
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
34
+ weight = layer.weight.data # [N, K]
35
+ N, K = weight.shape
36
+
37
+ scale_k = K // MXFP8_BLOCK_SIZE
38
+ weight_scale_2d = layer.weight_scale.data[:N, :scale_k].contiguous()
39
+ weight_scale_swizzled = swizzle_mxfp8_scale(weight_scale_2d, M=N, K=K)
40
+
41
+ layer.weight = Parameter(weight.contiguous(), requires_grad=False)
42
+ layer.weight_scale = Parameter(
43
+ weight_scale_swizzled.contiguous(), requires_grad=False
44
+ )
45
+
46
+ def apply_weights(
47
+ self,
48
+ layer: torch.nn.Module,
49
+ x: torch.Tensor,
50
+ bias: torch.Tensor | None = None,
51
+ ) -> torch.Tensor:
52
+ weight = layer.weight
53
+ weight_scale = layer.weight_scale
54
+ out_dtype = x.dtype
55
+ N, K = weight.shape
56
+
57
+ input_shape = x.shape
58
+ input_2d = x.view(-1, K)
59
+ M_orig = input_2d.shape[0]
60
+
61
+ min_dim = 128
62
+
63
+ assert min_dim <= K, (
64
+ f"mm_mxfp8 requires K >= {min_dim}, got K={K}. "
65
+ f"in_features is too small for mm_mxfp8."
66
+ )
67
+ assert K % MXFP8_BLOCK_SIZE == 0, (
68
+ f"mm_mxfp8 requires K to be divisible by {MXFP8_BLOCK_SIZE}, got K={K}."
69
+ )
70
+ assert min_dim <= N, (
71
+ f"mm_mxfp8 requires N >= {min_dim}, got N={N}. "
72
+ f"out_features is too small for mm_mxfp8."
73
+ )
74
+
75
+ M_padded = ((M_orig + min_dim - 1) // min_dim) * min_dim
76
+ if M_padded != M_orig:
77
+ pad_rows = M_padded - M_orig
78
+ input_2d = torch.nn.functional.pad(input_2d, (0, 0, 0, pad_rows))
79
+
80
+ input_mxfp8, input_scale = mxfp8_e4m3_quantize(
81
+ input_2d, is_sf_swizzled_layout=True
82
+ )
83
+
84
+ if not weight.is_contiguous():
85
+ weight = weight.contiguous()
86
+
87
+ output = vllm_flashinfer.mm_mxfp8(
88
+ input_mxfp8,
89
+ weight.t(),
90
+ input_scale,
91
+ weight_scale,
92
+ out_dtype=out_dtype,
93
+ backend="cutlass",
94
+ )
95
+
96
+ if M_padded != M_orig:
97
+ output = output[:M_orig, :]
98
+
99
+ if bias is not None:
100
+ output = output + bias
101
+
102
+ output_shape = (*input_shape[:-1], N)
103
+ return output.view(output_shape)
vllm/model_executor/kernels/linear/mxfp8/marlin.py
new
+53
@@ -0,0 +1,53 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
7
+
8
+
9
+class MarlinMxfp8LinearKernel(Mxfp8LinearKernel):
10
+ """MXFP8 W8A16 GEMM via Marlin (SM80+)."""
11
+
12
+ @classmethod
13
+ def is_supported(
14
+ cls, compute_capability: int | None = None
15
+ ) -> tuple[bool, str | None]:
16
+ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
17
+ is_fp8_marlin_supported,
18
+ )
19
+
20
+ if is_fp8_marlin_supported():
21
+ return True, None
22
+ return False, "Marlin FP8 not available"
23
+
24
+ @classmethod
25
+ def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
26
+ return True, None
27
+
28
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
29
+ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
30
+ prepare_mxfp8_layer_for_marlin,
31
+ )
32
+
33
+ prepare_mxfp8_layer_for_marlin(layer)
34
+
35
+ def apply_weights(
36
+ self,
37
+ layer: torch.nn.Module,
38
+ x: torch.Tensor,
39
+ bias: torch.Tensor | None = None,
40
+ ) -> torch.Tensor:
41
+ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
42
+ apply_mxfp8_marlin_linear,
43
+ )
44
+
45
+ return apply_mxfp8_marlin_linear(
46
+ input=x,
47
+ weight=layer.weight,
48
+ weight_scale=layer.weight_scale,
49
+ workspace=layer.workspace,
50
+ size_n=layer.output_size_per_partition,
51
+ size_k=layer.input_size_per_partition,
52
+ bias=bias,
53
+ )
vllm/model_executor/layers/quantization/fp8.py
+4
-4
@@ -517,10 +517,10 @@ class Fp8OnlineLinearMethod(Fp8LinearMethod):
517
518
# TODO: remove this check once the following RFC is resolved.
519
# https://github.com/vllm-project/vllm/issues/33314
520
- # This check is required because Mxfp8OnlineLinearMethod inherits from
521
- # Fp8OnlineLinearMethod but only calls super().create_weights(), so we must
522
- # skip the fp8_linear kernel creation.
523
- if hasattr(self, "mxfp8_linear"):
520
+ # Subclasses (e.g. Mxfp8OnlineLinearMethod) only need the weight
521
+ # registration above and manage their own kernel, so skip fp8_linear
522
+ # kernel creation for them.
523
+ if type(self) is not Fp8OnlineLinearMethod:
524
return
525
526
self.fp8_linear = init_fp8_linear_kernel(
vllm/model_executor/layers/quantization/modelopt.py
+4
-13
@@ -12,6 +12,7 @@ from vllm.config import get_current_vllm_config
12
from vllm.logger import init_logger
13
from vllm.model_executor.kernels.linear import (
14
init_fp8_linear_kernel,
15
+ init_mxfp8_linear_kernel,
16
init_nvfp4_linear_kernel,
17
)
18
from vllm.model_executor.layers.attention import Attention, MLAAttention
@@ -70,7 +71,6 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
71
MXFP8_BLOCK_SIZE,
72
MXFP8_SCALE_DTYPE,
73
MXFP8_VALUE_DTYPE,
73
- Mxfp8LinearOp,
74
mxfp8_e4m3_quantize,
75
)
76
from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -1576,7 +1576,7 @@ class ModelOptMxFp8LinearMethod(LinearMethodBase):
1576
"Dynamic quantization is not supported."
1577
)
1578
1579
- self.mxfp8_linear_op = Mxfp8LinearOp()
1579
+ self.kernel = init_mxfp8_linear_kernel()
1580
1581
def create_weights(
1582
self,
@@ -1658,7 +1658,7 @@ class ModelOptMxFp8LinearMethod(LinearMethodBase):
1658
f" got {layer.weight_scale.dtype}"
1659
)
1660
1661
- self.mxfp8_linear_op.process_weights(layer)
1661
+ self.kernel.process_weights_after_loading(layer)
1662
1663
def apply(
1664
self,
@@ -1666,16 +1666,7 @@ class ModelOptMxFp8LinearMethod(LinearMethodBase):
1666
x: torch.Tensor,
1667
bias: torch.Tensor | None = None,
1668
) -> torch.Tensor:
1669
- return self.mxfp8_linear_op.apply(
1670
- input=x,
1671
- weight=layer.weight,
1672
- weight_scale=layer.weight_scale,
1673
- out_dtype=x.dtype,
1674
- bias=bias,
1675
- workspace=getattr(layer, "workspace", None),
1676
- size_n=layer.output_size_per_partition,
1677
- size_k=layer.input_size_per_partition,
1678
- )
1669
+ return self.kernel.apply_weights(layer, x, bias)
1670
1671
1672
class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
vllm/model_executor/layers/quantization/mxfp8.py
+4
-14
@@ -9,6 +9,7 @@ import torch
9
from torch.nn import Module
10
11
from vllm.logger import init_logger
12
+from vllm.model_executor.kernels.linear import init_mxfp8_linear_kernel
13
from vllm.model_executor.layers.attention import Attention
14
from vllm.model_executor.layers.fused_moe import (
15
FusedMoE,
@@ -34,7 +35,6 @@ from vllm.model_executor.layers.quantization.fp8 import (
35
)
36
from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
37
MXFP8_BLOCK_SIZE,
37
- Mxfp8LinearOp,
38
mxfp8_e4m3_quantize,
39
)
40
from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -126,8 +126,7 @@ class Mxfp8OnlineLinearMethod(Fp8OnlineLinearMethod):
126
127
def __init__(self, quant_config: "Mxfp8Config"):
128
self.quant_config = quant_config
129
- self.out_dtype = torch.get_default_dtype()
130
- self.mxfp8_linear = Mxfp8LinearOp()
129
+ self.kernel = init_mxfp8_linear_kernel()
130
131
def create_weights(
132
self,
@@ -166,7 +165,7 @@ class Mxfp8OnlineLinearMethod(Fp8OnlineLinearMethod):
165
replace_parameter(layer, "weight", weight_fp8.data)
166
replace_parameter(layer, "weight_scale", weight_scale.data)
167
169
- self.mxfp8_linear.process_weights(layer)
168
+ self.kernel.process_weights_after_loading(layer)
169
170
layer._already_called_process_weights_after_loading = True
171
@@ -176,16 +175,7 @@ class Mxfp8OnlineLinearMethod(Fp8OnlineLinearMethod):
175
x: torch.Tensor,
176
bias: torch.Tensor | None = None,
177
) -> torch.Tensor:
179
- return self.mxfp8_linear.apply(
180
- input=x,
181
- weight=layer.weight,
182
- weight_scale=layer.weight_scale,
183
- out_dtype=self.out_dtype,
184
- bias=bias,
185
- workspace=getattr(layer, "workspace", None),
186
- size_n=layer.output_size_per_partition,
187
- size_k=layer.input_size_per_partition,
188
- )
178
+ return self.kernel.apply_weights(layer, x, bias)
179
180
181
class Mxfp8OnlineMoEMethod(Fp8OnlineMoEMethod):
vllm/model_executor/layers/quantization/utils/mxfp8_utils.py
-227
@@ -1,52 +1,16 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from enum import Enum
5
-
4
import torch
7
-from torch.nn.parameter import Parameter
5
9
-from vllm.logger import init_logger
10
-from vllm.utils import flashinfer as vllm_flashinfer
6
from vllm.utils.torch_utils import direct_register_custom_op
7
13
-logger = init_logger(__name__)
14
-
15
-
16
-class Mxfp8LinearBackend(Enum):
17
- EMULATION = "emulation"
18
- FLASHINFER_CUTLASS = "flashinfer-cutlass"
19
- MARLIN = "marlin"
20
-
21
-
8
# MXFP8 constants
9
MXFP8_VALUE_DTYPE = torch.float8_e4m3fn
10
MXFP8_SCALE_DTYPE = torch.uint8
11
MXFP8_BLOCK_SIZE = 32
12
13
28
-def select_mxfp8_linear_backend() -> Mxfp8LinearBackend:
29
- """Select the best MXFP8 linear backend for the current device.
30
-
31
- - SM100+ (Blackwell): FLASHINFER_CUTLASS (native MXFP8 W8A8 GEMM)
32
- - SM80+ (Ampere/Ada): MARLIN (MXFP8 W8A16 GEMM)
33
- - Otherwise: EMULATION (dequant to BF16 fallback)
34
- """
35
- from vllm.platforms import current_platform
36
-
37
- if current_platform.has_device_capability(100):
38
- return Mxfp8LinearBackend.FLASHINFER_CUTLASS
39
-
40
- from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
41
- is_fp8_marlin_supported,
42
- )
43
-
44
- if is_fp8_marlin_supported():
45
- return Mxfp8LinearBackend.MARLIN
46
-
47
- return Mxfp8LinearBackend.EMULATION
48
-
49
-
14
def swizzle_mxfp8_scale(sf: torch.Tensor, M: int, K: int) -> torch.Tensor:
15
"""Swizzle MXFP8 scales from row-major 2D to F8_128x4 layout."""
16
scaling_vector_size = MXFP8_BLOCK_SIZE # 32 for MXFP8
@@ -209,194 +173,3 @@ def xpu_mxfp8_quantize(
173
x: torch.Tensor, dtype: torch.dtype | None = None
174
) -> tuple[torch.Tensor, torch.Tensor]:
175
return torch.ops.vllm.xpu_mxfp8_quantize(x, dtype)
212
-
213
-
214
-class Mxfp8LinearOp:
215
- def __init__(self):
216
- self.backend = select_mxfp8_linear_backend()
217
- logger.info_once("Using %s backend for MXFP8 GEMM", self.backend)
218
-
219
- def process_weights(self, layer: torch.nn.Module) -> None:
220
- """Process MXFP8 weights after loading into backend-specific format."""
221
- if self.backend == Mxfp8LinearBackend.MARLIN:
222
- self._process_weights_marlin(layer)
223
- elif self.backend == Mxfp8LinearBackend.FLASHINFER_CUTLASS:
224
- self._process_weights_flashinfer_cutlass(layer)
225
- else:
226
- self._process_weights_emulation(layer)
227
-
228
- def _process_weights_emulation(self, layer: torch.nn.Module) -> None:
229
- """Keep scales as 2D uint8 for dequant-to-BF16 emulation."""
230
- weight = layer.weight.data # [N, K]
231
- N, K = weight.shape
232
- scale_k = K // MXFP8_BLOCK_SIZE
233
-
234
- weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous()
235
-
236
- layer.weight = Parameter(weight.contiguous(), requires_grad=False)
237
- layer.weight_scale = Parameter(weight_scale, requires_grad=False)
238
-
239
- def _process_weights_flashinfer_cutlass(self, layer: torch.nn.Module) -> None:
240
- """Swizzle scales to F8_128x4 layout for flashinfer CUTLASS."""
241
- weight = layer.weight.data # [N, K]
242
- N, K = weight.shape
243
-
244
- scale_k = K // MXFP8_BLOCK_SIZE
245
- weight_scale_2d = layer.weight_scale.data[:N, :scale_k].contiguous()
246
- weight_scale_swizzled = swizzle_mxfp8_scale(weight_scale_2d, M=N, K=K)
247
-
248
- layer.weight = Parameter(weight.contiguous(), requires_grad=False)
249
- layer.weight_scale = Parameter(
250
- weight_scale_swizzled.contiguous(), requires_grad=False
251
- )
252
-
253
- def _process_weights_marlin(self, layer: torch.nn.Module) -> None:
254
- """Repack MXFP8 weights and scales into Marlin kernel format."""
255
- from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
256
- prepare_mxfp8_layer_for_marlin,
257
- )
258
-
259
- prepare_mxfp8_layer_for_marlin(layer)
260
-
261
- def _apply_emulation(
262
- self,
263
- input: torch.Tensor,
264
- weight: torch.Tensor,
265
- weight_scale: torch.Tensor,
266
- out_dtype: torch.dtype,
267
- bias: torch.Tensor | None = None,
268
- ) -> torch.Tensor:
269
- if weight_scale.dtype != MXFP8_SCALE_DTYPE:
270
- raise ValueError(
271
- f"TORCH backend requires {MXFP8_SCALE_DTYPE} weight_scale dtype, "
272
- f"got {weight_scale.dtype}."
273
- )
274
- if weight_scale.ndim != 2:
275
- raise ValueError(
276
- f"TORCH backend requires 2D weight_scale, got {weight_scale.ndim}D. "
277
- f"Ensure process_weights_after_loading was called."
278
- )
279
-
280
- weight_bf16 = dequant_mxfp8_to_bf16(weight, weight_scale)
281
-
282
- output = torch.nn.functional.linear(input, weight_bf16, bias)
283
- return output.to(out_dtype)
284
-
285
- def _apply_flashinfer_cutlass(
286
- self,
287
- input: torch.Tensor,
288
- weight: torch.Tensor,
289
- weight_scale: torch.Tensor,
290
- out_dtype: torch.dtype,
291
- bias: torch.Tensor | None = None,
292
- ) -> torch.Tensor:
293
- N, K = weight.shape
294
-
295
- input_shape = input.shape
296
- input_2d = input.view(-1, K)
297
- M_orig = input_2d.shape[0]
298
-
299
- # Minimum dimension size for F8_128x4 block scaling layout
300
- min_dim = 128
301
-
302
- assert min_dim <= K, (
303
- f"mm_mxfp8 requires K >= {min_dim}, got K={K}. "
304
- f"in_features is too small for mm_mxfp8."
305
- )
306
- assert K % MXFP8_BLOCK_SIZE == 0, (
307
- f"mm_mxfp8 requires K to be divisible by {MXFP8_BLOCK_SIZE}, got K={K}."
308
- )
309
- assert min_dim <= N, (
310
- f"mm_mxfp8 requires N >= {min_dim}, got N={N}. "
311
- f"out_features is too small for mm_mxfp8."
312
- )
313
-
314
- M_padded = ((M_orig + min_dim - 1) // min_dim) * min_dim
315
- if M_padded != M_orig:
316
- pad_rows = M_padded - M_orig
317
- input_2d = torch.nn.functional.pad(input_2d, (0, 0, 0, pad_rows))
318
-
319
- input_mxfp8, input_scale = mxfp8_e4m3_quantize(
320
- input_2d,
321
- is_sf_swizzled_layout=True, # Swizzled for best accuracy
322
- )
323
-
324
- if not weight.is_contiguous():
325
- weight = weight.contiguous()
326
-
327
- output = vllm_flashinfer.mm_mxfp8(
328
- input_mxfp8,
329
- weight.t(),
330
- input_scale,
331
- weight_scale,
332
- out_dtype=out_dtype,
333
- backend="cutlass",
334
- )
335
-
336
- if M_padded != M_orig:
337
- output = output[:M_orig, :]
338
-
339
- if bias is not None:
340
- output = output + bias
341
-
342
- output_shape = (*input_shape[:-1], N)
343
- return output.view(output_shape)
344
-
345
- def _apply_marlin(
346
- self,
347
- input: torch.Tensor,
348
- weight: torch.Tensor,
349
- weight_scale: torch.Tensor,
350
- out_dtype: torch.dtype,
351
- bias: torch.Tensor | None = None,
352
- *,
353
- workspace: torch.Tensor,
354
- size_n: int,
355
- size_k: int,
356
- ) -> torch.Tensor:
357
- from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
358
- apply_mxfp8_marlin_linear,
359
- )
360
-
361
- return apply_mxfp8_marlin_linear(
362
- input=input,
363
- weight=weight,
364
- weight_scale=weight_scale,
365
- workspace=workspace,
366
- size_n=size_n,
367
- size_k=size_k,
368
- bias=bias,
369
- )
370
-
371
- def apply(
372
- self,
373
- input: torch.Tensor,
374
- weight: torch.Tensor,
375
- weight_scale: torch.Tensor,
376
- out_dtype: torch.dtype,
377
- bias: torch.Tensor | None = None,
378
- *,
379
- workspace: torch.Tensor | None = None,
380
- size_n: int = 0,
381
- size_k: int = 0,
382
- ) -> torch.Tensor:
383
- if self.backend == Mxfp8LinearBackend.EMULATION:
384
- return self._apply_emulation(input, weight, weight_scale, out_dtype, bias)
385
-
386
- if self.backend == Mxfp8LinearBackend.MARLIN:
387
- assert workspace is not None
388
- return self._apply_marlin(
389
- input,
390
- weight,
391
- weight_scale,
392
- out_dtype,
393
- bias,
394
- workspace=workspace,
395
- size_n=size_n,
396
- size_k=size_k,
397
- )
398
-
399
- assert self.backend == Mxfp8LinearBackend.FLASHINFER_CUTLASS
400
- return self._apply_flashinfer_cutlass(
401
- input, weight, weight_scale, out_dtype, bias
402
- )