From 20fec44ce1bcf96fac60b2c81fa0a72603abfe9c Mon Sep 17 00:00:00 2001 From: Jacob Date: Thu, 9 Jul 2026 22:27:42 -0500 Subject: [PATCH 1/5] [AMD][GLM5] Fuse shared-expert append into aiter grouped-topk (skip per-layer append kernel) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Auto-enabled (no env) fast path for the non-EP aiter grouped-topk route. A persistent topk buffer has its shared-expert columns pre-populated once (id = num_experts+i, weight = the shared-expert scale_factor); the aiter kernel writes only the routed columns via row stride, so the per-layer _fused_append_shared_experts kernel is removed. Applies to both prefill and decode ("pre-populated", not the LLM prefill phase). - Bit-identical to the plain append: the shared column uses the same constant fused_shared_experts_scaling_factor the aiter append writes (threaded into biased_grouped_topk_gpu), so it is correct for any shared-expert scaling. - Persistent buffer sized to the max prefill batch (chunked-prefill-size / max-prefill-tokens, capped); for token counts above it, and in _post_process_topk_ids, the plain append path is used (mirrored condition) — no crash on large prefill batches. - Active when non-EP (moe_ep_size == 1) and num_fused_shared_experts > 0. Validated on GLM-5.2-MXFP4 TP4 (MI355X): GSM8K 0.945 (vs 0.948 without the path, within fp8 noise); decode trace confirms _fused_append_shared_experts is gone while grouped_topk_kernel runs. Co-authored-by: Cursor --- python/sglang/srt/layers/moe/topk.py | 148 +++++++++++++++++++++++---- 1 file changed, 130 insertions(+), 18 deletions(-) diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index a707790f0b31..470e92e5ee27 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -138,6 +138,69 @@ def routing( # an accuracy run before becoming the default. _skip_hip_pad_mask = get_bool_env_var("SGLANG_MORI_NO_PAD_MASK", "False") +# ATOM-style shared-expert fusion for the aiter grouped-topk path: keep a +# persistent topk buffer whose shared-expert columns are pre-populated once (NOT +# related to the LLM prefill phase — this applies to both prefill and decode), and +# let the aiter kernel write only the routed columns (via row stride). This removes +# the per-layer _fused_append_shared_experts kernel. +# Auto-enabled (no env) when: non-EP aiter path (moe_ep_size == 1) and +# num_fused_shared_experts > 0. Bit-identical to the plain append: the shared column +# is filled with the same constant scale_factor the aiter append writes. The +# persistent buffer is sized to the max prefill batch (chunked-prefill-size); for +# token counts above it we fall back to the plain path (condition mirrored in +# _post_process_topk_ids). Mirrors ATOM's init_aiter_topK_meta_data. +# Hard upper bound on the persistent buffer (safety cap when chunked prefill is +# disabled / unexpectedly huge); the buffer only costs ~[MAX, topk+n_shared] * 8B. +_AITER_TOPK_FUSE_SHARED_MAX_TOKENS_CAP = 131072 +_aiter_topk_fuse_shared_max_tokens_cache = None +_aiter_topk_fuse_shared_bufs: dict = {} + + +def _get_aiter_topk_fuse_shared_max_tokens() -> int: + """Max per-forward token count the persistent buffer must cover. Sized to the + largest prefill batch (chunked-prefill-size / max-prefill-tokens); decode is + always tiny (bs * num_tokens_per_bs). Above this we fall back to the plain + path, so this only bounds the fast-path coverage, not correctness. Cached + (server args are fixed after startup).""" + global _aiter_topk_fuse_shared_max_tokens_cache + if _aiter_topk_fuse_shared_max_tokens_cache is None: + from sglang.srt.server_args import get_global_server_args + + sa = get_global_server_args() + cps = getattr(sa, "chunked_prefill_size", None) or 0 + mpt = getattr(sa, "max_prefill_tokens", None) or 0 + m = max(int(cps), int(mpt), 8192) # 8192 floor for tiny configs + if int(cps) <= 0 and int(mpt) <= 0: + # chunked prefill disabled -> use the safety cap + m = _AITER_TOPK_FUSE_SHARED_MAX_TOKENS_CAP + _aiter_topk_fuse_shared_max_tokens_cache = min( + m, _AITER_TOPK_FUSE_SHARED_MAX_TOKENS_CAP + ) + return _aiter_topk_fuse_shared_max_tokens_cache + + +def _get_aiter_topk_fuse_shared_buf( + topk_routed: int, n_shared: int, num_experts: int, shared_weight: float, device +): + """Persistent [MAX, topk_routed + n_shared] weight/id buffers whose shared + columns are pre-filled once (id = num_experts + i, weight = shared_weight). + Fixed max size (>= max prefill batch) so the tensor address is stable across + CUDA-graph replays.""" + key = (topk_routed, n_shared, num_experts, float(shared_weight), str(device)) + buf = _aiter_topk_fuse_shared_bufs.get(key) + if buf is None: + total = topk_routed + n_shared + M = _get_aiter_topk_fuse_shared_max_tokens() + w = torch.empty((M, total), dtype=torch.float32, device=device) + ids = torch.empty((M, total), dtype=torch.int32, device=device) + ids[:, topk_routed:] = torch.arange( + num_experts, num_experts + n_shared, dtype=torch.int32, device=device + ).unsqueeze(0) + w[:, topk_routed:] = shared_weight + buf = (w, ids) + _aiter_topk_fuse_shared_bufs[key] = buf + return buf + if _is_cuda: try: @@ -1317,6 +1380,7 @@ def biased_grouped_topk_gpu( num_fused_shared_experts: int = 0, routed_scaling_factor: Optional[float] = None, apply_routed_scaling_factor_on_output: Optional[bool] = False, + fused_shared_experts_scaling_factor: Optional[float] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: num_tokens = gating_output.shape[0] num_experts = gating_output.shape[1] @@ -1438,8 +1502,34 @@ def biased_grouped_topk_gpu( assert ( hidden_states.shape[0] == gating_output.shape[0] ), f"Number of tokens mismatch: hidden_states.shape[0] = {hidden_states.shape[0]}, gating_output.shape[0] = {gating_output.shape[0]}" - topk_weights = torch.empty((token, topk), dtype=torch.float32, device=device) - topk_ids = torch.empty((token, topk), dtype=torch.int32, device=device) + _shared_fuse = ( + num_fused_shared_experts > 0 + and get_parallel().moe_ep_size == 1 + and token <= _get_aiter_topk_fuse_shared_max_tokens() + ) + if _shared_fuse: + # Persistent buffer with pre-populated shared columns. The weight is the + # constant scale_factor the aiter _post_process append writes (default + # 1.0), so this is bit-identical for any shared-expert scaling. aiter + # writes only the routed columns [:, :topk] via row stride; the shared + # column stays intact. If token exceeds the buffer size we drop to the + # plain path below (and _post_process appends shared experts as usual) — + # the same condition is mirrored there so the two paths stay consistent. + _shared_w = ( + 1.0 + if fused_shared_experts_scaling_factor is None + else fused_shared_experts_scaling_factor + ) + full_w, full_ids = _get_aiter_topk_fuse_shared_buf( + topk, num_fused_shared_experts, num_experts, _shared_w, device + ) + topk_weights = full_w[:token, :topk] + topk_ids = full_ids[:token, :topk] + else: + topk_weights = torch.empty( + (token, topk), dtype=torch.float32, device=device + ) + topk_ids = torch.empty((token, topk), dtype=torch.int32, device=device) aiter_biased_grouped_topk( gating_output, correction_bias.to(dtype=gating_output.dtype), @@ -1450,6 +1540,10 @@ def biased_grouped_topk_gpu( renormalize, routed_scaling_factor if routed_scaling_factor is not None else 1.0, ) + if _shared_fuse: + # Return the full [token, topk + n_shared] view (routed just written, + # shared pre-populated). _post_process_topk_ids skips its append. + return full_w[:token], full_ids[:token] return topk_weights, topk_ids elif _is_musa and ( gating_output.shape[1] // num_expert_group <= 32 @@ -1590,6 +1684,9 @@ def biased_grouped_topk_cpu( num_fused_shared_experts: int = 0, routed_scaling_factor: Optional[float] = None, apply_routed_scaling_factor_on_output: Optional[bool] = False, + # Accepted for API parity with the GPU variant (aiter shared-expert fusion + # fast path); unused on CPU. + fused_shared_experts_scaling_factor: Optional[float] = None, ): return torch.ops.sgl_kernel.biased_grouped_topk_cpu( hidden_states, @@ -1828,25 +1925,39 @@ def _post_process_topk_ids( num_local_routed, ) elif _aiter_append: - M, N = router_logits.shape - scale_factor = ( - 1.0 - if fused_shared_experts_scaling_factor is None - else fused_shared_experts_scaling_factor + # Must mirror biased_grouped_topk_gpu's _shared_fuse condition EXACTLY + # (incl. the token<=MAX buffer bound): when the fast path fell back to the + # plain buffer (e.g. token > MAX during a large prefill), the shared experts + # were NOT pre-populated, so we must append them here. + _shared_fused_in_topk = ( + num_fused_shared_experts > 0 + and get_parallel().moe_ep_size == 1 + and router_logits.shape[0] <= _get_aiter_topk_fuse_shared_max_tokens() ) + if _shared_fused_in_topk: + # Shared experts were already appended in biased_grouped_topk_gpu via + # the persistent pre-populated topk buffer; nothing to do here. + pass + else: + M, N = router_logits.shape + scale_factor = ( + 1.0 + if fused_shared_experts_scaling_factor is None + else fused_shared_experts_scaling_factor + ) - # Lazy import to avoid circular-import issues - from sglang.kernels.ops.moe.fused_moe_triton_kernels import ( - fused_append_shared_experts, - ) + # Lazy import to avoid circular-import issues + from sglang.kernels.ops.moe.fused_moe_triton_kernels import ( + fused_append_shared_experts, + ) - topk_ids, topk_weights = fused_append_shared_experts( - topk_ids, - topk_weights, - num_fused_shared_experts, - scale_factor, - N, # base id for shared experts - ) + topk_ids, topk_weights = fused_append_shared_experts( + topk_ids, + topk_weights, + num_fused_shared_experts, + scale_factor, + N, # base id for shared experts + ) elif use_per_rank_shared_slots: # DeepEP/MegaMOE: remap to per-rank shared-slot layout where each @@ -1941,6 +2052,7 @@ def select_experts( num_fused_shared_experts=num_fused_shared_experts, routed_scaling_factor=routed_scaling_factor, apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output, + fused_shared_experts_scaling_factor=topk_config.fused_shared_experts_scaling_factor, ) elif torch_native and custom_routing_function is None: assert ( From abb96d95d8529c2853a7f560acb39bb6c849ed48 Mon Sep 17 00:00:00 2001 From: Jacob Date: Thu, 16 Jul 2026 04:23:39 -0500 Subject: [PATCH 2/5] fix: guard get_global_server_args in aiter topk shared-fuse sizing _get_aiter_topk_fuse_shared_max_tokens() called get_global_server_args() unconditionally, which raises ValueError("Global server args is not set yet!") when the process-wide args are not published (e.g. a unit test or offline init that reaches the aiter grouped-topk path before startup). That turned a fast-path sizing helper into a hard crash even though its result only bounds coverage, not correctness. Wrap the access in try/except and fall back to the safety cap without caching, so a later call (once args are set) still computes and caches the real value. Co-authored-by: Cursor --- python/sglang/srt/layers/moe/topk.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 470e92e5ee27..8499d2f52fae 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -166,7 +166,16 @@ def _get_aiter_topk_fuse_shared_max_tokens() -> int: if _aiter_topk_fuse_shared_max_tokens_cache is None: from sglang.srt.server_args import get_global_server_args - sa = get_global_server_args() + try: + sa = get_global_server_args() + except Exception: + # Global server args not published yet (e.g. a unit test or offline + # init that reaches the aiter grouped-topk path before startup). + # Degrade gracefully instead of crashing -- this value only bounds + # the fast-path coverage, not correctness (see docstring). Return the + # safety cap WITHOUT caching, so a later call (once args are set) + # still computes and caches the real value. + return _AITER_TOPK_FUSE_SHARED_MAX_TOKENS_CAP cps = getattr(sa, "chunked_prefill_size", None) or 0 mpt = getattr(sa, "max_prefill_tokens", None) or 0 m = max(int(cps), int(mpt), 8192) # 8192 floor for tiny configs From 1f4fe9267dd6413cc5cda0ec8a6fcc56bc6ec496 Mon Sep 17 00:00:00 2001 From: Jacob Date: Mon, 20 Jul 2026 19:46:29 -0500 Subject: [PATCH 3/5] apply review: use shape check for aiter topk shared-fuse detection Replace the mirrored biased_grouped_topk_gpu _shared_fuse condition in _post_process_topk_ids with a robust topk_ids.shape[1] == topk_config.top_k check. When the shared experts are pre-populated in the persistent buffer, topk_ids already has the full (routed + shared) width, so the shape check reliably detects the fused case without fragile condition mirroring. --- python/sglang/srt/layers/moe/topk.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 8499d2f52fae..447f18cea723 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -1934,14 +1934,18 @@ def _post_process_topk_ids( num_local_routed, ) elif _aiter_append: - # Must mirror biased_grouped_topk_gpu's _shared_fuse condition EXACTLY - # (incl. the token<=MAX buffer bound): when the fast path fell back to the - # plain buffer (e.g. token > MAX during a large prefill), the shared experts - # were NOT pre-populated, so we must append them here. + # Detect whether the shared experts were already fused/appended in + # biased_grouped_topk_gpu via the persistent pre-populated topk buffer. + # When fused, topk_ids already has the full width (routed + shared), i.e. + # topk_ids.shape[1] == topk_config.top_k; when the fast path fell back to + # the plain buffer (e.g. token > MAX during a large prefill), only the + # routed columns are present and we must append the shared experts here. + # Checking the tensor shape is robust to the exact fast-path conditions + # (no fragile mirroring of biased_grouped_topk_gpu's _shared_fuse check). _shared_fused_in_topk = ( num_fused_shared_experts > 0 and get_parallel().moe_ep_size == 1 - and router_logits.shape[0] <= _get_aiter_topk_fuse_shared_max_tokens() + and topk_ids.shape[1] == topk_config.top_k ) if _shared_fused_in_topk: # Shared experts were already appended in biased_grouped_topk_gpu via From 99cbde28047deb26b7e84e19655739ee6b0c994b Mon Sep 17 00:00:00 2001 From: Jacob0226 Date: Mon, 3 Aug 2026 21:29:45 -0500 Subject: [PATCH 4/5] fix: use runtime_context server args in topk shared-fuse sizing Use runtime_context.get_server_args() in the aiter topk shared-fuse sizing helper so it follows the legacy accessor ratchet guardrail. Keep the uninitialized-server-args fallback behavior while avoiding new get_global_server_args() call-sites. Co-authored-by: Cursor --- python/sglang/srt/layers/moe/topk.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 8d7b30158d7b..e17bc2ba2771 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -164,11 +164,11 @@ def _get_aiter_topk_fuse_shared_max_tokens() -> int: (server args are fixed after startup).""" global _aiter_topk_fuse_shared_max_tokens_cache if _aiter_topk_fuse_shared_max_tokens_cache is None: - from sglang.srt.server_args import get_global_server_args + from sglang.srt.runtime_context import get_server_args try: - sa = get_global_server_args() - except Exception: + sa = get_server_args() + except ValueError: # Global server args not published yet (e.g. a unit test or offline # init that reaches the aiter grouped-topk path before startup). # Degrade gracefully instead of crashing -- this value only bounds From 189b90ffca12d071b9de0ee02c6926ee5e1ec54c Mon Sep 17 00:00:00 2001 From: HaiShaw Date: Sun, 16 Aug 2026 18:02:02 -0700 Subject: [PATCH 5/5] remove unused arg parameter --- python/sglang/srt/layers/moe/topk.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index b3ddd472aec0..67742d208adb 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -1831,9 +1831,6 @@ def biased_grouped_topk_cpu( num_fused_shared_experts: int = 0, routed_scaling_factor: Optional[float] = None, apply_routed_scaling_factor_on_output: Optional[bool] = False, - # Accepted for API parity with the GPU variant (aiter shared-expert fusion - # fast path); unused on CPU. - fused_shared_experts_scaling_factor: Optional[float] = None, ): return torch.ops.sgl_kernel.biased_grouped_topk_cpu( hidden_states,