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 15, 01:00 UTC
Update fp4 quantize API (#21327)
Signed-off-by: Shu Wang <shuw@nvidia.com>
Shu Wang committed
Jul 22, 2025 at 01:40 UTC
9e23ad9655c890f24448551513c12dd611c7e071
3 files changed
+11
-11
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
+5
-5
@@ -181,12 +181,12 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
181
g2_alphas,
182
]
183
_ = flashinfer_cutlass_fused_moe(
184
- hidden_states,
185
- topk_ids.to(torch.int),
186
- topk_weights,
184
+ input=hidden_states,
185
+ token_selected_experts=topk_ids.to(torch.int),
186
+ token_final_scales=topk_weights,
187
# FlashInfer API requires weight to be long for nvfp4
188
- w1.view(torch.long),
189
- w2.view(torch.long),
188
+ fc1_expert_weights=w1.view(torch.long),
189
+ fc2_expert_weights=w2.view(torch.long),
190
output_dtype=out_dtype,
191
quant_scales=quant_scales,
192
input_sf=a1q_scale,
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 fp4_swizzle_blockscale
14
+from vllm.utils.flashinfer import 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 = fp4_swizzle_blockscale(a1q_scale, a1_m, a1_n * 2)
95
+ a1q_scale = block_scale_interleave(a1q_scale)
96
97
return a1q, a1q_scale, None, topk_ids, topk_weights
98
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
-fp4_swizzle_blockscale = _lazy_import_wrapper("flashinfer",
73
- "fp4_swizzle_blockscale")
72
+block_scale_interleave = _lazy_import_wrapper("flashinfer",
73
+ "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", "fp4_swizzle_blockscale"),
98
+ ("flashinfer", "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
- "fp4_swizzle_blockscale",
113
+ "block_scale_interleave",
114
"autotune",
115
"has_flashinfer_moe",
116
"has_flashinfer_cutlass_fused_moe",