Architecture Design and Training Stability for Million-Token Context MoE Models

1. Background: Extending Context to a Million Tokens

Long-context capability is shifting from a differentiator to a baseline requirement. Retrieval-augmented generation, long-document understanding, repository-level code analysis, and multi-turn agentic tasks all demand context windows far beyond the traditional 8K or 32K. However, extending the window directly to one million tokens triggers engineering problems on three fronts:

1. Attention computation efficiency: Standard full attention scales quadratically with sequence length, making memory and compute costs difficult to sustain. Linear attention variants reduce the asymptotic complexity to linear, but their actual throughput on GPUs often drops rather than improves.
2. Sparse expert training stability: MoE decouples capacity from inference cost by increasing total parameters while limiting activated parameters. Yet as the number of experts grows and routing becomes sparser, the risk of load imbalance and training collapse increases.
3. Long-horizon agentic reasoning: Within a million-token context, the model must maintain consistent state across many interaction turns, imposing higher requirements on positional extrapolation and deep representation propagation.

This article analyzes a representative technical solution and examines its key design choices. The specific parameters and structures cited come from public technical reports. The discussion focuses on why these designs were chosen and how they can be reused in practice.

2. Overall Architecture: Sparse MoE with Native Multimodality

The solution uses Mixture-of-Experts (MoE) as its backbone. According to the public technical report, it has 2.8 trillion total parameters and roughly 104 billion activated parameters, a typical high-capacity, low-activation configuration. A sparsity ratio of about 27:1 means each forward pass invokes only a small fraction of the total parameters. Inference cost is determined primarily by the activated portion, while model capacity is determined by the total parameter count.

Three architectural choices stand out:

• The attention module combines incremental linear attention (KDA) with gated multi-head latent attention (Gated MLA). The former enables efficient incremental computation over long sequences, while the latter preserves the accuracy advantages of standard attention. The two operate in a complementary manner.
• Cross-layer information sharing is implemented through attention residuals (AttnRes), which mitigate representation degradation in deep Transformer layers.
• The MoE side adopts a stable latent expert architecture (Stable LatentMoE) with 896 routed experts, combined with the SiTU-GLU activation function and a quantile balancing (QB) mechanism.

In addition, the model natively supports visual multimodal input and offers general capabilities such as file understanding, content writing, and data analysis, with multiple access options for end users. Behind these capability descriptions is the fact that context length and multimodal encoding are handled jointly during pretraining rather than patched together afterward.

3. Key Design 1: Linear Attention with Bounded-Decay Parameterization

3.1 The Problem: Why Diagonal Tiles Cannot Use Matrix Multiplication

Linear attention and Delta Net-style methods split sequences into chunks (typically 16 or 64 tokens). Within each chunk, explicit cumulative decay computation is performed, while recurrence handles transitions between chunks. The problem lies inside each chunk: the reciprocal of the cumulative decay factor Γ can easily overflow or underflow in BF16 precision, forcing the diagonal tile to follow a position-by-position explicit computation path that cannot invoke dense matrix multiplication on Tensor Cores.

This has two consequences. First, the diagonal tile becomes the throughput bottleneck of the entire attention operator. Second, the code path diverges: non-diagonal and diagonal tiles require two separate kernel implementations, increasing maintenance overhead.

3.2 The Mechanism: Converting Unbounded Decay to Bounded Decay

In conventional approaches, log-decay often uses the negative Softplus form, which has no lower bound. The decay factor can become arbitrarily small, making its reciprocal arbitrarily large. The solution changes the log-decay parameterization to a scaled Sigmoid form with a fixed lower bound:

# Illustrative pseudocode
g_min = -5.0
log_decay = g_min * sigmoid(e * A * h * z * t)   # Value falls within (g_min, 0)

Because log-decay is constrained to the (-5, 0) range, the reciprocal of the cumulative decay factor within any 16-token tile is kept within a safe numerical range, eliminating BF16 overflow. The diagonal tile can therefore use dense Tensor Core GEMM just like non-diagonal tiles.

3.3 Benefits and Reuse Value

The engineering value of this change is exceptionally high for three reasons:

• Minimal modification: It only affects the parameterization of the decay factor without changing the overall operator structure. It belongs to the category of changes where a single formula adjustment doubles throughput.
• Eliminates branching: All causal tiles, diagonal and non-diagonal, follow the same dense matrix multiplication path. Kernel implementation is simplified, and optimization opportunities at the CUDA/CUTLASS level are unlocked.
• Native support for positional extrapolation: Bounded decay, combined with the NoPE (no positional encoding) design, allows the one-million-token context to be extrapolated directly without additional positional interpolation.

For teams building their own linear attention or hybrid attention mechanisms, this parameterization technique is the highest-priority reuse item: it does not require rewriting operators, only re-examining the numerical range design of the decay factor.

4. Key Design 2: Auxiliary-Loss-Free Expert Load Balancing

4.1 The Problem: The Auxiliary Loss Dilemma

