From ce2aecc4dc42151d63f65d51c675edf28d40d671 Mon Sep 17 00:00:00 2001 From: JooHo Lee <96564470+BWAAEEEK@users.noreply.github.com> Date: Fri, 17 Jul 2026 22:53:48 +0900 Subject: [PATCH] [Performance] Use CuTe-DSL for FlashInfer MXFP4 quantization (#48417) Signed-off-by: BWAAEEEK --- vllm/model_executor/kernels/linear/mxfp4/flashinfer.py | 4 +++- vllm/utils/flashinfer.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/kernels/linear/mxfp4/flashinfer.py b/vllm/model_executor/kernels/linear/mxfp4/flashinfer.py index c0a5c86b0af..532a7d55e92 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/flashinfer.py +++ b/vllm/model_executor/kernels/linear/mxfp4/flashinfer.py @@ -56,7 +56,9 @@ class FlashInferMxFp4LinearKernel(MxFp4LinearKernel): out_shape = x.shape[:-1] + (layer.output_size_per_partition,) x_2d = x.reshape(-1, x.shape[-1]) - x_fp4, x_scale = flashinfer_mxfp4_quantize(x_2d.contiguous()) + x_fp4, x_scale = flashinfer_mxfp4_quantize( + x_2d.contiguous(), backend="cute-dsl" + ) out = flashinfer_scaled_fp4_mm( x_fp4, weight, diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 1334d110b49..8b999c84783 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -618,14 +618,16 @@ if has_flashinfer(): ) def flashinfer_mxfp4_quantize( a: torch.Tensor, + backend: str, ) -> tuple[torch.Tensor, torch.Tensor]: from flashinfer import mxfp4_quantize as _mxfp4_quantize - return _mxfp4_quantize(a) + return _mxfp4_quantize(a, backend=backend) @torch.library.register_fake("vllm::flashinfer_mxfp4_quantize") def flashinfer_mxfp4_quantize_fake( a: torch.Tensor, + backend: str, ) -> tuple[torch.Tensor, torch.Tensor]: m, k = a.shape sf_vec_size = 32