Triton 中的跨平台融合 MoE 分配：无需 CUDA 的可移植专家路由

混合专家（MoE）架构为大多数前沿大型语言模型提供动力，但其推理受益于不规则内存访问模式和专家路由开销的瓶颈。现有的优化 MoE 内核（Megablocks、Tutel、FasterMoE）是用 CUDA 实现的，仅限于 NVIDIA 硬件。我们提出 TritonMoE，一个完全使用 OpenAI Triton 编写的融合 MoE 分配内核，完成整个前向传递——包括路由评分、token 置换、专家 GEMMs 和加权输出组合——仅使用可移植的 Triton 原语。我们的关键优化是融合门+up GEMM 内核，通过共享 L2 缓存输入块的 in-register SiLU 激活，计算两个 SwiGLU 投影，从而消除 35\% 的全局内存流量。在 NVIDIA A100 上，TritonMoE 在 Mixtral-8x7B、DeepSeek-V3 和 Qwen2-MoE 配置下，在推理 batch size（≤512 token）范围内实现了 CUDA 优化 Megablocks 的 89-131\% 吞吐量。所有 162 个正确性测试在 NVIDIA A100 和 AMD MI300X 上均通过，无需代码更改，验证了跨平台可移植性。我们还表征了在 Zipfian 分布的专家分配下的路由不平衡敏感性，并确定了固定图块调度在极端不平衡下（64+ 个专家）不及 Megablocks 块稀疏布局的情形，这激励了动态块到专家分配作为未来工作。代码可在 https://github.com/bassrehab/triton-kernels 上获取。