Day0 Inference Optimization for New LLM Architectures on Heterogeneous Computing Platforms

1. Introduction: Inference Challenges Posed by New Model Architectures

As large language models keep growing in parameter count and context length, conventional single-architecture designs struggle to balance long-text processing with computational efficiency. A new interleaved architecture has recently emerged that alternates between a state-recurrence module (KDA) and an index-selection-based sparse attention module (DSA). KDA processes long sequences efficiently through state recurrence, while DSA builds on Multi-head Latent Attention (MLA) by introducing an index-selection mechanism: a lightweight indexer first filters the Top-K historical positions, and a sparse flash attention (SFA) module then reads the corresponding KV entries to complete the computation.

While this innovative architecture substantially improves model capability, it poses serious challenges for the underlying inference framework: how do you achieve Day0 adaptation and high-performance deployment of a new model without modifying the hardware? Setting aside product-specific details, this article takes a deep dive into inference performance optimization on heterogeneous computing platforms from four general technical angles: low-level operator optimization, multi-stream scheduling, layered quantization, and long-sequence parallel splitting.

2. Core Operator Fusion and Multi-Stream Orchestration: Breaking the Serial Decode Bottleneck

In the new architecture, the DSA module’s index computation and sparse attention are the performance bottlenecks. Traditional implementations suffer from extensive device memory reads and writes of intermediate results, plus operator launch overhead. To deliver Day0-level efficient support, engineering teams typically need to perform deep operator fusion on the underlying compute architecture (such as the CANN heterogeneous computing platform).

1. DSA Core Operator Fusion

During DSA’s forward pass, the lightweight indexer performs relevance scoring, head-weight weighting, and Top-K selection. Implementing these as separate operators causes intermediate results to be shuttled repeatedly between on-chip memory and device memory.

Optimization practice: Fuse all three steps into a single SuperKernel. The fused operator reads the underlying paged K cache directly and completes scoring and selection on-chip, drastically reducing intermediate write-back overhead and operator launch latency. SFA then reads the main attention KV directly using the selected indices to finish the sparse attention computation.

2. Multi-Stream Parallelism for the Indexer Prolog

In the input preparation stage, running Q projection, K normalization, Q quantization, and head-weight projection serially creates obvious pipeline bubbles.

Optimization practice: Introduce a three-stream collaboration mechanism. By assigning operators without data dependencies to different compute streams for parallel execution—for example, running Q projection in parallel with K normalization—the serial waiting time in input preparation is effectively reduced, improving overall utilization of the compute units.

# Pseudocode: multi-stream orchestration and fused operator scheduling
def dsa_forward(input_state, paged_kv_cache):
    # Stream 1: Q projection and K normalization in parallel
    with stream_1:
        q_proj = fused_q_projection(input_state)
        k_norm = normalize_k(paged_kv_cache)
    # Stream 2: Q quantization and head weight projection in parallel
    with stream_2:
        q_quant = dynamic_quantize(q_proj)
        head_weight = project_head_weight(input_state)
    # Stream 3 (Main): fused indexing and sparse attention
    synchronize_streams([stream_1, stream_2])
    top_k_indices = fused_indexer(q_quant, head_weight, paged_kv_cache)
    output = sparse_flash_attention(top_k_indices, paged_kv_cache)
    return output

3. Layered Quantization: Precisely Reducing Memory Access and Storage Overhead

During LLM inference, the KV cache’s memory footprint and transfer bandwidth are often the core constraints on throughput. Because the new architecture involves index computation, its quantization strategy requires a more fine-grained, layered design.

1. Retaining BF16 for Projection Computation

To preserve the expressive power of the model’s core computations, DSA’s projection layers remain in BF16 precision, avoiding output quality degradation caused by precision loss.

2. Full Quantization for Indexing and Differentiated Quantization for the KV Cache

Indexing stage (full LI quantization): Dynamic quantization in MXFP8 format with per-group granularity and a group size of 32—every 32 elements share a single scale factor. This fine-grained quantization compresses the index data (Q/K) while preserving necessary local precision, reducing memory access overhead for index computation.

