TL;DR

A.X K2는 688B 총파라미터, 토큰당 33B 활성파라미터의 대형 Mixture-of-Experts LLM으로 natively FP8 학습과 block-scaled FP8 체크포인트를 제공하여 FP8 서빙이 가능한 모델이다. 내부적으로 Sparse Gated Attention(SGA)와 Gated Norm(GN), Multi-head Latent Attention(MLA)을 결합해 긴 문맥에서 top-k 인덱싱과 head-specific gating으로 attention 연산을 희소화하고 attention sink를 억제함으로써 RULER에서 overall 94.6, NIAH에서 모든 컨텍스트·니들 깊이에서 100을 달성하는 성능을 보였다. Think-Fusion 레시피로 하나의 체크포인트에서 thinking 모드와 non-thinking 모드를 토큰 제어로 선택할 수 있고, FP8/NVFP4 운영 경로와 vLLM·Transformers 통합을 통해 실무적 배포 경로가 마련되어 있다.

핵심 역량

• A.X K2는 토큰당 8개의 활성화된 전문가로 라우팅되는 256 routed experts + 1 shared expert 구조로 작동하여 고용량 지식과 추론 능력을 유지하면서 활성 파라미터를 33B로 제한한다. 이 구조는 라우터가 각 입력을 스코어링하여 상위 experts만 계산에 참여시키는 동작을 통해 추론 비용을 통제한다. 대규모 파라미터 풀을 가진 상태에서 실용적 추론 효율을 달성하는 것이 목적이다.
• Sparse Gated Attention(SGA)는 경량 인덱서로 KV 후보에 점수를 매기고 top-k(=2048) 토큰을 선택한 뒤 MLA에서 latent attention을 수행하며, 각 attention head에 특화된 output gate가 attention 출력을 조절한다. 이 파이프라인은 긴 문맥에서 전체 attention 연산을 희소화하고 attention sink를 억제하여 TPOT와 TTFT를 낮추는 역할을 한다. 인덱서와 gate의 상호작용은 한정된 top-k 예산을 유의미한 위치에 집중시키는 메커니즘으로 작동한다.
• 네이티브 FP8(MXFP8, E4M3) 학습과 block-scaled FP8 체크포인트는 학습과 서빙에서 별도 양자화 단계 없이 저정밀 경로를 활용하게 만든다. FP8 forward/backward를 사용하면서 FP32 마스터 가중치와 그래디언트를 유지하여 안정성을 확보했고, 체크포인트는 128×128 블록 스케일을 사용해 NVFP4 같은 narrow 포맷에서도 블록 단위 활용률을 개선한다. 이 덕분에 디스크 상의 FP8 가중치는 약 647 GiB이며 로드 시 실제 GPU 집합 메모리 요구는 약 800 GiB로 산정되었다.

강점

• 수학과 한국어 성능에서 우수한 점수를 기록하여 도메인 전문성에서 강점을 보인다. 예를 들어 AIME26에서 97.1%를, KMMLU-Pro에서 80.5%를 달성하여 비교군 중 최고 점수를 보였다. 이 수치는 모델 아키텍처와 대규모 고품질 SFT·RL 파이프라인이 결합되어 고난이도 추론과 한국어 이해에서 우위에 있음을 시사한다.
• 긴 문맥 처리 능력이 탁월하여 RULER 평가에서 overall 94.6을 기록했고 256K까지 강한 성능을 유지했다. 또한 NIAH(needle-in-a-haystack) 평가에서 YaRN zero-shot 확장하에 모든 컨텍스트 길이와 니들 깊이에서 100을 달성하여 장문 내 사실 검색에서 완전한 회수 능력을 보였다. 이러한 결과는 SGA와 MLA 기반의 sparse attention 경로가 긴 시퀀스에서 실용적이고 정확한 정보 검색을 가능하게 했음을 나타낸다.
• FP8 네이티브 학습과 block-scaled FP8 체크포인트 덕분에 저정밀 서빙 경로를 직접 제공하여 메모리와 비용 효율을 확보했다. 체크포인트 크기는 약 647 GiB이며, 로드 시 KV 캐시·액티베이션을 포함해 약 800 GiB 이상의 GPU 집합 메모리가 필요하다고 명시되어 있다. 이러한 자원 수치는 FP8 기반 배포가 실제로 메모리 이득을 주면서도 운영상 큰 리소스 요구가 수반됨을 보여준다.

훈련 방식

A.X K2는 from-scratch로 네이티브 FP8(MXFP8, E4M3) forward/backward 학습을 수행한 대규모 MoE 모델이다. 학습은 세 단계 커리큘럼(6.4T 일반 지식 → 1.4T 고품질 추론 → 0.36T 장문 확장)으로 구성되어 데이터의 난이도와 비중을 단계별로 조정하였다. 이후 paired SFT(≈60.6B tokens)로 Think-Fusion 컨트롤 토큰을 학습시키고 multi-stage on-policy RL 및 안전성 red-teaming을 적용하여 agentic 행동과 안전 신호를 튜닝했다.

