SLAI T-Rex: Full-Parameter Post-Training of the DeepSeek-V4 Family on Ascend SuperPOD

Dongfang Li, Xiaodong Luo, Ruoyu Sun, Xuhui Chen, Linyuan Qiu, Jian Meng, Zhengxuan Lu, Yiting Wang, Yucheng Xie, Tao Guo, Tianxiang Fang, Jing Li, Sihang Chen, Shihao Hong, Chang Liu, Weihua Dai, Zirong Zeng, Ziwei Zhu, Zhuohan Wang, Zhengjun Yue, Igor Vasilyev, Min Liu, Weijian Sun, Xin Chen, Yingmeng Gao, Jinhua Zhou, Taolue Chen, Chenwei Wu, Dong Zhang, Wenlong Jin, Jinmin Xiang, Barkova Maria, Ushakov Anton, Xianfei Jin, Tian Ding, Zhihang Lin, Qian Chen, Linxin Yang, Mingzhe Yang, Bingwei Zhang, Hongzhang Yang, Fangxue Zhang, Shijun Qin, Jie Yu, Cuihua Hu, Tolstykh Vasiliy, Nosov Ivan, Abdullin Amir, Zhichen Zhou, Xin Zhang, Zhixiong Ning, Xutong Zhao, Junjie Huang, Jiajun Liu, Weiyan Kong, Zheng Zhang, Wenhan Luo, Lin Hu, Yangbo Guo, Li Zeng

SLAI T-Rex optimizes trillion-parameter MoE training on Ascend NPUs and enables solver-grounded OR specialization.

How can we optimize the full-parameter post-training of trillion-parameter MoE models on Ascend hardware to improve both training efficiency and downstream Operations Research (OR) capability?

Full-parameter training of trillion-parameter Mixture-of-Experts (MoE) models on SIMD-based hardware like Ascend NPUs is bottlenecked by irregular sparse-attention kernels, high communication overhead, and memory-intensive optimizer states. The authors introduce SLAI T-Rex, a hierarchical framework that combines expert-tensor parallelism, a double-buffered swap optimizer, and AuraKernel—an Operations Research-guided agent that automatically optimizes bottleneck kernels by formulating tiling as a constrained performance-modeling problem. This system achieves 34.22% Model FLOPs Utilization (MFU) on DeepSeek-V4-Pro, a 2.93× improvement over baseline, while enabling a specialized OR-domain workflow that outperforms GPT-5.4-Mini on zero-shot Pass@1 benchmarks.

Paper Primer

The system-level challenge stems from the mismatch between trillion-parameter MoE architectures and SIMD-based hardware. Sparse attention kernels and All-to-All communication patterns create irregular execution profiles that starve the matrix engines, while the massive optimizer state footprint forces inefficient memory management.

AuraKernel automates the optimization of these bottlenecks: it treats kernel tiling as a constrained search problem, using a stateful tree search to explore code-change candidates that minimize stage-level latency. By fusing fragmented eager-mode operators into unified AscendC kernels, the system keeps intermediate tensors in the Unified Buffer, drastically reducing global memory round trips.

SLAI T-Rex achieves a 2.93× improvement in Model FLOPs Utilization (MFU) for trillion-parameter MoE training.

Profiling of DeepSeek-V4-Pro on Ascend CloudMatrix384 SuperPOD compared to the open-source baseline recipe.

Continued Pre-Training (CPT) provides essential domain priors for Operations Research (OR) modeling.

Comparison of direct SFT versus CPT-initialized SFT on the B4O-ORGEval benchmark. CPT-initialized SFT reached 71.81% Pass@1, outperforming the base model by 11.27 percentage points.

Why is Operations Research (OR) a uniquely difficult domain for LLMs?

OR requires translating natural language into precise mathematical formulations and executable code that must satisfy strict solver constraints, making it a high-stakes test of both semantic reasoning and structural correctness.

How does AuraKernel differ from standard heuristic-based kernel tuning?

Unlike manual tuning or harness-based approaches, AuraKernel uses an OR-guided agent to systematically explore the tiling search space, formulating kernel optimization as a constrained performance-modeling problem to ensure long-horizon efficiency.

