A.X K2 Technical Report

Evaluation

• A.X K2 performs competitively with strong open-weight baselines, and Figure 1 reports leads in mathematics, Korean-language understanding, and agentic tool use.The figure names AIME26 and Apex, KMMLU-Pro and CLIcK, and τ2-Bench Telecom; the report also says A.X K2 matches or exceeds baselines on math and Korean tasks.Source: Figure 1 · p. 1, §1 · p. 2

Token-Efficient Pre-training

• With approximately 8.5T tokens—fewer than A.X K1’s ∼10T—A.X K2 improves across benchmarks, by over 30 percentage points on some, indicating greater token efficiency.The model has 688B total and 33B active parameters; its training mixture is smaller but higher-quality and oriented toward agentic workloads.Source: §1 · pp. 1–2

Gated Transformer Blocks for Stable Training and Efficient Long-Context Inference

• Sparse Gated Attention (SGA) limits each query to 2,048 positions, while Gated Norm (GN) suppresses hidden-state outliers to support stable, low-precision computation.Sparse warmup optimizes the indexer against its own sparse top-k selection; LongBench changes from 62.80 to 62.99 across adaptation, while GN improves low-bit accuracy including FP8 and FP4.Source: §1 · p. 2, Figure 2 · p. 3

Think-Fusion SFT and Multi-Stage Reinforcement Learning

• Think-Fusion trains one unified model to switch between thinking and non-thinking modes using explicit control tokens, while multi-stage reinforcement learning jointly optimizes instruction following, preference alignment, tool use, and safety.The modes let users allocate inference-time compute according to task complexity, from concise low-latency responses to extensive reasoning chains.Source: §1 · p. 2

Frontier-Scale Engineering Capability

• A.X K2 demonstrates end-to-end engineering feasibility for a 688B-parameter model through stable large-scale training, reproducible pipelines, and system-level optimizations.The report presents these capabilities as a foundation for scaling toward frontier-class model development.Source: §1 · p. 2

2.1 Model Configuration

• A.X K2 has 688B total and 33B active parameters, with 256 routed experts versus A.X K1’s 192 and 60 MoE layers among 61 transformer layers.The design uses MLA for KV-cache efficiency, head-specific output gating, QK normalization, and 64 attention heads with shared dense experts. It trains natively at 128K via ABF and reaches 256K at inference with increased YaRN scaling.Source: §2.1 · pp. 3–4, Table 1 · p. 3

2.2 Gated Transformer Blocks for Mitigating Attention Sinks and Outliers

• Gated Norm suppresses hidden-state outliers and stabilizes convergence, replacing A.X K1’s dual normalization with a single post-RMSNorm gate at an approximately 5% training-throughput cost.GN also improves activation distributions for NVFP4 serving; the ablation indicates that adding a post-MLP normalization layer yields no further convergence benefit.Source: §2.1 · p. 4, §2.2 · pp. 4–5, Figure 3 · p. 4
• 62.80 before and 62.99 after SGA adaptation on LongBench show that the sparse-attention adaptation preserved the reported score.Each query selects 2,048 positions; that fixed budget is 1.6% of positions at 128K context and 0.8% at 256K.Source: §2.2 · p. 4
• SGA combines pretraining gated attention with a sparse indexer that selects 2,048 top-k tokens per query to reduce long-context attention compute.The indexer is adapted after native 128K training; the output gate suppresses attention-sink mass and helps the indexer prioritize relevant positions.Source: §2.2 · p. 4

3.1 Pre-Training Dataset

• A.X K2 trains on approximately 8.2T tokens, refreshing web corpora and expanding agentic, software-engineering, multilingual, and long-context data.Later stages also blend in instruction-style data to strengthen instruction-following and task competence ahead of post-training.Source: §3.1 · p. 5

3.2 Quality Filtering and Difficulty-Aware Curation

• Curation selects a training-efficient subset from a roughly 16.2T-token pool, progressively tightening per-domain quality filters and introducing difficulty scoring in Stage 2.A distilled classifier scores educational value and difficulty from 0–4; thresholds are relaxed for long-tail fields to preserve rare content.Source: §3.2 · pp. 5–6

3.3 Pre-Training Process

• Pre-training follows a three-stage WSD schedule: 6.4T tokens for general knowledge, 1.4T for high-quality reasoning, then long-context extension.Stages 1 and 2 use 4,096-token sequences; Stage 2 begins learning-rate decay, and Stage 3 extends context to 32K and then 128K before sparse-attention adaptation.Source: §3.3 · p. 6

