[MoE Refactor] Split invoke_fused_moe_kernel#31050
Merged
vllm-bot merged 10 commits intovllm-project:mainfrom Jan 2, 2026
Merged
[MoE Refactor] Split invoke_fused_moe_kernel#31050vllm-bot merged 10 commits intovllm-project:mainfrom
invoke_fused_moe_kernel#31050vllm-bot merged 10 commits intovllm-project:mainfrom
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Comment on lines
839
to
+843
| num_tokens_post_padded, | ||
| mul_routed_weight, | ||
| top_k, | ||
| config["BLOCK_SIZE_M"], | ||
| config["BLOCK_SIZE_N"], | ||
| config["BLOCK_SIZE_K"], | ||
| bit, | ||
| config, | ||
| block_shape, |
There was a problem hiding this comment.
In dispatch_fused_moe_kernel the W8A16/W4A16 CUDA path calls invoke_fused_moe_triton_kernel_gptq_awq without the required compute_type, use_int8_w8a16, and use_int4_w4a16 arguments (the last positional argument is block_shape). The callee’s signature defined above requires those parameters, so when should_moe_wna16_use_cuda(...) returns true (grouped quantization with block_shape set), this branch raises a TypeError before running the kernel. Any WNA16 call with block quantization will therefore crash rather than execute.
Useful? React with 👍 / 👎.
|
Hi @zyongye, the pre-commit checks have failed. Please run: uv pip install pre-commit
pre-commit install
pre-commit run --all-filesThen, commit the changes and push to your branch. For future commits, Tip Is
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request refactors the fused MoE kernel invocation logic by introducing new specialized Triton kernel functions for different quantization types (WNA16, GPTQ/AWQ) and consolidating their dispatch through a new dispatch_fused_moe_kernel function. This dispatcher replaces previous direct calls to invoke_fused_moe_kernel across various fused_experts_impl and apply methods. Additionally, the unquantized_fused_moe_method.py file is updated to utilize a FusedMoEModularKernel abstraction, which now handles MoE preparation and expert execution. Review feedback indicates the need to remove a pdb.set_trace() debugging call and to optimize the instantiation of FusedMoEModularKernel by moving it outside the forward_cuda method to improve efficiency.
| @@ -1680,6 +1832,7 @@ def fused_experts( | |||
| allow_deep_gemm: bool = False, | |||
| allow_cutlass_block_scaled_grouped_gemm: bool = False, | |||
| ) -> torch.Tensor: | |||
| # import pdb; pdb.set_trace() | |||
Contributor
There was a problem hiding this comment.
A pdb.set_trace() call is present. This should be removed before merging.
Comment on lines
+329
to
+333
| self.kernel = mk.FusedMoEModularKernel( | ||
| MoEPrepareAndFinalizeNoEP(), | ||
| TritonExperts(self.moe_quant_config), | ||
| shared_experts=None, | ||
| ) |
Contributor
There was a problem hiding this comment.
The FusedMoEModularKernel is instantiated on every forward pass. This is inefficient. It should be instantiated once, for example in the process_weights_after_loading method, and stored as a class member to be reused across forward passes.
|
Hi @zyongye, the pre-commit checks have failed. Please run: uv pip install pre-commit
pre-commit install
pre-commit run --all-filesThen, commit the changes and push to your branch. For future commits, Tip Is
|
zyongye
changed the title
invoke_fused_moe_kernel
|
This pull request has merge conflicts that must be resolved before it can be |
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
zyongye
force-pushed
the
bf16_triton_refactor
branch
from
0c14373 to
4f69e85
Compare
jinzhen-lin
reviewed
Dec 26, 2025
|
|
||
|
|
||
| def invoke_fused_moe_kernel( | ||
| def invoke_fused_moe_triton_kernel_wna16( |
Contributor
There was a problem hiding this comment.
This function runs the cuda kernel, but not the triton kernel.
Maybe we can use invoke_fused_moe_wna16_cuda_kernel and invoke_fused_moe_wna16_triton_kernel for the two kernels of moe_wna16.
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Contributor
|
The rest of the PR LGTM. BTW, I just found the moe wna16 maybe used in SM75+ in some cases, for example, when |
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
zyongye
force-pushed
the
bf16_triton_refactor
branch
from
0f6a153 to
ff4af42
Compare
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
zyongye
force-pushed
the
bf16_triton_refactor
branch
from
986ce16 to
d206c47
Compare
robertgshaw2-redhat
approved these changes
Jan 1, 2026
Collaborator
|
Thanks for reviewing @jinzhen-lin |
mgoin
approved these changes
Jan 2, 2026
Contributor
|
Hi guys @zyongye @robertgshaw2-redhat this PR causes that Inference FAILS |
Member
Author
|
Thank you for the report! Could you paste your command here? Thank you! |
Contributor
|
@zyongye |
Contributor
|
@zyongye |
Member
Author
|
@JartX is that model on huggingface? How can I run it on my end? |
Contributor
|
@zyongye jart25/Qwen3-VL-30B-A3B-Instruct-AWQ-4bit https://huggingface.co/jart25/Qwen3-VL-30B-A3B-Instruct-AWQ-4bit |
Contributor
|
@zyongye It is qwen3 vl 30b compressed tensor awq int4 |
Contributor
Member
Author
|
@JartX Yes I am aware of the PR. But changing back defeats the purpose of MoE refactoring. I am looking into the code path to see what I am missing. Thank you for the reminder. |
Contributor
|
Sorry, I didn't mean to bother you, just wanted to offer some advice. I can test any code you have and give you feedback :). Cheers! |
LucasWilkinson
pushed a commit
to neuralmagic/vllm
that referenced
this pull request
Jan 6, 2026
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
yugong333
pushed a commit
to yugong333/vllm
that referenced
this pull request
Jan 9, 2026
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Member
|
@JartX one of your comments here included your huggingface token. I removed it from the comment, but I would also suggest revoking the token on huggingface if you haven't already. |
Contributor
|
@russellb Thanks for the heads-up; the token was edited before pasting the command to serve as an example. I also want to thank you for your dedication in personally caring for and supporting the community. I owe you two beers 🙂 |
Contributor
|
@russellb Speaking of something private, how could we upload the tests, LoRa tests by model? If I publish a LoRa for Qwen3vl 30B, could you tell me the correct way to do it? Thank you very much. |
akh64bit
pushed a commit
to akh64bit/vllm
that referenced
this pull request
Jan 16, 2026
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
dsuhinin
pushed a commit
to dsuhinin/vllm
that referenced
this pull request
Jan 21, 2026
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Signed-off-by: dsuhinin <suhinin.dmitriy@gmail.com>
ItzDEXX
pushed a commit
to ItzDEXX/vllm
that referenced
this pull request
Feb 19, 2026
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
8 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Currently
invoke_fused_moe_kernelfunction includes three kernels.Given that the marlin kernels widely support integer quantization starting SM75 (Turing), we isolated the widely used triton kernel with wna16 kernels so that later we can remove it cleanly when vllm drop support for SM70 GPUs.