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, 05:30 UTC
[NVIDIA] [Perf] Update to leverage flashinfer trtllm FP4 MOE throughput kernel (#26714)
Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
jiahanc committed
Oct 16, 2025 at 16:20 UTC
41d3071918bd6dcb6259bac8a617cfed45d42707
7 files changed
+25
-96
docker/Dockerfile
+2
-2
@@ -359,8 +359,8 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist
359
# Install FlashInfer pre-compiled kernel cache and binaries
360
# https://docs.flashinfer.ai/installation.html
361
RUN --mount=type=cache,target=/root/.cache/uv \
362
- uv pip install --system flashinfer-cubin==0.4.0 \
363
- && uv pip install --system flashinfer-jit-cache==0.4.0 \
362
+ uv pip install --system flashinfer-cubin==0.4.1 \
363
+ && uv pip install --system flashinfer-jit-cache==0.4.1 \
364
--extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
365
&& flashinfer show-config
366
docker/Dockerfile.nightly_torch
+2
-2
@@ -246,7 +246,7 @@ RUN pip install setuptools==75.6.0 packaging==23.2 ninja==1.11.1.3 build==1.2.2.
246
247
248
# build flashinfer for torch nightly from source around 10 mins
249
-# release version: v0.4.0
249
+# release version: v0.4.1
250
# todo(elainewy): cache flashinfer build result for faster build
251
ENV CCACHE_DIR=/root/.cache/ccache
252
RUN --mount=type=cache,target=/root/.cache/ccache \
@@ -254,7 +254,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
254
echo "git clone flashinfer..." \
255
&& git clone --recursive https://github.com/flashinfer-ai/flashinfer.git \
256
&& cd flashinfer \
257
- && git checkout v0.4.0 \
257
+ && git checkout v0.4.1\
258
&& git submodule update --init --recursive \
259
&& echo "finish git clone flashinfer..." \
260
&& rm -rf build \
requirements/cuda.txt
+1
-1
@@ -12,4 +12,4 @@ torchvision==0.23.0 # Required for phi3v processor. See https://github.com/pytor
12
# https://github.com/facebookresearch/xformers/releases/tag/v0.0.32.post1
13
xformers==0.0.32.post1; platform_system == 'Linux' and platform_machine == 'x86_64' # Requires PyTorch >= 2.8
14
# FlashInfer should be updated together with the Dockerfile
15
-flashinfer-python==0.4.0
\ No newline at end of file
15
+flashinfer-python==0.4.1
\ No newline at end of file
tests/kernels/moe/test_ocp_mx_moe.py
+7
-7
@@ -37,7 +37,7 @@ if TRTLLM_GEN_MXFP4_AVAILABLE:
37
trtllm_fp4_block_scale_moe,
38
)
39
from flashinfer.fp4_quantization import nvfp4_block_scale_interleave
40
- from flashinfer.fused_moe.core import _maybe_get_cached_w2_permute_indices
40
+ from flashinfer.fused_moe.core import get_w2_permute_indices_with_cache
41
42
43
@dataclass
@@ -319,7 +319,7 @@ def tg_mxfp4_moe(
319
if transpose_optimized:
320
for i in range(num_experts):
321
# w13 weight shuffling
322
- permute_indices = _maybe_get_cached_w2_permute_indices(
322
+ permute_indices = get_w2_permute_indices_with_cache(
323
_cache_permute_indices,
324
w13_weight[i].view(torch.uint8),
325
epilogue_tile_m,
@@ -330,7 +330,7 @@ def tg_mxfp4_moe(
330
.contiguous()
331
)
332
# w13 scale shuffling
333
- permute_sf_indices = _maybe_get_cached_w2_permute_indices(
333
+ permute_sf_indices = get_w2_permute_indices_with_cache(
334
_cache_permute_indices,
335
w13_weight_scale[i].view(torch.uint8),
336
epilogue_tile_m,
@@ -344,7 +344,7 @@ def tg_mxfp4_moe(
344
)
345
)
346
# w13 bias shuffling
347
- permute_bias_indices = _maybe_get_cached_w2_permute_indices(
347
+ permute_bias_indices = get_w2_permute_indices_with_cache(
348
_cache_permute_indices,
349
w13_bias[i].clone().reshape(-1, 1),
350
epilogue_tile_m,
@@ -356,7 +356,7 @@ def tg_mxfp4_moe(
356
.contiguous()
357
)
358
# w2 weight shuffling
359
- permute_indices = _maybe_get_cached_w2_permute_indices(
359
+ permute_indices = get_w2_permute_indices_with_cache(
360
_cache_permute_indices,
361
w2_weight[i].view(torch.uint8),
362
epilogue_tile_m,
@@ -367,7 +367,7 @@ def tg_mxfp4_moe(
367
.contiguous()
368
)
369
# w2 scale shuffling
370
- permute_sf_indices = _maybe_get_cached_w2_permute_indices(
370
+ permute_sf_indices = get_w2_permute_indices_with_cache(
371
_cache_permute_indices,
372
w2_weight_scale[i].view(torch.uint8),
373
epilogue_tile_m,
@@ -381,7 +381,7 @@ def tg_mxfp4_moe(
381
)
382
)
383
# w2 bias shuffling
384
- permute_indices = _maybe_get_cached_w2_permute_indices(
384
+ permute_indices = get_w2_permute_indices_with_cache(
385
_cache_permute_indices,
386
w2_bias[i].clone().reshape(-1, 1),
387
epilogue_tile_m,
vllm/model_executor/layers/fused_moe/trtllm_moe.py
+1
-28
@@ -11,7 +11,6 @@ from vllm.model_executor.layers.fused_moe.config import (
11
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
12
TopKWeightAndReduceNoOP,
13
)
14
-from vllm.utils import next_power_of_2
14
15
16
class TrtLlmGenExperts(mk.FusedMoEPermuteExpertsUnpermute):
@@ -65,30 +64,6 @@ class TrtLlmGenExperts(mk.FusedMoEPermuteExpertsUnpermute):
64
output = (M, K)
65
return (workspace1, workspace2, output)
66
68
- def _get_tile_tokens_dim(self, x: torch.Tensor, top_k: int, local_num_experts: int):
69
- # Number of tokens in the input tensor.
70
- num_tokens = x.shape[0]
71
- # Factor to account for the imbalance of the experts.
72
- # factor equals to the
73
- # max_real_num_tokens_per_expert / perfect_num_tokens_per_expert
74
- # 1.0 means perfect expert distribution.
75
- # > 1.0 means some experts have more tokens than the perfect
76
- # distribution.
77
- # < 1.0 does not make sense.
78
- imbalance_factor = 1.3
79
- # Calculate the number of tokens per expert assuming perfect
80
- # distribution.
81
- num_tokens_per_expert = (num_tokens * top_k) // local_num_experts
82
- # Apply the imbalance factor.
83
- num_tokens_per_expert = int(num_tokens_per_expert * imbalance_factor)
84
- # And pad the number to the next power of 2.
85
- tile_tokens_dim = next_power_of_2(num_tokens_per_expert)
86
- # Cap to 8-64 tokens per CTA tile as it's the range supported by the
87
- # kernel.
88
- tile_tokens_dim = min(max(tile_tokens_dim, 8), 64)
89
-
90
- return tile_tokens_dim
91
-
67
def apply(
68
self,
69
output: torch.Tensor,
@@ -148,9 +123,7 @@ class TrtLlmGenExperts(mk.FusedMoEPermuteExpertsUnpermute):
123
"local_expert_offset": local_expert_offset,
124
"local_num_experts": local_num_experts,
125
"routed_scaling_factor": None,
151
- "tile_tokens_dim": self._get_tile_tokens_dim(
152
- x_quant, topk, local_num_experts
153
- ),
126
+ "tile_tokens_dim": None,
127
"routing_method_type": 1,
128
"do_finalize": True,
129
"output": output,
vllm/model_executor/layers/quantization/modelopt.py
+4
-17
@@ -72,7 +72,6 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
72
)
73
from vllm.model_executor.parameter import ModelWeightParameter, PerTensorScaleParameter
74
from vllm.scalar_type import scalar_types
75
-from vllm.utils import next_power_of_2
75
from vllm.utils.flashinfer import (
76
flashinfer_scaled_fp4_mm,
77
has_flashinfer,
@@ -1125,16 +1124,6 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
1124
return out.view(*output_shape)
1125
1126
1128
-def _get_tile_tokens_dim(num_tokens: int, top_k: int, num_experts: int) -> int:
1129
- # Guess tokens per expert assuming perfect expert distribution first.
1130
- num_tokens_per_expert = (num_tokens * top_k) // num_experts
1131
- # And pad the number to the next power of 2.
1132
- tile_tokens_dim = next_power_of_2(num_tokens_per_expert)
1133
- # Cap to 8-64 tokens per CTA tile as it's the range supported by the kernel.
1134
- tile_tokens_dim = min(max(tile_tokens_dim, 8), 64)
1135
- return tile_tokens_dim
1136
-
1137
-
1127
class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1128
"""
1129
MoE Method for FP4 Quantization.
@@ -1332,8 +1321,8 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1321
):
1322
from flashinfer import nvfp4_block_scale_interleave
1323
from flashinfer.fused_moe.core import (
1335
- _maybe_get_cached_w2_permute_indices,
1324
_maybe_get_cached_w3_w1_permute_indices,
1325
+ get_w2_permute_indices_with_cache,
1326
)
1327
1328
"""Prepare quantized weights for kernel (done offline with weights)."""
@@ -1394,7 +1383,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1383
)
1384
)
1385
1397
- permute_indices = _maybe_get_cached_w2_permute_indices(
1386
+ permute_indices = get_w2_permute_indices_with_cache(
1387
self._cache_permute_indices,
1388
gemm2_weights_fp4[i].view(torch.uint8),
1389
epilogue_tile_m,
@@ -1405,7 +1394,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1394
.contiguous()
1395
)
1396
1408
- permute_sf_indices = _maybe_get_cached_w2_permute_indices(
1397
+ permute_sf_indices = get_w2_permute_indices_with_cache(
1398
self._cache_permute_indices,
1399
gemm2_scales_linear_fp4[i].view(torch.uint8),
1400
epilogue_tile_m,
@@ -1664,9 +1653,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1653
local_expert_offset=layer.ep_rank * layer.local_num_experts,
1654
local_num_experts=layer.local_num_experts,
1655
routed_scaling_factor=None,
1667
- tile_tokens_dim=_get_tile_tokens_dim(
1668
- x.shape[0], top_k, layer.local_num_experts
1669
- ),
1656
+ tile_tokens_dim=None,
1657
routing_method_type=routing_method_type,
1658
do_finalize=True,
1659
)[0]
vllm/model_executor/layers/quantization/mxfp4.py
+8
-39
@@ -50,7 +50,6 @@ from vllm.scalar_type import scalar_types
50
from vllm.utils import (
51
has_triton_kernels,
52
is_torch_equal_or_newer,
53
- next_power_of_2,
53
round_up,
54
)
55
from vllm.utils.flashinfer import has_flashinfer
@@ -97,12 +96,6 @@ def get_mxfp4_backend():
96
and has_flashinfer()
97
and envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8
98
):
100
- logger.info_once(
101
- "Using FlashInfer MXFP4 MXFP8 TRTLLM backend for SM100, "
102
- "for high concurrency throughput workloads consider setting "
103
- "VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS=1 for better "
104
- "performance"
105
- )
99
return Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM
100
elif current_platform.is_device_capability(100) and has_flashinfer():
101
logger.info_once(
@@ -357,7 +350,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
350
or self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16
351
):
352
from flashinfer.fp4_quantization import nvfp4_block_scale_interleave
360
- from flashinfer.fused_moe.core import _maybe_get_cached_w2_permute_indices
353
+ from flashinfer.fused_moe.core import get_w2_permute_indices_with_cache
354
355
layer.gemm1_alpha = Parameter(
356
torch.tensor([1.702] * self.num_experts, dtype=torch.float32).cuda(),
@@ -449,7 +442,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
442
epilogue_tile_m = 128 # FIXME: this depends on the kernel internals
443
for i in range(self.num_experts):
444
# w13 weight shuffling
452
- permute_indices = _maybe_get_cached_w2_permute_indices(
445
+ permute_indices = get_w2_permute_indices_with_cache(
446
self._cache_permute_indices,
447
w13_weight[i].view(torch.uint8),
448
epilogue_tile_m,
@@ -460,7 +453,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
453
.contiguous()
454
)
455
# w13 scale shuffling
463
- permute_sf_indices = _maybe_get_cached_w2_permute_indices(
456
+ permute_sf_indices = get_w2_permute_indices_with_cache(
457
self._cache_permute_indices,
458
w13_weight_scale[i].view(torch.uint8),
459
epilogue_tile_m,
@@ -476,7 +469,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
469
)
470
)
471
# w13 bias shuffling
479
- permute_bias_indices = _maybe_get_cached_w2_permute_indices(
472
+ permute_bias_indices = get_w2_permute_indices_with_cache(
473
self._cache_permute_indices,
474
w13_bias[i].clone().reshape(-1, 1),
475
epilogue_tile_m,
@@ -488,7 +481,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
481
.contiguous()
482
)
483
# w2 weight shuffling
491
- permute_indices = _maybe_get_cached_w2_permute_indices(
484
+ permute_indices = get_w2_permute_indices_with_cache(
485
self._cache_permute_indices,
486
w2_weight[i].view(torch.uint8),
487
epilogue_tile_m,
@@ -499,7 +492,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
492
.contiguous()
493
)
494
# w2 scale shuffling
502
- permute_sf_indices = _maybe_get_cached_w2_permute_indices(
495
+ permute_sf_indices = get_w2_permute_indices_with_cache(
496
self._cache_permute_indices,
497
w2_weight_scale[i].view(torch.uint8),
498
epilogue_tile_m,
@@ -515,7 +508,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
508
)
509
)
510
# w2 bias shuffling
518
- permute_indices = _maybe_get_cached_w2_permute_indices(
511
+ permute_indices = get_w2_permute_indices_with_cache(
512
self._cache_permute_indices,
513
w2_bias[i].clone().reshape(-1, 1),
514
epilogue_tile_m,
@@ -735,30 +728,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
728
else:
729
raise ValueError(f"Unsupported backend: {self.mxfp4_backend}")
730
738
- def _get_tile_tokens_dim(self, x: torch.Tensor, top_k: int):
739
- # Number of tokens in the input tensor.
740
- num_tokens = x.shape[0]
741
- # Factor to account for the imbalance of the experts.
742
- # factor equals to the
743
- # max_real_num_tokens_per_expert / perfect_num_tokens_per_expert
744
- # - 1.0 means perfect expert distribution.
745
- # - > 1.0 means some experts have more
746
- # tokens than the perfect distribution.
747
- # - < 1.0 does not make sense.
748
- imbalance_factor = 1.3
749
- # Calculate the number of tokens per expert
750
- # assuming perfect distribution.
751
- num_tokens_per_expert = (num_tokens * top_k) // self.num_experts
752
- # Apply the imbalance factor.
753
- num_tokens_per_expert = int(num_tokens_per_expert * imbalance_factor)
754
- # And pad the number to the next power of 2.
755
- tile_tokens_dim = next_power_of_2(num_tokens_per_expert)
756
- # Cap to 8-64 tokens per CTA tile
757
- # as it's the range supported by the kernel.
758
- tile_tokens_dim = min(max(tile_tokens_dim, 8), 64)
759
-
760
- return tile_tokens_dim
761
-
731
def get_fused_moe_quant_config(
732
self, layer: torch.nn.Module
733
) -> FusedMoEQuantConfig | None:
@@ -1037,7 +1006,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
1006
layer.ep_rank * layer.local_num_experts, # local_expert_offset
1007
self.num_experts, # local num experts
1008
None,
1040
- self._get_tile_tokens_dim(x, top_k),
1009
+ None,
1010
1 if renormalize else 0, # routing_method_type, renormalize
1011
True, # do finalize
1012
tune_max_num_tokens=self.max_capture_size,