For now, from a little testing in Colab, it looks like:
I think there is a real systems idea here.
In a small attention-only T4 microbenchmark, the 25% routing case eventually crossed over and became faster than an optimized dense causal SDPA baseline as the context grew. So I would not read what follows as “the routing idea does not work.”
What I would check before interpreting the current loss numbers as the effect of a learned causal router, though, are two very small invariants:
- Is the router itself actually receiving a learning signal?
- Is the routing decision itself causal, not only the final attention mask?
Both are cheap to test, and in my run both exposed something worth separating from the main sparsity idea.
After that, I think the highest-information experiment is probably just:
Dense
vs
Random K/V subset
vs
Fixed scorer
vs
Corrected learned + causal router
rather than jumping immediately to a large benchmark suite.
What I actually tested
I tested a small version of the actual RoutingGPT classes, then separately microbenchmarked the attention module on a Colab T4.
The source snapshot I tested corresponded to commit:
fde237343ad06fc2ac4be92212276088c7be7b31
So if main has changed since then, I would rerun these tiny checks before assuming they still apply.
For the GPU microbenchmark:
- Tesla T4
- PyTorch
2.11.0+cu128
- batch = 1
- embedding dim = 512
- 8 heads
- synthetic inputs
- pre-created inputs outside the timed region
- warm-up runs
- CUDA Events + synchronization for timings
This is deliberately not an end-to-end language-model training benchmark. It is just meant to separate the attention mechanism from tokenization/data loading/model-head costs.
1. First cheap check: is the router actually learning?
The current routing path appears to be conceptually:
x
│
├───────────────┐
│ │
router │
│ │
scores │
│ │
top-k │
│ │
indices │
│ │
└──── gather(x, indices)
│
selected x
│
K / V
│
attention
│
loss
The important distinction here is that this is not simply “topk is non-differentiable.”
PyTorch can propagate gradients through the values returned by topk. There is a nice small example in the PyTorch forum showing almost exactly this distinction: if the scoring model produces scores, but downstream computation uses only the integer top-k indices to gather the original input, there is no gradient path from the loss back into the scoring model.
In the current RoutingGPT path, the routing scores/values are not used in the resulting attention contribution; only the indices are used to gather x.
The smallest sanity check is therefore literally something like:
loss.backward()
for i, block in enumerate(model.blocks):
print(i, block.attn.router.weight.grad)
In my small run using the RoutingGPT classes:
router.weight.grad -> None
router.bias.grad -> None
q_proj.weight.grad -> nonzero
k_proj.weight.grad -> nonzero
v_proj.weight.grad -> nonzero
and after an AdamW step, the router weights had exactly zero change.
So, for this snapshot, I would describe the current result as:
the hidden representation is learning under a fixed router projection,
rather than:
the router parameters are learning which tokens are useful.
There is one subtle point here: that does not mean the selected token set is static.
I also checked that. While the router weights stayed unchanged, training the rest of the small model changed its hidden states enough that the same probe input later produced different selected indices.
So I would avoid calling this simply “random fixed token selection.” It is closer to:
fixed scoring projection
+
a representation that can move relative to that projection
which is actually a potentially interesting baseline by itself.
There are several ways to make the selector genuinely learnable; I would keep these as design branches
I do not think this implies there is one mandatory fix.
For example:
A. Keep hard top-k, but let the routing score affect the layer update
CoLT5 is a useful comparison here.
It also gives tokens learned routing scores and performs top-k routing, but the routed contribution is scaled using the routing score specifically so that the scoring parameters receive a learning signal.
That keeps discrete selection while leaving a differentiable path from the loss to the scoring function.
B. Use a differentiable sparse selector
SparseK Attention takes a different route: it combines a scoring network with a differentiable sparse top-k operator.
That is another design family rather than something I think this implementation necessarily has to adopt.
C. Give the selector its own objective
Another general option is to separate:
LM objective
from
routing / selection objective
and train the selector through its own signal.
The larger point is simply that “sparse top-k routing” and “how the router learns” are separable design choices.
For this experiment, even a very small correction is probably enough before testing more elaborate routing machinery.
2. Second cheap check: make the selection causal too
The current causal mask itself looks reasonable: after selecting K/V tokens, a selected key is usable only if its original position is not later than the query position.
The thing I would separate is what happens before that mask.
At the moment the sequence seems to be roughly:
all tokens in the sequence
↓
router scores
↓
global top-k
↓
selected K/V set
↓
causal mask
Suppose a future token gets a very high router score.
It can occupy one of the global top-k slots and therefore displace a past token from the selected K/V set.
The future token itself will later be masked for an earlier query, which is good — but the past token it displaced is already gone.
So this is not quite:
an earlier query directly attends to a future token.
It is instead:
future tokens can affect which past tokens survive routing before the causal mask is applied.
A very cheap regression test for that is prefix invariance.
Keep:
- the prefix identical,
- the total sequence length identical,
- therefore also the routing budget identical,
and change only the suffix.
Conceptually:
prefix = fixed_prefix
x1 = concat(prefix, suffix_A)
x2 = concat(prefix, suffix_B)
logits1 = model(x1)
logits2 = model(x2)
compare(
logits1[:, :prefix_len],
logits2[:, :prefix_len],
)
For a strictly autoregressive model, changing tokens that have not happened yet should not change the earlier positions.
In my small RoutingGPT-class test, changing only the suffix changed both:
- the first-layer selected token set, and
- some prefix logits.
I would treat the exact frequency/magnitude from a random tiny model as irrelevant; the useful part is the property test itself.
Possible causal-routing branches
Again, I do not think there is only one correct redesign.
One possibility is simply to make the routing candidate set causal:
for query position q:
candidates = positions <= q
top-k only inside candidates
A blockwise version can make that much more practical.
MoBA is useful as a reference point here: it uses sparse block routing in an autoregressive attention setting, and causality is part of the routing/attention structure rather than something added only after arbitrary future blocks have competed for the same sparse budget.
Another very reasonable branch would be:
always-on local causal attention
+
routed long-range attention
There is no requirement that the router carry the entire causal-attention responsibility by itself.
That might also make the early-token behavior nicer, since every query can retain at least a local causal neighborhood while routing is used only for expensive long-range access.
I would probably keep the current global scalar scorer initially rather than immediately adding query-dependent/head-dependent routing. A shared notion of token utility is a perfectly reasonable simple baseline. Query-dependent routing can be a later axis if the simpler version first shows a measurable learned-selection benefit.
3. After those two checks, the control experiment can stay very small
Before a large long-context evaluation, I think this four-way comparison would tell a lot:
| Variant |
What it helps isolate |
| Optimized dense causal attention |
quality/speed reference |
| Random K/V subset |
effect of sparsity alone |
| Fixed random scorer |
effect of representation adapting around a fixed scoring projection |
| Learned + causal router |
added value from learned selection |
Keep the important training conditions matched:
- same dataset
- same context length
- same optimizer
- same token budget
- same model size
- preferably multiple seeds, or at least a fixed reported seed
and evaluate on held-out data rather than comparing only the final training step.
This gives useful outcomes in either direction.
If learned > random/fixed
That is much cleaner evidence that the router learned a useful selection rule.
If learned ~= random/fixed
That is also informative.
It could mean that most of the gain comes from sparsity itself, or that the main Transformer learns representations that adapt around whichever selector it is given.
A cheap follow-up diagnostic in that case would be to freeze the backbone and train only the router, just to separate:
router learning
from
backbone adaptation to the routing pattern
If all sparse variants lose some quality but gain useful speed
That is still a perfectly legitimate result.
Then the result is a quality / latency Pareto trade-off, rather than a claim that sparse routing dominates dense attention everywhere.
4. The T4 speed result was actually encouraging — but the crossover matters
For FP16 forward attention, I got approximately:
| Context |
optimized dense |
routed 50% |
routed 25% |
| 256 |
0.241 ms |
0.525 ms |
0.509 ms |
| 512 |
0.323 ms |
0.498 ms |
0.493 ms |
| 1024 |
0.505 ms |
0.888 ms |
0.698 ms |
| 2048 |
1.268 ms |
1.766 ms |
1.110 ms |
| 4096 |
3.532 ms |
5.356 ms |
3.026 ms |
So on this particular T4/runtime:
- at short context, routing overhead dominates;
- 25% routing starts winning around the 2K–4K region;
- 50% routing did not beat the optimized dense baseline through 4K.
That seems like a useful result for the basic idea.
It suggests that the routing ratio is not only a quality knob. It is also a systems crossover knob:
attention saved
vs
router + top-k + gather + mask overhead
At high enough sparsity / context length, the saved attention work begins to dominate the routing overhead.
One important detail: ratio=1 is a semantic dense control, but not a performance dense control
I also compared the current routed attention with top_k_ratio=1.0 against a plain dense causal SDPA module using matched Q/K/V/output projection weights.
The outputs were numerically essentially identical:
max absolute difference ≈ 1.2e-7
So ratio=1 looks like a good semantic dense-equivalence check.
But it is not a good performance dense baseline, because the routed version still does:
router
top-k
gather
position gathering
explicit mask construction
SDPA
At context 4096 in the same FP16 T4 test I got roughly:
optimized dense causal SDPA : 3.53 ms
routed ratio=1 : 9.95 ms
So if the Dense number in a benchmark is implemented simply as top_k_ratio=1, I would keep that baseline for correctness but add a separate optimized dense is_causal=True implementation for performance.
5. Fewer attention interactions did not automatically mean lower peak memory
One result that surprised me a little was memory.
At T=4096 / FP16, the incremental peak allocated CUDA memory in the attention-module probe was roughly:
optimized dense : 20.0 MiB
routed 50% : 38.1 MiB
routed 25% : 23.1 MiB
So the 25% version was faster at that length, but it did not have lower measured peak allocation than optimized dense in this implementation.
I would not infer a single root cause from that measurement, but there are several extra objects on the routed path:
router scores
top-k indices
selected_x
gathered K/V
explicit boolean causal mask
and modern dense SDPA is already heavily optimized.
PyTorch’s scaled_dot_product_attention can select among optimized CUDA implementations depending on the inputs, dtype, shape, mask, and hardware.
So I think these are best treated as separate metrics:
Q×K interaction count
FLOPs
wall-clock latency
peak VRAM
rather than assuming one determines all the others.
Backend detail from this one T4 runtime
I also forced individual SDPA backends as a sanity check.
For the tested dense and routed shapes on that runtime:
Memory-Efficient / Efficient Attention : available
Math : available
Flash Attention : unavailable when forced
cuDNN Attention : unavailable when forced
That only tells me backend eligibility under those test inputs; I would not claim that the default dispatcher necessarily chose a particular backend without profiling it.
The main practical point is just that kernel choice is another reason the real crossover can look different from the simple theoretical T² interaction count.
6. I would slightly change the timing setup for future benchmarks
One small measurement detail:
if the README timing comes from the current training loop, the timed region appears to include next(dataset), and dataset refill can include tokenization/batch preparation.
That means the number can mix:
data/tokenization
+
CPU work
+
device transfer
+
forward/backward
+
optimizer
Also, CUDA execution is asynchronous.
PyTorch’s CUDA semantics documentation explicitly notes that GPU timings without synchronization are inaccurate and recommends either torch.cuda.synchronize() or CUDA Events.
So for the model-compute benchmark I would use something like:
pre-tokenize / pre-create batch
↓
warm up
↓
CUDA Event start
↓
forward (+ backward if desired)
↓
CUDA Event end
↓
synchronize
↓
repeat and report median
Then separately measure end-to-end throughput if data-pipeline performance matters.
A repository I found useful for seeing this kind of decomposition in practice is TokenButler: its benchmark code separates QKV projection, attention, predictor/scorer work, top-k selection, K/V gathering, etc., and also includes random/contiguous oracle baselines.
That is probably more useful as a benchmark-design reference than as a claim that RoutingGPT should use the same architecture.
A compact decision tree for the next experiment
I think the next steps can stay pretty small:
1. One backward pass
|
+-- router grad is None
| -> connect the selector to a learning signal
|
+-- router grad is nonzero
-> inspect the actual gradient path
2. Same prefix, different suffix
|
+-- prefix output changes
| -> make routing eligibility causal
|
+-- prefix invariant
-> move on to evaluation
3. Dense / Random / Fixed / Learned comparison
|
+-- Learned clearly wins
| -> evidence for useful learned routing
|
+-- Learned ~= Random/Fixed
| -> inspect backbone adaptation / sparsity effect
|
+-- Sparse trades quality for speed
-> characterize the Pareto frontier
4. Only then scale context aggressively
|
+-- synchronized latency
+-- peak VRAM
+-- tokens/sec
+-- held-out quality
5. Once that is stable:
-> long-context task evaluation
That seems like a fairly high information-gain path without requiring a large reimplementation.
A few later-stage long-context caveats
These are not blockers for the current experiment, but they may matter once the context sweep gets much larger.
Fixed ratio is still quadratic
If the routing ratio is a fixed r, the attention interaction count is roughly:
T × (rT)
so it is still quadratic in sequence length, just with a smaller constant.
That can still be very useful — the T4 crossover suggests it can be — but if the eventual goal is a different asymptotic scaling regime, a fixed K, block budget, or other sublinear routing budget would be another design axis.
SparseK, for example, explicitly explores selecting a constant number of KV pairs per query.
Learned positional embeddings also set a context boundary
The current model uses a learned positional embedding table.
So a checkpoint trained with a maximum context of 512 cannot simply be evaluated at 4K/16K without also changing/training the positional representation.
For context-length comparisons, I would therefore make the positional setup part of the controlled experiment rather than changing only the routing ratio.
Match training tokens, not only training steps
With fixed batch size, a training step at 16K context processes far more tokens than a training step at 512.
So if comparing models trained at different sequence lengths, I would report whether the experiment matches:
steps
tokens
or total compute
because those are not equivalent.
Attention may stop being the dominant memory bottleneck
At very long training sequences, materializing full vocabulary logits [B, T, V] can itself become huge.
So if routing eventually succeeds at reducing the attention bottleneck, the LM head / cross-entropy path may become the next memory wall.
That is independent of whether the routing idea works.
KV caching is connected to the causality invariant
A standard append-only KV cache relies on past representations not changing when future tokens are appended.
So fixing prefix invariance is also useful preparation if KV caching is added later.
Evaluation after the basic mechanism is stable
I would probably not start with a giant long-context benchmark.
A cheaper progression would be:
held-out LM loss / perplexity
↓
small associative-recall / retrieval diagnostic
↓
position + context-length sweep
↓
larger long-context suite
Once the mechanism itself is behaving as intended, RULER becomes useful because it goes beyond a single needle-in-a-haystack test and includes retrieval, multi-hop/tracing, aggregation and QA-style synthetic tasks.
Its own documentation also makes a useful caution: good vanilla NIAH performance does not necessarily imply robust long-context behavior, and RULER itself is still a diagnostic suite rather than a replacement for realistic downstream tasks.
So, overall, my current read would be:
the sparse-K/V idea looks worth continuing.
The small T4 test actually made me more interested in the efficiency side, because the 25% case showed a real crossover against optimized dense attention once the context got long enough.
I would just separate two things before drawing conclusions from the quality numbers:
"does sparse K/V selection save useful compute?"
from
"did a causal router actually learn which K/V tokens to select?"
The first one already seems to have some promising evidence.
The second one should be much easier to interpret after the two tiny checks above — router gradient and prefix invariance — followed by the Dense / Random / Fixed / Learned control.
That seems like a much cheaper next step than rebuilding the architecture, and whichever way that comparison turns out should give a useful result.