3.4 Long-Context Adaptation

• A.X K2 reaches 128K context natively and achieves 100 NIAH retrieval at every position at 256K; sparse warmup trains its indexer against sparse top-k selection.The 128K extension uses an unchanged RoPE base of 106; sparse warmup freezes non-indexer parameters, with each query selecting 2,048 tokens, and is faster than dense warmup with negligible downstream-performance difference.Source: §3.4 · pp. 7–8, Figure 5 · p. 8

3.5 Checkpoint Merging

• The merged base model is consistently stronger than any individual checkpoint in its roughly 8B-token averaging window, at negligible cost.The merge averages the last six checkpoints with greater weight on more recent ones, including MoE router weights.Source: §3.5 · p. 8

3.6 Hyperparameters

• A.X K2 uses global load-balancing loss plus a learnable expert-bias term, with balancing strength and bias updates adjusted across training stages.During SGA indexer warmup, the expert bias is briefly frozen; it resumes updating afterward.Source: §3.6 · p. 9

3.7 Parallelism and Training Efficiency

• Training used 512 NVIDIA B200 GPUs for approximately 70 days, with pipeline and expert parallelism of 8 and no tensor parallelism.The released checkpoint uses blockwise FP8 weights and dynamic per-token activation scaling, so it serves in FP8 without separate post-hoc quantization.Source: §3.7 · pp. 9–10

4.1 Supervised Fine-Tuning

• Paired SFT examples and explicit control tokens enable users to select thinking or non-thinking responses, despite an approximately 13:1 thinking-to-non-thinking token ratio.Each prompt has both response modes, and mode confusion is tracked through post-training to verify that control remains effective after RL.Source: §4.1 · p. 10

4.1.1 Data

• SFT pairs reasoning and instruction examples, and extends safety data with evolved specifications and a larger English–Korean prompt pool to target unsafe compliance and unnecessary over-refusal.Safety samples undergo rubric-based filtering and meta-review for unsafe residue, language mismatch, formatting leakage, and excessive refusal.Source: §4.1.1 · pp. 10–11

4.1.2 Training

• The two SFT stages train on packed 128K-token sequences: an initial 0.4M-example stage updates the model and indexer, then main SFT fixes the indexer.Length-aware First-Fit Decreasing packing sorts samples from longest to shortest to reduce unused bin space and padding waste.Source: §4 · p. 10, §4.1.2 · p. 11

4.2 Reinforcement Learning Infrastructure

• End-to-end blockwise FP8 avoids the reward collapse seen with trainer–rollout precision mismatch, while asynchronous TP1/DP8 rollout with EP8 overlaps generation and optimization.Matching trainer and rollout quantization supports stable learning; the rollout layout replicates attention across eight data-parallel replicas while keeping expert sharding aligned.Source: §4.2 · pp. 12–13, Figure 6 · p. 12

4.3 Multi-Stage Reinforcement Learning

• Every RL stage jointly optimizes instruction following, human-preference alignment, and agentic tool use; later stages add safety, with data reweighted or replaced to address weak behaviors.Thinking and non-thinking modes are sampled in roughly equal proportion throughout; the shared stages avoid isolated capability optimization that can regress other behaviors.Source: §4.3 · p. 13
• RL combines verifiable instruction and tool-call rewards with judged preferences and safety, using CISPO and separately normalized GDPO rewards without a KL penalty.Long-horizon tool tasks pass through a validity gate before reference-based judging; safety rewards include harmfulness and principle-based preference assessments.Source: §4.3 · p. 13

5.1 Evaluation Settings

• Evaluations compare A.X K2 with A.X K1 and open-weight baselines across Math, Korean, Code, Science & Knowledge, General, and Agentic tasks.Unless otherwise noted, models are evaluated in thinking mode; A.X K2 also supports non-thinking mode through Think-Fusion.Source: §5.1 · p. 14

5.2 Evaluation Results

• A.X K2 leads on AIME26, Apex, Apex-shortlist, KMMLU-Pro, and CLIcK, with major gains over A.X K1 across all six evaluation categories.It scored 35/42 on IMO 2025 and correctly proved all 8 KMO26 second-round problems. In Agentic evaluation, it led on τ2-Bench Telecom and scored 1031 Elo on GDPval, but had remaining gaps on Terminal Bench and BrowseComp.Source: §5.2 · p. 14, Table 6 · p. 15

