[SGLANG] Enable test_fused_moe.py in the MoE suite - #7886
Open
wdziurdz wants to merge 1 commit into
Open
Conversation
sgl-kernel is intentionally removed from the XPU environment by install-sglang.sh, but the Triton MoE path imports it unconditionally in its `elif _is_xpu:` branches, so the test fails at collection with ModuleNotFoundError. Guarding the imports alone is not enough: moe_align_block_size and topk_softmax resolve to AOT sgl_kernel ops with no Triton or JIT backend registered in sglang, so the test would import and then fail at call time. Patch the imports and add fallbacks: - activation.py: guard silu_and_mul / gelu_and_mul / gelu_tanh_and_mul and fall back to the same math forward_native uses. - triton_utils/fused_moe.py: clear _is_xpu on ImportError so the existing native silu_and_mul and moe_sum_reduce_triton paths are selected. - triton_utils/moe_align_block_size.py: torch expert grouping producing the same layout, sentinel padding and per-block expert ids. Raises NotImplementedError for the expert-parallel -1 sentinel instead of silently miscomputing. - moe/topk.py: torch topk_softmax shim, leaving topk_sigmoid undefined as the non-XPU branch already does. Also patches the fp8 capability check in the test, which assumes a CUDA-style get_device_capability(), and adds test_fused_moe.py to the sglang_moe suite. No Triton backend change is needed: the kernels the test exercises (fused_moe_kernel, moe_sum_reduce_triton) pass the test's full m/n/k/e/topk/dtype sweep on XPU. Signed-off-by: Witold Dziurdz <witold.dziurdz@intel.com>
wdziurdz
force-pushed
the
sglang-enable-fused-moe
branch
from
September 1, 2026 11:28
26c8676 to
97671a7
Compare
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
1 participant
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.
Enables the upstream
test_fused_moe.pyin--sglang-moe, which until now could not even import on XPU becausesrt/layers/activation.pyimportssgl_kernelat module level. The patch adds a native fallback for that import and makes the fp8 capability check skip rather than raise on XPU, whereget_device_capability()returns(None, None).