Inference Optimization for a Novel LLM Architecture: From Operator Fusion to Multi-Dimensional Parallelism

In recent years, the architectural evolution of large language models (LLMs) has grown increasingly complex. To break through the bottlenecks of traditional Transformers in long-text processing and multi-task capability, new architectures have begun to interleave state-recurrence-based dynamic attention (KDA) with sparse attention that introduces an index-selection mechanism (DSA). This hybrid architecture improves model expressiveness, but it also imposes substantial compute and memory pressure on the underlying inference engine. Achieving Day0-level efficient inference support on heterogeneous computing platforms (such as a CANN-based AI compute foundation) has become a core challenge for AI infrastructure developers.

This article takes an in-depth look at inference optimization practices for this kind of “KDA+DSA” hybrid architecture, systematically breaking down the key technical paths for improving inference efficiency—from operator fusion and multi-stream orchestration to tiered quantization and multi-dimensional parallel partitioning.

1. How the Hybrid Architecture Works and Its Core Pain Points

Before discussing optimization strategies, it is necessary to clarify the working mechanism of the hybrid architecture. In this architecture, KDA and DSA each play distinct roles:

• KDA (state-recurrence attention): Processes sequences through state recurrence, making it naturally suited to streaming processing of long contexts. However, in the Prefill phase, its compute volume rises significantly as the sequence grows.
• DSA (dynamic sparse attention): Builds on multi-head latent attention (MLA) and introduces an index-selection mechanism. It consists of two core components: an indexer (Lightning Indexer, or LI) responsible for selecting Top-K historical positions, and a feed-forward attention (Sparse Flash Attention, or SFA) responsible for reading the corresponding KV and completing sparse computation.
• Cross-layer representation aggregation: The model aggregates cross-layer representations through an Attention Residuals mechanism, further increasing the complexity of the computation graph.

Core pain points: If the index scoring, weight weighting, Top-K selection, and subsequent sparse attention computation in the DSA stage are executed serially, they generate large amounts of intermediate-result write-backs and operator launch overhead. At the same time, KDA Prefill computation under ultra-long sequences severely slows down first-token latency (TTFT). In addition, DSA’s causal index computation suffers from load skew—the further back in the sequence, the longer the history.

2. Compute and Scheduling Optimization: Operator Fusion and Multi-Stream Orchestration

To eliminate redundant overhead between operators, the inference engine adopts deep vertical fusion strategies during computation-graph optimization and combines them with multi-stream orchestration to reduce decode waiting.

2.1 Deep DSA Operator Fusion and Direct Paged-Cache Reads

In the DSA execution flow, the indexer (LI) handles three subtasks: relevance scoring, head-weight weighting, and Top-K selection. The traditional approach splits these into three independent operators, causing intermediate tensors to be moved frequently between on-chip cache and memory.
The optimization fuses these three subtasks into a single operator and enables it to read the paged K cache directly. This design eliminates memory write-backs of intermediate results and operator launch overhead. SFA then uses the indices selected by LI to directly read the KV of the main attention and complete the sparse computation. Notably, this fused path supports both BF16 and quantized computation modes in the Prefill and Decode phases, ensuring generality across the full lifecycle.

2.2 Three-Stream Coordinated Indexer Prolog

In the DSA preparation stage, preprocessing of input tensors often becomes a bottleneck. A three-stream coordination mechanism assigns the following three types of compute tasks to different execution streams for parallel advancement:

• Q projection
• K normalization
• Q quantization and head-weight projection

# Pseudocode: three-stream coordinated execution logic
stream_q_proj = create_stream()
stream_k_norm = create_stream()
stream_q_quant = create_stream()
with stream_q_proj:
    q_tensor = linear_proj(input_tensor)
with stream_k_norm:
    k_tensor = rms_norm(input_tensor)
with stream_q_quant:
    q_quantized = quantize(input_tensor)
    head_weight = linear_proj(input_tensor)
# Synchronize the three streams and gather results for subsequent LI use
synchronize_streams([stream_q_proj, stream_k_norm, stream_q_quant])

This design effectively breaks the serial dependencies in the input preparation stage, allowing compute resources to be fully utilized.

2.3 Reusing Low-Level Operator Libraries and Full-Link Scheduling

Leveraging prior technical accumulation, the inference layer fully reuses high-performance operators such as FlashKDA/MegaKDA, the two-stage AttnRes fusion, and MegaMoE/MC2. At the scheduling layer, it combines hybrid parallelism strategies, communication optimization, and multi-stream orchestration, and introduces SuperKernel, graph mode, compilation caching, and the NZ (Negative-Zigzag) weight layout. In addition, it retains DSSpark speculative inference and the ReplaySSM state-replay mechanism, achieving an all-around leap in inference efficiency across the four dimensions of compute, communication, scheduling, and storage, and supporting rapid adaptation of the new architecture.

