[Performance] Use CuTe-DSL for FlashInfer MXFP4 quantization (#48417)
Signed-off-by: BWAAEEEK <jooho414@gmail.com>
JooHo Lee committed
Jul 17, 2026 at 22:53 UTC
ce2aecc4dc42151d63f65d51c675edf28d40d671
2 files changed
+6
-2
vllm/model_executor/kernels/linear/mxfp4/flashinfer.py
+3
-1
@@ -56,7 +56,9 @@ class FlashInferMxFp4LinearKernel(MxFp4LinearKernel):
56
out_shape = x.shape[:-1] + (layer.output_size_per_partition,)
57
x_2d = x.reshape(-1, x.shape[-1])
58
59
- x_fp4, x_scale = flashinfer_mxfp4_quantize(x_2d.contiguous())
59
+ x_fp4, x_scale = flashinfer_mxfp4_quantize(
60
+ x_2d.contiguous(), backend="cute-dsl"
61
+ )
62
out = flashinfer_scaled_fp4_mm(
63
x_fp4,
64
weight,
vllm/utils/flashinfer.py
+3
-1
@@ -618,14 +618,16 @@ if has_flashinfer():
618
)
619
def flashinfer_mxfp4_quantize(
620
a: torch.Tensor,
621
+ backend: str,
622
) -> tuple[torch.Tensor, torch.Tensor]:
623
from flashinfer import mxfp4_quantize as _mxfp4_quantize
624
624
- return _mxfp4_quantize(a)
625
+ return _mxfp4_quantize(a, backend=backend)
626
627
@torch.library.register_fake("vllm::flashinfer_mxfp4_quantize")
628
def flashinfer_mxfp4_quantize_fake(
629
a: torch.Tensor,
630
+ backend: str,
631
) -> tuple[torch.Tensor, torch.Tensor]:
632
m, k = a.shape
633
sf_vec_size = 32