Engineering Inference for Ultra-Large MoE Models on Ascend Platforms

1. Background: When Model Scale Crosses a New Threshold

Large model capability evolution follows two main trajectories: continuous expansion of parameter scale and steady growth of context window length. A flagship model recently released by a leading team reaches approximately 2.8 trillion parameters, natively supports vision understanding, and offers a context window of up to one million tokens. While this class of models delivers significant capability gains in code generation, 3D creation, and financial research, it also places stringent demands on inference infrastructure. The core challenge is not whether the model can run at all, but whether it can run and remain usable on the first day of release. This article breaks down the key technical aspects of inference deployment for ultra-large MoE models based on a typical day-one adaptation effort on Ascend platforms.

2. Model Architecture Analysis: Computational Characteristics of Hybrid Linear Attention

2.1 Combining KDA and Sparse Attention

The model adopts a hybrid linear attention architecture that combines KDA (Kernelized Delta Attention, a linear attention variant based on kernelization techniques; related ideas can be found in work by Ma et al., 2024) with DSA (Dynamic Sparse Attention). Linear attention layers have computational complexity that scales linearly with sequence length, avoiding the O(n²) cost of standard attention. KDA uses a kernelized state matrix to implement RNN-like recurrent updates, enabling token-by-token progression during autoregressive decoding. However, recurrent precision can degrade over extremely long sequences, so engineering practice typically sets a precision tolerance threshold (e.g., per-step error below 1e-4) for monitoring. Sparse attention layers use filtering mechanisms to compute only over the historical tokens most relevant to the current query, preserving long-range dependency modeling while controlling computational cost.

These two computational graphs impose different requirements on operator fusion, memory layout, and cross-card communication. The inference engine must simultaneously manage recurrent state updates and sparse computation with filtering logic.

2.2 Orchestration Challenges in Multimodal Inference

The model natively supports vision understanding, which introduces image encoding and cross-modal alignment stages into the inference pipeline. The engineering difficulty lies in coordinating visual feature preprocessing, encoder forward computation, and the main language decoding loop. If encoder computation blocks the decoding loop, end-to-end latency rises sharply. Practical mitigation approaches include using pipeline parallelism to execute the encoder and decoder in separate stages and configuring asynchronous scheduling at the framework level to ensure visual features are ready before decoding starts.

3. Operator-Level Co-Optimization: From Generic Support to Path Customization

3.1 Early-Engagement Model Analysis

The key to day-one adaptation is starting early. Building on adaptation experience from the previous generation of models in the same series, the inference team conducted structural analysis before weight release, identifying two incremental areas: the sparse attention path and the multimodal branch. Based on this analysis, they prioritized operator development. Specific work included compiling an operator inventory, comparing old and new differences, evaluating the reuse rate of the existing operator library, and determining the set of new operators requiring development. Given that the Ascend platform already had KDA and MoE support, the incremental work centered primarily on sparse attention filtering and computation paths.

3.2 Phased Optimization of Core Operators

This adaptation covered four critical stages:

• Historical token filtering: A memory-intensive task. Optimization focused on compact index layout and parallel threshold computation to reduce computational bubbles caused by branch divergence.
• Sparse attention computation: Requires efficient batch processing of variable-length sequences. The underlying implementation must provide good load balancing.
• Preprocessing and residual computation: Operator fusion merges multiple memory accesses to reduce bandwidth pressure.
• State matrix update: Linear attention recurrence requires maintaining state across layers, with memory footprint decoupled from sequence length. This must be coordinated with cache management.

3.3 Framework-Level Inference Feature Support

The framework layer needs differentiated KV cache strategies for the hybrid architecture: linear attention layers maintain fixed-size state matrices, while sparse attention layers maintain on-demand retrievable key-value caches. Memory allocation and reclamation for these two cache types must be designed separately.

4. Cache and Communication Optimization for Long-Context Inference

4.1 Cache Management at Million-Token Scale

Million-token contexts put pressure on caches in both capacity and access efficiency. Engineering strategies include:

