Day0 Inference Adaptation and Performance Optimization for a Sparse Hybrid Architecture LLM

1. Background: Inference Challenges Posed by a New Architecture

A new class of large language model architecture has recently emerged with a sparse hybrid attention design. During decoding, these models interleave two attention mechanisms. One processes sequence information through state recurrence. The other builds on multi-head latent attention (MLA) and introduces index selection: a lightweight Lightning Indexer (LI) first selects the Top-K historical positions, then Sparse Flash Attention (SFA) reads the corresponding KV entries to complete sparse attention computation. At the same time, the model continues to aggregate cross-layer representations through attention residuals.

The main value of this architecture is reducing computation and memory usage in long-sequence scenarios by making attention computation sparse. At the same time, it increases engineering complexity. Operator fusion, quantization strategy, parallel partitioning, and state management all need targeted adaptation. This article reviews Day0 inference support and performance optimization for this type of architecture from four dimensions: computation, communication, scheduling, and storage.

2. Sparse Attention Fusion and Multi-Stream Orchestration

2.1 Fused Execution of the Indexer and Sparse Attention

In a conventional implementation, index selection and sparse attention reads are separate stages: first compute relevance scores, then perform Top-K selection, and finally read KV according to the selected indices. This sequential approach introduces substantial intermediate result write-back and operator launch overhead.

The optimized approach fuses relevance scoring, head weighting, and Top-K selection into a single operator that directly reads the paged K cache, avoiding intermediate materialization. SFA then reads the main attention KV according to the selected indices and completes sparse attention computation. Both prefill and decode support BF16 and quantized paths, allowing developers to choose based on precision requirements.

2.2 Input Preparation with Three-Stream Coordination

The input preparation stage contains several parallelizable subtasks. With three-stream coordination, Q projection and K normalization are placed in one stream, while Q quantization and head weight projection are placed in another stream for parallel execution. This reduces serial waiting and is especially effective in the prefill stage, where input preparation accounts for a larger share of total latency.

# Pseudo-code: three-stream input preparation
with stream_q_proj:
    q = q_projection(hidden_states)
with stream_k_norm:
    k = k_normalize(k_cache)
with stream_quant:
    q_quant, head_weight = quantize_and_weight(q, k)
# Synchronize streams before attention computation
attention_output = sparse_attention(q_quant, k, head_weight)

3. Hierarchical Quantization: Reducing Cache and Data Movement Overhead

3.1 Quantization Granularity Design

The memory overhead of this sparse architecture comes mainly from two parts: indexer Q/K data and main attention KV cache. A hierarchical quantization strategy is applied according to data characteristics.

• Indexer path: Full quantization uses MXFP8 with per-group dynamic quantization for index Q/K. The group size is 32, meaning every 32 elements share one scaling factor. This fine-grained quantization preserves accuracy while reducing index data storage to half of BF16.
• Main attention KV cache: Uses FP8 with per-token-head-tile quantization. The tile size is 128, meaning quantization is performed in blocks of 128 elements within each token and each head. This granularity balances memory access efficiency and precision.

3.2 Keeping BF16 for Projection Computation

Notably, projection computation in the DSA path remains in BF16. This is because projection layers are sensitive to numerical precision, and quantization may introduce non-negligible error accumulation. The core idea of hierarchical quantization is to keep high precision for precision-sensitive computation while quantizing storage- and bandwidth-sensitive cache data.

The combined effect is reduced storage and memory access overhead for index data and KV cache. This is especially important for long sequences because KV cache size grows linearly with sequence length, and quantization can directly halve memory usage.

4. Sequence Parallelism and Load Balancing

4.1 CP×TP Combined Deployment

For long-sequence inputs, CP (Context Parallelism) × TP (Tensor Parallelism) combined deployment is supported. Long inputs are partitioned along the sequence dimension, and sequence parallelism is used within TP groups to reduce per-device activations and intermediate tensors. This combination covers both sequence-dimension partitioning and attention-head-dimension partitioning, making it suitable for very long sequences.

4.2 Zigzag Layout for Causal Load Imbalance

Causal index computation in sparse attention naturally creates load imbalance: positions later in the sequence have longer histories and therefore require more relevance score computation. If the sequence is simply partitioned in order, later devices will have much higher compute load than earlier devices.

