ATOM (AiTer Optimized Model) is a lightweight vLLM-like implementation, focusing on integration and optimization based on AITER.

📢 News

• [2026/08] ATOM now supports Qwen3.8 (Qwen3_5MoeForCausalLM, text-only MoE). The amd/Qwen3.8-2.4T-A95B-Quark-MXFP4 checkpoint, whose MoE (expert) weights are quantized to MXFP4, runs directly on a single 8-rank MI355 machine.
• [2026/07] Featured AMD Developer Article: Day 0 Kimi-K3 Inference Deployment with ATOM on AMD Instinct MI355X GPUs walks through Day 0 deployment of the 2.78T-parameter Kimi-K3 (KDA + Gated MLA) on a single 8x MI355X node with TP8 — why the ~1.56 TB checkpoint fits, how the weights are distributed under TP8, and how to bring the model up with ATOM and run a minimal correctness check. See Kimi-K3 recipe.
• [2026/07] ATOM now supports DeepSeek-V4-Pro DSpark speculative decoding — a semi-autoregressive block drafter (parallel backbone + Markov head + confidence head) with confidence-scheduled ragged verification, FP8 KV cache, DP attention, and PIECEWISE CUDA graphs. See DSpark recipe.
• [2026/06] ATOM now supports MiniMax-M3 inference on the native OpenAI-compatible server path, including MXFP4/MXFP8 checkpoints, FP8 KV cache, and EAGLE3 speculative decoding. See MiniMax-M3 recipe.
• [2026/06] Featured ROCm Blog: DP Attention and TBO for DeepSeek-V4 on MI355X highlights how ATOM optimizes DeepSeek-V4 inference on AMD Instinct MI355X GPUs with DP Attention using all-gather/reduce-scatter and Two-Batch Overlap, achieving strongly competitive DeepSeek-V4 inference performance.
• [2026/06] Featured ROCm Blog: ATOMesh: Unlocking AMD Hardware for Scalable LLM Serving explains how ATOMesh orchestrates distributed inference on AMD GPUs with ATOM, AITER, MORI, and RCCL.
• [2026/06] Featured ROCm Blog: ATOM: Unlocking Extreme AMD Instinct Inference with Software-Hardware Co-Optimization covers ATOM architecture, feature scope, model coverage, and benchmark dashboard usage.
• [2026/06] Experimental Navi 4 (RDNA4 / gfx1201) support — AMD Radeon RX 9070 / RX 9070 XT and Radeon AI PRO R9700. See the Qwen3-8B-FP8 and Ministral-3-8B recipes.
• [2026/06] ATOM now supports GLM-5.2 (glm_moe_dsa) in FP8, including the new IndexShare DSA schedule (shared layers reuse the preceding full layer's indexer). See GLM-5.2 recipe.
• [2026/05] Featured ROCm Blog: vLLM-ATOM: Unlocking Native AMD Performance in the vLLM Ecosystem shows how ATOM integrates with vLLM as an AMD-optimized plugin path.
• [2026/05] ATOM now supports Qwen3.5 multimodal image+text inference on the native engine and OpenAI-compatible chat API. See Qwen3.5 multimodal recipe.
• [2026/05] ATOM now supports online quantization — re-quantize unquantized or FP8-block source checkpoints to PTPC-FP8 / MXFP4 mixed precision at load time via --online_quant_config, no offline re-packing required. See online quantization guide.
• [2026/05] Dissecting DeepSeek V4 Compressor — interactive animation visualizing how the CSA/HCA compressor state cache works (overlap mechanism, prefill vs decode, bulk compression vs sequential accumulation).
• [2026/05] DeepSeek V4-Pro PD disaggregation — Prefill/Decode separation now supports DeepSeek V4-Pro with Mooncake RDMA KV cache transfer. See V4 recipe and PD guide.
• [2026/05] ATOM now supports Prefill/Decode (P/D) disaggregation with Mooncake RDMA push-mode KV cache transfer. See PD disaggregation guide.
• [2026/03] ATOM now supports Prefill/Decode (P/D) disaggregation — run prefill and decode on separate GPU nodes with RDMA-based KV cache transfer via MORI-IO. See disaggregation docs.

🚀 Features

• ROCm Optimized: Built on AMD's ROCm platform with AITER kernels (ASM, CK, Triton)
• OpenAI-Compatible API: Drop-in server with /v1/chat/completions and /v1/completions endpoints
• Piecewise torch.compile: 4 compilation levels with CUDA graph capture for low-latency decode
• Multi-GPU Parallelism: Tensor parallelism (TP), data parallelism (DP), and expert parallelism (EP) with MORI all-to-all
• Two-Batch Overlap (TBO): Following DeepSeek's system design, TBO splits each batch into two micro-batches and pipelines them across compute and communication streams. Effectively hiding expert-parallel communication latency and reducing peak memory usage. See recipe
• Quantization: FP8, MXFP4, INT8, INT4 with auto-detection from HuggingFace configs
• Speculative Decoding: Multi-Token Prediction (MTP) with EAGLE proposer
• Prefix Caching: xxhash64-based KV cache block sharing across sequences

Supported Models

| Model Family | HF Architecture | Dense/MoE | Notes | |---|---|---|---| | Llama | LlamaForCausalLM | Dense | Llama 2, Llama 3, Llama 3.1 | | Qwen3 | Qwen3ForCausalLM | Dense | | | Qwen3-MoE | Qwen3MoeForCausalLM | MoE | 128 experts, top-8 routing | | Qwen3-Next | Qwen3NextForCausalLM | MoE | Hybrid full attention + Gated DeltaNet | | Qwen3.8 | Qwen3_5MoeForCausalLM | MoE | 2.4T total / A95B active, 512 experts top-10, hybrid full attention + Gated DeltaNet, 1-layer MTP. MXFP4 MoE weights | | DeepSeek V2/V3 | DeepseekV3ForCausalLM | MoE | MLA attention, MTP speculative decoding | | Mixtral | MixtralForCausalLM | MoE | 8 experts, top-2 routing | | GLM-4-MoE | Glm4MoeForCausalLM | MoE | | | GLM-5 / GLM-5.2 | GlmMoeDsaForCausalLM | MoE | MLA + DSA sparse attention, similar to DeepSeek V3.2; GLM-5.2 adds IndexShare. See recipe | | GPT-OSS | GptOssForCausalLM | MoE | Sliding window + attention sinks | | Kimi-K2 | via --trust-remote-code | MoE | See recipe | | MiMo V2/V2.5 | MiMoV2ForCausalLM | MoE | Hybrid full + SWA attention, 3-layer MTP. See recipe |

📋 Requirements

• AMD GPU with ROCm support
• Docker

🛠️ Installation

Option A: Nightly Image (Recommended)

Pre-built image with AITER + ATOM ready to use:

docker pull rocm/atom-dev:latest
docker run -it --network=host \
  --device=/dev/kfd \
  --device=/dev/dri \
  --group-add video \
  --cap-add=SYS_PTRACE \
  --security-opt seccomp=unconfined \
  -v $HOME:/home/$USER \
  -v /mnt:/mnt \
  -v /data:/data \
  --shm-size=16G \
  --ulimit memlock=-1 \
  --ulimit stack=67108864 \
  rocm/atom-dev:latest

Option B: Build from Base ROCm Image

1. Pull and run the base image

docker pull rocm/pytorch:rocm7.0.2_ubuntu24.04_py3.12_pytorch_release_2.8.0
docker run -it --network=host \
  --device=/dev/kfd \
  --device=/dev/dri \
  --group-add video \
  --cap-add=SYS_PTRACE \
  --security-opt seccomp=unconfined \
  -v $HOME:/home/$USER \
  -v /mnt:/mnt \
  -v /data:/data \
  --shm-size=16G \
  --ulimit memlock=-1 \
  --ulimit stack=67108864 \
  rocm/pytorch:rocm7.0.2_ubuntu24.04_py3.12_pytorch_release_2.8.0

2. Install AITER and ATOM inside the container

pip install amd-aiter
git clone https://github.com/ROCm/ATOM.git && pip install ./ATOM

💡 Usage

Basic Example

Before running the example, please install ninja and the Hugging Face CLI, and log in to your account.

pip install ninja
pip install -U "huggingface_hub"
hf auth login

The default optimization level is 3 (piecewise torch.compile with CUDA graphs).

python -m atom.examples.simple_inference --model meta-llama/Meta-Llama-3-8B --kv-cache-dtype fp8

Note: First-time execution may take approximately 10 minutes for model compilation.

Flags accept both spellings: every long flag can be written kebab-case or

snake_case interchangeably (e.g. --kv-cache-dtype == --kv_cache_dtype).

Serving

Start an OpenAI-compatible server:

# Single GPU
python -m atom.entrypoints.openai_server --model Qwen/Qwen3-0.6B --kv-cache-dtype fp8
Multi-GPU with tensor parallelism
python -m atom.entrypoints.openai_server --model deepseek-ai/DeepSeek-R1 --kv-cache-dtype fp8 -tp 8
With MTP speculative decoding
python -m atom.entrypoints.openai_server --model deepseek-ai/DeepSeek-R1 --kv-cache-dtype fp8 -tp 8 \
  --method mtp --num-speculative-tokens 3

📊 Performance

Live Benchmark Dashboard

rocm.github.io/ATOM/benchmark-dashboard

The dashboard tracks nightly performance across models and configurations:

• Interactive vs Throughput — tok/s/user vs tok/s/gpu tradeoff across concurrency levels
• Throughput & Latency trends — Output throughput, TTFT, TPOT over time, grouped by model
• Regression detection — Automatic alerts when throughput drops >5% or latency increases >10%
• Profiler trace collection — On regression, automatically re-runs with PyTorch profiler and uploads traces

Online Serving Throughput

!DS R1 Performance

For more information, visit InferenceX.

Benchmarking

Run an online throughput benchmark against a running server:

python -m atom.benchmarks.benchmark_serving \
  --model=deepseek-ai/DeepSeek-R1 --backend=vllm --base-url=http://localhost:8000 \
  --dataset-name=random \
  --random-input-len=1024 --random-output-len=1024 \
  --random-range-ratio=0.8 \
  --num-prompts=1280 --max-concurrency=128 \
  --request-rate=inf --ignore-eos \
  --save-result --percentile-metrics="ttft,tpot,itl,e2el"

Profiling & Trace Analysis

Collect a Trace

Launch the server with --torch-profiler-dir and --mark-trace:

python -m atom.entrypoints.openai_server \
  --model deepseek-ai/DeepSeek-R1 --kv-cache-dtype fp8 -tp 8 \
  --torch-profiler-dir ./trace --mark-trace

Collect traces via benchmark --profile flag (auto start/stop):

python -m atom.benchmarks.benchmark_serving \
  --model=deepseek-ai/DeepSeek-R1 --backend=vllm --base-url=http://localhost:8000 \
  --dataset-name=random --random-input-len=1024 --random-output-len=1024 \
  --num-prompts=128 --max-concurrency=128 \
  --request-rate=inf --ignore-eos --profile

Or control profiling manually on a running server:

curl -X POST http://127.0.0.1:8000/start_profile
... run your workload ...
curl -X POST http://127.0.0.1:8000/stop_profile

Analyze the Trace

# Kernel breakdown per layer → Excel
python tools/parse_trace.py ./trace/rank_0/DeepSeek-R1_ts_*.json.gz
Performance summary → Markdown report
python tools/analyze_trace_summary.py ./trace/rank_0/DeepSeek-R1_ts_*.json.gz

| Output | Description | |---|---| | prefill_breakdown.xlsx | Per-kernel duration, call count, pct%, module grouping, cross-layer averages | | decode_breakdown.xlsx | Same for decode phase, with CUDAGraph kernel mapping | | performance_summary.md | Prefill/decode/draft step timing, iteration breakdown |

Accuracy Validation

pip install lm-eval[api]
Start server, then run evaluation
lm_eval --model local-completions \
  --model_args model=meta-llama/Meta-Llama-3-8B,base_url=http://localhost:8000/v1/completions,num_concurrent=64,max_retries=3,tokenized_requests=False \
  --tasks gsm8k --num_fewshot 5

📚 Documentation

Full documentation: rocm.docs.amd.com/projects/atom/en/latest/

| Topic | Description | Guide | |---|---|---| | Architecture | System overview, request lifecycle, component design | Architecture Guide | | Configuration | Config classes, CLI arguments, environment variables | Configuration Guide | | Model Support | Supported models, weight loading, adding new architectures | Model Support Guide | | Model Operations | AITER kernel integration, linear/attention/MoE/norm wrappers | Model Ops Guide | | Scheduling & KV Cache | Batch scheduling, block allocation, prefix caching | Scheduling Guide | | Compilation | torch.compile levels, CUDA graphs, piecewise compilation | Compilation Guide | | Distributed | Tensor/data/expert parallelism, multi-GPU deployment | Distributed Guide | | Serving & Benchmarks | OpenAI API server, benchmarking, profiling, speculative decoding | Serving Guide | | Observability (experimental) | Opt-in request tracing, configuration and test tools | OpenTelemetry Guide | | Environment Variables | All ATOM_* variable definitions | Env Vars |

Deployment Recipes:

• DeepSeek-V4.1-Flash AgentX — TP2/TP4, no EP, DSpark5, fixed AL=3.51 benchmark commands and measured operating points
• DeepSeek-R1 — FP8/MXFP4 with MTP speculative decoding on 8 GPUs
• Qwen3-235B-A22B — TP8 + EP with FP8 KV cache
• Qwen3-Next — Hybrid GDN + MoE architecture
• Qwen3.8-Flash-Next — Hybrid GDN + sparse attention (QSA) + n-gram memory, multimodal;
• Kimi-K2-Thinking — MXFP4 MoE on 4 GPUs
• GLM-5 — FP8 MoE with MLA on 8 GPUs
• GPT-OSS-120B — Single GPU or DP+EP on 2 GPUs
• TBO (Two-Batch Overlap) — Compute-communication overlap for MoE models with DP attention
• PD Disaggregation (Mooncake) — Prefill/Decode separation with RDMA KV cache transfer (DeepSeek-R1, DeepSeek-V4-Pro)

• vLLM Plugin Backend — ATOM as the out-of-tree plugin backend for vLLM

Acknowledgements

This project was adapted from nano-vllm.

Support & Reporting Issues

We welcome issues and contributions! Please use the GitHub Issues page to report bugs or request features: https://github.com/ROCm/ATOM/issues