AI agents are rapidly scaling: longer sessions, deeper reasoning chains, more autonomy. But the Transformer architecture underneath them still pays a memory tax that grows with the number of tokens it processes. As agents become more capable, that cost scales with them.
Now, there’s a proliferation of hybrid architectures that use some combination of methods to achieve a better balance of memory, speed, and quality. Those methods include sparse attention, which selects a subset of past tokens to consider when selecting the next token; sliding-window attention, which preserves only the most recent tokens; and state space models, which compress information from the full history into a fixed-size state.
These hybrid architectures can be understood as points in a sequence-model design space: to differing degrees, they mix information while traversing a sequence and passing through different states. What distinguishes them from each other is how they retain and use the past.
Until now, exploring this design space has required training every new architecture from scratch, using trillions of tokens and hundreds of thousands of GPU-hours. This has made controlled comparisons of hybrid architectures extremely expensive.
In a new paper, “Priming: Hybrid state space models from pre-trained Transformers,” our team at Amazon Web Services addresses this problem. We introduce priming, a method that leverages pretrained Transformer weights as a starting point to reduce the pretraining cost of hybrid architectures by more than 99%.
With priming, model developers can ask questions that were previously too expensive to answer at scale: Which attention layers can be replaced by more-efficient layers? How many attention layers can be replaced without sacrificing quality? Which state space layer works best? And what does the trade-off between speed and accuracy look like?
In our paper, we compare hybrid models that use three different types of state space models (Gated KalmaNet, Gated DeltaNet, and Mamba-2). Some of our findings:
- Not all state space models are equal. Under controlled conditions, Gated KalmaNet (GKA) performs best, followed by Gated DeltaNet (GDN) and then Mamba-2. GKA uniquely offers variable-length inference computation without retraining and gives the user direct control over its position on the speed-accuracy frontier.
A 32-billion-parameter hybrid model reaches each accuracy threshold faster than an equivalently sized Transformer, converging to similar final accuracy in less wall clock time. The dataset is the AIME 2025 math competition benchmark (hard subset). Accuracy at time t is the fraction of rollouts for which a correct final answer has been produced by time t. - A good lossy summary of everything beats a perfect recording of the most recent past. State space models consistently outperform sliding-window attention (SWA), even when SWA gets twice the memory budget.
- Same quality, twice the throughput. A 32-billion-parameter GKA-based hybrid model matches its source Transformer on reasoning benchmarks while decoding up to twice as quickly with 128,000 tokens of context and reaching correct answers in as little as two-thirds of the wall-clock time.
To foster further research on hybrid architectures, we have released a zoo of primed hybrid models at the eight-billion-parameter and 32-billion-parameter scales for long-context reasoning and instruction following, together with training code, optimized kernels, and a vLLM-serving plugin under the Apache 2.0 License.
The attention mechanism
A Transformer’s attention mechanism — which determines how relevant previously seen tokens are to the token currently being processed — relies on three different vectors, which take their names from database design: the query, the key, and the value. For each token, the query captures what that token is “looking for,” the key captures what it “offers,” and the value carries the actual information it will pass along. Relevance is computed by comparing the current token’s query against every other token’s key; the closer the match, the more that token’s value contributes to the output of attention.
Each of these vectors is produced by multiplying the token by a learned weight projection matrix. The Transformer outputs one token at a time, and each output token becomes a new input, so the context grows one token at a time as well. To save computation, the Transformer thus stores the key and value projections of each new input in a “KV cache”.
Two kinds of memory layers: Eidetic versus fading
The attention mechanism gives Transformers what we call eidetic memory: verbatim recall of every past token. The possibility of recalling over an LLM’s entire context window is what made in-context learning, long-document reasoning, and retrieval-augmented generation possible. But it comes at a cost that scales with the number of input tokens.
However, not all attention layers do the same job. Some perform long-range retrieval across tens of thousands of tokens. Others operate more locally, tracking the current clause, copying nearby tokens, and the like. The local ones don't need a growing cache: a compact, fixed-size summary of the past — what we call fading memory — suffices.
Hybrid models mix both: attention layers (eidetic, growing cache) for the layers that need global recall and state space layers (fading, fixed size) for the rest. The result is a model with a reduced memory footprint that still has access to verbatim recall where it matters.
Priming: Knowledge isn't locked to the architecture
Exploring hybrids’ design space has been impractical because each hybrid configuration would normally require pretraining from scratch. Priming sidesteps this by reusing the stored knowledge in a pretrained Transformer's weights.
Priming is possible because a Transformer is itself a special case of a state space realization. The outputs of both attention and SSM layers can be expressed as the result of multiplying the value sequence by a “mixing” matrix that decides how much each past token contributes to each current token. Attention computes the mixing matrix by multiplying the query and key matrices; an SSM produces a similar output through a recurrent state of fixed dimension that all past information must pass through.
Priming takes advantage of this structural similarity: if an SSM with a fixed-dimensional state can produce a mixing matrix close enough to the attention layer's, the query, key, and value projections from the source Transformer can be reused as the SSM's own input projections, with only the SSM-specific parameters needing to be learned.
Whether the approximation holds depends on whether the attention layer performs complex, long-range retrievals or simple, local operations. We distinguish between the two using a cheap empirical proxy: replace each layer with an SWA — which preserves only a fraction of the context — and measure the degradation.
Layers that tolerate this substitution are natural candidates for SSM conversion: their mixing matrices don’t need to capture long-range relationships, so the state of an SSM can faithfully represent what they do. This proxy connects to a formal quantity from dynamical-systems theory, the Hankel rank of the layer's mixing matrix, which measures how much information must pass through the state bottleneck and therefore the minimum SSM state dimension required to realize the layer exactly.
Priming proceeds in three stages:
- Layer selection: identify which layers can be faithfully compressed into a fixed-size state while retaining the rest as attention;
- Weight transfer and alignment: initialize SSM layers from the corresponding attention weights, then run a short distillation process aligning each SSM to the behavior of the layer it replaces;
- End-to-end adaptation: fine-tune the full hybrid model on next-token prediction, letting all layers coadapt.
The entire pipeline consumes less than 0.5% of the source model's pretraining token budget, and it works across Transformer families (Qwen, Llama, Mistral), model classes (dense or MoE), and scales (eight billion to 32 billion and more).
With priming we can now explore the design space, the layer type, the hybrid ratio, and layer placement cheaply and systematically, because every variant starts from the same strong initialization and requires minimal training, compared to pretraining.
A controlled comparison of fading-memory layers
For instance, consider this question: given the same data, the same training budget, and the same ratio of eidetic to fading layers, which SSM layer works best in a hybrid model?
In our paper, we evaluate three families of state space models, each representing a different approach to using a fixed-size state:
- Mamba-2 applies a uniform scalar decay to the entire state at each step, a running average that forgets everything at the same rate. When a new token arrives, it is written into the state, while all previous entries are uniformly attenuated. Mamba-2 is simple and efficient but indiscriminate.
- Gated DeltaNet (GDN) adds targeted updates on top of uniform decay. Rather than decaying everything uniformly, it can selectively erase the value previously associated with a specific key before writing the new one (the delta rule). This makes it more expressive than Mamba-2: it can revise a specific association without uniformly decaying the rest.
- Gated KalmaNet (GKA) goes further. Its erase direction depends not only on the current key but on statistics of all past keys, a history-aware update that accounts for the cumulative structure of everything the model has seen.
The theoretical expressiveness hierarchy predicts the empirical one: GKA > GDN > Mamba-2. The gap widens on long-context reasoning, where compression quality matters most, and narrows on short-context tasks. Furthermore, at the 32-billion-parameter (32B) scale, hybrid GKA is within 1% of its source Transformer trained on the same data, despite requiring half the KV cache at inference.
Fading memory beats verbatim windowed memory
An interesting finding in our study comes from a direct comparison of hybrid models using SSM layers and those using sliding-window attention (SWA), which retain a fixed window of recent tokens verbatim rather than compressing the past into a state. One might expect that verbatim recall, even over a limited window, would outperform lossy compression. After all, no information is lost within the sliding window’s range. And stacking multiple SWA layers increases the “receptive field” of the end-to-end model.
Instead, every SSM-based hybrid we tested outperformed the corresponding SWA hybrid on long-context benchmarks, even when the window gets two to four times as much memory as the state.
A runtime knob for speed versus accuracy: Variable computation at inference time
GKA has a property that neither Mamba-2 nor GDN offers: variable test-time computation.
Mamba-2 and GDN have the same computation budget for every input token, regardless of its content. GKA, by contrast, solves a regression problem iteratively, and the iteration count is adjustable at serving time without retraining.
In practice, serving GKA at 10 iterations instead of its usual 30 sacrifices roughly 1% of its accuracy on AIME 2025 while delivering 10% faster decoding. A single trained model can serve multiple points on the speed-accuracy frontier, so model providers can offer different latency tiers without maintaining separate model variants.
This connects to a principle we have discussed in earlier work: intelligence is the ability to allocate effort commensurate with difficulty. GKA's variable-computation mechanism is an architectural embodiment of this principle: a deployment that demands speed can run fewer iterations, while one that demands accuracy can run more.
Faster reasoning, not just faster tokens
Raw per-token throughput tells only part of the story. A model that is cheaper per token can still be slower end to end if its reasoning is weaker and it needs more tokens to compensate. The deployment-relevant question is how much more rapidly, in wall clock terms, the hybrid reaches a target accuracy.
We measured this on AIME 2025, a challenging math competition benchmark, comparing our GKA-Primed-HQwen3-32B-Reasoner against the source Qwen3-32B trained on the same reasoning recipe. Both models generate similar-length thinking traces (a median of about 14,000 tokens on the full set, 28,000 on the hardest problems) — that is, they produce equally in-depth responses — but the hybrid has a much higher throughput.
At realistic concurrency levels, where the Transformer's KV cache saturates GPU memory, the hybrid reaches target accuracy thresholds 1.2 to 1.4 times as quickly on the full benchmark and up to 1.6 times as quickly on the hardest problems. This is a consequence of the SSM layers' fixed-size state, which roughly halves the per-sequence memory footprint, enabling more concurrent operations on the same hardware.
The throughput advantage grows with context length: 1.50× at 16K tokens, scaling to 2.06× at 128K for the 32B model (GKA at 10 iterations). With the 8B parameter models, the gains are even greater: 1.94× at 16K to 2.29× at 128K. At long context, KV cache reads dominate the forward pass through the Transformer, and halving the cache roughly doubles the effective batch size at memory saturation.
This has implications beyond inference serving. In reinforcement learning (RL) post-training pipelines, throughput gates how rapidly the policy generates correct trajectories. The hybrid's advantage translates directly to a 20% to 100% increase in RL training data rate at a fixed GPU budget.
Beyond the training window: Context extension with state composition
Our primed hybrids natively support 128,000 tokens. But agent sessions, codebases, and long legal documents routinely exceed that. Extending context further would normally require another round of continued pretraining on curated data. Can the hybrid architecture itself generalize to longer contexts without additional training?
Existing training-free methods for expanding the context window fall into two classes: those for pure Transformers (positional-embedding rescaling schemes like YaRN) and those for pure SSMs, such as our PICASO, which we presented last year at the International Conference on Learning Representations (ICLR). Neither addresses the hybrid setting, where KV caches and recurrent states must be reconciled jointly.
In our new arXiv paper, we describe a method of extending hybrid models’ context beyond their native limits without training. The procedure is simple:
- partition a long input into chunks shorter than the native context length and prepend a shared prefix to each chunk;
- process all chunks independently in parallel;
- concatenate the KV caches across chunks while merging the SSM states via the weighted-averaging scheme we introduced in our PICASO paper.
This approach enables us to extend the context length of our primed hybrid models from 128K to 256K tokens while preserving substantial accuracy on both retrieval (RULER NIAH) and reasoning (BABILong) tasks at lengths the model was never trained on. In our experiments we show that hybrid state composition preserves the SSM expressivity hierarchy (GKA > GDN > Mamba-2).
The hybrid design space is now open
All artifacts discussed in this blog are released under the Apache 2.0 License in our Hybrid Model Factory on GitHub and our model zoo on HuggingFace:
- A collection of primed hybrids for long-context reasoning and instruction following (8B and 32B);
- The priming training code;
- Point-to-point sequence parallelism algorithms for SSM layers (up to 2.3 times as fast as all-to-all-based alternatives at multinode scale);
- Universal sequence parallelism for prototyping novel layers on long sequences without custom kernels or distributed-training expertise;
- Optimized Triton kernels for GKA decoding and training;
- A vLLM serving plugin supporting primed hybrid architectures with continuous batching, chunked prefill, and tensor parallelism for inference deployment at scale.
If you serve long-context workloads, the 32B GKA reasoning model with the vLLM plugin gives you higher throughput with no quality loss on long context reasoning tasks.
If you're scaling RL post-training, the hybrid's throughput advantage means more correct rollouts per unit time at a fixed GPU budget, for faster policy iteration without changing your training recipe.
If you're researching new architectures, priming lets you evaluate any layer type under controlled conditions at a fraction of the pretraining cost while starting from a strong initialization. The universal SP module means you don't need custom distributed kernels to train at long context.
Priming is not tied to a specific Transformer family or model class. Extending it to Llama, Mistral, or MoE models requires no algorithmic changes. We hope these tools enable the community to explore architectures beyond what we've tested and to develop new algorithms and training recipes for building scalable next-generation architectures for AI agents.
For full details, see our paper, "Priming: Hybrid state space models from pre-trained Transformers". Code, models, and serving plugin are available on GitHub, and the models are also on Hugging Face.
