---
title: 'TritonMoE: Portable Fused Mixture-of-Experts'
url: https://www.emergentmind.com/topics/tritonmoe
type: topic
---

# TritonMoE: Portable Fused Mixture-of-Experts

Searching arXiv for TritonMoE and related MoE kernel papers to ground the article in current literature.
TritonMoE is a fused Mixture-of-Experts dispatch implementation introduced in "Cross-Platform Fused MoE Dispatch in Triton: Portable Expert Routing Without CUDA" [2605.23911]. It targets the MoE inference path in Transformer blocks, where irregular memory access patterns and expert routing overhead commonly dominate latency. The system replaces the dense feed-forward sublayer with an $E$-expert MoE layer and implements the complete inference forward pass in OpenAI Triton rather than CUDA, with the stated objective of preserving portability across NVIDIA and AMD accelerators while retaining performance competitive with CUDA-optimized MoE kernels such as Megablocks, Tutel, and FasterMoE [2605.23911].

## 1. Architectural position and problem setting

TritonMoE is situated at the intersection of sparse expert routing and kernel fusion. In the formulation described by the paper, the MoE layer performs router scoring, token-to-expert assignment, expert-specific feed-forward computation, and output recombination within a single Triton-based execution pipeline [2605.23911]. The immediate technical problem is that MoE inference is bottlenecked not only by GEMM throughput but also by routing-induced permutation overhead and nonuniform expert loads.

The paper frames TritonMoE against a hardware-software landscape in which existing optimized MoE kernels are implemented in CUDA and locked to NVIDIA hardware [2605.23911]. TritonMoE therefore addresses two coupled questions: whether the forward path of an MoE block can be implemented entirely in portable Triton primitives, and whether such an implementation can remain competitive with specialized CUDA kernels at inference batch sizes typical of LLM serving.

A central design choice is that the implementation covers the *entire inference forward pass* rather than only isolated expert GEMMs. This gives the system a dispatch-centric character: routing, permutation, grouped expert computation, and recombination are treated as one fused execution problem rather than as a sequence of loosely connected library calls. A plausible implication is that the contribution is as much about end-to-end scheduling as about individual kernel micro-optimizations.

## 2. Five-kernel forward-pass realization

The forward pass is implemented in five Triton kernels, with a fixed kernel-launch count independent of the number of experts $E$ [2605.23911]. The paper contrasts this with a naive loop-over-experts implementation requiring $(3E+4)$ launches.

| Stage | Operation |
|---|---|
| 1 | Router scoring and top-$k$ selection |
| 2 | Permutation to expert-contiguous layout |
| 3 | Fused gate+up projection |
| 4 | Down projection as grouped GEMM |
| 5 | Unpermutation and weighted combination |

The first stage computes router probabilities as
$$
\mathbf{s} = \mathrm{softmax}(\mathbf{W}_r \mathbf{X})\in\mathbb{R}^{B\times E},
$$
followed by iterative top-$k$ expert selection with masking [2605.23911]. The second stage scatters tokens from the original layout $(B,d)$ to an expert-contiguous layout $(T,d)$, where $T = B\times k$. This reordering is the structural precondition for the grouped GEMM stages that follow.

The third stage is the fused gate+up projection for SwiGLU experts. The fourth stage performs the down projection as a block-scheduled grouped GEMM over the intermediate activations $(T,\,d_{\mathrm{ffn}})$ and per-expert down weights. The fifth stage gathers outputs back to $(B,k)$ slots, multiplies by routing weights, and sums over the $k$ expert contributions per token [2605.23911].

This five-stage decomposition is important because it preserves explicit MoE semantics while constraining launch overhead. The paper’s formulation suggests that TritonMoE is not merely a custom expert matmul but a complete MoE dispatch substrate for inference.

## 3. Fused gate+up GEMM and the memory-traffic model

The key optimization is a fused gate+up GEMM that computes both SwiGLU projections from shared L2-cached input tiles, with in-register SiLU activation [2605.23911]. For an input $\mathbf{x}\in\mathbb{R}^d$, the expert FFN is written as
$$
G(\mathbf{x}) = \mathbf{x}W_g,\quad
U(\mathbf{x}) = \mathbf{x}W_u,\quad
\mathrm{SwiGLU}(\mathbf{x}) = \mathrm{SiLU}\bigl(G(\mathbf{x})\bigr)\odot U(\mathbf{x}),
$$
with $\mathrm{SiLU}(z)=z\cdot\sigma(z)$.

In TritonMoE’s fused realization, a tile of the permuted input and the corresponding weight tiles for $W_g$ and $W_u$ are loaded, two accumulators are updated in registers,
$$
\mathrm{acc}_g \;\mathrel{+}=\;\mathbf{A}\,\cdot\,W_g,\quad
\mathrm{acc}_u \;\mathrel{+}=\;\mathbf{A}\,\cdot\,W_u,
$$
and after the $K$ loop the kernel computes $\mathrm{SiLU}(\mathrm{acc}_g)\odot\mathrm{acc}_u$ entirely in FP32 registers before writing a single $(M\times d_{\mathrm{ffn}})$ tile back [2605.23911].

