[Kernel] Add FlashInferCutedslMxfp8LinearKernel (cute-dsl mm_mxfp8) (#46393)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Yongye Zhu committed
Jun 23, 2026 at 12:45 UTC
11b56b2ff28eae32a1e65c051c24b50f3c39e03d
3 files changed
+87
-1
vllm/config/kernel.py
+1
-1
@@ -199,7 +199,7 @@ class KernelConfig:
199
- "auto": Automatically select the best backend based on model and hardware
200
- "cutlass": Use CUTLASS-based kernels
201
- "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels
202
- - "flashinfer_cutedsl": Use FlashInfer with CuteDSL kernels
202
+ - "flashinfer_cutedsl": Use FlashInfer with CuTe-DSL kernels (NVFP4, MXFP8)
203
- "flashinfer_trtllm": Use FlashInfer with TensorRT-LLM kernels
204
- "flashinfer_cudnn": Use FlashInfer with cuDNN kernels
205
- "flashinfer_b12x": Use FlashInfer b12x CuteDSL NVFP4 GEMM (SM120+)
vllm/model_executor/kernels/linear/__init__.py
+4
@@ -88,6 +88,7 @@ from vllm.model_executor.kernels.linear.mxfp8.emulation import (
88
EmulationMxfp8LinearKernel,
89
)
90
from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
91
+ FlashInferCutedslMxfp8LinearKernel,
92
FlashInferCutlassMxfp8LinearKernel,
93
)
94
from vllm.model_executor.kernels.linear.mxfp8.marlin import (
@@ -212,6 +213,7 @@ _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = {
213
},
214
"flashinfer_cutedsl": {
215
FlashInferCuteDslNvFp4LinearKernel,
216
+ FlashInferCutedslMxfp8LinearKernel,
217
},
218
"flashinfer_trtllm": {
219
FlashInferTrtllmNvFp4LinearKernel,
@@ -385,6 +387,7 @@ _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = {
387
# in priority/performance order (when available)
388
_POSSIBLE_MXFP8_KERNELS: dict[PlatformEnum, list[type[Mxfp8LinearKernel]]] = {
389
PlatformEnum.CUDA: [
390
+ FlashInferCutedslMxfp8LinearKernel,
391
FlashInferCutlassMxfp8LinearKernel,
392
MarlinMxfp8LinearKernel,
393
EmulationMxfp8LinearKernel,
@@ -1036,6 +1039,7 @@ __all__ = [
1039
"MxFp4LinearLayerConfig",
1040
"FlashInferMxFp4LinearKernel",
1041
"MarlinMxFp4LinearKernel",
1042
+ "FlashInferCutedslMxfp8LinearKernel",
1043
"FlashInferCutlassMxfp8LinearKernel",
1044
"MarlinMxfp8LinearKernel",
1045
"XPUMxFp8LinearKernel",
vllm/model_executor/kernels/linear/mxfp8/flashinfer.py
+82
@@ -11,6 +11,7 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
11
)
12
from vllm.platforms import current_platform
13
from vllm.utils import flashinfer as vllm_flashinfer
14
+from vllm.utils.flashinfer import has_flashinfer_cutedsl
15
16
from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
17
@@ -91,3 +92,84 @@ class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel):
92
93
output_shape = (*input_shape[:-1], N)
94
return output.view(output_shape)
95
+
96
+
97
+class FlashInferCutedslMxfp8LinearKernel(Mxfp8LinearKernel):
98
+ """MXFP8 W8A8 GEMM via FlashInfer CuTe-DSL (SM100/SM103)."""
99
+
100
+ @classmethod
101
+ def is_supported(
102
+ cls, compute_capability: int | None = None
103
+ ) -> tuple[bool, str | None]:
104
+ if not (
105
+ current_platform.is_cuda()
106
+ and current_platform.is_device_capability_family(100)
107
+ ):
108
+ return False, "requires sm_100/sm_103 (Blackwell)"
109
+ if not has_flashinfer_cutedsl():
110
+ return False, "requires FlashInfer CuTe-DSL module"
111
+ return True, None
112
+
113
+ @classmethod
114
+ def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
115
+ return True, None
116
+
117
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
118
+ weight = layer.weight.data # [N, K]
119
+ N, K = weight.shape
120
+
121
+ scale_k = K // MXFP8_BLOCK_SIZE
122
+ weight_scale_2d = layer.weight_scale.data[:N, :scale_k].contiguous()
123
+ weight_scale_swizzled = swizzle_mxfp8_scale(weight_scale_2d, M=N, K=K)
124
+
125
+ # Store weight column-major [K, N] as mm_mxfp8 expects for operand B.
126
+ layer.weight = Parameter(weight.contiguous().t(), requires_grad=False)
127
+ layer.weight_scale = Parameter(
128
+ weight_scale_swizzled.contiguous(), requires_grad=False
129
+ )
130
+
131
+ def apply_weights(
132
+ self,
133
+ layer: torch.nn.Module,
134
+ x: torch.Tensor,
135
+ bias: torch.Tensor | None = None,
136
+ ) -> torch.Tensor:
137
+ weight = layer.weight # [K, N], column-major
138
+ weight_scale = layer.weight_scale
139
+ out_dtype = x.dtype
140
+ K, N = weight.shape
141
+
142
+ input_shape = x.shape
143
+ input_2d = x.view(-1, K)
144
+ min_dim = 128
145
+
146
+ assert min_dim <= K, (
147
+ f"mm_mxfp8 requires K >= {min_dim}, got K={K}. "
148
+ f"in_features is too small for mm_mxfp8."
149
+ )
150
+ assert K % MXFP8_BLOCK_SIZE == 0, (
151
+ f"mm_mxfp8 requires K to be divisible by {MXFP8_BLOCK_SIZE}, got K={K}."
152
+ )
153
+ assert min_dim <= N, (
154
+ f"mm_mxfp8 requires N >= {min_dim}, got N={N}. "
155
+ f"out_features is too small for mm_mxfp8."
156
+ )
157
+
158
+ input_mxfp8, input_scale = mxfp8_e4m3_quantize(
159
+ input_2d, is_sf_swizzled_layout=True
160
+ )
161
+
162
+ output = vllm_flashinfer.mm_mxfp8(
163
+ input_mxfp8,
164
+ weight,
165
+ input_scale,
166
+ weight_scale,
167
+ out_dtype=out_dtype,
168
+ backend="cute-dsl",
169
+ )
170
+
171
+ if bias is not None:
172
+ output = output + bias
173
+
174
+ output_shape = (*input_shape[:-1], N)
175
+ return output.view(output_shape)