• Tiered caching: Recent tokens stay in high-speed memory, while older tokens are offloaded to large-capacity storage, with dynamic migration based on access frequency.
• Sparse retention: Combined with sparse filtering mechanisms, full caches are kept only for historical segments likely to be retrieved; the rest are stored as compressed summaries.
• Distributed cache pool: In multi-card and multi-node scenarios, cache shards are distributed across device memory and pulled on demand over high-speed interconnect.

It is important to note that cache eviction operations may cause data loss in the offloaded tier. In production environments, the offloaded tier should be periodically persisted to external storage, with data consistency checks and recovery mechanisms in place.

The following runnable pseudocode illustrates tiered cache management:

class TieredKVCache:
    def __init__(self, hot_capacity, cold_store):
        self.hot = {}          # Recent cache in high-speed memory
        self.cold = cold_store # Historical cache in the offloaded tier
        self.access_freq = {}
    def get(self, token_range):
        # Use a tuple rather than a list so it can serve as a dict key
        key = tuple(token_range) if isinstance(token_range, list) else token_range
        if key in self.hot:
            self.access_freq[key] = self.access_freq.get(key, 0) + 1
            return self.hot[key]
        return self.promote_if_needed(key)
    def promote_if_needed(self, key):
        if self.access_freq.get(key, 0) > THRESHOLD:
            data = self.cold.fetch(key)
            self.evict_lru_hot()  # Note: this may drop hot cache entries
            self.hot[key] = data
            return data
        return self.cold.peek(key)
    def evict_lru_hot(self):
        # Note: production environments should add persistence and recovery logic
        pass

4.2 Cross-Card Communication Co-Design

Ultra-large MoE inference involves expert parallelism and tensor parallelism. Optimization directions include overlapping communication with computation, compressed transmission of sparse index data, and leveraging high-bandwidth interconnect on super-node hardware to reduce aggregation latency.

5. Day-One Deployment Validation Workflow

The complete validation workflow includes the following stages, each with recommended expected metric ranges:

1. Environment provisioning: Align framework versions, operator libraries, drivers, and firmware in advance so work can start immediately once weights are available.
2. Precision alignment: Compare against the HuggingFace reference implementation or official weight extraction results using small-scale inputs, validating numerical consistency layer by layer, with an expected per-layer error below 1e-3.
3. Long-context functional validation: Construct million-token inputs to verify cache management and sparse retrieval stability, with end-to-end latency within a reasonable range depending on hardware configuration.
4. Multimodal inference validation: Verify end-to-end output quality for mixed image-text inputs.
5. Serving stress testing: Measure throughput and latency under concurrent load to evaluate scheduling and cache strategy effectiveness.

These stages are not strictly sequential. They are advanced in parallel based on early model analysis, compressing the overall timeline.

6. Practical Insights and Outlook

Several lessons can be drawn from this adaptation effort:

First, architecture-aware optimization is irreplaceable. Only by deeply understanding the recurrent characteristics of KDA and the sparse filtering logic can correct trade-offs be made at the operator and framework levels.

Second, cache management is becoming a core competitive capability for long-context inference. In million-token scenarios, cache efficiency can contribute 30%–50% of end-to-end latency improvements depending on specific hardware. The combination of tiered caching, sparse retention, and distributed cache pools is a pragmatic path forward.

Third, day-one adaptation reflects systematic engineering maturity. It depends on the integrated coordination of model analysis, operator reuse, framework evolution, and validation workflows. Platform selection should consider how quickly a platform can respond to new architectures and its level of engineering maturity.

Fourth, the combination of multimodal and long-context capabilities introduces orchestration challenges. When visual encoding coexists with long-sequence decoding, the framework layer must provide more fine-grained memory allocation and batch scheduling capabilities.

As model scale and context length continue to evolve, the center of gravity in inference optimization will shift from operator-level acceleration to system-level coordination across computation, caching, communication, and scheduling. For developers, reserving architectural headroom during technology selection delivers more long-term value than chasing single-point performance numbers.