Large Model Localization Adaptation: From Inference Migration to Unified Training and Inference

Background: Large Model Deployment Enters the “Deep Adaptation” Phase

As large models surpass the trillion-parameter scale, raw compute stacking alone can no longer meet deployment requirements. Whether a model can run efficiently and reliably on a specific hardware platform depends on the deep alignment between the software stack and the model’s characteristics. Recently, an open-source model with 3 trillion parameters and support for a 1M-token context window completed public adaptation, and its technical characteristics and engineering path are highly representative: the model underwent large-scale post-training and reinforcement learning, performing strongly in code generation, 3D creation, financial research, knowledge Q&A, and long-horizon agent tasks, with some code and agent benchmark scores approaching those of top closed-source models.

The value of this case lies not only in the model’s capabilities but also in the engineering methodology revealed by its “Day-0 adaptation”—completing inference and training support on the day the model is officially released. This requires an efficient collaboration mechanism among hardware vendors, software stacks, and model developers. This article dissects this process from a technical perspective and explores a general paradigm for adapting ultra-large models on domestic intelligent compute cards.

1. Understanding the Target Model’s Computational Profile

Before starting any adaptation work, the first task is to build a computational profile of the target model. For a 3-trillion-parameter sparse model, its core computational characteristics can be summarized as follows:

1. Parameter Distribution and Memory Pressure

Even at FP16 precision, 3 trillion parameters require approximately 6 TB of memory for static weights, far exceeding single-card capacity. Therefore, the model must adopt a combination of Expert Parallelism (EP) and Tensor Parallelism (TP), distributing different expert networks across multi-card nodes. This distribution means that one focus of adaptation work is the efficient implementation of communication primitives—the bandwidth utilization of cross-card All-to-All communication directly affects overall throughput.

2. KV Cache Challenges from Long Context

Supporting a 1M-token context means KV Cache memory usage grows linearly. For MoE architectures, although fewer parameters are activated, the KV Cache is still non-shared. During adaptation, one must consider pooled memory management for KV Cache, memory defragmentation, and necessary quantization compression schemes (such as INT8/INT4 KV Cache). This is not only an operator-level optimization but also involves the runtime’s memory scheduling strategy.

3. Dynamics During Inference

MoE models exhibit a clear “token routing imbalance” phenomenon during inference—certain experts are accessed at high frequency, creating compute hotspots on individual cards. Adaptation must allow the runtime to dynamically adjust expert loading strategies, or even implement an “expert pre-warming” mechanism, to avoid severe inference latency fluctuations.

2. Adaptation Core: Inference Path Migration and Operator Mapping

The essence of “Day-0 adaptation” is mapping the model’s native computational graph onto the target hardware instruction set. Taking an intelligent compute card and its MUSA software stack as an example, adaptation work follows the principle of “layered mapping, point-by-point verification.”

1. Computational Graph Parsing and Operator Lowering

First, through framework-level interception (such as PyTorch’s custom backend), the model’s computational graph is converted into a hardware-agnostic intermediate representation (IR). At this stage, high-frequency operators in the model are identified, such as:

• GEMM series, including standard GEMM, batched GEMM (for multi-head attention), and grouped GEMM (for expert routing);
• Normalization operators (LayerNorm/RMSNorm) and activation functions (SwiGLU);
• FlashAttention-style fused operators for long sequences;
• TopK and Softmax operators for routing decisions.

For each operator category, the software stack must provide hand-optimized kernels or auto-tuned templates. The key is: do not simply call CPU implementations or generic GPU kernels, but rewrite them for the target card’s SRAM capacity, register file, and thread scheduling model. For example, for grouped GEMM, the tiling strategy must be dynamically adjusted based on the expert dimension size to avoid bank conflicts and tail effects.

2. Lightweight Code Example: Operator Mapping Illustration

The following pseudocode demonstrates the process of mapping standard attention computation to a custom kernel:

# Pseudocode: Registering a custom attention kernel in the MUSA software stack
import musa as ms
@ms.custom_kernel(backend="musa")
def flash_attn_fwd(q, k, v, output, scale, causal=True):
    # Actual kernel code must be written using MUSA C/C++ extensions
    # Only the mapping logic is shown here
    bsz, heads, seq_len, head_dim = q.shape
    # Tiling based on SRAM capacity
    block_m = 128 if seq_len > 2048 else 64
    tid = ms.get_local_id(0)
    # ... core loop and data movement logic omitted
    pass
