언어 모델 코드 하나로 짜서 파이토치, JAX, MLX, vLLM 어디서든 돌리는 새 언어 Axon
arXiv:2608.198892026-08-21
Write Once, Run Everywhere: The Axon DSL for Shape-Safe and Framework-Agnostic LLM Architectures
언어 모델 코드 하나로 짜서 파이토치, JAX, MLX, vLLM 어디서든 돌리는 새 언어 Axon
오픈소스 언어모델 생태계는 사실상 하나의 플랫폼(허깅페이스)에 의존하고 있어, 새 하드웨어나 프레임워크로 옮기려면 모델을 사람이 일일이 다시 짜야 했다. 연구진은 하스켈 문법을 닮은 강타입 전용 언어 Axon을 만들어, 모델 구조를 한 번만 정의하면 컴파일러가 파이토치, 트라이톤, JAX, MLX, vLLM용 코드를 자동으로 만들어내게 했다. 1억 3500만에서 320억 파라미터급 모델 467개 실험에서 기존 트랜스포머스 라이브러리 대비 대체로 더 빠른 속도를 보였다.
METAL MEDIA 해설 도표
언어 모델 코드 하나로 짜서 파이토치, JAX, MLX, vLLM 어디서든 돌리는 새 언어 Axon
01문제: 언어모델 코드는 특정 플랫폼(허깅페이스)에 묶여 있어 다른 프레임워크로 옮길 때마다 손으로 다시 구현해야 하고, 이 과정에서 최적화가 누락되거나 어긋나는 일이 잦다
02해법: Axon이라는 전용 프로그래밍 언어를 만들어 모델 구조(레이어, 텐서 모양, 파라미터 위치)를 한 번만 기술하면, 컴파일러가 이를 중간표현(Graph IR)으로 바꾼 뒤 다섯 개 백엔드용 독립 실행 코드를 자동 생성한다
03타입 시스템이 텐서의 모양(shape)을 미리 검사해 계산 도중 모양이 안 맞는 오류를 원천 차단한다
04결과: 파이토치 대비 중간값 7% 빠름, 트라이톤 결합 파이토치 12%, JAX 91%, MLX 107% 빠름, vLLM 배포 시 트랜스포머스 대비 중간값 58% 빠름(467회 추론 벤치마크 기준)
05학습 실험에서도 Axon으로 만든 파이토치 모델이 기존 트랜스포머스 구현과 거의 동일한 손실 곡선을 보이면서 스텝당 약 9.6% 더 빨랐다
METAL MEDIA이 원문을 바탕으로 재구성한 해설 도표이며, 논문 저자의 원문 figure가 아닙니다.
무엇을 했나
문제: 언어모델 코드는 특정 플랫폼(허깅페이스)에 묶여 있어 다른 프레임워크로 옮길 때마다 손으로 다시 구현해야 하고, 이 과정에서 최적화가 누락되거나 어긋나는 일이 잦다
해법: Axon이라는 전용 프로그래밍 언어를 만들어 모델 구조(레이어, 텐서 모양, 파라미터 위치)를 한 번만 기술하면, 컴파일러가 이를 중간표현(Graph IR)으로 바꾼 뒤 다섯 개 백엔드용 독립 실행 코드를 자동 생성한다
타입 시스템이 텐서의 모양(shape)을 미리 검사해 계산 도중 모양이 안 맞는 오류를 원천 차단한다
결과: 파이토치 대비 중간값 7% 빠름, 트라이톤 결합 파이토치 12%, JAX 91%, MLX 107% 빠름, vLLM 배포 시 트랜스포머스 대비 중간값 58% 빠름(467회 추론 벤치마크 기준)
학습 실험에서도 Axon으로 만든 파이토치 모델이 기존 트랜스포머스 구현과 거의 동일한 손실 곡선을 보이면서 스텝당 약 9.6% 더 빨랐다
Figure 1: Write-once, run everywhere. Axon DSL compiles axon definitions (.axon) to standalone model definitions for PyTorch, JAX, MLX and vLLM. All that is needed is a standard safetensors checkpoint.
Table 1: Per-backend comparison of Axon’s autoregressive generation performance with decoder-only models against PyTorch, Triton, and JAX. “Axon ≤1×” = count (and %) of checkpoints where Axon is faster than Transformers. “Axon >1×” = count (%) of checkpoints where Axon is slower. Median and mean report runtime ratio (lower means Axon is faster).
Backend
Checkpoints
Axon ≤1×
Axon >1×
Median ratio
Mean ratio
<4B
PyTorch
76
48 (63%)
28 (37%)
0.903
0.925
Triton
75
60 (80%)
15 (20%)
0.843
0.878
JAX
74
64 (86%)
10 (14%)
0.481
1.122
<4B total
225
172 (76%)
53 (24%)
0.804
0.974
4–32B
PyTorch
87
55 (63%)
32 (37%)
0.987
0.980
Triton
87
67 (77%)
20 (23%)
0.924
0.931
JAX
68
55 (81%)
13 (19%)
0.589
0.981
4–32B total
242
177 (73%)
65 (27%)
0.908
0.963
Figure 2: Autoregressive generation performance with decoder-only models with up to 4B parameters. Log-log plot where the diagonal line indicates equal performance. Points above the parity line indicate Axon is faster, color-fill denotes backends and border-line denotes dtype.
Table 2: Per-backend comparison of Axon’s Forward performance with Encoder-Only and Encoder-Decoder Models against PyTorch, Triton, and JAX backends. “Axon ≤1×” = count (and %) of checkpoints where Axon is faster than Transformers. “Axon >1×” = count (%) of checkpoints where Axon is slower. Median and mean report runtime ratio (lower means Axon is faster).
Backend
Checkpoints
Axon ≤1×
Axon >1×
Median ratio
Mean ratio
≤4B
PyTorch
26
15 (58%)
11 (42%)
0.883
1.786
Triton
26
9 (35%)
17 (65%)
1.332
2.169
JAX
26
7 (27%)
19 (73%)
1.079
1.260
≤4B total
78
31 (40%)
47 (60%)
1.084
1.738
4–32B
PyTorch
20
13 (65%)
7 (35%)
0.990
1.506
Triton
20
8 (40%)
12 (60%)
1.133
1.546
JAX
14
0 (0%)
14 (100%)
3.747
3.356
4–32B total
54
21 (39%)
33 (61%)
1.175
2.001
Figure 3: Benchmarking Autoregressive models between 4B and 32B models. Log-log plot where the diagonal parity line indicates equal performance. Generally, Axon produces comparable or superior throughput on the bigger models.
Table 3: Autoregressive generation comparison of Axon (vLLM native) against Transformers. “Axon ≤1×” = number (and %) of checkpoints where Axon is faster than Transformers. “Axon >1×” = number (and %) of checkpoints where Axon is slower than Transformers. Median and mean report Axon-to-Transformers runtime ratio (lower is faster).
Model size
Checkpoints
Axon ≤1×
Axon >1×
Median ratio
Mean ratio
Small (≤4B)
54
38 (70%)
16 (30%)
0.656
1.430
Large (4B–32B)
34
27 (79%)
7 (21%)
0.609
2.889
Total
88
65 (74%)
23 (26%)
0.631
1.994
Figure 4: vLLM: Axon (vLLM native) vs. Transformers generation throughput on 88 checkpoints. Points above the parity line indicate Axon is faster. 74% fall above.
Table 4: Breakdown of Axon performance by precision, model type, and sequence length. “Axon ≤1×” = number (and %) of checkpoints where Axon is faster than Transformers. “Axon >1×” = number (and %) of checkpoints where Axon is slower than Transformers. Median and mean report Axon-to-Transformers runtime ratio (lower = Axon is faster).
Group
Points
Axon ≤1×
Axon >1×
Median ratio
Mean ratio
BF16
63
61 (97%)
2 (3%)
0.423
0.500
FP32
63
58 (92%)
5 (8%)
0.503
0.588
causal_lm
84
82 (98%)
2 (2%)
0.390
0.437
seq2seq_lm
42
37 (88%)
5 (12%)
0.574
0.752
len=64
42
42 (100%)
0 (0%)
0.417
0.448
len=128
42
39 (93%)
3 (7%)
0.483
0.525
len=256
42
38 (90%)
4 (10%)
0.494
0.652
Total
126
119 (95%)
7 (5%)
0.483
0.544
Figure 5: MLX. Benchmarking conventional HF models against Axon derived standalone model definitions. Axon yields some considerable speed-ups across the board.
Table 5: Axon Compiler Phase and Stage with Invariants
Phase
Representation
Primary invariant
Parse
one-file AST
Syntactic structure and explicit MAIN pragma insertion
Load
loaded AST set
Imports and builtins located without rewriting semantics
Materialize
one-file AST
Optional checkpoint/config specialization for generic models
Resolve/validate-closed
closed AST
No unresolved imports or names; unreachable definitions pruned from MAIN
Normalize
normalized AST
Call syntax, pipes, path sugar, and zero-arg call/name distinctions made explicit
Elaborate/validate-elaborated
elaborated AST
Default arguments filled and call arguments positionalized
Flatten/validate-flat
flat AST
Explicit evaluation order; flat calls and binds accepted by typecheck and Graph IR lowering
Typecheck/validate-typed
Typed flat AST
expression types, arities, dimensions, and primitive rules applied to a fixpoint
Optimize-ast
typed flat AST
Optional conservative AST cleanup with retype/validation
Graph lowering/validation
Graph IR
Typed graph modules, multi-output nodes, structured paths, constraints, and metadata
Optimize-graph/validation
Graph IR
Optional graph cleanup, specialization, backend-neutral rewrites, and opt-in backend intrinsics
Backend
generated/runtime code
Executable tensor program consuming the validated Graph IR contract
Figure 12: Benchmarking Encoder-only and Encoder–Decoder up to 4B models. Log-log plot where the diagonal parity line indicates equal performance. Here, we see the majority of the points in the vicinity below the parity line.
Table 6: Wall-clock time and GPU idle for autoregressive generation (112 tokens for Qwen 0.5B, 96 for Pleias 3.8B). “—” indicates the metric is not directly measurable: JAX has no torch.profiler equivalent for per-kernel timing, and GPU active/idle is a synchronous-execution concept that does not apply to JAX’s async dispatch model.
Qwen2.5-0.5B (0.67B)
Pleias-3b-Preview (3.8B)
Metric
Transformers
Axon-Torch
Axon-JAX
Transformers
Axon-Torch
Axon-JAX
Wall-clock (ms)
1380
1185
541
1091
1071
1160
Speedup vs Transformers
1.0×
0.86×
0.39×
1.0×
0.98×
1.07×
GPU active (ms)
373
370
—
557
561
—
GPU idle (ms)
1007
815
—
534
510
—
GPU idle %
73.0%
68.8%
0.0%
48.9%
47.6%
0.0%
CUDA kernels
129,189
120,987
—
101,935
104,123
—
Figure 13: Benchmarking Encoder-only and Encoder-Decoder between 4B and 32B models. Log-log plot where the diagonal parity line indicates equal performance.
Table 7: Python function call counts and cProfile cumulative time. “—” indicates the category does not apply: JAX fuses all per-op dispatches into a single XLA compilation per step, so individual F.linear and SDPA calls do not exist as Python-level dispatches.
Qwen2.5-0.5B (0.67B)
Pleias-3b-Preview (3.8B)
Metric
Transformers
Axon-Torch
Axon-JAX
Transformers
Axon-Torch
Axon-JAX
Total Python calls
455,769
508,390
137,805
359,909
330,929
82,210
Calls vs Transformers
1.0×
1.12×
0.30×
1.0×
0.92×
0.23×
cProfile time (s)
1.13
0.95
0.54
0.87
0.78
1.16
Per-op dispatch (Python calls per generate pass):
nn.Module.__call__
35,616
0
0
28,032
0
0
F.linear
18,928
18,928
—
14,880
14,880
—
SDPA
2,688
2,688
—
2,112
2,112
—
rope_apply
0
5,376
—
0
4,224
—
forward (jit dispatch)
—
—
112
—
—
96
Figure 14: Training of Gemma 3 270M on summarization task for 2000 steps, demonstrating identical training behaviour between conventional Transformers-definition and Axon derived PyTorch model. Note the Axon and Transformers loss curves are on top of each other together.
왜 중요한가
지금은 소수 플랫폼이 사실상 표준을 쥐고 있어 그 플랫폼이 문제를 겪으면 생태계 전체가 흔들리는 구조인데, Axon은 코드 명세 자체를 공유 표준으로 삼아 이런 종속을 줄이려는 시도다. 연구자나 소규모 팀이 특정 회사의 최적화 인프라 없이도 자체 하드웨어에서 빠른 모델을 돌릴 수 있게 해준다는 점에서 실무적 의미가 크다.
이 논문의 용어
DSL(도메인 전용 언어) · 특정 목적(여기서는 신경망 구조 기술)에 맞춰 설계된 프로그래밍 언어
강타입 · 값의 종류와 형태를 컴파일 시점에 엄격히 검사하는 성질
Graph IR(중간표현) · 여러 백엔드가 공통으로 사용하는 컴파일 중간 단계 데이터 구조
PagedAttention/KV-cache · vLLM이 사용하는 메모리 관리 기법으로, 생성 중인 토큰들의 캐시를 효율적으로 다루어 서빙 속도를 높인다
top-1 token parity · 서로 다른 구현이 매 단계 동일한 다음 토큰을 예측하는지 확인하는 정확성 기준
저자 · Jacob Nielsen, Danial Namazifard, Lukas Galke Poech, Peter Schneider-Kamp