알아두면 좋은 것

• A.X K2는 688B 총파라미터와 토큰당 33B 활성파라미터를 가진 대형 Sparse MoE 모델로, 256 routed experts와 1 shared expert 구조에서 토큰당 8 experts가 활성화된다. Think-Fusion 레시피로 thinking 모드와 non-thinking 모드를 하나의 트랙에서 학습하여 요청별로 품질과 비용을 교환할 수 있게 설계되었다. 체크포인트는 block-scaled FP8(E4M3)으로 공개되어 FP8 네이티브 서빙을 지원한다.
• Sparse Gated Attention(SGA)는 top-k 인덱서와 head-specific gated MLA를 결합해 긴 문맥에서 attention 연산을 희소화하고 attention sink를 억제하는 구조로 동작한다. Gated Norm(GN)은 RMSNorm 직후에 input-dependent gate를 적용하여 대형 활성화의 외란을 억제하고 저정밀 연산에서 안정성을 높인다. YaRN을 통해 128K에서 256K, 512K로의 zero-shot 확장이 가능하며 needle-in-a-haystack(NIAH) 평가에서 모든 위치에서 100점을 기록했다.
• 학습 데이터는 약 8.2T 토큰이며 단계별 커리큘럼(6.4T 일반 지식 → 1.4T 고품질 추론 → 0.36T 장문 확장)을 사용했다. 공개 데이터는 Nemotron-CC-v2.1, Nemotron-Pretraining-Code-v2, FineWeb2 등이며 한국어는 내부 크롤링으로 약 1.37T 토큰이 포함되어 있다. SFT와 RL 단계를 거쳐 Think-Fusion paired SFT(≈60.6B tokens)와 후속 RL이 적용되었다.

기술적 특징

• Sparse Gated Attention(SGA)는 경량 indexer로 KV 후보의 position score를 산출하고 selector가 top-k(=2048)를 선택한 뒤 Multi-head Latent Attention(MLA)에서 latent KV에 대해 attention을 계산하는 파이프라인으로 동작한다. 각 head에는 output gate가 적용되어 attention sink를 억제하고 인덱서 예산을 중요 위치에 집중시키는 상호보완적 메커니즘으로 설계되었다. 이 구조는 긴 문맥에서 attention 연산을 희소화하면서도 학습 중 손실 수렴을 개선하도록 고안되었다.
• Gated Norm(GN)은 RMSNorm 이후에 input-dependent gate를 적용하여 정규화 출력의 외란을 억제하고 레지듀얼 경로로 들어가는 활성화 분포를 안정화한다. 이 기법은 A.X K1에서 사용하던 dual-normalization을 대체하여 아키텍처 복잡도를 줄이면서도 동일하거나 우수한 학습 안정성을 확보했다. 외란 억제된 활성화 분포는 블록 스케일 FP8과 NVFP4 같은 좁은 저정밀 포맷에서 블록 활용률과 정밀도를 개선하는 효과를 냈다.
• 네이티브 FP8 학습은 MXFP8 E4M3 형식을 전·후방에 적용하였고 FP32 마스터 가중치와 그래디언트를 유지하여 안정성을 확보했다. 체크포인트는 128×128 블록 단위로 스케일이 기록되어 NVFP4와 같은 전문가 전용 W4A4 양자화 경로에서도 재현 가능한 블록 스케일을 제공한다. Transformers 통합과 vLLM fork에서 FP8 경로가 지원되어 별도의 포스트 훈련 양자화 없이도 FP8 서빙이 가능하도록 구성되었다.

차별점

• 단일 모델에서 Think-Fusion을 통해 reasoning(trace 포함)과 non-thinking(간결 응답) 모드를 토큰 기반 제어로 모두 제공하는 점이 차별화 요소이다. paired SFT 데이터에서 각 프롬프트에 대해 thinking과 non-thinking 응답을 함께 학습시키고 제어 토큰으로 모드를 전환하게 했기 때문에 런타임에서 호출별로 품질·지연 트레이드오프를 조절할 수 있다. 이 접근은 별도 모델이나 dual-track 파이프라인 없이 하나의 체크포인트로 두 가지 운영 패턴을 지원하게 만든다.
• 네이티브 FP8 학습과 block-scaled FP8 체크포인트 배포로 FP8 서빙 경로를 기본 제공하는 점이 운영 측면에서 차별화된다. 많은 대형 모델이 후처리 양자화에 의존하는 반면 A.X K2는 학습부터 FP8을 적용하여 양자화로 인한 추가 단계와 정밀도 저하를 줄였다. 이로 인해 FP8 환경에서의 메모리 절감과 함께 NVFP4 등 협소 포맷으로의 전환성도 확보되었다.

벤치마크

이미지 분석