The memory-traffic argument is explicit. Letting $T=Bk$, $F=d_{\mathrm{ffn}}$, hidden dimension $d$, and assuming 2 bytes per FP16 element, the unfused formulation has total traffic approximately
$$
8TF + 4Td,
$$
while the fused formulation has total traffic
$$
2TF + 2Td.
$$
The stated savings are
$$
(6TF + 2Td)/(8TF+4Td)\approx 35\%
$$
for typical $F\gg d$ [2605.23911]. In the paper’s interpretation, this elimination of intermediate-buffer traffic is the principal reason the Triton implementation can compete with CUDA baselines despite using only portable primitives.

## 4. Scheduling, routing, and Triton-level execution strategy

The down projection and related grouped GEMMs use a block-scheduled grouped GEMM schedule in which expert offsets are precomputed, token counts per expert are converted into block ranges, and each Triton program block selects one $(\text{expert}, \text{offset})$ pair for execution [2605.23911]. This transforms expert imbalance into a block scheduling problem rather than a sequence of expert-local launches.

The router kernel uses in-register top-$k$ selection and stable softmax. The reported sequence is: compute scores with `tl.dot`, subtract the rowwise maximum, exponentiate, normalize, and then iteratively extract top-$k$ experts by `argmax` while masking selected indices [2605.23911]. The paper further notes that all of this is done in Triton, with no external library calls except that the initial $XW_r$ may use cublas for small $E$.

The fused gate+up kernel is described as operating entirely through Triton primitives, with L2-cached input reuse and in-register activation. In combination with the expert-contiguous permutation, this establishes a consistent execution model: routing creates a packed expert-major layout, grouped GEMMs operate on that packed representation, and unpermutation restores token order at the output.

A plausible implication is that TritonMoE’s portability is inseparable from this kernel structure. By avoiding inline CUDA and vendor-specific intrinsics, the implementation constrains itself to a common subset of operations that Triton’s backend can lower across architectures.

## 5. Reported performance and correctness

The evaluation is reported on NVIDIA A100 and AMD MI300X, with benchmarks against Megablocks and a PyTorch reference [2605.23911]. On A100, the end-to-end throughput table for batch sizes at or below 512 tokens includes the following values:

- **Mixtral-8×7B**: 131% relative throughput at 32 tokens, 122% at 128 tokens, and 89% at 512 tokens.
- **Qwen2-MoE-57B**: 103% at 32 tokens, 104% at 128 tokens, and 93% at 512 tokens.
- **Mixtral-8×22B**: 118% at 32 tokens, 118% at 128 tokens, and 87% at 512 tokens.
- **DeepSeek-V3**: 11.53 ms at 32 tokens, 16.74 ms at 128 tokens, and 20.16 ms at 512 tokens, with no Megablocks baseline reported in the table [2605.23911].

The batch-size sensitivity is stated directly: at small batches (32, 128) TritonMoE outperforms Megablocks by 18–31%, whereas at 512 tokens it achieves 89–93% across models [2605.23911]. This pattern places the implementation in the latency-sensitive regime of MoE serving rather than in a maximal-throughput regime associated with larger batches.

Correctness and portability are validated through 162 tests, all of which pass on both NVIDIA A100 and AMD MI300X with zero code changes [2605.23911]. The paper explicitly presents this as evidence of cross-platform portability. Performance tuning on AMD is identified as future work, so functional parity is established more strongly than performance parity.

## 6. Routing imbalance, limitations, and future directions

The paper includes a targeted study of routing imbalance using synthetic expert assignments with Zipfian shapes $\alpha=1.2$ and $\alpha=2.0$ [2605.23911]. At 512 tokens, Mixtral-8×7B and Mixtral-8×22B remain stable across skew, with speedup over Megablocks staying approximately $1.2\times$. The contrasting case is Qwen2-MoE-57B, which has 64 experts and top-4 routing: its speedup drops from $1.03\times$ under uniform assignments to approximately $0.70\times$ at $\alpha=2.0$.

The paper stresses that this degradation is *not* because TritonMoE slows down: its latency remains 3.18 ms in all cases. Rather, Megablocks accelerates under skew because its block-sparse layout merges dominant-expert tokens into large blocks that its CUDA kernels handle more efficiently [2605.23911]. This is the clearest reported limitation of the fixed BLOCK\_M schedule.

The identified underperformance regime is explicit: 64+ experts under extreme skew [2605.23911]. The proposed response is dynamic block-to-expert assignment as future work. Additional future directions listed by the paper are backward-pass kernels for training, persistent-kernel variants to fuse down projection plus unpermutation, and multi-GPU expert-parallel dispatch with all-to-all communication [2605.23911].

A common misconception would be to interpret portability as implying uniform performance across all routing regimes. The reported results do not support that interpretation. Instead, they show that a fully portable MoE inference kernel can match or exceed CUDA-optimized implementations at typical serving batch sizes while still exposing a specific sensitivity to extreme load imbalance. Through Triton’s LLVM backend, the same sources compile to NVIDIA PTX and AMD GCN/CDNA, and the implementation achieves this without inline CUDA or vendor-specific intrinsics [2605.23911].

Source: https://www.emergentmind.com/topics/tritonmoe