Inference Acceleration for a New Domestic LLM Architecture: From Day-0 Adaptation to Full-Pipeline Performance Optimization

As large model architectures continue to evolve, hybrid attention and sparsification designs have become a mainstream approach to balancing quality and inference cost. A new-generation model builds on its predecessor’s accumulated work by introducing a hybrid structure that interleaves KDA and DSA, and achieved Day-0 inference support and performance optimization on heterogeneous compute platforms at launch. This article walks through that adaptation process—covering architectural characteristics, operator optimization, scheduling and orchestration, quantization strategy, and long-sequence parallelism—to outline an engineering methodology for inference acceleration targeting complex new architectures.

1. Architecture Analysis: Why KDA and DSA Require Dedicated Adaptation

The core change in this model lies in the hybrid design of the attention layer. KDA processes sequence information through state recurrence, exhibiting linear complexity characteristics and suiting long-range dependencies. DSA builds on multi-head latent attention (MLA) by introducing an index selection mechanism: a lightweight indexer (LI) first scores historical positions for relevance and selects the Top-K positions, after which sparse flash attention (SFA) reads only the KV of the selected positions to complete computation. In addition, the model continues to aggregate representations across layers via attention residuals.

This combination of “dense projection + sparse read + state recurrence” imposes three new requirements on the inference engine:

• Finer operator granularity: Index scoring, Top-K selection, and sparse reads form separate stages; launching operators one by one would significantly amplify scheduling overhead.
• More complex data paths: The index cache and the main attention KV cache use different precisions and layouts, with markedly different memory access patterns.
• More uneven load distribution: Causal index computation is inherently imbalanced—the further along the sequence, the longer the history—making long tails likely when splitting across multiple cards.

Day-0 support is therefore not simple operator integration but requires coordinated design across computation, communication, scheduling, and storage.

2. Leveraging Accumulated Work: Operator Fusion and Execution Orchestration

Rapid adaptation of the new model relied heavily on optimization components accumulated during the predecessor phase. The engineering effort mainly reused the following capabilities:

1. State-recurrence operator optimization: For KDA’s recurrent computation, a block-fused operator implementation merges state updates with matrix operations, reducing intermediate tensor writes to memory.
2. Two-stage fused attention residuals: Cross-layer representation aggregation is split into local aggregation and global reduction stages, lowering the frequency of cross-layer data movement.
3. Mixture-of-experts and communication fusion: Through operator-level fusion and communication overlap, expert-computation load balancing and collective communication are orchestrated within the same pipeline.
4. Execution-layer optimization: Combining hybrid parallelism strategies, communication optimization, and multi-stream orchestration with SuperKernel, graph-mode execution, compilation caching, and NZ weight layout shortens the end-to-end execution path.
5. Speculative inference and state replay: Speculative sampling accelerates decoding, while a state replay mechanism ensures consistency of recurrent states under batching and parallel partitioning.

Together these capabilities form the “infrastructure” for adapting to the new architecture, allowing the team to focus on the incremental optimization introduced by DSA.

3. DSA Integration and Multi-Stream Coordination: Reducing Decoding Latency

DSA’s computation chain can be summarized as a two-stage hop: “index filtering → sparse read.” The optimization focus is on reducing intermediate result write-backs and operator launch counts.

On the index side, LI fuses relevance scoring, head-weight weighting, and Top-K selection into a single execution that reads the paged K cache directly, avoiding intermediate tensors at each stage. After fusion, the index stage converges from multiple operator dispatches into a single kernel launch—particularly critical for decoding, where each step has only a millisecond-level budget. On the read side, SFA locates the main attention KV based on the selected indices and performs attention computation only on necessary positions, tying computation volume to the actual effective context. Both the prefill and decoding stages support BF16 and quantized paths for the indexer, allowing flexible switching based on precision requirements.

