Skip to content

[Kernels][Bugfix] Use torch op for all kernels in FusedMoE forward. Add additional testing for cudagraphs. #19717

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 18 commits into from
Jun 25, 2025
Merged

[Kernels][Bugfix] Use torch op for all kernels in FusedMoE forward. Add additional testing for cudagraphs. #19717

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
18 commits
Select commit Hold shift + click to select a range
3d288bf
turn try_get_optimal_moe_config into an op so it can be torch.compiled
bnellnm Jun 16, 2025
385e0c5
lint
bnellnm Jun 16, 2025
c98ffbe
torch.compile tests
bnellnm Jun 16, 2025
c1c362a
add tests
bnellnm Jun 16, 2025
776ad95
add compiler + cudagraph tests
bnellnm Jun 16, 2025
961b5e8
tests
bnellnm Jun 16, 2025
bd9bd37
reduce number of compile/cudagraph tests
bnellnm Jun 16, 2025
23f26c9
lint
bnellnm Jun 16, 2025
5d564f6
fix lint
bnellnm Jun 16, 2025
06b4583
fix lint
bnellnm Jun 17, 2025
463ccaa
replace import that lint removed
bnellnm Jun 17, 2025
4ab6af7
fixes
bnellnm Jun 17, 2025
695203d
lint
bnellnm Jun 17, 2025
287a204
opify at a higher level
bnellnm Jun 18, 2025
1c9fd39
de-opify deepgemm kernels
bnellnm Jun 18, 2025
79a1962
remove cruft
bnellnm Jun 18, 2025
8b5492c
fix up deep gemm tests
bnellnm Jun 18, 2025
de1d096
fix test_mixtral_moe test
bnellnm Jun 24, 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
11 changes: 10 additions & 1 deletion tests/kernels/moe/test_cutlass_moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,10 @@
(224, 1024, 1536),
(224, 3072, 1024),
(224, 3072, 1536),
(1024 * 128, 1024, 1024),
(32768, 1024, 1024),
# These sizes trigger wrong answers.
Copy link
Collaborator

Choose a reason for hiding this comment

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

@bnellnm could you share some e2e numerical testing results here? this means the current approach still has correctness problems?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

This is an issue with the cutlass kernels themselves (separate from chunking). I'm not sure if the tolerances need to be changed or if there's a real problem with cutlass. @ElizaWszola can probably provide some more insight here.

#(7232, 2048, 5120),
#(40000, 2048, 5120),
]

vllm_config = VllmConfig(parallel_config=ParallelConfig(
Expand Down Expand Up @@ -232,8 +235,10 @@ def test_cutlass_moe_8_bit_no_graph(
topk: int,
per_act_token: bool,
per_out_ch: bool,
monkeypatch,
):
current_platform.seed_everything(7)
monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")
with set_current_vllm_config(vllm_config):
mt = MOETensors8Bit.make_moe_tensors_8bit(m, k, n, e, per_act_token,
per_out_ch)
Expand Down Expand Up @@ -274,8 +279,10 @@ def test_cutlass_moe_8_bit_cuda_graph(
topk: int,
per_act_token: bool,
per_out_ch: bool,
monkeypatch,
):
current_platform.seed_everything(7)
monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")
with set_current_vllm_config(vllm_config):
dtype = torch.half

Expand Down Expand Up @@ -329,8 +336,10 @@ def test_cutlass_moe_8_bit_EP(
per_act_token: bool,
per_out_channel: bool,
ep_size: int,
monkeypatch,
):
current_platform.seed_everything(7)
monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")
with set_current_vllm_config(vllm_config):
mt = MOETensors8Bit.make_moe_tensors_8bit(m, k, n, e, per_act_token,
per_out_channel)
Expand Down
Loading