Architecture Deep Dive: Scaling Trillion-Parameter MoE Models for Million-Token Contexts

1. Background: When Model Scale Meets Context Length

The capability boundaries of large models are expanding along two simultaneous axes: parameter scale and context window size. Parameter scale is moving from dense architectures to Mixture-of-Experts (MoE), with total parameter counts climbing continuously. Context windows are advancing from a few thousand tokens to the 100K and 1M range. The intersection of these two axes introduces three intertwined engineering challenges.

First is the computational efficiency of ultra-long contexts. The complexity of standard attention mechanisms grows quadratically with sequence length. At the million-token level, both memory consumption and computational overhead become unmanageable. Linear attention and its variants reduce complexity to linear via chunked recursion. However, in actual hardware implementations, the cumulative product of decay factors is highly prone to overflow or underflow under limited precision. This forces the diagonal blocks to take a special position-by-position computation path, preventing the use of high-throughput tensor core matrix multiplications.

Second is the training stability of massive MoE. When the number of routed experts reaches the hundreds or nearly a thousand, and the activation ratio is pushed to extremely low levels, load distribution easily collapses: a few experts receive the vast majority of tokens, while the rest rarely receive effective gradients. The traditional approach introduces an Auxiliary Loss to force balancing. However, auxiliary loss has gradient conflicts with the main task objective, requires high tuning costs for weight parameters, and exhibits unpredictable convergence behavior at large scales.

Third is long-range agentic reasoning. Long-range tasks require models to maintain consistent planning and memory across interaction histories spanning hundreds of thousands of tokens. This places new demands on positional encoding extrapolation, cross-layer information flow efficiency, and reinforcement learning signal design in post-training.

These three challenges are not independent. The longer the context, the more prominent the numerical stability issue of attention becomes. The sparser the experts, the more training stability relies on fine-grained load regulation. The quality of long-range reasoning directly depends on whether the former two can be stably optimized during training. This article deconstructs how a trillion-parameter MoE multimodal architecture addresses these challenges layer by layer.

2. Overall Architecture: Combining Sparse Experts with Hybrid Attention

The overall form of this architecture can be summarized in three parts: a hybrid attention module fusing two attention mechanisms, an extremely sparse MoE feedforward layer, and cross-layer residual pathways connecting the layers.

On the attention side, the model combines incremental linear attention with Gated Multi-Head Latent Attention (MLA). The former efficiently aggregates historical information via recursion within chunks, while the latter compresses the KV cache through low-rank latent representations, reducing memory pressure on long sequences. The two are complementary: the linear path handles long-range decaying memory, and the latent attention path retains precise local retrieval capabilities.

On the MoE side, the routed expert scale reaches 896. Coupled with expert parameterization in the latent space and specific activation function designs, it maintains training stability under extreme sparsity. The activated parameter count is controlled at the hundred-billion level, decoupling total parameters from the computation required per token.

For inter-layer connectivity, a chunked attention residual mechanism is used to share information across layers. The value of this design is particularly evident in deep networks—as layers deepen, the representational capacity of the residual stream easily hits a bottleneck. Cross-layer direct connections can alleviate deep gradient attenuation and reduce unnecessary memory and communication overhead.

Pre-training adopts a multi-stage context progression strategy: starting from 8K, expanding to 64K, then 256K, and finally reaching 1M. This progressive window expansion allows the model to fully converge on short sequences first, then gradually adapt to the positional distribution of long sequences, avoiding the optimization difficulties of training directly on ultra-long contexts. The post-training phase stacks million-token-level agentic reinforcement learning, specifically optimizing for long-range planning and tool-calling scenarios.

3. Key Technology 1: Bounded Decay Parameterization

The Problem: The Diagonal Block Dilemma Under Limited Precision

When linear attention and Delta Net-like architectures perform chunked computation, the cumulative decay factor within a chunk is a continuous product. When the decay rate is small, the reciprocal of this factor rapidly expands; when the decay rate is large, the factor itself approaches zero. Under limited precision formats like BF16, this bidirectional numerical extremization causes severe precision loss.

The engineering consequence is that to implement causal masking, the diagonal blocks within a chunk (where queries and keys reside in the same chunk) must bypass the standard matrix multiplication path and use explicit position-pair element-wise computation. This special branch not only has low throughput but also breaks the continuity of operator fusion, making it difficult for the entire attention kernel to achieve ideal implementation efficiency on CUDA/CUTLASS.

Mechanism: Changing Unbounded Decay to Bounded Scaling

The solution is to change the log decay parameter from an unbounded negative Softplus form to a bounded Sigmoid scaling form. Let the decay parameter be g. After applying a Sigmoid mapping with a fixed lower bound, the value of g is constrained to a safe interval (e.g., a lower bound of -5). Thus, within any fixed-length tile (e.g., 16 tokens), the reciprocal scaling result of the decay factor falls within a representable numerical range without hitting precision limits.

Formally, the original unbounded form can be written as a log decay mapped through negative Softplus, while the new form changes to Sigmoid scaling with a fixed lower bound:

# Original form: Unbounded, reciprocal prone to overflow
log_decay = -softplus(x)          # Range (-inf, 0)
# New form: Bounded, numerically safe
log_decay = g_min * sigmoid(x)    # g_min = -5, Range (-5, 0)

Benefit: Unified Dense Matrix Multiplication