Furthermore, the index pre-stage achieves parallelism through three-stream coordination: Q projection and K normalization on one stream, Q quantization and head-weight projection on another, with multiple streams advancing in parallel to significantly shorten serial waiting in the input preparation stage. Illustrative pseudocode:

# Three-stream coordinated index pre-stage (illustrative)
with stream_qk, stream_quant, stream_weight:
    stream_qk.launch(q_proj_and_k_norm, x)
    stream_quant.launch(q_quantize, q)
    stream_weight.launch(head_weight_proj, w)
sync([stream_qk, stream_quant, stream_weight])
li_out = fused_lightning_indexer(q_quant, k_cache_paged, head_weight)
topk_idx = topk_select(li_out, k)
sfa_out = sparse_flash_attention(q, kv_cache, topk_idx)

4. Layered Quantization: Reducing Cache and Data Movement Overhead

The benefits of sparse attention come not only from reduced computation but also from lower memory access pressure, and quantization is key to amplifying this benefit. The approach adopts a layered strategy:

• Projection computation retains BF16: The indexer’s projection part is sensitive to numerical precision, so high precision is maintained to avoid degradation in selection quality.
• Index Q/K uses MXFP8 dynamic quantization: Grouped per-group with a group size of 32—every 32 elements share one scaling factor—compressing index data volume while preserving index scoring discrimination.
• Main attention KV cache uses FP8: Quantization granularity is per-token-head-tile with a tile size of 128, meaning each token and each head is quantized in blocks of 128 elements, balancing memory access efficiency and precision stability.

The point of layering is to “allocate precision by sensitivity”: precision-sensitive projections stay in BF16, while bandwidth-sensitive cache and index data are aggressively compressed, reducing storage footprint and data movement overhead without significantly sacrificing quality.

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

In long-input scenarios, per-card activations and intermediate tensors can easily become bottlenecks. This support combines context parallelism with tensor parallelism: long inputs are split along the sequence dimension, and sequence parallelism is coordinated within tensor-parallel groups to reduce per-card memory usage.

To address the load imbalance in DSA causal index computation, a Zigzag layout is used: two sequence segments paired head-to-tail are assigned to the same card, so that “longer history” and “shorter history” computation volumes offset each other, balancing cross-card index computation. After computation, the original sequence order is restored to ensure correctness. The essence of this strategy is trading data layout for load balance, preventing individual cards from becoming pipeline bottlenecks.

6. KDA Prefill Optimization for Ultra-Long Sequences

In ultra-long sequence scenarios, KDA’s prefill computation time grows linearly with sequence length, directly affecting first-token latency. When splitting tensor parallelism only by attention head, each card still processes the full sequence, yielding limited splitting benefit. To address this, the solution implements a context-parallel version of KDA based on a domain-specific language (DSL): inputs are further split along the sequence dimension, distributing long-sequence computation across multiple cards, reducing first-token latency while also lowering per-card activation and intermediate tensor usage.

7. Practical Takeaways

Several general engineering lessons for new architectures can be distilled from this adaptation process:

• Day-0 support depends on accumulation, not ad-hoc adaptation: Long-term investment in operator libraries, parallelism strategies, graph execution, and compilation caching determines whether a new model can launch quickly.
• Fusion should target “launch counts” and “intermediate write-backs”: Decoding bottlenecks often lie not in compute power but in scheduling and memory access.
• Quantization should be layered, not one-size-fits-all: Allocating precision by module sensitivity is the only way to obtain both compression ratio and stability.
• Parallelism must solve both splitting and balancing: Sequence splitting addresses capacity, while layout strategies address long tails—both are indispensable.

For teams looking to deploy complex hybrid-attention models on heterogeneous compute platforms, the combination of “operator fusion + multi-stream orchestration + layered quantization + context parallelism” offers strong reusability. As model structures continue to evolve, the competitiveness of inference engines will increasingly depend on this kind of systemic, coordinated optimization capability rather than single-point operator performance.