The core test: at the same compute and parameters, does choosing blocks from the recurrent state beat choosing them another way (or not routing at all)?
Multi-Query Associative Recall: store key→value pairs, recall them later. The four arms share the same blocks, same k/B (50% active), same seed; only the router rule differs. stateful reads ht; stateless reads the token (MoE-Mamba analog); random-k routes randomly; dense runs all blocks.
| accuracy | router arm | k / B | active | params | read |
|---|
Accuracy and loss at every epoch. MQAR recall is acquired in a sudden phase transition, not gradually — so the curve, not the final number, tells you whether (and how fast) a config solves the task. A run whose loss keeps dropping while accuracy stays flat hasn’t transitioned yet; a late or absent transition flags a harder task or a degraded setup.
Each line is one run; colour identifies the run. Filter by router / difficulty / seed / variant to focus the view, and choose how seeds combine (mean ± band / median ± band / per seed). Shown by default: kv8, the recommended recipe mag25a95e80 vs dense (the baseline, always drawn), as mean ± band; widen to other difficulties, variants, and seeds with the dropdowns. (Default config = strongest found so far in the annealing study, not claimed optimal.)
Routing soft→blocky anneal study (2026-06-03) — strongest stateful-routing training config
found so far on MQAR kv8/v256 (not claimed optimal): linear “magnet” anneal,
anneal_frac=0.25, soft-floor alpha_max=0.95,
route_select_ema=0.8,
fp32, fused kernel — the curve labelled
mag25a95e80. Across 7 seeds it lifts mean final accuracy
0.825→0.942 and beats dense on 5/5 seeds with a baseline, fixing both the
“never-transitions” and “late-collapse” failures of the no-EMA recipe
(mag25a95). Window / long-anneal / router-freeze / router-lr variants were tried and
rejected. Full method: docs/routed_ssm_annealing_algorithm.md; evidence:
docs/annealing_research_log.md.
Here we let capacity vary: as more blocks are switched on, recall should have more room to climb.
The full MQAR sweep across block type (ssm / mamba / elm), router, and k/B. Unlike the matched A1 panel, these vary capacity. Reads the recall capacity axis: more active blocks → more recall headroom.
| best acc | router | block | k / B | router net | d / N | last | active | epochs |
|---|
Everything that makes a training epoch fast: the GLA chunked-matmul kernel we ship for routed training (B1a), the exact parallel-scan it falls back to on long/narrow shapes (B1b), and the three-kernel bake-off + ncu profiles that say which wins where (B1c). Every × is fwd+bwd vs torch.compile on identical math (graph-vs-graph), parity-verified first.
Our fastest training kernel, timed against torch.compile running the exact same math.
The GLA chunked-matmul kernel with a fused Triton backward. Compared head-to-head against the torch.compile implementation of the same math (graph-vs-graph). All three sub-panels are parity-verified against the sequential reference before any speed is reported.
Lead column = × faster than torch.compile. Both fwd+bwd, median, same shape (Bsz8 H8 P16 N16 chunk64).
| × vs compile | L | mem saved | fused ms | compile ms | fused MB | compile MB |
|---|
End-to-end epoch wall-time on the actual model. Lead = × faster epoch vs torch.compile. (Speed bench — accuracy still early.)
| × vs compile | L | mem saved | triton ms | compile ms |
|---|
When the router is live, the GLA path vs the scan path. Lead = GLA × vs scan; both already beat eager.
| × GLA vs scan | L | k | mem saved | gla ms | scan ms |
|---|
Full model training step (embed + 4 layers + tied 50257-vocab readout): dense vs the routed gla2 dispatch kernel at different k. Lead = × vs dense (this page's usual convention, >1 = faster), derived from the raw paired-interleaved time ratio (raw ratio column). Error bars = interquartile range over 40 paired iters. E2E is vocab-head-diluted — the shared vocab readout costs the same on every arm, so the routed-region-only speedup is larger than this full-step number.
| × vs dense | arm | raw ratio | step ms | source |
|---|
Replacing the step-by-step time loop with one parallel-scan kernel — and the win grows as sequences get longer.
The block recurrence is a Python for t in range(T) loop — T sequential launches torch.compile
cannot flatten. A linear recurrence is an associative scan, so one Triton kernel does it
with depth T→log T. Lead = × vs torch.compile (capped at T≤256 — compile loop-unroll explodes above that, shown as n/a).
| × vs compile | config | triton ms | compile ms | (vs eager) |
|---|
A reference timing of all three kernels as sequences get longer — lower is faster.
Forward+backward total time for each kernel as L grows. The dense line here is torch.compile’d causal self-attention — an O(L²) tensor-core floor, NOT a semantic alternative (it just shows how close each linear kernel gets to raw tensor-core throughput). For the meaningful vs-torch.compile kernel win, see B1. Lower ms = faster.
| total ms | kernel | L | fwd ms | peak MB | source |
|---|
K1 scan is memory-bound (~90% DRAM): the only win is moving fewer bytes → chunked SRAM (K2/K3). K3 SSD fwd is latency-bound (small grid, register pressure).
| kernel | shape | SM % | Mem % | occ achieved | bound |
|---|
Autotuning launch config does not beat the original scan (already ~90% bandwidth-bound). The scan’s win at large T is structural (depth log T), not a tuning trick. Lead = tuned × vs torch.compile.
| × vs compile | shape | T | tuned/orig | tuned ms | orig ms | compile ms |
|---|
Decode one token at a time (B2a) and prefill a whole prompt (B2b), routed vs dense; then the hand-tuned top-1 CUDA kernel and the compressed router that lets it scale to many blocks (B2c). Every × is vs torch.compile dense.
At single-token generation, running only the active blocks beats running all of them — once the model is wide enough.
Single-step inference. Routing computes only the k active blocks. The dense diagonal-SSM baseline is memory-bound (not a tensor-core GEMM), so block-sparsity genuinely wins above a width threshold. Lead = × vs torch.compile dense; ideal B/k is the FLOP ceiling.
| × vs dense | config | best k | ideal B/k | routed ms | dense ms |
|---|
Processing a whole prompt at once: scanning only the active blocks vs scanning all of them.
| × vs dense | config | best k | ideal K/k | per-k |
|---|
Computing only the single active block per token, in one fused kernel — and making the router cheap so it scales to many blocks.
All × vs torch.compile dense, parity-exact, fair graph-vs-graph. Full write-up: docs/kernel_acceleration_research_log.md. Curated relative speedups (a graph-timing artifact inflated some raw absolute-ms in logs; the × figures shown are the vetted same-harness values).
After coalescing W_router + a warp-shuffle router reduction (+1.4–1.6× over the naive version).
| × vs dense | batch | hand µs | dense µs |
|---|---|---|---|
| 3.56× | 32 | 0.58 | 2.04 |
| 3.95× | 128 | 0.62 | 2.43 |
| 1.89× | 512 | 1.92 | 3.63 |
| 2.69× | 2048 | 6.69 | 17.99 |
P=D/K. The full router's all-state read grows ∝K² (faster than dense block work ∝K), so adding blocks makes top-1 worse — the reason a compressed router is needed.
| × vs dense | config | D | K | router/block |
|---|---|---|---|---|
| 1.45× | real | 128 | 8 | 7.3 |
| 2.24× | D×4 | 512 | 8 | 3.8 |
| 1.98× | D×8 | 1024 | 8 | 2.0 |
| 0.80× | K=16 | 512 | 16 | 15.3 |
| 0.37× | K=32 | 1024 | 32 | 31.5 |
Router reads a per-block scalar summary, not the full state ⇒ input 4224→136 (31×). Global-state kernel runs K=8..128. Speed result (random weights); recall needs a retrain — gated router_pool code is ready.
| × vs full-router | K | router_in full→comp | runs? |
|---|---|---|---|
| 2.02× | 8 | 4224→136 | yes (clean) |
| ✓ | 16 | 4352→272 | yes (full-router collapses) |
| ✓ | 32 | 16896→544 | yes (global-state) |
| ✓ | 64 | 33792→1088 | yes |
| ✓ | 128 | 67584→2176 | yes |
topk1_kernel_verdict, topk_descent_win.Which router can exploit which kernel form. This sweep is still running: only the dense rows and a few compiled rows are populated so far (the eager baselines have not landed, so ×-vs-eager is not yet computable). Shown for transparency; treat numbers as preliminary. Lead = raw step ms until the baseline run completes.
| step ms | router | kernel | L | tok/s | peak MB | note |
|---|
Honest negative result: at this small scale, the recall advantage does not yet show up as better language modeling.
Small selective-SSM LM, GPT2 tokenizer, tied embeddings. Experts = SSM blocks. state-routed (ours) vs token-routed (BlackMamba analog) vs dense. Lower perplexity is better. Pure-PyTorch backend (no mamba_ssm kernel on Windows). anneal = soft→hard route schedule.
| test ppl | arm | anneal | val ppl | active | tok/s | VRAM MB | params |
|---|
Toy diagnostic — confirms the router can switch the active sub-network mid-sequence; not a headline result.
A transient cue sets the active context; later steps must apply it. Solving it needs holding context in the recurrent state and routing to the matching block — a memory-switch sanity check (a different axis from the A1 recall result above). The control that matters: an MLP router does not rescue stateless — the win is the state, not router capacity. chance = 0.125.
| accuracy | router | block | k / B | router net | active | params |
|---|