Long-Context Architecture Design for Trillion-Parameter MoE Models: KDA, Quantile Balancing, and Block Attention Residual

Background: When Context Windows Reach a Million Tokens

As LLM applications move beyond single-turn Q&A into long-document understanding, repository-level code analysis, and long-horizon agent reasoning, context windows have advanced rapidly from the early 8K and 64K to 256K and even 1M tokens. Meanwhile, model scale keeps expanding along the Mixture-of-Experts (MoE) route: total parameters in the trillions with roughly hundred-billion-level activated parameters per forward pass has become a common configuration in frontier research.

Scaling size and context simultaneously creates three intertwined engineering challenges:

1. Computational efficiency over ultra-long contexts. The compute and memory cost of attention grows rapidly with sequence length. Full attention is nearly infeasible at the million-token scale, making linear attention and chunked attention mandatory—yet both carry inherent weaknesses in kernel efficiency and numerical stability.
2. Training stability of massive-expert MoE. When the number of routed experts reaches several hundred or close to a thousand and the activation ratio is extremely low, load imbalance can starve some experts while overloading others, and training can easily diverge.
3. Representation bottlenecks in deep Transformers. As networks get deeper, gradients and information flow less efficiently across layers, deep representations tend to saturate, and cross-layer communication and memory overhead rise significantly.

This article examines the architecture of a trillion-parameter multimodal MoE model and breaks down the three techniques it uses to tackle these challenges: KDA lower-bound decay parameterization, Quantile Balancing (QB), and Block Attention Residual (Block AttnRes). All three are highly reusable and require relatively contained engineering changes.

Architecture Overview

The model adopts a MoE architecture with total parameters at the trillion scale and roughly hundred-billion-level activated parameters per forward pass. It natively supports visual multimodal input and 1M-token ultra-long context. Its attention module combines incremental attention (KDA) with Gated Multi-head Latent Attention (Gated MLA), and shares information across layers through an attention residual mechanism. On the MoE side, it uses a sparse architecture with hundreds of routed experts, together with specific activation functions and load balancing mechanisms to keep training stable.

Training follows a multi-stage context extension strategy:

Pretraining context extension: 8K → 64K → 256K → 1M
Post-training: million-token-level Agentic RL

This “scale up step by step” strategy lets the model fully adapt to the current sequence length at each stage, avoiding the optimization difficulties of jumping straight from short to ultra-long contexts. The three core innovations are detailed below.

Technique 1: KDA Lower-Bound Decay Parameterization

The Challenge: The Numerical Dilemma of Linear Attention

Linear attention and Delta Net-style methods accumulate a decay factor Γ within each chunk during chunked computation. Under limited precision (e.g., BF16), the inverse of the accumulated decay factor can easily overflow or underflow. As a result, diagonal tiles cannot take the standard dense matrix multiplication path and must fall back on explicit position-pair computation, which fails to fully exploit Tensor Core’s matrix multiplication acceleration and noticeably hurts throughput.

How It Works: From Unbounded to Bounded Decay

The core idea of KDA is to change the log-decay from the unbounded Negative-Softplus form to a lower-bounded Scaled Sigmoid form, fixing a decay lower bound (e.g., -5) so that the inverse scaling factor within any 16-token tile stays locked in a safe numerical range:

g_t = g_min · Sigmoid(e^{A_h z_t})
where g_min = -5, and z_t is an input-dependent gating quantity

With an explicit lower bound on the decay factor, the inverse of the accumulated intra-chunk decay no longer diverges sharply, and the numerical range stays controllable.

The Payoff: No Special Branches, Full Dense GEMM

The most direct benefit is that diagonal tiles no longer need a special computation branch: all causal tiles (diagonal and non-diagonal) can directly invoke dense Tensor Core GEMM, which greatly simplifies the kernel implementation and significantly improves throughput. In addition, this parameterization naturally supports direct extrapolation to 1M-token contexts under a NoPE (no positional encoding) setting, reducing the precision loss caused by positional encoding extrapolation.

From a reuse perspective, the change has minimal engineering invasiveness—it essentially replaces only the parameterization of the decay parameter—yet delivers gains in both numerical stability and kernel efficiency, making it a good candidate for incremental adoption in existing linear attention implementations.

Technique 2: Quantile Balancing (QB)

The Challenge: Load Balancing in Ultra-Large-Scale MoE

