Skip to content

Commit 5eb81a7

Browse files
Minor modifications
Signed-off-by: Jinyang Yuan <[email protected]>
1 parent bfd9a65 commit 5eb81a7

File tree

2 files changed

+6
-2
lines changed

2 files changed

+6
-2
lines changed

tensorrt_llm/_torch/custom_ops/torch_custom_ops.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,7 @@ def forward(
8888
):
8989
x, fc1_expert_weights, fc1_expert_biases, fc2_expert_weights, fc2_expert_biases = inputs
9090
self.fused_moe_runner.run_gemm_profile(
91-
x.contiguous(),
91+
x,
9292
fc1_expert_weights,
9393
fc1_expert_biases,
9494
fc2_expert_weights,
@@ -148,6 +148,8 @@ def fused_moe(
148148
tuner = AutoTuner.get()
149149
MoERunner.refine_tuning_config(tune_max_num_tokens)
150150

151+
# Only the non-alltoall case is considered for profiling in the warmup phase.
152+
# Therefore, to get the correct tactics during the actual inference, the inputs to the tuner should be the same as when not using alltoall.
151153
if enable_alltoall:
152154
assert tuner_num_tokens is not None
153155
assert tuner_top_k is not None

tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -438,7 +438,9 @@ def forward_chunk(
438438
# If alltoall is disabled, we need also disable use_postquant_alltoall
439439
use_postquant_alltoall = self.use_postquant_alltoall and use_all_to_all
440440

441-
# Prepare additional information for profiling in case padding is applied in all-to-all
441+
# Prepare additional information for profiling in case padding is applied when using alltoall.
442+
# Only the non-alltoall case is considered for profiling in the warmup phase.
443+
# Therefore, to get the correct tactics during the actual inference, the inputs to the tuner should be the same as when not using alltoall.
442444
if use_all_to_all:
443445
if all_rank_num_tokens is not None:
444446
tuner_num_tokens = sum(all_rank_num_tokens)

0 commit comments

Comments
 (0)