Day-0 Inference Adaptation and Performance Optimization for New-Architecture Large Language Models
New-architecture large language models that pair recurrent state updates with indexed sparse attention offer strong capabilities, but they also introduce heterogeneous compute paths, unbalanced long-sequence workloads, and heavy cache and data-movement overhead. This article describes a practical Day-0 adaptation playbook: reuse proven optimization assets from the previous model generation and apply targeted increments across operators, quantization, parallelism, and scheduling. It covers fusing the Top-K indexer into a single kernel, three-stream orchestration for input preparation, BF16/FP8 tiered quantization for index data and the KV cache, and Zigzag-layout sequence parallelism, and closes with actionable engineering recommendations for teams adapting similar architectures on the Ascend platform.
Background: Inference Challenges for New-Architecture Models
Recently, a class of new-architecture large models built on hybrid attention mechanisms has been released and open-sourced. While retaining cross-layer representation aggregation, these models introduce two complementary sequence-modeling paths: one processes sequence information through state recursion, and the other extends multi-head latent attention with an index selection mechanism—a lightweight indexer first filters out the Top-K historical positions, and a sparse attention module then reads the corresponding KV entries to complete the computation.
This design delivers accuracy benefits, but it creates three challenges for inference deployment:
- Heterogeneous compute paths: the state recursion branch and the sparse attention branch have very different computation graphs, making them hard to cover with a unified set of operators.
- Unbalanced long-sequence workloads: in causal index computation, the later a position appears in the sequence, the longer its history, so load skews naturally when sequences are split across devices.
- High cache and data-movement overhead: index data and the KV cache grow rapidly with sequence length, making memory access the bottleneck.
To achieve Day-0 support on release day and let developers deploy and try the model quickly on the Ascend platform, the key is to reuse existing optimization assets rather than rewrite from scratch. The following sections break this down across four layers: operators, quantization, parallelism, and scheduling.
Fast Adaptation by Building on Existing Assets
The new model does not appear in isolation; it shares much of its structural DNA with the previous-generation architecture. The core adaptation strategy is therefore to migrate the proven optimization assets from the previous generation and build incrementally only for the new structures.
Specifically, the reused optimizations include:
- Fused operators for the state recursion branch (merging multi-step recursion into a single large-kernel computation);
- The two-stage fusion strategy for cross-layer representation aggregation;
- Operator optimizations and communication fusion for Mixture-of-Experts (MoE) layers;
- The hybrid parallelism strategy, communication optimizations, and multi-stream orchestration;
- Large-kernel fusion, graph-mode execution, compilation caching, and weight layout optimization;
- Speculative inference and state replay mechanisms.
Together, these optimizations improve efficiency across compute, communication, scheduling, and storage, allowing the new architecture to reach inference efficiency close to the previous generation without re-tuning.
Fused Execution of Index Selection and Sparse Attention
Fusing the Indexer to Reduce Intermediate Write-Backs
The lightweight indexer originally involves three steps: relevance scoring, head-weight weighting, and Top-K selection. When executed step by step, each step writes intermediate results back to device memory and reads them back in, and operator launch overhead accumulates.
The optimization fuses the three steps into a single execution: it reads the paged K cache directly and performs scoring and selection inside the kernel, avoiding intermediate tensors being written to memory. The sparse attention module then reads the main attention KV using the selected indices and completes the sparse computation.
Notably, both the prefill and decode phases support two paths for the indexer—BF16 and quantized—so it runs under different precision requirements.
Three-Stream Orchestration for Index Preparation
The input preparation stage also suffers from serial waiting: Q projection, K normalization, Q quantization, and head-weight projection depend on one another. With three-stream orchestration, Q projection runs in parallel with K normalization, and Q quantization runs in parallel with head-weight projection. This turns the originally serial preparation stage into a parallel pipeline and reduces the wait before the first decode step.
Tiered Quantization: Cutting Cache and Data-Movement Overhead
The performance bottleneck of sparse attention often lies not in compute power but in the storage and memory access of index data and the KV cache. To account for differences in importance across data types, a tiered quantization strategy is applied:
| Data | Precision policy | Quantization granularity |
|---|---|---|
| Projection computation | BF16 | — |
| Index Q/K | Fully quantized | Per-group dynamic quantization, with every 32 elements sharing one scale factor |
| Main attention KV cache | FP8 | Per-token-head-tile, quantized in blocks of 128 elements |
This “keep precision for compute, shrink size for cache” approach significantly reduces the storage and memory-access overhead of index data and the KV cache with almost no accuracy loss.
Illustrative pseudocode:
# Per-group dynamic quantization for index Q/Kfor group in split(index_qk, group_size=32):scale = max(abs(group)) / quant_maxgroup_q = round(group / scale)# Block-wise quantization of the KV cache by token-head-tilefor tile in split(kv_cache, tile_size=128):scale = max(abs(tile)) / quant_maxtile_q = round(tile / scale)
Sequence Parallelism: Long-Sequence Splitting with Load Balancing
Combined Deployment to Reduce Per-Device Footprint
This release supports combined sequence parallelism and tensor parallelism: long inputs are split along the sequence dimension, while sequence parallelism is applied within tensor parallel groups, reducing per-device activation and intermediate tensor memory usage.
Zigzag Layout to Balance Causal Index Workloads
Sparse index computation is causal—the later a position is in the sequence, the longer the history it must retrieve. With a naive sequential split, the later devices would carry far more computation than the earlier ones.
The solution is a Zigzag layout: the first and last segments of the sequence are paired onto the same device, so each device handles both a “short history” and a “long history” portion, balancing the index computation across devices. The original order is restored after computation to guarantee correct results.
Original order: [Segment 1][Segment 2][Segment 3][Segment 4]Zigzag assignment:Card 0 <- Segment 1 + Segment 4Card 1 <- Segment 2 + Segment 3
Sequence Splitting for the State Recursion Branch on Ultra-Long Inputs
In ultra-long-sequence scenarios, prefill time for the state recursion branch grows with sequence length, directly affecting time to first token (TTFT). When tensor parallelism splits only across attention heads, each device still processes the full sequence.
Splitting the input further along the sequence dimension distributes the long-sequence computation across devices, which reduces TTFT for ultra-long sequences as well as per-device activation and intermediate tensor usage. This capability is implemented on a domain-specific DSL: each device first computes its local states in parallel, and the global state is then merged through cross-device communication.
Engineering Recommendations
For teams adapting similar new architectures on their own platforms, the following practice order is a useful reference:
- Align the operator checklist first: map out the structural differences between the new model and existing models, and develop incrementally only for the new structures.
- Prioritize fusion: merge short operators such as scoring, weighting, and selection to reduce launches and write-backs.
- Divide and conquer with quantization: keep precision on compute paths, compress size on cache paths, and choose granularity that matches the memory access pattern.
- Balance parallelism: long-sequence splitting must account for causal load skew; a paired layout is a low-cost, effective technique.
- Make scheduling multi-stream: break dependent serial stages into parallel streams to compress preparation waits.
Conclusion
Day-0 inference support for a new-architecture large model is essentially an asset-reuse problem, not a rewrite problem. By building on existing operator optimizations, introducing index fusion and tiered quantization, and adopting Zigzag sequence parallelism and multi-stream orchestration, gains can be achieved simultaneously across compute, communication, scheduling, and storage, enabling developers to quickly deploy and experience new models on the Ascend platform. As sparse attention and state-recursion structures continue to evolve, this “reuse plus incremental adaptation” methodology will remain the backbone of inference engineering efficiency.
