Architecture Design for Trillion-Parameter MoE Models with Ultra-Long Context

1. Background: When “Larger Models” Meet “Longer Contexts”

The evolution of large model capabilities now follows two clear trajectories: one is parameter scale, where sparse Mixture-of-Experts (MoE) architectures raise the ceiling of model capacity while controlling inference costs; the other is context length, expanding from thousands of tokens to hundreds of thousands or even millions, enabling repository-level code understanding, long-document analysis, and long-horizon agent tasks. Once these two trajectories converge, the problem is no longer a matter of isolated technical optimization—it becomes a set of coupled systems engineering challenges, commonly referred to as three “walls”:

1. The attention computation and memory wall. Standard full attention scales quadratically with sequence length, and KV Cache memory grows linearly with length. At the million-token scale, simply adding more GPUs is no longer sufficient; the solution must come from restructuring at the operator and architecture level.
2. The training stability wall with massive experts. When MoE routed experts scale from dozens to hundreds or even nearly a thousand, problems such as load imbalance, expert collapse in the later stages of training, and interference between auxiliary losses and the primary objective are dramatically amplified. Traditional balancing methods begin to fail.
3. The long-horizon agentic reasoning wall. Agent task contexts consist of multi-round tool calls, intermediate results, and environmental feedback. Their length is unpredictable, dependency spans are extremely long, and the model must maintain goal consistency and stable output formatting under ultra-long dependencies.

Recent public frontier technical reports are generally designed around these three walls. This article takes a general technical perspective to break down the key mechanisms and engineering considerations involved, helping developers understand the construction logic behind “trillion-parameter, million-token” models.

2. Attention Side: Making Ultra-Long Contexts Computable and Storable

2.1 Lower-Bounded Decay Parameterization: A Small Change for Large Speedups

Linear attention and incremental attention methods, which replace quadratic-complexity attention matrices with fixed-size recurrent states, represent an important direction for long-context handling. However, they have two long-standing bottlenecks in engineering deployment:

• Numerical overflow/underflow: In the chunked parallel computation paradigm, the reciprocal scaling term of the cumulative decay factor is needed within each chunk. Under traditional parameterization, the decay factor is unbounded on the negative half-axis, and its reciprocal easily overflows under limited precision such as BF16, leading to unreliable computation results.
• Diagonal blocks cannot use matrix multiplication paths: Due to the numerical risks above, the diagonal block portion in chunked computation can only degrade to position-by-position pairwise computation, unable to leverage low-level matrix multiplication acceleration units, greatly reducing hardware utilization.

A highly reusable solution is to change the logarithmic decay parameterization from “unbounded” to “lower-bounded”: instead of using the unbounded negative half-axis softplus form, adopt a scaled sigmoid form and fix the lower bound at a safe magnitude (e.g., -5). An illustration follows:

# Traditional parameterization (illustrative): unbounded negative half-axis, numerical risk
g_t = -softplus(raw_t)              # g_t ∈ (-∞, 0), reciprocal scaling factor may overflow
# Lower-bounded parameterization (illustrative): bounded, numerically safe
g_t = g_min * sigmoid(raw_t)        # g_min = -5, g_t ∈ [-5, 0)
scale_t = exp(g_t / tile_len)       # reciprocal scaling factor within tile locked to safe range

This change locks the reciprocal scaling factor within any 16-token tile into a safe numerical range, allowing diagonal blocks to safely use matrix multiplication implementations. The entire operator can fully leverage hardware efficiency on high-performance GPU operators. Minimal engineering change, maximum benefit—a classic example of “one formula saves an entire optimization effort.”

2.2 Gated Multi-Head Latent Attention and Hybrid Layout

The second key mechanism is the gating of Multi-Head Latent Attention (MLA). MLA significantly reduces KV Cache memory footprint through low-rank latent compression, making it one of the mainstream approaches for reducing memory pressure in long-context inference. By introducing a learnable gate on its output path, the model can dynamically control how much attention information flows into the residual stream, mitigating the “attention dilution” problem commonly seen with ultra-long sequences.

In terms of architectural layout, the industry-standard approach is layered hybrid rather than uniform full attention:

• Local layers: Sliding window attention, responsible for near-distance dependencies at constant cost;
• Linear layers: Incremental attention, maintaining global information with a fixed state;
• A small number of global layers: Full attention, responsible for anchoring critical cross-length information.

This “local + linear + global” hybrid recipe achieves an engineering-acceptable balance between efficiency and capability, and is the de facto choice for most current long-context models.

2.3 Block-Level Attention Residuals: Opening “Information Shortcuts” for Deep Networks

As model depth continues to increase, information must be passed layer by layer, gradient paths grow longer, and the representational capacity growth of deep networks exhibits diminishing marginal returns. The idea of block-level attention residuals is to inject shallow-layer attention outputs into deep layers across layers, effectively providing explicit “information shortcuts” for deep networks:

• Breaking the representational bottleneck caused by depth stacking, allowing deeper layers to produce real gains;
• Sharing intermediate results across layers, reducing memory and communication overhead from redundant computation;
• Achieving more stable training dynamics, conducive to stable convergence of ultra-large-scale models.