System Challenges in MoE Training

We expose why trillion‑parameter MoE training on Ascend stalls and how the paper tackles it.

Training trillion‑parameter Mixture‑of‑Experts models on Ascend NPUs hits three systemic roadblocks: fragmented memory‑bound kernels, stalled inter‑node communication, and low Model FLOPs Utilization (MFU). Existing large‑scale MoE pipelines are built for CUDA or TPU, leaving SIMD‑based hardware largely untapped. Moreover, Operations Research—critical for real‑world optimization—has received little attention in LLM post‑training, creating a gap this work aims to fill.

**Figure 1.** (a) MFU improvements during DeepSeek-V4-Pro (1.6T) training on Ascend SuperPOD. Starting from open-sourced recipes, optimized parallelism, NPU-CPU co-optimization, and kernel optimization increase MFU from 11.67% to 34.22%. (b) Overall accuracy, computed as the mean of NL4OPT (Ramamonjison et al., 2022), OptiBench (Yang et al., 2025c), Bench4Opt (B4O)-Feasible (Wang et al., 2025b), and Bench4Opt (B4O)-ORGEval (Wang et al., 2025b). SLAI T-Rex achieves the highest overall score while using a smaller disclosed activated-parameter budget.

To address these challenges we build the SLAI T‑Rex framework, which couples Ascend‑native system optimizations (including the AuraKernel OR‑guided tiling agent) with a solver‑grounded CPT‑SFT pipeline tailored for OR tasks. The resulting stack not only raises MFU to 34.22 % but also delivers a 2.93× speedup over the default recipe, while the CPT‑initialized SFT further lifts OR benchmark scores, most notably B4O‑ORGEval.

MFU more than triples, surpassing one‑third of peak compute, unlocking practical trillion‑parameter MoE training on Ascend.

Infrastructure and Bottlenecks

We expose the Ascend SuperPOD bottlenecks and resolve them with AuraKernel‑driven kernel fusion and system‑level scheduling tricks.

Training a 1.6 T‑parameter MoE model on the Ascend CloudMatrix384 SuperPOD surfaces three acute pain points: (1) fragmented, memory‑bound kernels that starve the matrix engine, (2) frequent expert‑routing all‑to‑all communication that cannot be fully overlapped, and (3) a long tail of tiny launch‑fragmented operators that inflate kernel‑launch overhead.

**Figure 2.** Latency breakdown of training trillion-parameter-class LLM model on Ascend SuperPOD NPU cluster. (a) Partition of latency in the end-to-end LLM training pipeline. (b) Averaged end-to-end latency breakdown, including computation, communication, and NPU idle bubble.

The profiling trace isolates two kernel families: (i) architecture‑intrinsic kernels (sparse‑attention backward, dense MatMulV3, grouped MoE GEMM) that dominate device time, and (ii) a launch‑fragmented tail of casts, views, and moves that together launch over 200 k kernels but contribute little compute.

AuraKernel treats a kernel like a car engine: it first measures which piston (stage) is throttling performance, then rewrites the engine map (tiling, buffering, loop order) to keep every cylinder firing efficiently.

How does AuraKernel differ from a conventional hand‑tuned kernel?

Traditional tuning picks a single tile size and hopes it works across all input shapes. AuraKernel continuously profiles, predicts the limiting stage, and automatically explores a constrained search tree until it finds a configuration that strictly improves the measured cycle count, then records that pattern for future reuse.

Choose an L0 tile of $4\times4$ (fits 2 KB in UB), leaving $4\times8$ and $8\times4$ as the outer dimensions.

Tile the $K$ dimension into two $4$‑element slices; each slice is loaded once into UB, reused for the $M\times N$ block.

Schedule a ping‑pong buffer: while slice 1 is being multiplied, slice 2 is pre‑loaded, eliminating load‑compute gaps.

Compute the $4\times4$ block on the CUBE core, then write the result back to GM.

Repeat for the second $K$ slice, reusing the same UB buffers.

By matching tile size to UB capacity, the kernel eliminates idle cycles caused by repeated GM‑UB traffic, turning a memory‑bound kernel into a compute‑bound one.

