Engineering Practices for Domestic GPU Software Stacks: Rapid Adaptation of Ultra-Large-Scale Sparse Attention Models

I. Background: Compute and Software Stack Challenges in Ultra-Large-Scale Model Inference

As large language models (LLMs) scale toward the trillion-parameter mark, model inference poses a dual challenge to computing hardware and software ecosystems. On one hand, the growth in model parameters and context length drives an exponential increase in demand for memory capacity, compute throughput, and communication bandwidth. On the other hand, the rapid evolution of model architectures—such as sparse attention and mixed-precision expert weights—requires underlying compute platforms to complete adaptations immediately. Otherwise, even the most advanced hardware cannot deliver practical value.

Take next-generation models characterized by million-level context windows and multi-trillion parameters as an example. The inference pipeline involves the efficient loading of model weights, efficient compilation of compute graphs, specialized implementation of attention mechanisms, and fine-grained management of expert parallelism strategies. Any missing or inefficient component can prevent the model from starting or cause inference performance to fall severely below standards. Consequently, building a software stack capable of keeping pace with model evolution and possessing rapid adaptation capabilities has become a key competitive frontier in the AI infrastructure domain.

II. Core Objectives and Overall Strategy of Day-0 Adaptation

Day-0 adaptation refers to the ability of a compute platform to complete the entire inference pipeline on the exact day the model weights are publicly released, eliminating the waiting period where “the model is released but the hardware cannot run it.” This objective requires the software stack to possess the following core capabilities:

1. Automation and Standardization of Model Integration: By maintaining compatibility with mainstream inference frameworks (such as widely used community serving engines), the platform can ingest upstream model implementation code, thereby reducing redundant development efforts.
2. Rapid Compute Graph Migration: Mapping various operators within the model to the underlying hardware instruction set to ensure the highly efficient execution of critical compute paths.
3. Timely Support for New Architectural Features: Providing specialized operator libraries or compilation optimization strategies tailored to entirely new structures introduced by the model, such as sparse attention and mixed-precision storage.

The overall strategy for achieving Day-0 adaptation is layered decoupling: the upper layer interfaces with community model code, the middle layer handles compute graph conversion via the compiler and operator libraries, and the lower layer drives hardware execution. This layered architecture ensures that model developers do not need to understand underlying hardware details, while the hardware platform can continuously refine its compilers and operator libraries to cover new model features.

III. Deep Dive into Key Technical Components

3.1 Model Loading and Compute Graph Migration

The weight files of ultra-large-scale models often reach hundreds of gigabytes or even terabytes. How to efficiently and accurately load and initialize distributed inference services is the first challenge. The software stack must implement a sharded loading mechanism, distributing weights across multiple compute cards according to model parallel dimensions. Simultaneously, compute graph migration requires converting dynamic graphs described by frameworks like PyTorch into static execution graphs, enabling the underlying compiler to perform deep optimizations such as operator fusion and memory reuse.

In the specific practices discussed in this article, the implementation team leveraged a mainstream community inference framework to ingest the model implementation. They then progressively replaced critical compute paths in the model code with highly efficient underlying hardware implementations using a self-developed compiler and operator libraries. This process is not merely an API-level substitution but a re-evaluation and optimization of the compute graph topology. For instance, consecutive matrix multiplication operations can be fused into a single hardware instruction. For operator sequences with data dependencies, pipelines must be carefully designed to hide memory access latency.

3.2 Adaptation Strategies for Sparse Attention Mechanisms

3.2.1 Index Selection and KV Cache Layout Optimization

Unlike traditional dense attention, sparse attention mechanisms (such as DSA, Dynamic Sparse Attention) compute attention scores only between queries and a small number of key-value pairs, thereby reducing computational complexity from O(n²) to nearly O(n). However, realizing this advantage depends on a prerequisite: the ability to efficiently determine which key-value pairs are important.

The primary task of adaptation is migrating the index selection logic. The model typically maintains an index structure for each query head to locate the key-value positions required for computation. The underlying hardware must design an efficient memory access pattern for this index structure to ensure that index query operations do not become a performance bottleneck.

An even more critical optimization lies in the layout design of the KV (Key-Value) cache. Because the range of key-values focused on varies significantly across different queries, the physical storage layout of the KV cache must support flexible and efficient access. A common strategy is to design the KV cache as a block structure, where each block manages a fixed-length sequence of key-values and dynamically combines the required blocks based on the index. This design not only improves cache hit rates but also simplifies memory management and reduces fragmentation. Furthermore, optimal block sizes and layout strategies differ between compute phases (Prefill and Decode), requiring separate tuning.

3.2.2 Differentiated Compute Paths for Prefill and Decode Stages

LLM inference can be divided into two distinct phases:

