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 11, 15:00 UTC
[Perf] Optimize deepgemm experts initialization, 3.9% TTFT improvement (#30494)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: li-jinpeng <3332126450@qq.com> Co-authored-by: youkaichao <youkaichao@gmail.com>
Wentao Ye committed
Dec 11, 2025 at 17:28 UTC
c817b1415121cf88178af1e4e78f651d802df4da
1 file changed
+17
-6
vllm/model_executor/layers/fused_moe/deep_gemm_utils.py
+17
-6
@@ -84,10 +84,16 @@ def _fwd_kernel_ep_scatter_1(
84
m_indices_start_ptr = m_indices + cur_expert_start
85
off_expert = tl.arange(0, BLOCK_E)
86
87
+ # any rows in the per-expert aligned region that do not correspond to
88
+ # real tokens are left untouched here and should remain initialized to
89
+ # -1 so DeepGEMM can skip them
90
for start_m in tl.range(0, cur_expert_token_num, BLOCK_E, num_stages=4):
91
+ offs = start_m + off_expert
92
+ mask = offs < cur_expert_token_num
93
tl.store(
89
- m_indices_start_ptr + start_m + off_expert,
94
+ m_indices_start_ptr + offs,
95
cur_expert,
96
+ mask=mask,
97
)
98
99
@@ -366,12 +372,17 @@ def deepgemm_moe_permute(
372
(M_sum, H // block_k), device=device, dtype=torch.float32
373
)
374
369
- maybe_has_empty_blocks = (expert_tokens_meta is None) or (
370
- expert_tokens_meta.expert_num_tokens_cpu is None
375
+ # DeepGEMM uses negative values in m_indices (here expert_ids) to mark
376
+ # completely invalid / padded blocks that should be skipped. We always
377
+ # initialize expert_ids to -1 so any row that is not explicitly written
378
+ # by the scatter kernel will be treated as invalid and skipped by
379
+ # DeepGEMM's scheduler.
380
+ expert_ids = torch.full(
381
+ (M_sum,),
382
+ fill_value=-1,
383
+ device=device,
384
+ dtype=torch.int32,
385
)
372
- expert_ids_init = torch.zeros if maybe_has_empty_blocks else torch.empty
373
-
374
- expert_ids = expert_ids_init((M_sum), device=device, dtype=torch.int32)
386
inv_perm = torch.empty(topk_ids.shape, device=device, dtype=torch.int32)
387
388
expert_num_tokens = None