3. MoE Side: Enabling Nearly a Thousand Experts to Collaborate Stably

3.1 Fine-Grained Experts and High Sparsity Ratios

The clear trend in current MoE design is “many small experts“: the number of routed experts expands from dozens to hundreds or even nearly a thousand, while a small number of shared experts are retained to carry general knowledge. At inference time, only Top-k experts are activated. Fine-grained decomposition causes the combinatorial expression space of experts to grow exponentially, while the activated parameter ratio can still be kept to single-digit percentages of total parameters—this is precisely the source of the typical configuration of “trillion-level total parameters, tens-of-billions-level activated parameters.”

3.2 Quantile-Based Load Balancing: Moving Beyond Auxiliary Losses

Traditional Auxiliary Loss-based load balancing has three persistent issues: gradient conflicts with the primary objective, sensitivity to hyperparameters that makes tuning difficult, and failure as a “soft constraint” under extreme sparsity. Quantile Balancing offers an alternative approach: instead of adding a new loss term, it performs dynamic redistribution directly at the routing level. During training, it continuously tracks the load distribution of each expert and, by quantile, downweights routing scores for overheated experts and upweights those for underutilized experts:

# Language-agnostic pseudocode: quantile-based load balancing (illustrative)
load = ema_expert_load()                 # exponential smoothing statistics of recent load per expert
q_hi = quantile(load, 0.9)               # high quantile threshold
q_lo = quantile(load, 0.1)               # low quantile threshold
score = router_affinity(x)               # routing scores
score[load > q_hi] *= shrink_factor      # downweight overheated experts
score[load < q_lo] *= boost_factor       # upweight underutilized experts
selected = topk(score, k)

Its advantages: zero auxiliary loss hyperparameters, takes effect immediately at every training step, and maintains stable balance even under extreme sparsity configurations with nearly a thousand experts. For teams developing their own smaller-scale MoE models, this approach is equally applicable.

3.3 Activation Functions and Numerical Stability Details

On the FFN side, gated activation functions such as SwiGLU and their improved variants, combined with numerical clipping of routing scores and stable initialization of expert weights, jointly ensure that training does not diverge under extreme sparsity. These details may seem insignificant individually, but when scaled to nearly a thousand experts, any numerical hazard can be amplified into training collapse.

4. Training Strategy: “Growing” a Million-Token Context

4.1 Multi-Stage Context Extension

Ultra-long context capability is not achieved in a single training run. The standard approach is stepwise progression: 8K → 64K → 256K → 1M. Each stage requires synchronizing four things:

1. Adjusting positional encoding parameters (e.g., RoPE base frequency) to ensure effective length extrapolation;
2. Shifting data mixture toward long documents, code repositories, and multi-turn long conversations;
3. Switching parallel strategy combinations (different ratios of sequence parallelism, context parallelism, and expert parallelism);
4. Performing length generalization validation at stage boundaries to avoid “training collapse from one-shot extension.”

4.2 Agentic RL Post-Training

Pre-training addresses “knowledge”; post-training addresses “behavior.” The uniqueness of agent scenarios lies in: a single trajectory can span dozens of interaction rounds with million-level context length; rewards are sparse and delayed (success only counts upon task completion); and the model must simultaneously master tool-call formatting, task planning, and self-reflection. Therefore, performing reinforcement learning on million-token-level long-horizon tasks, using task success rate as the reward signal, and teaching the model “long-horizon goal consistency” across complete multi-round trajectories has become the standard post-training configuration for top open-source models.

5. Practical Takeaways for Developers

1. Evaluate KV Cache curves when selecting models: For workloads with long-document or long-session requirements, prioritize hybrid attention architectures and measure actual memory curves at different lengths rather than only looking at nominal context length.
2. Consider dropping auxiliary losses for load balancing: Even with a moderate number of experts, quantile-based clipping-style balancing is worth trying, with significantly lower tuning costs.
3. Numerical safety is the first lesson of linear attention: During deployment, explicitly verify the value range of scaling factors within tiles, and prefer bounded forms for decay parameterization.
4. Extend length in stages: Whether continuing pre-training or fine-tuning, context length should be advanced and validated level by level.
5. Leverage mature infrastructure: Large-scale training and inference involve compute scheduling, parallel strategies, monitoring and alerting, and other systems engineering concerns. Teams can rely on GPU compute clusters and model training/inference service platforms—such as those provided by Baidu AI Cloud—to handle the engineering foundation, concentrating team efforts on algorithm iteration.

6. Conclusion

Looking back at the technical reports of this generation of frontier models, a clear paradigm can be distilled: co-design of algorithms and systems. Lower-bounded decay parameterization solves “whether operators can saturate hardware”; quantile-based balancing solves “whether nearly a thousand experts can train stably”; block-level residuals solve “whether deep networks can continue to benefit”; multi-stage extension and agentic post-training solve “whether capabilities can truly emerge.” For developers, these mechanisms need not all be adopted wholesale, but the engineering philosophy behind them—“ensure numerical and system feasibility first, then pursue capability ceilings”—is worth internalizing for every team working on large model technology.