**Figure 3.** Overview of the proposed full-stack optimization framework, which includes three cooperating layers: parallel execution strategy, kernel optimization, and hardware-affinity tuning.

The double‑buffered swap optimizer overlaps host‑to‑device state transfers with AdamW computation: while chunk $i$ is being updated on‑chip, chunk $i+1$ is prefetched and chunk $i-1$ is written back, effectively hiding the swap latency behind useful work.

**Figure 6.** Bottleneck anatomy of one representative DeepSeek-V4-Pro training step on Ascend 910C (206,503 compute kernels, 39.84 s device-side duration). Panel (a): the heavy head of architecture-intrinsic operators by core-busy share (the tool’s “Task Duration” column)—sparse-attention backward, dense matrix multiplication, and grouped MoE multiplication. Panel (b): the complementary launch-fragmentation view, contrasting the launch share of non-matmul eager operators with their much smaller share of device duration, and isolating the dense sub-20 $\mu s$ tail. Panel (c): duration-weighted vector-pipeline ratios—vector compute well below the MTE2 load and MTE3 store ratios, whose combined memory movement is about 7.4$\times$ the vector-compute share. Panel (d): the resulting division of work—architecture-intrinsic heavy-head kernels routed to kernel-internal optimization by AuraKernel, and the launch-fragmented, memory-bound, and generic-Triton tails routed to re-expression as UB-resident AscendC kernels.

AuraKernel and Post-Training Pipeline

Optimizes fragmented kernels via AuraKernel fusion and builds an OR‑CPT pipeline for domain‑specific post‑training.

Fragmented, memory‑bound kernels dominate the training bottleneck on Ascend hardware, causing repeated global‑memory round trips that dwarf compute cost.

AuraKernel fuses a whole operator‑level dataflow into a single UB‑resident kernel, preserving the original mathematical contract while eliminating intermediate memory moves.

How does AuraKernel differ from a naïve “merge‑all‑ops” approach that simply concatenates kernels?

A naïve merge would still emit separate memory buffers for each sub‑operation, preserving the original data‑movement pattern. AuraKernel, by contrast, rewrites the entire dataflow so that only the input and final output buffers exist; every intermediate tensor lives in on‑chip registers or UB memory, which the hardware can stream without leaving the chip.

Tile over rows: each core receives two tokens (rows 0‑1 and 2‑3).

Within a tile, broadcast the $N=2$ weight vectors across $D=8$ features using a single Brcb instruction.

Perform the weighted expansion $x[t,d]\cdot h[t,n]$ for each $t$ and $n$, accumulating into a local buffer.

Apply the final BF16 cast and write the result back to global memory in one contiguous write.

The entire sequence runs without allocating any intermediate global buffers, so the only memory traffic is the initial read of $x$ and the final write of the fused result.

AuraKernel iterative optimization loop.

**Figure 7.** The architecture of AuraKernel. Starting from a functionally correct baseline kernel, AuraKernel generates an optimized operator in three stages. (1) **OR-based tiling optimization** diagnoses the operator's dominant regime and instantiates a hardware-aware parameter space, yielding an initial configuration and a set of promising optimization directions. (2) **Hardware-grounded iterative optimization** runs a closed modify-verify-profile loop: a harness compiles each candidate, gates it on correctness with the return d profiling, while K-Search uses this feedback to explore a branchable, rollbackable tree of optimization directions and commits a checkpoint on every verified improvement. (3) **Skill distillation** converts each verified gain into a reusable skill that is retrieved to guide hypothesis generation in later episodes.

The OR‑CPT pipeline injects domain‑specific optimization knowledge into the model by generating solver‑verified synthetic instances, converting them to natural‑language problems, and filtering them through a strict contract‑checking loop.

Generate a structured optimization instance with a parameterized generator.

Solve the instance using Gurobi to obtain feasibility, objective value, and variable assignments.

Render the instance into a natural‑language problem statement, applying industry‑specific templates.

Run the model on the problem statement to produce decision variables, constraints, and Gurobi‑style code.

Validate the generated code by executing it; compare the resulting objective to the reference solution.

