Repository navigation
[AMD] Fuse shared-expert append + DeepEP remap into one Triton kernel - #28450
Conversation
Add fused_append_remap_shared_experts_deepep, which performs the shared-expert append and the DeepEP interleaved remap in a single Triton pass on the aiter/DeepEP-class path. This replaces the fused_append_shared_experts() + eager _remap_topk_for_deepep() pair, collapsing ~6 launch-bound elementwise kernels per MoE layer (div_floor / add / arange / fill / copy) into the one append kernel that already runs. The shared-expert weight is 1.0 on this aiter-only path. Also fuse the HIP padded-token masking into a single _fill_padded_rows Triton kernel (CUDA/HIP-graph safe), remove the redundant first HIP weight-zeroing pass, and add the SGLANG_MORI_NO_PAD_MASK gate. Add test/registered/moe/test_fused_append_remap_deepep.py (registered for CUDA and AMD CI) covering golden-reference equivalence, bit-identical parity with the eager append+remap pair, the aiter shared-weight contract, and the no-shared-experts no-op. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
|
/tag-and-rerun-ci |
Resolve topk.py conflict in _mask_topk_ids_padded_region: keep both the AMD fused-region branch (_can_fuse_padded_region) and upstream's new NPU branch (_is_npu); mutually exclusive platform paths.
Main's get_parallel() refactor (sgl-project#28567) removed the get_moe_expert_parallel_world_size/rank imports and migrated its own call sites, but the fused shared-expert deepep block added by this PR still referenced the old names, leaving them undefined after the merge. Switch to get_parallel().moe_ep_size / .moe_ep_rank to match.
| _HAS_TRITON = True | ||
| except ImportError: | ||
| _HAS_TRITON = False | ||
|
|
There was a problem hiding this comment.
Please simplify this, our environ has triton, move import triton to the top, as python/sglang/srt/layers/quantization/fp8_utils.py
There was a problem hiding this comment.
Pull request overview
This PR optimizes the AMD aiter + DeepEP-class MoE routing path by fusing the shared-expert append step and the DeepEP interleaved remap into a single Triton kernel, and reduces launch overhead in HIP padded-token masking by adding a fused row-fill Triton kernel with an optional env-gated skip.
Changes:
- Add a fused Triton kernel + wrapper to append shared experts and apply the DeepEP routed-id remap in one pass.
- Update MoE top-k post-processing to use the fused append+remap path on aiter + DeepEP-class backends, and add a Triton
_fill_padded_rowshelper for padded-region masking/zeroing (with an env-gated HIP skip). - Add GPU unit tests covering fused append+remap behavior and equivalence against a reference / eager sequence.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| test/registered/moe/test_fused_append_remap_deepep.py | New GPU tests for fused append+remap correctness and equivalence. |
| python/sglang/srt/layers/moe/topk.py | Wire fused append+remap into _post_process_topk_ids; add Triton _fill_padded_rows and HIP pad-mask env gate. |
| python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe_triton_kernels.py | Add fused append+remap Triton kernel + Python wrapper. |
Comments suppressed due to low confidence (1)
python/sglang/srt/layers/moe/topk.py:1255
- In the NPU branch, this function assigns
topk_ids = torch.where(...), which only rebinds the local variable and does not mutate the caller’s tensor. This effectively makes padded-region masking a no-op on NPU. It also hardcodes-1instead of usingfill_value. Use an in-place update (e.g.,copy_) and respectfill_value.
_fill_padded_rows(topk_ids, num_token_non_padded, fill_value)
elif _is_npu:
# On NPU, bool-indexed scatter `topk_ids[bool_mask, :] = -1` lowers
# to aclnnNonzeroV2 and can trigger an aicore timeout under long
# workloads; `torch.where` avoids that nonzero scan.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| with ( | ||
| patch.object( | ||
| topk_module, | ||
| "get_moe_expert_parallel_world_size", | ||
| return_value=ep_size, | ||
| ), | ||
| patch.object( | ||
| topk_module, | ||
| "get_moe_expert_parallel_rank", | ||
| return_value=ep_rank, | ||
| ), | ||
| ): |
| Routed IDs: e -> e + e // num_local_routed (insert gaps for shared slots) | ||
| Shared IDs: shared_id_base + s (route to home rank) | ||
| Shared wgt: scale_factor (1.0 on aiter; 1/rsf otherwise) | ||
| """ |
| # Shared weight is 1.0 here because this branch is aiter-only: | ||
| # aiter_biased_grouped_topk folds routed_scaling_factor into the routed | ||
| # weights and forward_deepep skips the post-MoE multiply for _use_aiter, | ||
| # so the always-on shared expert must contribute 1.0x (matches the | ||
| # _remap_topk_for_deepep aiter branch, see PR #28237). |
| assert isinstance( | ||
| num_token_non_padded, torch.Tensor | ||
| ), "num_token_non_padded must be a torch.Tensor" | ||
| assert num_token_non_padded.numel() == 1, ( | ||
| "num_token_non_padded must be a single-element tensor, got shape " | ||
| f"{tuple(num_token_non_padded.shape)}" | ||
| ) | ||
| assert ( | ||
| not num_token_non_padded.dtype.is_floating_point | ||
| ), f"num_token_non_padded must be an integer tensor, got {num_token_non_padded.dtype}" | ||
| assert ( | ||
| num_token_non_padded.device == x.device | ||
| ), "num_token_non_padded and x must be on the same device" |
…functions The eager path now reads EP topology via get_parallel().moe_ep_size / moe_ep_rank (main's sgl-project#28567 refactor), so patch.object() on the removed get_moe_expert_parallel_world_size/rank names raised AttributeError. Force EP size/rank with the get_parallel().override() context manager.
Per review: the runtime environment always ships triton, so the try/except ImportError guard and _HAS_TRITON flag are unnecessary. Import triton/triton.language at the top like other core layers, drop the _HAS_TRITON gating on the kernel definition and _can_fuse_padded_region.
- Clarify shared-id docstring: shared_id_base + arange(S), not a constant. - Reword the aiter shared-weight comment: _remap_topk_for_deepep has no separate aiter branch; it unconditionally sets 1/routed_scaling_factor. - _fill_padded_rows: replace assert-based input validation with explicit raises so checks are not stripped under python -O.
|
@amd-bot ci-status |
CI Status for PR #28450Merge verdict: ❌ Do not merge — PR CI is incomplete and the changed code was never tested. All 10 red checks on HEAD Caution PR CI did not run. Every vendor pipeline (AMD, NVIDIA/base, Extra, MUSA, Xeon, XPU, NPU, Arm64) fast-failed at the gate because the author triggered CI again at 06:38 within the 120-min cooldown of the prior 04:52 trigger. The PR's changed kernel ( Changed files: Executed CI failure attribution: AMD: 0 executed failures · Others: 0 executed failures · all 10 red checks = gate/finish cascade from rate-limit cooldown (not tested). Gate cascade (root cause of all red checks)
Each Details / what to do before merge
Generated by amd-bot using Claude Code CLI |
|
@rbrugaro-amd could you please double confirm the result from this gsm8k command? |
|
@hubertlu-tw @HaiShaw Re-ran the full gsm8k eval with Server launch (env + command)# --- Server environment (PR recipe, incl. SGLANG_MORI_NO_PAD_MASK=1) ---
export PYTHONPATH=/workdir/code/sglang/python
export SGLANG_USE_AITER=1 SGLANG_AITER_MOE=1 AITER_FUSE_STAGE1_QUANT=1
export SGLANG_AITER_AR=1 SGLANG_AITER_MLA_PERSIST=1 SGLANG_ROCM_FUSED_DECODE_MLA=1
export SGLANG_USE_ROCM700A=1 TORCH_BLAS_PREFER_HIPBLASLT=1 TORCH_NCCL_BLOCKING_WAIT=1
export TRITON_ENABLE_PRELOAD_KERNELS=0 PYTORCH_ALLOC_CONF=expandable_segments:True
export ROCM_QUICK_REDUCE_QUANTIZATION=INT4 ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB=512
export MORI_SHMEM_MODE=ISOLATION MORI_SHMEM_HEAP_SIZE=16G
export SGLANG_MORI_DISPATCH_DTYPE=auto SGLANG_MORI_COMBINE_DTYPE=fp8
export SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK=4096 SGLANG_MORI_NO_PAD_MASK=1
# --- Server command ---
python -m sglang.launch_server \
--model-path /data/workloads-inference/models/DeepSeek-R1-0528-MXFP4/ \
--tp-size 8 --ep-size 8 --dp-size 8 \
--moe-a2a-backend mori --deepep-mode normal --moe-dense-tp-size 1 \
--enable-dp-attention --enable-dp-lm-head --enforce-shared-experts-fusion \
--mem-fraction-static 0.96 --kv-cache-dtype fp8_e4m3 \
--max-running-requests 22400 --max-prefill-tokens 4500000 \
--schedule-conservativeness 3.35 --num-continuous-decode-steps 192 \
--triton-attention-num-kv-splits 2 \
--disable-radix-cache --disable-chunked-prefix-cache \
--trust-remote-code --watchdog-timeout 6000 --host 127.0.0.1 --port 30000Eval command: python -m sglang.test.few_shot_gsm8k \
--num-questions 1319 --num-shots 5 \
--parallel 1319 \
--host 127.0.0.1 --port 30000Result: |
…sgl-project#28450) Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
Co-authored-by: Brayden Zhong <brayden@radixark.ai>
Motivation
On the AMD aiter / DeepEP-class MoE path, every layer runs the shared-expert
append (
fused_append_shared_experts) immediately followed by the eager DeepEPinterleaved remap (
_remap_topk_for_deepep). The remap is a sequence of small,launch-latency-bound elementwise ops (floor-div, add, arange, fill, copy) that
each cost a kernel launch. With one shared-expert append + remap per MoE layer
across many layers, this launch overhead is pure dispatch cost on top of work the
append kernel is already doing.
This PR fuses the append and the remap into the single Triton kernel that
already iterates over the topk rows, so the remap math runs on the rows already
resident in registers instead of as extra eager launches. It also fuses the HIP
padded-token masking into one Triton kernel.
The change is numerics-preserving by construction: the fused kernel produces
bit-identical
topk_ids/topk_weightsto the eager append+remap pair itreplaces (proven by the unit test below).
Modifications
python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe_triton_kernels.py_fused_append_remap_shared_experts_deepep_kerneland thefused_append_remap_shared_experts_deepepwrapper. In one pass it:e -> e + e // num_local_routed),shared_id_base + s),Host code computes the scalar remap params so the kernel stays branch-free.
python/sglang/srt/layers/moe/topk.py_post_process_topk_ids: when shared experts arefused and the backend is DeepEP-class on aiter, call the new fused kernel
instead of
fused_append_shared_experts()+_remap_topk_for_deepep(). Theaiter-only / deepep-only branches are preserved unchanged.
1.0: on aiter,aiter_biased_grouped_topkfoldsrouted_scaling_factorinto the routedweights and
forward_deepepskips the post-MoE multiply, so the always-onshared expert must contribute 1.0x. (This matches the aiter behavior of
_remap_topk_for_deepepaddressed in [AMD] fix(moe): correct fused shared-expert scaling on aiter/DeepEP path (mori all-to-all) #28237.)_fill_padded_rows/_fill_padded_rows_kernel: a single Triton launchthat zeroes/fills the CUDA/HIP-graph padded rows (
row >= num_token_non_padded,read from device memory), replacing the eager
arange + (>=) + index_put_sequence. Used by
_mask_topk_ids_padded_regionand_zero_topk_weights_padded_regionwith an eager fallback for non-2D /non-contiguous tensors.
re-zeros after the shared-expert append/remap; zeroing is idempotent).
SGLANG_MORI_NO_PAD_MASKenv gate (default off) to optionally skipthe HIP padded-mask pass entirely. This is numerics-affecting and must be
validated with an accuracy run before being enabled, hence off by default.
test/registered/moe/test_fused_append_remap_deepep.py(new, registered forboth CUDA and AMD CI)
test_matches_golden_reference— kernel output vs a pure-torch reference of thedocumented remap/append math, across shapes and EP ranks.
test_equivalence_with_eager_append_then_remap— bit-identical tofused_append_shared_experts()+_remap_topk_for_deepep().test_shared_weight_is_one_on_aiter_path.test_no_shared_experts_is_noop.Accuracy Tests
End-to-end (gsm8k): standard sglang gsm8k eval against a MoRI-EP DeepSeek-R1
server with this change, including the
SGLANG_MORI_NO_PAD_MASK=1mask-skip pathactive (8xMI355X, fp8 MoRI dispatch/combine, mem_fraction_static 0.96):
95.0% / 0% invalid on the full test set confirms both the fused append+remap and
the (default-off)
SGLANG_MORI_NO_PAD_MASK=1path are accuracy-clean onDeepSeek-R1.
Speed Tests and Profiling
The fused kernel collapses the per-MoE-layer append + eager-remap sequence
(~6 launch-bound elementwise kernels: floor-div / add / arange / fill / copy)
into the one append kernel that already runs, removing those launches per layer.
Before:
After
End-to-end
bench_servingA/B (8xMI355X, DeepSeek-R1 MoRI-EP, fp8 MoRIdispatch/combine,
SGLANG_MORI_NO_PAD_MASK=1;--dataset-name random --num-prompts 1000 --random-input-len 1024 --random-output-len 1024 --request-rate inf --seed 42). BEFORE = the pre-fix base image (eagerfused_append_shared_experts()+_remap_topk_for_deepep()); AFTER = with thefused kernel. The padded-token mask fusion is present in both images, so this
isolates the fused append+remap kernel. Mean of 3 runs each:
Per-run total throughput — BEFORE: 25613 / 25877 / 26181 tok/s; AFTER:
26448 / 27175 / 27179 tok/s. The improvement is consistent across every metric
and run.
Reproduction — the server is launched with the same env +
sglang.launch_serverflags shown in the Accuracy Tests section above (BEFORE from the pre-fix base
image, AFTER with the fused kernel; identical flags), then:
Checklist
CI States
Latest PR Test (Base): ❌ Run #28067472170
Latest PR Test (Extra): ❌ Run #28067472020