Optimizing LLM Inference for New Architectures: Day-0 Adaptation, Sparse Attention, and Long-Sequence Parallelism

1. Background: Inference Challenges Posed by New Architectures

Recent years have seen rapid architectural evolution in large language models. From dense to sparse attention, from conventional Transformers to structures that interleave state recurrence with attention—each innovation improves model capability while complicating inference deployment. Typical problems developers face include:

• New operators lack ready-made efficient implementations, and stitching them together from basic primitives causes excessive intermediate write-backs and wasted memory traffic;
• The irregularity of sparse-attention index computation easily leads to load imbalance across compute units;
• In long-sequence scenarios, Prefill compute and activation memory grow quickly with sequence length, hurting time-to-first-token latency;
• Mixture-of-experts (MoE) mechanisms add communication and scheduling complexity, making parallelism strategy selection harder.

The model architecture discussed in this article is characterized by interleaving a state-recurrence module (KDA) with a sparse attention module (DSA). KDA handles long-range dependencies through state recurrence. DSA builds on multi-head latent attention and introduces an index-selection mechanism: a lightweight indexer (LI) first scores historical positions by relevance and selects the Top-K positions, then the sparse attention compute module (SFA) reads the KV cache at those positions and performs the sparse attention computation. In addition, the model aggregates cross-layer representations via Attention Residuals, further increasing the structural complexity of the computation graph.

For architectures like this, the core engineering problem is how to complete inference adaptation and performance optimization quickly without writing every operator from scratch. This article addresses that problem from several angles: operator reuse, fusion optimization, quantization strategy, parallel splitting, and automated code generation.

2. The Core Idea Behind Day-0 Adaptation: Reuse and Abstraction

“Day-0 support” means having a working inference deployment on the target hardware platform as soon as a model is released or open-sourced, rather than waiting weeks or months for adaptation. The key is to abstract existing optimization experience into reusable compute patterns instead of starting from zero for every new model.

Specifically, rapid adaptation is achieved at several levels:

Operator-level reuse. Although the new DSA structure introduces index selection, its underlying attention computation, state recurrence, and feed-forward modules share substantial structural similarity with existing implementations. The KDA portion in particular retains essentially the same core compute pattern as previous versions, so existing fused kernels (such as FlashKDA and MegaKDA) can be reused directly, avoiding duplicate development.

Pattern-level abstraction. Attention residual aggregation looks like a cross-layer data dependency in the graph, but at its core it is “a fusion operation over multiple branch outputs.” Once abstracted into a two-stage fusion strategy, the optimization becomes independent of specific layer counts or dimensions.

Scheduling-level reuse. Communication optimization and multi-stream orchestration under MoE are model-agnostic. By parameterizing these scheduling strategies, different expert counts and sequence lengths can be supported quickly.

This “abstract–reuse–adapt” approach turns adaptation work from “developing new operators” into “composing existing capabilities and debugging new paths,” dramatically shortening the window from model release to usable deployment.

3. Key Sparse-Attention Optimizations: Fused Indexing and Multi-Stream Pipelining

The DSA module is the most distinctive part of the new architecture. Its execution flow is as follows: LI scores historical positions against the query vector, applies head-weight weighting, and selects the Top-K indices; SFA then reads the corresponding KV cache entries and performs the sparse attention computation. Implementing each step naively in semantic order incurs multiple global memory round-trips and kernel launches, and intermediate index results must be written back and re-read, noticeably increasing latency.

Fused index computation. The core optimization fuses three steps inside LI—relevance scoring, head-weight weighting, and Top-K selection—into a single kernel. The fused kernel reads the paged K cache directly, completes computation in registers or shared memory, and outputs only the final Top-K indices, eliminating global write-back of intermediate tensors. This fusion works for both Prefill and Decode, and supports both BF16 and quantized paths.

Three-stream coordination to reduce serial waiting. Before index computation, preparation steps such as query projection (Q Projection) and K normalization must run. These steps have dependencies, but executing them sequentially creates a clear serial bottleneck. The solution is a three-stream orchestration: Q projection and K normalization run on one stream, while Q quantization and head-weight projection run on another, both proceeding in parallel and joining at an aggregation point. This converts serial waiting in the input-preparation stage into parallel execution, with particularly visible per-step latency improvements in Decode.

Note that “stream” here refers to concurrent execution queues on the compute device. By assigning independent operators to different queues and inserting synchronization events at dependency points, hardware parallelism can be fully exploited without changing compute semantics.

4. Layered Quantization: Balancing Accuracy and Overhead

Sparse attention introduces two categories of data: the query and key data LI uses for index selection, and the SFA KV cache. Their access patterns, precision sensitivity, and sizes differ, so a layered quantization strategy is applied.

MXFP8 for LI data. The indexer’s core job is relative comparison, not precise retrieval. LI dynamically quantizes each query and key in per-group fashion, with one scale factor shared by every 32 elements. The fine granularity keeps accuracy loss within acceptable bounds while sharply cutting memory traffic. For a memory-bound lightweight operator like index computation, the benefit of moving less data usually outweighs the extra dequantization cost.

FP8 for the SFA KV cache. The KV cache is far larger than the index data and is one of the main contributors to memory footprint. When quantizing to FP8, the granularity is per-token-head-tile: within each token and each head, elements are quantized in blocks of 128. This granularity reflects the access pattern of attention computation: the KV cache is typically accessed per head, and the block size a compute unit processes relates to the tile size. Quantizing at this granularity maximizes locality during dequantization and avoids cross-tile data dependencies.

