Skip to content

Update cutlass_moe.py #8545

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 1 commit into from
Jul 30, 2025
Merged

Update cutlass_moe.py #8545

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
2 changes: 1 addition & 1 deletion python/sglang/srt/layers/moe/cutlass_moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -209,7 +209,7 @@ def cutlass_fused_experts_fp8(
)

result = torch.empty((m, k), device=device, dtype=out_dtype)
apply_shuffle_mul_sum(c2, result, c_map, topk_weights)
Copy link
Contributor

Choose a reason for hiding this comment

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

high

There is a potential dtype mismatch here. The apply_shuffle_mul_sum kernel expects c2, result, and topk_weights to have the same dtype (out_dtype). However, topk_weights is passed without an explicit cast, which can lead to incorrect results or silent errors if its dtype does not match. It should be cast to out_dtype before being passed to the kernel.

Suggested change
apply_shuffle_mul_sum(c2, result, c_map, topk_weights)
apply_shuffle_mul_sum(c2, result, c_map, topk_weights.to(out_dtype))

apply_shuffle_mul_sum(c2, result, c_map, topk_weights.to(out_dtype))
Copy link
Collaborator

Choose a reason for hiding this comment

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

Can we also add a dtype check in c++ side void get_apply_shuffle_mul_sum_caller() to avoid any misusage in the future?

Copy link
Collaborator

Choose a reason for hiding this comment

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

It seems we want to make sure the c2, result, topk_weights to have the same dtype, right?

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

We have the checker hence the need to cast to out_dtype.
Also the kernel is very specifically written for MoE weighted sum, it will likely only appear in this file only.

sglang/sgl-kernel/csrc/moe/prepare_moe_input.cu

Line 351 in 9effeb5

TORCH_CHECK(factors_opt->dtype() == output_tensor.dtype(), "Factors must match output dtype");

Copy link
Collaborator

Choose a reason for hiding this comment

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

I see. Just curious why this mismatch wasn't captured by the previous commit.

return result


Expand Down
Loading