Adapting and Optimizing Inference for a New Model Architecture on Ascend

1. Background: Inference Adaptation Challenges in a Rapidly Evolving Architecture Landscape

Large language model architectures are evolving from single attention mechanisms toward hybrid designs. A recently released next-generation model introduces a notable shift from its predecessor by interleaving two attention paths. One path handles sequence information through state recurrence. The other builds on multi-head latent attention and adds an index selection mechanism: a lightweight indexer selects the Top-K historical positions, and a sparse attention module then reads the corresponding KV entries to complete the computation. The model also continues to use an attention residual mechanism to aggregate cross-layer representations.

These architectural innovations raise a core engineering question: How can Day-0 inference support and usable performance be achieved on the Ascend platform within a limited adaptation window? The answer is not to build from scratch. Instead, it lies in systematically reusing operator optimizations, parallel strategies, and scheduling mechanisms accumulated from the previous-generation model, then making targeted enhancements for the new architecture’s differentiating characteristics.

2. Reuse and New Development: Overall Adaptation Framework

2.1 Reusable Optimization Assets

Adapting the new model does not mean starting over. Several optimizations from the previous generation can be carried over directly:

• State recurrence operators: The fused FlashKDA and MegaKDA implementations have already been optimized at the operator level for the recurrence path.
• Attention residual fusion: The two-stage fusion approach can continue to reduce memory access overhead in cross-layer aggregation.
• Mixture-of-experts optimizations: MegaMoE and communication-computation overlap mechanisms remain applicable to the new model.
• Compilation and scheduling layer: SuperKernel, graph mode, compilation cache, and NZ weight layout are platform-level capabilities that can be inherited naturally.
• Speculative inference and state replay: The DSpark speculative inference mechanism and ReplaySSM state replay continue to provide acceleration during decoding.

Reusing these assets compresses the adaptation effort from full development to incremental adjustment, which is the key prerequisite for achieving Day-0 support.

2.2 New Adaptation Points

The new architecture’s differentiating features are concentrated in the index selection and sparse attention paths. Specifically:

• The indexer needs to fuse relevance scoring, head-weighted scoring, and Top-K selection into a single execution flow.
• The sparse attention module must read the main attention KV entries based on the index results and perform sparse computation.
• Both paths need to support BF16 and quantized precision modes in both prefill and decode stages.

The adaptation effort therefore focuses on indexer fusion, sparse attention integration, and the corresponding quantization strategy.

3. Operator Fusion and Multi-Stream Orchestration: Reducing Decode Latency

3.1 Fused Indexer Execution

The indexer is responsible for the selection step: computing relevance scores, applying per-head weighting, and selecting the Top-K positions. A naive implementation would perform these three steps with multiple intermediate write-backs and operator launch overheads. The optimized approach fuses all three into a single execution unit and reads the paged K cache directly, avoiding explicit materialization of intermediate tensors.

# Illustrative pseudocode: fused indexer execution logic
def fused_indexer(query, paged_k_cache, head_weights, topk):
    # Relevance scoring (reads paged cache directly to avoid materializing full K)
    scores = matmul(query, paged_k_cache, transposed=True)
    # Head weighting + Top-K selection in a fused execution
    weighted = scores * head_weights
    topk_indices = topk_select(weighted, k=topk)
    return topk_indices

The sparse attention module then reads the main attention KV entries according to the selected indices and performs attention computation only on the Top-K positions. This reduces the computation from the full sequence length to a sparse subset.

3.2 Three-Stream Coordination with the Indexer as the Lead

The input preparation stage has serial dependencies: K normalization can only proceed after Q projection completes, and head weight projection can only start after Q quantization finishes. The optimized approach splits this into three parallel streams:

• Stream A: Q projection
• Stream B: K normalization
• Stream C: Q quantization and head weight projection

These three streams execute concurrently, allowing the originally serial input preparation stage to overlap and reducing the fixed overhead in time-to-first-token (TTFT).