In practice, this layered strategy reduces both index-data memory traffic and KV-cache storage overhead, allowing longer contexts or larger batch sizes on the same hardware.

5. Long-Sequence Prefill Parallelism: Context Parallelism and Load Balancing

For long sequences, Prefill compute time directly determines the time-to-first-token that users perceive. Traditional tensor parallelism (TP) splits along the attention-head dimension, so every device still processes the full sequence. When sequence length reaches hundreds of thousands or even millions of tokens, per-device compute and activation memory pressure become severe.

Sequence-dimension splitting. Context Parallelism (CP) splits inputs along the sequence dimension. When CP is combined with TP, the long sequence is partitioned across CP groups, and each group further applies TP internally. This combination reduces both per-device activation memory and compute, making ultra-long-sequence Prefill feasible within acceptable time.

Fine-grained load balancing. DSA’s causal index computation has a notable property: the later a position is in the sequence, the longer its visible history and the greater its index computation cost; earlier positions cost less. Naive sequential splitting produces large load gaps across devices—the “straggler effect,” where total time is dictated by the slowest device.

The fix is a Zigzag layout strategy: pair head and tail segments of the sequence on the same device. For example, the beginning of the sequence (short history, low compute) is paired with the end (long history, high compute) so that each device’s total index workload is roughly balanced. After computation, the original sequence order is restored for downstream processing. This strategy effectively mitigates load imbalance in causal sparse attention without changing compute semantics.

6. Automated Code Generation for Operator Development

Optimizing KDA Prefill for ultra-long sequences requires a CP-aware fused kernel. Hand-writing low-level kernel code takes long and iterates slowly. In this work, a DSL-based automated code generation approach was adopted (tools like CANNBot-DSL, a class of operator development tools targeting specific hardware architectures).

The workflow can be summarized as:

1. Describe the compute logic. The developer defines the KDA operator’s semantics under CP mode at a high level, including the sequence splitting scheme, boundary handling for state recurrence, and the placement of cross-device communication points.
2. Generate kernel code automatically. The tool combines the compute description with the target hardware’s memory hierarchy and instruction characteristics to produce an optimized low-level implementation.
3. Validate and tune. The generated code is verified for correctness and performance in real tests; the description is adjusted and regenerated as needed.

The core value of this approach: when the model architecture evolves again, updating the high-level description yields new operator implementations, greatly lowering adaptation cost. The generator can also internalize hardware-specific optimization experience (such as cache block sizes, pipeline depth, and memory-access coalescing strategies) into code-generation rules, so even non-experts can achieve performance close to hand-tuned kernels.

7. Coordinated Scheduling and Storage Optimizations

Beyond compute-level optimizations, inference efficiency also depends on scheduling and storage strategies.

At the scheduling level, multi-stream orchestration overlaps computation with communication. Expert dispatch and result aggregation in MoE involve cross-device communication; placing these communication operations on parallel streams alongside local computation hides communication latency. The ordering between lightweight operators (such as index computation) and heavyweight ones (such as attention) also needs careful arrangement so that compute units do not idle waiting for data.

At the storage level, optimized weight layouts (such as NZ layout) are used to match the hardware’s preferred memory-access pattern for matrix multiplication. Inference involves extensive weight-activation matrix multiplications, and how weights are arranged in memory directly affects access efficiency and compute-unit utilization. Re-arranging weights into the target hardware’s preferred blocking scheme yields substantial performance gains without changing compute semantics.

Additionally, compile caching and graph-mode optimization help at inference startup. Caching already-compiled computation graphs avoids repeated graph optimization and code generation overhead, which is particularly suitable for online services that require frequent cold starts.

8. Speculative Decoding and State Replay: Further Decode Optimization

Decode generates tokens step by step—each step involves relatively little computation, but there are many steps and latency is sensitive. Two optimization strategies target this characteristic.

Speculative decoding. The basic idea is to speculatively compute multiple plausible next tokens at each step; when the speculation is correct, several tokens are accepted at once, reducing the effective number of steps. This strategy is particularly well suited to models using sparse attention, because the inherent uncertainty of sparse computation provides more tolerance for speculation.

State replay. In the state-recurrence module, each step’s state update depends on the previous state. If state snapshots are saved at key steps, then when a rollback or recomputation is needed, execution can resume directly from a snapshot instead of recursing from the beginning. Used together with speculative decoding, this mechanism enables fast recovery after failed speculation, reducing the cost of wrong paths.

9. Summary and Outlook

This article has presented, from an engineering perspective, an inference optimization scheme for a new model architecture combining sparse attention, state recurrence, and mixture-of-experts mechanisms. The key takeaways are:

• Day-0 adaptation hinges on abstraction and reuse, not ground-up development. Organizing existing optimizations by compute pattern rather than by specific model drastically shortens adaptation time for new models.
• Operator fusion is an effective way to cut memory traffic, especially in multi-step pipelines like sparse attention, where fusing index computation eliminates global round-trips of intermediate results.
• Layered quantization should be differentiated by data characteristics. The differences in precision sensitivity and access patterns between index data and the KV cache call for different quantization formats and granularities.
• Parallel strategies must be designed around compute-pattern unevenness. Simple uniform splitting can cause severe load imbalance in causal sparse attention scenarios.
• Automated code generation is becoming an important direction for operator development, allowing optimization experience to accumulate and propagate as rules rather than relying on the manual work of a few experts.

As model architectures continue to evolve, inference optimization methodology must keep updating as well. But the core principle remains unchanged: understand the structural characteristics of the compute pattern, then choose or design the optimization strategy that best matches it. We hope the practices shared in this article provide a valuable reference for readers facing similar challenges.