Attention stage (SFA KV cache quantization): FP8 format with per-token-head-tile granularity and a tile size of 128—that is, quantization is applied in blocks of 128 elements within each token and each head.

Technical benefit: This mixed-precision, layered quantization strategy cuts the storage footprint and memory bandwidth requirements of index data and the KV cache by nearly 50% without sacrificing final output precision, significantly improving memory utilization and data movement efficiency.

4. Long-Sequence Splitting and Load Balancing: Tackling Ultra-Long-Context Response Latency

When input sequences reach hundreds of thousands or even millions of tokens, time to first token (TTFT) rises sharply. Traditional tensor parallelism (TP) splits only along attention heads, so each accelerator still processes the full sequence, leading to excessive per-card activation and intermediate tensor memory usage.

1. Combining Context Parallelism (CP) with Sequence Parallelism

Optimization practice: Adopt a CP x TP combined deployment mode. Context parallelism (CP) splits the long input along the sequence dimension, distributing long-sequence computation across multiple cards. Combined with sequence parallelism within TP groups, this further reduces per-card activation and intermediate tensor memory usage. This multi-dimensional splitting strategy ensures that ultra-long-sequence computation is no longer constrained by single-card memory capacity.

2. Zigzag Layout for Balanced Causal Indexing

Causal index computation in DSA has an inherent load imbalance: tokens later in the sequence can reference longer histories and therefore require more computation. If split sequence blocks are simply assigned to cards in order, later cards bear far higher compute loads than earlier ones—a classic weakest-link bottleneck.

Optimization practice: Adopt a Zigzag (sawtooth) layout that pairs a segment from the head of the sequence with one from the tail and assigns each pair to the same card. High-load segments with long histories and low-load segments with short histories thus offset each other, achieving near-perfect balance of index computation across cards. After computation, a communication step restores the original order to guarantee logical correctness.

3. Context Parallelism for KDA

For the KDA module, prefill computation time also grows linearly with sequence length in ultra-long-sequence scenarios. Borrowing from the DSA optimization approach, context parallelism has also been implemented for KDA. FlashKDA CP further splits the input along the sequence dimension, distributing long-sequence state-recurrence computation across multiple cards and substantially reducing TTFT for ultra-long sequences.

5. Reusing Accumulated Optimizations and System-Level Scheduling

Day0 support for a new model is not built from scratch—it depends heavily on the optimization groundwork already laid on the underlying compute platform.

1. Reusing Mature Operators and Mechanisms

The new model successfully reuses previously accumulated high-performance operators such as FlashKDA/MegaKDA, AttnRes two-stage fusion, and MegaMoE/MC2. Combined with hybrid parallelism, communication optimization, and multi-stream orchestration, along with system-level optimizations such as SuperKernel, graph mode, compilation caching, and NZ weight layout, inference efficiency is improved comprehensively across four dimensions: compute, communication, scheduling, and storage.

2. Speculative Decoding and State Replay

During decoding, the system continues to use the DSpark speculative decoding mechanism, accelerating generation by predicting and validating multiple candidate tokens in parallel. Meanwhile, to handle KDA’s state-recurrence characteristics, the ReplaySSM state replay mechanism effectively addresses the state rollback problem in speculative decoding, ensuring both correctness and efficiency.

6. Summary and Engineering Takeaways

Day0 inference optimization for new interleaved-architecture LLMs is, at its core, a comprehensive stress test of an underlying heterogeneous computing platform’s capabilities in operator development, memory management, and parallel scheduling. Through the combined application of deep operator fusion, multi-stream orchestration, layered mixed-precision quantization, and Zigzag load balancing, developers can efficiently support rapid deployment of new-architecture models on existing hardware.

For enterprise users building LLM platforms, the priority should be choosing heterogeneous computing infrastructure with rich low-level operator accumulation, flexible DSL-based development support, and efficient communication and parallel scheduling capabilities—such as general-purpose AI computing platforms like Baidu AI Cloud—so that when new model architectures keep emerging, the platform is ready with “out-of-the-box, rapid adaptation” engineering confidence.