# Replace the original operator in the model graph optimization
model.graph.replace_op(operator="aten::scaled_dot_product_attention",
                       custom_kernel=flash_attn_fwd)

The challenge of this mapping process is: numerical behavior must remain consistent with the original operator (especially floating-point rounding modes), while fully utilizing the target hardware’s special instructions (such as vectorized loads and matrix fused compute units).

3. Memory Optimization for Long-Context Processing

For a 1M context, even with FlashAttention, the total KV Cache can reach hundreds of GB. During adaptation, a “memory offloading + hierarchical access” strategy was adopted: low-frequency historical KV states are stored in host-side large-capacity memory, while only the currently active window’s KV states reside in device memory. This requires the runtime to have precise page scheduling prediction capability—prefetching KV blocks about to be accessed by analyzing attention patterns. In implementation, asynchronous copy engines are typically used to overlap data transfer with computation, masking PCIe or NVLink access latency.

3. Architectural Considerations for Unified Training and Inference

The case mentions “AI training-inference unified intelligent compute cards,” which reveals an important trend in modern AI infrastructure: training and inference are not isolated phases but share the same underlying resource pool.

1. Resource Orchestration for Unified Training and Inference

During training, model parameters need frequent updates; during inference, they are relatively static. A unified card must support rapid switching between the two modes. At the software level, this manifests as fine-grained management of “Streams” and “Contexts.” For example, training tasks can use high-occupancy, long-duration job scheduling; inference tasks require low-latency, high-concurrency priority inversion mechanisms.

2. Mixed Precision and Numerical Stability

Training typically uses BF16 mixed precision, while inference may use FP8 or INT8 for speed. The same card must handle switching between precision modes gracefully. During adaptation, special attention was paid to scaling factor management under the FP8 format—because the output distributions of different experts in an MoE model differ significantly, scaling factors need to be set individually for each expert rather than shared across the entire model. This detail directly affects the model’s accuracy loss after quantization.

4. Rigor in Performance Validation and Benchmarking

An objective evaluation system must be established around adaptation effectiveness. The model scored 65.7% on Terminal-Bench 4.0 and 61.2% on long-horizon tasks. This requires the adaptation team to focus not only on throughput (Tokens/s) but also on end-to-end task success rates.

1. Layered Evaluation Dimensions

• Low-level metrics: Model FLOPs Utilization (MFU), memory bandwidth utilization, communication ratio;
• Model metrics: Perplexity (PPL) drift, accuracy difference before and after quantization;
• Business metrics: Benchmark scores, long-task completion success rate, average latency jitter (P99).

2. Reliability and Long-Duration Stability Testing

“Ghost failures” after ultra-large model deployment (such as performance degradation caused by ECC correction) are often difficult to trace. Day-0 adaptation must include at least 72 hours of long-duration stress testing, with monitoring items including but not limited to: GPU temperature curves, memory RAS events, and communication timeout retransmission rates. Chaos engineering is used to inject certain anomalies to verify the effectiveness of runtime recovery mechanisms.

5. Key Path from Adaptation to Scaled Deployment

Day-0 adaptation is only the starting point. To achieve scaled deployment in production environments, the following issues must also be addressed:

1. Multi-tenant Isolation: When multiple inference tasks share the same physical intelligent compute card, how can transparency of memory and compute isolation be guaranteed? Current approaches favor leveraging MUSA’s MPS (Multi-Process Service)-like capabilities to partition physical devices into multiple compute instances, each with independent memory address spaces and stream priorities.

2. Automated Tuning Toolchain: Manually writing operators requires significant human effort. It is recommended to build a “performance profiler” tool that automatically scans model hotspots and suggests operator replacements. For example, if a GEMM operation is detected to account for more than 30% of execution time, the tool automatically recommends using a grouped matrix multiplication kernel.

3. Integration with the Cloud-Native Ecosystem: In Kubernetes clusters, intelligent card resources are exposed to containers through the Device Plugin mechanism, combined with custom schedulers to achieve “topology-aware” Pod distribution, ensuring minimal inter-node latency for tensor parallel communication.

Conclusion

The completion of adaptation for a 3-trillion-parameter model is a microcosm of domestic AI compute maturing. Its value lies not in the performance numbers of a single piece of hardware, but in validating a “model-agnostic, hardware-adaptive” software stack methodology. For developers, understanding the operator mapping logic, memory scheduling strategies, and evaluation systems in this process is far more important than chasing specific model numbers. In the future, as model scales continue to expand and inference costs become more sensitive, the ability to perform software-hardware co-design starting from model structure will become a core competency for AI engineering teams.