Adapting Ultra-Large Sparse Models on Day-0: Key Software Stack Modifications

1. The Core Challenge: Why Day-0 Adaptation Is Difficult

When parameter counts reach the trillion level and context windows extend to millions of tokens, the inference bottleneck shifts from “is there enough compute?” to “can the software stack handle it?” Day-0 adaptation means that on the day model weights are publicly released, an inference service can load them, start up, and run the basic pipeline end to end. What is being tested is not isolated compute capability, but the entire software stack’s ability to absorb a new architecture.

A typical combination of technical challenges is sparse attention, ultra-long context, and low-precision expert weights appearing together. Sparse attention changes the computational graph structure of attention. Long context amplifies the memory pressure of the KV cache. Low-precision expert weights require the runtime to support flexible numerical conversion. If any one of these pieces is missing, the model cannot perform end-to-end inference.

This article draws on an adaptation effort for a trillion-parameter model and abstracts a general engineering methodology. It discusses the key modifications an inference software stack needs to undergo.

2. Overall Architecture: From Community Model Code to a Working Inference Service

The central idea of the adaptation effort is to take model code implemented by the upstream community and map it onto target hardware through an operator library, a domain-specific language (DSL), and a compiler, producing a runnable inference service. The pipeline can be divided into four layers:

• Inference framework layer: receives the upstream model definition and handles scheduling, batching, KV cache management, and serving interfaces.
• Operator layer: provides key operators such as attention, normalization, matrix multiplication, and expert routing, with dedicated implementations for sparse patterns.
• DSL / compiler layer: translates custom computation logic in the model, such as special state management or residual connections, into hardware-executable kernels.
• Hardware and runtime layer: manages memory allocation, stream scheduling, and low-precision numerical support.

In this layered structure, the most likely failure points are not standard operators but new computation paths written by model authors. Therefore, the first step in adaptation is usually to identify which computation paths are new, then decide whether to reuse existing operators or fill the gaps quickly through the DSL.

3. Three-Stage Modification for Sparse Attention

The difference between dynamic sparse attention and traditional dense attention is not merely “computing fewer positions.” It introduces runtime dynamic decision-making. The adaptation work can be split into three stages.

3.1 Index Selection: Turning Dynamic Decisions into Operators

The first step of sparse attention is selecting, for each query, the set of key/value positions that should participate. This selection process itself is computationally intensive and involves irregular memory access. Common implementations rely on a combination of sorting, top-k, and gather. In the software stack, it needs to be encapsulated as a standalone operator, with two key considerations:

1. Avoid synchronization points: if index selection introduces synchronization between the device and the host, it will severely disrupt the pipeline. Selection should be completed on the device side and consumed directly whenever possible.
2. Support variable-length outputs: different queries may select different numbers of indices. The operator needs to support variable-length or padded outputs and align them with the subsequent gather step.

The following pseudocode illustrates how index selection and attention computation connect:

# Pseudocode: index selection and consumption for sparse attention
def sparse_attention(q, k, v, block_table, topk):
    # 1. Index selection is done on device, producing variable-length indices and weights
    idx, w = select_indices(q, k, block_table, topk)   # No host synchronization
    # 2. Gather KV by index and perform grouped attention
    k_sel = gather(k, idx)
    v_sel = gather(v, idx)
    out = grouped_attention(q, k_sel, v_sel, w)
    return out

3.2 KV Cache Layout: Reorganizing Memory for Sparse Access

Dense attention typically uses a contiguous KV cache layout, ordered by layer and head. Sparse attention requires organizing the cache by blocks or pages so that index selection results can be mapped directly to memory addresses. The adaptation requires redesigning the cache layout:

• Use a paged layout where each page holds a fixed number of tokens, and index selection outputs page numbers plus intra-page offsets.
• Keep addresses within a page contiguous to reduce fragmented memory access during gather.
• Coordinate with batch scheduling so that page tables for different requests can be managed independently, enabling memory reclamation under long contexts.

The benefit of this step is twofold: it makes memory access patterns friendlier for sparse access and keeps memory usage manageable at million-token context lengths.

3.3 Separating Prefill and Decode Paths

