执行摘要

• 一句话：移除 DeepSeek V4 的 AMD MegaMoE 集成
• 推荐动作：值得阅读以了解清理大规模遗留代码的模式。此 PR 展示了如何系统性地识别和移除不再维护的功能块，同时保持核心逻辑稳定。对于关注 AMD 后端 DeepSeek 模型实现的工程师，应仔细检查是否还有其他遗留引用。

功能与动机

根据 PR 描述，这是一个简单的清理：A simple cleanup: Remove MegaMoE in deepseek_v4/amd/。可能由于 MegaMoE 在 AMD 上不再维护或已被标准 MoE 替代，移除可降低维护成本。

实现拆解

1. 在 vllm/models/deepseek_v4/amd/model.py 中移除 Triton kernel _deepseek_v4_stage_mega_moe_inputs_kernel 和 _stage_deepseek_v4_mega_moe_inputs，以及相关的工具函数 make_deepseek_v4_expert_params_mapping、_ue8m0_uint8_to_float 等。同时删除自定义 MoE 类 DeepseekV4MegaMoEExperts 及其 __init__、_map_global_expert_id、weight_loader 方法。
2. 清理不再需要的导入，包括 triton、tl、direct_register_custom_op、fused_topk_bias 等，减少模块依赖。
3. 在 vllm/models/deepseek_v4/amd/mtp.py 中，移除对 make_deepseek_v4_expert_params_mapping 的导入，删除 load_weights 中基于 use_mega_moe 的条件分支，直接使用 FusedMoE.make_expert_params_mapping。
4. 删除 finalize_mega_moe_weights 方法及其调用逻辑，因为 MegaMoE 权重最终化不再需要。
5. 整体减少约 645 行代码，使得 AMD DeepSeek V4 实现与 NVIDIA 版本保持一致的标准 MoE 实现。

关键文件：

• vllm/models/deepseek_v4/amd/model.py（模块 AMD模型；类别 source；类型 core-logic；符号 _deepseek_v4_stage_mega_moe_inputs_kernel, _stage_deepseek_v4_mega_moe_inputs, make_deepseek_v4_expert_params_mapping, DeepseekV4MegaMoEExperts）: 核心模型文件，移除了 MegaMoE Triton kernel、自定义专家映射类 DeepseekV4MegaMoEExperts 及相关辅助函数，清理了导入依赖，简化了 MoE 层
• vllm/models/deepseek_v4/amd/mtp.py（模块 MTP层；类别 source；类型 data-contract；符号 finalize_mega_moe_weights）: MTP 模型文件，移除对 make_deepseek_v4_expert_params_mapping 的依赖，删除 finalize_mega_moe_weights 方法，统一使用 FusedMoE 专家映射

关键符号：_deepseek_v4_stage_mega_moe_inputs_kernel, _stage_deepseek_v4_mega_moe_inputs, make_deepseek_v4_expert_params_mapping, DeepseekV4MegaMoEExperts.init, _map_global_expert_id, weight_loader, _ue8m0_uint8_to_float, finalize_mega_moe_weights

关键源码片段

vllm/models/deepseek_v4/amd/mtp.py

MTP 模型文件，移除对 make_deepseek_v4_expert_params_mapping 的依赖，删除 finalize_mega_moe_weights 方法，统一使用 FusedMoE 专家映射

# vllm/models/deepseek_v4/amd/mtp.py (head 版本 )
​
# 导入部分：移除了 make_deepseek_v4_expert_params_mapping 的导入，
# 现在仅从 .model import DeepseekV4DecoderLayer
​
class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
    # ... 其他方法
​
    def load_weights(self, weights):
        # ... 前面的代码
​
        # Pre-compute expert mapping ONCE.
        # 原来此处有条件分支判断 use_mega_moe，现在直接使用 FusedMoE 映射
        expert_mapping = FusedMoE.make_expert_params_mapping(
            self,
            ckpt_gate_proj_name='w1',
            ckpt_down_proj_name='w2',
            ckpt_up_proj_name='w3',
            num_experts=self.config.n_routed_experts,
        )
​
        # ... 权重加载循环
​
        # 原来会调用 self.finalize_mega_moe_weights()，该方法已被删除
        logger.info_once('MTP draft model loaded: %d params', len(loaded_params))
        return loaded_params

评论区精华

没有人工审查评论。仅由自动代码审查机器人 gemini-code-assist[bot] 进行了一次总结性评论，确认变更为清理性质，未提出具体问题或建议。

• 自动代码审查总结 (other): 无实质讨论，变更可安全合并。

风险与影响

• 风险：移除大量专用代码可能导致潜在兼容性问题，如果某些配置或外部调用仍然引用 MegaMoE 特性（如 use_mega_moe 标志）。但审查确认所有 MegaMoE 相关代码已在两个文件中完全移除，且标准 MoE 路径已能完全等效替代。风险较低。建议在 AMD ROCm CI 环境中运行 DeepSeek V4 模型的回归测试，以确认推理正确性。
• 影响：影响范围限于 AMD ROCm 平台上的 DeepSeek V4 模型推理。用户将无法使用 MegaMoE 特有的优化（如果之前有使用），但标准 FusedMoE 路由会继续正常工作。代码维护负担降低，未来开发不需要维护两套 MoE 实现。
• 风险标记：核心路径变更, 功能移除

关联脉络

• PR #43690 [DSv4] Drop _get_compressed_kv_buffer in DeepseekCompressor: 同一作者（WoosukKwon）对 DeepSeek V4 的另一项清理工作，移除不再使用的 KV 缓冲区管理代码，体现了对 DeepSeek V4 代码的持续简化。