In MoE training, routing collapse, where a small number of experts receive the vast majority of tokens, is the most common failure mode. The mainstream industry solution is to introduce an auxiliary load-balancing loss. However, this auxiliary loss is a double-edged sword: if its weight is too small, balancing is ineffective; if too large, it interferes with the main task gradient. It also requires repeated tuning across different scales. Under routing conditions with 896 experts and extreme sparsity, the problem is further amplified.

4.2 The Mechanism: Dynamically Adjusting Routing Bias via Quantiles

Quantile Balancing circumvents gradient-level intervention by applying dynamic bias to routing logits. It tracks the quantile of token counts assigned to each expert within the current batch and adjusts the routing bias for each expert accordingly, steering token distribution toward uniformity. The entire process introduces no additional loss term and does not depend on tuning the strength of an auxiliary gradient.

Its engineering advantages include:

• No hyperparameter tuning: There is no auxiliary loss weight hyperparameter, reducing uncertainty in large-scale training.
• Fast response: The bias is updated dynamically per batch, allowing it to keep pace with drift in expert load distribution during training.
• Decoupled from sparsity: When the number of experts scales from tens to hundreds or more, the mechanism itself does not need to be redesigned.

For teams building large-scale sparse models, the significance of this strategy is that it turns load balancing from a loss function design problem into a runtime statistical correction problem, substantially reducing the complexity of maintaining training stability.

5. Key Design 3: Cross-Layer Attention Residuals

The classic problem of deep Transformers is the representation bottleneck caused by depth. As layer count increases, update magnitude in the residual stream tends to saturate, the marginal contribution of deeper layers to the final output decreases, and memory and communication overhead across layers continues to rise.

Block Attention Residuals introduce block-wise cross-layer residual connections between attention sublayers, allowing information to connect directly across layers in blocks rather than relying solely on a single residual stream passing through every layer. The effect is twofold: it mitigates deep representation degradation so that increasing depth still yields benefits, and it reduces the memory footprint and communication volume of cross-layer connections through block aggregation.

This design is especially critical for extremely long context scenarios: with a million tokens, activation memory is already tight, and any cross-layer mechanism must operate with minimal overhead. Blocking serves exactly this purpose.

6. Training Pipeline: Staged Context Progression and Agentic Post-Training

6.1 Pretraining: Staged Context Length Extension

The solution employs a multi-stage pretraining strategy progressing from 8K to 64K to 256K to 1M tokens. The rationale is that short-context stages focus on fundamental language and multimodal alignment, while long-context stages gradually introduce long-range dependencies, avoiding inefficient training on extremely long sequences from the start. Each stage transition requires synchronized adjustments to parallelism strategies, sequence partitioning, and memory layout, representing a typical case of coordinated engineering and algorithm design.

Stage	Context Length	Primary Objective
Stage 1	8K	Foundational language modeling and multimodal alignment
Stage 2	64K	Establishing medium-to-long-range dependencies
Stage 3	256K	Modeling long documents and long code structures
Stage 4	1M	Converging on extremely long context capability

6.2 Post-Training: Million-Token Agentic Reinforcement Learning

After pretraining, the solution adds agentic reinforcement learning post-training at the million-token level. Unlike conventional instruction fine-tuning, agentic tasks require the model to maintain goal consistency across extremely long trajectories, invoke tools correctly, and retrieve early information after many interaction turns. This is precisely the real stress test for long-context capability, and it in turn drives the implementation of attention efficiency and positional extrapolation designs.

7. Engineering Reuse Checklist

Overall, the reuse value of the three techniques ranks as follows:

1. Bounded-decay parameterization (high): The modification is confined to the parameterization form. It directly improves implementation efficiency of linear attention operators on Tensor Cores and simplifies kernel branching.
2. Quantile load balancing (high): It removes the auxiliary loss hyperparameter dimension and provides a more stable training path for extremely large-scale sparse MoE models.
3. Block attention residuals (medium-high): It addresses deep representation bottlenecks while controlling memory and communication overhead, making it suitable for models with substantial depth.

For teams planning to build long-context sparse models, the recommended implementation order is: first resolve numerical stability and kernel path unification for the attention operator, then address the MoE routing balance mechanism, and finally handle deep representation design and training pipeline orchestration. The three have dependencies: attention efficiency determines whether the long-context stage can run at all, routing stability determines whether model scale can be expanded, and cross-layer design determines whether additional depth remains cost-effective.

8. Summary

The combination of million-token context and trillion-parameter sparse MoE is difficult not because of any single innovation, but because multiple engineering constraints must be satisfied simultaneously: numerical precision, operator throughput, routing balance, memory overhead, and training stability. The value of the solution discussed here lies in the fact that each design targets a specific engineering bottleneck, with clear modification boundaries, independent evaluability, and standalone reusability. For developers, the trade-off logic behind these small-change, high-impact architectural decisions deserves more attention than the pursuit of parameter scale itself.