• Prefill Stage: Processes all tokens of the input prompt and computes attention scores for all positions in parallel. This phase is highly compute-intensive and requires fully utilizing hardware compute power.
• Decode Stage: Generates output token-by-token, computing attention for only one position at each step. This phase has significantly higher memory access intensity than compute intensity, and system throughput is limited by memory bandwidth.

Under the sparse attention framework, these two phases present different challenges. The Prefill stage exhibits lower sparsity because the input sequence typically contains a large number of tokens that need attention. The Decode stage exhibits higher sparsity because, upon generating a new token, it may only need to focus on a limited number of historical positions.

In the adaptation solution, the software stack designs distinct compute paths for both phases. The Prefill stage focuses on efficient parallelization of the compute graph and operator optimization for large-scale matrix multiplications. The Decode stage emphasizes optimizing memory access patterns, utilizing indices to directly locate KV blocks and avoid loading irrelevant data. Some platforms also introduce standalone sparse attention operators capable of dynamically adjusting the compute pipeline based on input indices, thereby achieving near-peak theoretical compute utilization in both phases.

3.3 Online Processing of Mixed-Precision Expert Weights

New-generation ultra-large-scale models universally adopt the Mixture-of-Experts (MoE) architecture, where each token activates only a portion of the expert networks. To reduce memory footprint and bandwidth pressure, expert weights tend to be stored in low-precision formats (such as the 4-bit floating-point MXFP4 format). However, using low-precision weights directly for computation leads to precision loss, failing to meet inference quality requirements.

Industry-standard solutions fall into two categories:

1. Offline Conversion: Before starting the inference service, all expert weights are converted from low-precision to high-precision (such as FP16 or BF16) and then loaded into memory. This method is simple and reliable but has significant drawbacks: the converted weights occupy 4 times the memory space, and for trillion-level models, the conversion process can take hours, severely slowing down service startup.
2. Online Conversion: During computation, only the currently activated expert weights are converted from low-precision to high-precision in real-time and temporarily stored in a workspace. This strategy distributes the conversion overhead across the inference process, avoiding the startup latency associated with full conversion. However, it requires allocating additional workspace memory and executing conversion instructions, which may introduce bubbles in the compute pipeline.

A more refined approach in practice is “cross-layer workspace reuse.” Given that the sets of activated experts across different layers in an MoE model often do not overlap, workspaces can be dynamically allocated at a block granularity, allowing the same workspace block to serve multiple expert computations across different layers sequentially. By carefully designing the lifecycle management of the workspace, the additional memory overhead can be kept to a minimum. Concurrently, leveraging conversion pipeline technology enables precision conversion and matrix multiplication computation to overlap, achieving the ideal effect of “fast startup, low overhead, and transparent conversion costs.”

IV. Engineering Implementation Highlights and Optimization Cases

4.1 State Management Component Collaboration Mechanism

In a real inference system, modules such as sparse attention index selection, KV cache management, and cross-layer residual connections do not exist in isolation; they form a tightly coupled data flow network. For example, the results of index selection directly affect the transfer addresses of KV cache blocks; intermediate results of cross-layer residual connections need to be temporarily stored in specific buffers; and visual attention branches (unique to multimodal models) and text attention branches need to be processed separately before fusion. The core of engineering implementation lies in designing a clear state management mechanism, defining data dependencies and synchronization points between modules, and ensuring orderly execution through events or synchronization primitives.

4.2 Inference Service Stability Assurance

Service startup is only the first step; production-grade inference services also require long-term stable operation without crashes or memory leaks. Addressing the dynamic index characteristics of sparse attention requires special precautions against out-of-bounds access and wild pointers. For the online precision conversion workspace, it is necessary to ensure that allocation and release operations in multi-threaded environments are thread-safe. Some platforms introduce mechanisms such as deadlock detection and watchdog timers. Once an anomaly is detected, the inference process can be automatically restarted, minimizing the impact on online business operations.

V. Outlook: The Future Form of AI Infrastructure

The continuous breakthroughs in large model capabilities and the iteration of underlying hardware acceleration technologies are showing a trend of synchronized resonance. For compute platforms, merely providing high-performance hardware is far from sufficient. True competitiveness lies in the completeness of the software stack and the responsiveness of the ecosystem. Platforms capable of completing adaptations on the day a model is released will earn users valuable time advantages and build insurmountable ecosystem barriers.

Future software stacks should continue to evolve in the following directions: higher-level automated model compilation entry points (reducing manual operator migration efforts), smarter sparse pattern recognition (automatically discovering potential sparse structures within models), more efficient heterogeneous scheduling capabilities (simultaneously managing different types of compute resources), and more comprehensive end-to-end debugging and profiling tools. With sustained efforts in these directions, the performance ceiling of LLM inference will be continuously pushed higher, and the deployment costs of AI applications will continue to decline.