5.3 Long-Context Evaluation

• A.X K2 averages 94.6 on RULER through 256K while each query reads only 2,048 positions.It retains 95.1% of its 4K score at 128K and 88.8% at 256K. NIAH retrieval was perfect at 128K, 256K, and 512K without length-specific fine-tuning, but 512K is not a supported serving configuration.Source: §5.3 · p. 15, Table 7 · p. 15
• LongBench v2 scores are 63.9 for A.X K2, 64.2 for Qwen3.5, and 63.2 for GLM-5.1 on the filtered set.On all 503 items, A.X K2 scores 62.2 and ranks third; the reported filtering changes rankings. Dense-to-sparse adaptation leaves LongBench v1 unchanged within noise (62.80 versus 62.99), while long-context throughput rises as input length grows.Source: §5.3 · p. 16, Table 8 · p. 16, Table 9 · p. 16

5.4.1 KS-Eval

• On KS-Eval, A.X K2 improves over A.X K1 in all four categories and scores higher on all 26 tasks.The suite covers 26 tasks and 1,800 items designed around Korean linguistic and cultural contexts; its results inform an iterative evaluation, analysis, training, and re-evaluation cycle.Source: §5.4.1 · p. 17
• A.X K2 ranks first in KS-Eval Common Sense (57.45%) and Instruction Following (71.67%), while GLM-5.1 leads Long Context and Qwen3.5 leads STEM.A.X K2 scores 66.85% in Long Context and 75.47% in STEM. The category gains over A.X K1 are 19.47 pp in Common Sense, 27.98 pp in Instruction Following, 27.11 pp in Long Context, and 13.46 pp in STEM.Source: §5.4.1 · p. 17

5.4.2 Manufacturing Benchmark

• A.X K2 has the highest Manufacturing Benchmark overall F1 at 69.0%, driven by 61.1% abstention F1 despite fourth-place answerable F1.It also leads under misleading cues: 67.8% when a user asserts a wrong option and 41.4% when a document supports one. The document-assertion score is substantially lower, identifying resistance to misleading supplied references as a priority.Source: §5.4.2 · pp. 18–19, Table 11 · p. 18

5.4.3 Red-Teaming Benchmark

• A.X K2’s red-team ASR is 12.1% in non-thinking mode, the lowest among compared models, but rises to 33.9% in thinking mode.Qwen3.5 has the lowest thinking-mode ASR at 21.0%; A.X K2 remains ahead of A.X K1 and GLM-5.1. The mode effect varies across models, so extended reasoning is not inherently safer or less safe.Source: §5.4.3 · p. 19, Table 12 · p. 19

6.1 Low-Precision Robustness

• NVFP4 averages 78.19 vs. 78.95 for FP8 across eleven tasks, retaining 99.0% of FP8 accuracy with a 0.76-point drop.Nine of eleven tasks change by less than 2 points, three improve, and the largest regressions are GSM8K (−2.50) and MATH (−2.42); NIAH remains perfect at 128K, 256K, and 512K. On dedicated NPU hardware, A.X K2 reaches 107% performance-per-watt relative to an NVIDIA L40S. GN and the SGA output gate suppress outliers that could reduce block-scaled NVFP4 precision.Source: Table 14 · p. 20, §6.1 · p. 20
• NVFP4 reduces minimum serving memory to roughly 57% of FP8, freeing 276 GB and bringing A.X K2 to about 36% of A.X K1’s BF16 requirement.The freed memory can be redirected to KV cache, raising sustainable concurrency at 128K–256K inputs.Source: §6.1 · p. 20, Table 13 · p. 20

6.2 Long-Context Serving Throughput

• At 64K, matched-FP8 A.X K2 first exceeds A.X K1 (+2.7%, 10.6K total tok/s), reaching +35.4% (12.3K vs. 9.1K) at 120K.A.X K2 throughput keeps rising through 120K while A.X K1 declines after peaking at 64K; at 32K, A.X K2 is still 18.2% behind.Source: §6.2 · p. 20, Figure 7 · p. 21
• EAGLE3 advances A.X K2’s crossover against A.X K1 from 64K to 32K and reaches +65.9% (15.1K total tok/s) at 120K.It raises throughput by 23–30% across input lengths; its drafter commits an average of 2.24 tokens per step and requires no additional GPUs. Gains depend on acceptance rates and may differ on real traffic.Source: §6.2 · p. 21, Figure 8 · p. 22
• At 120K, NVFP4 A.X K2 reaches 20.3K total tok/s (+123.7% vs. A.X K1) and exceeds A.X K1 from 1K through 120K inputs.Its gains over A.X K1 reach +37.8% at 32K and +69.4% at 64K. Halved weight memory enables B200 FP4 acceleration and adds KV-cache headroom.Source: §6.2 · pp. 21–22, Figure 9 · p. 22