4. Layered Quantization: Reducing Cache and Data Movement Overhead

The quantization strategy needs to distinguish between computation-sensitive and memory-sensitive paths.

4.1 BF16 for Projection Computation

Projection computation in the sparse attention path is relatively sensitive to numerical precision. It remains in BF16 to prevent quantization error from being amplified during sparse selection.

4.2 Full Quantization for the Index Path

Index Q/K tensors use MXFP8 with per-group dynamic quantization and a group size of 32, meaning every 32 elements share a single scaling factor. This granularity balances precision and compression ratio, significantly reducing storage and memory access overhead for index data.

4.3 FP8 Quantization for KV Cache

The main attention KV cache uses FP8 with per-token-head-tile quantization granularity and a tile size of 128. Each token and each head is quantized in blocks of 128 elements. This strategy keeps attention evaluation stable while substantially reducing KV cache memory usage.

Path	Precision	Quantization Granularity	Main Benefit
Projection computation	BF16	—	Ensures numerical stability
Index Q/K	MXFP8	per-group, size=32	Reduces index memory access
Main attention KV	FP8	per-token-head-tile, size=128	Compresses cache usage

5. Parallel Strategy: Long-Sequence Partitioning and Load Balancing

5.1 Combined CP and TP Deployment

This adaptation supports combined CP×TP deployment: long inputs are partitioned along the sequence dimension, combined with sequence parallelism within tensor parallel groups. This reduces per-device activation memory and intermediate tensor usage. For extremely long-sequence scenarios, this combination effectively relieves per-device memory pressure.

5.2 Zigzag Layout for Balancing Causal Index Load

Causal index computation in sparse attention has a natural load imbalance: positions later in the sequence have longer visible histories and therefore higher computational cost. With contiguous partitioning, the devices handling the later portion become the bottleneck.

The optimized approach uses a Zigzag layout: the first and last paired segments of the sequence are assigned to the same device, so each device handles both a long-history segment and a short-history segment. This balances the cross-device index computation load. After computation, the original order is restored to ensure correctness.

Contiguous partitioning (imbalanced):  [Device 0: 1..N/2] [Device 1: N/2..N]  → Device 1 has higher load
Zigzag layout (balanced):              [Device 0: 1..N/4 + 3N/4..N] [Device 1: N/4..3N/4] → Balanced load

5.3 Sequence Parallelism for the State Recurrence Path

For extremely long sequences, the prefill computation time of the state recurrence path grows linearly with sequence length, directly affecting TTFT. When only TP is used to partition by attention heads, each device still processes the full sequence.

A sequence parallelism approach based on domain-specific language support further partitions the input along the sequence dimension, distributing long-sequence computation across multiple devices. This reduces TTFT while also lowering per-device activation and intermediate tensor usage. Combined with the CP strategy described above, this forms a complete parallel solution for extremely long sequences.

6. Summary and Best Practices

Day-0 support for the new model on the Ascend platform is fundamentally an engineering methodology centered on reuse first, incremental development second:

1. Inventory reusable assets first: Platform-level capabilities such as operator fusion, compilation cache, and scheduling mechanisms should be inherited before anything else.
2. Focus on differentiating paths: Direct adaptation resources toward newly introduced structures such as index selection and sparse attention.
3. Prefer fusion over splitting: The indexer’s scoring, weighting, and selection should be executed in a fused manner to reduce operator launches and intermediate write-backs.
4. Apply quantization in layers: Keep BF16 for computation-sensitive paths and apply fine-grained quantization to memory-sensitive paths.
5. Match parallel strategies to load characteristics: Use Zigzag for causal index load balancing and sequence parallelism for distributing extremely long-sequence computation.

For developers, the value of this methodology lies not only in deployment efficiency for a single model, but in forming a transferable adaptation pattern. When the next architectural innovation arrives, teams can quickly determine what can be reused and what must be built new, enabling sustained Day-0 support and performance targets.