执行摘要

• 一句话：新增 DeepSeek V4 w4a4 MegaMoE 推理支持
• 推荐动作：建议阅读此 PR 以了解如何在 SGLang 中新增 DeepGEMM 后端的量化选项。环境变量转导模式（_apply_mega_moe_dg_env）是一个简洁的跨库配置传递方式，值得借鉴。团队应跟进 sgl-deep-gemm 版本发布，并考虑在更多模型上验证 w4a4 效果。

功能与动机

该 PR 是 #24444 在 main 分支的重新提交。主要目的在于启用 DeepSeek V4-Flash 的 w4a4 量化推理，通过 FP4 激活存储和 MXF4 计算主循环在保持精度的同时提升推理效率。需要新版本 sgl-deep-gemm（#31）配合使用。

实现拆解

实现分为以下步骤：

1. 环境变量声明：在 python/sglang/srt/environ.py 的 Envs 类中新增 SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS 和 SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_MXF4_KIND 两个布尔环境变量，默认均为 False。
2. 环境变量转发：在 python/sglang/srt/layers/moe/mega_moe.py 中新增 _apply_mega_moe_dg_env() 函数，通过 os.environ.setdefault 将 SGLANG 的 FP4/MXF4 标志转设为 DG_USE_FP4_ACTS 和 DG_USE_MXF4_KIND，供 DeepGEMM 的 fp8_fp4_mega_moe 和 get_symm_buffer_for_mega_moe 读取。该函数在 _get_mega_moe_symm_buffer 开头调用，保证首次分配对称缓冲区时就设置好。
3. 预调度分支：在 _run_mega_routed 函数中，根据 use_fp4_acts 标志分流：若为 True，调用 deep_gemm.mega_moe_pre_dispatch（处理 E2M1 打包）；否则保持原有 mega_moe_pre_dispatch（JIT 版本，仅 FP8）。
4. 依赖升级：在 python/pyproject.toml 中将 sgl-deep-gemm 从 0.0.1 改为 0.1.0，以包含新的 mega_moe_pre_dispatch 绑定。
5. 测试重构：将原先位于 test/registered/dsv4/test_deepseek_v4_flash_fp4_b200.py 中的 TestDSV4FlashFP4B200MegaMoE 类及其相关环境配置（_MEGAMOE_ENV）删除，迁移至新增的 test/registered/dsv4/test_deepseek_v4_flash_fp4_megamoe_b200.py。新测试文件包含两个测试类：TestDSV4FlashFP4B200W4A8MegaMoE（FP8 激活基准）和 TestDSV4FlashFP4B200W4A4MegaMoE（新增 FP4 激活变体），均通过 GSM8K 精度门限（>0.93）。

关键文件：

• test/registered/dsv4/test_deepseek_v4_flash_fp4_megamoe_b200.py（模块 集成测试；类别 test；类型 test-coverage；符号 _gsm8k_check, TestDSV4FlashFP4B200W4A8MegaMoE, setUpClass, tearDownClass）: 新增的 CI 测试文件，直接验证 W4A8 和 W4A4 两种 MegaMoE 推理路径的 GSM8K 精度，是 w4a4 特性的质量门禁。
• python/sglang/srt/layers/moe/mega_moe.py（模块 MoE层；类别 source；类型 core-logic；符号 _apply_mega_moe_dg_env）: 核心逻辑文件，新增 _apply_mega_moe_dg_env 函数将 SGLANG 环境变量转发至 DeepGEMM，并修改预调度分支支持 FP4 路径。
• test/registered/dsv4/test_deepseek_v4_flash_fp4_b200.py（模块 集成测试；类别 test；类型 test-coverage；符号 TestDSV4FlashFP4B200MegaMoE, setUpClass, tearDownClass, test_gsm8k）: 旧测试文件删除了 MegaMoE 相关测试类和环境配置，将测试迁移至新文件，是测试重构的一部分。
• python/sglang/srt/environ.py（模块 配置；类别 source；类型 core-logic）: 声明新的环境变量 SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS 和 SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_MXF4_KIND，是特性配置入口。
• python/pyproject.toml（模块 依赖管理；类别 config；类型 configuration）: 升级 sgl-deep-gemm 依赖版本至 0.1.0，确保包含新的预调度绑定。

关键符号：_apply_mega_moe_dg_env, _get_mega_moe_symm_buffer, _run_mega_routed

关键源码片段

python/sglang/srt/layers/moe/mega_moe.py

核心逻辑文件，新增 _apply_mega_moe_dg_env 函数将 SGLANG 环境变量转发至 DeepGEMM，并修改预调度分支支持 FP4 路径。