The Zigzag layout pairs head and tail sequence blocks on the same device: device i processes the i-th block from the front and the i-th block from the back. Each device therefore contains one short-history block and one long-history block, balancing the cross-device index computation. After computation, the sequence order is restored, keeping the process transparent to upper-level logic.

Sequence: [0, 1, 2, 3, 4, 5, 6, 7]
Device 0: [0, 7]  Device 1: [1, 6]  Device 2: [2, 5]  Device 3: [3, 4]
Each device contains one head block and one tail block for balanced load

5. Sequence Parallelism for the State Recurrence Mechanism

5.1 TTFT Bottleneck for Very Long Sequences

The prefill latency of state-recurrence attention, such as KDA, grows with sequence length and directly affects time to first token (TTFT). When using only TP to partition by attention head, each device still processes the full sequence, so the sequence-dimension computation is not distributed.

5.2 FlashKDA CP Implementation

A sequence-parallel version of state-recurrence attention, FlashKDA CP, is implemented using a domain-specific language (DSL). The key challenge is that state recurrence has sequential dependencies, so the sequence cannot simply be partitioned and computed independently.

The implementation consists of three steps:

1. Parallel precomputation: Each device first performs parallel precomputation on its own sequence chunk to obtain a local state summary.
2. State exchange and composition: AllGather is used to exchange local state summaries across devices. The contribution of preceding chunks to the current chunk is composed in order, producing the correct starting state for each chunk.
3. Parallel completion: Each device completes the full computation for its chunk based on the correct starting state.

# Pseudo-code: FlashKDA CP three stages
# Stage 1: local precomputation
local_summary = precompute(local_chunk)
# Stage 2: AllGather + state composition
all_summaries = all_gather(local_summary)
start_state = compose_prefix_states(all_summaries, rank)
# Stage 3: parallel completion for this chunk
output = compute_with_state(local_chunk, start_state)

This approach distributes long-sequence computation across multiple devices to reduce TTFT for very long sequences while also lowering per-device activations and intermediate tensors.

6. System-Level Optimization: From Operators to Graph Compilation

In addition to the architecture-level optimizations above, several accumulated system-level techniques are reused:

• Operator fusion: Optimizations such as FlashKDA/MegaKDA, two-stage attention residual fusion, and MegaMoE/MC2 reduce operator launch overhead and intermediate data movement.
• Hybrid parallelism and communication optimization: Combined hybrid parallelism and communication optimization reduce overhead in distributed training and inference.
• Multi-stream orchestration: Computation and communication overlap is achieved through multi-stream scheduling.
• SuperKernel and graph mode: Multiple operators are fused into SuperKernel, combined with graph mode to reduce scheduling overhead.
• Compilation cache and NZ weight layout: Compilation cache accelerates repeated compilation scenarios, and NZ weight layout improves matrix multiplication efficiency.
• Speculative inference and state replay: Existing mechanisms such as DSpark speculative inference and ReplaySSM state replay are reused to improve inference efficiency across computation, communication, scheduling, and storage.

7. Practical Recommendations and Summary

For developers planning to deploy this type of sparse hybrid architecture model on Ascend, the following points are worth attention:

1. Quantization strategy: Choose between BF16 and quantized paths according to precision requirements. MXFP8 per-group quantization is recommended for the indexer path, and FP8 per-token-head-tile quantization for KV cache.
2. Parallel configuration: For long sequences, prioritize CP×TP combined deployment and enable Zigzag layout for load balancing. For very long sequences, enable the sequence-parallel version of state-recurrence attention.
3. Multi-stream orchestration: Make full use of three-stream coordination in the prefill stage to reduce serial waiting during input preparation.
4. Combining system-level optimizations: Operator fusion, graph mode, and compilation cache can be stacked for cumulative benefits.

Day0 support and optimization for this architecture are essentially systematic engineering across four dimensions: computation sparsification, storage quantization, parallel partitioning, and state management. Hierarchical quantization reduces memory access overhead, Zigzag layout and sequence parallelism address load balancing, and fused operators with multi-stream orchestration reduce scheduling stalls. Together, these techniques deliver acceptable TTFT and throughput in very long-sequence scenarios. As sparse hybrid architectures continue to evolve, this optimization framework also has the potential to be transferred to more architecture variants.