[XPU][MOE] Optimize XPU MoE activation by skipping invalid rows#401
[XPU][MOE] Optimize XPU MoE activation by skipping invalid rows#401ranzhejiang wants to merge 1 commit into
Conversation
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
Adds MoE-specific activation variants that use a device-side valid_rows/valid_tokens scalar to skip padded rows on XPU, avoiding CPU↔GPU synchronization.
Changes:
- Introduced
*_moeTorch ops (bindings + headers) acceptingvalid_rowsas an int32 device scalar to guard kernels on-device. - Added SYCL kernel variants for SiLU/GELU/ReLU2/SwiGLU* that early-exit when
token_idx >= *valid_rows. - Updated Python MoE execution paths to compute
valid_tokenson-device and call the new*_moeactivation ops.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| vllm_xpu_kernels/fused_moe_interface.py | Computes valid_tokens on-device and routes activation through new MoE-specific ops. |
| csrc/torch_bindings.cpp | Registers new *_moe activation ops with valid_rows parameter. |
| csrc/ops.h | Declares new MoE activation op entry points. |
| csrc/activation.cpp | Implements SYCL kernels and launch macros that guard work-groups by valid_rows. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| #define LAUNCH_SWIGLUOAI_AND_MUL_MOE(KERNEL, ALPHA, LIMIT) \ | ||
| int d = input.size(-1) / 2; \ | ||
| int64_t num_tokens = input.numel() / input.size(-1); \ | ||
| sycl::range<1> grid(num_tokens); \ | ||
| sycl::range<1> block(std::min(d, 1024)); \ | ||
| at::DeviceGuard device_guard(input.device()); \ | ||
| auto& queue = vllm::xpu::vllmGetQueue(); \ | ||
| VLLM_DISPATCH_FLOATING_TYPES( \ | ||
| input.scalar_type(), "swigluoai_and_mul_kernel_moe", [&] { \ | ||
| queue.submit([&](sycl::handler& cgh) { \ | ||
| cgh.parallel_for( \ | ||
| sycl::nd_range<1>(grid * block, block), \ | ||
| vllm::swigluoai_and_mul_kernel_moe< \ | ||
| scalar_t, KERNEL<scalar_t>>( \ | ||
| out.data_ptr<scalar_t>(), \ | ||
| input.data_ptr<scalar_t>(), \ | ||
| valid_rows.data_ptr<int32_t>(), \ | ||
| d, \ | ||
| ALPHA, \ | ||
| LIMIT)); \ | ||
| }); \ | ||
| }); | ||
|
|
||
| #define LAUNCH_SWIGLUSTEP_AND_MUL_MOE(KERNEL, LIMIT) \ | ||
| int d = input.size(-1) / 2; \ | ||
| int64_t num_tokens = input.numel() / input.size(-1); \ | ||
| sycl::range<1> grid(num_tokens); \ | ||
| sycl::range<1> block(std::min(d, 1024)); \ | ||
| at::DeviceGuard device_guard(input.device()); \ | ||
| auto& queue = vllm::xpu::vllmGetQueue(); \ | ||
| VLLM_DISPATCH_FLOATING_TYPES( \ | ||
| input.scalar_type(), "swiglustep_and_mul_kernel_moe", [&] { \ | ||
| queue.submit([&](sycl::handler& cgh) { \ | ||
| cgh.parallel_for( \ | ||
| sycl::nd_range<1>(grid * block, block), \ | ||
| vllm::swiglustep_and_mul_kernel_moe< \ | ||
| scalar_t, KERNEL<scalar_t>>( \ | ||
| out.data_ptr<scalar_t>(), \ | ||
| input.data_ptr<scalar_t>(), \ | ||
| valid_rows.data_ptr<int32_t>(), \ | ||
| d, \ | ||
| LIMIT)); \ | ||
| }); \ | ||
| }); |
| void silu_and_mul_moe( | ||
| torch::Tensor& out, // [num_moe_inputs, d] | ||
| torch::Tensor& input, // [num_moe_inputs, 2*d] | ||
| torch::Tensor& valid_rows) // int32 device scalar = rows_per_expert.sum() | ||
| { | ||
| VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "silu_and_mul_moe", [&] { | ||
| LAUNCH_ACTIVATION_GATE_KERNEL_VEC_MOE(vllm::silu_kernel, true); | ||
| }); | ||
| } | ||
|
|
||
| void gelu_and_mul_moe( | ||
| torch::Tensor& out, | ||
| torch::Tensor& input, | ||
| torch::Tensor& valid_rows) | ||
| { | ||
| VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "gelu_and_mul_moe", [&] { | ||
| LAUNCH_ACTIVATION_GATE_KERNEL_VEC_MOE(vllm::gelu_kernel, true); | ||
| }); | ||
| } |
| elif activation == "swigluoai" or ("SWIGLUOAI" in str(activation)): | ||
| torch.ops._C.swigluoai_and_mul_moe(act_output, gemm1_output, | ||
| valid_tokens, 1.702, 7.0) |
| void swigluoai_and_mul_moe(torch::Tensor& out, torch::Tensor& input, | ||
| torch::Tensor& valid_rows, double alpha, | ||
| double limit); | ||
|
|
||
| void swiglustep_and_mul_moe(torch::Tensor& out, torch::Tensor& input, | ||
| torch::Tensor& valid_rows, double limit); |
127aa6b to
b6c5156
Compare
|
fix pre-commit |
|
I this if this code change is merged, modes other than EP mode will result in a performance decline, because other modes don't need this. |
ok，I will test it with performance data |
faec530 to
0f55b6f
Compare
1558c65 to
a7ea252
Compare
I add
Tested with Qwen3-30B-A3B using In theory, the time savings should become more significant as EP increases and sequence length becomes longer. |
|
Please fix pre-commit. |
|
BTW, can you please provide the performance data with tp=4 and without ep? |
be45eb3 to
0ba9bca
Compare
Signed-off-by: ranzhejiang <zhejiang.ran@intel.com>
Tested with Qwen3-30B-A3B using |
|
@Liangliang-Ma |
Summary
This PR optimizes the XPU fused MoE activation path for EP prefill by passing the valid routed row count into fusion and skipping invalid trailing rows on device.
In EP, the real routed rows can be smaller than num_rows * topk. Without this information, activation still runs on the full upper-bound shape even though grouped GEMM already uses the real routed rows. Carrying the valid row count into fusion avoids redundant activation work without introducing Python-side host synchronization.
This path is only enabled for EP prefill. Decode keeps the existing behavior, and TP is unaffected because it does not use this EP-specific valid-row path
In https://github.com/Tencent/hpc-ops, they also use valid_row_range_ptr to skip unused rows