Pith. sign in

REVIEW 2 major objections 5 minor 4 references

Cross-Platform Fused MoE Dispatch in Triton: Portable Expert Routing Without CUDA

T0 review · 2 major / 5 minor · reviewed 2026-07-13 · grok-4.5

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

Pith's one-line read A fully portable Triton MoE dispatch kernel matches CUDA Megablocks at inference batch sizes and runs unchanged on AMD.

desk verdict Clean portable Triton MoE dispatch that is competitive with Megablocks at serving batch sizes and actually runs unchanged on MI300X; the fixed-tile schedule is the known soft spot the authors already measure. read the letter →

arxiv 2605.23911 v1 pith:P6UKMCFN submitted 2026-04-07 cs.DC cs.ETcs.PF

Subhadip Mitra This is my paper · ORCID
classification cs.DCcs.ETcs.PF
keywords Mixture-of-ExpertsTritonMoEdispatchcross-platformGPUfusedGEMMSwiGLUinferencekernelsblock-sparsescheduling
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

Mixture-of-Experts layers are the backbone of most frontier language models, but their inference path is messy: tokens must be scored, reordered onto experts, run through variable-size GEMMs, then scattered back with weights. Almost every high-performance implementation of that path is written in CUDA and therefore locked to NVIDIA GPUs. This paper shows that the same complete forward path can be written entirely in OpenAI Triton, using only portable primitives. The resulting TritonMoE kernel reaches 89–131 percent of the throughput of the leading CUDA baseline (Megablocks) at the batch sizes typical of interactive serving, while every correctness test passes on both NVIDIA A100 and AMD MI300X with zero source changes. The decisive engineering move is a fused gate-plus-up projection that reuses the same L2-resident input tiles and finishes the SiLU activation in registers, cutting roughly 35 percent of global memory traffic. The work therefore argues that vendor-specific CUDA is no longer a prerequisite for competitive MoE inference kernels.

What carries the argument

The fused gate+up GEMM: both SwiGLU projections are computed from a single shared load of each input tile, with the SiLU activation performed in registers so that the two intermediate buffers never touch global memory; the rest of the pipeline is a block-scheduled grouped GEMM that maps fixed-size program blocks onto (expert, token-offset) pairs.

What would settle it

Measure end-to-end latency of TritonMoE versus Megablocks on a 64-expert top-4 configuration under Zipfian α=2.0 routing at 128–512 tokens; if the Triton kernel’s relative speedup stays near or above 1.0 instead of collapsing toward 0.7 imes, the fixed-tile assumption holds.

Watch

Extended reading notes

Core claim

TritonMoE, a five-kernel MoE dispatch pipeline written solely in portable Triton, delivers 89–131 percent of CUDA Megablocks throughput at inference batch sizes (≤512 tokens) across Mixtral, DeepSeek-V3 and Qwen2-MoE configurations, and the identical source passes all 162 correctness tests on both NVIDIA A100 and AMD MI300X.

Load-bearing premise

That a compile-time fixed tile size and a static expert-to-block schedule remain competitive when real production traffic is heavily skewed toward a few experts.

Editorial extensions

If this is right

  • Interactive LLM serving stacks can ship a single MoE kernel binary that runs on both NVIDIA and AMD accelerators without platform forks.
  • At the small-to-medium batch sizes that dominate chat serving, portable Triton kernels can already match or beat hand-tuned CUDA dispatch.
  • The 35 percent memory-traffic reduction from gate+up fusion is available to any SwiGLU MoE that can share input tiles.
  • Beyond roughly 64 experts under heavy skew, fixed tile schedules become a liability, so future kernels must adopt dynamic block-to-expert assignment.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • Once the same source is proven correct on both vendors, the remaining performance gap on AMD becomes pure optimization work rather than a portability research problem.
  • The same fused-tile pattern should transfer directly to other dual-projection activations (e.g., GEGLU) with only the element-wise op changed.
  • If dynamic block assignment closes the skew gap, portable Triton could become the default MoE backend even for training-scale expert counts.
Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

2 major / 5 minor

Summary. The paper presents TritonMoE, a complete MoE forward-pass dispatch pipeline (router scoring, token permutation, expert GEMMs, weighted unpermutation) implemented solely in OpenAI Triton. The main technical contributions are a block-scheduled grouped GEMM that maps program blocks to (expert, token-offset) pairs (Algorithm 1) and a fused gate+up SwiGLU kernel that shares L2-cached input tiles and performs SiLU in registers, with a byte-count argument claiming ~35% less global memory traffic. On A100 the system reports 89–131% of Megablocks throughput at inference batch sizes (≤512 tokens) across Mixtral-8x7B/22B, Qwen2-MoE and (unfused/fused) DeepSeek-V3 configurations; all 162 correctness tests pass unchanged on AMD MI300X. A controlled Zipfian-skew study (§4.7) quantifies the regime (64+ experts, extreme skew) where fixed BLOCK_M scheduling loses to Megablocks’ block-sparse layout.