Infrastructure Constraints on Scaling

• Expert parallelism was limited to the eight-GPU intra-node NVLink domain (EP = 8), restricting exploration of multi-node parallelism that requires tighter interconnect integration.This kept expert-parallel communication off slower inter-node links but limited finer-grained expert partitioning.Source: §7 · p. 22

Optimization under Resource Constraints

• Fixed time and GPU budgets led to modest performance shortfalls versus similarly sized models trained with larger compute budgets, despite empirically guided configuration choices.The authors expect more compute to relax these constraints and improve performance.Source: §7 · p. 23

Scope of Modalities

• A.X K2 is text-only and lacks native multimodal understanding, which the authors plan to add in future work.The passage notes strong reasoning performance while identifying this modality limitation.Source: §7 · p. 23

8 Conclusion

• A.X K2, a 688B-parameter MoE model, is competitive with strong open-weight baselines across English and Korean mathematics, coding, and general-knowledge benchmarks.The model was planned using compute-budget-based trade-offs between model scale and training tokens.Source: §8 · p. 23
• Think-Fusion combines paired thinking and non-thinking responses with multi-stage on-policy reinforcement learning to enable user-controlled mode switching within one model.Source: §8 · p. 23
• Training on 512 NVIDIA B200 GPUs for approximately 70 days involved iterative optimization to address throughput gaps from communication overhead, load imbalance, and execution inefficiencies.The approach supported stable convergence and predictable training behavior under fixed resource and timeline constraints.Source: §8 · p. 23
• Future work will prioritize native multimodal capabilities and scaling toward the trillion-parameter regime while advancing training and inference efficiency.Source: §8 · p. 23

B.1 Pre-train Data Category Mixture

• Percentages describe each stage’s designed mixture, not necessarily the tokens actually consumed.Stage 1 consumed the first 6.4T tokens of an approximately 8T schedule; the other reported stages used their full schedules.Source: §B.1 · p. 25
• Stages 3B and 3C share one data mixture and are reported together.Stage 2, Stage 3A, and Stages 3B&3C used schedules of 1.4T, 0.22T, and 0.14T tokens, respectively.Source: §B.1 · p. 25
• A dash means a category was not separately tracked in that stage, while “Etc.” aggregates remaining minor categories.The taxonomy was refined across stages; “Etc.” includes encyclopedia, news, and articles.Source: §B.1 · p. 25
• “Instruction” denotes instruction-style SFT-format data introduced during pre-training, while “STEM” covers science and engineering text distinct from academic-paper sources.Source: §B.1 · p. 25
• “Academic Papers” is predominantly academic papers, with some specialized long-form domain text such as legal documents and patents.Categories are grouped by data source or domain and by the target capability they primarily serve; percentages are rounded and may not sum exactly to 100%.Source: §B.1 · p. 25
• The pre-training mixture is reported by language and by content category across stages.The supplied table captions identify the dimensions reported but provide no composition values.Source: Table 16 · p. 25, Table 17 · p. 25

• Tables 18 and 19 present statistics for all SFT samples, whose data was supplemented periodically based on training performance assessments.Source: §B.2 · p. 25

B.3 Safety Reasoning Data Construction

• Building on ESSA, A.X K2 adds a specification-evolution and data-synthesis pass tailored to its training format.A.X K1 had used ESSA safety data in its post-training mixture.Source: §B.3 · pp. 25, 27
• The pipeline collects English and Korean prompts across harmful requests, safety-adjacent benign requests, and over-refusal-sensitive cases, then uses taxonomy labels to select among 50 domain–task specifications.Labels are used only for specification selection, not as supervision targets.Source: §B.3 · p. 27
• A reasoning-capable teacher uses the selected specification to answer in reasoning mode, explaining boundaries in ordinary language and offering safe alternatives without exposing internal labels or rules.Outputs are converted from ESSA’s explicit thought/response markers into the target chat template’s standard chat-sample format.Source: §B.3 · p. 27
• Validation filters unsafe residue, proxy harmful artifacts, privacy leakage, language mismatch, malformed reasoning boundaries, and unnecessary refusals, favoring useful safe artifacts when benign goals can be served.Raw request–response logs and source metadata are preserved before conversion; flagged samples may be excluded, regenerated, or sent for meta-review.Source: §B.3 · p. 27

