Skip to content

Route the missing parameter for trtllm_fp8_per_tensor_scale_moe_op #3094

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
aleozlx merged 4 commits into from
Apr 18, 2026
Merged

Route the missing parameter for trtllm_fp8_per_tensor_scale_moe_op #3094

Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
8928433
Fix vLLM CI failure for 0.6.8
pavanimajety Apr 16, 2026
480ee88
Fix unintended change
pavanimajety Apr 16, 2026
311ef08
Fix the positional arg issue
pavanimajety Apr 16, 2026
27e9b21
cherry pick 1125dbc (allow post release)
aleozlx Apr 3, 2026
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
109 changes: 59 additions & 50 deletions flashinfer/fused_moe/core.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1428,31 +1428,31 @@ def trtllm_bf16_moe_op(

# Call the C++ function with the selected tactic
intermediate_output = moe_op.trtllm_bf16_moe(
routing_logits,
routing_bias,
topk_ids,
expert_weights,
hidden_states,
gemm1_weights,
gemm2_weights,
output,
num_experts,
top_k,
n_group,
topk_group,
intermediate_size,
local_expert_offset,
local_num_experts,
routed_scaling_factor,
routing_method_type,
use_shuffled_weight,
weight_layout,
do_finalize,
enable_pdl,
[-1, -1] if tactic == -1 else tactic,
activation_type,
norm_topk_prob,
routing_replay_out,
routing_logits=routing_logits,
routing_bias=routing_bias,
topk_ids=topk_ids,
expert_weights=expert_weights,
hidden_states=hidden_states,
gemm1_weights=gemm1_weights,
gemm2_weights=gemm2_weights,
output=output,
num_experts=num_experts,
top_k=top_k,
n_group=n_group,
topk_group=topk_group,
intermediate_size=intermediate_size,
local_expert_offset=local_expert_offset,
local_num_experts=local_num_experts,
routed_scaling_factor=routed_scaling_factor,
routing_method_type=routing_method_type,
use_shuffled_weight=use_shuffled_weight,
weight_layout=weight_layout,
do_finalize=do_finalize,
enable_pdl=enable_pdl,
config_index=[-1, -1] if tactic == -1 else tactic,
activation_type=activation_type,
norm_topk_prob=norm_topk_prob,
routing_replay_out=routing_replay_out,
)
if do_finalize:
return [output]
Expand Down Expand Up @@ -1601,30 +1601,31 @@ def trtllm_fp8_per_tensor_scale_moe_op(
)
# Call the C++ function
intermediate_output = moe_op.trtllm_fp8_per_tensor_scale_moe(
routing_logits,
routing_bias,
hidden_states,
gemm1_weights,
output1_scales_scalar,
output1_scales_gate_scalar,
gemm2_weights,
output2_scales_scalar,
output,
num_experts,
top_k,
n_group,
topk_group,
intermediate_size,
local_expert_offset,
local_num_experts,
routed_scaling_factor,
use_routing_scales_on_input,
routing_method_type,
do_finalize,
enable_pdl,
[-1, -1] if tactic == -1 else tactic,
activation_type,
norm_topk_prob,
routing_logits=routing_logits,
routing_bias=routing_bias,
hidden_states=hidden_states,
gemm1_weights=gemm1_weights,
output1_scales_scalar=output1_scales_scalar,
output1_scales_gate_scalar=output1_scales_gate_scalar,
gemm2_weights=gemm2_weights,
output2_scales_scalar=output2_scales_scalar,
output=output,
num_experts=num_experts,
top_k=top_k,
n_group=n_group,
topk_group=topk_group,
intermediate_size=intermediate_size,
local_expert_offset=local_expert_offset,
local_num_experts=local_num_experts,
routed_scaling_factor=routed_scaling_factor,
use_routing_scales_on_input=use_routing_scales_on_input,
routing_method_type=routing_method_type,
do_finalize=do_finalize,
enable_pdl=enable_pdl,
config_index=[-1, -1] if tactic == -1 else tactic,
activation_type=activation_type,
norm_topk_prob=norm_topk_prob,
routing_replay_out=routing_replay_out,
)
if do_finalize:
return [output]
Expand Down Expand Up @@ -2620,6 +2621,13 @@ def trtllm_fp8_per_tensor_scale_moe(
otherwise, returns the intermediate results (gemm2_output, expert_weights, expanded_idx_to_permuted_idx) that need further processing.
"""
_validate_routing_replay_out(routing_replay_out, top_k)
num_tokens = hidden_states.shape[0]
hidden_size = hidden_states.shape[-1]
if enable_pdl is None:
enable_pdl = device_support_pdl(hidden_states.device)
output = torch.empty(
num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device
)
Comment thread
pavanimajety marked this conversation as resolved.
Outdated
result = get_trtllm_moe_sm100_module().trtllm_fp8_per_tensor_scale_moe(
routing_logits,
routing_bias,
Expand All @@ -2629,6 +2637,7 @@ def trtllm_fp8_per_tensor_scale_moe(
output1_scales_gate_scalar,
gemm2_weights,
output2_scales_scalar,
output,
num_experts,

@gemini-code-assist gemini-code-assist Bot Apr 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

This change introduces a critical bug due to a positional argument mismatch. The target function trtllm_fp8_per_tensor_scale_moe_op does not include output in its signature (defined at line 1503). By inserting output at this position, all subsequent arguments are shifted: num_experts will be passed to the top_k parameter, top_k to n_group, and so on. This will cause a TypeError or incorrect kernel execution. This line should be removed to restore the correct positional mapping.

Suggested change
output,
num_experts,
num_experts,

top_k,
n_group,
Expand All @@ -2641,7 +2650,7 @@ def trtllm_fp8_per_tensor_scale_moe(
routing_method_type,
do_finalize,
enable_pdl,
tune_max_num_tokens,
[-1, -1],
Comment thread
pavanimajety marked this conversation as resolved.
Outdated
activation_type,
norm_topk_prob,
routing_replay_out,
Expand Down
Loading