The direct benefit of numerical safety is the normalization of the computation path: diagonal blocks no longer require special branches. All causal tiles—whether diagonal or non-diagonal—can directly invoke dense tensor core GEMMs. This eliminates scheduling overhead and operator fragmentation caused by branching, significantly improving throughput. Furthermore, bounded decay allows the model to naturally support direct extrapolation to million-token contexts under the NoPE (No Positional Encoding) setting. Because the decay itself carries relative positional information, the model can maintain reasonable behavior on sequences longer than the training length without relying on explicit positional encodings.

From a reuse perspective, the engineering invasiveness of this change is minimal—simply replacing the parameterization function of the decay parameter yields efficiency gains on existing linear attention operators, making it a low-cost optimization item before large-scale training.

4. Key Technology 2: Auxiliary-Loss-Free Load Balancing

The Problem: The Inherent Contradiction of Auxiliary Loss

MoE load balancing has long relied on auxiliary loss: adding a penalty term measuring expert load variance to the training objective to force the routing distribution toward uniformity. However, auxiliary loss is fundamentally an optimization direction different from the language modeling objective, creating gradient-level tension between the two. If the auxiliary loss weight is too large, main task convergence suffers; if too small, load collapse persists. Under the setting of 896 experts and an extremely low activation ratio, this tuning difficulty is further amplified.

Mechanism: Dynamic Calibration of Routing Bias via Quantile Estimation

The concept of Quantile Balancing transforms load balancing from a “loss term” to a “dynamic calibration of routing bias”. The core is to maintain a bias value for each expert and update this bias based on the quantile statistics of routing scores within the current batch: the bias for heavily loaded experts is adjusted downward, while the bias for lightly loaded experts is adjusted upward. Because the update is based on quantiles rather than means, it is more sensitive to the long tail of the distribution and can quickly flatten the load even under extreme sparsity.

Since balancing is achieved through the routing decision itself, the training objective no longer requires an auxiliary loss term, keeping the main task gradient pure. The advantages of this mechanism are: no need to introduce additional hyperparameter weights, as the balancing strength is adaptively determined by quantile statistics; the update overhead only involves lightweight statistics and bias correction, adding no significant computational burden; and it responds quickly during early training when loads fluctuate drastically, without introducing extra noise in later stages when it stabilizes.

For any training system employing large-scale sparse experts, this strategy holds high reuse value—it shifts load balancing from an “objective function design problem” to a “routing control problem”, reducing tuning dimensions.

5. Key Technology 3: Chunked Attention Residual

The representation bottleneck of deep Transformers is a repeatedly observed phenomenon: as the number of layers increases, the incremental effective information in the residual stream diminishes, and the contribution of deep layers to the final output is diluted. In ultra-long context scenarios, this problem is compounded by memory and communication costs—if direct connections to deep or output layers are simply established for every layer, the volume of cross-device activation transmission will expand linearly with layers.

The chunked attention residual approach groups layers into several blocks and establishes residual pathways at the block granularity, rather than fully connecting layer by layer. Standard residual structures are maintained within blocks, and information is selectively transmitted between blocks via attention-style weighted aggregation. The benefits are twofold: representationally, deep layers can directly access high-level semantics from early blocks, alleviating the depth bottleneck; engineering-wise, the granularity of cross-block communication shifts from “layers” to “blocks”, meaning memory consumption and inter-device communication volume scale with the number of blocks rather than layers, ensuring scalability in distributed training.

The reuse scenarios for this technology are not limited to ultra-long context models—any architecture with deep layers where cross-device training communication becomes a bottleneck can benefit from a block-granularity residual redesign.

6. Engineering Insights and Selection Recommendations

Looking at these technologies together yields several experiences of universal value for large-scale model training.

Numerical stability should be a first-class citizen in architecture design. The case of bounded decay parameterization shows that a seemingly minor parameterization choice can amplify into differences in throughput and scalability via hardware execution paths. When designing linear attention-type operators, priority should be given to checking the dynamic range of cumulative product-type factors under the target precision format. Proactively introducing bounded constraints when necessary is better than patching with special branches downstream.

Prioritize control-based solutions for load balancing. Auxiliary loss is a “soft constraint”, while quantile bias calibration is a “hard control”. The latter achieves a more direct balancing effect without polluting the main task gradient, making it especially suitable for scenarios with large expert scales and low activation ratios. In engineering implementation, the key is the efficient approximation of quantile statistics and the asynchronous safety of bias updates.

Adopt a progressive strategy for context window expansion. The multi-stage progression from 8K to 1M essentially allows the model’s attention distribution to gradually adapt to longer dependency spans. Training directly at the target length is not only costly but also prone to optimization instability due to sudden changes in positional distribution. Progressive window expansion combined with NoPE extrapolation capabilities is a proven viable path.

Design cross-layer connections by block, not by layer. Block granularity provides an adjustable compromise between representational gain and communication costs. The choice of block size should be determined in conjunction with device topology and bandwidth budgets.

7. Conclusion

The combination of a trillion-parameter MoE and million-token contexts tests not just the advancement of single-point algorithms, but the synergy among numerical precision, routing control, inter-layer communication, and training curriculum design. Bounded decay parameterization solves the execution efficiency problem of linear attention on tensor cores. Quantile balancing transforms MoE load balancing from a loss design problem to a routing control problem. Chunked attention residual strikes a balance between deep representation and distributed communication. Together, these three support progressive pre-training from 8K to 1M and million-token-level agentic post-training, providing a reusable engineering paradigm for the architectural selection of similar large-scale sparse models.