执行摘要

• 一句话：修复 MXFP8 MoE 权重更新后 GPU 缓存索引脏数据问题
• 推荐动作：值得精读。该 PR 揭示了 GPU 缓存与动态权重卸载之间的典型交互问题，修复方案简洁但覆盖面完整。review 中关于性能权衡的讨论对设计类似缓存失效策略有参考价值。

功能与动机

DeepSeek-V4 MXFP8 在 flashinfer_trtllm_routed MoE 路径上，RL 权重更新后 train_rollout_logprob_abs_diff 从 ~0.06 跳升到 ~3.83，严重影响训练收敛。原因是 align_mxfp8_moe_weights_for_flashinfer_trtllm 缓存的 shuffle index 是 GPU 驻留的，权重更新时内存被释放重分配，缓存指向陈旧内存，导致后续排列使用脏数据。

实现拆解

1. 在 python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py 中新增 clear_mxfp8_shuffle_index_cache() 函数，用于清空模块级 GPU 缓存字典 _flashinfer_trtllm_shuffle_row_indices_cache_mxfp8。
2. 在 python/sglang/srt/layers/moe/fused_moe_triton/layer.py 的两个权重加载入口 _weight_loader_impl 和 weight_loader_fused 中，当 method 是 Fp8MoEMethod 且后端为 flashinfer_trtllm 或 flashinfer_trtllm_routed 时，调用 clear_mxfp8_shuffle_index_cache()，确保每次重载都重新计算索引。
3. 在 test/registered/rl/test_update_weights_from_disk_blackwell.py 中将 update_timeout 从 240 秒降低到 120 秒，以更紧的预算捕捉性能回归。
4. 在 hash_topk.py 中实现了 apply_routed_scaling_factor_on_output（该文件变更未在当前上下文中展示），启用 V4 的 routed 缩放因子。

关键文件：

• python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py（模块 MoE后端；类别 source；类型 core-logic；符号 clear_mxfp8_shuffle_index_cache）: 核心修复：新增 clear_mxfp8_shuffle_index_cache() 函数，提供清除 GPU 缓存的能力。
• python/sglang/srt/layers/moe/fused_moe_triton/layer.py（模块 融合MoE；类别 source；类型 dependency-wiring；符号 _weight_loader_impl, weight_loader_fused）: 权重加载入口：在两个主要权重加载函数中条件调用缓存清除，覆盖 colocate 和分布式 EP 更新路径。
• test/registered/rl/test_update_weights_from_disk_blackwell.py（模块 RL测试；类别 test；类型 test-coverage）: 测试调整：降低 update_timeout 从 240s 到 120s，以便更早发现性能退化。

关键符号：clear_mxfp8_shuffle_index_cache, _weight_loader_impl, weight_loader_fused

关键源码片段

python/sglang/srt/layers/moe/fused_moe_triton/layer.py

权重加载入口：在两个主要权重加载函数中条件调用缓存清除，覆盖 colocate 和分布式 EP 更新路径。

# 在 _weight_loader_impl 中，加载权重时检测 FP8 MoE 和 flashinfer_trtllm 后端
elif isinstance(method, Fp8MoEMethod) and (
    get_moe_runner_backend().is_flashinfer_trtllm_routed()
    or get_moe_runner_backend().is_flashinfer_trtllm()
):
    # 丢弃陈旧的 GPU shuffle 索引缓存，确保下次 align 使用新计算值
    from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
        clear_mxfp8_shuffle_index_cache,
    )
    clear_mxfp8_shuffle_index_cache()

评论区精华

• 条件强化（Fridge003）：建议只对 mxfp8 且 trtllm 后端时调用 clear，避免不必要开销。作者已修复，加入 isinstance(method, Fp8MoEMethod) and (get_moe_runner_backend().is_flashinfer_trtllm_routed() or ...) 判断。
• 性能权衡（zianglih）：指出 cache 是 #21280 有意添加的性能优化，每次清除重新计算可能导致几分钟延迟。提出两种方案：要么仅在 colocate RL offload 触发时清除（Option 1），要么只清除一次并允许后续复用（Option 2）。作者实现为全局清除但限制后端，随后 zianglih 通过基准测试验证无额外更新开销（<10s）。

• 测试超时调整（zianglih）：建议将 update_timeout 从 240 秒改为 10 秒以捕获性能退化。作者认为 10 秒太紧，设置为 120 秒足够通过测试。

• 条件加强 (design): 采纳，在条件中加入 isinstance 和 backend 判断。

• 性能影响 (performance): 采用当前实现，基准测试确认无性能退化。
• 测试超时设置 (testing): 折中设为 120 秒。

风险与影响

• 风险：
  • 性能回归风险：每次权重重载都清除缓存并重新计算 shuffle index，可能引入额外延迟。但 zianglih 在 DeepSeek-V3.2-MXFP8 TP8 上实测验证更新时间 <10s，与优化前一致，风险可控。
  • 影响范围：仅影响 Fp8MoEMethod 且后端为 flashinfer_trtllm 或 flashinfer_trtllm_routed 的路径，不会波及非 FP8 或其他后端。
  • 兼容性：缓存清除逻辑与现有 memory 管理机制（pause/resume）配合，不会引入新竞态条件。
• 影响：
  • 用户：使用 DeepSeek-V4 MXFP8 进行 RL 训练的用户将不再遇到权重更新后数值炸裂问题，rollout logprob 差异稳定在 ~0.06。
  • 系统：影响所有启用 flashinfer_trtllm 后端的 FP8 MoE 模型权重更新路径，但开销极低。
  • 团队：修复了 RL 训练管线中的关键阻塞性 bug，降低后续训练调优的成本。
  • 风险标记：GPU缓存脏数据, 核心权重加载路径, 性能回归已验证无

关联脉络

• PR #21280 Add MXFP8 shuffle index cache for flashinfer trtllm: 引入的缓存正是本 PR 修复的根源，清除操作撤销了该缓存导致的脏数据问题。