moe-triton

A Triton top-k MoE router layer for transformers: the softmax, the top-k and the renormalisation in one kernel.

The three library calls it replaces read only a few tens of KB, so at decode row counts they are launch-bound rather than bandwidth-bound: on Qwen3-30B-A3B at tp4 they cost 0.77 ms of a 15.5 ms decode step.

Layers

Correctness

Against softmax + topk + renormalisation over 33 shapes (1 to 4096 rows, 8/64/128 experts, top-k 2/4/8, renormalisation on and off): max delta on the selected scores 3.6e-07, selected indices identical everywhere. On a full Qwen3-30B-A3B forward the logits are bit-identical to the reference path.

The kernel falls back to the reference implementation for any input it is not written for (not on CUDA, not 2-D, or not row-contiguous), so the layer is always safe to install.