Skip to content

[torch.compile] Unwrap fused_marlin_moe custom op#26739

Merged
mgoin merged 1 commit intovllm-project:mainfrom
neuralmagic:varun/unwrap-fused-marlin-moe
Oct 14, 2025
Merged

[torch.compile] Unwrap fused_marlin_moe custom op#26739
mgoin merged 1 commit intovllm-project:mainfrom
neuralmagic:varun/unwrap-fused-marlin-moe

Conversation

@varun-sundar-rabindranath
Copy link
Copy Markdown
Contributor

@varun-sundar-rabindranath varun-sundar-rabindranath commented Oct 13, 2025

Purpose

Unwrap unnecessary CustomOP'ing of fused_marlin_moe.

fused_marlin_moe was first wrapped into a custom op in the PR #9632 as fused_marlin_moe was doing file reads to get optimal triton config. This function has evolved since and no longer needs file reads and could be unwrapped.

Test Plan

CI

Test Result

Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
@varun-sundar-rabindranath
Copy link
Copy Markdown
Contributor Author

cc @youkaichao @ProExpertProg @bnellnm PTAL. Thanks 🙌

Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

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

Code Review

This pull request correctly refactors the fused_marlin_moe function by removing its torch.ops custom operator wrapper. The changes are consistent across all modified files, replacing the custom op calls with direct function calls and removing the associated registration code. This is a clean and well-executed refactoring that achieves its goal of making the function more amenable to torch.compile by exposing its internal logic. The changes are sound and I found no issues.

Copy link
Copy Markdown
Member

@mgoin mgoin left a comment

Choose a reason for hiding this comment

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

Nice find, LGTM

@mgoin mgoin added quantization ready ONLY add when PR is ready to merge/full CI is needed labels Oct 14, 2025
@mgoin mgoin enabled auto-merge (squash) October 14, 2025 02:22
@mgoin mgoin merged commit 8ae1692 into vllm-project:main Oct 14, 2025
60 checks passed
1994 pushed a commit to 1994/vllm that referenced this pull request Oct 14, 2025
Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
Signed-off-by: 1994 <1994@users.noreply.github.com>
Dhruvilbhatt pushed a commit to Dhruvilbhatt/vllm that referenced this pull request Oct 14, 2025
Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
Signed-off-by: Dhruvil Bhatt <bhattdbh@amazon.com>
bbartels pushed a commit to bbartels/vllm that referenced this pull request Oct 16, 2025
Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
Signed-off-by: bbartels <benjamin@bartels.dev>
lywa1998 pushed a commit to lywa1998/vllm that referenced this pull request Oct 20, 2025
Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
alhridoy pushed a commit to alhridoy/vllm that referenced this pull request Oct 24, 2025
Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
0xrushi pushed a commit to 0xrushi/vllm that referenced this pull request Oct 26, 2025
Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
Signed-off-by: 0xrushi <6279035+0xrushi@users.noreply.github.com>
0xrushi pushed a commit to 0xrushi/vllm that referenced this pull request Oct 26, 2025
Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
Signed-off-by: 0xrushi <6279035+0xrushi@users.noreply.github.com>
rtourgeman pushed a commit to rtourgeman/vllm that referenced this pull request Nov 10, 2025
Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
devpatelio pushed a commit to SumanthRH/vllm that referenced this pull request Nov 29, 2025
Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ready ONLY add when PR is ready to merge/full CI is needed

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants