Architecture Design and Engineering Practices for Trillion-Scale Sparse LLMs

I. Why Trillion-Scale Sparse Models Are Necessary

The scaling law of large language models reveals a fundamental truth: the coordinated growth of parameter scale, data volume, and computation brings predictable capability improvements. However, when dense model parameters approach the hundreds of billions, the cost curve for training and inference becomes unsustainable—every forward pass requires activating all parameters, causing inference latency and memory consumption to scale linearly with parameter count.

The Mixture-of-Experts (MoE) architecture offers an alternative path: exchanging sparse activation for capacity expansion. The total model parameters can reach the trillion scale, but a single token only passes through a few expert networks during each forward step, making the actual computation (FLOPs) far lower than that of a dense model with an equivalent parameter count. This partial decoupling of “parameter count” and “computation cost” has made MoE the mainstream technical route for ultra-large-scale models.

To understand these models, one must look beyond the single metric of “total parameters” and break it down into three levels: total parameter capacity, activated parameter scale, and the engineering complexity introduced by routing and parallel strategies. The following sections expand on these three aspects.

II. Overall Architecture: Base Models and Systematic Delivery

From an engineering perspective, a trillion-scale model product is typically not a single weight file, but a systematically layered delivery system:

• Base Model: A general foundation obtained through large-scale pre-training, possessing complete language modeling and knowledge storage capabilities. It is the starting point for all subsequent capabilities.
• Alignment and Specialization Layer: Built upon the Base model through supervised fine-tuning and preference optimization to form a usable configuration for scenarios like dialogue, tool calling, and long-context processing.
• Inference and Serving Layer: Includes parallel partitioning strategies, batch scheduling, KV Cache management, quantization, and operator optimization. This layer determines whether the model can run stably and cost-effectively in real-world applications.

The value of this layered design lies in the separation of concerns: the pre-training phase pursues “capacity and knowledge density,” the alignment phase pursues “controllable behavior,” and the serving phase pursues “throughput and latency.” Mixing these optimization targets leads to conflicting engineering decisions.

III. Core Mechanisms of Sparse Activation

3.1 Basic Flow of Expert Routing

The MoE layer replaces the traditional single Feed-Forward Network (FFN) with multiple parallel expert FFNs and introduces a router/gate network to select the Top-K experts for each token:

# Simplified Top-K routing illustration (pseudo-code)
def moe_forward(x, experts, router, k=2):
    logits = router(x)                 # [num_tokens, num_experts]
    probs = softmax(logits, dim=-1)
    topk_probs, topk_idx = topk(probs, k)   # Select k experts per token
    topk_probs = topk_probs / topk_probs.sum(-1, keepdim=True)  # Normalize
    out = zeros_like(x)
    for i in range(k):
        out += topk_probs[:, i] * experts[topk_idx[:, i]](x)
    return out

Key design parameters include: total number of experts, number of activated experts per token (K), and the expert capacity factor. The total number of experts determines the capacity limit, K determines the computation per token, and the capacity factor determines the maximum number of tokens an expert can process per batch—the latter directly impacts the token drop rate and training stability.

3.2 Load Balancing: The Most Underestimated Challenge in Sparse Models

Routing is “learned,” making it naturally susceptible to the Matthew Effect: selected experts receive more gradients, perform better, and are more likely to be selected again. Without intervention, a few experts will handle the vast majority of tokens, while the rest remain chronically under-trained, leading to “expert collapse.”

Common countermeasures include:

1. Auxiliary Load Balancing Loss: Adding a variance penalty term for expert usage frequency to the training loss to encourage uniform distribution.
2. Capacity Constraints and Token Dropping: Setting a capacity limit for each expert. Excess tokens are dropped or passed directly through the residual connection to prevent single-expert overload.
3. Router Noise and Randomness Injection: Adding noise to the logits during training to increase exploration and mitigate early routing solidification.
4. Communication Optimization in Expert Parallelism: Experts are distributed across different devices, requiring All-to-All communication for token dispatch. Load imbalance directly amplifies communication hotspots.

These measures require joint tuning: an excessive load balancing loss weight degrades model quality, while too little loses its constraining effect; a capacity factor that is too small leads to massive token drops, while too large wastes memory and compute. In practice, intervals are typically determined through small-scale ablation experiments, and expert utilization distributions are monitored during large-scale training.

3.3 Training Stability of Sparse Models

Training stability for trillion-scale MoE models presents far greater challenges than dense models, primarily stemming from:

• Routing Discontinuity: Top-K selection is a discrete operation. Gradients can only backpropagate through the weighted paths of selected experts, and the learning signal for the router itself is weak.
• Numerical Precision Issues: Weighted summation of expert outputs and reduction operations during All-to-All communication easily accumulate errors under low precision. High-precision accumulation is often required on critical paths.
• Consistency in Expert Parallelism: When overlapping cross-device communication and computation, gradient synchronization correctness must be ensured to prevent routing drift caused by asynchronous operations.

Common practices include: computing router logits with higher precision, clipping gating probabilities, adopting more conservative learning rates and lower capacity pressure early in training, and gradually relaxing them once the routing distribution stabilizes.

IV. Parallel Strategies: Fitting Trillion Parameters into a Cluster

Trillion-scale parameters cannot fit into a single machine, necessitating the combination of multiple parallel methods:

Parallel Method	Partitioning Target	Primary Problem Solved
Data Parallelism	Batch	Increases throughput; requires global gradient synchronization
Tensor Parallelism	Weight matrices within a single layer	A single layer is too large for one GPU
Pipeline Parallelism	Partitioned by layers	Reduces per-GPU memory footprint; requires handling pipeline bubbles
Expert Parallelism	Different experts distributed across devices	Unique to MoE; accompanied by All-to-All communication

Expert parallelism is a key component unique to MoE: since each token only accesses a few experts, tokens must be “redispatched” across devices to where the target experts reside, and returned after computation. This introduces two engineering imperatives:

1. Communication-Computation Overlap: Pipelining All-to-All communication with local expert computation to hide latency.
2. Grouped Routing (Grouped GEMM): Batching tokens routed to the same expert and using Grouped Matrix Multiplication to improve operator efficiency, avoiding the overhead of per-token scheduling.

A typical training step can be abstracted as: routing computation -> token reshuffling and All-to-All dispatch -> local grouped expert computation -> result reduction and return -> residual and layer normalization. The bottleneck of this entire chain is often not floating-point compute power, but communication bandwidth and scheduling efficiency.

V. Inference: Sparse Does Not Mean Cheap

A common misconception is that “fewer activated parameters mean inference is necessarily cheap.” In reality, the cost structure of MoE inference is more complex:

• Memory Footprint is Dictated by Total Parameters: All experts must reside in memory or be swapped in on-demand. The memory pressure is far greater than that of a dense model with an equivalent number of activated parameters.
• Sensitivity to Batch Efficiency: Under small batch sizes, expert utilization is low—many experts are triggered by few tokens, leading to obvious compute waste. Only larger batches can amortize routing and communication overhead.
• Routing Overhead: Gating computation, token reshuffling, and cross-device communication account for a rising proportion of the total time in latency-sensitive scenarios.

Therefore, inference optimization focuses on: hierarchical caching and on-demand loading of expert weights, MoE-oriented continuous batching scheduling, and operator fusion for grouped GEMMs. In latency-sensitive scenarios, trade-offs must be made between “activating more experts for better quality” and “reducing expert switching to lower costs.”

VI. The Significance of Naming: From Symbols to Technical Propositions

Model naming often carries dual information regarding technical routes and product positioning. Breaking it down:

• The Numeric Hierarchy in “K2”: Marking generations with version numbers signals to users that “this is an architectural leap” rather than a minor patch. Numeric naming reduces communication costs and facilitates unified references in documentation, evaluations, and ecosystems.
• The Semantics of “K”: It can be understood as Knowledge (knowledge capacity) or Key (entry point for key capabilities). This ambiguity leaves narrative space—emphasizing capability density over a single metric.
• “Kimi” as a Product Name: Decoupled from the underlying model name, it forms a two-tier structure of “brand layer” and “model layer.” The brand layer communicates user experience and scenarios, while the model layer communicates architecture and capability boundaries to developers. Decoupling allows both to evolve independently.
• Base vs. Non-Base Distinction: Clarifying the boundary between the “base model” and the “aligned model” reflects technical transparency—it tells developers which part is the general capability foundation and which is scenario-specific processing, enabling correct choices during secondary development.

The essence of naming is reducing comprehension costs: a good name should allow the listener to establish a preliminary expectation of “scale + technical route + capability positioning” within a single sentence. From this perspective, the combination of numbers, letters, and product names is not arbitrary, but a compressed expression of technical propositions.

VII. Engineering Practice Recommendations

For teams planning to use or adopt similar architectures, the following points can be extracted:

1. Clarify Optimization Goals First: Are you pursuing the upper limit of capability per unit of compute, or the lower limit of cost per unit request? The former leans towards larger total parameters and higher sparsity; the latter leans towards smaller activated scales and stronger quantization.
2. Treat Load Balancing as a First-Class Citizen: From day one of training, monitor expert utilization, token drop rates, and communication hotspots. Build visualization dashboards rather than fixing issues retroactively.
3. Validate Parallel Configurations at Small Scale: Use a small model on the target cluster to validate combinations of tensor parallelism, expert parallelism, and pipeline partitioning. Measure communication proportions before scaling up.
4. Co-design Batching and Caching for Inference: The cost-effectiveness of MoE highly depends on batch size. Scheduling strategies should match business request distributions to avoid the gap of “advanced architecture, underutilized in production.”
5. Maintain Restraint in Precision Strategies: Low precision significantly reduces costs, but routing and reduction paths should retain sufficient precision. Stability takes priority over extreme compression.

Conclusion

The value of trillion-scale sparse large models lies not in the parameter number itself, but in the engineering compromise it makes between the conflicting goals of “capacity expansion” and “computation cost.” What truly determines deployment effectiveness lies in the “unsexy” details like routing balance, parallel communication, and inference scheduling. Understanding the trade-off logic behind the architecture holds more long-term value than memorizing any single metric.