The prefill phase is throughput-oriented, with a large number of queries, so sparse attention index selection can be parallelized in batches. The decode phase is latency-oriented, with only a few queries per step, but it requires reading the historical KV cache. Since the computational characteristics differ significantly, the two should be optimized separately:

Phase	Main bottleneck	Optimization focus
Prefill	Compute throughput and index selection parallelism	Large-batch index selection, operator fusion
Decode	KV read bandwidth and latency	Page table caching, reducing kernel launch overhead

After the two phases are connected to the sparse attention operators separately, they must be coordinated with the model’s other new structures to form a complete execution path.

4. Assembling the Complete Execution Path

Getting a single operator to run does not mean the model will run. Sparse attention often appears together with other new structures, and adaptation must ensure they connect correctly within the same execution graph:

• State management modules: some new architectures introduce learnable state caches that require dedicated memory regions and update timing. These must be coordinated with KV cache allocation and deallocation order to avoid memory fragmentation.
• Cross-layer residuals: cross-layer connections mean intermediate results must be preserved across multiple layers. The software stack needs explicit intermediate tensor lifecycle management to prevent premature reuse.
• Vision attention branches: multimodal inputs introduce a separate attention path whose tensor shapes differ from the text path. Shape dispatch and operator selection need to be handled at the scheduling layer.

These structures together form a complete execution path. In practice, it is advisable to describe dependencies explicitly in the framework layer using a computational graph and let the compiler decide fusion and scheduling, rather than relying on implicit execution order.

5. Runtime Handling of Low-Precision Expert Weights

Trillion-parameter models commonly adopt a mixture-of-experts (MoE) structure and store expert weights in low-precision formats such as MXFP4. Loading these weights directly introduces two problems: memory usage remains high, and compute units have inconsistent support for mixed precision.

A pragmatic approach is online conversion plus workspace reuse:

1. Online conversion: at load time, convert MXFP4 weights on demand into a higher-precision format better suited to compute units, such as FP8. Conversion is done on the device side to avoid host round trips.
2. Cross-layer workspace reuse: expert computation is triggered layer by layer and token by token, with only a small number of experts active at any moment. Therefore, a fixed-size FP8 workspace can be allocated once and reused across layers, rather than persisting high-precision copies for every expert in every layer.
3. No offline pre-generation required: this eliminates the step of pre-generating full high-precision weight files, significantly shortening deployment preparation time and making weight updates more flexible.

The pseudocode is as follows:

# Pseudocode: online conversion of expert weights with workspace reuse
fp8_workspace = alloc_workspace(max_active_experts, expert_shape)
def expert_forward(x, expert_id, mxfp4_weight):
    w = convert_mxfp4_to_fp8(mxfp4_weight, out=fp8_workspace)  # On-device online conversion
    return matmul(x, w)

It is worth noting that online conversion introduces additional conversion overhead. The implementation should therefore overlap conversion with the compute pipeline as much as possible and keep conversion results confined to the workspace, discarding them after use.

6. Practical Engineering Checklist

Based on the analysis above, here is an actionable checklist for similar adaptation efforts:

• Get the pipeline running first, then optimize performance: prioritize end-to-end availability for loading, startup, and basic text and image inference before profiling individual segments.
• Identify new computation paths: compare against the upstream model code and list all non-standard operators. For each one, decide whether to reuse, modify, or write from scratch.
• Eliminate host synchronization first: index selection, dynamic shapes, and variable-length outputs are common sources of synchronization and should be the first optimization targets.
• Plan memory up front: KV cache, state cache, intermediate tensors, and expert workspace requirements should be planned together at the design stage.
• Keep low-precision strategy fallback-ready: the online conversion path should retain an off switch for precision debugging.
• Establish layered regression testing: cover operator-level, model-level, and service-level tests separately to prevent issues from drifting between layers.

7. Summary

Day-0 adaptation of trillion-parameter sparse models is fundamentally a question of whether the software stack can absorb new computational graph structures. Sparse attention requires modifications in three areas: index selection, KV cache layout, and phase-specific computation paths. Low-precision expert weights require the runtime to support online conversion and workspace reuse. State management, cross-layer residuals, and multimodal branches require the framework layer to provide explicit dependency and lifecycle management. By decomposing these pieces according to the layered structure and verifying them against a checklist, community model code can be reliably turned into a working inference service.