执行摘要
- 一句话：GLM5 共享专家 append 融合进 aiter topk，消除每层内核
- 推荐动作：值得精读，重点看三点设计：① 用持久预填充 buffer + row stride 部分列写入消除每层 kernel 启动，并与 CUDA graph 地址稳定性耦合；② 容量超限降级回 fallback 的“只影响覆盖、不影响正确性”设计哲学；③ review 中“条件镜像易碎、改用 shape 检测”的健壮性权衡，以及“server args 未初始化时优雅降级”的工程细节。同时应把“最终被 PR#35105 回滚”作为教训：特殊硬件路径优化应配套直接针对该分支的单元测试与更小粒度的 CI 验证，避免只依赖 nightly 精度套件。
功能与动机
PR body 指出：在非 EP aiter grouped-topk MoE 路径下，每个 decoder layer 都要运行独立的 _fused_append_shared_experts 内核把 shared expert 追加进 top-k ids/weights。这个 PR 的目标是“pre-populate a persistent top-k buffer's shared-expert columns once and let the aiter kernel write only the routed columns via row stride, so the per-layer append kernel is removed”，并且强调该做法与 plain append“Bit-identical”——shared 列使用与 aiter append 相同常数 fused_shared_experts_scaling_factor，因此对任何 shared-expert scaling 都正确。性能动机方面，PR 给出了 MI355X TP4 上 i1024/o1024 与 i8192/o1024 两组吞吐/TPOT 对比，结论是“removes one kernel launch per layer (3 kernels → fewer) and is bit-identical, so it is a clean simplification with a slight net-positive tendency at low concurrency”。
实现拆解
- 新增持久 buffer 与容量计算 helper（topk.py 顶部）：新增
_AITER_TOPK_FUSE_SHARED_MAX_TOKENS_CAP = 131072、_aiter_topk_fuse_shared_max_tokens_cache 和全局 dict _aiter_topk_fuse_shared_bufs。_get_aiter_topk_fuse_shared_max_tokens() 从 runtime_context.get_server_args() 读取 chunked_prefill_size / max_prefill_tokens（下限 8192、上限 131072）决定持久 buffer 的行数，结果缓存；未初始化时降级返回安全上限且不缓存，避免把“仅影响覆盖范围”的 sizing helper 变成硬崩溃。_get_aiter_topk_fuse_shared_buf() 按 (topk_routed, n_shared, num_experts, shared_weight, device) 缓存 [M, topk_routed + n_shared] 的 weight/id 张量，shared 列一次性预填充，固定 M 保证 CUDA graph 重放时地址稳定。
- 改造
biased_grouped_topk_gpu 的 aiter 分支：新增参数 fused_shared_experts_scaling_factor。满足 num_fused_shared_experts > 0、moe_ep_size == 1、token <= max_tokens 时启用 _shared_fuse：从持久 buffer 取切片 full_w[:token, :topk] 交给 aiter_biased_grouped_topk 只写 routed 列，随后返回全宽视图 full_w[:token], full_ids[:token]；不满足条件时仍走原 torch.empty 临时 buffer 路径，由后续 append 补齐。该分支同时作用于 prefill 与 decode。
_post_process_topk_ids 用 shape 检测跳过 append：原 _aiter_append 分支无条件调用 fused_append_shared_experts；现改为先判断 topk_ids.shape[1] == topk_config.top_k——若持久 buffer 已返回全宽（routed + shared）则跳过 append，否则照旧追加。采用 shape 检查而非镜像上游 _shared_fuse 条件（含 token <= MAX 边界），避免条件漂移导致静默缺少 shared expert。- 配套与测试情况：PR 早期新增的
get_global_server_args() 调用点触发 NVIDIA lint guardrail test_legacy_global_ratchet.py 失败，经 commit 99cbde2 改用 runtime_context.get_server_args() 后 CUDA/AMD CI 全绿；CPU 变体 biased_grouped_topk_cpu 曾为 API 对齐添加未使用参数，最终 commit 189b90ff 删除。本 PR 未附带直接单元测试，融合路径主要依赖 AMD nightly 精度套件（GSM8K）验证。
关键文件：
python/sglang/srt/layers/moe/topk.py（模块 专家路由；类别 source；类型 core-logic；符号 _get_aiter_topk_fuse_shared_max_tokens, _get_aiter_topk_fuse_shared_buf, biased_grouped_topk_gpu, _post_process_topk_ids）: 唯一变更文件，MoE 路由核心模块。新增持久 topk buffer 与容量计算 helper，改造 biased_grouped_topk_gpu 的 aiter 融合分支，并在 _post_process_topk_ids 中通过 shape 检测跳过 per-layer append，整体消除了 AMD aiter 路径每层一次的内核启动。
关键符号：_get_aiter_topk_fuse_shared_max_tokens, _get_aiter_topk_fuse_shared_buf, biased_grouped_topk_gpu, _post_process_topk_ids
关键源码片段
python/sglang/srt/layers/moe/topk.py
唯一变更文件，MoE 路由核心模块。新增持久 topk buffer 与容量计算 helper，改造 biased_grouped_topk_gpu 的 aiter 融合分支，并在 _post_process_topk_ids 中通过 shape 检测跳过 per-layer append，整体消除了 AMD aiter 路径每层一次的内核启动。
# python/sglang/srt/layers/moe/topk.py
# 持久 topk buffer 的上限与缓存：仅用于 aiter 非 EP 融合路径的容量，不影响正确性
_AITER_TOPK_FUSE_SHARED_MAX_TOKENS_CAP = 131072
_aiter_topk_fuse_shared_bufs: dict = {}
def _get_aiter_topk_fuse_shared_buf(
topk_routed: int, n_shared: int, num_experts: int, shared_weight: float, device
):
"""按 (topk_routed, n_shared, num_experts, shared_weight, device) 缓存持久 buffer。
buffer 形状固定为 [M, topk_routed + n_shared]，shared 列只预填充一次：
id = num_experts + i，weight = shared_weight。固定 M 保证 CUDA graph 重放时
张量地址稳定，这是该优化能与 CUDA graph 配合的关键前提。
"""
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)
# 预填充 shared 列：id 从 num_experts 开始连续编号，weight 为统一常数
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
# biased_grouped_topk_gpu 的 aiter 分支（节选）
elif _use_aiter:
assert not apply_routed_scaling_factor_on_output, "Not implemented"
token = gating_output.shape[0]
device = gating_output.device
# 仅非 EP + 有 fused shared experts + token 数在持久 buffer 容量内时启用融合
_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:
# shared 列权重必须与 _post_process_topk_ids 里 append 内核写入的常数一致，
# 才能保证与旧路径 bit-identical
_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
)
# aiter 内核只写 routed 列（[:, :topk]），shared 列保持预填充值不受扰动
topk_weights = full_w[:token, :topk]
topk_ids = full_ids[:token, :topk]
else:
# 大 prefill 或非适用场景回退普通临时 buffer，后续由常规 append 路径补 shared
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),
topk_weights,
topk_ids,
num_expert_group,
topk_group,
renormalize,
routed_scaling_factor if routed_scaling_factor is not None else 1.0,
)
if _shared_fuse:
# 返回全宽视图（routed 已写入 + shared 已预填充），_post_process 据此跳过 append
return full_w[:token], full_ids[:token]
return topk_weights, topk_ids
# _post_process_topk_ids 的 aiter append 分支（节选）
elif _aiter_append:
# 用 shape 判断 shared expert 是否已在 biased_grouped_topk_gpu 中融合：
# 融合时 topk_ids 已是全宽（routed + shared），shape[1] == top_k；
# 回退路径只有 routed 列，必须在此处 append。相比镜像上游 _shared_fuse 条件，
# shape 检查对调用链变化更健壮，不会因条件漂移产生静默缺共享专家的问题。
_shared_fused_in_topk = (
num_fused_shared_experts > 0
and get_parallel().moe_ep_size == 1
and topk_ids.shape[1] == topk_config.top_k
)
if not _shared_fused_in_topk:
M, N = router_logits.shape
scale_factor = (
1.0
if fused_shared_experts_scaling_factor is None
else fused_shared_experts_scaling_factor
)
# 延迟导入避免循环依赖
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
)
评论区精华
风险与影响
- 风险：
- 已被回滚（最大现实风险）：本 PR 合并后，PR#35105 以“回滚 #31323，恢复逐层 append，解除 CI 阻塞”为由 revert 了本改动，说明该优化在 main 上引发了 CI 阻塞（Issue 评论区 hzh0425 也指出某次 run 疑似被本 PR 弄挂）。恢复逐层 append 后风险解除。
- 全局可变状态缓存：
_aiter_topk_fuse_shared_bufs 是模块级 dict，进程生命周期内不释放；每个 key 对应约 [131072, topk + n_shared] * 8B 的显存，多配置组合（不同 topk/n_shared/device/shared_weight）会累积占用。
- shape 检查的隐式契约：
_post_process_topk_ids 用 topk_ids.shape[1] == topk_config.top_k 判断是否已融合，隐含“任何上游路径都不会先改变 topk_ids 宽度”的假设；未来若引入 remap/reshape 路径可能误判。
- CUDA graph 地址稳定依赖：持久 buffer 固定 M 是地址稳定的前提，但 fast path 覆盖的 graph 捕获尺寸若超过 buffer 行数会静默走 fallback，依赖
_post_process 的 shape 检查保持两路径一致，缺少直接测试锁定。
- 测试覆盖缺口：无直接单元测试；
biased_grouped_topk_gpu 的 aiter 分支与 _post_process_topk_ids 的跳过逻辑没有针对融合/回退/EP 模式的分支级验证。
- 影响：影响范围限定在 AMD gfx950 / MI355X + aiter + 非 EP（
moe_ep_size == 1）+ 带 fused shared experts（GLM-5.2）的 MoE 路由路径：消除每层一次 append 内核启动，TPUT +0.4%–0.9%、TPOT −0.1%–1.4%（PR 自测数据，clintg6 独立复测约 1%）。其他架构（CUDA/MUSA/CPU）、EP 模式及非适用路径一律 fallback 到原有 per-layer append，数值与行为不受影响（PR 声称 bit-identical）。对团队而言，该 PR 暴露了 AMD 专属路径优化与全局 CI 质量门（legacy accessor ratchet、NVIDIA GPU 阶段）之间的摩擦：一轮 guardrail 修复、一次 rerun、最终被回滚，说明这类改动需要更强的 CI 覆盖与更小的 main 集成面。
- 风险标记：已被回滚（PR#35105）, 缺少直接测试覆盖, 全局缓存增加显存, CI 门禁摩擦, shape 检查依赖隐式契约
关联脉络
- PR #35105 Revert "[AMD] [GLM5] Fuse shared-expert append into aiter grouped-topk (skip per-layer append kernel)": 直接回滚本 PR，恢复逐层 append 以解除 main CI 阻塞，是本 PR 合并后的关键后续。
- PR #31324 [AMD] [GLM5] Skip DSA decode indexer when kv_len <= index_topk (dense k-only fast path): 同为 GLM5 在 AMD 上的性能优化线，展示了该模型系列持续通过消除冗余内核/路径来降延迟的演进方向。
参与讨论