Accept the sample only if the objective matches within a tolerance and all contracts (syntax, API usage) are satisfied.

Deduplicate accepted samples and add them to the CPT corpus.

Why does the OR‑CPT pipeline resemble a standard “synthetic data” approach, yet achieve higher fidelity?

Standard synthetic data often skips solver verification, so generated problems may be ill‑posed or inconsistent. OR‑CPT closes the loop by solving each instance first, then enforcing strict contract checks on the model’s reconstruction, guaranteeing that every retained sample is both semantically correct and executable.

Generator creates the LP: minimize $c\,x$ subject to $x \ge 0$ with $c=5$.

Gurobi solves it, yielding optimal $x^\* = 0$ and objective $0$.

Natural‑language rendering produces “Find the smallest non‑negative value of $x$ that satisfies the constraint $x \ge 0$.”

Model reconstructs the formulation, emits Gurobi Python code, and the code is executed to confirm $x^\* = 0$.

This tiny loop demonstrates the full verification cycle; scaling to $N\approx 10$ variables follows the same pattern.

**Figure 13.** Pipeline of OR-CPT data construction. Parameterized optimization generators produce structured instances, which are verified with Gurobi, rendered into business-oriented problem statements, reconstructed into executable formulations, and filtered by contract checks, solver execution, and objective-value agreement before export.

Mixed‑precision dataflow in the mHC pre‑projection path previously promoted BF16 activations to FP32 early, inflating global‑memory traffic and turning a compute‑light stage into a memory‑bound bottleneck.

The fused backward kernel computes the gradient of the limited SwiGLU activation in closed form, eliminating the eager chain of casts, masks, and concatenations that previously caused a 401 µs memory‑bound trace.

Why can the limited SwiGLU backward be fused when a standard SiLU backward cannot?

Standard SiGLU backward lacks the explicit clamp limits that bound the activation; without those limits the gradient expression cannot be expressed as a single closed‑form kernel. The limits introduce piecewise conditions that are amenable to a single fused implementation.

Instead of exposing slice assignments that generate backward scatter chains, the unified operator returns the full rotated row as a single tensor, enabling both forward and inverse rotations with one kernel that uses real‑arithmetic pairwise rotations.

How does the unified rotary embedding differ from a naïve “apply complex rotation” implementation?

A naïve implementation would construct a complex datatype, perform a complex multiplication, and then convert back to real tensors, incurring extra memory moves and type conversions. The unified operator stays in real arithmetic, processes pairs directly, and therefore avoids the intermediate complex construction entirely.

AuraKernel fuses entire operator dataflows into a single UB‑resident kernel, and the OR‑CPT pipeline injects solver‑verified domain knowledge, together eliminating the dominant memory‑bound stalls that cripple trillion‑parameter MoE training on Ascend hardware.

Performance and Ablation Studies

Key speedups and downstream gains demonstrate the full‑stack impact of AuraKernel and the post‑training pipeline.

The earlier sections showed that fragmented kernels and eager‑chain overhead cripple trillion‑parameter MoE training on Ascend hardware; this section quantifies how AuraKernel and the post‑training pipeline resolve those bottlenecks.

CPT+SFT lifts the 0‑shot Pass@1 on B4O‑ORGEval from 34.26 % to 59.39 %, a +25.13 pp gain over the Base model.

Table 13 reports the end‑to‑end OR benchmark scores for Base, SFT, and CPT+SFT; the B4O‑ORGEval column shows 34.26 % → 59.39 %.

**Figure 15:** The AI-Assisted Quality Gates workflow. This multi-stage validation process uses the Cleaner, Reviewer, and code execution Validation to ensure the quality and auditability of the final dataset.

**Figure 16:** AuraKernel optimization of the shared-key-value sparse-attention kernels on Ascend 910C. (a) Per-call forward task duration (1.23x at compression ratio 128, 1.09x at compression ratio 4). (b) Per-call backward task duration of SparseAttnSharedkvGrad (1.24x and 1.21x). (c) AIV pipeline-time decomposition of the backward kernel at compression ratio 4: the reductions concentrate in the memory-movement pipes (MTE2 load and MTE3 store) rather than in vector compute, confirming that the gain comes from cutting global-memory traffic.

