Engineering Linear Attention and Large-Scale MoE: Practical Optimizations

As context windows grow from hundreds of thousands toward the million-token scale, and as model parameters shift from dense architectures to very large mixture-of-experts (MoE) designs, the engineering efficiency of low-level operators and training paradigms is becoming a key variable in whether a model can actually be deployed. Rather than restating the promotional claims of any specific model, this article breaks down three highly reusable technical approaches from an engineering perspective: lower-bounded decay parameterization for linear attention, dynamic load balancing for MoE, and block attention residuals in deep networks. These correspond to three different bottleneck layers—operator, routing, and representation—but they share a common direction: under constrained memory, bandwidth, and numerical precision, push as much computation as possible into high-throughput dense matrix multiplication paths.

Why Linear Attention Looks Fast on Paper but Runs Slow in Practice

The theoretical advantage of linear attention and Delta Net-style architectures is reducing sequence modeling complexity from quadratic to linear. But at the actual GPU execution level, one frequently underestimated detail is that cumulative decay factors inside chunked computation break numerical stability.

Specifically, when computing intra-chunk contributions, chunked linear attention needs to take the reciprocal of the cumulative decay factor Γ for normalization or rescaling. When decay strength is large and sequences are long, the reciprocal of Γ easily overflows or underflows under limited precision such as BF16. Once the numerics become uncontrollable, compilers and operator libraries can no longer treat the diagonal block—the part where queries and keys fall within the same chunk—as an ordinary dense matrix multiplication, and must fall back to an explicit position-pair element-wise path.

This fallback is expensive: the diagonal block cannot use Tensor Core dense GEMM and can only go through inefficient element-wise or small-scale reduction paths, producing the classic pattern of “linear in theory, slowed by the diagonal block in practice.” Worse, this special-case branch forces the diagonal and non-diagonal blocks under a causal mask onto two different computation graphs, increasing kernel implementation complexity and tuning cost.

The Lower-Bounded Decay Parameterization Approach

The root cause is the unbounded nature of the decay parameter. Conventional approaches use an unbounded negative Softplus mapping for log-decay, so the output range can in theory be arbitrarily small and the reciprocal arbitrarily large. The core change in lower-bounded decay parameterization is very restrained: replace the unbounded negative Softplus mapping of log-decay with a lower-bounded scaled Sigmoid, namely

g = g_min * sigmoid(e + A * h + z)

with a fixed lower bound of g_min = -5. The engineering significance is that the decay factor is locked into a known numerical interval, so within any fixed-length chunk (for example, 16 tokens), the reciprocal scaling factor stays in a safe range and does not trigger precision overflow.

The resulting chain of benefits is substantive:

• Elimination of the diagonal-block special case: diagonal and non-diagonal blocks can uniformly use dense Tensor Core GEMM, causal chunking no longer needs two paths, and kernel branching is greatly simplified;
• Throughput improvement: the diagonal portion previously slowed by the element-wise path returns to matrix multiplication, raising overall operator utilization;
• Natural fit for NoPE extrapolation: because decay behavior is bounded and predictable, numerical behavior during long-context extrapolation is more stable, making it easier to extend context directly to the million-token scale without relying on explicit positional encoding.

The reuse value of this technique lies in its minimal engineering change and broad benefit surface: it does not alter the model’s main structure, only adjusts one parameterization mapping, yet simultaneously improves operator implementation efficiency and long-sequence numerical stability—a typical high-return optimization.

Load Balancing for Very Large MoE: From Auxiliary Loss to Quantile Balancing

One of MoE’s core engineering challenges is routing collapse: if a few experts are over-allocated tokens while the rest sit idle for long periods, serious compute waste and communication hotspots follow. The common industry approach is to introduce an auxiliary loss to penalize unbalanced allocation, but the auxiliary loss itself interferes with the main-task gradient, requires careful tuning, and at large-scale training tends to pull against the primary objective.

Quantile balancing offers a different path: it does not rely on an auxiliary loss, but instead achieves dynamic balance by splitting routing scores at quantiles. The intuition is that rather than “softly guiding” the routing distribution with a penalty term, you directly determine each expert’s capacity boundary by quantile at the allocation stage, making token allocation approximately uniform in a statistical sense.

The engineering benefits appear in several areas:

1. No tuning burden from auxiliary loss: the training objective is cleaner and the main-task gradient is not polluted by an extra term;
2. Dynamic adaptation: quantile thresholds can adjust as the routing distribution changes during training, avoiding the mismatch of fixed capacity factors at different training stages;
3. Fit for very large scale: when the number of experts is huge and a single GPU cannot hold all experts, balanced allocation directly determines how regular cross-device communication is—the more uniform the allocation, the smoother the all-to-all load.

It is worth emphasizing that deploying such a strategy depends on stable statistics of routing scores and low-overhead threshold estimation. Implementation must avoid introducing extra synchronization points on the critical path, otherwise the gains from balancing will be offset by synchronization overhead.

Representation Bottlenecks in Deep Transformers and Block Attention Residuals

As model depth keeps increasing, a frequently discussed but hard-to-quantify problem emerges: representations in deep networks tend toward homogeneity, the difference between shallow and deep hidden states shrinks, and the marginal benefit of stacking more layers declines. At the same time, standard attention at depth must store large amounts of intermediate activations, and memory and cross-device communication costs grow linearly with depth.

The idea behind block attention residuals is to change residual connections from “layer-by-layer direct connection” to “aggregate by block, then inject.” The structural intuition is to treat several consecutive layers as a block, keep conventional computation within the block, and pass information across blocks through attention-style residual aggregation. The effect is two-sided:

• At the representation level: cross-block information pathways are explicitly modeled, alleviating deep representation homogeneity and preserving distinguishability across depths;
• At the system level: residual aggregation happens per block, reducing the number of intermediate activations that must reside long-term, and lowering the frequency and bandwidth pressure of cross-device residual transfer.

Such changes place high demands on the training framework: the granularity of residual aggregation must align with how pipeline parallelism is partitioned, otherwise misalignment between block boundaries and pipeline-stage boundaries introduces extra waiting. In practice, block size and the number of pipeline stages usually need to be co-designed.

The Shared Engineering Philosophy of These Three Techniques

Viewed together, these three techniques reveal a clear throughline: all of them try to eliminate “special paths” so that computation falls as much as possible into regular, dense channels that hardware executes efficiently.

• Lower-bounded decay parameterization eliminates the special computation branch of the diagonal block, letting causal attention uniformly use dense GEMM;
• Quantile balancing eliminates the auxiliary loss as an extra optimization objective, returning routing allocation to a direct statistical decision;
• Block attention residuals eliminate the activation expansion and communication fragmentation caused by layer-by-layer residuals, making cross-layer information transfer regular at the block level.

For developers, the lesson from these techniques is not to copy specific parameters but to establish a judgment criterion: when an optimization is more elegant in theory yet requires many special cases and branches in implementation, first consider whether parameterization or structural constraints can eliminate the special cases. Numerical stability, load regularity, and representation distinguishability are essentially three facets of the same engineering goal—keep hardware fully utilized and training predictable.

As million-token contexts and very large MoE become the norm, the value of these “small change, large return” low-level optimizations often exceeds fine-tuning at the model-structure level. Understanding their mechanisms and boundary conditions has more long-term value than memorizing any specific configuration.