Day-0 Inference Adaptation and Performance Optimization for New LLM Architectures

1. Background: Inference Challenges Posed by New Model Architectures

The large language model field has recently seen a significant architectural iteration. The new generation of models introduces notable innovations in attention mechanisms, interleaving two distinct attention paths: one that processes sequence information through state recurrence, and another that builds on multi-head latent attention by introducing an index selection mechanism — a lightweight indexer first selects Top-K historical positions, and then a sparse attention module reads the corresponding key-value cache to complete the computation. Meanwhile, the model continues to aggregate cross-layer representations through an attention residual mechanism.

This kind of architectural innovation places new demands on inference systems: they must quickly adapt to new operators while maintaining efficiency across computation, communication, scheduling, and storage. This article systematically outlines how to achieve Day-0 inference support and performance optimization for the new architecture from a technical solution perspective.

2. Core Approach to Rapid Adaptation: Reuse and Extension

2.1 Building on Existing Optimization Accumulations

When facing a new architecture, the most efficient adaptation strategy is not to start from scratch, but to carry forward the optimization accumulations from the previous model generation. Specifically, the following capabilities already refined in prior models can be directly reused or slightly extended:

• Fused implementations of state-recurrence operators: Including the previously optimized fused operator solutions, which can continue to handle the core sequence processing computation in the new model;
• Two-stage fusion of attention residuals: The computation pattern for cross-layer representation aggregation is consistent with the previous generation, and the fusion strategy can be directly transferred;
• Mixture-of-experts and communication overlap solutions: The computation-communication co-optimization for large-scale expert parallelism is equally applicable in the new architecture;
• Hybrid parallelism strategies: Combinations of tensor parallelism, pipeline parallelism, and data parallelism, along with communication optimization and multi-stream orchestration;
• Compiler and scheduling layer optimizations: Including super kernels, graph-mode compilation, compilation caching, and weight layout optimization;
• Speculative inference and state replay mechanisms: Existing solutions for accelerating the decoding phase and state recovery.

By reusing these capabilities, the adaptation cycle for the new model can be greatly shortened, and efficiency levels across computation, communication, scheduling, and storage can quickly reach a usable state.

2.2 Targeted Optimization for New Operators

On top of reuse, the sparse attention path introduced by the new architecture requires dedicated optimization. This path contains two core components:

Indexer module: Responsible for relevance scoring, head-weight weighting, and Top-K selection. The optimization strategy is to fuse these three steps into a single execution, directly reading the paged key cache, reducing intermediate result write-backs and operator launch overhead.

Sparse attention module: Reads the main attention keys and values based on the positions selected by the indexer to complete sparse computation. Both the preprocessing stage and the decoding stage support BF16 and quantized paths.

In addition, the indexer preprocessing stage is further optimized through three-stream coordination: query projection and key normalization, query quantization and head-weight projection are executed in parallel, reducing serial waiting time during the input preparation stage.

3. Quantization Strategy: Layered Design to Reduce Memory Access Overhead

Quantization is a key means of reducing inference costs, but different modules have different sensitivity to precision, requiring a layered design:

3.1 Projection Computation Retains High Precision

The projection computation in the sparse attention path retains BF16 precision, ensuring that the computation quality of attention weights is not affected by quantization.

3.2 Full Quantization of the Indexer

The indexer’s queries and keys use the MXFP8 format with per-group dynamic quantization, where the group size is set to 32 — that is, every 32 elements share a single scaling factor. This fine-grained quantization maintains the accuracy of index selection while compressing data.

3.3 Key-Value Cache Quantization

The key-value cache of the sparse attention module uses the FP8 format with a per-token-head-tile quantization granularity, where the tile size is 128 — that is, within each token and each head, elements are quantized in blocks of 128. This strategy significantly reduces the storage and memory access overhead of index data and key-value caches, which is especially critical for long-sequence scenarios.

4. Long-Sequence Parallelism: Splitting Strategies and Load Balancing

4.1 Combined Deployment of Context Parallelism

For long-sequence inputs, this work supports a combined deployment of context parallelism and tensor parallelism. Long inputs are split along the sequence dimension, combined with sequence parallelism within tensor parallel groups, reducing per-GPU activation and intermediate tensor memory usage.

4.2 Load Balancing for Causal Index Computation

Causal index computation in sparse attention has an inherent problem: the further back a position is in the sequence, the longer the history it needs to retrieve, resulting in significant computation differences across sequence segments. If simply split in order, this causes cross-GPU load imbalance.

The solution is to adopt a Zigzag layout: pairing segments from the beginning and end of the sequence and assigning them to the same GPU, so that the total computation on each GPU tends to be balanced. After computation is complete, the original order is restored to ensure result correctness.

4.3 Sequence Splitting for State-Recurrence Operators

In ultra-long sequence scenarios, the preprocessing computation time of the state-recurrence path grows linearly with sequence length, directly affecting time-to-first-token latency. When only tensor parallelism is used to split by attention heads, each GPU still needs to process the full sequence, and cannot share the pressure brought by sequence length.

To address this, a sequence parallelism solution implemented on a domain-specific language further splits the input along the sequence dimension, distributing long-sequence computation across multiple GPUs. Each GPU first computes local states in parallel, then aggregates global states through cross-GPU communication, ultimately reducing time-to-first-token latency while also decreasing per-GPU activation and intermediate tensor memory usage.

5. Summary and Outlook

Day-0 inference adaptation for new model architectures is essentially an engineering problem of “reuse plus extension.” The core methodology can be summarized as:

1. Maximize reuse of existing optimizations: General capabilities such as operator fusion, parallelism strategies, and compiler scheduling should not be reinvented;
2. Targeted optimization for new operators: Apply fusion and multi-stream coordination to newly introduced computation paths such as sparse attention;
3. Layered quantization: Different modules adopt different precision strategies to strike a balance between precision and efficiency;
4. Long-sequence parallelism: Control time-to-first-token latency and memory usage through sequence splitting and load balancing strategies.

This methodology is not only applicable to the current generation of models but also serves as a reference for new architectures that may emerge in the future. As large model architectures continue to evolve, the ability to rapidly adapt inference systems will become an important component of platform competitiveness.