**Figure 17:** AuraKernel optimization of the weight-free RMS-normalization kernels (tp=2). (a) Per-call task duration speedup: forward 11.9x and backward 3.4x. (b,c) Duration-weighted AIV pipeline ratios before and after optimization: the vector-compute ratio rises (forward 27.8% → 67.5%, backward 17.8% → 32.0%) while the scalar ratio falls (forward 31.5% → 17.2%, backward 34.8% → 18.0%), showing that 2D head blocking, multi-row processing, and explicit FP32 accumulation turn a scalar/address-bound reduction loop into a vector-compute-dominated kernel.

Beyond individual kernels, collapsing eager operator chains into fused AscendC kernels removes launch overhead, intermediate tensors, and dtype round‑trips.

**Figure 18:** AuraKernel optimization of the sparse lightning-indexer gradient on Ascend 910C (a) Per-call backward task duration speedup, 1.16x on the full CFA path (cmp. ratio 0) and 1.21x on the SCFA path (cmp. ratio 4). (b,c) AIV pipeline-time decomposition before and after optimization at the two ratios: the kernel is scalar- and load-bound, and the reductions concentrate in the scalar and MTE2-load pipes rather than in vector compute, consistent with cutting address-computation and global-memory-load traffic.

For kernels originally written in Triton, a direct rewrite to AscendC further shifts the compute profile from movement‑bound to vector‑compute‑bound.

Training stability on Ascend 910C was verified with a 50K self‑distilled SFT run (GBS = 128, MBS = 1, `SEQ_LEN` = 8192) that completed 800 steps without NaNs.

**Figure 19:** Module task duration before (left bar) and after (right bar) operator fusion and rewrite for the seven targets, decomposed by operator category. Category durations are re-aggregated from the per-operator trace times of the pre- and post-optimization profiles. Each bar pair represents kernel-module durations (Table 8), with the module speedup annotated above. Across targets, the baseline implementations are dominated by memory movement, dtype casts, and small elementwise/mask fragments, whereas the proposed fusion retain achieves significant efficiency improvements.

SFT scaling experiments reveal a trade‑off: larger distilled datasets improve code‑protocol tasks but can degrade natural‑language optimization scores.

Applying contract‑aware cleaning and chain‑of‑thought (CoT) enhancement restores natural‑language performance while further boosting code tasks.

Initializing the model with Continued Pre‑Training (CPT) before SFT yields consistent gains across all OR benchmarks.

Summary and Future Directions

We wrap up the SLAI T‑Rex framework, its system gains, OR improvements, and artifact releases.

The ablation study confirms that interleaving domain‑specific and general data during CPT is crucial: the Balanced mixture yields strong OR and general benchmark performance, while the Pure OR variant harms overall capability.

**Figure 21.** Ablation on data mixture under the end-to-end (CPT→SFT) pipeline.

At the system level, our full‑stack Ascend SuperPOD optimizations deliver up to 34.22 % MFU, demonstrating that trillion‑parameter MoE training can be efficiently realized on non‑GPU hardware.

In the post‑training phase, CPT initialization adds substantial OR priors: B4O‑Feasible rises from 65.93 % (SFT only) to 71.22 %, and B4O‑ORGEval improves from 48.73 % to 59.39 %.

Scaling self‑distilled data from 10 K to 50 K does not guarantee gains; targeted contract‑aware cleaning and concise chain‑of‑thought enhancements on the 10 K set consistently boost NL4OPT, OptiBench, and ORGEval performance.

**Table 15.** Performance of various models on operations research benchmarks. The best result in each column is in bold, and the second best is underlined.

Future work will extend the pipeline to DeepSeek‑V4‑Pro and open‑source the optimized AscendC kernels, prompt templates, and monitoring tools to enable reproducibility and further research.

We also release the full set of reproducibility artifacts—including prompt templates, format‑contract specifications, curriculum scheduling configurations, and monitoring callbacks—so that the community can inspect, adapt, and build upon our data‑cleaning and evaluation pipeline.

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers