Blackwell 上的 DeepGEMM FusedMoE（sglang）

• 基准版本：SGLang 最新 release v0.5.21。
• DeepGEMM 版本 sgl-deep-gemm==0.2.0

TL;DR

上篇博客我们介绍了 vLLM 上 DeepSeek V4 的 MoE 层 DeepGemm 相关实现，这篇博客我想走读一下 sglang 上 deepgemm MoE 层的 DeepGemm 使用。在 sglang 中，当 runner backend 选中 deep_gemm 后，整个前向是一套固定的流程，本文按 六大阶段 组织：

1. 权重预处理阶段（加载期一次性）：权重与权重 scale 整理成 DeepGEMM 最终消费的布局；
2. GEMM1 前：路由（TopK）→ dispatch → pre_permute，把 token 打成 grouped 布局并量化激活；
3. GEMM1 计算：w13 的 grouped GEMM（gate/up 投影）；
4. GEMM1 与 GEMM2 间：SiLU/SiTU 激活 + 下游激活量化（融合 kernel）；
5. GEMM2 计算：w2 的 grouped GEMM（down 投影）；
6. GEMM2 后处理：post_permute → combine → 模型尾部融合与归约。

四个可以直接对号入座的量化数据类型（统一由 DeepGemmMoeQuantInfo 表达，scale_recipes 决定 GEMM 行为）：

数据类型	w13 / w2 dtype	权重布局 (E,N,K)	block_shape	scale recipes (A/B)	scale 类型
BF16	bf16	(E,2I,K) / (E,K,I)	—	—	—
FP8 block（含 per-tensor 展开）	float8_e4m3fn	(E,2I,K) / (E,K,I)	[128,128]	(None,None)	UE8M0 int32（if DEEPGEMM_SCALE_UE8M0）或 fp32
MXFP8	float8_e4m3fn	同上	[1,32]	(1,32)/(1,32)	必须 UE8M0
MXFP4 experts（W4A8，Blackwell 默认）	int8（packed 2×fp4）	(E,2I,K//2) / (E,K,I//2)	[128,128]	(1,128)/(1,32)	必须 UE8M0

如何走到 deep_gemm：

方法一：--moe-runner-backend deep_gemm 直接显示指示走 deep gemm

方法二：--moe-a2a-backend ∈ {deepep, mooncake, nixl} + ENABLE_JIT_DEEPGEMM
fp8.py:1334-1356

方法三：--moe-a2a-backend deepep_v2 会强制使用 deep_gemm，runner.py:85-94

Blackwell 相关环境变量：

常量	成立条件	影响
ENABLE_JIT_DEEPGEMM	sm≥90 ∧ SGLANG_ENABLE_JIT_DEEPGEMM（默认 True）∧ 包可导入；SM120/121 额外要求探测到 fp8_einsum 等 API（configurer.py:17-29,44-45）	全链路总开关
DEEPGEMM_BLACKWELL	ENABLE_JIT_DEEPGEMM ∧ is_sm100（capability[0]==10，含 100/103）	Blackwell 专属路径
DEEPGEMM_SCALE_UE8M0	ENABLE ∧ (sm100 ∨ sm120/121)	scale 以 UE8M0 打包为 int32 存储
DEEPGEMM_NEED_TMA_ALIGNED_SCALES	not (UE8M0 ∨ MUSA)	非 UE8M0 时 scale 需 TMA 对齐

注：SGLang 中 is_sm100 即 is_sm100_supported（capability[0]==10 且 CUDA ≥ 12.8，utils/common.py:298-302），把 SM100 与 SM103 都算进来；is_sm120()/is_sm121() 分别对应 (12,0)/(12,1)（utils/common.py:326-334）。本文行文中的“SM100 / Blackwell”均指前者。

1. 全景：从 Decoder 到 DeepGEMM kernel

MoE 层的全局计算：每个 token 同时经过 routed 与 shared 两条分支，最后相加：

图：DeepSeek V4 MoE 层全局计算流（V4 默认路径）。每个 token 先经 gate 打分、Top-K（V4 默认 6/256）选出专家；routed 分支的 token 经 dispatch 按目标专家分组（EP 部署下即 all2all），在被选中的专家上依次执行 gemm1（w13，gate/up 投影）→ SiLU/SiTU 激活 → gemm2（w2，down 投影），最后 combine 按 topk 权重加权合并（含 routed_scaling_factor）；shared experts 分支对全量 token 始终计算一遍 MLP，两条分支的结果相加即为 MoE 层输出。V4 默认专家权重 MXFP4、激活 FP8（W4A8）。下面 §2 的六大阶段与 §3 的 DeepGEMM 接口，都对应这张图里的具体环节。

1.1 六阶段速览

阶段	位置（v0.5.21）	关键函数 / 类	输入 → 输出
① 权重预处理（加载期）	layers/quantization/fp8.py:2004, 2743-2801；mxfp4.py:1590-1606	process_weights_after_loading_block_quant / _process_mxfp8_moe_weights / create_moe_runner；运行期 apply 组装 DeepGemmMoeQuantInfo	checkpoint 权重 → deep_gemm 消费布局 + runner 选定
② GEMM1 前	models/deepseek_v2.py:916；fused_moe_triton/layer.py:1562；moe_runner/runner.py:195；moe_runner/deep_gemm.py:970/1194/1312/1354/1698	TopK → dispatcher.dispatch → pre_permute_*	hidden [M,H] bf16 → DeepGemmRunnerInput（masked 或 contig 激活 + scale）
③ GEMM1 计算	moe_runner/deep_gemm.py:455-462 / 771-779	grouped_gemm_nt_f8f8bf16_{contig,masked}	激活 + w13 → gateup (·,2I) bf16
④ GEMM1 与 GEMM2 间	moe_runner/deep_gemm.py:470-588 / 790-824	SiLU/SiTU + sglang_per_token_group_quant_fp8 / *_post_quant 融合	gateup bf16 → down 输入 fp8 + scale
⑤ GEMM2 计算	moe_runner/deep_gemm.py:597-621 / 830-894	第二次 grouped_gemm_nt_f8f8bf16_*	down 输入 + w2 → (·,K) bf16
⑥ GEMM2 后处理	moe_runner/deep_gemm.py:1236/1288/1338/1457/1847；fused_moe_triton/layer.py:1589；deepseek_v2.py:1368-1393	post_permute_* → dispatcher.combine → 尾部融合 + post_experts_all_reduce	→ hidden [M,H] bf16

1.2 调用图

DeepseekV4DecoderLayer.forward()                        models/deepseek_v4.py:2604
 ├── self.self_attn(...)                                 （DP attention 时 MoE 前向另有同步处理）
 └── self.mlp(...) = DeepseekV2MoE                       models/deepseek_v2.py:2644 / :916
      ├── self.topk(hidden_states) → TopKOutput
      └── 分支：
           ├── forward_deepep()        （a2a 使能）       :1450
           └── forward_normal()                            :1236
                ├── self.experts(hidden_states, topk_output)      ← MoE 实现
                │     = FusedMoE / DeepEPMoE（get_moe_impl_class 决定）
                │     └── forward_impl()                   fused_moe_triton/layer.py:1562
                │          ├── ① dispatch_output = self.dispatcher.dispatch(...)
                │          │      = StandardDispatcher | DeepEPDispatcher(normal/ll/v2) | FlashinferDispatcher
                │          ├── ② runner_output = self.run_moe_core(dispatch_output)
                │          │      └── MoeRunner.run()      moe_runner/runner.py:195
                │          │           ├── fused_func（deep_gemm 无，跳过）
                │          │           ├── pre_permute_*()  → DeepGemmRunnerInput
                │          │           │     moe_runner/deep_gemm.py: 970 / 1194 / 1312 / 1354 / 1698
                │          │           ├── DeepGemmRunnerCore.run()      :359
                │          │           │     ├─ _run_contiguous_gemm()        :395
                │          │           │     ├─ _run_bf16_contiguous_gemm()   :628
                │          │           │     ├─ _run_masked_gemm()            :695
                │          │           │     └─ _run_masked_bf16_gemm()       :898
                │          │           │           └── deep_gemm_wrapper.entrypoint.grouped_gemm_nt_*
                │          │           │                 └── deep_gemm.*  （sgl-deep-gemm JIT kernel）
                │          │           └── post_permute_*() → *CombineInput
                │          └── ③ combined = self.dispatcher.combine(runner_output)
                └── 尾部：routed-scale / shared-expert 融合 + post_experts_all_reduce   :1368-1393

2. 六大阶段拆解

2.1 ① 权重预处理阶段

这一块只在权重加载期执行一次（外加运行时对 quant_info 组装），目的是让 deep_gemm 在运行期只做纯 GEMM，不再碰权重布局。

2.1.1 谁触发了 deep_gemm runner

两个入口：一个在权重加载期，一个在runner 创建期。

1. 加载期：Fp8MoEMethod.process_weights_after_loading_block_quant 的 CUDA 部分实现里，会先判断 deepgemm 是否会被使用 will_use_deepgemm = is_deepgemm_moe_runner_backend_enabled()。这个布尔值决定权重与 scale 是否按 DeepGEMM 所需布局预处理，（具体细节详见下文 2.1.2）。
2. 运行期（runner 创建）：Fp8MoEMethod.create_moe_runner）：
   • auto：若 is_deepgemm_moe_runner_backend_enabled() 为真 → MoeRunnerBackend.DEEP_GEMM，否则 TRITON（:2764-2774）；
   • 显式 deep_gemm 等后端直接触发

is_deepgemm_moe_runner_backend_enabled（fp8.py:1334-1356）判定：

• moe_runner_backend == DEEP_GEMM → True（显式指定）；
• AUTO ∧ moe_a2a_backend ∈ {DEEPEP, MOONCAKE, NIXL} ∧ ENABLE_JIT_DEEPGEMM → True；
• 否则 False。

另外两个来源的强制 / 例外：

• deepep_v2 强制：runner.py:85-94 校验 deepep_v2 必须搭配 deep_gemm runner；
• V4 例外（arg_groups/model_overrides/deepseek_v4.py:66-88）：nvfp4_moe_meta 时改为 flashinfer_trtllm_routed；fp4 experts ∧ a2a=none ∧ SM90/100/120 时改为 flashinfer_mxfp4；
• SGLANG_DSV4_FP4_DEQUANT=1（fp8.py:414-418）：只允许 auto 起步（fp4 反量化成 fp8 运行的对照路径）。

2.1.2 四条权重路径

现在我们仔细看一下 Fp8MoEMethod.process_weights_after_loading_block_quant 的 CUDA 实现细节：

class Fp8MoEMethod(FusedMoEMethodBase):
  def process_weights_after_loading_block_quant():
    will_use_deepgemm = self.is_deepgemm_moe_runner_backend_enabled()
    if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 and will_use_deepgemm:
        from deep_gemm import transform_sf_into_required_layout

        for scale_param, weight_param in [
            (layer.w13_weight_scale_inv, layer.w13_weight),
            (layer.w2_weight_scale_inv, layer.w2_weight),
        ]:
            num_experts, n, _ = scale_param.data.shape
            k = weight_param.shape[2] * 2
            scale_param.data = transform_sf_into_required_layout(
                scale_param.data,
                mn=n,
                k=k,
                recipe=(1, 32),
                num_groups=num_experts,
                disable_ue8m0_cast=False,
            )
        layer.w13_weight_scale_inv.format_ue8m0 = True
        layer.w2_weight_scale_inv.format_ue8m0 = True

第一行 will_use_deepgemm = is_deepgemm_moe_runner_backend_enabled() 做判断，然后按权重类型分四条路。

路径 A：MXFP4 experts（Blackwell 默认路径，is_fp4_expert=True）

1. 先把权重 view 成 int8（marlin / NPU / megamoe 各有早退分支）。
2. if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 and will_use_deepgemm:，对 w13/w2 分别调用：

scale_param.data = transform_sf_into_required_layout(
    scale_param.data,
    mn=n,                      # w13: 2I；w2: K
    k=k,                       # 需要传入真实的 K = weight.shape[2] * 2，packed int8 还原
    recipe=(1, 32),            # MXFP4 权重 scale：K 方向每 32 个元素一份
    num_groups=num_experts,
    disable_ue8m0_cast=False,  # 允许就地 cast 成 UE8M0
)

即把 (E, N, K/32) 的 fp32 scale 转成 DeepGEMM 要求的 MN-major、TMA 对齐、UE8M0 打包布局，并打 format_ue8m0=True。MXFP4 的 scale 本来就是 2 的幂，cast 无精度损失。输出的权重本体保持 int8 packed。

路径 B：FP4 反量化（SGLANG_DSV4_FP4_DEQUANT=1 ）为 FP8

class Fp8MoEMethod(FusedMoEMethodBase):
  def process_weights_after_loading_block_quant():
    if self.is_fp4_expert and self.dequant_fp4_to_fp8:
        for weight_param, scale_param in [
            (layer.w13_weight, layer.w13_weight_scale_inv),
            (layer.w2_weight, layer.w2_weight_scale_inv),
        ]:
            num_experts = weight_param.shape[0]
            new_weights = []
            new_scales = []
            for e in range(num_experts):
                w, s = cast_e2m1fn_to_e4m3fn(
                    weight_param.data[e], scale_param.data[e]
                )
                new_weights.append(w)
                new_scales.append(s)
            weight_param.data = torch.stack(new_weights)
            scale_param.data = torch.stack(new_scales).float()
            scale_param.format_ue8m0 = False
        self.is_fp4_expert = False
        logger.warning_once("Dequantized FP4 expert weights to FP8.")

代码会逐个 expert 调 cast_e2m1fn_to_e4m3fn(w[e], s[e])，把 fp4 无损转 fp8（e2m1 → e4m3 是精度扩展，最多补 6 个尾数位，MAX_OFFSET_BITS=6），它会先将 int8 中的 mxfp4 数据单独提取出来，分别做 dequant，然后将量化组改为 128 个一组，再重新扩展到 FP8 的量级 + offset。随后 format_ue8m0=False、self.is_fp4_expert=False，交回路径 C ，做通用 block-FP8 逻辑继续处理。

路径 C：普通 block-FP8（is_fp4_expert=False）

class Fp8MoEMethod(FusedMoEMethodBase):
  def process_weights_after_loading_block_quant():
    if not self.is_fp4_expert:
        weight_block_size = self.quant_config.weight_block_size
        for weight, weight_scale in (
            (layer.w13_weight, layer.w13_weight_scale_inv),
            (layer.w2_weight, layer.w2_weight_scale_inv),
        ):
            requant_block_scale_ue8m0_for_deepgemm(
                weight,
                weight_scale,
                weight_block_size,
                use_deepgemm_runner=will_use_deepgemm,
                output_dtype=torch.bfloat16,
                weight_shape=weight.shape[-2:],
            )

对 w13/w2 分别调用 requant_block_scale_ue8m0_for_deepgemm(...)。该函数的进入条件是四连与：

• use_deepgemm_runner（确实选了 DeepGEMM runner）；
• weight_block_size == [128, 128]（权重量化 kernel 只支持这一种）；
• not format_ue8m0（尚未重排过）；
• should_deepgemm_weight_requant_ue8m0(...)（加载器侧判定：输出 dtype、形状对齐是否满足）。

满足则 requant_weight_ue8m0_inplace → 对 3D MoE 权重做 requant，见 requant_weight_ue8m0，并走 _requant_weight_ue8m0_grouped，逐 expert 三步：

1. block_quant_dequant：fp8 权重 + scale 还原为 bf16；
2. quant_weight_ue8m0 → per_block_cast_to_fp8：重新按 128×128 block 量化成 fp8，scale 保持 fp32；
3. transform_scale_ue8m0(out_s, mn=n)：先 index_select(-2, arange(mn)//128) 把 mn 方向 128 行共享一份 scale（MN-major packing 的粒度要求），再 get_mn_major_tma_aligned_packed_ue8m0_tensor 将 fp32 scale 打包成 UE8M0 int32、MN-major、TMA 对齐；返回时修复 DLPack 在 K ≤ 128×4 时塌缩的尾维 stride。

完成后 weight_scale.format_ue8m0 = True，运行时无需再转换。注意这是数值重量化（block 量化→反量化→重量化），目的是为了将 FP8 128 blockwise 的数据重新变为 DeepGemm 要求的 FP8 quant 类型，scale 粒度依然是 128x128，但 scale 存储会变为 MN-major packed 的 int32 存储方式。

路径 D：MXFP8（use_mxfp8=True）

最后我们来看 mxfp8 量化的处理方法，在 _process_mxfp8_moe_weights(layer, quantize=not is_checkpoint_fp8_serialized)。先平台断言：要求 SM100（或 ROCm gfx95）。随后按来源分两种：

• quantize=True（checkpoint 非 fp8 序列化，需要权重在线量化）：
  • deep_gemm 分支 _quantize_for_deepgemm 把 (E, N, K) 拍平成 2D (E·N, K)
  • mxfp8_group_quantize 调用 triton 内函数，沿 K 轴每 32 个元素一组做标准 MX 量化
  • _ue8m0_to_fp32 把 8 位指数放回 fp32 的 exponent 位，得到"值为 2 的幂的 fp32"
  • _pack_moe_scale_for_deepgemm DEEPGEMM_SCALE_UE8M0 时走 get_mn_major_tma_aligned_packed_ue8m0_tensor 打包成 MN-major 的 int32 TMA 行对齐的 scale。
  • 一句话，就是将原先的权重变为了 deepgemm 需要的 FP8 quant 格式。
• quantize=False（checkpoint 已是 MXFP8）：deep_gemm 分支 _convert_ue8m0_scales_for_deepgemm（:2342-2350 → :2264-2271）：u8 scale → fp32 → 同样打包。

结尾 _copy_or_rebind（形状/dtype 一致则原地 copy，否则重绑参数）并设置 format_ue8m0=True 表示权重已经完成处理。MXFP8 权重布局保持 (E, 2I, K) fp8 e4m3，仅 scale 变为打包布局。

注：MXFP8 的 weight_block_size 通常都是 [1, 32] 的颗粒度

2.1.3 运行时组装：DeepGemmMoeQuantInfo

每次前向 Fp8MoEMethod.apply 在 runner_backend.is_deep_gemm() 分支）组装量化信息：

• block_quant（block-FP8 / MXFP8）：block_shape = quant_config.weight_block_size，scale 直接引用 w13_weight_scale_inv / w2_weight_scale_inv（2.1.2 已就绪）；
• per-tensor（weight_block_size is None）时：把 per-tensor scale 展开成 block 伪 scale ——

  w13_scale = w13_weight_scale.unsqueeze(1).repeat_interleave(ceil(N/128), dim=1)
                              .unsqueeze(2).repeat_interleave(ceil(K/128), dim=2)

  这段代码关键，它将形状 (E, ceil(N/128), ceil(K/128))、block_shape=[128,128]。数值上每个 128×128 block 共享同一 per-tensor scale，这样做纯粹是为了让同一套 GEMM 接口吃下 per-tensor 检查点。

DeepGemmMoeQuantInfo 字段与约束：

class DeepGemmMoeQuantInfo(MoeQuantInfo):
  w13_weight, w2_weight           # (E, 2I, K) / (E, K, I)，或 packed 变体
  use_fp8                         # False → bf16 路径
  w13_scale, w2_scale
  block_shape                     # [128,128] / [1,32] / None
  is_fp4_experts: bool = False    # DSV4 MXFP4 布局旗标
  use_mxfp8: bool = False

__post_init__ 断言：use_mxfp8 必须 block_shape == [1, 32] 且 DEEPGEMM_SCALE_UE8M0=True。

关键方法 scale_recipes(...) 把字段翻译为 GEMM 的 A/B（激活/权重）recipe：

情形	返回 (A, B)	说明
use_mxfp8	((1, activation_block_size), (1, 32))	断言 ceil_div(K, ab*4) == activation_scale_width（打包后列数匹配）
is_fp4_experts	((1, 128), (1, 32))	激活 fp8 每 128 一份；权重 fp4 每 32 一份
其余	(None, None)	交给 deep_gemm 默认 128×128 处理

MXFP4 若走 Mxfp4MoEMethod.apply（mxfp4.py:1578-1606），则等价地硬编码 DeepGemmMoeQuantInfo(block_shape=[128,128], is_fp4_experts=True)，但实际 sglang 实现中没有加入 DeepGemm FP4 x FP4 的处理计算逻辑

另外记几个后续阶段会用到的 MoeRunnerConfig 字段（参考代码 moe_runner/base.py:37-69）：

• swiglu_limit gemm1_clamp_limit float，对激活函数的输出值做一些限制，防止产生极大异常值甚至溢出。在低精度量化下，MoE 路由或矩阵乘法容易产生极大的异常值（Outliers）。如果不加限制，这些异常值会导致量化缩放因子（Scale）失调，甚至引发 FP8/FP4 计算溢出（Overflow）。强制钳位可以截断极端值，保障底层 MoE Kernel 的数值稳定性
• gemm1_alpha float，SwiGLU 激活函数中 SiLU 部分的缩放系数
• silu_mul_keep_fp32 bool，表示是否需要用 FP32 来表示 SiLU 结果
• top_k int，MoE 路由后，决定每个 Token 被激活并分配给的 Experts 的数量
• num_local_experts int，在当前并行策略下，当前卡分配到的本地 Experts 数量
• intermediate_size_per_partition：int，当前并行策略下，Experts 的中间隐藏层大小

2.2 ② GEMM1 前：路由、dispatch 与进料

2.2.1 模型层：路由与三段式的入口

DeepSeek V4 的 Decoder 层仍然是变种的 transformer，先跑 attention，再调用 self.mlp(hidden_states) 计算 MoE/MLP 层（deepseek_v4.py:2604，self.mlp = deepseek_v2.DeepseekV2MoE）。

若使用 data parallel attention，DP attention 场景下 MoE 前向会被包进 _run_moe_ffn_dp_sync，调用点发生在 DP attention 与 EP MoE 之间做输入/输出 gather 或 scatter（详细见 §2.6.3）。

进入 DeepseekV2MoE.forward：

1. 一般地，self.topk(hidden_states) 产出 TopKOutput（topk_ids / topk_weights）。针对 DeepSeek V4 的路由参数，有特别的 override ：use_grouped_topk=False、sqrtsoftplus_log1p、fused_gate_packed_ids（model_type == "deepseek_v41" 时），并把 quant_config.is_fp4_experts 透传给 TopK）——让路由输出格式对齐 FP4 MoE 后端的约束。
2. 选择走向：
   • should_use_mega_moe → MegaMoE 融合路径（不在本文范围）；
   • _enable_a2a_moe，a2a ∈ {deepep, mooncake, nixl, mori, ascend_fuseep, flashinfer, deepep_v2}）→ forward_deepep）走 EP 并行
   • 否则，forward_normal

experts 由 get_moe_impl_class(quant_config)）决定（构造于 deepseek_v2.py:665-679）：a2a ∈ {mori, deepep, deepep_v2, mooncake, nixl, pplx} 时是 DeepEPMoE，否则是 FusedMoE。V4 走 DeepGEMM 时通常 a2a=deepep，因此实例是 DeepEPMoE。

前向尾部 forward_normal 主要处理残差和 allreduce 处理，下文 §2.6.3 统一说明。

2.2.2 FusedMoE 的三段式骨架

MoE 实现类的前向主流程在 FusedMoE.forward_impl 中实现（fused_moe_triton/layer.py）：

# 与源码对应的伪代码
dispatch_output = self._dispatch_with_pre_quant(hidden_states, topk_output)
runner_output   = self.run_moe_core(dispatch_output)
combined        = self.dispatcher.combine(runner_output)
if self.reduce_results and (self.moe_tp_size > 1 or self.moe_ep_size > 1):
    combined = tensor_model_parallel_all_reduce(combined)

• _dispatch_with_pre_quant：dispatch 前可选的 pre-quant（是否预量化取决于 dispatcher 与配置）；
• run_moe_core调用 MoeRunner.run，执行真正 MoE 计算
• dispatcher 由 create_moe_dispatcher）按 a2a backend 五分支选出：standard / deepep（normal+LL 共用 MaybeTboDeepEPDispatcher）/ deepep_v2 / flashinfer。

而 DeepEPMoE）在这里的角色是：当 deep_gemm runner + Fp8Config 组合下为 True（其余条件还包括 deepep_v2 / humming / aiter / npu 等）时， DeepEPMoE.forward_impl）会使用 super().forward_impl(...)，复用上面的现代三段式；否则才走它自带的旧式 dispatch/combine 实现。这就是为什么 DeepSeek V4 的 EP 形态最终与 FusedMoE 汇到同一段代码。

2.2.3 MoeRunner：pre_permute → core → post_permute 的编排

MoeRunner.run：略去本文不关心的代码后，主体结构就是三句话：

class MoeRunner
  def run(
      self, dispatch_output: DispatchOutput, quant_info: MoeQuantInfo, lora_info=None
  ) -> CombineInput:
    pre_permute_fn, post_permute_fn = PermuteMethodPool.get((dispatch_format, runner_backend))
    runner_input  = pre_permute_fn(dispatch_output, quant_info, runner_config, running_state)
    runner_output = self.runner.run(runner_input, quant_info, running_state)      # DeepGemmRunnerCore
    combine_input = post_permute_fn(runner_output, quant_info, runner_config, running_state)

PermuteMethodPool 以 (dispatch_format, runner_backend) 为字典的方式注册），如表：

dispatch_format	pre_permute（moe_runner/deep_gemm.py）	触发场景
standard	pre_permute_standard_to_deep_gemm	无 a2a / TP-MoE / SM120
flashinfer	pre_permute_flashinfer_to_deep_gemm	FlashInfer one-sided A2A
deepep_ll	pre_permute_deepep_ll_to_deep_gemm	DeepEP 低延迟模式
deepep_normal	pre_permute_deepep_normal_to_deep_gemm	DeepEP 普通模式
deepep_v2	pre_permute_deepep_v2_to_deep_gemm	DeepEP v2（强制 deep_gemm）

五种都在做同样三件事：激活量化（如未量化）→ 布局重排 → 写 running_state。产物统一是 DeepGemmRunnerInput：

class DeepGemmRunnerInput(RunnerInput):
  hidden_states: torch.Tensor                 # grouped 后的激活（Gemm1 的 lhs）
  hidden_states_scale: torch.Tensor
  use_masked_gemm: bool
  masked_m: Optional[torch.Tensor]            # masked 专用：(E,) int32
  expected_m: Optional[int]                   # masked 专用：调优提示
  m_indices: Optional[torch.Tensor]           # contig 专用：(M,) int32 行→expert 映射
  hidden_states_scale_tma_aligned: bool
  activation_scale_block_size: Optional[int]  # 激活 K 方向的量化组大小（与实际输入量化一致）

2.2.4 五种 pre permute 方式

1）standard：

• 先处理 SM120：deep_gemm_sm120.maybe_pre_permute(...) 若命中直接返回；命中不了则 assert 拒绝，因为 SM120 DeepGemm 的 grouped GEMM 只能支持 standard-layout activation 布局。

• masked 布局，处理后会直接喂给 DeepGEMM masked grouped GEMM，所谓的 masked grouped GEMM 接口可以看第三节对 DeepGEMM 的简介，简单来说 masked 接口是用内存换简单（小 batch / capture 默认），它给每个 Experts 分一个简单固定的内存空间，用不上的打 masked 遮掩，当条件允许时，moe_ep_deepgemm_preprocess 会一步完成“按 expert 对齐 slab + 激活量化”，产出 masked_m / src2dst / hidden / scale；以及 expected_m。他们分别指代：

  • masked_m：每个 local experts 分给真实有效行数，即 experts 对输入 在 3D 布局 [E_local, m_max, K] 中真正需要计算的行数。通过 fused_moe_dispatch_index_triton_kernel 计算得到
  • expected_m：每个 experts 在当前 hidden_states 的大小下分配到的 token 数的期望，调优 gemm 时使用
  • src2dst 记录 permute 过程中当前 rank 如何重排 token 的表，同样也是 fused_moe_dispatch_index_triton_kernel 计算得到

• compact 布局 处理 DeepGEMM contiguous grouped GEMM，contiguous grouped GEMM 后面我们也有说到。在处理更大的 token 时，contiguous gemm 会更省内存（因为它只会保留真实有效的 token），token 也是按照 [M_total, K] 专家连续平铺。

  • get_contiguous_layout_alignment，DeepGEMM contiguous 布局要求的 MK 对齐（每个专家段的起始和长度都必须是它的倍数）。_get_compact_all_tokens(...)，这是在计算最坏路由下的 padding 上界，保证不会越界。
  • fused_moe_dispatch_index 统计每 expert 需要的 token 数，但还需要 pad_expert_counts 对齐到 contiguous 的要求和 all_tokens 的对齐要求
  • 量化：sglang_per_token_group_quant_fp8 做 per-token-group 量化。使用 UE8M0 时 (scale_width, all_tokens) 使用 int32 4个一组，列主序布局方式；否则 (all_tokens, scale_width) 使用 fp32；
  • ep_scatter(...)，expert_alignment=block_e）散到 expert 连续段；
  • 返回 running_state：src2dst、all_tokens、contiguous_layout_alignment=block_e

  何时 masked、何时 compact 由 _should_use_masked_standard_layout 决定：先看 clamped/swizzled 激活是否被支持 _masked_activation_unsupported_reason：swiglu_limit 或 megamoe 场景要求 D//8 ≥ E、group_size==128、D % (4*group)==0）、再看 SGLANG_DEEPGEMM_STANDARD_LAYOUT 是否强制、否则按内存预算 _estimate_masked_standard_layout_peak_bytes vs free_memory × fraction）自动选择（capture 模式禁 masked）。详见附录 B。

2）flashinfer 部分不细聊了，它仅接受 BF16 未量化 dispatch，后转交 standard 实现。

接下来是针对 DeepEP 模式的处理，先回顾一个前提：三种 DeepEP 模式的差别，本质是 dispatch（all-to-all 收发）交给 pre_permute 的"半成品"长什么样。pre_permute 要做的，就是把这个半成品整理成 DeepGEMM 能直接使用的 DeepGemmRunnerInput——要么是 masked [E_local, m_max, K]，每 expert 一块固定空间、用 masked_m 标注真实行数，要么是 contiguous 平铺 [all_tokens, K]，各 expert 段紧挨着、用 m_indices 标注每行属于哪个 expert。

3）deepep_ll（低延迟模式）：dispatch 直接把结果摆成了 masked 布局，并一同计算了 masked_m 和 expected_m。所以这个函数几乎什么都不做，原样透传，设置 use_masked_gemm=True，activation_scale_block_size 固定写 128（LL dispatch 出口就是 128 分组的 FP8）即可。

4）deepep_normal（普通模式）：与 LL 相反，dispatch 只给"按到达顺序堆在一起的 token"，外加一个 num_recv_tokens_per_expert（每个 expert 各收到多少条）。这种显然与 contiguous 布局更接近，但需要自己把这些乱序 token 按 expert 重排成连续段：

• 先算总量 all_tokens = sum(num_recv_tokens_per_expert)），据此分配输出激活缓冲 (all_tokens, K)
• scale 形状分两种：Blackwell 的 UE8M0 走 int32 打包 + 列主序 (ceil_div(K//128,4), all_tokens).transpose(0,1)；其余走普通 fp32 行主序 (all_tokens, K//128)
• ep_scatter 一步完成重排：把乱序 token 按 recv 计数散进连续段，同时产出 m_indices（行→expert）和 output_index（记下行原来在哪，回程时靠它把结果发回，存入），最后设置 use_masked_gemm=False。

5）deepep_v2：v2 dispatcher 在 GPU 上直接算好了每个 expert 的前缀和 每个 experts 要接收到的 tokens 数量 psum_num_recv_tokens_per_expert，并带两个开关：is_expanded（token 是否已在 dispatch 阶段展开到 topk 份、基本就位）和 use_masked_gemm。按这两个开关分三条路，目标都是凑出 masked 或 contig 输入：

• is_expanded || use_masked：已展开的 token 只差填入到 masked 槽这一步，expand_to_masked_slab 直接把前缀和描述的数据变为 masked 3D 布局，产出 masked GEMM 输入
• is_expanded || ¬masked：同样已展开，但要 contiguous 布局，于是只需给每行标 expert 号——ep_expand_init_m_indices_from_psum 按前缀和填 m_indices，尾部没用上的行留 -1
• ¬expanded：token 还是乱序到达，退化成和 deepep_normal 一样的重排流程，只是散开时改用 v2 的前缀和 ep_scatter_from_psum，scale 缓冲的 UE8M0 / fp32 分配公式与 deepep_normal 完全相同。

到这一步，两种 GEMM 形态（masked / contig）与四种量化配方（bf16 / block-fp8 / mxfp8 / mxfp4）的输入都已备好，进入 GEMM1。

2.3 ③ GEMM1 计算

入口 DeepGemmRunnerCore.run（deep_gemm.py:359-393）：

class DeepGemmRunnerCore(MoeRunnerCore):
  def run(
      self,
      runner_input: DeepGemmRunnerInput,
      quant_info: DeepGemmMoeQuantInfo,
      running_state: dict,
      hooks: Optional[Any] = None,
    ) -> DeepGemmRunnerOutput:
      alignment = running_state.get("contiguous_layout_alignment") if not use_masked_gemm else None
      with deep_gemm_wrapper.contiguous_layout_alignment_scope(alignment):   # :374 设置 deep_gemm 全局 M/K 对齐
          # 四分支（use_masked_gemm × weight_dtype）

use_masked_gemm	权重 dtype	方法
False	bf16	_run_bf16_contiguous_gemm
False	fp8 / int8-packed	_run_contiguous_gemm
True	bf16	_run_masked_bf16_gemm
True	fp8 / int8-packed	_run_masked_gemm

注意 MXFP4 的权重经 view(torch.int8) 后 dtype 是 int8，因此和 fp8 共用“量化分支”；contiguous_layout_alignment_scope 只对 contig 生效（masked 行对齐由 slab 自身保证）。

2.3.1 contig 形态的 GEMM1 _run_contiguous_gemm

我们重点看 FP8 和 FP4 的 GEMM 计算。首先请各位读者注意，MXFP8、MXFP4 激活为 32 一组，FP8 block wise 激活为 128 一组；记住这一点后，我们再看一下权重 scale 部分：
quant_info.scale_recipes 中写死了 MXFP8 的 weight recipe 为 32，而 activation recipe 则是 follow 激活函数的量化方式，通常也是FP8 blockwise（1，128）或者 MXFP8 （1,32）； MXFP4 时（is_fp4_experts=1）因为
固定为 weight recipe 为 32（MXFP4），而 activation recipe 则是 FP8 Blockwise，所以是（1，128）

class DeepGemmMoeQuantInfo(MoeQuantInfo):
    def __post_init__(self):
        if self.use_mxfp8:
            assert self.block_shape == [
                1,
                32,
            ], f"MXFP8 requires block_shape [1, 32]"
  
    def scale_recipes(
        self,
        *,
        activation_block_size: Optional[int],
        hidden_size: int,
        activation_scale_width: int,
    ) -> tuple[Optional[tuple[int, int]], Optional[tuple[int, int]]]:
        """Return DeepGEMM A/B scale recipes from explicit layout metadata."""
        if self.use_mxfp8:
            weight_recipe = (self.block_shape[0], self.block_shape[1]) # (1, 32)
            activation_block_size = activation_block_size or self.block_shape[1]
            return (self.block_shape[0], activation_block_size), weight_recipe
        if self.is_fp4_experts:
            return (1, 128), (1, 32)
        return None, None

然后接着往下，会有一段针对 Hopper 架构的 TMA 对齐代码：tma_align_input_scale(hidden_states_scale)：

    if (
        deep_gemm_wrapper.DEEPGEMM_NEED_TMA_ALIGNED_SCALES
        and not runner_input.hidden_states_scale_tma_aligned
    ):
        hidden_states_scale = tma_align_input_scale(hidden_states_scale)

def tma_align_input_scale(input_scale: torch.Tensor):
    m, k_div_block_size = input_scale.shape
    padd_m = get_tma_aligned_size(m, input_scale.element_size())
    output = torch.empty(
        (k_div_block_size, padd_m), dtype=input_scale.dtype, device=input_scale.device
    )

    grid_m = min(m, 8192)
    BLOCK_SIZE_K = triton.next_power_of_2(k_div_block_size)

    _tma_align_input_scale_kernel[(grid_m,)](
        input_scale_ptr=input_scale,
        output_ptr=output,
        m=m,
        k_div_block_size=k_div_block_size,
        input_scale_stride_m=input_scale.stride(0),
        input_scale_stride_k=input_scale.stride(1),
        output_stride_m=output.stride(1),  # Note: these are swapped
        output_stride_k=output.stride(0),  # for column-major
        BLOCK_SIZE_K=BLOCK_SIZE_K,
    )
    return output.t()[:m]

它将 input scale 按照 MN-major （列主序）方式排布，16 字节对齐，并且做 M 方向 的padding。目的是为了与 DeepEP v2 对齐。

然后就是调用 deep_gemm 做一个 GEMM1 计算：

deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_contig(
    (hidden_states, hidden_states_scale),   # lhs: (all_tokens, K) fp8 + scale
    (w13_weight, w13_scale),                # rhs: (E, 2I, K) fp8 / (E, 2I, K//2) int8-fp4
    gateup_output,                          # out: (all_tokens, 2I) bf16
    m_indices,
    recipe_a=recipe_a, recipe_b=recipe_b,
)

GEMM1 计算结束后，还有一个比较令人在意的细节是，程序在非 CUDA graph 时对 hidden states 做了 dispose 操作（CUDA Graph 时地址会被固化进入 graph，不能随意变动）

dispose_tensor(hidden_states)
dispose_tensor(hidden_states_scale)

这是考虑到 GEMM1 计算完成后，hidden states 这片内存空间中的数值其实已经用不上了，它只被 gate-up GEMM 消费，而且它又很占地方，所以通过 dispose_tensor 立刻将显存归还给 caching allocator，供之后的计算函数复用空间，从而压低单层 MoE 的峰值显存。x.set_(torch.empty((0,), device=x.device, dtype=x.dtype)) 这个语句操作很聪明，set_() 把 Tensor 的 storage 换成 0 字节空张量，于是可以让指向这片地址的张量别名同时失效，存储确定性地归还 allocator，但直接用 del 就做不到这一点。

2.3.2 masked 形态的 GEMM1 _run_masked_gemm

GEMM1 前的 scale 收拾：UE8M0 时若 scale 还是 fp32（dtype != torch.int），先断言各维 4 对齐，再使用 _cast_to_e8m0_with_rounding_up：fp32 将指数位提取出来，然后根据值变换为 uint8，再做 transpose(1,2).contiguous().transpose(1,2) 变成 MN-major 对齐布局；非 UE8M0 走 get_mn_major_tma_aligned_tensor，同样也是 MN-major 布局，但要 TMA 对齐，该函数功能见 deep gemm 3.6.1 章节。

def _cast_to_e8m0_with_rounding_up(x: torch.Tensor) -> torch.Tensor:
    temp = x.to(torch.float32).view(torch.int32)
    exp = torch.bitwise_right_shift(temp, 23)
    mant = torch.bitwise_and(temp, 0x7FFFFF)
    is_ru = torch.logical_and(
        torch.logical_and((mant > 0), (exp != 0xFE)),
        ~torch.logical_and((exp == 0), (mant <= 0x400000)),
    )
    exp = torch.where(is_ru, exp + 1, exp)
    new_x = exp.to(torch.uint8).view(torch.int)
    return new_x.transpose(1, 2).contiguous().transpose(1, 2)

然后就是使用 GEMM1 masked gemm 计算完成，其余都和 contiguous 的一样。

deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_masked(
    (hidden_states, hidden_states_scale),   # (E, m, K) fp8 + (E, m, G/4) int32 或 (E, m, G) fp32
    (w13_weight, w13_scale),
    gateup_output,                          # (E, m, 2I) bf16
    masked_m, expected_m,
    recipe_a=recipe_a, recipe_b=recipe_b,
)

OOM 时报错信息会建议 SGLANG_OPT_DG_MASKED_M_CAP=1 收缩 M 上限。该环境变量通过将 pad 上限改为 round_up(max(masked_m), 256) 而非 round_up(rank_tokens, 256)，即从按照 rank 的全部 token 预留内存（最坏情况就是全部 token 都集中到一个 experts 上），改为当下最热 experts 实际接收到的 token 数取 256 对齐，但代价是每层 MoE 多一次 D2H 同步和 fused_moe_dispatch_index 操作。

2.3.3 BF16

• contig bf16：grouped_gemm_nt_bf16_contig ×2 + _legacy_silu_and_mul（无量化）；
• masked bf16：grouped_gemm_nt_bf16_masked ×2 + silu_and_mul_masked_fwd；不支持 overlap

2.4 ④ GEMM1 与 GEMM2 间：激活与量化

GEMM1 返回的 gateup_output 是 (all_tokens, 2I) 或 (E, m, 2I) 的 bf16，而在一般量化情况下（FP8 / MXFP8 / MXFP4），GEMM2 的 A 侧不能是 BF16，必须已经量化处理后的。所以这一阶段只有两件事：做 gated 激活（把 gate / up 两半成 I 宽），以及做激活和量化。

为什么这两件事要挤进同一个 kernel？2I 因为如果 先算激活写回 bf16、再读回来量化 两步走，内存就需要多搬两遍；而融合 kernel 只读一遍 bf16、写一遍 fp8 与 scale，这也是为什么大家都在尽量将 kernel 融合的原因。

再强调一个贯穿全节的约束是 scale 的产出布局。GEMM2 的 A 侧 scale 由 deep_gemm 直接消费：在 UE8M0 类型下必须是 packed int32、MN-major、TMA 对齐（三条生产线的来历见附录 A）。

2.4.1 contig 形态：三个互斥分支

contig 形态的激活与量化分三条互斥支路，先给一张对照表：

触发条件	激活实现	scale 产出
activation == "situ"（Kimi K3）	UE8M0：Triton _situ_mul_quant_contig_kernel；否则纯 torch	UE8M0：_cast_to_e8m0_with_rounding_up；否则行主序 fp32
use_swizzle or silu_mul_keep_fp32	DSV4 JIT silu_and_mul_contig_post_quant	create_per_token_group_quant_fp8_output_scale 按消费布局直接分配
其余（含 DSV4 默认）	silu_and_mul_clamp 或 _legacy_silu_and_mul	sglang_per_token_group_quant_fp8，UE8M0 时列主序 packed

1）situ 分支，这是给 Kimi K3 用的 SiTU 激活：

$$situ(x)=\left[\beta_s \tan\left(\frac{W_gx}{\beta_s}\right) \odot sigmoid(W_gx)\right] \odot \left[\beta_l \tanh\left(\frac{W_ux}{\beta_l}\right)\right]$$

UE8M0 平台先跑 Triton 的 _situ_mul_quant_contig_kernel 把 SiTU、amax 统计与分组量化一次做完，再用 _cast_to_e8m0_with_rounding_up 把 fp32 scale 向上取整到 2 的幂。为什么取整方向不能反：deep_gemm 的既定约定是 sf = ceil_to_ue8m0(amax / 448) 再 x / sf（fp8_utils.py:1972-1974），向上取整保证 x / sf ≤ 448；一旦向下取整，amax / sf 就超过了 448，量化那一步会直接饱和。这个 Triton kernel 的顺序是“先用精确 amax 量化、再把 scale 向上取整”，结果只会让反量化值略微偏大，同样不会溢出。

非 UE8M0 平台则退回纯 torch 计算，最后交给 sglang_per_token_group_quant_fp8(column_major_scales=False, ...) 出 fp32 行主序 scale：

gate = gateup_output[:, : N // 2].float()
up = gateup_output[:, N // 2 :].float()
gate = situ_beta * torch.tanh(gate / situ_beta) * torch.sigmoid(gate)
up = situ_linear_beta * torch.tanh(up / situ_linear_beta)
down_input = (gate * up).to(torch.bfloat16)

2）融合量化分支（use_swizzle or silu_mul_keep_fp32）：先用 create_per_token_group_quant_fp8_output_scale 按 UE8M0 的消费布局把 scale 缓冲分配好，再调 DSV4 JIT 的 silu_and_mul_contig_post_quant 一次完成 SiLU、钳位与分组量化：

down_input_scale = create_per_token_group_quant_fp8_output_scale(
    x_shape=(all_tokens, N // 2), device=..., group_size=scale_block_size,
    column_major_scales=UE8M0, scale_tma_aligned=UE8M0, scale_ue8m0=UE8M0,
)
silu_and_mul_contig_post_quant(
    input=gateup_output, output=down_input_fp8, output_scale=down_input_scale,
    quant_group_size=scale_block_size,
    scale_ue8m0=UE8M0, transposed=UE8M0,
    swiglu_limit=self.swiglu_limit, swizzle=use_contig_swizzle,
)

分配函数在 UE8M0 下按 (K_g/4, aligned_m) 分配 int32、再转成 (m, K_g/4) 的视图：MN 方向补齐到 4、K 方向 4 个 uint8 打包进一个 int32，内存形态就是 deep_gemm 要的 MN-major，省掉了后续任何一次转置。

use_swizzle 在 __init__ 里定义为 megamoe ∧ deep_gemm_sm120.use_swizzle()（deep_gemm.py:355-357），DeepSeek V4 + DeepEP 下恒为 False，所以这条支路实际由 silu_mul_keep_fp32 触发。

为什么要 silu_mul_keep_fp32 这个开关：通用融合量化 kernel 会把 SiLU 的中间结果先 round 到 bf16 再量化，等于多做了一次 bf16 舍入；silu_mul_keep_fp32 让 kernel 在 fp32 里算完 SiLU 直接量化，少一次精度损失。它只对 FP8 专家权重开放（deep_gemm.py:367-368）。

3）默认分支：先激活算得 BF16，再调通用量化 kernel。有 swiglu_limit 时用 silu_and_mul_clamp（把钳位融进激活，语义与 _apply_swiglu_limit 的兜底实现一致：up 双向钳、gate 只钳上限），否则走 _legacy_silu_and_mul；随后 sglang_per_token_group_quant_fp8(column_major_scales=UE8M0, scale_tma_aligned=UE8M0, scale_ue8m0=UE8M0)——三个 flag 绑在一起，是因为“列主序”与“TMA 对齐”在 UE8M0 下本来就是同一个 int32 packed 布局的两面。DSV4 的 contig 主路径落在这里：它既不满足 swizzle 条件、默认也没开 silu_mul_keep_fp32，只能走这条默认分支；swiglu_limit 是否生效取决于模型 config 里带不带这个字段（deepseek_v2.py:677 是 getattr(config, "swiglu_limit", None)）：带就用 silu_and_mul_clamp，不带就走 _legacy_silu_and_mul。

2.4.2 masked 形态：varlen 激活量化

masked 形态（deep_gemm.py:790-824）多一层麻烦：输入是 (E, m_pad, 2I) 的 slab，每个 expert 只有前 masked_m[e] 行是真实数据。激活 kernel 必须以 masked_m 为上界逐个 expert 跳行，也就是代码里的 varlen。

• situ 走 _varlen_deep_gemm_situ_mul_quant，底层是 CUDA JIT 的 situ_and_mul_masked_post_quant，量化 group 在这里被直接改写成 128（deep_gemm.py:791-800）， SiTU 分支的激活量化粒度固定为 128；
• 其余走 _varlen_deep_gemm_silu_mul_quant。进函数之前先算 num_real_tokens（deep_gemm.py:802-812）：只有 use_mxfp8、DEEPGEMM_SCALE_UE8M0、topk_ids、src2dst 四者齐备时才把真实 token 数传进去，因为只有“token 已展开且知道源行映射”的布局才真的数得清真实行；这个数用于把 kernel grid 限制在真实行上，避免按 padded_m 空转。

_varlen_deep_gemm_silu_mul_quant 内部还有三个子分支（deep_gemm.py:1559/1600/1634，打包细节见附录 A）：

1. gemm1_alpha（oai-swiglu）：Triton 的 silu_and_mul_masked_post_quant_fwd，因为通用量化 kernel 还没有 activation-kind 这个维度，oai-swiglu 暂时只能留在 Triton 实现上。use_packed 需要同时满足 UE8M0、num_real_tokens 非空、G % 4 == 0、D % (group*4) == 0；packed 时先分配 (E, G//4, N) 的 int32，再 transpose(-1,-2) 把内存形态转成 MN-major。
2. swiglu_limit / swizzle / silu_mul_keep_fp32（DSV4 主战场）：DSV4 JIT 的 silu_and_mul_masked_post_quant。这里有三条硬断言——N % 4 == 0、G % 4 == 0、D // 8 >= E（deep_gemm.py:1601-1604）：前两条是 packed scale 与向量化的对齐要求，最后一条来自 kernel 的线程编排——block 里开 D / 8 个线程，而 get_work 要用前 E 个线程各读一个 expert 的 masked_m 做前缀和，所以必须 D / 8 >= E。不满足就在激活量化处直接报错，没有退回其它分支的后路。scale 同样 packed + transpose(-1,-2)。
3. 默认：通用的 per_token_group_quant（fuse_silu_and_mul=True、masked_m=...、expected_m=ceil_div(num_real_tokens*topk, E)、column_major_scales=True）。它自己分配输出，scale 直接落成 deep_gemm 消费的列主序布局，所以调用方在 GEMM2 前的 get_mn_major 变换会短路。注意这个通用 kernel 会把 SiLU 中间值 round 到 bf16——这正是上面 silu_mul_keep_fp32 存在的理由。

2.5 ⑤ GEMM2 计算

GEMM2 和 GEMM1 调的是同一族 grouped 接口，差别只是 A 侧换成了刚量化好的 down 输入、B 侧换成 (E, K, I) 的 w2、N 方向从 2I 变回 K。真正值得单独拎出来说的只有三处：输出进对称内存、A 侧 recipe 要重算、masked 形态还要承担 combine 的 overlap。

2.5.1 contig 形态（deep_gemm.py:597-626）

with use_symmetric_memory(get_parallel().tp_group, disabled=not is_allocation_symmetric()):
    down_output = torch.empty((all_tokens, K), device=..., dtype=torch.bfloat16)  # 唯一进池的中间张量

if deep_gemm_wrapper.DEEPGEMM_NEED_TMA_ALIGNED_SCALES:
    down_input_scale = tma_align_input_scale(down_input_scale)   # 非 UE8M0 平台再对齐一次

recipe_a_down, _ = quant_info.scale_recipes(
    activation_block_size=scale_block_size, hidden_size=down_input_fp8.shape[-1], ...
)   # K 从 H 变成 I，A 侧 recipe 必须按新 K 重算

deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_contig(
    (down_input_fp8, down_input_scale), w2_weight_fp8, down_output, m_indices,
    recipe_a=recipe_a_down, recipe_b=recipe_b,
)

三处安排各有各的理由：

• 对称内存：down_output 是 GEMM2 的输出，也是 core 前向里唯一进 NCCL symmetric memory pool 的张量（注释就写在 deep_gemm.py:593-596）：它是这条计算链里唯一要参与下游通信的结果，放进池里就能让 all-reduce / reduce_scatter 走低延迟的对称路径。gateup_output、down_input_fp8 这些中间 buffer 留在默认 allocator 上，是为了不让 pool 的占用被中间激活撑大。回程的 (M, H) 输出同理——post_permute_deep_gemm_to_standard 也把它分配在池里（deep_gemm.py:1260-1265）。
• scale 再对齐：非 UE8M0（Hopper）平台刚量化出来的 down_input_scale 是行主序 fp32；GEMM1 前那次 tma_align_input_scale 只处理了输入激活，这里的 A 侧 scale 是新产出的，必须照同样规格转成列主序加 16 字节对齐。UE8M0 平台因为量化时已经直接产出消费布局，这一步自然跳过。
• recipe 重算：scale_recipes 给出的 A 侧 recipe 依赖 K——GEMM1 的 K 是 hidden H，GEMM2 的 K 是 intermediate I，两者不同。对 MXFP8 而言它还顺带断言 ceil_div(K, 32*4) == activation_scale_width（打包后列数必须对得上），这个断言本身也得按新的 K 重新算一遍。

2.5.2 masked 形态（deep_gemm.py:830-896）

masked 形态同样把 (E, m, K) 的 down_output 放进对称内存（deep_gemm.py:855-860）：

with use_symmetric_memory(
    get_parallel().tp_group, disabled=not is_allocation_symmetric()
):
    down_output = torch.empty(
        (num_groups, m, n), device=hidden_states_device, dtype=torch.bfloat16
    )

并且额外多做了两件事：

• scale 的最后一跳打包：use_mxfp8 且 DEEPGEMM_SCALE_UE8M0 且 down_input_scale 还不是 int32 时，此时需要调整 sacle 为 UE8M0 的格式，需要调 deep_gemm.get_mn_major_tma_aligned_packed_ue8m0_tensor（deep_gemm.py:833-844）。正常路径上各激活分支产出的已经是 packed int32，这一步是兜底；非 UE8M0 平台则走 wrapper 的 get_mn_major_tma_aligned_tensor（deep_gemm.py:845-848）。

if (
    use_mxfp8
    and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0
    and down_input_scale.dtype != torch.int32
):
    import deep_gemm.utils.layout

    down_input_scale = (
        deep_gemm.utils.layout.get_mn_major_tma_aligned_packed_ue8m0_tensor(
            down_input_scale
        )
    )
elif deep_gemm_wrapper.DEEPGEMM_NEED_TMA_ALIGNED_SCALES:
    down_input_scale = deep_gemm_wrapper.get_mn_major_tma_aligned_tensor(
        down_input_scale
    )

recipe_a_down, _ = quant_info.scale_recipes(
    activation_block_size=scale_block_size,
    hidden_size=down_input.shape[-1],
    activation_scale_width=down_input_scale.shape[-1],
)

• combine 与 down GEMM 的双流 overlap：down_gemm_overlap_args 非空时先 start_event.record()，再根据平台不同选择合适的 block 数值，具体为 max_block_n = 160 if (H20 且 expected_m ≤ 64) else 256，连同 overlap_args 一起传给 GEMM。GEMM 分块计算，combine 在另一条流上等第一批 block 完成就开跑；

  这种 overlap 模式是将 GEMM2 和之后的 allreduce 做了通算重合，此时 GEMM2 的返回值是 (block_m, threshold)，这里原样写回的 meta_overlap_args，供 DeepEP 的 low_latency_combine 判断“前多少行已经算完、可以开始合并发送”。H20 上把 max_block_n 压到 160，这应该是一个经验值，H20 上计算的 N 方向的 tile 更小（因为算力小），于是第一批 block 就能更早完成，那么通算的 overlap 更容易真正叠起来。

down_gemm_overlap_args = running_state.get("down_gemm_overlap_args", None)
if down_gemm_overlap_args is None:
    gemm_overlap_args_dict = {}
else:
    down_gemm_overlap_args.start_event.record()
    max_block_n = (
        160 if (_DEEPGEMM_ON_H20 and runner_input.expected_m <= 64) else 256
    )
    gemm_overlap_args_dict = {
        "overlap_args": down_gemm_overlap_args,
        "max_block_n": max_block_n,
    }

deep_gemm_return_value = deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_masked(
    (down_input, down_input_scale),
    (w2_weight, w2_scale),
    down_output,
    masked_m,
    expected_m,
    recipe_a=recipe_a_down,
    recipe_b=recipe_b,
    **gemm_overlap_args_dict,
)
if trace_deepep_v2_masked:
    torch.cuda.synchronize()
    logger.warning("DeepEP v2 masked runner down GEMM returned")
meta_overlap_args = running_state.get("meta_overlap_args", None)
# Returns (block_m, threshold) only with down-gemm overlap, else None;
# meta_overlap_args may be set without overlap, so guard the unpack.
if meta_overlap_args is not None and deep_gemm_return_value is not None:
    block_m, threshold = deep_gemm_return_value
    meta_overlap_args["block_m"] = block_m
    meta_overlap_args["threshold"] = threshold

return down_output

两个形态还有一个共同收尾：GEMM1 用完的输入激活在 GEMM1 之后就被 dispose_tensor 掉了（contig deep_gemm.py:467-468、masked deep_gemm.py:783-784），也就是 §2.3.1 讲过的“把 storage 换成 0 字节空张量、立刻归还 caching allocator”。所以峰值显存基本只包含当前这次 GEMM 的输入与输出，而不是整条链路的中间激活之和。

2.6 ⑥ GEMM2 后处理

2.6.1 post_permute：五种回程

post_permute 和 pre_permute 是镜像关系：pre_permute 把 dispatch 的“半成品”整理成 grouped 布局，post_permute 则要把 grouped 结果还原成各 dispatcher 的 combine 能消费的形态。但回程比进料多了一道题——top-k 加权与 routed_scaling_factor 到底在哪一层计算。五种回程的差别，本质上就是这道题的五种答案：

dispatch_format	post_permute	加权折叠位置
standard	post_permute_deep_gemm_to_standard	post_reorder_deepgemm 内，随 src2dst 累加一并乘
flashinfer	post_permute_deep_gemm_to_flashinfer	复用 standard，再转成 bf16 payload
deepep_ll	post_permute_deep_gemm_to_deepep_ll	留给 combine
deepep_normal	post_permute_deep_gemm_to_deepep_normal	ep_gather 内
deepep_v2	post_permute_deep_gemm_to_deepep_v2	expanded 分支在这里，非 expanded 在 ep_gather 内

standard 最典型：post_reorder_deepgemm 中使用 post_reorder_deepgemm_triton_kernel 按 src2dst 把同一 token 的 top-k 份结果累加回 (M, H)，并在写回时把 routed_scaling_factor 一并算进去，所以模型层不需要再单独做一次比例缩放。输出同样走对称内存，写完立刻 dispose_tensor 掉 GEMM 的大 buffer。

def post_reorder_deepgemm_triton_kernel(...):
    sum_vec = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
    for idx in range(topk):
        dst_idx = tl.load(token_src2dst_ptr + idx).to(tl.int64)
        if dst_idx >= 0:
            weight_scale = tl.load(token_topk_weights_ptr + idx).to(tl.float32)
            load_ptr_offs = down_output_ptr_offs + dst_idx * hidden_size
            in_data = tl.load(load_ptr_offs, mask=mask).to(tl.float32)
            sum_vec += in_data * weight_scale
    sum_vec *= routed_scaling_factor
    store_ptr_offs = output_ptr_offs + src_idx * hidden_size
    tl.store(store_ptr_offs, sum_vec.to(OutDtype), mask=mask)

但在 SM120 上，：deep_gemm_sm120.maybe_post_permute 命中就直接返回，交给它自己的重排实现 ep_gather。

flashinfer 只是把 standard 的结果再包一层：它要求 payload 是 bf16，不是就直接 TypeError；转成 FlashinferCombineInput 之后，由 combine 侧的 A2A 完成回程。

deepep_ll 最短：dispatch 阶段就是 masked 布局、combine 阶段本来就要消费 topk_ids / topk_weights，所以这里原样透传，一个权重都不乘。

deepep_normal 把“收回 + 加权 + 归约”压进一次 ep_gather(hidden, topk_ids, topk_weights, output_index, gather_out)（deep_gemm.py:1477），output_index 就是 pre_permute 时 ep_scatter 顺手记下的回程映射。

deepep_v2 要看 dispatch 时 token 有没有展开，分三种情况：

• 展开且 masked：masked_slab_to_expand 把 (E, m, H) 的 slab 按前缀和摊回每个 route 一行。这里还分两种模式：no_combine 时只做展开、不乘权重，把 topk_weights 一并交给 routewise finalizer；否则在 kernel 里把权重乘掉，combine 只负责求和：

  def _fwd_kernel_masked_slab_to_expand(...):
      for j in tl.range(worker, count, NUM_WORKERS):
          src = (e * MAX_M + j).to(tl.int64)
          dst = (start + j).to(tl.int64)
          v = tl.load(input_tensor_ptr + src * input_tensor_stride0 + off, mask=mask)
          if HAS_W:
              w = tl.load(weight_ptr + dst)
              v = (v.to(tl.float32) * w).to(v.dtype)
          tl.store(output_tensor_ptr + dst * output_tensor_stride0 + off, v, mask=mask)

• 展开但不 masked：token 已经就位，而 expanded combine 不消费权重，所以这里显式做 hidden * topk_weights.unsqueeze(-1)（deep_gemm.py:1891-1896）。

• 未展开：与 deepep_normal 同构。no_combine 时按 output_index >= 0 把结果沿 route 维散回 (M, top_k, H)、非本地 route 填 0（deep_gemm.py:1902-1919），否则就是 ep_gather。

2.6.2 combine 与模型尾部

回到 FusedMoE.forward_impl 的最后一步：dispatcher.combine(combine_input) 完成跨 rank 的合并与再通信，把结果还原成 (M, H)（fused_moe_triton/layer.py:1586-1589）。这一步同样要使用 NCCL Symm 技术，包裹在 use_symmetric_memory 里，combine 的输出直接落进 pool，然后 all-reduce 才能能走对称通信；出来之后再按入口记录的 origin_hidden_states_dim 切回原始宽度并 contiguous()（fused_moe_triton/layer.py:1592-1594）。

模型层的尾部（deepseek_v2.py:1368-1393）做三件事：

1. routed 与 shared 相加：maybe_fuse_routed_scale_and_shared_add。这里要澄清一点：对 DeepGEMM 路径而言，routed_scaling_factor 已经在 post_reorder_deepgemm 里折进写回了，这个函数实际只做 routed += shared 的原地相加；函数里那套 alpha=routed_scaling_factor 的融合分支，是给 flashinfer trtllm 这类“把 scaling 融进 top-k 权重”的后端准备的。

   final_hidden_states = maybe_fuse_routed_scale_and_shared_add(
           self.experts,
           final_hidden_states,
           None if self._shared_expert_tp1 else shared_output,
           self.routed_scaling_factor,
       )

2. mhc 统计点与 post-experts 归约：DeepSeek V4 且 TP>1 时，会在归约前先给 mhc post fusion 打一个 start_stats_before_all_reduce() 统计点；随后 post_experts_all_reduce 把 EP 组与 MoE-TP 组的归约收在一次里（两者都生效时合并成一次 TP all-reduce，deepseek_v2.py:1385）。这一步会被 should_skip_post_experts_all_reduce(is_tp_path=True) 跳过——那正是 §2.6.3 里 mlp_reduce_scatter scoped 标记的用途：DP attention 下由 combine 之后的 reduce_scatter 一并完成归约，避免双归约。

   if (
       self.is_deepseek_v4
       and self.tp_size > 1
       and not should_skip_post_experts_all_reduce(is_tp_path=True)
   ):
       from sglang.srt.layers.moe.mhc_post_fusion import current_mhc_post_fusion
   
       mhc = current_mhc_post_fusion()
       if mhc is not None:
           mhc.start_stats_before_all_reduce()
       final_hidden_states = post_experts_all_reduce(final_hidden_states)        

3. TP1 shared 的后置相加：shared expert 没做 TP 切分时（_shared_expert_tp1），每个 rank 算出的 shared 输出完全相同，所以它被安排在 all-reduce 之后才相加，否则会被归约重复累加 tp_size 次。注意第 1 步传给 maybe_fuse_routed_scale_and_shared_add 的 shared 这时是 None（deepseek_v2.py:1371），正是为了把这次相加推迟到这里。

   if (
       shared_output is not None
       and self._shared_expert_tp1
       and should_add_replicated_moe_output()
   ):
       final_hidden_states += shared_output

为什么归约放在模型层，而不是 MoE 层内部？可以看到， FusedMoE.reduce_results 的传参默认是 False，forward_impl 里的兜底 all_reduce 只在显式配置时才触发（fused_moe_triton/layer.py:1596-1597）。这是因为，MoE 层“身在此山中”，它要处理太多地细节，看不清上层用的是 TP-MoE 还是 EP，更不知道 attention 侧的并行形态（EP 的跨 rank 归约已经在 combine 里做完，DP-attn 有 reduce_scatter ）；故只能把“要不要归约、归约哪一组”交给最了解并行形态的 deepseek_v2 模型层决定，才不会出现 EP 与 all_reduce 各做一遍、或者 DP-attn 与 all_reduce 重复归约的错误。

2.6.3 V4 的 DP attention 变体：_run_moe_ffn_dp_sync

之前小节对 MoE 的叙述有一个隐含前提：本 rank 手上的 token 就是这层 MoE 要算的全部 token。然而，在 DeepSeek V4 在 attention 数据并行（attn DP）× MoE 张量/专家并行下打破了这个前提——当每个 rank 只有自己算出来的部分 token 时，expoerts 不论是按 intermediate 切在 TP 上（单卡只算出半个 MLP），还是按 EP 切在不同 rank 上，都是没办法准确计算了。于是 DeepSeek V4 在每层 MoE 前向上又加了一层抽象，将 MoE 整体收进 _run_moe_ffn_dp_sync，由它决定“进 MoE 前怎么把 token 凑齐、出 MoE 后怎么把结果还回去”。

三条形态互斥，判定都带 not _use_cp：

_use_cp = self.dsa_enable_prefill_cp and dsa_use_prefill_cp(forward_batch)
_use_tp_moe_gather = (
    not _use_cp
    and get_parallel().attn_dp_size > 1
    and get_moe_a2a_backend().is_none()
)
_use_tp_attn_a2a_scatter = (
    not _use_cp
    and get_parallel().attn_tp_size > 1
    and not get_moe_a2a_backend().is_none()
)

形态	MoE 之前	MoE 之后	归约落在哪
TP-MoE + attn DP（_use_tp_moe_gather）	dp_gather_replicate 拼全局 token	reduce_scatterv / 等分片 reduce_scatter / dp_scatter	回程的 reduce_scatter；MoE 内部 all_reduce 被 mlp_reduce_scatter=True 跳过
EP + attn TP（_use_tp_attn_a2a_scatter）	按 attn_tp_rank 把 token tensor_split 一份给本 rank	attn_tp_all_gather 拼回	EP 的 all2all 里；尾部只 gather、不再 reduce
CP（_use_cp）	a2a==none 走 dsa_cp_gather_hidden_states，否则直接 all2all	dsa_cp_reduce_scatter_hidden_states	CP 分片的 reduce_scatter

三者都不命中时，可以是纯 TP-MoE，函数退化成一次裸的 self.mlp(...)，前后不加任何搬运——因为这时 token 就全在本 rank，各 TP rank 算完自己的专家分片后交给 post_experts_all_reduce 求和即可。如果 EP + DP，且 attn_tp_size=1，那没有 gather、没有 split、没有 reduce_scatter。因为在这个配置下每个 rank 手上的 token 是自己的、且完整的，不同 rank 的 token 互不重叠。

1）TP-MoE + attn DP：先 gather 全局 token，再 reduce_scatter 还回去

TP-MoE 下专家按 intermediate 维切在 TP 组上，单卡只算出 MLP 的一部分和，所以必须“所有 token 都过一遍所有 rank 上的 experts”，然而因为 attn DP ，每个 rank 都只有一部分 token，所以方法就是先把 token gather 都拿回来：

elif _use_tp_moe_gather:
    hidden_states, local_hidden_states = (
        get_global_dp_buffer(get_parallel().tp_group),
        hidden_states,
    )
    # ……（SGLANG_DP_SHARED_EXPERT_LOCAL 时在此先算本地 shared，见下文）
    # self_attn has already reduced across attention TP, so these hidden
    # states are replicated and must not be summed by a partial gather.
    dp_gather_replicate(hidden_states, local_hidden_states, forward_batch)

注意这里必须是“复制”语义的 gather 而不是 partial-sum：因为 self_attn 已经在 attention TP 上 all_reduce 过，各 rank 的 hidden 其实是同一份 tensor 的不同部分。

在结束阶段， reduce_scatterv 则是同一个道理的反向：一次算子完成“跨 TP reduce 求和，再把 token 切片各自拷贝到对应 shard 回去”，好让各个 rank 继续下一层的计算，这两个操作合并到一起就是 reduce scatter。当然，此时 MoE 内部的 post-experts all_reduce 必须被跳过，否则就多算了一次。

mlp_reduce_scatter = (
    _use_cp
    or _use_reduce_scatterv
    or _use_reduce_scatter
    or (_use_tp_moe_gather and should_use_dp_reduce_scatterv())
)

with get_forward().scoped(mlp_reduce_scatter=mlp_reduce_scatter):   # :3861
    hidden_states = self.mlp(hidden_states, forward_batch, ...)

hidden_states, global_hidden_states = (                             # :3874
    get_local_dp_buffer(get_parallel().tp_group),
    hidden_states,
)
if should_use_dp_reduce_scatterv() or _use_reduce_scatterv:          # :3879
    get_parallel().tp_group.reduce_scatterv(
        global_hidden_states, output=hidden_states, sizes=get_dp_global_num_tokens(),
    )

mlp_reduce_scatter 正是 §2.6.2 里 should_skip_post_experts_all_reduce(is_tp_path=True) 读到的那个 per-forward 标志——一边写、一边读，靠 get_forward() 的 scoped 上下文传递。

三条回程按能力从强到弱排：

1. reduce_scatterv：变长的 gatherv/scatterv 配对，分片长度由 get_dp_global_num_tokens() 给出。前提是 SGLANG_DP_USE_GATHERV（默认关）、SUM_LEN padding、attn_tp_size == 1 且 tp_size == attn_dp_size（is_dp_gatherv_active，dp_attention.py:793-810），否则两边长度对不上，直接调用 pynccl_comm.reduce_scatter(output, input_, sizes=sizes)
2. 等分片 reduce_scatter：MAX_LEN（各 rank 等长 padding、gatherv 不生效）时走这条，开关是 SGLANG_DP_USE_REDUCE_SCATTER，要求 tp_size == attn_dp_size， 调用 pynccl_comm.reduce_scatter(output, input)
3. dp_scatter：兜底。这时 mlp_reduce_scatter 为 False，MoE 内部 all_reduce 照常做，那就相当于先 all_reduce，再拷贝本 rank 切片。

顺带一个可选优化：SGLANG_DP_SHARED_EXPERT_LOCAL（deepseek_v4.py:368）打开且 shared expert 是 replicated（TP1）时，shared 的输出改成在 gather 之前只用本地 token 算（deepseek_v4.py:3837-3838），等回程之后再加到本 rank 切片上（deepseek_v4.py:3902-3907）。shared expert 是逐 token 的 MLP，“在 M_local 行上算”与“在 M_global 行上算完再取切片”结果一致，但前者只要 1/dp_size 的行数——这正是 §2.6.2 第 3 点“TP1 shared 在归约之后相加”那份代价的补丁。

2）EP + attn TP：token 均分给各 rank，回程 all_gather

EP 下 experts 切在不同 rank 上，token 由 dispatch/combine 的 all2all 自己找到对应的 rank，不需要 gather 全局；但 attn TP 组内每个 rank 握着同一批 token ，让每个 rank 都把全量 token 来跑一遍 MoE 纯属浪费，为此，按 attn_tp_rank 把 token 均分成 attn_tp_size 份、每个 rank 只留一部分 token shard（deepseek_v4.py:3843-3846），回程 attn_tp_all_gather 拼回（deepseek_v4.py:3908-3912）。

if _use_tp_attn_a2a_scatter:
    s, r = get_parallel().attn_tp_size, get_parallel().attn_tp_rank
    _a2a_scatter_chunks = list(hidden_states.tensor_split(s))
    hidden_states = _a2a_scatter_chunks[r].contiguous()

if _use_tp_attn_a2a_scatter:
    assert _a2a_scatter_chunks is not None
    gathered = [torch.empty_like(t) for t in _a2a_scatter_chunks]
    attn_tp_all_gather(gathered, hidden_states.contiguous())
    hidden_states = torch.cat(gathered)

回程只 gather、不 reduce，是因为每份切片算出的都是完整结果（该 token 的 top-k 专家输出已由 combine 求和），只是分散在不同 rank 上。

细节：input_ids / input_ids_global 要跟着一起切（deepseek_v4.py:3849-3852）。V4 有按 token id 路由的 hash MoE 层（deepseek_v2.py:1871-1877），路由依据和 token 一旦错位就全错；而 DSpark 这类不带 id 的层则是“有就切、没有就跳过”。

3）CP：用 CP 专用 gather 或直接走 all2all

CP（prefill context parallel）把一条序列按 context 切成几段放在不同 rank 上，所以 _use_cp 为真时上面两条都不成立，MoE 前必须把 token 汇聚：a2a == none 用专用的 gather dsa_cp_gather_hidden_states；走 all2all 则要求 moe_a2a_backend ∈ {deepep, megamoe, mori}（断言 deepseek_v4.py:3824-3831）——因为只有这几个后端有“按 token 送去专家”的能力。

reduce scatter 部分也有专用的 dsa_cp_reduce_scatter_hidden_states。专用的 all gather 和 reduce scatter 都是按照 rank 拼接到了 hidden states tensor 内。每个 CP rank 的真实行都落在自己那段 slot 里，不是整个缓冲的前缀，所以在 a2a==none 时，因为 CP 下 gather 后的缓冲是按 CP rank 拼的，真实行是所有的行，而非仅仅前 N 行，所以必须在 MoE 期间把有效行数设成整块缓冲大小，否则 router 的前缀掩码会误删后面 rank 的真实 token。

所以对 DeepSeek V4 在 Blackwell 上用 DeepEP 的典型部署（attn 侧 TP、MoE 侧 EP）：先切掉 attn TP 组内重复的 token 副本，再交给 DeepEP 做 all2all，最后 all_gather 拼回来；而 DeepEP 这一层内部实际执行的仍然是 §2.4/§2.5 讲的 DeepGEMM MoE 中的三段式。

总结一下，不同并行方式下 MoE 相关的 token 通信和内存行为：

并行方式	MoE 前	dispatch	combine	MoE 后
attn TP × TP-MoE 纯 TP	—	本地重排	本地加权	1 次 all_reduce
attn DP × TP-MoE	gather 成全量	本地重排	本地加权	1 次 reduce_scatter
attn TP × EP	切掉重复副本	all2all 发出	all2all 收回	1 次 all_gather
attn DP × EP	—	all2all 发出	all2all 收回	—

3. DeepGEMM 接口详解：参数与职责

调用链分两层：

SGLang MoE/GEMM 代码
  └── sglang/layers/deep_gemm_wrapper/entrypoint.py        ← 本文的“wrapper 接口”
        └── deep_gemm 包（sgl-deep-gemm wheel，JIT 编译）   ← “库侧接口”
              └── TVM-FFI 扩展 _C（kernel 选择 / 编译 / 发射）

wrapper 层做参数适配、sanity check、对齐 scope 与 num_sms 管理；库侧负责 JIT 编译与 kernel 派发。下面先给 wrapper 接口总表，再逐个展开。

3.1 wrapper 接口总表

wrapper（entrypoint.py）	职责
grouped_gemm_nt_f8f8bf16_masked	masked fp8/fp4 grouped GEMM（支持 down-GEMM overlap）
grouped_gemm_nt_bf16_masked	masked bf16 grouped GEMM
grouped_gemm_nt_f8f8bf16_contig	contig fp8/fp4 grouped GEMM（m==0 提前返回）
grouped_gemm_nt_bf16_contig	contig bf16 grouped GEMM
gemm_nt_f8f8bf16	非 grouped fp8 稠密 GEMM（shared expert down_proj 等）
gemm_nt_mxfp8_f8f8bf16	非 grouped MXFP8 稠密 GEMM
gemm_nt_bf16bf16f32	非 grouped bf16 稠密 GEMM
tf32_hc_prenorm_gemm	V4 的 HC 预归一化 GEMM（TF32）

四个 grouped 接口与 runner 四分支一一对应（§2.3/§2.5 的调用点）；四个非 grouped 接口服务于模型里的稠密层（deepseek_v2.py:397、kernels/ops/gemm/fp8_kernel.py:48、kernels/ops/gemm/bf16_fp32.py:157、deepseek_v4.py:340 / :2892）。

3.2 grouped_gemm_nt_f8f8bf16_masked：masked fp8/fp4 GEMM

完整签名：

grouped_gemm_nt_f8f8bf16_masked(
    lhs: Tuple[Tensor, Tensor],      # 激活：值 + scale
    rhs: Tuple[Tensor, Tensor],      # 权重：值 + scale
    out: Tensor,
    masked_m: Tensor,                # (E,) int32，每 expert 有效行数
    expected_m: int,                 # masked M 的调优提示
    overlap_args: Optional[Any] = None,
    max_block_n: int = 256,
    recipe_a: Optional[Tuple[int, int]] = None,   # 激活 recipe
    recipe_b: Optional[Tuple[int, int]] = None,   # 权重 recipe
):

逐参数：

参数	形状 / 类型	含义
lhs[0]	(E, m, K) fp8 e4m3（或 fp4 packed）	激活值；m 是每 expert slab 行数（对齐后的上界），K = hidden_size
lhs[1]	(E, m, G/4) int32（UE8M0 packed）或 (E, m, G) fp32	激活 scale；G = K / group（fp8 128，mxfp8 32）
rhs[0]	(E, N, K) fp8；MXFP4：(E, N, K//2) int8	权重值（fp4 以 2/half 打包）
rhs[1]	(E, N/128, K/32) int32 等	权重 scale（MN-major、TMA 对齐、UE8M0 packed）
out	(E, m, N) bf16	输出，由调用方分配（对称内存可选）
masked_m	(E,) int32	第 e 个 expert 只算前 masked_m[e] 行，其余行属 padding
expected_m	int	调优提示：expect_m 用于库内选 kernel 配置（MoE 侧传 ceil_div(M*topk, 全局E)）
overlap_args	DownGemmOverlapArgs{num_sms, signal, start_event}	非空时开启 down-GEMM overlap（见下）
max_block_n	int = 256	库侧 N 方向 tile 上限；overlap 场景下调小可降低单 tile 延迟
recipe_a / recipe_b	(int,int) 或 None	仅非 None 时作为 kwargs 透传；(None, None) 时库按默认 128×128 解读

内部依次做四件事（:59-97）：

1. 形状探测：num_groups, _, k = lhs[0].shape、_, n, _ = rhs[0].shape；
2. _sanity_check_input(lhs/rhs)（:314-328）：UE8M0 模式下校验 fp32 scale 已是 2 的幂（ceil_to_ue8m0(x) == x，已打包的 int32 直接跳过）；_ensure_cuda（:100-106）保证 tensor 在 CUDA 上；
3. compile_utils.deep_gemm_execution_hook(expected_m, n, k, num_groups, kernel_type)（:69-71 → compile_utils.py:436-451）：运行期遇到未预热形状时记录/补编（附录 D）；
4. configure_deep_gemm_num_sms(overlap_args.num_sms)（:72-74 → :301-311）：with 块内 set_num_sms，退出恢复。overlap 模式下把一部分 SM 让给通信；

最后调用 deep_gemm.fp8_m_grouped_gemm_nt_masked(lhs, rhs, out, masked_m, expected_m, **fp4_kwargs, **overlap_kwargs)（:81-97），其中 overlap_kwargs = dict(enable_overlap=True, max_block_n=max_block_n, signal=overlap_args.signal) 仅在 overlap 时附加——这三个是 sgl-deep-gemm wheel 在 masked 接口上的 overlap 扩展（上游 v0.2.0 tag 的等价签名不含它们）。

OOM 时的辅助通道：_run_masked_gemm 在报错信息里提示 SGLANG_OPT_DG_MASKED_M_CAP=1（把每 expert 容量从 round_up(rank_tokens,256) 收缩到 round_up(max(masked_m),256)，代价是每层一次 D2H 同步，deep_gemm.py:761-770）。

3.3 grouped_gemm_nt_f8f8bf16_contig：contig fp8/fp4 GEMM

完整签名（:132-139）：

grouped_gemm_nt_f8f8bf16_contig(
    lhs: Tuple[Tensor, Tensor],      # 激活：值 + scale
    rhs: Tuple[Tensor, Tensor],      # 权重：值 + scale
    out: Tensor,
    m_indices: Tensor,               # (M,) int32，行 → expert
    recipe_a: Optional[Tuple[int, int]] = None,
    recipe_b: Optional[Tuple[int, int]] = None,
):

与 masked 的差异：

• 激活是 (M, K)（M = all_tokens，expert 连续段 + 对齐 padding）而非 (E, m, K)；m_indices 告诉 kernel 每一行属于哪个 expert（grouped layout 的描述符）；
• 没有 masked_m / expected_m / overlap 参数——contig 形态目前不参与 down-GEMM overlap；
• m == 0 直接提前返回（:144-145），纯 EP 且本 rank 无 token 时是一个纯 no-op；
• M/K 对齐由外层 scope 保证：DeepGemmRunnerCore.run 用 contiguous_layout_alignment_scope(alignment)（deep_gemm.py:374）把 pre_permute 算出的 block_e 设为全局对齐，deep_gemm 的 M/K 对齐断言在该 scope 内成立。

3.4 BF16 两个接口

grouped_gemm_nt_bf16_masked(a, b, d, masked_m, expected_m)（:109-129）与 grouped_gemm_nt_bf16_contig(a, b, d, m_indices)（:162-170）：

• 输入是裸 Tensor（没有 scale 元组），分别下推 deep_gemm.m_grouped_bf16_gemm_nt_masked / m_grouped_bf16_gemm_nt_contiguous；
• 只做 execution hook（warmup 记账），没有 sanity check / num_sms / recipe；
• masked bf16 不支持 overlap（deep_gemm.py:961 注释）。

3.5 非 grouped 接口：MoE 之外的四条 GEMM

MoE 之外，wrapper 还暴露 4 个非 grouped 接口，服务模型其他位置的矩阵乘。它们共享同一骨架——「_sanity_check_input → deep_gemm_execution_hook → 库调用」，没有 grouped 维度与 recipe 组装：

3.5.1 gemm_nt_f8f8bf16（:219-237）——普通 fp8 GEMM

gemm_nt_f8f8bf16(
    lhs: Tuple[Tensor, Tensor],   # (M, K) fp8 + scale
    rhs: Tuple[Tensor, Tensor],   # (N, K) fp8 + scale
    out: Tensor,                  # (M, N) bf16
)

• 形状探测 m, k = lhs[0].shape、n, _ = rhs[0].shape，num_groups = 1；
• 两侧 sanity check 后下推 deep_gemm.fp8_gemm_nt(lhs, rhs, out)（:233-237）；
• 调用方：dense MLP 的 down_proj（deepseek_v2.py:397），即非 MoE 的 block-fp8 线性层。

3.5.2 gemm_nt_mxfp8_f8f8bf16（:240-263）——MXFP8 GEMM

签名同上。与 3.5.1 的差异：

• 显式传 recipe_a=(1, 32), recipe_b=(1, 32) 并下推 deep_gemm.fp8_fp4_gemm_nt；
• disable_ue8m0_cast = lhs[1].dtype == torch.int and rhs[1].dtype == torch.int（:253）：两侧 scale 都已是打包 int32 时跳过库内 UE8M0 转换；
• 调用方：mxfp8 权重矩阵乘分派（kernels/ops/gemm/fp8_kernel.py:48）。

3.5.3 gemm_nt_bf16bf16f32（:266-277）——bf16→fp32 高精度 GEMM

gemm_nt_bf16bf16f32(lhs, rhs, out)   # (M,K)bf16 × (N,K)bf16 → (M,N)fp32

• 裸 Tensor、无 scale，下推 deep_gemm.bf16_gemm_nt；
• 调用方：kernels/ops/gemm/bf16_fp32.py:157——_linear_bf16_fp32_algo == "deep_gemm" 时走这条，否则回退 cuBLAS。

3.5.4 tf32_hc_prenorm_gemm（:280-289）——DSV4 HyperConnection 的 prenorm

tf32_hc_prenorm_gemm(x, fn, out, sqrsum, num_splits: Optional[int])

• x @ fn 的同时把 rows 平方和累加进 sqrsum（后续 rsqrt 用），fp32 中间精度（tf32 输入）；
• x.shape[0] == 0 直接返回；无 execution hook、无 sanity check，纯直通（:287-289）——它不参与 JIT 预热记账；
• 调用方：deepseek_v4.py:340（prefill 侧 mix 投影）与 :2892（decode 侧），DSV4 独有 HyperConnection 结构的一部分。

3.6 辅助与配置接口

3.6.1 get_mn_major_tma_aligned_tensor

把普通权重/激活 scale 转成 MN-major、TMA 对齐布局，非 UE8M0 平台激活侧调用。

docstring 里记录了仓库作者遇到的 use-after-free 事故并做了防御（:25-36）：sgl-deep-gemm 的快速路径在「输入已是目标布局」时会跨 TVM-FFI 返回一个 非持有（non-owning）的 from_blob 别名；而调用方习惯写 x = get_mn_major_tma_aligned_tensor(x)——重绑结果后原存储被释放，GEMM 仍透过别名读取，一旦 allocator 复用该块就是 UAF（表现为 NaN logits 或 CUDA graph capture 期的 “pointer resides on host memory”）。wrapper 因此检测 out.data_ptr() == sf.data_ptr() 时直接返回原 tensor 以保住所有权。

3.6.2 get_contiguous_layout_alignment / contiguous_layout_alignment_scope（:173-194 / :197-216）

• getter：读库侧 get_mk_alignment_for_contiguous_layout()（兼容返回 tuple 的版本）；若库提供了 get_theoretical_mk_alignment_for_contiguous_layout，则按 per_group_m = ceil_div(expected_m, num_groups) 求理论对齐值——只要理论值 < 全局值就下调，反之保持全局值（:184-194）；
• scope：contextmanager，setter(alignment) 进入、setter(previous) 退出恢复；库侧缺 getter/setter 或 alignment 为 None 时退化为 no-op；
• 消费方：DeepGemmRunnerCore.run（deep_gemm.py:374）把 pre_permute 算出的对齐包在整个 GEMM 区段外，contig 路径的 M/K 对齐断言在该 scope 内成立；JIT 预热侧由 compile_utils 以相同 scope 编译对齐一致的 kernel。

3.6.3 configure_deep_gemm_num_sms（:301-311）

contextmanager：num_sms 为 None 或 JIT 关闭时 no-op；否则块内 set_num_sms(num_sms)、退出 set_num_sms(original)。overlap 场景用它给通信让出 SM（deep_gemm.py:72-74）。

3.6.4 update_deep_gemm_config（:292-298）

启动期一次性：SGLANG_DEEPGEMM_PDL 开启且库有 set_pdl 时调用 deep_gemm.set_pdl(True)（注释强调它会在 fork 后、绑定 GPU 之后执行，避免提前初始化 CUDA 状态），然后转交 compile_utils.update_deep_gemm_config(gpu_id)。

3.6.5 _sanity_check_input（:314-328）

由 SGLANG_DEEPGEMM_SANITY_CHECK 控制（默认关，:44）。逻辑三层：已经是 int（打包 UE8M0）直接放过；非 DEEPGEMM_SCALE_UE8M0 放过；否则断言 ceil_to_ue8m0(x_scale) == x_scale——fp32 scale 必须已是 2 的幂，否则库内 round-up 会静默改变数值。

3.7 库侧签名对照（sgl-deep-gemm 0.2.0）

SGLang 依赖 SGL 维护的 DeepGEMM 发行版（import deep_gemm）。以下签名摘自上上游 GitHub v0.2.0 tag 的 deep_gemm/__init__.py，作为 wrapper 调用面的底账：

# ---- grouped（MoE）族 ----
m_grouped_fp8_fp4_gemm_nt_masked(a, b, d, masked_m, expected_m,
                                 recipe=None, recipe_a=None, recipe_b=None,
                                 compiled_dims='nk', disable_ue8m0_cast=False)
m_grouped_fp8_fp4_gemm_nt_contiguous(a, b, d, grouped_layout,
                                     recipe=None, recipe_a=None, recipe_b=None,
                                     compiled_dims='nk', disable_ue8m0_cast=False,
                                     use_psum_layout=False, ensure_zero_padding=True,
                                     expected_m_for_psum_layout=None)
m_grouped_bf16_gemm_nt_masked(a, b, d, masked_m, expected_m, compiled_dims='nk')
m_grouped_bf16_gemm_nt_contiguous(a, b, d, grouped_layout, compiled_dims='nk',
                                  use_psum_layout=False, ensure_zero_padding=True,
                                  expected_m_for_psum_layout=None)

# ---- dense 族 ----
fp8_fp4_gemm_nt(a, b, d, c=None, recipe=None, recipe_a=None, recipe_b=None,
                compiled_dims='', disable_ue8m0_cast=False, alpha=None)
bf16_gemm_nt(a, b, d, c=None, compiled_dims='', alpha=None)

# ---- scale / layout 工具 ----
transform_sf_into_required_layout(sf, mn, k, recipe, num_groups=None, is_sfa=None,
                                  disable_ue8m0_cast=False, psum_layout=None)
get_mk_alignment_for_contiguous_layout() / set_mk_alignment_for_contiguous_layout(v)
get_num_sms() / set_num_sms(n)

同名别名（fp8 与 fp4 共用实现，按输入 dtype 分派）：fp8_fp4_gemm_nt = fp8_gemm_nt = fp4_gemm_nt；m_grouped_fp8_fp4_gemm_nt_masked = fp8_m_grouped_gemm_nt_masked = m_grouped_fp8_gemm_nt_masked = m_grouped_fp4_gemm_nt_masked；contiguous 版同理三别名。

参数语义：

参数	语义
a / b / d	激活 / 权重 / 输出；a、b 可为 (tensor, scale) 元组或裸 Tensor——库内 _parse_tensor_or_tuple 把裸 Tensor 补成 (t, [1.0]) 的假 scale，所以 bf16 接口在 wrapper 里可以裸传
masked_m	masked 布局每 expert 有效行数 (E,) int32
expected_m	平均每 expert 行数期望，供库内选 kernel 配置（MoE 侧传 ceil_div(M*topk, 全局E)）
grouped_layout	contig 布局的行 → expert 映射 (M,) int32，需满足 get_mk_alignment_for_contiguous_layout 对齐
recipe / recipe_a / recipe_b	scale 分组规则，(block_m, block_n) 或 (block_m, block_n, block_k)（二元组库内补 None）；A 侧=激活、B 侧=权重
compiled_dims	JIT 编译的特化维名（'nk' 表示按 N/K 维特化；dense 版默认空串）
disable_ue8m0_cast	True 跳过库内「fp32 scale → UE8M0」转换；输入已是 packed int32 时必须 True，否则二次 cast
use_psum_layout / ensure_zero_padding / expected_m_for_psum_layout	v0.2.0 的 psum（partial-sum）布局选项；wrapper 未启用，走默认
c / alpha	dense 版：可选累加目标 C 与输出缩放 alpha
sf / mn / k / is_sfa / psum_layout	transform_sf_into_required_layout：把任意 scale 整理为目标布局；is_sfa 标明是激活侧

wrapper → 库侧对应表（未列参数均走库侧默认）：

wrapper 接口	库侧调用	显式传参
grouped_gemm_nt_f8f8bf16_masked	m_grouped_fp8_gemm_nt_masked	recipe_a/b（非 None 时）；overlap 时追加 enable_overlap / max_block_n / signal *
grouped_gemm_nt_f8f8bf16_contig	m_grouped_fp8_gemm_nt_contiguous	recipe_a/b（非 None 时）
grouped_gemm_nt_bf16_masked	m_grouped_bf16_gemm_nt_masked	—
grouped_gemm_nt_bf16_contig	m_grouped_bf16_gemm_nt_contiguous	—
gemm_nt_f8f8bf16	fp8_gemm_nt	—
gemm_nt_mxfp8_f8f8bf16	fp8_fp4_gemm_nt	recipe_a/b=(1,32)、disable_ue8m0_cast
gemm_nt_bf16bf16f32	bf16_gemm_nt	—
tf32_hc_prenorm_gemm	tf32_hc_prenorm_gemm	num_splits

* 这三个是 sgl-deep-gemm wheel 在 masked 接口上的 overlap 扩展，上游 v0.2.0 tag 的等价签名不含它们（§3.2）。

3.8 MoE 数据布局一览

hidden_states_scale 的 K 维宽度由激活量化 group（block fp8 为 128、mxfp8 为 32）决定，与权重 scale 相互独立——DeepGemmRunnerInput.activation_scale_block_size（deep_gemm.py:271-273）专门记录这一点。UE8M0 模式下，K 维的每 4 个 fp32 scale 会用 _cast_to_e8m0_with_rounding_up 打包成 1 个 int32（打包后宽度 ÷4）。

张量	masked 布局	contig 布局
激活值	(E, m_pad, K) fp8	(M_all, K) fp8
激活 scale	(E, m_pad, K/bs)（UE8M0 时 K/bs/4 的 int32）	(M_all, K/bs) 同理
权重	(E, 2I, K) / (E, K, I) fp8；MXFP4 为 (E, 2I, K//2) / (E, K, I//2) int8	同左
权重 scale	由 recipe 决定（fp8 block128：K/128；MXFP4：(1,128)/(1,32)），MN-major + TMA 对齐 + UE8M0 打包 int32	同左
分组描述	masked_m：(E,) int32	m_indices：(M_all,) int32
GEMM1 输出	(E, m_pad, 2I) bf16	(M_all, 2I) bf16

其中 bs = 激活量化 group；m_pad 是每 expert 的对齐容量（round_up(rank_tokens, 256)，或 SGLANG_OPT_DG_MASKED_M_CAP=1 时收缩到 round_up(max(masked_m), 256)）；M_all 为 routed 行数 + 对齐 padding。

4. 环境与平台要求

前面所有接口都有隐式前提：它们只在满足特定平台、版本、开关时才会被选中或允许运行。本节把「量化方案 × 并行方式 × runner backend × 环境变量」四个维度的硬性要求汇总，每条都给出代码依据。

4.1 量化方案

方案	平台 / 硬件	scale 约定	依据
BF16	DeepGEMM runner 本身要求 SM90+	无 scale	utils/common.py:298-321
FP8 block（fp32 scale）	SM90+	MN-major + TMA 对齐（DEEPGEMM_NEED_TMA_ALIGNED_SCALES）	configurer.py:61
FP8 block（UE8M0）	SM100 / SM120 / SM121	UE8M0 packed int32，无需再 TMA 对齐	configurer.py:58-60
MXFP8	SM100（CUDA）或 ROCm gfx95；gfx942 在加载期转成 block-fp8	必须 UE8M0，weight_block_size 强制 [1, 32]	fp8.py:2126-2134、fp8.py:304-308
MXFP4 experts（W4A8）	DEEPGEMM_BLACKWELL（SM100）；DSV4 默认开启	UE8M0；recipe (1,128)/(1,32)	environ.py:1525、deep_gemm.py:332-333
cutlass fp8（对照）	SM90 / SM100 / SM120；CUDA 12.0+ (SM90) 或 12.4+ (SM89)	block quant	fp8.py:1322-1331

几个值得展开的点：

• 量化方法如何声明平台门槛：get_min_capability（fp8.py:318-328）返回 NPU→0、MUSA→31、MXFP8+gfx95→95、MXFP8-需转换→94、默认 block-fp8→80 / MXFP8→100；SGLang 启动时用它过滤当前 GPU 可用的量化方法；
• gfx942 特例：_mxfp8_to_block_fp8_required 为真时 MXFP8 权重在加载期直接转为 block-fp8（门槛 94），而不是运行期支持 MXFP8；
• MXFP8 MoE 的硬断言（fp8.py:2126-2134）：(_is_cuda and is_sm100) or (_is_hip and _is_gfx95_supported)，否则直接 RuntimeError；
• MXFP4 的加载路径：SGLANG_DSV4_FP4_EXPERTS（默认 True）决定 DSV4 专家走 MXFP4；SGLANG_DSV4_FP4_DEQUANT=True 时运行期反量化回 FP8（environ.py:1525-1527）。

4.2 并行方式

并行 / 通信形态	要求与约束	依据
EP + deepep（v1）	每 rank 缓冲 SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK（默认 128）；internode_ll 用 FINISHED_SUM_TAG=1024，单 rank 发往另一 rank 的 token 数必须 < 1024；NVSHMEM_QP_DEPTH 自动抬到 max(1024, 2·(max+1))	deepep.py:78-85、:393-398
deepep_v2	强制 deep_gemm runner（BF16 或 FP8 blockwise checkpoint）；ElasticBuffer 自适应通信 SM/QP（SGLANG_DEEPEP_V2_NUM_SMS=0）；FP8 dispatch 需要 SGLang FP8 量化 kernel 可导入	runner.py:85-94、deepep_v2.py:91-119
mooncake	SGLANG_MOONCAKE_EP_NUM_MAX_DISPATCH_TOKENS_PER_RANK（默认 128），断言 ≤ 1024（同 FINISHED_SUM_TAG 约束）	mooncake.py:156-162
nixl	SGLANG_NIXL_EP_NUM_MAX_DISPATCH_TOKENS_PER_RANK，断言 ≤ 1024	nixl.py:228-233
DP attention（no-EP，TP-MoE）	gatherv 路径需 SGLANG_DP_USE_GATHERV（默认 off）；reduce_scatter 路径需 SGLANG_DP_USE_REDUCE_SCATTER（HIP 默认开）且 tp_size == attn_dp_size	deepseek_v4.py:3749-3913
CP（prefill context parallel）	a2a == none 走 dsa_cp_gather/reduce_scatter；否则断言 a2a ∈ {deepep, megamoe, mori}	deepseek_v4.py:3819-3831

DP attention 无 EP 时的 combine 三形态（deepseek_v4.py:3775-3802）：reduce_scatterv（变长、gatherv 对称逆操作，会跳过 MoE 内部 post-experts all_reduce 改由它一并完成 SUM，避免双归约）→ reduce_scatter（MAX_LEN 等长场景）→ dp_scatter（兜底）；三者的共同前提都是 mlp_reduce_scatter=True 的 scoped 标记。

4.3 runner backend

backend	要求	依据
deep_gemm	需 ENABLE_JIT_DEEPGEMM：sm≥90 或 MUSA≥31；SM120/121 还要探测到 fp8_einsum、m_grouped_fp8_fp4_gemm_nt_contiguous、transform_sf_into_required_layout 三个 API 才开启	configurer.py:17-29,55
hpc_ops	仅 SM90（HPC-Ops kernel 只发 sm90a）；且只支持 FP8 blockwise / per-tensor 量化 MoE，量化方法回退到别的 runner 时启动即报错	runner.py:76-83、:150-159
deepep_v2（a2a 维度）	只与 deep_gemm runner 兼容（见 §4.2）	runner.py:85-94
silu_mul_keep_fp32（config 开关）	仅 deep_gemm runner 支持；要求 gated SiLU、无 gemm1_alpha、FP8 专家权重	runner.py:64-68、deep_gemm.py:340-347,367-368

4.4 环境变量与开关

变量	默认	作用
SGLANG_ENABLE_JIT_DEEPGEMM	True	全链路总开关（environ.py:1172）
SGLANG_DEEPGEMM_STANDARD_LAYOUT	“auto”	masked GEMM 布局决策（:1176）
SGLANG_DEEPGEMM_MASKED_MEMORY_BUDGET_FRACTION	0.25	masked 布局显存预算比例（:1177）
SGLANG_OPT_DG_MASKED_M_CAP	False	masked 每 expert 容量收缩到 round_up(max(masked_m), 256)，OOM 缓解（:1183）
SGLANG_OPT_DG_COMPACT_EAGER	False	宽 DP eager prefill 用 compact 路由存储（:1186）
SGLANG_OPT_MASK_DP_PAD_MOE	False	DP attention MAX_LEN pad 行从 MoE dispatch 剔除（topk_ids → -1）（:1192）
SGLANG_JIT_DEEPGEMM_PRECOMPILE	True	启动期 JIT 预热（:1193）
SGLANG_JIT_DEEPGEMM_FAST_WARMUP	False	快速预热（只编高频形状）（:1194）
SGLANG_JIT_DEEPGEMM_COMPILE_WORKERS	4	编译并行进程数（:1195）
SGLANG_DG_CACHE_DIR	~/.cache/sglang/deep_gemm	JIT 产物缓存目录（:1198）
SGLANG_DG_USE_NVRTC	False	改用 NVRTC 编译后端（:1199）
SGLANG_DEEPGEMM_SANITY_CHECK	False	UE8M0 scale 合法性校验（:1201）
SGLANG_DEEPGEMM_PDL	True	PDL（programmatic dependent launch）（:1202）
SGLANG_DEEPGEMM_ON_H20	off	H20 特化：masked GEMM 的 max_block_n 降到 160（deep_gemm.py:89,866-869）
SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK	128	DeepEP v1 per-rank 通信缓冲（:1253）
SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK	128	DeepEP v2 per-rank 通信缓冲（:1255）
SGLANG_DEEPEP_V2_NUM_SMS	0	v2 通信 SM 数（0 = ElasticBuffer 自动选择）（:1257）
SGLANG_DSV4_FP4_EXPERTS	True	DSV4 专家按 MXFP4 加载（:1525）
SGLANG_DSV4_FP4_DEQUANT	False	运行期把 FP4 专家反量化回 FP8（:1527）
SGLANG_DP_USE_REDUCE_SCATTER	HIP 平台默认开	DP-attn MAX_LEN combine 用 reduce_scatter（:876）
SGLANG_DP_USE_GATHERV	off	DP-attn gather 用 all_gatherv（dp_attention.py:656 读取）

5. 并行形态对照

同一条三段式骨架，因 a2a backend 与 attention 并行方式不同，进料（pre_permute）与回程（post_permute）有六种组合：

场景	dispatcher	pre_permute	GEMM 形态	回程与归约
纯 TP（a2a=none，无 attn DP）	StandardDispatcher	standard	masked 或 compact	post_reorder 加权归并；模型层 post_experts_all_reduce（TP>1）
TP-MoE + attn DP（a2a=none）	StandardDispatcher	standard（外包 §2.6.3 的 gather/scatter）	masked 或 compact	可选 reduce_scatterv / reduce_scatter 替代 all_reduce
FlashInfer one-sided A2A	FlashinferDispatcher	flashinfer（转 standard 实现）	compact	转 FlashinferCombineInput，combine 再通信
DeepEP Normal（EP）	DeepEPDispatcher normal	deepep_normal（ep_scatter + output_index）	compact	ep_gather 按 output_index 收回
DeepEP LL	DeepEPDispatcher LL	deepep_ll（masked 直通）	masked	combine 内加权
DeepEP v2	DeepEPDispatcher v2	deepep_v2（三分支）	masked / contig / expanded	masked_slab_to_expand / ep_gather / routewise

两个容易混淆的维度：

• TP-MoE 与 EP 的切分方式不同：TP-MoE 时每个 rank 持有全部 expert 的一部分 intermediate 维（N 切分），路由后仍要走 TP all_reduce（或 §2.6.3 的 reduce_scatter）；EP 时每个 rank 持有一部分 expert 的全部权重，token 通过 all2all 找专家，归约在 ep_gather / combine 里完成，无需 all_reduce。
• masked 与 compact 是两种内存布局，与 dispatcher 无关（选择逻辑见附录 B）：masked 是 (E, m, K) 的「每 expert 一段对齐 slab」；compact 是 (all_tokens, K) 的「expert 连续段 + padding」。前者走 *_masked 接口并以 masked_m 描述有效行，后者走 *_contig 接口并以 m_indices 做行到专家的映射。

6. 量化配方对照

四个配方在 §2 六个阶段中的行为正面对比（平台门槛见 §4.1）：

维度	FP8 block128	MXFP8	MXFP4 experts（V4 默认）	BF16
权重加载处理	requant_block_scale_ue8m0_for_deepgemm（dequant→requant）	_process_mxfp8_moe_weights	transform_sf_into_required_layout(recipe=(1,32))	—
权重值	(E, 2I, K) e4m3	(E, 2I, K) e4m3	(E, 2I, K//2) int8（packed 2×fp4）	(E, 2I, K) bf16
权重 scale	(E, 2I/128, K/128) packed int32 / fp32	(E, 2I/128, K/32) packed	(E, 2I/128, K/32) packed	—
激活量化（运行期）	128 组；UE8M0 列主序 packed	32 组	128 组（激活仍是 fp8）	无
GEMM recipe (A/B)	(None, None)	(1,32)/(1,32)	(1,128)/(1,32)	—

运行期差异三条：

• UE8M0 打包动作只在 DEEPGEMM_SCALE_UE8M0 下发生：_cast_to_e8m0_with_rounding_up、transpose(-1,-2)、scale 宽度 ÷4 这类操作；非 UE8M0 平台反而要把 scale 经 get_mn_major_tma_aligned_tensor 对齐（deep_gemm.py:750-753）；
• _DEEPGEMM_ON_H20 特判：masked down GEMM overlap 的 max_block_n 从 256 降到 160（deep_gemm.py:866-869）；
• SM120 早退：standard pre_permute 的 maybe_pre_permute 直接返回，只消费 standard-layout 激活，且 masking/overlap 相关分支不可用（deep_gemm.py:1246-1250）。

7. 小结

• 一副骨架：dispatch → pre_permute → DeepGemmRunnerCore.run → post_permute → combine。DeepGEMM runner 的全部差异都收敛在 pre/post_permute 的布局转换与 core 的四个 GEMM 分支里；
• 两个正交开关：布局（masked / compact）× 配方（bf16 / fp8 / mxfp8 / mxfp4）。dispatcher 决定 pre/post_permute 的实现；权重类型与 DeepGemmMoeQuantInfo 决定 core 的走向；
• 接口分三层：模型层（DeepGemmRunnerInput/QuantInfo）→ runner core（shape 推导、recipe、激活量化）→ wrapper（entrypoint.py 八个接口：sanity check、execution hook、num_sms 管理）→ deep_gemm 库（JIT 编译的 grouped/dense GEMM）。每层的参数职责在 §3 逐一展开；
• Blackwell 的实质变化：UE8M0 int32 packed scale（MN-major TMA）与 FP4/MXFP8 配方。SGLang 在加载期就把权重与权重 scale 整理成 deep_gemm 最终消费的布局（requant / transform / pack），运行期只负责激活侧量化与布局重排（§2.1、附录 A）；
• V4 的特有之处：MoE 复用 DeepseekV2MoE；is_fp4_experts 默认打开（recipe (1,128)/(1,32)）；swiglu_limit 存在（走 V4 JIT 激活）；DP attention 下有 _run_moe_ffn_dp_sync 协调 gather/scatter 与 all_reduce 的替代关系（§2.6.3）。

附录 A：UE8M0——从 fp32 到 int32 packed 的三条生产线

UE8M0 是只保留指数的 8-bit 浮点（无尾数、无符号，值域即 2 的幂），deep_gemm 要求 scale 以 MN-major、TMA 对齐、4 个 uint8 打包进 1 个 int32 的形式存储。样本分别来自三条生产线：

1. 权重加载期（fp4 → UE8M0）：transform_sf_into_required_layout(recipe=(1,32), disable_ue8m0_cast=False)，两条调用线——fp8 量化方法的 is_fp4_expert 分支（fp8.py:2056-2074）与纯 MXFP4 量化方法（mxfp4.py:650-674，先经 float8_e8m0fnu → float32 的精确往返）；库内完成 cast + 布局重排；
2. 权重加载期（fp8-block / mxfp8）：requant_block_scale_ue8m0_for_deepgemm（fp8_utils.py:1719-1750）→ transform_scale_ue8m0（fp8_utils.py:1854-1878）；调用点 fp8.py:2085-2098；mxfp8 的 _pack_moe_scale_for_deepgemm（fp8.py:2248-2262）同款；
3. 运行期激活量化：
   • contig silu 融合：silu_and_mul_contig_post_quant(scale_ue8m0=UE8M0, transposed=UE8M0)，scale 缓冲由 create_per_token_group_quant_fp8_output_scale(column_major_scales=UE8M0, scale_tma_aligned=UE8M0, scale_ue8m0=UE8M0) 直接按 packed 布局分配（deep_gemm.py:530-552）；
   • masked silu：per_token_group_quant(scale_ue8m0=UE8M0, column_major_scales=True) 自带 packed 布局输出（deep_gemm.py:1635-1643，注释说明调用方的 get_mn_major transform 会短路）；
   • masked situ 三分支：JIT kernel 输出后补 transpose(-1,-2)（deep_gemm.py:1528 / 1595 / 1627）；
   • 兜底：_cast_to_e8m0_with_rounding_up（deep_gemm.py:96-106；调用点 :501 contig、:747 masked）。

_cast_to_e8m0_with_rounding_up 的位运算：把 fp32 看成 int32，取指数位 exp = bits >> 23、尾数 mant = bits & 0x7FFFFF；当「尾数非零且指数非 0xFE（inf/nan 保护），或指数为 0 时尾数足够大」时指数 +1（round up），再取低 8 位作为 uint8，最后 transpose(1,2).contiguous().transpose(1,2) 转 MN-major 对齐。这也是 _sanity_check_input 断言「fp32 scale 必须已是 2 的幂」的原因——非 2 的幂会被 round-up cast 改变数值。

附录 B：masked 还是 compact——自动选择的过滤链

_should_use_masked_standard_layout（deep_gemm.py:196-249）依次过滤：

1. SGLANG_OPT_DG_COMPACT_EAGER=1 且非 capture → 强制 compact（逃生舱）；
2. _masked_activation_unsupported_reason（:175-193）给出理由 → compact。理由只在「clamped/swizzled 激活」（swiglu_limit 存在或 megamoe）下出现，条件为 D//8 < E、group_size != 128、D % (4*group) != 0 之一（DSV4 JIT kernel 的硬约束）；
3. 环境变量 SGLANG_DEEPGEMM_STANDARD_LAYOUT：masked / compact 直接定夺；auto 继续；
4. auto：CUDA Graph capture 中 → compact（capture 内不能查显存）；否则估算 masked 峰值（_estimate_masked_standard_layout_peak_bytes :132-167：E × padded_m × 单行峰值字节，单行峰值取三项差的最大值）与预算 min(free_memory) × SGLANG_DEEPGEMM_MASKED_MEMORY_BUDGET_FRACTION 比较——放得下就 masked，放不下就 compact。

几何对比：masked 用 (E, m, K) slab + masked_m/expected_m，workspace 是 E × padded_m 量级；compact 用 (all_tokens, K) + m_indices/src2dst，all_tokens 由 _get_compact_all_tokens（:252-259）给出上界：block_e × (非空 expert 数 + 剩余 tokens 摊入的块数)。

附录 C：ep_scatter 家族与三个索引

布局重排的底层 kernel（kernels/ops/moe/ep_moe_kernels.py）按场景分三支：

• ep_scatter：standard compact / deepep_normal 使用——按 tokens_per_expert（或 recv 计数）把 token 与激活 scale 散进 expert 连续段，同时产出 m_indices（行→expert）与 src2dst/output_index（回程映射）；expert_alignment=block_e 保证每 expert 段按 GEMM 对齐要求 padding；
• ep_scatter_from_psum：deepep_v2 非 expanded 分支，用原生 psum 前缀和替代计数数组；
• moe_ep_deepgemm_preprocess（masked）：一步融合「路由分桶 + slab 填充 + 激活量化」，直接产出 (E, m, K) 与 src2dst。

三个索引的分工（易混）：

索引	形状	语义	消费者
m_indices	(M,) int32	grouped 行 → expert id	contig GEMM
src2dst	(M,) int32	输出行 → 原 token 槽位	post_reorder_deepgemm
output_index	(M,) int32	回程映射（v2 中 -1 = 非本地）	ep_gather / v2 的 route 还原

附录 D：启动期 JIT 预热（compile_utils.py）

deep_gemm 是 JIT 编译的库，SGLang 在启动/捕获前做形状预热：

• 缓存与编译后端：DG_JIT_CACHE_DIR（:46）、DG_JIT_USE_NVRTC（:51）、update_deep_gemm_config（:54-105）设置 SGLANG_DG_CACHE_DIR / SGLANG_DG_USE_NVRTC 等库侧全局；
• _local_rank_compile_lock（:120-140）：flock 文件锁防止多进程重复编译；
• _maybe_compile_deep_gemm_one_type_all（:144-186）：对 _BUILTIN_M_LIST（:37，一组预置 M）逐类型编译，contig 类型按 m % alignment == 0 过滤（:201-206）；
• 显存预算不足时收缩 max_m（:213-234）；
• compile mode 区分 warmup/运行模式（:241-252）；
• 多组 warmup executor 覆盖 contig/masked × fp8/bf16（:263 / 338 / 352 / 368 / 384 / 404 / 414）；
• deep_gemm_execution_hook（:436-451）在运行期遇到未预热形状时记录/补编；
• PP 场景 pp_parallel_deep_gemm_warmup（:454-542）。

附录 E：平台差异速查

项	SM90	SM100（含 103）	SM120/121
ENABLE_JIT_DEEPGEMM	✓	✓	✓（额外探测 fp8_einsum 等 API）
DEEPGEMM_BLACKWELL	✗	✓	✗
DEEPGEMM_SCALE_UE8M0	✗	✓	✓
DEEPGEMM_NEED_TMA_ALIGNED_SCALES	✓	✗	✗
本文主路径	FP8/BF16	全部四配方	独立 sm120 模块（standard layout only）

（is_sm100 在 SGLang 中即 capability[0] == 10，覆盖 100/103；平台判定见 utils/common.py:326-334。注：MUSA 平台在 configurer.py:55 也参与 ENABLE_JIT_DEEPGEMM 判定（capability ≥ 31），且 get_min_capability 对 MUSA 返回 31（fp8.py:318-328）；但 DEEPGEMM_BLACKWELL / DEEPGEMM_SCALE_UE8M0 仅按 capability 判定，MUSA 上不成立。）

附录 F：代码位置速查

主题	位置（v0.5.21 本地检出）
三段式骨架	layers/moe/fused_moe_triton/layer.py:1562-1599（forward_impl）、:1642-1647（run_moe_core）
dispatcher 工厂	同上 :162-228（create_moe_dispatcher）
DeepEPMoE 退化开关	layers/moe/ep_moe/layer.py:106-153（deprecate_flag）
触发判定	layers/quantization/fp8.py:1334-1356、:2743-2801、:2004
权重预处理	fp8.py:2004-2098（fp4 与 block 两路）、:2126-2417（MXFP8）、fp8_utils.py:1719-1878（requant/pack）、fp8.py:205-247（cast_e2m1fn_to_e4m3fn）
quant_info	moe_runner/deep_gemm.py:289-334（字段 + scale_recipes）；fp8.py:3064-3073（组装）；mxfp4.py:1597-1605
五种 pre_permute	moe_runner/deep_gemm.py:970 / 1194 / 1312 / 1354 / 1698
runner core	同上 :337-393；四个分支 :395 / :628 / :695 / :898
激活量化	同上 :530-552（contig silu 融合）、:1533-1643（silu 三分支）、:1486-1530（situ）
wrapper 接口	deep_gemm_wrapper/entrypoint.py:48 / 109 / 132 / 162（grouped 四接口）、:219 / 240 / 266 / 280（dense 四接口）、:173-216（alignment）、:301-328（num_sms / sanity）
平台开关	deep_gemm_wrapper/configurer.py:17-61
JIT 预热	deep_gemm_wrapper/compile_utils.py（:54-105、:144-186、:436-451、:454-542）
DP attention 同步	models/deepseek_v4.py:3749-3913
模型层尾部归约	models/deepseek_v2.py:1368-1393（含 post_experts_all_reduce :1385）

---

至此，DeepSeek V4 在 Blackwell 上走 DeepGEMM MoE 的完整路径已经闭合：从 DeepseekV2MoE 的路由，到五种 dispatch 布局的进料转换、四个 GEMM 分支的 core 执行、八个 wrapper 接口的最终下推，以及 post_permute 与 combine 的回程归约。文中所有行号基于本地 release/v0.5.21 检出逐一核对；后续版本若有重构，请以 moe_runner/deep_gemm.py 与 deep_gemm_wrapper/entrypoint.py 的最新实现为准。