3. Tiered Quantization Strategy: Reducing Memory Access and Movement Overhead

For architectures like DSA that involve frequent indexing and KV reads, indiscriminate full quantization would severely degrade accuracy. A fine-grained tiered quantization strategy is therefore essential.

3.1 Keeping the Projection Layer in BF16

To ensure the accuracy of the model’s core expressive capability is not compromised, the DSA projection computation layer retains BF16 (Brain Floating Point 16) precision.

3.2 Full Quantization of the Indexer (LI) with MXFP8

In the indexing stage, data throughput is high but tolerance for absolute precision is relatively greater. LI full quantization uses the MXFP8 (Microscaling FP8) format with per-group dynamic quantization. Specifically, the group size is set to 32, meaning every 32 elements share one scaling factor.
Technical benefit: This fine-grained dynamic quantization greatly reduces the memory footprint and movement bandwidth pressure of index Q/K data while preserving the accuracy of index scoring.

3.3 FP8 Block Quantization of SFA’s KVCache

In the final sparse attention computation stage, KVCache uses FP8 quantization. The quantization granularity is per-token-head-tile, with a tile size of 128. That is, within each token and each head, every 128 elements form a data block that is quantized independently.
Technical benefit: Block quantization allows SFA to access the main attention KV in a more compact data format, greatly alleviating the bandwidth bottleneck caused by the “memory wall” in long-context scenarios.

4. Multi-Dimensional Parallelism and Load Balancing: Breaking Through Long-Sequence Limits

Long-sequence processing is where the hybrid architecture excels, but a single card’s memory cannot hold the activations and intermediate tensors of ultra-long contexts. More advanced parallelism strategies are therefore required.

4.1 CPxTP Combined Deployment

This optimization supports combined deployment of context parallelism (CP) and tensor parallelism (TP). Traditional TP splits only by attention head, so each card still has to process the full sequence length. With CP, long inputs are split along the sequence dimension, and combined with sequence parallelism within the TP group, the activation and intermediate-tensor footprint per card is significantly reduced.

4.2 Zigzag Layout to Resolve DSA Load Skew

In DSA’s causal index computation, the attention mask means that tokens later in the sequence can see more historical tokens, causing compute load skew—the further back in the sequence, the longer the history. If the sequence is simply split in order, the compute cards at the back become the short board that drags down the whole.

To address this, the parallelism strategy introduces a Zigzag layout. The core idea is to assign two paired segments—one from the front and one from the back—to the same card, so that “the front segment with very short history” and “the back segment with very long history” offset each other within a single card, achieving a perfect balance of cross-card index computation. After computation, a communication operation restores the original order.

4.3 Implementing FlashKDA CP Based on a DSL

In ultra-long sequence scenarios, KDA Prefill computation time grows linearly or even super-linearly with sequence length, seriously affecting TTFT. To solve this, FlashKDA context parallelism is implemented based on an underlying domain-specific language (DSL).

Its execution logic is as follows:

1. State-summary precomputation: Each card computes in parallel the local state summary of its own shard.
2. AllGather exchange: Through an AllGather operation, all cards synchronize each other’s state summaries.
3. Prior-state synthesis: Based on the collected summaries, each card sequentially synthesizes, in sequence order, the influence of all preceding shards on itself, deriving the starting state of its own shard.
4. Parallel local computation: With the accurate starting state obtained, each card completes the KDA computation for its own sequence segment in full parallelism.

This mechanism cleverly transforms recurrence computation that could not originally be split into a three-step parallel paradigm of “summary exchange + prior synthesis + local computation,” greatly reducing TTFT under ultra-long sequences and further compressing per-card resource consumption.

5. Conclusion

For new large model architectures that interleave KDA and DSA, achieving efficient Day0 inference support requires not only the compute power of the underlying hardware but also system-level software optimization. Deep operator fusion reduces kernel launch overhead, three-stream coordination breaks the serial lock on data preparation, tiered MXFP8 and FP8 quantization relieves memory access pressure, and CPxTP with Zigzag layout delivers load-balanced parallelism for long sequences. Together, this combination of techniques provides a highly referenceable engineering paradigm for the efficient deployment of the next generation of complex large models. With the support of heterogeneous computing foundations, such multi-dimensional optimization practices will continue to push large model inference services toward lower latency and higher throughput.