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 14, 23:45 UTC
Update flashinfer CUTLASS MoE Kernel (#21408)
Signed-off-by: Shu Wang. <shuw@nvidia.com>
Shu Wang committed
Jul 24, 2025 at 10:13 UTC
1b25f1fe757b0f0f877004f299f30dc079f81881
3 files changed
+8
-8
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py
+2
-2
@@ -11,7 +11,7 @@ from vllm.forward_context import get_forward_context
11
from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
12
from vllm.model_executor.layers.fused_moe.utils import (
13
extract_required_args, moe_kernel_quantize_input)
14
-from vllm.utils.flashinfer import block_scale_interleave
14
+from vllm.utils.flashinfer import nvfp4_block_scale_interleave
15
16
17
def get_local_sizes(local_tokens):
@@ -92,7 +92,7 @@ class FlashInferCutlassMoEPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
92
dim=0,
93
sizes=get_local_sizes(local_tokens))
94
a1_m, a1_n = a1q.shape
95
- a1q_scale = block_scale_interleave(a1q_scale)
95
+ a1q_scale = nvfp4_block_scale_interleave(a1q_scale)
96
97
return a1q, a1q_scale, None, topk_ids, topk_weights
98
vllm/model_executor/layers/quantization/modelopt.py
+2
-2
@@ -1254,8 +1254,8 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1254
x, layer.w13_weight, layer.w2_weight), (
1255
"Flashinfer CUTLASS Fused MoE not applicable!")
1256
1257
- a1_gscale = torch.min(layer.w13_input_scale_quant)
1258
- a2_gscale = torch.min(layer.w2_input_scale_quant)
1257
+ a1_gscale = layer.w13_input_scale_quant
1258
+ a2_gscale = layer.w2_input_scale_quant
1259
extra_expert_args = {
1260
'g1_alphas': layer.g1_alphas,
1261
'g2_alphas': layer.g2_alphas,
vllm/utils/flashinfer.py
+4
-4
@@ -69,8 +69,8 @@ flashinfer_trtllm_fp8_block_scale_moe = _lazy_import_wrapper(
69
flashinfer_cutlass_fused_moe = _lazy_import_wrapper("flashinfer.fused_moe",
70
"cutlass_fused_moe")
71
fp4_quantize = _lazy_import_wrapper("flashinfer", "fp4_quantize")
72
-block_scale_interleave = _lazy_import_wrapper("flashinfer",
73
- "block_scale_interleave")
72
+nvfp4_block_scale_interleave = _lazy_import_wrapper(
73
+ "flashinfer", "nvfp4_block_scale_interleave")
74
75
# Special case for autotune since it returns a context manager
76
autotune = _lazy_import_wrapper(
@@ -95,7 +95,7 @@ def has_flashinfer_cutlass_fused_moe() -> bool:
95
required_functions = [
96
("flashinfer.fused_moe", "cutlass_fused_moe"),
97
("flashinfer", "fp4_quantize"),
98
- ("flashinfer", "block_scale_interleave"),
98
+ ("flashinfer", "nvfp4_block_scale_interleave"),
99
]
100
101
for module_name, attr_name in required_functions:
@@ -110,7 +110,7 @@ __all__ = [
110
"flashinfer_trtllm_fp8_block_scale_moe",
111
"flashinfer_cutlass_fused_moe",
112
"fp4_quantize",
113
- "block_scale_interleave",
113
+ "nvfp4_block_scale_interleave",
114
"autotune",
115
"has_flashinfer_moe",
116
"has_flashinfer_cutlass_fused_moe",