Skip to content

feat: autotune tile_tokens_dim in trtllm-gen MOE #1980

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 15 commits into from
Oct 29, 2025
Merged

feat: autotune tile_tokens_dim in trtllm-gen MOE #1980

Show file tree
Hide file tree
Changes from 14 commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
5656c2a
add autotune tile N
jiahanc Oct 25, 2025
a8ebb37
opt autotune logic and add deprecate log
jiahanc Oct 27, 2025
cdf366f
rename
jiahanc Oct 27, 2025
e835e02
fix format
jiahanc Oct 28, 2025
43e912d
add FP8 autotune and update deprecate warning
jiahanc Oct 28, 2025
5bdffad
fix autotune and apply fp8 auto tune to test
jiahanc Oct 28, 2025
9894e03
add include set
jiahanc Oct 28, 2025
1c2bcc5
add log_once and update deprecation log
jiahanc Oct 28, 2025
1284c44
add 128 to FP8 PT
jiahanc Oct 28, 2025
884ea82
fix test
jiahanc Oct 28, 2025
7062a5d
uncomment
jiahanc Oct 28, 2025
b68c853
fix test
jiahanc Oct 28, 2025
075f096
fix
jiahanc Oct 28, 2025
0316720
update test
jiahanc Oct 29, 2025
f8a75ff
pre-commit
yzh119 Oct 29, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 2 additions & 5 deletions benchmarks/bench_trtllm_gen_fused_moe_autotuner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
from flashinfer.fused_moe import trtllm_fp4_block_scale_moe
from flashinfer.autotuner import autotune
from flashinfer.testing.utils import bench_gpu_time
from flashinfer.utils import device_support_pdl, calculate_tile_tokens_dim
from flashinfer.utils import device_support_pdl


def bench_trtllm_gen_fused_moe_autotuner(
Expand Down Expand Up @@ -99,9 +99,6 @@ def bench_trtllm_gen_fused_moe_autotuner(
bias13 = torch.randn(num_experts, intermediate_size * 2, device=device) * 10
bias2 = torch.randn(num_experts, intermediate_size * 2, device=device) * 10

tile_tokens_dim = calculate_tile_tokens_dim(
num_tokens, num_experts, top_k, 64 if quant_mode == "MxFP4xBf16" else 128
)
output1_scale_scalar = torch.tensor(
[hidden_states_global_scale * w13_global_scale] * num_experts, device=device
)
Expand Down Expand Up @@ -136,7 +133,7 @@ def bench_trtllm_gen_fused_moe_autotuner(
0, # local_expert_offset
num_experts,
None, # routed_scaling_factor
tile_tokens_dim,
None, # tile_tokens_dim
RoutingMethodType.Renormalize.value,
True,
enable_pdl,
Expand Down
Loading
Loading