Super Apriel: One Checkpoint, Many Speeds

1. Problem Statement

As Large Language Models (LLMs) scale to longer context windows and higher throughput requirements, Full Attention (FA) becomes a significant bottleneck. FA requires $O(n^2)$ time and linearly growing KV-cache memory, making inference memory-bound and slow for long sequences.

While hybrid architectures (mixing FA with efficient alternatives like Sliding Window Attention or Linear Recurrence) offer speedups, current approaches suffer from rigidity:

• Fixed Placement: Models typically have a single, fixed arrangement of attention types (e.g., "FA in layers 1-10, Mamba in 11-48") determined before training or via a one-time search.
• Deployment Inflexibility: A fixed architecture forces a compromise between speed and quality. It cannot adapt to varying workloads (e.g., short prompts needing high quality vs. long contexts needing speed) without retraining or deploying separate models.
• Search Complexity: The combinatorial space of placing different mixers across 48+ layers is vast ($4^{48}$), making exhaustive search impossible.

Goal: Create a single model checkpoint that supports multiple speed-quality operating points at runtime without reloading weights, allowing dynamic switching between different architectural configurations based on workload needs.

2. Methodology

A. The Supernet Architecture

The authors introduce Super Apriel, a 15B-parameter "supernet" derived from the Apriel 1.6 teacher model (48 layers).

• Mixer Vocabulary: Every decoder layer contains weights for four distinct token mixers:
  1. Full Attention (FA): Standard attention ($O(n^2)$).
  2. Sliding Window Attention (SWA): Local window attention ($O(n \cdot w)$).
  3. Kimi Delta Attention (KDA): Recurrent with channel-wise gating ($O(n)$).
  4. Gated DeltaNet (GDN): Recurrent with scalar gating ($O(n)$).
• Parameter Sharing: All layers share the same Feed-Forward Networks (FFNs), embeddings, and normalization layers. Only the mixer blocks differ.
• Active Parameters: In any specific configuration ("placement"), only one mixer type is active per layer. The total checkpoint size is ~25B parameters, but any single inference run uses ~15B active parameters.

B. Training Recipe

The training process consists of two main stages:

1. Stochastic Distillation (S1):
   • The supernet is trained to mimic the frozen Apriel 1.6 teacher.
   • Local Sampling: At each training step, a mixer is randomly selected for each layer (uniformly from the 4 types). This ensures all mixer combinations receive gradient updates.
   • Loss: A composite loss involving activation matching (MSE), Reverse KL, and Forward KL divergence.
   • Data: High-quality reasoning traces and structured data are crucial; raw pretraining data caused reasoning collapse.
2. Supervised Fine-Tuning (S2):
   • The shared parameters (FFNs, etc.) remain frozen to maintain a stable reference.
   • Only the mixer weights are fine-tuned on instruction data.
   • Targeted Sampling: To optimize specific high-value configurations, the authors use a "targeted" strategy where training focuses on a curated set of Pareto-optimal placements identified by the optimizer, rather than uniform random sampling.

C. Placement Optimization Framework

To navigate the $4^{48}$ search space, the authors developed a Surrogate-Guided Optimization pipeline:

• Cost Model: Inference cost is modeled as a linear function of the count of each mixer type.
• Cluster Expansion Surrogate: Borrowing from statistical physics, the quality score $S(x)$ of a placement is approximated as a sum of interactions (unary, pairwise, etc.) between layers.
  • $S(x) = \sum V^{(1)}_i + \sum V^{(2)}_{ij} + \dots$
• Exact Optimization: A Bayesian linear model is fitted to the surrogate. Because the interactions are local (short-range), the optimal placement for any cost budget can be found exactly using Dynamic Programming (DP), avoiding heuristic search.
• Iterative Refinement: The system iteratively evaluates candidates, updates the surrogate, and refines the Pareto frontier (Quality vs. Speed).

D. Inference & Serving

• Runtime Switching: A single checkpoint can serve different placements. Switching between presets (e.g., from "All-FA" to "Hybrid") takes 5–15 seconds as only mixer weights are swapped between CPU/GPU; shared weights remain in memory.
• Speculative Decoding: The supernet enables speculative decoding without a separate draft model. The All-FA placement acts as the verifier, while any efficient placement (e.g., All-GDN) acts as the draft model, leveraging the shared weights for high acceptance rates.

3. Key Contributions

1. First Token-Mixer Supernet: A single 15B checkpoint offering four trained mixer choices per layer, enabling flexible runtime placement.
2. Surrogate-Based Optimization: A novel application of cluster expansion and dynamic programming to solve the combinatorial layer-placement problem exactly, identifying the Pareto frontier of speed vs. quality.
3. Training Dynamics Analysis: Controlled ablations on 0.5B and 15B models reveal that while general rankings stabilize early, Pareto-optimal placements exhibit higher volatility at larger scales (15B), cautioning against extrapolating small-scale findings.
4. Open Source Release: Release of the supernet weights, Fast-LLM training code, vLLM serving extensions, and the placement optimization toolkit.

4. Key Results

Performance vs. Throughput

The authors identified several "presets" (configurations) that offer different trade-offs:

• All-FA Preset: Matches the Apriel 1.6 teacher exactly (74.2 avg score vs. 73.9 teacher).
• Hybrid Presets (Reg|Lklhd–26): 2.9× speedup with 96% quality retention.
• Aggressive Presets (Reg|Lklhd–10): 10.7× decode throughput speedup with 77% quality retention.
• Linear Pareto Frontier: The trade-off between speed and quality is smooth and predictable, validating the linear cost model.

Context Length Scaling

A critical finding is that efficiency gains compound with context length:

• Moving from 16k to 32k context, efficient placements gain 80–155% relative speedup.
• External hybrid baselines (with fixed placements) only gain 5–46% in the same range because their remaining FA layers still incur quadratic memory overhead.

Comparison to State-of-the-Art

• vs. Apriel-H1: Super Apriel offers superior flexibility and quality at the 15B scale.
• vs. External Hybrids (e.g., Nemotron, Qwen): While larger models (27B–30B) achieve higher absolute quality, Super Apriel provides multiple operating points from a single 15B checkpoint, offering better efficiency (e.g., 6.2× speedup) than many larger models.

Speculative Decoding

Using the supernet for speculative decoding (All-GDN as draft, All-FA as verifier) achieves the highest speedup, confirming that the joint training maintains high distributional similarity between placements.

5. Significance and Implications

• Operational Flexibility: Organizations no longer need to train and deploy multiple models for different latency/quality requirements. A single artifact can serve high-quality short-context requests and high-throughput long-context requests by switching presets at runtime.
• Cost Efficiency: The ability to switch to aggressive hybrid presets during peak load and revert to All-FA for quality-critical tasks optimizes infrastructure costs.
• Methodological Advance: The paper demonstrates that stochastic distillation combined with exact surrogate optimization is a viable path to discovering optimal hybrid architectures, moving beyond fixed heuristics or expensive post-hoc search.
• Scalability Warning: The finding that 15B models exhibit different landscape dynamics (higher volatility in frontier placements) than 0.5B models highlights the necessity of large-scale validation for architectural search strategies.

In summary, Super Apriel transforms the speed-quality trade-off from a static, pre-deployment decision into a dynamic, runtime control surface, offering a new paradigm for efficient LLM serving.