Architecture Deep Dive: Building a Trillion-Parameter Sparse MoE with Million-Token Context

    This article explores the architectural trade-offs in building a trillion-parameter, ultra-sparse Mixture-of-Experts (MoE) multimodal large model capable of handling a one-million-token context window. It breaks down three core technical paradigms: lower-bounded decay parameterization for linear attention (KDA) to maximize Tensor Core utilization, quantile balancing (QB) for auxiliary-loss-free expert load distribution, and block attention residuals (Block AttnRes) to overcome deep representation saturation. Together, these approaches solve the computational, stability, and depth challenges of extreme-scale models.

    Architecture Deep Dive: Building a Trillion-Parameter Sparse MoE with Million-Token Context

    When a model’s total parameters reach the trillion scale, its context window expands to a million tokens, and its expert count approaches a thousand, the foundational engineering assumptions of traditional dense Transformers break down. The computational and memory overhead of attention grows quadratically with sequence length, MoE routing frequently collapses under extreme sparsity, and the representational capacity of deep networks gradually saturates through residual stacking. This article examines a sparse MoE multimodal large model designed for million-token contexts, systematically analyzing its architectural trade-offs across three main areas: attention, MoE routing, and residual connections. We also extract three high-value technical paradigms for developers.

    1. Overall Architecture: Design Constraints Under Three Extremes

    The model’s scale can be defined by three extremes:

    • Extreme parameter sparsity: With approximately 2.8 trillion total parameters and around 104 billion activated parameters per forward pass, the activation ratio is under 4%, making it a typical ultra-sparse MoE.
    • Extreme context length: It natively supports 1 million tokens, progressively scaling during pre-training in four stages: 8K → 64K → 256K → 1M.
    • Native modality fusion: Visual inputs are not patched via an external encoder but are integrated into the unified modeling from the pre-training phase.

    The combination of these three factors means no subsystem can be optimized in isolation. If the attention layer cannot maintain high throughput for long sequences, a million-token context is unachievable. If MoE routing unbalances with a thousand experts, trillion-parameter capabilities degrade into noise. If residual paths cannot share information across layers, deep stacking yields diminishing returns. Therefore, the core design philosophy is: every component must simultaneously serve both “efficiency” and “stability.”

    On the attention side, the model fuses incremental linear attention (KDA) with Gated Multi-head Latent Attention (Gated MLA) within a single module, balancing long-range decay modeling with global retrieval. On the MoE side, it employs a Stable LatentMoE architecture with 896 routed experts, paired with a SiTU-GLU activation function and Quantile Balancing (QB) mechanism. For cross-layer connections, it introduces Block Attention Residuals (Block AttnRes) to replace naive residual accumulation. Let’s break these down.

    2. KDA Lower-Bound Decay Parameterization: Running Linear Attention at Full Tensor Core Capacity

    The Problem: Numerical Overflow in Diagonal Blocks Forces Special Branches

    Linear attention and Delta Net-like structures typically use chunked computation in their engineering implementation: the sequence is split into tiles of length 16, performing local attention with causal masking within the chunk, and passing recursive states between chunks. The issue lies in the cumulative decay factor Γ within the chunk. In traditional implementations, the logarithmic form of the decay often uses unbounded negative Softplus, causing the inverse of Γ to severely overflow or underflow in BF16 precision. To prevent numerical explosions, implementations are forced to route diagonal blocks through an explicit “position-pair” path. This path cannot invoke dense matrix multiplication and degrades into element-wise operations.

    As a result, non-diagonal blocks run at full Tensor Core speed, but diagonal blocks become a throughput bottleneck. This bottleneck scales linearly with sequence length—making it unacceptable for million-token contexts.

    The Mechanism: Replacing Unbounded Decay with a Lower-Bounded Scaled Sigmoid

    KDA’s solution is straightforward: replace the unbounded negative Softplus log-decay parameterization with a lower-bounded scaled Sigmoid. By fixing a decay lower bound (e.g., -5), the decay factor satisfies:

    1. g_t = g_min * sigmoid(e_A(h_z, t)) # g_min is the fixed lower bound
    2. Γ_ij = exp(Σ g) # Cumulative decay, its inverse is locked in a safe numerical range

    The key benefit here is that the inverse scaling factor within any 16-token tile is constrained to a range safely representable by BF16. Diagonal blocks no longer require a special computational branch. All causal tiles—whether diagonal or non-diagonal—can uniformly invoke dense Tensor Core GEMM operations.

    The Payoff: Dual Dividends in Throughput and Length Extrapolation

    This modification has minimal engineering intrusion but yields an order-of-magnitude improvement in operator efficiency. Eliminating the branching means the kernel can be fully vectorized, intra-warp divergence disappears, and memory access patterns are unified. A more subtle benefit is that because the decay is explicitly parameterized and numerically stable, the model inherently supports direct extrapolation to million-token contexts without Positional Encoding (NoPE). Positional information is implicitly carried by the decay structure rather than relying on explicit Rotary Position Embedding interpolation.

    For developers, this is a highly reusable paradigm: any chunked recursive linear attention implementation can check if its decay parameterization is bounded, thereby determining whether the special path for diagonal blocks can be eliminated.

    3. Quantile Balancing (QB): Auxiliary-Loss-Free Load Balancing for Thousand-Expert Models

    The Problem: Contradiction Between Auxiliary Loss and Extreme Sparsity

    The classic challenge of ultra-sparse MoE is routing collapse: a few experts are over-selected while most remain idle, especially during early training. A common industry practice is adding an auxiliary load-balancing loss. However, an auxiliary loss essentially imposes a soft constraint on top of the main task objective, which has two side effects. First, the weight is difficult to tune—too strong damages language modeling quality, too weak fails to balance. Second, it only provides a gradient signal and does not guarantee actual allocation ratios.

    When the number of experts reaches 896, the problem is amplified: the routing distribution has a very long tail, and the adjustment granularity of the auxiliary loss is insufficient.

    The Mechanism: Dynamic Threshold Calibration via Quantiles

    Quantile Balancing (QB) takes a different approach. Instead of soft guidance at the gradient level, it performs hard calibration at the routing decision level. Its core maintains selection frequency statistics for each expert and dynamically adjusts each expert’s selection threshold based on a target load ratio (e.g., global uniformity). Essentially, during each routing step, it uses historical quantiles to decide if a current token crosses an expert’s threshold, forcing the actual allocation ratio to approach the target distribution.

    This can be understood as an “adaptive Top-K truncation.” The threshold is not fixed but slides based on load statistics, achieving rapid dynamic balancing without introducing auxiliary losses. The significance for training stability is that the main loss function remains pure, and balancing is guaranteed by the routing mechanism itself, avoiding objective tug-of-war between the auxiliary loss and the main task.

    The Payoff: Synchronized Improvement in Training Stability and Expert Utilization

    Under extreme sparsity (activation ratio under 4%), the QB mechanism ensures no large-scale expert idling throughout training. Gradients cover all 896 experts, significantly improving parameter utilization. Additionally, because there is no auxiliary loss term, the hyperparameter search space is drastically reduced, lowering tuning costs during early training. This provides tangible engineering benefits for large-scale pre-training.

    4. Block Attention Residuals (AttnRes): Breaking the Deep Representation Bottleneck

    The Problem: Representation Saturation in Naive Residuals

    Residuals in deep Transformers are accumulated layer by layer: the input to layer L is the sum of the outputs of all preceding layers. As depth increases, this cumulative residual causes representations to gradually “homogenize.” New information added by each layer is drowned in the cumulative sum, and the effective representational capacity of deep networks tends to saturate. Concurrently, layer-by-layer accumulation means each layer must read the full residual stream, causing memory bandwidth and cross-device communication overhead to grow linearly with depth.

    The Mechanism: Cross-Layer Attention Aggregation at the Block Level

    Block Attention Residuals (Block AttnRes) divides the network into blocks of several layers. Standard residuals are maintained within a block, but between blocks, an attention mechanism performs weighted aggregation of historical block representations instead of simple addition. Cross-layer information sharing shifts from “indiscriminate accumulation” to “selective retrieval”: each block can attend to preceding blocks as needed, rather than passively receiving the entire cumulative sum.

    The engineering benefits are twofold. First, at the representation level, deep layers are no longer drowned out by shallow information, maintaining the marginal returns of deep stacking. Second, at the overhead level, block-grained aggregation reduces the size of the residual stream that needs to reside in memory and be transmitted, lowering memory footprint and cross-device communication volume. This is critical for multi-card training of trillion-parameter models.

    5. Multi-Stage Pre-training and Agentic Post-Training: Operationalizing Capabilities

    Beyond architecture, the training pipeline is equally noteworthy. Pre-training adopts a multi-stage context progression strategy (8K → 64K → 256K → 1M), gradually extending sequences so the model fully converges across different length scales, avoiding the instability of jumping directly to long-sequence training. This is followed by Agentic Reinforcement Learning (RL) post-training at the million-token level, aligning long-context capabilities with agentic behaviors like multi-step tool calling and long-range planning.

    The design logic is clear: a million-token context is merely “capacity.” To transform capacity into “capability,” long-range tasks must activate it during post-training. Agentic RL fulfills this role—it trains the model to maintain state consistency over ultra-long trajectories, allocate attention appropriately, and recall early information even after multiple rounds of interaction.

    6. Reusable Technology Checklist and Adoption Recommendations

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

    1. KDA Lower-Bound Decay Parameterization (Reuse Value: High). Minimal engineering changes are required—just adjust the boundary form of the decay parameterization to eliminate the special branch for linear attention diagonal blocks. This allows all causal tiles to invoke dense GEMM while gaining NoPE long-context extrapolation. Suitable for any team developing in-house linear or hybrid attention.
    2. Quantile Balancing QB (Reuse Value: High). Auxiliary-loss-free dynamic load balancing acts directly on routing decisions rather than the loss function. It suits MoE training scenarios with large expert counts and high sparsity, significantly reducing the tuning burden.
    3. Block Attention Residuals AttnRes (Reuse Value: Medium-High). Requires modifying network topology and the training framework, involving substantial changes. However, it provides clear benefits for representation preservation and communication optimization in ultra-deep, ultra-large models, making it suitable for teams with deep customization capabilities.

    It is crucial to emphasize that these three technologies are not isolated: KDA solves “computability,” QB solves “training stability,” and AttnRes solves “learning depth.” Together, they form the engineering foundation that makes a trillion-parameter sparse MoE trainable, convergent, and deployable under million-token contexts. For most teams, we recommend adopting them progressively in the order of “attention operators first, then MoE routing, and finally cross-layer structures.” Each step should be validated with measurable throughput and convergence metrics rather than executing a one-time rewrite.