Day-0 Adaptation of Trillion-Parameter LLMs on Domestic AI Accelerators

Introduction

As LLM parameter counts surpass the trillion mark, successfully running the latest open-source weights on emerging AI accelerators has become a core challenge for computing deployment. Recently, a general-purpose AI accelerator and its accompanying software stack successfully completed a Day-0 adaptation for a 3-trillion-parameter model, achieving a closed-loop pipeline from model loading to multimodal inference. This article strips away specific brand details to reveal the key optimization paths for hardware-software co-design from a purely technical perspective.

1. Core Challenges in Ultra-Large Model Adaptation

When model parameters scale to 3 trillion, single-card inference becomes entirely unfeasible, making multi-card, multi-node cluster collaboration the only viable option. The goal of Day-0 adaptation is to quickly validate the basic inference pipeline for newly released open-source weights on a specific hardware architecture. This process faces two major challenges: memory management pressure caused by long contexts and sparse architectures, and the native support for new low-precision formats like MXFP4 at the underlying operator level.

In this practice, the target model supports a context length of one million tokens and utilizes sparse attention mechanisms alongside MXFP4 format expert weights. The adaptation team not only had to solve the massive KV cache storage problem but also needed to reconstruct the underlying computation paths to accommodate non-standard precision formats and dynamic routing logic.

2. Building a Compatibility Compilation and Operator Mapping System

To efficiently run standard open-source model code on specialized hardware, the primary task is establishing a highly compatible compilation and operator mapping system.

1. Full-Stack Compilation and Operator Ecosystem
The adaptation team leveraged open-source inference frameworks to inherit upstream model implementations. However, the underlying computations did not directly reuse community code; instead, they were reconstructed using a proprietary compiler and multi-layered operator libraries. Specifically, the team utilized underlying operator libraries, Triton-based custom operator libraries, low-level tensor description languages, and a specialized compute compiler. This combination lowers the high-level Python computation graph into efficient machine instructions executable by the hardware, completing the transformation from community code to a practical inference service.

2. Integration Logic for Sparse Attention Operators
In actual engineering, developers must rewrite the Attention mechanism to adapt to the target hardware. The following pseudocode illustrates the logical abstraction when integrating sparse attention into the underlying operator library, with annotations for the framework adaptation layer:

# Adaptation layer: mapped to hardware SDK via torch.compile backend
def custom_sparse_attention(query, key, value, block_size):
    # 1. Index selection: calculate Mask based on Block sparse matrix (corresponds to hardware Sparse Select operator)
    sparse_mask = select_topk_blocks(query, key, block_size)
    # 2. Call hardware sparse computation operator (e.g., MATE backend's sparse_sdpa interface)
    # Avoid traditional SDPA full computation, reducing complexity from O(n²) to O(n·k)
    attn_output = hardware_mate_sparse_op(
        query, key, value, sparse_mask
    )
    return attn_output

3. Restructuring KV Cache Layout and Deep Optimization of Sparse Attention

For sparse attention mechanisms, the software stack must be deeply customized to ensure inference throughput is not bottlenecked by long contexts.

1. PagedAttention-Style Block Management to Prevent OOM
The memory bottleneck for long-text inference often lies in the KV cache. Traditional continuous allocation strategies easily trigger OOM when facing one million tokens. This adaptation borrowed the paged management concept from PagedAttention, adopting a non-continuous block allocation method: dividing the KV cache into fixed-size blocks (e.g., 128 tokens per block) and dynamically allocating physical memory pages on demand. This strategy can reduce memory fragmentation by over 60%, effectively preventing OOM during the prefill stage. Combined with the index selection logic of sparse attention, the system only loads the KV blocks required for the current token computation during both Prefill and Decode stages, fundamentally avoiding the bandwidth loss caused by full access.

2. Binding KV Cache State Management for Full-Chain Integration
Optimizing a single operator is insufficient to support the entire model. This adaptation deeply bound the sparse attention operator with the KV Cache state management mechanism, jointly tuning cross-layer residual information transmission and Attention computation for visual modalities. This integration from single-point operators to the full chain ensures state consistency across multiple cards for multimodal inputs like text and images, maintaining precise tensor and pipeline parallelism even when the model is split across different accelerators.

4. Online Conversion and Computation Strategy for MXFP4 Expert Weights

The Mixture of Experts (MoE) architecture is key to efficient inference for trillion-parameter models. The target model uses MXFP4 (4-bit floating-point) format expert weights, placing special requirements on underlying matrix multiply-accumulate units.

1. Engineering Trade-offs: Abandoning Offline Conversion
Some industry solutions uniformly convert weights to FP8 or INT8 formats before loading, known as “offline conversion.” However, facing TB-scale weights, offline conversion is not only time-consuming but also requires massive additional storage, severely contradicting the Day-0 rapid adaptation goal. It is also crucial to configure an OOM protection threshold and data verification mechanism when loading TB-scale weights to prevent weight corruption caused by loading interruptions.

2. Online Conversion and Cross-Layer Workspace Reuse
The adaptation team designed an “online conversion” mechanism: at the dynamic nodes of expert computation, the system unpacks MXFP4 weights in real-time and converts them to hardware-supported FP8 format for calculation. It is necessary to specifically evaluate the impact of FP8 de-quantization on inference accuracy, and it is recommended to use a calibration dataset for precision verification.

A particularly clever strategy is “cross-layer FP8 workspace reuse”: since different layers in a MoE architecture activate different experts at different time steps, the system allocates only a single fixed-size FP8 workspace in memory (calculated by the maximum number of activated experts × single expert weight size). When computation flows from the current layer to the next, this workspace is overwritten in-place, completely eliminating the need to pre-generate full offline FP8 weights and drastically reducing model cold start time.

Expert Weight Online Computation Optimization Logic:

[Startup Phase]
  -> Directly load native MXFP4 weights to VRAM (zero extra storage overhead)
[Inference Phase]
  -> Receive activated expert IDs for the current layer
  -> Online de-quantize and compute within the shared FP8 Workspace
  -> Output expert hidden states
  -> Release Workspace lock, hand over to the next layer for reuse

5. Day-0 Adaptation Engineering Checklist and Industry Outlook

To help developers quickly replicate the practices from this implementation, the following is a five-step core checklist for Day-0 adaptation:

1. Weight Loading Verification: Confirm the loading integrity and verification mechanism for non-standard format weights like MXFP4/FP8.
2. Operator Integration Testing: Core operators like sparse attention and MoE expert routing must pass unit tests and end-to-end precision alignment.
3. KV Cache Layout: Adopt paged strategies like PagedAttention, set reasonable OOM thresholds, and validate million-token scenarios.
4. Precision Calibration: Use calibration datasets to evaluate inference quality loss during precision conversion steps like FP8 de-quantization.
5. Multi-Card Parallel Debugging: Verify state consistency and communication overhead under tensor and pipeline parallelism.

This adaptation not only validates the maturity of domestic computing power in hardware-software co-design ecosystems but also provides a rapid pathway for developers: when new open-source models emerge, deployment can be completed using standard open-source framework interfaces and underlying software stacks without waiting for lengthy official operator adaptations.

As large models accelerate towards ultra-long contexts, multimodality, and low precision, the underlying computing power competition has shifted from pure computing metric comparisons to the maturity of software stack ecosystems. In the future, general-purpose computing infrastructure supporting larger parameter scales and more complex inference pipelines will become the core engine driving AI application deployment. The Day-0 adaptation in the era of ten-thousand-card clusters is an indispensable piece of this grand ecosystem puzzle.