[Dont merge]Enale FusedMoe kernel V2#448
Conversation
There was a problem hiding this comment.
Pull request overview
This PR adds a new “single-pass” SYCL implementation of the fused MoE path, exposes it as a new _moe_C.fused_moe custom op, and wires a Python-side dispatch switch (env-gated) to use it under specific constraints.
Changes:
- Add a new env-gated dispatch path in
XpuFusedMoe.apply()to call a new_moe_C.fused_moeop (intended for bf16/fp16, non-EP, non-quantized). - Introduce and register a new XPU custom op
fused_moe(...) -> Tensorand its C++/SYCL implementation (fused_moe.cpp+fused_moe_impl.hpp). - Update the
_moe_Cextension build to compile with CUTLASS/Cute flags and include directories needed by the new kernel.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
vllm_xpu_kernels/fused_moe_interface.py |
Adds env-controlled routing to the new fused MoE kernel path. |
csrc/moe/torch_bindings.cpp |
Registers the new fused_moe custom op schema and XPU implementation binding. |
csrc/moe/moe_ops.h |
Declares the new fused_moe entry point. |
csrc/moe/fused_moe.cpp |
Implements the new fused MoE op and supporting top-k sorting on XPU/SYCL. |
csrc/moe/fused_moe_impl.hpp |
Contains the SYCL kernel implementation (tiling, GEMMs, and reduction). |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
|
||
| def _use_sycl_kernel_v2(self) -> bool: | ||
| """Return True when the new single-pass SYCL kernel should be used.""" | ||
| if not _is_env_enabled(FUSED_MOE_KERNEL_V2_ENV, "1"): |
| elif self._use_sycl_kernel_v2(): | ||
| self._apply_sycl_kernel(output, hidden_states, | ||
| topk_weights, topk_ids) |
| static constexpr int COUNT_WG = 256; | ||
| static constexpr int WARP_SIZE = 16; | ||
| static constexpr int MAX_LOCAL = 32; |
| const int H = (int)hidden_size; | ||
| const int I = (int)inter_size; | ||
| const int E = (int)num_experts; | ||
| const int NT = (int)input_tokens.size(0); | ||
| const int K = (int)topk_ids.size(1); |
| torch::Tensor out; | ||
| if (output_opt.has_value()) | ||
| out = output_opt.value(); | ||
| else | ||
| out = torch::zeros_like(input_tokens); | ||
|
|
|
What's your motivation? |
|
The sole benefit of this PR is the reduction of intermediate variable movement. |
d0bf749 to
f377efc
Compare
Signed-off-by: Yihua Xu <yihua.xu@intel.com>
yihuaxu
changed the title
To optimize the performance based on current MOE implementation. |
Please refer to the test result. |
To optimize the performance for fused_moe
Purpose
Try to improve the performance based on current split moe kernels to save the io cost.
Test Plan
Test Result
UT:
Decode:
Prefill:
(Optional) Documentation Update