Architectural Design of MoE Large Models for Million-Token Contexts
This article examines the architectural design choices behind a trillion-parameter sparse Mixture-of-Experts model with million-token context support. It focuses on three core innovations: bounded decay parameterization for linear attention, quantile-based load balancing for extremely sparse MoE training, and block-level attention residuals for deep transformer representations. The article also covers the multi-stage pretraining strategy and agentic reinforcement learning post-training used to turn long-context capacity into practical task performance.
Architectural Design of MoE Large Models for Million-Token Contexts
1. Background: When Context Length Crosses the Million-Token Threshold
Long-text processing has evolved from a differentiator into a baseline capability for large models. When input scale jumps from a few thousand tokens to millions, the traditional Transformer architecture faces three simultaneous pressures: the quadratic complexity of attention computation causes compute costs to balloon; cross-layer information transfer in deep networks becomes less efficient, creating a representation bottleneck; and in Mixture-of-Experts (MoE) architectures with very large expert counts and extreme activation sparsity, uneven routing loads can trigger training instability or even divergence.
This article examines a multimodal MoE large-model architecture designed for regular users and supporting file understanding, content generation, and data analysis. It analyzes the design trade-offs in three areas: long-context efficiency, MoE training stability, and long-horizon agentic reasoning. The model adopts a very large sparse MoE structure with total parameters at the trillion level and activated parameters in the hundreds of billions per forward pass. It natively supports visual multimodal inputs and a million-token context window, and is served through multiple access methods. It is worth emphasizing that this article focuses on reusable architectural ideas and engineering patterns rather than a specific product evaluation.
2. Overall Architecture: Coordinated Design of Attention and MoE
The model backbone consists of two key module types. On the attention side, it combines Key-Decay Attention (KDA) with Gated Multi-head Latent Attention (Gated MLA), and uses attention residuals (AttnRes) to share information across layers. On the MoE side, it employs a Stable LatentMoE architecture with hundreds of routed experts, paired with the SiTU-GLU activation function and a Quantile Balancing (QB) mechanism to ensure training stability under extreme sparsity conditions.
This parallel “attention + MoE” design is not a simple stacking of components. Rather, the two sides are jointly optimized for long-context scenarios: the attention modules control compute and memory growth over very long sequences, while the MoE modules expand parameter capacity while keeping per-token compute cost manageable. Together they determine the feasibility and economic viability of the model at million-token contexts.
3. Core Innovation 1: Lower-Bounded Decay Parameterization in KDA
3.1 The Problem: Numerical Overflow Under Limited Precision
When linear attention and its variants (such as Delta Net-style structures) are computed in chunks, the reciprocal of the intra-chunk cumulative decay factor Γ can exceed the representable range of BF16. Once overflow or underflow occurs, the diagonal blocks can no longer use the standard matrix multiplication path and must instead fall back to explicit position-pair computation. This means diagonal and off-diagonal blocks require two separate computational branches, increasing implementation complexity and preventing diagonal blocks from fully utilizing dense tensor-core matrix multiplication acceleration.
3.2 The Mechanism: Converting Unbounded Decay to Bounded Decay
The solution is to change the logarithmic decay from the unbounded Negative-Softplus form to a scaled Sigmoid form with a fixed lower bound. By setting a decay floor (for example, -5), the decay factor is constrained within a safe interval:
g_t = g_min * Sigmoid(e^A * h_z * t) # g_min is a fixed negative lower bound
As a result, within any fixed-size tile (such as 16 tokens), the reciprocal scaling factor stays within a bounded numerical range and no longer triggers precision anomalies.
3.3 Benefits: A Unified Computation Path and Length Extrapolation
The direct effect of this modification is that diagonal and off-diagonal blocks now follow the same causal tile computation path. Both can directly invoke dense tensor-core GEMM operations, eliminating the throughput loss caused by special-case branching. Additionally, because the decay parameterization is inherently position-insensitive, the model natively supports million-token context extrapolation in a NoPE (no positional encoding) setting, avoiding the extra engineering burden of position interpolation. This change has minimal engineering invasiveness yet significantly improves the implementation efficiency of the linear attention operator on mainstream GPU programming frameworks, making it highly reusable.
4. Core Innovation 2: Quantile Balancing (QB) for Load Balancing
4.1 The Problem: Auxiliary Loss Versus Very Large Expert Counts
Traditional MoE training usually relies on an auxiliary loss to constrain expert load balancing. However, in extremely sparse settings with hundreds of routed experts and very low activation ratios, the auxiliary loss can compete with the main task loss in terms of gradients. Slight misadjustment of the loss weight can degrade model quality, while prolonged load imbalance leads to expert collapse and training instability.
4.2 The Mechanism: Dynamic Thresholds via Quantiles
The Quantile Balancing mechanism abandons the auxiliary loss approach entirely. Instead, it uses quantiles of routing scores as dynamic thresholds to allocate expert capacity: at each training step, the quantile is computed from the current routing score distribution, and tokens exceeding the threshold are assigned to the corresponding expert. This achieves rapid dynamic load balancing without introducing an additional loss term. Its advantages are twofold: the balancing constraint is decoupled from the task objective, so it does not pollute the main gradients; and the threshold adjusts dynamically during training, adapting to distribution drift across different stages.
4.3 Benefits: Stability and Efficiency Combined
This strategy maintains training stability under very large expert configurations and avoids the repeated trial-and-error cost of tuning auxiliary loss weights. Because it eliminates the extra loss computation and backward pass overhead, training throughput also improves. For teams planning to build large-scale sparse MoE models, this is an engineering path worth prioritizing for evaluation.
5. Core Innovation 3: Block-Level Attention Residuals
5.1 The Problem: Representation Bottlenecks in Deep Transformers
As the number of layers grows, information in the residual stream becomes progressively diluted, reducing the efficiency with which deeper layers reuse shallow features. At the same time, if cross-layer information sharing is implemented as full connectivity, it introduces substantial memory usage and communication overhead, which becomes especially problematic under long-context and large-scale parallel training.
5.2 Mechanism and Benefits
The block-level attention residual approach organizes residual sharing by block, aggregating and propagating cross-layer information at block granularity rather than through full pairwise layer interactions. This breaks the deep representation bottleneck, improves gradient and information flow, and keeps memory and communication overhead within acceptable limits. For deep backbone networks that must maintain stable training on long sequences, this is a cost-effective structural enhancement.
6. Training Paradigm: Multi-Stage Pretraining and Agentic Reinforcement Learning Post-Training
In terms of data and training scale, the model adopts a progressive context extension strategy:
Pretraining context progression: 8K → 64K → 256K → 1M
The benefit of staged extension is that short-context phases first allow the model to fully learn foundational language and multimodal alignment capabilities. The sequence length is then gradually increased, letting the model adapt to longer dependency structures and reducing both convergence difficulty and compute waste compared with training directly on very long sequences. After pretraining, million-token-level agentic reinforcement learning (Agentic RL) post-training is applied to optimize long-horizon reasoning, tool calling, and multi-step planning behaviors, transforming the ability to “read long documents” into the ability to “complete complex tasks within long documents.”
This paradigm closes the loop on the three core challenges: KDA and sparse attention address the computational efficiency of very long contexts; QB and the stable expert architecture address training stability for massive expert counts; and Agentic RL addresses the upper bound of long-horizon agentic reasoning.
7. Engineering Reuse Recommendations
For teams planning to build or optimize long-context models, the following technologies can be evaluated in order of priority:
- Prioritize lower-bounded decay parameterization: The changes are concentrated in the decay factor computation, requiring minimal engineering effort. It unifies the diagonal block computation path, unlocks tensor-core performance, and provides positional-encoding-free length extrapolation as a side benefit.
- Evaluate quantile balancing as an auxiliary loss replacement: In scenarios with many experts and low activation ratios, start with small-scale ablation comparisons to confirm balancing effectiveness and training stability before full deployment.
- Introduce block-level attention residuals as needed: Models with deeper layers and tight memory constraints during long-sequence training benefit more. Start with a subset of layer blocks as a pilot.
- Maintain progressive context extension: Avoid pretraining directly on million-token sequences. Staged progression significantly reduces convergence risk and compute cost.
8. Summary
The value of this architecture lies not in stacking individual metrics but in providing a decomposable, reusable engineering combination: bounded decay enables linear attention to fully utilize hardware; quantile balancing stabilizes training of very large-scale MoE models; block-level residuals improve deep representations; and multi-stage pretraining with agentic post-training turns long-context capacity into practical task performance. For applications serving regular users with file understanding, content generation, and data analysis capabilities, these design choices collectively form the technical foundation for a million-token context experience.