import os
from sglang.srt.environ import envs
​
_MEGA_MOE_DG_ENV_APPLIED = False
​
def _apply_mega_moe_dg_env() -> None:
    """将 SGLANG 的 FP4/MXF4 标志转发到 DeepGEMM 环境变量。
​
    DeepGEMM 在运行时读取 DG_USE_FP4_ACTS 和 DG_USE_MXF4_KIND。
    使用 setdefault 确保外部显式设置的变量优先。
    """
    global _MEGA_MOE_DG_ENV_APPLIED
    if _MEGA_MOE_DG_ENV_APPLIED:
        return
    if envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS.get():
        os.environ.setdefault("DG_USE_FP4_ACTS", "1")
    if envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_MXF4_KIND.get():
        os.environ.setdefault("DG_USE_MXF4_KIND", "1")
    _MEGA_MOE_DG_ENV_APPLIED = True
​
# 在 _get_mega_moe_symm_buffer 开头调用 _apply_mega_moe_dg_env()
# ... ( 省略 SymmBuffer 获取逻辑 )
​
# _run_mega_routed 中的分支：
use_fp4_acts = envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS.get()
if use_fp4_acts:
    # FP4 路径：调用 DeepGEMM 的原生 pre_dispatch，处理 E2M1 打包
    deep_gemm.mega_moe_pre_dispatch(
        hidden_states, topk_ids_in, topk_weights_in,
        buf.x, buf.x_sf, buf.topk_idx, buf.topk_weights,
        num_tokens=num_tokens, group_size=32, use_fp4_acts=True)
else:
    # FP8 路径：使用 JIT kernel 实现
    mega_moe_pre_dispatch(
        hidden_states, topk_ids_in, topk_weights_in,
        buf.x, buf.x_sf, buf.topk_idx, buf.topk_weights,
        quant_group_size=32)

评论区精华

PR 主要由作者 Fridge003 驱动，review 仅包含作者自己的评论：

• 在 environ.py 的 diff 上建议引入 SGLANG_OPT_MEGA_MOE_FUSED_PRE_DISPATCH 环境变量（未出现在最终代码中）。

• 在 pyproject.toml 上建议使用正式版本 sgl-deep-gemm==0.1.0 而非 0.1.0rc0（最终采用正式版本）。
  此外，作者在 issue 评论中发布了详细的验证信息和基准测试结果。没有外部 reviewer 参与实质讨论。

• 是否添加 SGLANG_OPT_MEGA_MOE_FUSED_PRE_DISPATCH 环境变量 (design): 未采纳，最终代码未包含该变量，保持现有设计。

• 依赖版本选择：rc0 还是正式版本 (design): 采纳，最终版本使用 sgl-deep-gemm==0.1.0。

风险与影响

• 风险：
  1. 外部依赖风险：需要 sgl-deep-gemm 0.1.0 发布并包含 #31 的更改，若该库发布延迟则此特性被阻塞。
  2. 环境变量转发准确性：_apply_mega_moe_dg_env() 仅在首次调用 SymmBuffer 时设置一次，后续若动态更改 SGLANG 变量不会生效。但该变量设计为静态配置，因此风险较低。
  3. 硬件限定：FP4 路径仅在 B200 GPU 上经过验证，其它 GPU（如 H100）上的行为未测试。
  4. 测试覆盖：新增测试仅运行 GSM8K 精度门限（200 样本），未覆盖极端负载或长时间运行稳定性。
  5. 内存占用：FP4 激活可减小 SymmBuffer 占用一半，但仍需要准确的 NUM_MAX_TOKENS_PER_RANK 配置，过大可能导致 OOM。
• 影响：影响范围：
  • 对用户：使用 deepseek-ai/DeepSeek-V4-Flash 模型的用户可通过设置 SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS=1 和 SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_MXF4_KIND=1 启用新量化路径，获得潜在的性能提升和显存节省。
  • 对系统：默认行为不变（两个新变量默认 False），不影响现有用户。
  • 对团队：需要维护新的测试文件和依赖版本，CI 中加入 B200 特定测试。
  • 影响程度：中等，仅限于特定模型和硬件配置。
  • 风险标记：依赖外部库新版本, 仅 B200 验证, 环境变量转导静态单次设置, 测试覆盖有限

关联脉络

• PR #24884 [MoE] Decouple Mega MoE from DeepEP backend: 本 PR 在 #24884 引入的 MegaMoE 框架基础上增加了 FP4 激活支持，是对该功能的扩展。
• PR #24444 DeepSeek V4 w4a4 MegaMoE (original): 本 PR 是 #24444 在 main 分支的重基，原始 PR 的继续。