Architecting Long-Context LLMs: From Sparse Attention to Stable MoE in Production

1. Background: Three Interlocking Challenges for Long-Context LLMs

As agents, long-document understanding, and multimodal applications mature, the context length demanded of large language models has grown rapidly from tens of thousands of tokens to the million-token scale. Under this trend, the conventional Transformer architecture runs into three interrelated bottlenecks.

First, attention computational complexity and numerical stability. Standard self-attention scales quadratically in both compute and memory as sequence length grows. Linear attention reduces this to linear complexity through kernel decomposition, but under limited numerical precision, the accumulation of recursive decay factors can cause underflow or overflow, making it difficult to reliably apply Tensor Core matrix multiplication acceleration. This problem is especially pronounced in 16-token block computations—the reciprocal of the decay factor can exceed the representable range of BF16, forcing implementations back to a position-by-position computation path.

Second, training stability at ultra-large MoE scale. To expand model capacity under constrained inference cost, Mixture-of-Experts architectures are widely adopted. However, when the number of routed experts reaches the hundreds or even close to a thousand, load balancing becomes the core difficulty. Traditional auxiliary-loss-based strategies converge slowly, are hard to tune, and tend to conflict with the primary task objective once expert count grows.

Third, representational bottlenecks and communication overhead in deep Transformers. As network depth increases, the risk of vanishing gradients and representational collapse rises, and cross-layer information flow becomes less efficient. Meanwhile, in long-sequence training, activation memory and communication overhead accumulate linearly with depth, constraining further model scaling.

Kimi K3, an open-source multimodal large model released in 2026, presents a complete engineering solution to these problems in its technical report. This article uses its core designs as a blueprint, analyzing them from both architectural principles and engineering implementation.

2. Numerical Stability in Linear Attention: Lower-Bounded Decay Parameterization

2.1 Root Cause

In block-wise computation for linear attention and Delta networks, each token block requires computing the reciprocal of decay factors. Let the decay factor be γ; the scaling factor within a block is the cumulative power of 1/γ. When the log value of γ is unbounded—for example, through negative Softplus parameterization—log-decay can approach negative infinity, causing 1/γ to exceed the BF16 representation limit.

Specifically, the conventional parameterization is:

γ = exp(-softplus(w))

where w is a learnable parameter. Since softplus has no upper bound on its output, the log value of γ can theoretically approach negative infinity. In a 16-token block computation, the reciprocal of γ must be raised to the 15th power; once logγ drops below approximately -5, 1/γ¹⁵ overflows under BF16.

2.2 Lower-Bounded Decay as an Engineering Solution

The core modification in KDA (Kimi Delta Attention) is to impose a hard lower bound on log-decay at the parameterization level. The unbounded negative Softplus is replaced with a bounded scaled Sigmoid:

g_th = g_min · Sigmoid(e · A · h · z / t · h)

where g_min is a fixed lower bound (taken as -5 in the report), and the Sigmoid output lies in (0,1), so the theoretical range of g_th is confined to (g_min, 0). This means the reciprocal scaling factor within any 16-token block is locked to:

1/γ^15 ≤ exp(-g_min × 15) = exp(75)

This value falls within the representable range of BF16, ensuring that diagonal block computation can stably follow the Tensor Core matrix multiplication path without degrading to explicit scalar position-pair computation.

2.3 Impact on Implementation Efficiency

The reuse value of this change lies in its extremely low engineering invasiveness—only the parameterization function of the decay factor needs modification, with no changes to the overall attention framework. But the payoff is significant: the diagonal block matrix multiplication path achieves far higher compute utilization than the scalarized position-pair path. In CUDA or CUTLASS implementations, this means highly optimized GEMM kernels can be used uniformly, eliminating warp divergence caused by numerical branching.

3. Auxiliary-Loss-Free MoE Load Balancing: Quantile Balancing

3.1 Limitations of the Auxiliary Loss Approach

Traditional MoE load balancing relies on adding an auxiliary loss term to the training objective that penalizes imbalance across experts. This approach works reasonably well with a small number of experts, but at the scale of 896 routed experts it exposes three problems:

• Tuning complexity: The auxiliary loss coefficient must be carefully traded off against the primary task loss—too large harms model quality, too small leaves severe imbalance.
• Convergence speed: The gradient signal from the auxiliary loss must propagate through the router’s soft assignment, and the feedback path lengthens as expert count grows.
• Optimization conflict: The auxiliary loss is fundamentally a different optimization direction from the primary task objective, causing persistent interference especially in late training.

3.2 How Quantile Balancing Works

Quantile Balancing (QB) takes a completely different approach: instead of penalizing at the loss function level, it intervenes directly at the routing decision level.

The core logic maintains a statistical distribution of per-expert load. When the router assigns tokens, it adjusts assignment probabilities based on the current load quantile. When an expert’s load exceeds a configured quantile threshold, its selection probability is correspondingly reduced; likewise, underloaded experts receive higher routing priority. This adjustment is immediate and deterministic, producing no additional gradient signal.

3.3 Synergy with Stable LatentMoE