Serving setup

• Open-weight baselines use publisher-only OpenRouter serving at xhigh reasoning effort, provider-default generation settings, and a benchmark-specific maximum output length.The output limit uses the model page’s stated maximum when available, or the context limit minus 100K tokens if that limit is hit; Qwen3.5-397B-A17B is fixed at 65,536 tokens, while all others receive at least 128K.Source: §B.4 · p. 27

Metrics

• Table 6 reports pass@1 averaged across multiple generations per sample, with eight generations for every mathematics benchmark.AA-Omniscience follows the AAII v4.1 percentage formula: (2 × accuracy + (1 − hallucination rate))/3.Source: §B.4 · p. 27

BrowseComp

• BrowseComp permits only Brave Search’s LLM Context API, with no additional page fetching or summarization and at most 10 searches per problem.Source: §B.4 · p. 27

B.5 Details of Evaluation Benchmarks

• Across the benchmark suite, responses are generated to the model’s context limit, and any response reaching that limit is graded incorrect.Evaluation prompt templates largely follow A.X K1.Source: §B.5 · p. 27
• Math benchmarks use final-answer correctness with chat-based chain-of-thought prompting and pass@1 averaged over eight generations; CLIcK uses a standard zero-shot protocol.The math suite is AIME26, Apex and Apex-shortlist, and KMO26 (1st round); Korean benchmarks are KMMLU-Pro, KoBALT, and CLIcK.Source: §B.5 · pp. 27–28
• The suite covers executable coding, scientific programming, graduate-level science and advanced knowledge, factuality, instruction following, and long-context coherence.It uses LiveCodeBench v6, SciCode, Terminal Bench v2.1, GPQA Diamond, Humanity’s Last Exam, AA-Omniscience, IFBench, and AA-LCR; LiveCodeBench is limited to its February–May 2025 subset.Source: §B.5 · p. 28
• Agentic evaluation includes professional-task performance, conversational agents in dual-control environments, and browsing agents locating hard-to-find web information.The reported measures are GDPval v2 arena Elo, τ2-Bench Telecom and τ3-Bench Banking domain scores, and BrowseComp; the BrowseComp search-tool configuration is described in Appendix B.4.Source: §B.5 · p. 28

B.6 Multilingual Evaluation

• A.X K2 scores 87.56 on Global-MMLU-Lite, ahead of DeepSeek-V4 Flash (86.81) and MiniMax-M2.7 (87.31), while its cross-language score range falls from A.X K1’s 21.50 to 4.00.The evaluation uses 1,600 four-way multiple-choice items across English, Spanish, Chinese, and Japanese, in thinking mode with zero-shot exact-match scoring.Source: §B.6 · pp. 28–29, Table 20 · p. 29

B.7 LongBench v2: Full-Set and Category-Level Results

• A.X K2’s LongBench v2 overall score rises from 62.2 to 63.9 (+1.7) after filtering, moving from third on the full set to second on the filtered set, 0.3 points behind Qwen3.5.A.X K2 is the only model whose overall score increases after filtering; the paper reports both sets because filtering affects models unevenly.Source: Table 21 · p. 29, §B.7 · p. 29
• Filtering raises A.X K2’s overall score by 1.7 points but lowers its Long score by 2.1; Code Repository Understanding and Long Structured Data Understanding fall by 14.0 and 10.6 points.Long-dialogue History Understanding is unchanged because none of its items are removed.Source: §B.7 · p. 29, Table 22 · p. 30
• A.X K2 leads Multi-Document QA on both sets (60.8, 63.9), ties filtered Single-Document QA (65.8), and leads filtered Long In-context Learning (75.0), but trails on two other categories.After filtering, Code Repository Understanding and Long Structured Data Understanding have only 30 and 22 items, respectively, so category gaps are broad indicators rather than precise measurements.Source: Table 23 · p. 30, §B.7 · p. 30