Significance. If the reported numbers hold, the work shows that a pure-Triton MoE dispatch stack can match or beat a mature CUDA baseline at the batch sizes that matter for interactive serving, while remaining portable to AMD with zero source changes. That combination is practically useful: it lowers the barrier to multi-vendor MoE inference and supplies a self-contained, open implementation (code link provided) rather than a framework-internal kernel. The fusion ablation, roofline breakdowns and explicit skew characterization are concrete engineering contributions that other portable-kernel efforts can build on. The fixed-tile limitation under extreme imbalance is already scoped by the authors and does not erase the small-batch / near-uniform result.

major comments (2)
  1. §4.2 / Tables 2–3: the headline claim of 89–131% of Megablocks is supported for Mixtral and Qwen2, but DeepSeek-V3 (256 experts) has no Megablocks baseline at all—only an unfused-vs-fused Triton comparison. The abstract and introduction still list DeepSeek-V3 among the configurations for which the relative-throughput claim is made. Either obtain a runnable Megablocks (or equivalent CUDA) number for that configuration, or rephrase the claim so that the 89–131% range applies only to the configurations that actually have a Megablocks comparator.
  2. §4.6: cross-platform validation is correctness-only (162 tests on MI300X). The central portability claim is therefore only half-demonstrated; without even a single latency or bandwidth number on AMD it is impossible to know whether the same fused schedule remains competitive once the memory hierarchy and tensor-core mapping change. A minimal performance table (or an explicit statement that AMD performance is out of scope and why) is needed before the portability result can be treated as load-bearing.
minor comments (5)
  1. §3.3 memory-traffic analysis: the unfused formula counts 8TF + 4Td bytes; a short parenthetical confirming the exact read/write sequence (including the final intermediate write) would make the 35% figure easier to reproduce.
  2. Table 5: d_ffn is scaled down as expert count rises “to approximate constant total compute.” State the exact FLOP target used so readers can verify the scaling.
  3. Figure 1 is referenced as a two-panel roofline but the manuscript text does not include axis units or peak lines; ensure the camera-ready figure is self-contained.
  4. Related Work: vLLM’s Triton fused MoE is mentioned; a one-sentence note on whether its kernel is available outside the serving stack (or how TritonMoE differs in schedule/fusion) would help readers place the contribution.
  5. Limitations (5): the host-side block-schedule construction is listed; quantify the CPU time or the resulting host–device sync cost relative to kernel time so the overhead is not left qualitative.

Circularity Check

0 steps flagged · score 0.0 of 10

No circularity: measured throughput and byte-count savings are independent of any fitted parameter or self-justifying definition.

full rationale

TritonMoE is an engineering systems paper. Its central claims are (1) end-to-end latency numbers against the external Megablocks and PyTorch baselines on public model shapes (Tables 2–3), (2) a 35% memory-traffic reduction obtained by direct byte counting of fused vs. unfused SwiGLU traffic (Section 3.3), and (3) zero-change correctness of 162 tests on AMD MI300X. None of these quantities is obtained by fitting a free parameter and then re-predicting a related quantity; the memory-traffic formulas are pure arithmetic (8TF+4Td vs. 2TF+2Td). Algorithm 1 (block schedule) and the fixed BLOCK_M constraint are design choices whose limitations under Zipfian skew are measured and scoped in §4.7 rather than defined away. There are no self-citations that carry the load of a uniqueness theorem, no ansatz imported from prior author work, and no renaming of a known empirical pattern. The derivation chain is therefore self-contained against external benchmarks; circularity score is 0.

Assumptions & free parameters 2 free parameters · 3 assumptions · 0 invented entities

The work is an engineering systems paper. It inherits standard GPU programming and MoE definitions; the only free choices are engineering knobs (tile sizes, number of kernel launches) that are either fixed by the schedule or autotuned. No new physical entities or fitted scientific constants are introduced.

free parameters (2)
  • BLOCK_M
    Fixed at compile time because it must match the precomputed block schedule; only BLOCK_N and BLOCK_K are autotuned. Choice directly affects load-imbalance sensitivity shown in §4.7.
  • number of Triton kernel launches (5)
    Design decision that trades fusion depth against launch overhead; claimed to be lower than Megablocks’ multi-stage dispatch at small batches.