The QB mechanism works in concert with the Stable LatentMoE architecture to form a complete training stability solution. Stable LatentMoE introduces a latent-space routing representation in addition to routed experts, so that routing decisions are made in a low-dimensional latent space, reducing the parameter count and computational cost of the routing network. The combined effect is:

• Load balancing no longer depends on the loss function, eliminating optimization competition between the primary task and the balancing objective.
• The dynamic adjustment mechanism adapts to data distribution shifts during training, with no need to manually tune balancing strength.
• 896 routed experts can still maintain a relatively uniform load distribution under extremely sparse activation conditions.

4. Cross-Layer Information Sharing: Block Attention Residuals

4.1 The Depth Bottleneck in Deep Transformers

Conventional Transformer residual connections pass information only between adjacent layers; information flow across multiple layers must accumulate layer by layer. When network depth reaches dozens or even hundreds of layers, deep representations are prone to degradation and information redundancy. At the same time, standard cross-layer information passing grows quadratically in compute with depth, which is unacceptable in long-sequence scenarios.

4.2 Design Motivation for Block-Level Residuals

Block AttnRes (Block Attention Residual) extends the scope of attention residual computation from a single layer to the block level. Specifically, a block containing several consecutive layers is treated as the basic unit, and all layers within the block share the same set of attention residual information. This design yields three benefits:

Representation level: Intra-block information sharing breaks the representation degradation bottleneck in deep layers, enabling deeper networks to effectively leverage information from shallow-layer attention.

Memory level: Shared attention residuals reduce the number of activations that must be stored layer by layer. In training scenarios with million-token contexts, the memory savings are significant.

Communication level: In distributed training, cross-layer information passing requires cross-device communication. Block-level sharing reduces communication rounds and lowers bandwidth pressure.

4.3 Relationship to Gated Multi-Head Latent Attention

Gated MLA (Gated Multi-Head Latent Attention) addresses attention computation within a single layer, reducing inference cost through low-rank compression of the KV cache. AttnRes, by contrast, provides an additional information pathway along the cross-layer dimension. The two optimize attention mechanisms along orthogonal dimensions: MLA compresses the cache overhead of single-layer computation, while AttnRes improves cross-layer information utilization efficiency.

5. Multi-Stage Context Extension and Agentic Reinforcement Learning

5.1 The Necessity of Progressive Context Extension

Training a model directly from 8K context to 1M context faces two difficulties: first, a mismatch in training data length distribution leads to unstable optimization; second, compute cost grows linearly with sequence length, making early iterations extremely inefficient if 1M context is used from the start.

K3 adopts a four-stage progression: 8K → 64K → 256K → 1M. Each stage uses a data distribution matched to the current context length for pretraining, and only after the model converges at that length does training switch to the next stage. This progressive strategy makes full use of compute resources at each stage while avoiding the instability of early long-context training.

5.2 Design Considerations for Million-Token Agentic RL

The reinforcement learning stage of post-training likewise requires million-token-level context support. The key difference between Agentic RL and conventional RLHF is that reward signal acquisition requires the model to complete multi-step interactions and decisions in long-horizon environments, so the token count of a single trajectory far exceeds that of conventional dialogue scenarios.

This demands the following capabilities from RL training infrastructure:

• Long-sequence inference: The rollout stage must efficiently handle generation processes at the million-token scale.
• Multi-step reward modeling: Decomposing long-horizon tasks into verifiable intermediate steps and providing dense reward signals at different positions along the trajectory.
• Training throughput guarantees: Maintaining sufficient sample throughput under long-sequence conditions to ensure the statistical efficiency of RL training.

5.3 The Complementary Role of the SiTU-GLU Activation Function

Under the dual pressure of MoE and long context, the numerical properties of activation functions directly affect training stability. SiTU-GLU introduces the SiTU activation within the Gated Linear Unit framework. Compared with Swish or GELU, it exhibits smoother saturation behavior and more stable gradient propagation characteristics, making it particularly suitable for maintaining training signal transmission under extremely sparse MoE routing conditions.

6. Open-Source Benchmark Performance and Practical Engineering Takeaways

K3’s performance on open-source large model benchmarks validates the effectiveness of the designs above. From an engineering deployment perspective, three techniques are most worth attention and adoption:

KDA lower-bounded decay parameterization: Minimal change, clear benefit. It can be applied directly in existing linear attention implementations and is especially suited to scenarios where Tensor Core utilization matters.

Quantile Balancing mechanism: For any system requiring large-scale MoE, QB offers a new path around the auxiliary loss dilemma. Its dynamic, deterministic load regulation logic is more controllable in engineering implementation.

Block attention residuals: In deep model training, block-level information sharing is an effective means of reducing memory and communication costs, with direct value especially for long-sequence training scenarios.

Together, these techniques point in one direction: the engineering core of long-context large models lies not in stacking more parameters or longer training time, but in making precise architectural adjustments to underlying problems such as numerical stability, load balancing, and cross-layer information flow. For R&D teams building or optimizing large model training infrastructure, these solutions offer validated technical paths that can be adopted independently.