Our RG-LRU GPU audit

Legacy differentiable training path

The pre-existing inference/prefill file benchmark/triton_rglru.py already fused gate activation, transition reconstruction, write and recurrence after the two gate projections. The missing optimization was the complete differentiable training path in RGLRUMixer._terms.

That old training path constructed logical full-shape tensors for two gate activations, log_a, transition a, square-root multiplier and write, and used torch.cat to install the first-token reset. It then called a custom serial or chunk scan. The chunk scan additionally allocated FP32 transition and write summaries plus chunk input states. torch.compile could potentially fuse some pointwise producers, but the custom scan boundary consumed explicit a and write tensors, so their logical interface remained.

Strengthened fused path

benchmark/triton_fused_rglru_training.py accepts recurrent input, two gate logits, a_param, optional segment_pos and optional initial state. It contains:

a fully fused state-stationary serial forward/backward;
exact two-level chunk forward/backward for K=8/16/32/64;
arbitrary resets, including multiple resets within a chunk;
nonzero initial state and its gradient;
partial final chunks;
clipped square-root backward;
BF16 transition/write boundaries and FP32 recurrent accumulation;
backward recomputation of pointwise gate/transition/write terms;
partial a_param gradients followed by a reduction.

The fused paths retain the projection outputs because the two official block-diagonal projections run before the custom operator. They do not create logical global [B,L,D] activation, transition, multiplier or write allocations. Outputs and saved logits/states remain state-sized, and chunk paths allocate [B,ceil(L/K),D] FP32 summaries. “No logical global transition” is not a claim of zero HBM traffic or guaranteed register residency.

The final strengthened baseline also has a linear-work real-affine prefix: group-local serial composition, an outer group prefix, and local correction. It preserves the noncommutative affine composition order and is used in both forward and reverse mode. The selected H800 backend is grouped-prefix32 with K=16 and a shape-calibrated mode block. SAMU therefore is not compared with the old serial or Hillis–Steele RG-LRU path where this grouped path is faster.

Exact real chunk algebra

Each token is an affine map h -> a_t h + b_t. For a chunk, the exact summary is h_out = P h_in + Q, updated left to right by

P <- a_t P
Q <- a_t Q + b_t.

A reset is encoded by a_t=0 and the canonical reset write, so it naturally erases the prefix before it without inexact post-processing. The last partial chunk applies identity/no-write for invalid lanes. Reverse mode uses the same affine closure on the adjoint recurrence and replays each chunk to recompute pointwise derivatives.

Attribution

Fusing RG-LRU pointwise preparation into a serial Triton scan, keeping a running carry and recomputing activations in backward are Fattori prior art. The reset-aware exact chunk forward/backward, partial chunks, initial-state gradient, grouped affine hierarchy, and H800 dispatch in this repository are our implementation work. Fattori’s repository is described as an open-source Hawk/RG-LRU implementation, not as an official implementation.

Validation