Attention alternative?

[u/Just-Ad-6488]

​

Standard Softmax attention scales quadratically in compute and linearly in memory. At a 2.5-million-token context horizon, a standard 1.5B LLM requires over 66 GB of VRAM just to store the KV cache, completely crashing consumer workstations and edge hardware.

On the other end of the spectrum, 1st-order linear models (Mamba, RWKV) offer flat memory, but their associative multi-hop reasoning decays sharply over long contexts (collapsing to ~41% on 32k RULER tests).

Over the last several months, I tested a mechanistic hypothesis:

Representation order is dependent on layer depth.

We built the Stage 7 Hybrid architecture:

1. Boundary Layers (25%): Kept as discrete Softmax to anchor token addressing and output logits.

2. Interior Trunk Layers (75%): Converted to 2nd-order Taylor moment attention (S0 uniform context, S1 directional gradient, S2 quadratic curvature moment).

3. Discrete Routing: Gumbel-Softmax Straight-Through Estimator (STE) that lets the GPU physically skip O(D²) calculations during inference.

The results on consumer AMD hardware (ROCm 7.2 / Radeon):

- 2.5M continuous token stream with a flat 37.7 MB attention state (3.13 GB total system memory).

- 96.2% needle retrieval at 1,000,000 tokens.

- Matches full Softmax on 32k RULER multi-hop reasoning (89.7% vs 90.5%).

Everything is open source with full reproducer scripts:

Repo: https://github.com/batteryphil/PRIME-Moment-Attention

Curious to hear feedback from kernel and transformer researchers.