Inference Adaptation and Performance Optimization for a Domestically Developed LLM's New Architecture

1. Background: Inference Challenges Posed by the New Architecture

A large model team recently released and open-sourced a new-generation flagship model. Architecturally, it continues the accumulation of its predecessor while introducing a more structural combination of attention mechanisms: it interleaves two attention paths—one that processes sequences via state recurrence (hereafter “recurrent attention”), and another that builds on Multi-head Latent Attention (MLA) with an index selection mechanism (hereafter “sparse index attention”). The latter uses a lightweight indexer (Lightning Indexer, LI) to select the Top-K historical positions, after which Sparse Flash Attention (SFA) reads the corresponding KV to complete the sparse attention computation. The model also continues to aggregate cross-layer representations through Attention Residuals.

The common characteristics of this class of architecture are: the longer the sequence, the more uneven the load of index computation; the storage and memory-access overhead of KV cache and index data become bottlenecks; and prefill latency grows significantly with sequence length. For an inference platform, whether adaptation can be completed on the day the model is open-sourced (i.e., Day0 support) depends on whether reusable operator libraries, parallel strategies, and scheduling mechanisms have been accumulated beforehand.

This article starts from four stages—computation, communication, scheduling, and storage—and systematically reviews the adaptation path and optimization methods for this architecture on a domestic AI accelerator platform, as a reference for developers engaged in large model inference deployment.

2. Operator-Layer Reuse: From Accumulated Optimization to Rapid Adaptation

The new model does not start from scratch. Its predecessor had already accumulated a set of key operators and fusion strategies during training and inference, and the new version directly reuses this accumulation, greatly shortening the adaptation cycle. Reusable optimizations mainly include:

• Recurrent attention operator family: FlashKDA and MegaKDA implementations for the state-recurrence structure, covering forward computation and state updates;
• Two-stage fusion of attention residuals: splitting the aggregation of cross-layer representations into two fusion stages to reduce intermediate tensor write-backs;
• Mixture-of-experts and communication fusion operators: implementations such as MegaMoE and MC2, which combine expert routing, computation, and collective communication into a single orchestration;
• Parallelism and scheduling infrastructure: hybrid parallel strategies, communication optimization, multi-stream orchestration, as well as SuperKernel, graph mode, compilation cache, and NZ weight layout;
• Speculative inference and state replay: the DSpark speculative inference mechanism and the ReplaySSM state replay mechanism, used to reduce effective latency in the decoding stage.

These capabilities were not developed temporarily for the new model but are general-purpose components validated on the predecessor. The core work of adapting to the new architecture is to correctly connect the newly added sparse index path with the existing operator system and to fill in the partitioning strategy along the parallel dimensions.

3. Fusion and Multi-Stream Orchestration for Sparse Index Attention

The computation of sparse index attention can be divided into two stages: the index stage and the main attention stage. In the index stage, LI performs relevance scoring, head-weight weighting, and Top-K selection; in the main attention stage, SFA reads the main attention KV according to the selected indices and completes the sparse computation.

3.1 Operator Fusion for LI

In a naive implementation, scoring, weighting, and selection each produce intermediate results and launch operators multiple times, incurring considerable write-back and launch overhead. The optimization fuses the three into a single execution unit: LI directly reads the paged K cache and completes scoring and Top-K selection within the fused kernel, avoiding write-back of intermediate results to global memory. This fusion applies to both the prefill and decode stages and supports both the BF16 and quantized paths.

3.2 Three-Stream Coordination in the Indexer Prolog

Before index computation formally begins, there is a segment of input preparation: Q projection, K normalization, Q quantization, and head-weight projection. These operations have data dependencies among them, and serial execution creates a “preparation-stage wait.” The optimization splits them into three parallel streams: one for Q projection, one for K normalization, and one for Q quantization and head-weight projection. With three-stream coordination, the serial wait in the input preparation stage is significantly shortened, and the index kernel can start earlier.

3.3 Overall Approach to Multi-Stream Orchestration

The value of multi-stream orchestration extends beyond the prolog. In the decoding stage, there are natural overlap windows among attention computation, expert computation, and communication operations. By assigning different operators to independent streams and setting reasonable event synchronization points, part of the communication and memory-access latency can be hidden without changing computation semantics. This mechanism is consistent with the multi-stream scheduling framework accumulated by the predecessor, and the new model only needs to register a new operator dependency graph to reuse it.

4. Tiered Quantization: Compressing Index Data and KV Cache

In long-sequence scenarios, the size of the KV cache and index data directly determines memory-access overhead and memory footprint. The solution adopts a “tiered processing, on-demand quantization” strategy:

Data Type	Precision Strategy	Quantization Granularity
DSA projection computation	Retain BF16	—
LI index Q/K	MXFP8 dynamic quantization	per-group, group size 32
SFA KV cache	FP8	per-token-head-tile, tile size 128