assumptions (3)
  • domain assumption Triton compiles the same source to both NVIDIA PTX and AMD GCN/CDNA via its LLVM backend, preserving numerical semantics for the primitives used.
    Stated in §2.3 and validated by the 162-test suite on MI300X; required for the zero-code-change portability claim.
  • domain assumption Standard MoE forward equations (softmax/sigmoid router, top-k, SwiGLU FFN) as used by Mixtral, DeepSeek-V3 and Qwen2-MoE.
    Equations (1)–(4) in §2.1; taken as given from the model literature.
  • standard math A100 peak bandwidth (2039 GB/s) and FP16 peak (312 TFLOPS) for roofline efficiency calculations.
    Hardware constants used in §4.5; not fitted.

how reviews work

0 comments
Cite this review

Pith. "Pith review of Cross-Platform Fused MoE Dispatch in Triton: Portable Expert Routing Without CUDA." pith.science (2026). https://pith.science/paper/P6UKMCFN

@misc{pith2026260523911,
  author       = {Pith},
  title        = {Pith review of: Cross-Platform Fused MoE Dispatch in Triton: Portable Expert Routing Without CUDA},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/P6UKMCFN}},
  note         = {Machine review of arXiv:2605.23911}
}
read the original abstract

Mixture-of-Experts (MoE) architectures power the majority of frontier large language models, but their inference is bottlenecked by irregular memory access patterns and expert routing overhead. Existing optimized MoE kernels (Megablocks, Tutel, FasterMoE) are implemented in CUDA and locked to NVIDIA hardware. We present TritonMoE, a fused MoE dispatch kernel written entirely in OpenAI Triton that performs the complete forward pass -- router scoring, token permutation, expert GEMMs, and weighted output combination -- using only portable Triton primitives. Our key optimization is a fused gate+up GEMM kernel that computes both SwiGLU projections from shared L2-cached input tiles with in-register SiLU activation, eliminating 35% of global memory traffic. On an NVIDIA A100, TritonMoE achieves 89-131% of the throughput of the CUDA-optimized Megablocks at inference batch sizes (<= 512 tokens) across Mixtral-8x7B, DeepSeek-V3, and Qwen2-MoE configurations. All 162 correctness tests pass on both NVIDIA A100 and AMD MI300X with zero code changes, validating cross-platform portability. We additionally characterize sensitivity to routing imbalance under Zipfian-skewed expert assignments and identify the regime -- 64+ experts under extreme skew -- where our fixed-tile scheduling underperforms Megablocks' block-sparse layout, motivating dynamic block-to-expert assignment as future work. Code is available at https://github.com/bassrehab/triton-kernels.

Figures

Figures reproduced from arXiv: 2605.23911 by the authors.

Figure 1
Figure 1. Per-stage roofline analysis on A100-SXM4-80GB. Expert FFN stages (compute-bound for [PITH_FULL_IMAGE:figures/full_fig_p007_1.png] view at source ↗
Figure 2
Figure 2. MoE dispatch latency at 512 tokens under three routing distributions. The Mixtral config [PITH_FULL_IMAGE:figures/full_fig_p009_2.png] view at source ↗
Figure 3
Figure 3. Speedup of our Triton fused kernel relative to Megablocks dMoE across the full skew [PITH_FULL_IMAGE:figures/full_fig_p012_3.png] view at source ↗

Discussion (0). Continue with ORCID to comment.

Reference graph

Works this paper leans on

4 extracted references · 4 linked inside Pith

  1. [1]

    Deepseek-v3 technical report.arXiv preprint arXiv:2412.19437,

    DeepSeek-AI. Deepseek-v3 technical report.arXiv preprint arXiv:2412.19437,

    arXiv

  2. [2]

    Mixtral of experts.arXiv preprint arXiv:2401.04088,

    Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts.arXiv preprint arXiv:2401.04088,

    arXiv

  3. [3]

    Scattermoe: Efficient mixture-of-experts with scatter and gather operations.arXiv preprint arXiv:2403.08245,

    Shawn Tan, Yikang Shen, Zhenfang Chen, Aaron Courville, and Chuang Gan. Scattermoe: Efficient mixture-of-experts with scatter and gather operations.arXiv preprint arXiv:2403.08245,

    arXiv

  4. [4]

    Qwen2 technical report.arXiv preprint arXiv:2407.10671,

    An Yang, Baosong Yang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Zhou, Chengpeng Li, Chengyuan Li, Dayiheng Liu, Fei Huang, et al. Qwen2 technical report.arXiv preprint arXiv:2407.10671,

    arXiv

Pith tools

Reviewed July 13, 2026 · model on record in the stance chip above.