[Kernel] Add FlashInfer TRTLLM MXFP8 linear backend (#52204)
Signed-off-by: seonjinn <sna@nvidia.com> Signed-off-by: Misha Goin <mgoin64@gmail.com> Co-authored-by: Misha Goin <mgoin64@gmail.com>
Seonjin committed
Aug 20, 2026 at 07:07 UTC
bd8865a299c4a68cff9b6443b9fd795f4c4735f6
5 files changed
+265
docs/features/quantization/modelopt.md
+5
@@ -33,6 +33,11 @@ following `quantization.quant_algo` values:
33
[Engine Arguments](../../configuration/engine_args.md) page and shown by
34
`vllm serve --help=KernelConfig`.
35
36
+!!! note
37
+ For models quantized to MXFP8 with BF16 activations on SM100-family GPUs,
38
+ use `--linear-backend flashinfer_trtllm` to select FlashInfer's TensorRT-LLM
39
+ GEMM backend.
40
+
41
## Quantizing HuggingFace Models with PTQ
42
43
You can quantize HuggingFace models using the example scripts provided in the Model Optimizer repository. The primary script for LLM PTQ is typically found within the `examples/llm_ptq` directory.
tests/kernels/quantization/test_flashinfer_mxfp8_trtllm.py
new
+123
@@ -0,0 +1,123 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+import torch.nn.functional as F
7
+
8
+from vllm.model_executor.kernels.linear import (
9
+ FlashInferTrtllmMxfp8LinearKernel,
10
+ Mxfp8LinearLayerConfig,
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
15
+
16
+if not (
17
+ current_platform.is_cuda()
18
+ and current_platform.is_device_capability_family(100)
19
+ and has_flashinfer()
20
+):
21
+ pytest.skip(
22
+ reason="FlashInfer TRTLLM MXFP8 requires an SM100-family GPU",
23
+ allow_module_level=True,
24
+ )
25
+
26
+
27
+def _make_layer(weight: torch.Tensor) -> torch.nn.Module:
28
+ from flashinfer import SfLayout, mxfp8_quantize
29
+
30
+ weight_mxfp8, weight_scale = mxfp8_quantize(
31
+ weight,
32
+ sf_swizzle_layout=SfLayout.layout_linear,
33
+ )
34
+ layer = torch.nn.Module()
35
+ layer.weight = torch.nn.Parameter(weight_mxfp8, requires_grad=False)
36
+ layer.weight_scale = torch.nn.Parameter(
37
+ weight_scale.view(weight.shape[0], weight.shape[1] // 32),
38
+ requires_grad=False,
39
+ )
40
+ return layer
41
+
42
+
43
+@pytest.mark.parametrize("shape", [(1, 130, 256), (7, 256, 512), (128, 130, 768)])
44
+@torch.inference_mode()
45
+def test_flashinfer_trtllm_mxfp8_linear_numerics(
46
+ shape: tuple[int, int, int],
47
+) -> None:
48
+ torch.manual_seed(0)
49
+ m, n, k = shape
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())
54
+ kernel.process_weights_after_loading(layer)
55
+
56
+ output = kernel.apply_weights(layer, x)
57
+ reference = torch.mm(x, weight.t())
58
+ similarity = F.cosine_similarity(
59
+ output.float().flatten(), reference.float().flatten(), dim=0
60
+ )
61
+
62
+ assert output.shape == (m, n)
63
+ assert output.is_contiguous()
64
+ assert similarity.item() > 0.98
65
+
66
+
67
+@torch.inference_mode()
68
+def test_flashinfer_trtllm_mxfp8_custom_ops() -> None:
69
+ x = torch.randn((7, 512), dtype=torch.bfloat16, device="cuda")
70
+ weight = torch.randn((256, 512), dtype=torch.bfloat16, device="cuda")
71
+ layer = _make_layer(weight)
72
+ kernel = FlashInferTrtllmMxfp8LinearKernel(Mxfp8LinearLayerConfig())
73
+ kernel.process_weights_after_loading(layer)
74
+
75
+ torch.library.opcheck(
76
+ torch.ops.vllm.flashinfer_mxfp8_quantize_8x4.default,
77
+ (x,),
78
+ )
79
+ x_mxfp8, x_scale = vllm_flashinfer.flashinfer_mxfp8_quantize_8x4(x)
80
+ # SchemaCheckMode compares inputs with allclose, which CUDA does not
81
+ # implement for float8. The numerical tests above guard input mutation.
82
+ torch.library.opcheck(
83
+ torch.ops.vllm.mm_mxfp8.default,
84
+ (
85
+ x_mxfp8,
86
+ layer.weight.t(),
87
+ x_scale,
88
+ layer.weight_scale,
89
+ torch.bfloat16,
90
+ "trtllm",
91
+ True,
92
+ ),
93
+ test_utils=(
94
+ "test_autograd_registration",
95
+ "test_faketensor",
96
+ "test_aot_dispatch_dynamic",
97
+ ),
98
+ )
99
+
100
+
101
+@torch.inference_mode()
102
+def test_flashinfer_trtllm_mxfp8_linear_cuda_graph() -> None:
103
+ torch.manual_seed(0)
104
+ m, n, k = 7, 130, 512
105
+ weight = torch.randn((n, k), dtype=torch.bfloat16, device="cuda")
106
+ layer = _make_layer(weight)
107
+ kernel = FlashInferTrtllmMxfp8LinearKernel(Mxfp8LinearLayerConfig())
108
+ kernel.process_weights_after_loading(layer)
109
+
110
+ static_x = torch.randn((m, k), dtype=torch.bfloat16, device="cuda")
111
+ kernel.apply_weights(layer, static_x)
112
+ torch.accelerator.synchronize()
113
+
114
+ graph = torch.cuda.CUDAGraph()
115
+ with torch.cuda.graph(graph):
116
+ graph_output = kernel.apply_weights(layer, static_x)
117
+
118
+ new_x = torch.randn_like(static_x)
119
+ static_x.copy_(new_x)
120
+ graph.replay()
121
+ eager_output = kernel.apply_weights(layer, new_x)
122
+
123
+ torch.testing.assert_close(graph_output, eager_output, rtol=0, atol=0)
vllm/model_executor/kernels/linear/__init__.py
+4
@@ -115,6 +115,7 @@ from vllm.model_executor.kernels.linear.mxfp8.emulation import (
115
from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
116
FlashInferCutedslMxfp8LinearKernel,
117
FlashInferCutlassMxfp8LinearKernel,
118
+ FlashInferTrtllmMxfp8LinearKernel,
119
)
120
from vllm.model_executor.kernels.linear.mxfp8.humming import (
121
HummingMxfp8LinearKernel,
@@ -267,6 +268,7 @@ _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = {
268
FlashInferCutedslMxfp8LinearKernel,
269
},
270
"flashinfer_trtllm": {
271
+ FlashInferTrtllmMxfp8LinearKernel,
272
FlashInferTrtllmNvFp4LinearKernel,
273
},
274
"flashinfer_cudnn": {
@@ -507,6 +509,7 @@ _POSSIBLE_MXFP8_KERNELS: dict[PlatformEnum, list[type[Mxfp8LinearKernel]]] = {
509
B12xMxfp8LinearKernel,
510
EmulationMxfp8LinearKernel,
511
HummingMxfp8LinearKernel,
512
+ FlashInferTrtllmMxfp8LinearKernel,
513
],
514
PlatformEnum.ROCM: [
515
# Native CDNA4 (gfx950) MX linear; is_supported() gates to gfx95x and
@@ -1216,6 +1219,7 @@ __all__ = [
1219
"MarlinMxFp4LinearKernel",
1220
"FlashInferCutedslMxfp8LinearKernel",
1221
"FlashInferCutlassMxfp8LinearKernel",
1222
+ "FlashInferTrtllmMxfp8LinearKernel",
1223
"MarlinMxfp8LinearKernel",
1224
"XPUMxFp8LinearKernel",
1225
"EmulationMxfp8LinearKernel",
vllm/model_executor/kernels/linear/mxfp8/flashinfer.py
+100
@@ -177,3 +177,103 @@ class FlashInferCutedslMxfp8LinearKernel(Mxfp8LinearKernel):
177
178
output_shape = (*input_shape[:-1], N)
179
return output.view(output_shape)
180
+
181
+
182
+class FlashInferTrtllmMxfp8LinearKernel(Mxfp8LinearKernel):
183
+ """MXFP8 W8A8 GEMM via FlashInfer's TensorRT-LLM wrapper."""
184
+
185
+ @classmethod
186
+ def is_supported(
187
+ cls, compute_capability: int | None = None
188
+ ) -> tuple[bool, str | None]:
189
+ if not (
190
+ current_platform.is_cuda()
191
+ and current_platform.is_device_capability_family(100)
192
+ ):
193
+ return False, "requires SM100-family GPU"
194
+ if not has_flashinfer():
195
+ return False, "requires FlashInfer"
196
+ return True, None
197
+
198
+ @classmethod
199
+ def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
200
+ return True, None
201
+
202
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
203
+ from flashinfer import shuffle_matrix_a, shuffle_matrix_sf_a
204
+
205
+ if hasattr(layer, "_mxfp8_trtllm_output_size") and layer.weight_scale.ndim == 1:
206
+ return
207
+
208
+ weight = layer.weight.data # [N, K]
209
+ N, K = weight.shape
210
+ if K % 256 != 0:
211
+ raise ValueError(
212
+ f"FlashInfer TRTLLM MXFP8 requires K to be divisible by 256, got K={K}."
213
+ )
214
+
215
+ scale_k = K // MXFP8_BLOCK_SIZE
216
+ weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous()
217
+ padded_n = ((N + 127) // 128) * 128
218
+ if padded_n != N:
219
+ padded_weight = weight.new_zeros((padded_n, K))
220
+ padded_weight[:N] = weight
221
+ weight = padded_weight
222
+
223
+ padded_scale = weight_scale.new_zeros((padded_n, scale_k))
224
+ padded_scale[:N] = weight_scale
225
+ weight_scale = padded_scale
226
+ else:
227
+ weight = weight.contiguous()
228
+
229
+ layer.weight = Parameter(
230
+ shuffle_matrix_a(weight, 128).reshape(padded_n, K),
231
+ requires_grad=False,
232
+ )
233
+ layer.weight_scale = Parameter(
234
+ shuffle_matrix_sf_a(
235
+ weight_scale,
236
+ 128,
237
+ num_elts_per_sf=MXFP8_BLOCK_SIZE,
238
+ ).reshape(-1),
239
+ requires_grad=False,
240
+ )
241
+ layer._mxfp8_trtllm_output_size = N
242
+
243
+ def apply_weights(
244
+ self,
245
+ layer: torch.nn.Module,
246
+ x: torch.Tensor,
247
+ bias: torch.Tensor | None = None,
248
+ ) -> torch.Tensor:
249
+ assert x.dtype == torch.bfloat16, (
250
+ f"FlashInfer TRTLLM MXFP8 requires bfloat16 activations, got {x.dtype}."
251
+ )
252
+
253
+ weight = layer.weight # shuffled [padded N, K]
254
+ weight_scale = layer.weight_scale
255
+ _, K = weight.shape
256
+ output_size = layer._mxfp8_trtllm_output_size
257
+ input_shape = x.shape
258
+ input_2d = x.view(-1, K)
259
+
260
+ input_mxfp8, input_scale = vllm_flashinfer.flashinfer_mxfp8_quantize_8x4(
261
+ input_2d
262
+ )
263
+ output = vllm_flashinfer.mm_mxfp8(
264
+ input_mxfp8,
265
+ weight.t(),
266
+ input_scale,
267
+ weight_scale,
268
+ out_dtype=x.dtype,
269
+ backend="trtllm",
270
+ use_8x4_sf_layout=True,
271
+ )
272
+ if output.shape[-1] != output_size:
273
+ output = output[:, :output_size].contiguous()
274
+
275
+ if bias is not None:
276
+ output = output + bias
277
+
278
+ output_shape = (*input_shape[:-1], output_size)
279
+ return output.view(output_shape)
vllm/utils/flashinfer.py
+33
@@ -726,6 +726,36 @@ if has_flashinfer():
726
rounded_m, rounded_n, dtype=torch.uint8, device=a.device
727
)
728
729
+ @torch.library.custom_op(
730
+ "vllm::flashinfer_mxfp8_quantize_8x4",
731
+ mutates_args=[],
732
+ device_types="cuda",
733
+ )
734
+ def flashinfer_mxfp8_quantize_8x4(
735
+ a: torch.Tensor,
736
+ ) -> tuple[torch.Tensor, torch.Tensor]:
737
+ from flashinfer import SfLayout
738
+ from flashinfer import mxfp8_quantize as mxfp8_quantize_
739
+
740
+ return mxfp8_quantize_(
741
+ a,
742
+ backend="cuda",
743
+ sf_swizzle_layout=SfLayout.layout_8x4,
744
+ )
745
+
746
+ @torch.library.register_fake(
747
+ "vllm::flashinfer_mxfp8_quantize_8x4",
748
+ )
749
+ def flashinfer_mxfp8_quantize_8x4_fake(
750
+ a: torch.Tensor,
751
+ ) -> tuple[torch.Tensor, torch.Tensor]:
752
+ m, k = a.shape
753
+ scale_size = cdiv(m, 8) * 8 * cdiv(k // 32, 4) * 4
754
+ return (
755
+ torch.empty(m, k, dtype=torch.float8_e4m3fn, device=a.device),
756
+ torch.empty(scale_size, dtype=torch.uint8, device=a.device),
757
+ )
758
+
759
@torch.library.custom_op(
760
"vllm::mm_mxfp8",
761
mutates_args=[],
@@ -738,6 +768,7 @@ if has_flashinfer():
768
B_scale: torch.Tensor,
769
out_dtype: torch.dtype,
770
backend: str = "cutlass",
771
+ use_8x4_sf_layout: bool = False,
772
) -> torch.Tensor:
773
from flashinfer import mm_mxfp8 as mm_mxfp8_
774
@@ -749,6 +780,7 @@ if has_flashinfer():
780
out=None,
781
out_dtype=out_dtype,
782
backend=backend,
783
+ use_8x4_sf_layout=use_8x4_sf_layout,
784
)
785
786
@torch.library.register_fake(
@@ -761,6 +793,7 @@ if has_flashinfer():
793
B_scale: torch.Tensor,
794
out_dtype: torch.dtype,
795
backend: str = "cutlass",
796
+ use_8x4_sf_layout: bool = False,
797
) -> torch.Tensor:
798
# A is [m, k], B is [k, n] -> output [m, n]
799
return torch.empty(A.shape[0], B.shape[1], dtype=out_dtype, device=A.device)