Specifically, the Q/K on the index side uses MXFP8 with per-group dynamic quantization, where every 32 elements share one scaling factor, balancing precision and compression ratio; the main attention KV cache uses FP8 with a quantization granularity of “blocks of 128 elements within each token and each head,” i.e., within each head of each token, a scaling factor is computed independently for each 128-element block. Projection computation retains BF16 to avoid accumulation of quantization error in the projection layers.

The benefit of this tiered strategy is twofold: on the one hand, it reduces the storage footprint of index data and KV cache, allowing longer contexts to reside in memory; on the other hand, it reduces memory-access bandwidth pressure, which is often the actual bottleneck in the decoding stage of sparse attention.

5. Prefill Sequence Parallelism: Splitting Long Inputs and Balancing Load

Long-input prefill is the main source of time to first token (TTFT). The solution supports combined deployment of context parallelism and tensor parallelism (CP×TP), splitting long inputs along the sequence dimension and combining sequence parallelism within the TP group to reduce per-card activation and intermediate tensor footprint.

5.1 Load Imbalance in Causal Indexing

The index computation of sparse index attention is causal: the later a position is in the sequence, the longer the history it can reference, and the larger the index computation. If the sequence is simply split into contiguous intervals, cross-card imbalance occurs—“early segments leave cards idle, later segments overload cards.”

5.2 Zigzag Layout

The solution is to adopt a Zigzag layout: pair the head and tail of the sequence, assigning “a small early segment” and “a large later segment” to the same card, so that the index computation borne by each card is roughly equal. After computation, the positions of the segments are restored in the original order to ensure consistency with subsequent layers and the output. This layout has no effect on the semantics of attention computation; it only changes the distribution of data across devices.

5.3 Splitting Ultra-Long Sequences for Recurrent Attention

The prefill-stage time of recurrent attention also grows with sequence length. When splitting only by attention head for TP, each card still needs to process the full sequence, so TTFT improvement is limited for long sequences. To address this, the solution implements sequence parallelism for recurrent attention (FlashKDA CP) based on a domain-specific language toolchain: the input is further split along the sequence dimension, spreading long-sequence computation across multiple cards, reducing TTFT while also lowering per-card activation and intermediate tensor footprint. For ultra-long context scenarios, this capability is key to the first-token experience.

6. Scheduling and Storage: Making the Optimizations Truly Land

Once operators and parallel strategies are determined, system-level mechanisms determine whether the actual gains can be realized:

• Graph mode and compilation cache: orchestrate operators such as attention, expert, and communication into computation graphs to reduce host-side launch overhead; the compilation cache allows requests with repeated shapes to skip recompilation.
• NZ weight layout: adapt to the matrix computation preferences of the accelerator platform and reduce format conversion during weight loading.
• SuperKernel: aggregate fine-grained operators into large kernels to reduce kernel launch and scheduling frequency.
• Speculative inference and state replay: speculative inference improves decoding throughput with low-cost drafts; the state replay mechanism is used to quickly restore recurrent state when speculation fails or branches switch, avoiding repeated computation.

These mechanisms complement the operator optimizations above: operator fusion reduces per-execution overhead, graph mode and SuperKernel reduce scheduling overhead, speculation and state replay reduce decoding steps, and quantization and sequence parallelism reduce memory-access and memory footprint pressure. Only the combination of all four constitutes a complete path to inference efficiency improvement.

7. Practical Recommendations and Conclusion

For developers planning to deploy models with this kind of hybrid attention architecture on a domestic AI accelerator platform, the following order can be followed:

1. First verify operator availability: confirm whether core operators such as recurrent attention, sparse index attention, and attention residual aggregation already have platform implementations, and prioritize reusing mature operator libraries.
2. Then establish the parallel combination: choose the CP×TP combination according to sequence length and number of cards; for long-sequence scenarios, be sure to enable sequence-dimension splitting and Zigzag load balancing.
3. Set quantization strategy layer by layer: retain high precision for projections, and use group quantization and block quantization for the index and KV cache respectively, avoiding precision loss caused by one-size-fits-all quantization.
4. Finally tune scheduling parameters: enable graph mode, compilation cache, and multi-stream orchestration, and use speculative inference to evaluate the balance point between end-to-end throughput and time to first token.

Overall, the Day0 inference support for this model is not a single-point breakthrough but the combined result of “operator reuse + fusion of new paths + parallel补齐 + system scheduling.” Its methodology has general reference value for similar hybrid attention architectures: decompose the new architecture into computation patterns that can be mapped onto the existing optimization system, then make targeted reinforcements for the new sparse index and long-sequence load characteristics, so that usable deployment and experience capabilities are available on the day of open-sourcing.