One of the core challenges in MoE is routing load balance. Traditional solutions mostly rely on auxiliary losses to penalize unbalanced routing, but under extreme sparsity—hundreds of experts with a very low activation ratio—auxiliary loss hyperparameters are hard to tune, and the loss can conflict with the main task loss at the gradient level, hurting both training stability and model quality.

How It Works: Dynamic Routing Correction with Quantiles

Quantile Balancing (QB) is an auxiliary-loss-free dynamic load balancing strategy. The basic idea: maintain a routing score threshold for each expert and adjust it dynamically based on the quantile of routing scores within the current batch, so that tokens are allocated to experts in real time during training and inference, driving the load toward balance. Because balancing is achieved through threshold correction rather than an additional loss term that “squeezes” the routing distribution, the main optimization objective remains unpolluted.

It can be understood as a form of “post-hoc correction” for routing:

• Measure the distribution of routing scores each expert receives in the current batch;
• Determine the quantile threshold according to the target load ratio;
• Offset routing decisions by the threshold to achieve load rebalancing.

The Payoff: Stability and Efficiency Together

QB keeps training stable under extreme sparsity, avoiding the tuning burden and objective conflicts introduced by auxiliary losses. And because balancing happens dynamically, it reduces dependence on static assumptions about expert capacity, cutting the waste caused by token drops and padding. For MoE architectures with many experts and low activation ratios, this strategy offers high reuse value.

Technique 3: Block Attention Residual (Block AttnRes)

The Challenge: Representation Bottlenecks in Deep Transformers

As networks get deeper, the information transfer efficiency of standard residual connections declines and deep networks tend to suffer representation saturation. Meanwhile, if cross-layer information sharing adopts a fully connected design, memory and communication overhead grow rapidly with the number of layers, becoming a bottleneck for both training and inference.

How It Works: Block-Wise Aggregation of Cross-Layer Information

Block Attention Residual partitions the network into layer blocks and aggregates cross-layer representations through attention within and across blocks, instead of having every layer fully interact with all preceding layers. This block-based design preserves cross-layer information sharing while constraining interaction complexity to the block level, significantly reducing memory usage and communication volume.

The Payoff: Breaking the Depth Bottleneck

Block AttnRes delivers value on two levels. At the representation level, cross-layer attention aggregation alleviates the representation bottleneck of deep networks, so that adding depth still yields gains. At the engineering level, the block structure reduces memory and communication overhead, making the practical training of deep large models more feasible. For trillion-parameter models with dozens or even hundreds of layers, this design is a key foundation for depth scaling.

How the Three Techniques Work Together

The three techniques target three different bottlenecks, but they support one another in a real system:

Technique	Problem Addressed	Layer
KDA lower-bound decay parameterization	Numerical stability and kernel efficiency of linear attention	Attention computation
Quantile Balancing (QB)	Load balancing and training stability of massive-expert MoE	MoE routing
Block Attention Residual (AttnRes)	Deep representation bottleneck and memory/communication overhead	Cross-layer information transfer

The efficient and stable attention module provides the compute foundation for long-context pretraining and million-token-level Agentic RL post-training; stable MoE routing enables effective training at trillion-parameter scale; and the cross-layer residual design keeps depth scaling from being undermined by representation bottlenecks. Together, the three support the overall training paradigm of “multi-stage context extension + long-horizon agent post-training.”

Reusable Engineering Takeaways

For teams looking to apply these ideas in their own architectures, several points deserve attention:

1. Fix numerical stability before optimizing kernels. The performance bottleneck of linear attention often comes from numerical issues before the algorithm itself. Adjustments like lower-bound decay parameterization—small changes with large gains—should be evaluated first.
2. Load balancing can bypass auxiliary losses. At very large expert scales, quantile-based dynamic threshold correction is a low-coupling alternative that reduces hyperparameter tuning and objective conflicts.
3. Control the complexity of cross-layer interaction. Block-wise cross-layer attention aggregation constrains overhead growth while preserving information sharing—a worthwhile trade-off when scaling depth.
4. Extend context in stages. Multi-stage pretraining from 8K to 1M, combined with long-horizon post-training, is more robust than stretching the context in one shot.

Overall, the core methodology behind this architecture design is to decompose the systemic challenges of scale expansion into sub-problems that can be optimized independently yet work together, and to choose solutions for each sub-problem that are minimally invasive and broadly applicable. This is the key engineering experience that makes trillion-parameter MoE long-context models practical.