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 5, 21:00 UTC
[Perf] Use torch compile to fuse pack topk in trtllm moe (#37695)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com> Signed-off-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com>
Wei Zhao committed
Mar 27, 2026 at 19:30 UTC
b69bf2f0b170ac5b43f72f4dd4139c5388fa5de8
3 files changed
+18
-8
vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
+3
-5
@@ -15,6 +15,7 @@ from vllm.model_executor.layers.fused_moe.config import (
15
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
16
TopKWeightAndReduceNoOP,
17
)
18
+from vllm.model_executor.layers.fused_moe.utils import trtllm_moe_pack_topk_ids_weights
19
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
20
activation_to_flashinfer_int,
21
)
@@ -152,11 +153,8 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
153
import flashinfer
154
from flashinfer.fused_moe import Fp8QuantizationType
155
155
- # Pack topk_ids and topk_weights into single tensor
156
- # Format: (expert_id << 16) | (weight_bf16.view(int16))
157
- packed_topk_ids = (topk_ids << 16) | topk_weights.to(torch.bfloat16).view(
158
- torch.int16
159
- )
156
+ # Pack topk ids and weights into format expected by the kernel.
157
+ packed_topk_ids = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)
158
159
# trtllm_fp8_block_scale_routed_moe does not support autotuning
160
# so skip this kernel during dummy run for autotuning.
vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+2
-3
@@ -15,6 +15,7 @@ from vllm.model_executor.layers.fused_moe.config import (
15
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
16
TopKWeightAndReduceNoOP,
17
)
18
+from vllm.model_executor.layers.fused_moe.utils import trtllm_moe_pack_topk_ids_weights
19
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
20
activation_to_flashinfer_int,
21
)
@@ -183,9 +184,7 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
184
assert self.quant_config.w2_scale is not None
185
186
# Pack topk ids and weights into format expected by the kernel.
186
- packed_tensor = (topk_ids.to(torch.int32) << 16) | topk_weights.to(
187
- torch.bfloat16
188
- ).view(torch.int16)
187
+ packed_tensor = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)
188
189
# trtllm_fp4_block_scale_routed_moe does not support autotuning
190
# so skip this kernel during dummy run for autotuning.
vllm/model_executor/layers/fused_moe/utils.py
+13
@@ -323,3 +323,16 @@ def normalize_batched_scales_shape(
323
@functools.cache
324
def disable_inplace() -> bool:
325
return is_torch_equal_or_newer("2.9")
326
+
327
+
328
+@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
329
+def trtllm_moe_pack_topk_ids_weights(
330
+ topk_ids: torch.Tensor, topk_weights: torch.Tensor
331
+) -> torch.Tensor:
332
+ """
333
+ Pack topk_ids and topk_weights into a single int32 tensor.
334
+ Format: (expert_id << 16) | weight_bf16.view(int16)
335
+ """
336
+ return (topk_ids.to(torch.int32) << 16) | topk_weights.to(torch.bfloat16).view(
337
+ torch.int16
338
+ )