A.X-K2

srcset="https://huggingface.co/skt/A.X-K2/resolve/main/assets/A.X_K2_BI_Negative.png">

Model Summary

A.X K2 is a large-scale Mixture-of-Experts (MoE) language model trained from scratch as a high-performance, agentic foundation model, and the successor to A.X K1. The model contains 688 billion total parameters, with 33 billion active parameters, delivering strong reasoning and instruction-following performance while maintaining practical inference efficiency.

Through a Think-Fusion training recipe, a single unified model supports both a thinking mode for complex problem solving and a non-thinking mode for concise, low-latency responses, allowing the user to trade quality for cost on a per-request basis.

A.X K2 is developed as part of the Korean government's Sovereign AI foundation model project, aiming to build a frontier-scale model with deep understanding of the Korean language and culture.

Key Features

• Large-Scale Sparse MoE (688B / 33B Active) Employs a high-capacity Mixture-of-Experts architecture (256 routed experts + 1 shared expert) that activates only 8 experts per token, enabling strong reasoning performance with practical inference efficiency. Relative to A.X K1, the number of routed experts increased from 192 to 256 while keeping the activated parameter count at 33B.

• Hybrid Reasoning Control (Think / Non-Think) Supports user-controllable reasoning depth via the Think-Fusion recipe, allowing explicit multi-step reasoning or concise low-latency responses within a single unified model.

• Sparse Gated Attention (SGA) for Efficient Long-Context Inference Combines a lightweight sparse-attention indexer (top-k token selection) with head-specific gated attention on top of Multi-head Latent Attention (MLA). This prunes attention computation at long context while suppressing attention sinks, delivering higher throughput and lower latency at long sequence lengths.

• Gated Norm (GN) for Stable Training and Low-Precision Serving Applies an input-dependent gate immediately after RMSNorm to suppress massive activations and hidden-state outliers. This stabilizes large-scale training and produces an activation distribution with suppressed outliers, which is favorable for low-precision (FP8 / NVFP4) deployment.

• Native FP8 Training and FP8 Checkpoint A.X K2 was trained from scratch natively in FP8 (MXFP8, E4M3) for both the forward and backward passes, and the released checkpoint is distributed in block-scaled FP8 (E4M3). This halves the memory footprint relative to BF16 and enables efficient FP8 serving out of the box, without a separate post-hoc quantization step.

• Tokenizer Optimized for Multilingual and Code Data Inherits the large-vocabulary BBPE-based tokenizer from A.X K1, optimized for token efficiency across five languages (English, Korean, Chinese, Japanese, and Spanish), with a strong emphasis on source code, structured text, and programming-related patterns.

Model Details

• Architecture: Decoder-only Transformer with Mixture-of-Experts (MoE)
• Total parameters: 688B (256 routed experts + 1 shared expert)
• Active parameters: 33B per token (8 experts + 1 shared expert)
• Number of layers: 61 (1 dense + 60 MoE)
• Number of attention heads: 64
• Hidden size: 7168
• Intermediate size: 18432 (dense layer)
• Expert intermediate size: 2048
• Attention: Multi-head Latent Attention (MLA) with head-specific output gate and sparse attention (SGA)
• Normalization: Gated Norm (gating applied after RMSNorm) with QK-normalization
• Vocab size: 163,840
• Context length: 262,144 tokens (256K) — 128K trained natively, extended to 256K via YaRN scaling
• Training precision: Native FP8 (MXFP8, E4M3) forward and backward passes; FP32 master weights and gradients
• Checkpoint precision: Block-scaled FP8 (E4M3, 128×128 weight blocks with dynamic activation scaling)

Architecture Highlights

A.X K2 is built from Gated Transformer Blocks, which integrate two gating mechanisms into the MoE transformer block: Sparse Gated Attention (SGA) in the attention sublayer and Gated Norm (GN) on the normalization path. Together they keep large-scale training stable while making long-context inference and low-precision serving efficient.

Sparse Gated Attention (SGA)

A.X K2 combines two mechanisms into SGA to make long-context inference efficient while keeping training stable:

• Gated attention: A head-specific output gate is applied throughout pre-training on top of MLA. It introduces non-linearity into the attention output, mitigates attention sinks, and improves loss convergence.
• Sparse attention: A lightweight indexer scores and ranks key–value candidates and selects a small set of top-k tokens (k = 2048) per query, substantially reducing attention compute at long sequence lengths.

The two mechanisms are mutually reinforcing: by suppressing the attention-sink mass on uninformative tokens, the output gate yields a better-calibrated signal so the indexer spends its limited budget on genuinely relevant positions. In long-context serving, as the input length grows into the tens of thousands of tokens, A.X K2 with SGA sustains substantially higher token throughput, lower time per output token (TPOT), and lower time-to-first-token (TTFT) than A.X K1.

Gated Norm (GN)

A.X K2 adopts Gated Norm, in which a learned, input-dependent gate modulates the normalized activation before it enters the residual stream. By preventing the normalized output from fitting outlier scalars, GN suppresses massive activations—hidden units orders of magnitude larger than the rest that persist across layers and are a documented source of both training instability and low-precision quantization error.

With this design, training remained sufficiently stable without the dual-normalization scheme used in A.X K1: the single post-normaliza