Performance summary → Markdown report

May 7, 2026 · View on GitHub

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

📢 News

[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/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
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	`GlmMoeDsaForCausalLM`	MoE	MLA attention, similar to DeepSeek V3.2. See recipe
GPT-OSS	`GptOssForCausalLM`	MoE	Sliding window + attention sinks
Kimi-K2	via `--trust-remote-code`	MoE	See recipe
MiMo-V2-Flash	`MiMoV2FlashForCausalLM`	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.

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

Models tracked: DeepSeek-R1-0528 (FP8 & MTP3), GLM-5-FP8, gpt-oss-120b

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 --layer 3

# 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.github.io/ATOM/docs

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
Environment Variables	All `ATOM_*` variable definitions	Env Vars

Deployment Recipes:

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
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

Framework Integration:

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