ATOM serving & benchmarking guide#
ATOM (AiTer Optimized Model) is AMD’s lightweight LLM inference engine built on AITER kernels for ROCm/HIP GPUs. This guide covers the OpenAI-compatible serving API, programmatic engine usage, benchmarking tools, profiling, and speculative decoding.
Quick reference#
# Start the OpenAI-compatible server
python -m atom.entrypoints.openai_server --model <model_name_or_path> --kv_cache_dtype fp8
# Run the online serving benchmark
python -m atom.benchmarks.benchmark_serving \
--backend vllm --model <model_name_or_path> \
--base-url http://localhost:8000 \
--dataset-name random --random-input-len 1024 --random-output-len 128 \
--num-prompts 1000 --request-rate inf --ignore-eos
# Simple inference example
python -m atom.examples.simple_inference --model <model_name_or_path> --kv_cache_dtype fp8
# Offline profiling
python -m atom.examples.profile_offline --model <model_name_or_path> --kv_cache_dtype fp8
# Accuracy validation with lm-eval
lm_eval --model local-completions \
--model_args model=<model>,base_url=http://localhost:8000/v1/completions,num_concurrent=64,max_retries=3,tokenized_requests=False \
--tasks gsm8k --num_fewshot 5
OpenAI-compatible server#
The server is implemented in atom/entrypoints/openai_server.py using FastAPI
and Uvicorn. It exposes OpenAI-compatible HTTP endpoints so that existing
clients (curl, OpenAI SDK, lm-eval) work without modification.
Endpoints#
Method |
Path |
Description |
|---|---|---|
|
|
Chat completion (ChatCompletionRequest -> ChatCompletionResponse) |
|
|
Text completion (CompletionRequest -> CompletionResponse) |
|
|
List available models |
|
|
Health check (returns |
|
|
Start torch profiler on the engine |
|
|
Stop torch profiler and flush traces |
Request models#
ChatCompletionRequest fields:
Field |
Type |
Default |
Description |
|---|---|---|---|
|
|
|
Model name (validated against the loaded model) |
|
|
|
List of chat messages ( |
|
|
|
Alias for |
|
|
|
Sampling temperature |
|
|
|
Nucleus sampling threshold |
|
|
|
Maximum tokens to generate |
|
|
|
Stop strings |
|
|
|
Ignore end-of-sequence token |
|
|
|
Enable server-sent events streaming |
|
|
|
Random seed |
CompletionRequest fields:
Field |
Type |
Default |
Description |
|---|---|---|---|
|
|
|
Model name |
|
|
(required) |
Text prompt |
|
|
|
Sampling temperature |
|
|
|
Nucleus sampling threshold |
|
|
|
Maximum tokens to generate |
|
|
|
Stop strings |
|
|
|
Ignore end-of-sequence token |
|
|
|
Enable SSE streaming |
Response models#
Both ChatCompletionResponse and CompletionResponse include:
id— unique request identifier (e.g.chatcmpl-<uuid>orcmpl-<uuid>)object—"chat.completion"or"text_completion"created— Unix timestampmodel— model namechoices— list of generated completionsusage— token counts (prompt_tokens,completion_tokens,total_tokens) plusttft_s,tpot_s, andlatency_stiming fields
Streaming responses use the SSE (Server-Sent Events) protocol with
data: [DONE]\n\n as the termination signal.
Delivery under load#
The API server is a single Python process, so at high concurrency the fixed per-chunk cost of delivering tokens (detokenize, coroutine wakeup, JSON encode, socket write) can cap throughput before the GPU does. Two things keep that cost down:
Backlog merging. Each request’s chunks land in a
StreamOutputCollector(atom/entrypoints/openai/streaming_dispatch.py), which holds at most one chunk per stream: anything arriving behind an unread one merges into it. Nothing is held back waiting for more, so a consumer that keeps up sees exactly one chunk per engine step.msgspec frame encoding (
atom/entrypoints/openai/sse.py), roughly 5.8x cheaper per frame thanjson.dumps.
A token can be delivered later than the engine produced it, by a bounded
amount. Two stages downstream of the collector read the text for markers —
the reasoning channel’s delimiters
(atom/entrypoints/openai/reasoning.py) and the opening tags of whichever
tool-call format this model uses (atom/entrypoints/openai/tool_parser/) —
and neither may hand out a byte that could turn out to be the first character
of one. Both ask the same
question through MarkerScanner
(atom/entrypoints/openai/marker_scanner.py): release everything except the
longest suffix of the buffer that is a prefix of some marker. The wait is
therefore bounded by the longest marker a format declares, a few dozen bytes,
and is usually zero — a chunk whose tail cannot begin a marker is released
whole.
This is worth stating because it used to be unbounded. The rule was “hold
everything once a marker’s first character appears anywhere in the buffer”,
which one < in an ordinary answer — if (a < b) — satisfied forever, and
the buffer was never cleared while it held. The whole answer then arrived in
a single frame at end of stream, indistinguishable from a hang to a streaming
client, and the scan over that ever-growing buffer made the cost quadratic in
the response length.
Two waits are longer than that. Text inside the reasoning channel is held until
its end marker — not a stall: it is reasoning, and it is delivered as
reasoning_content as it arrives. And once a marker that opens a tool-call
region appears, everything from it onward belongs to the format until it can
parse the region.
What is held, and for how long#
The rule is that a byte is buffered only while its destination field is
undecided — content or tool_calls — because an SSE frame cannot be taken
back. Four cases, and only two of them wait:
Where the byte is |
Held? |
Until |
|---|---|---|
Before any start marker |
no |
— |
Inside a region, before the name is legible |
yes |
the format can name a declared tool, usually the first 30–70 characters |
Inside a region, after that — the argument values |
yes |
the region closes |
After the region closed |
no |
— released as it arrives |
The name goes out as soon as it is legible and the request declared it
(tool_call_start), so a client learns which tool is being called at chunk
7–24 rather than after the whole payload.
Argument values wait, deliberately. vLLM and SGLang stream them as JSON
fragments for their JSON-shaped formats; a response cut off by max_tokens
then leaves the client accumulating an object it cannot parse. Buffering them
means five of the six formats here hand back valid JSON even for a truncated
call — the half-written value becomes a string — which streaming fragments
gives up. (Kimi-K2 is the exception: it passes the model’s bytes through, so a
truncated call already yields {"city": "Par. That is a separate defect.)
Both of those engines buffer for their tag-shaped formats too, K3 included, so
this is not a gap against them.
The region closes on the call’s own closer, not on the wrapper’s. Every
format’s grammar lists the call closer among the terminators of an argument
value — </function> for Qwen, </invoke> for DSML and MiniMax,
<|close|>call for K3 — so a model writing one inside a parameter ends the
parameter, and the literal can never hide in a value. The wrapper closer
(</tool_call>) can, which is why it serves only as a trigger to look and
never as the answer. Declared as CALL_SELF_CLOSERS; REGION_END_MARKERS
overrides it where a region is larger than one call, which is Kimi-K2’s
section.
Getting that wrong is expensive and was: while only Kimi declared a region end, everything a model wrote after its tool call waited for end of stream — 0 of 397 characters streamed on five of six formats, and “call a tool, then explain the result” is the ordinary agentic shape.
A region that never closes is still held to end of stream. An answer that
merely quotes its own opener is the case: measured on a GLM answer naming the
tag at character 29 and then explaining for 1234 more, 98% arrives in one
frame at EOS. Nothing is lost — the region is released verbatim once it turns
out not to be a call — and atom:stream_longest_silence_seconds reports the
wait while it happens. vLLM and SGLang both do the same here, and SGLang’s K3
detector drops the text rather than releasing it.
A probe that gave up on a region producing nothing after N bytes was written
for this and reverted. It rests on acceptance being monotone in how many bytes
have arrived, and that is false: MiniMax gates its in-progress test on the
first tag being in the declared schema, and DSML’s wrapper-less and
direct-JSON branches match no prefix at all — so real calls over N bytes were
delivered as raw text with finish_reason: stop on three of the six formats.
It was quadratic besides, because giving up re-fed bytes that immediately
reopened a region with a fresh budget: 1.19 ms to 18.2 s on a 250 KB answer,
in the request coroutine. Fixing the latency needs the format to say “this
can no longer become a call”, which is a different question from “does not
parse yet” and one no format answers today.
A second version was written and reverted for the same reason: bounded to the
256-byte peek window, asked once, and restricted to tool_choice: "none"
where nothing would be dispatched anyway. It measured clean on every shape the
corpus carries — and the corpus carries one form per format, because it is
generated from render_call. DSML’s direct-JSON body needs the whole object,
so a real call with a payload past the window is invisible to any head-sized
peek and would have gone out as text.
TestAGiveUpProbeStaysReverted now carries that shape, with a positive
control asserting DSML actually accepts it: the first draft of that test
invented both shapes from this paragraph rather than from the parser, DSML
accepted neither, and it therefore proved nothing.
“Opens a region” is asked of the format, not assumed of every marker it
declares. Kimi-K3 declares 16 and only two of them mean a tool call; the
rest are channel framing that wraps every answer it gives, including
<|open|>response<|sep|> at the very start. Treating those as a handover meant
a K3 response streamed nothing — measured, 324 of 324 characters in one frame
at EOS — which was the common path for that model rather than an edge case.
A start marker is not a promise, and that applies to the handover markers too.
An answer quoting one opens a region that then parses to no call, and every
format releases that region verbatim rather than deleting it. K3 was the one
without such a branch: it cut the answer at a quoted call opener and lost 62
characters with no event and finish_reason still stop.
The tool’s name does not wait for its arguments. A region is buffered until it closes, so on a 20 KB file write the client learned which tool was being called only after 5030 of 5040 tokens. Four of the six formats can recognise a call that has not finished arriving, and for those the name is sent as soon as the region reveals it — chunk 11–21 instead of 225–248.
The name is read out of parse_region itself, over the region so far and with
at_end=False. That is what makes the early name and the parsed call agree:
same function, same enumeration, the second run seeing a superset of the
first’s bytes. Every format used to answer this with a regex of its own, and
four of the five that had one disagreed with their own parse — Qwen’s peek
accepted </tool_call>, which closes the outer wrapper and leaves the
<function= block open; DeepSeek-V4’s skipped a self-closing
<invoke name="x"/> its parse returned first, putting three tool calls on the
/v1/messages wire for a response containing two. There is no separate peek
now, so there is nothing left to disagree.
at_end is the only difference between the two questions. With it — the
region has closed, or the stream has — a token cut off part-way through is all
there will ever be, so a prefix counts, which is what a call truncated by
max_tokens looks like. Without it a prefix means “not yet”, and accepting
one let a chunk boundary landing one character into <br> name a tool for
prose. Same bytes, announced at one chunk size and silent at another.
A name only goes out for a tool the request declared. Prose can name a real tool, so that alone is not enough — the follower test above is the other half. SGLang’s cursor parsers announce with neither check and will emit a call named after whatever follows the tag.
The read is bounded to Region.head and stops once that prefix has gone by
without a name. Running the format’s regex over the whole region on every
chunk is quadratic in the response — 3.0 → 9.8 → 36 → 137 ms across
2k/4k/8k/16k tokens, the shape marker_scanner exists to retire, one layer up.
Kimi-K2 and Kimi-K3 do not name a call whose arguments are still arriving: a
K2 entry is invisible until <|tool_call_end|> and a K3 call until
<|close|>call, so on a large payload the name arrives with the arguments.
A call short enough to fit inside Region.head is named early by all six,
because the whole call is in the window and parse_region sees a finished one.
Nothing declares which formats are which – the property suite measures both
facts independently (can the parse read a call in progress; did the name land
before the arguments on an 800-byte payload) and asserts they agree. A class
attribute used to stand in for this, outlived its only reader when the
give-up probe below was reverted, and took these two paragraphs false with it.
Arguments still wait for the region to close. SGLang streams those too, as
JSON fragments; a response cut short then leaves the client holding an
unterminated object. On /v1/chat/completions a name with no arguments is
harmless — clients accumulate by index and wait for finish_reason, which
keys on the arguments. On /v1/messages it is not: a content_block_start of
type tool_use carries "input": {} and is, on its own, a complete
zero-argument call, with no frame for “the name is known, the arguments are
coming”. So that endpoint opens the block when the arguments arrive, not when
the name does.
Which tool-call format a model uses is decided at startup, not from its
output. --tool-call-parser defaults to auto, which renders the model’s
chat template with a tools payload — the template’s own instructions for
calling one — and runs the _DETECT_ORDER cascade on the result. It reads a
Jinja template or a model-side Python encoder (<model>/encoding/encoding_*.py,
which is how DeepSeek-V4 ships its), and logs the format it chose. When nothing
is recognised it says so and tool calls are delivered as plain text. There is
no fallback to reading the output — not on either path, which is the point: the
non-streaming path used to run the cascade over the response whenever no format
had been resolved, so an answer that merely quoted another format’s section
token had everything from the token onward deleted with stream=false and
arrived whole with stream=true. A guess is silent, and it is also two
different answers to one request.
stream=false and stream=true deliver the same text, and not because a
test compares them: stream=false is read_whole, which is the streaming
engine over a single chunk. There is no second implementation to disagree
with. A format used to be read twice — once by a parse taking the whole
output, once by a process/flush state machine of its own — and both had to
decide where content ends, whether an unclosed tag is a call, what a region
that parses to nothing means, and which bytes are framing. Six formats, two
copies, four rules; three rounds of review found the copies disagreeing about
all four.
A format now declares only what is particular to it: the literals that must
not be split (START_MARKERS), which of those hand the stream over
(opens_region, the rest being framing the reader drops), and what one
region’s bytes mean (parse_region, returning the calls and the two offsets
that bracket its own markup). Everything else — reading ahead, releasing
content, the rule that a start marker is not a promise, stamping call indices,
and handing back the answer that follows the markup — is the engine’s, once.
The content comes back byte-for-byte except for markers the format declares.
Whitespace is not one — every format used to .strip(), which cost a
code-block answer its trailing newline on one path only. Text after a call
is not one either: five of the six deleted it, and the property suite now
holds every registered format to delivering it, at four chunk sizes.
The reasoning split is held to the same rule one stage earlier, and was not.
Two ways: </think> was matched only at position 0, so a model that answers,
opens a <think> block and answers again had it extracted when streamed and
handed over as literal tags with the chain of thought inside content when
not — and both halves were then .strip()ed, which is the trailing-newline
bug above, in the stage before it. A model writes </think>\n\nThe answer.;
stream=true delivers "\n\nThe answer." at every real chunk size and
stream=false delivered "The answer.". Measured over 12544 (dialect, shape,
chunking) comparisons, the two agreed byte-for-byte on 50% of them; they now
agree on all of them, and the property that says so is byte-exact rather than
word-level.
The streaming filter also stopped eating the newline after its end marker. It only ever saw what happened to be buffered when the marker arrived, so the same answer kept those bytes at one chunk size and lost them at another — there was no chunk-invariant behaviour on that whitespace for the other path to match even if it had wanted to.
Which reasoning dialect a model speaks is decided at startup too, from the
same evidence as the tool-call format and by the same kind of function
(resolve_dialect, on the chat-template source). It used to be decided twice
per response and differently each time: the non-streaming split tried each
registered dialect in order and took the first that matched, while the
streaming filter carried no dialect at all and closed the channel on the
union of every dialect’s end markers. A <think> model answering a question
about Kimi’s wire format therefore ended its chain of thought at the quoted
<|open|>response<|sep|> when streamed and at the real </think> when not —
24 characters of the answer filed as reasoning on one path, a raw </think>
shipped to the user on the other. ReasoningChannel now carries the dialect
and whether the output begins inside the channel, with one accessor per
delivery mode, so the two cannot be handed different answers. A template that
names no dialect falls back to the inline-<think> one, which is a no-op for
a model that never writes the tag.
Whether the output begins inside the channel is per request, not per model.
A request that switches reasoning off renders a prompt that does not open it,
and the model-level fact — a template that closes a block it never opens, as
DeepSeek-R1’s does — used to be OR-ed in regardless. On such a model an
ordinary answer to a request that had asked for no thinking came back entirely
as reasoning_content, with content empty.
tool_choice: "none" suppresses the call, not the answer. It used to be
enforced where the events are sent — twelve places across two endpoints —
while the parser went on consuming the region, so the model’s own words were
deleted and nothing took their place: 89 characters of a 95-character answer,
no event, finish_reason: stop. The rule now lives at the one place the
parser is asked, as suppress_calls. What that suppresses is dispatch: the
region is read exactly as a permitted call’s would be and only the calls are
dropped, so the answer around them survives and the model’s raw wire markup
does not reach the client. A reply that was nothing but a forbidden call
therefore has empty content — the model produced no answer.
Not by using no parser at all, which is where the first fix for that went.
Dropping the parser drops everything else a parser does, and a format whose
framing wraps every answer then leaks it: Kimi-K3’s Hello there. arrived
as <|open|>response<|sep|>Hello there.<|close|>response<|sep|><|end_of_msg|>
the moment a request said none. The format is still read; only dispatch is
suppressed. /v1/messages reads the field too, in Anthropic’s
{"type": "none"} spelling; it previously parsed it off the request and used
it nowhere, so a client that forbade tool calls got tool_use blocks and
stop_reason: tool_use anyway.
Forwarding it to the chat template is a separate step, and one that used to be
a 500. The handler passes template controls it cannot know the model reads —
response_format, tool_choice, thinking_effort, and whatever a client puts
in chat_template_kwargs. A Jinja template silently ignores a kwarg it does
not read; a model-shipped Python encoder raises TypeError. So on DeepSeek-V4
and Kimi-K3, which ship encoders instead of templates, any request carrying one
of those was an unhandled exception. The adapter now reads the encoder’s
signature once at startup and passes on only what it can take.
thinking is answered in the prompt, not in the response. On
/v1/messages, thinking: {"type": "disabled"} sets the chat template’s own
reasoning switch, so the model emits no chain of thought — there is then none
to separate, none to discard, and none for the tool parser to misread.
Separation stays unconditional, exactly as on /v1/chat/completions: the
tool parser is a second reader of the same text, so a chain of thought left in
it is one the tool parser will try to parse.
That ordering is the whole of it. Handling an unwanted chain of thought after
generating it fails three different ways — discarding it returns an empty
message for a reasoning model stopped at max_tokens; relabelling it as text
hands the client the thing it declined; and leaving it unseparated feeds it to
the tool parser, which is a second reader of the same text and read one model’s
musing about <function=NAME> as a call to a tool named NAME. SGLang answers
the same field the same way (apply_reasoning_enabled), and vLLM gets it
structurally by having no such field: its reasoning parser runs unconditionally
and include_reasoning only suppresses the result after the split.
Which kwarg carries the switch is resolved at startup by rendering the template
twice and comparing, because a template silently ignores a kwarg it does not
read. On this box: Qwen3/Qwen3.5 enable_thinking, Kimi-K3 thinking,
MiniMax-M3 thinking_mode="disabled", DeepSeek-V4 thinking_mode="chat".
A model whose template has no switch is named in the startup log. Its reasoning
cannot be prevented, so thinking: {"type": "disabled"} is answered the only
way left: the text is still separated, and the thinking blocks are withheld.
That is the one downstream suppression there is, and it is reached only when
the prompt could not carry the answer — without it an explicit opt-out was
honoured at neither layer. A response that was nothing but reasoning then
ends on an empty text block, which is the honest reply to “do not think”.
Two details that bite: {"type": "disabled"} is a non-empty object, so testing
the field for truthiness read the standard off-switch as on; and an absent
thinking leaves the model’s own default alone rather than switching reasoning
off, at both layers or neither, so an existing caller’s answers do not change.
A stalled response is visible while it is stalled — on the OpenAI server.
Every SSE frame from openai_server leaves through _client_stream, which
times the gap before each one and registers it,
and atom:stream_longest_silence_seconds reports the age of the oldest gap
in flight. Zero when every stream has just been served; non-zero and growing is
a response whose client is receiving nothing. A gap longer than 30 seconds also
logs a line naming the request — the gauge cannot see a stall that has already
recovered by scrape time. Neither costs a timer: asyncio.wait_for measured
1.38 us per frame per stream against 0.07 us for a timestamp and a dict entry.
This exists because the symptom that started this work was ten minutes of
silence with every metric looking healthy.
The atomesh standalone entrypoint has none of it. Its frames leave through
ChatCompletionStreamState.drain / CompletionStreamState.drain, polled by
the Rust router, which builds no FrameWait; and AtomMetricsExporter is
constructed only by openai_server, so that deployment exposes no /metrics
route at all. A stalled atomesh stream is therefore invisible rather than
reported as zero.
Measured at the frame and not at StreamOutputCollector.get, which is where it
started and which cannot see the thing it was built for. The collector is where
a stream waits for the engine, but the reasoning read-ahead and the tool-call
read-ahead sit between it and the socket, and while either withholds, the
collector wakes on every token. Measured: an answer quoting a tool marker fed
126 tokens and sent the client 6 frames, and the gauge read zero. At the frame
it reads the silence.
The wait for the first frame is still excluded, and moving out did not
change that — a claim this paragraph made and did not hold. Every response
generator awaits the collector before yielding anything, so that wait is
admission, queueing and prefill: timing it put 0.2 s on the gauge for a
request 200 ms into a queue with no token yet produced, which is
atom:requests_waiting under another name, and past the threshold would log a
line per admitted request blaming the read-ahead.
One consequence matters when reading benchmark output. ITL is sampled once per
received SSE chunk (backend_request_func.py, benchmark_serving.py), so
merging N tokens into one chunk removes N-1 samples and stretches the gaps that
remain: every ITL statistic - mean, median and p99 alike - inflates by roughly
the merge factor, without any token being delivered later. Measured on
Qwen3.5-27B-FP8 tp4 at concurrency 2048, mean ITL read 191.8 ms against a TPOT
of 126.6 ms, while the same workload with merging disabled read 122.9 ms against
a TPOT of 123.3 ms.
Compare TPOT, not ITL, whenever merging is active. It is the only
token-normalized latency in the report (latency - ttft over output_len - 1),
so it stays honest at any merge factor. The ratio ITL/TPOT is itself the useful
number: it is the merge factor, and a value near 1.0 means the frontend is
keeping up and nothing ever merged.
Server startup#
python -m atom.entrypoints.openai_server \
--model <model_name_or_path> \
--kv_cache_dtype fp8 \
--host 0.0.0.0 \
--server-port 8000
Server-specific CLI arguments:
Argument |
Default |
Description |
|---|---|---|
|
|
Bind address |
|
|
HTTP port (note: |
|
|
Seconds an idle keep-alive connection is held. Pooling clients hold their end longer (aiohttp defaults to 15s), so a caller that pauses for longer than this reuses a socket the server already closed and has to re-send. Raise it past the caller’s idle window to avoid that |
|
off |
Stop uvicorn logging a line per HTTP request. It copies a |
|
|
Tool-call wire format. |
All EngineArgs arguments are also accepted (see Section 7 for the full list).
Example: curl#
# Non-streaming chat completion
curl http://localhost:8000/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "deepseek-ai/DeepSeek-R1",
"messages": [{"role": "user", "content": "Hello!"}],
"max_tokens": 128
}'
# Streaming text completion
curl http://localhost:8000/v1/completions \
-H "Content-Type: application/json" \
-d '{
"prompt": "The capital of France is",
"max_tokens": 64,
"stream": true
}'
Programmatic API (LLMEngine)#
The LLMEngine class in atom/model_engine/llm_engine.py provides a
Python-native interface for inference without running an HTTP server.
Initialization#
from atom import LLMEngine, SamplingParams
engine = LLMEngine(model="deepseek-ai/DeepSeek-R1", kv_cache_dtype="fp8",
tensor_parallel_size=8)
LLMEngine.__init__(model, **kwargs) accepts all Config field names as
keyword arguments (e.g. tensor_parallel_size, kv_cache_dtype,
max_model_len, data_parallel_size, gpu_memory_utilization).
SamplingParams#
Defined in atom/sampling_params.py:
@dataclass
class SamplingParams:
temperature: float = 1.0
max_tokens: int = 64
ignore_eos: bool = False
stop_strings: Optional[list[str]] = None
Core methods#
Method |
Signature |
Description |
|---|---|---|
|
|
Synchronous batch generation; blocks until all prompts complete |
|
|
Submit requests for asynchronous processing |
|
|
Retrieve completed sequences |
|
|
Check whether all pending requests have completed |
|
|
Start torch profiler on all workers |
|
|
Stop torch profiler and write traces |
|
|
Print speculative decoding acceptance statistics |
Synchronous generation example#
from atom import LLMEngine, SamplingParams
engine = LLMEngine(model="meta-llama/Meta-Llama-3-8B", kv_cache_dtype="fp8")
params = SamplingParams(temperature=0.6, max_tokens=256)
outputs = engine.generate(["Explain quantum computing in simple terms."], params)
for out in outputs:
print(out["text"])
Each output dictionary contains: text, token_ids, latency,
finish_reason, num_tokens_input, num_tokens_output, ttft, and tpot.
Asynchronous / streaming usage#
engine.add_request(
prompt_or_tokens_list=["Hello world", "How are you?"],
sampling_params_list=SamplingParams(temperature=0.8, max_tokens=128),
stream_callback=my_callback, # called per-token with RequestOutput
)
while not engine.is_finished():
completed = engine.step()
# process completed sequences
Simple inference#
The atom/examples/simple_inference.py script provides a quick way to validate
model loading and generation.
Usage#
python -m atom.examples.simple_inference \
--model meta-llama/Meta-Llama-3-8B \
--kv_cache_dtype fp8 \
--temperature 0.6
What it does#
Parses all
EngineArgsplus--temperature(default0.6).Creates an
LLMEngineviaEngineArgs.from_cli_args(args).create_engine().Applies the model’s chat template to four built-in prompts (English and Chinese) with
enable_thinking=True.Runs a warmup generation, then generates completions for the batch.
Calls
llm.print_mtp_statistics()to report speculative decoding stats (if MTP is enabled).
Benchmarking#
ATOM ships a comprehensive online serving benchmark in
atom/benchmarks/benchmark_serving.py (adapted from vLLM’s benchmarking
tooling).
Metrics#
The BenchmarkMetrics dataclass tracks:
Metric |
Abbreviation |
Description |
|---|---|---|
Time to First Token |
TTFT |
Latency from request submission to the first generated token |
Time per Output Token |
TPOT |
Average latency per output token (excluding the first) |
Inter-Token Latency |
ITL |
Latency between successive output tokens |
End-to-End Latency |
E2EL |
Total latency from request send to full response receipt |
Request Throughput |
– |
Completed requests per second |
Output Token Throughput |
– |
Generated tokens per second |
Total Token Throughput |
– |
(input + output) tokens per second |
Request Goodput |
– |
Requests per second meeting SLO targets |
Concurrency |
– |
Average in-flight requests (sum of per-request end-to-end latency / benchmark duration) |
Accept Length |
– |
Speculative decoding only: mean tokens per model forward (1 + accepted draft tokens), from |
Acceptance Rate |
– |
Speculative decoding only: fraction of drafted tokens accepted (accepted / drafted), from |
For each latency metric, mean, median, standard deviation, and configurable percentiles (default: P99) are reported.
Key CLI arguments#
Argument |
Default |
Description |
|---|---|---|
|
|
Backend type. Choices: |
|
(required) |
Model name or path |
|
|
Server base URL (e.g. |
|
|
Server host (used when |
|
|
Server port (used when |
|
|
API endpoint path |
|
|
Dataset type: |
|
|
Path to dataset file or HuggingFace dataset ID |
|
|
Number of prompts to benchmark |
|
|
Requests per second ( |
|
|
Burstiness factor (1.0 = Poisson process) |
|
|
Maximum concurrent requests |
|
|
Ignore EOS token in generation |
|
|
Save results to JSON |
|
|
Directory for result JSON files |
|
|
Custom filename for results |
|
|
Comma-separated metrics to report percentiles for |
|
|
Comma-separated percentile values (e.g. |
|
|
SLO targets as |
|
|
Enable torch profiler during the benchmark run |
|
|
Custom tokenizer name or path |
|
|
Random seed |
Random dataset options:
Argument |
Default |
Description |
|---|---|---|
|
|
Input token length |
|
|
Output token length |
|
|
Length variation ratio |
|
|
Fixed prefix token length |
|
|
Apply chat template to random prompts |
Backend request functions#
Defined in atom/benchmarks/backend_request_func.py:
Backend Key |
Function |
Protocol |
|---|---|---|
|
|
OpenAI Completions API (streaming) |
|
|
OpenAI Completions API (streaming) |
|
|
OpenAI Chat Completions API (streaming) |
|
|
TGI |
|
|
TRT-LLM |
|
|
DeepSpeed-MII |
|
|
OpenAI Completions API |
|
|
OpenAI Completions API |
|
|
OpenAI Completions API |
Each function uses RequestFuncInput and returns a RequestFuncOutput with
timing data (ttft, itl, latency, tpot).
Full benchmark example#
# 1. Start the server
python -m atom.entrypoints.openai_server \
--kv_cache_dtype fp8 -tp 8 --model deepseek-ai/DeepSeek-R1
# 2. Run benchmark
MODEL=deepseek-ai/DeepSeek-R1
ISL=1024
OSL=1024
CONC=128
PORT=8000
RESULT_FILENAME=Deepseek-R1-result
python -m atom.benchmarks.benchmark_serving \
--model=$MODEL --backend=vllm --base-url=http://localhost:$PORT \
--dataset-name=random \
--random-input-len=$ISL --random-output-len=$OSL \
--random-range-ratio 0.8 \
--num-prompts=$(( $CONC * 10 )) \
--max-concurrency=$CONC \
--request-rate=inf --ignore-eos \
--save-result --percentile-metrics="ttft,tpot,itl,e2el" \
--result-dir=./ --result-filename=$RESULT_FILENAME.json
Profiling#
ATOM supports PyTorch profiling via environment variables, HTTP endpoints, and the programmatic API.
Configuration#
Mechanism |
Description |
|---|---|
|
CLI arg to set the trace output directory |
|
Sets the default |
|
Enables detailed profiling: |
|
Overrides the |
|
Appends attention FLOP aggregates ( |
When a profiler directory is configured, each worker saves traces to a rank-specific subdirectory:
Multi-GPU with DP:
{profiler_dir}/dp{dp_rank}_tp{rank}/Single-GPU / TP-only:
{profiler_dir}/rank_{rank}/
Traces are saved in gzip-compressed TensorBoard format and can be viewed with
tensorboard --logdir <profiler_dir> or Chrome’s chrome://tracing.
Online profiling (HTTP)#
While the server is running, start and stop profiling with HTTP requests:
# Start profiling
curl -s -S -X POST http://127.0.0.1:8000/start_profile
# ... run your workload ...
# Stop profiling and flush traces
curl -s -S -X POST http://127.0.0.1:8000/stop_profile
The server must be started with --torch-profiler-dir or with
ATOM_TORCH_PROFILER_DIR set for these endpoints to produce traces.
For large traces, set ATOM_PROFILER_TIMEOUT higher before starting the server.
Programmatic profiling#
engine = LLMEngine(model="Qwen/Qwen3-0.6B", torch_profiler_dir="./traces")
engine.start_profile()
outputs = engine.generate(prompts, sampling_params)
engine.stop_profile()
# Traces written to ./traces/rank_0/
Offline profiling script#
atom/examples/profile_offline.py provides a self-contained offline profiling
workflow:
python -m atom.examples.profile_offline \
--model Qwen/Qwen3-0.6B \
--kv_cache_dtype fp8 \
--torch-profiler-dir ./profiler_traces \
--input-length 128 \
--output-length 32 \
--bs 4
Script-specific arguments:
Argument |
Default |
Description |
|---|---|---|
|
|
Approximate input prompt length in tokens |
|
|
Output generation length in tokens |
|
|
Batch size (number of parallel requests) |
|
|
Use random token input instead of predefined text |
If --torch-profiler-dir is not specified, the script defaults to
./profiler_traces.
Profiling during benchmarks#
The benchmark tool can trigger profiling automatically via --profile:
python -m atom.benchmarks.benchmark_serving \
--model <model> --backend vllm \
--base-url http://localhost:8000 \
--dataset-name random --num-prompts 100 \
--profile
This sends POST /start_profile before the benchmark and
POST /stop_profile after completion.
CUDA-graph capture traces#
During CUDA-graph capture (server bring-up), ATOM can emit one trace file per
captured batch size instead of a single combined blob. This makes each graph’s
capture cost easy to inspect in isolation and keeps individual trace files
small. Capture-trace profiling is gated on --mark-trace (with
--torch-profiler-dir/ATOM_TORCH_PROFILER_DIR set).
Each file covers one full iteration of the capture loop: the warmup forward
followed by the graph capture itself. Both are needed — inside
torch.cuda.graph(...) the stream is in capture mode, so kernel launches are
recorded as graph nodes rather than dispatched, and a trace of that region
alone has an empty GPU track. The warmup forward is where the kernels actually
run.
The traces are written to:
{profiler_dir}/capture_traces/bs_<bs>_q_<max_q_len>_rank<rank>.json.gz
where <bs> is the captured batch size, <max_q_len> the query-length bucket
(1 without speculative decoding, mtp_k + 1 with a drafter, and one file per
bucket when DSpark expands them — see
Speculative decoding), and <rank> the worker
rank. Each file is a gzip-compressed Chrome trace viewable with
chrome://tracing or TensorBoard.
Like the run-phase profiler, these traces carry record_shapes, with_stack,
and profile_memory only when ATOM_PROFILER_MORE=1. Leave it unset unless you
need the shapes or Python stacks — stack capture runs on every rank and
noticeably stretches server bring-up.
To additionally annotate the run-phase prefill[]/decode[] labels with the
attention FLOP aggregates used for roofline analysis, set
ATOM_ENABLE_DETAILED_ANNOTATION=1 (see Configuration). The added
fields are sqsq (Σ N_Q²), sqsk (Σ N_Q·N_KV), and sk (Σ N_KV), summed over
every request in the forward. These are attention-quadratic terms only — a full
roofline still requires GEMM FLOPs and bytes moved.
Speculative decoding (MTP)#
ATOM supports Multi-Token Prediction (MTP) for DeepSeek models using the Eagle-style speculative decoding framework.
Architecture#
EagleProposer (
atom/spec_decode/eagle.py): Loads and runs the draft (MTP) model to propose speculative tokens. Supports theDeepSeekMTPModelarchitecture viaDeepSeekMTP.RejectionSampler (
atom/model_ops/rejection_sampler.py): Implements greedy rejection sampling with a Triton kernel. Compares draft token IDs against target model argmax and accepts matching prefixes; appends a bonus token if all drafts are accepted.
Configuration#
Enable MTP via CLI arguments:
python -m atom.entrypoints.openai_server \
--model deepseek-ai/DeepSeek-R1 \
--kv_cache_dtype fp8 -tp 8 \
--method mtp \
--num-speculative-tokens 1
Argument |
Default |
Description |
|---|---|---|
|
|
Speculative method: |
|
|
Number of draft tokens per iteration (draft model runs this many autoregressive steps) |
|
|
Path or HF repo of the speculative draft model. Required for |
|
|
Benchmark-only: force a mean acceptance length in |
|
|
The same knob as a rate in |
MTP statistics#
ATOM tracks acceptance statistics at runtime:
total_draft_tokens: Total number of draft tokens proposed
total_accepted_tokens: Number of draft tokens accepted by rejection sampling
acceptance_rate: Ratio of accepted to draft tokens
Statistics are logged every 1000 draft tokens and can be printed on demand:
engine.print_mtp_statistics()
Example output:
MTP Statistics:
Total draft tokens: 5000
Accepted tokens: 4250
Acceptance rate: 85.00%
How rejection sampling works#
The draft model generates
num_speculative_tokenstoken predictions autoregressively using argmax.The target model verifies all draft tokens in a single forward pass.
The
rejection_greedy_sample_kernel(Triton) compares each draft token against the target model’s argmax:If they match, the token is accepted.
On the first mismatch, the target model’s token replaces it and all subsequent draft tokens are discarded.
If all draft tokens match, a bonus token from the target model is appended.
Forced acceptance length#
Speculative throughput is dominated by how many tokens each target forward
emits, so a run cannot be compared against another engine unless both accept at
the same rate. --spec-decode-acceptance-length pins that number: the sampler
stops comparing draft against target and instead accepts draft tokens with a
fixed per-position probability, hitting the requested mean acceptance length.
It exists to benchmark the serving system while a draft head is still training,
and to replay a published acceptance-length figure such as an
InferenceX golden AL.
python -m atom.entrypoints.openai_server \
--model /models/Kimi-K3 \
--draft-model /models/Kimi-K3-DSpark \
--method dspark \
--num-speculative-tokens 7 \
--spec-decode-acceptance-length 3.78
Acceptance length counts the target’s own guaranteed token, matching vLLM’s
synthetic_acceptance_length and SGLang’s SGLANG_SIMULATE_ACC_LEN, so a
published figure goes in unchanged. The budget is spent on the earliest
positions — length 3.78 over 7 draft slots accepts 2 tokens always and a 3rd
with probability 0.78 — which is the minimum-variance schedule vLLM and
SGLang also use, so the accepted-length distribution matches and not just its
mean. Read the realized value back from average_tokens_per_forward on
/debug/mtp_stats (or the atom:mtp_average_tokens_per_forward metric).
Two caveats:
Generated text is meaningless, because tokens are accepted without agreeing with the target. Never run an accuracy evaluation with this enabled.
It cannot be combined with the DSpark confidence scheduler (
--dspark-config '{"confidence_schedule": true}'), which picks each request’s verify length at runtime; a short one silently caps acceptance below the requested length, so the combination is rejected at startup.
The full reference — the resolved schedule, the rate-based spelling, and how to
replay a golden AL curve — is in
forced_acceptance_length.md.
Deployment examples#
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
Docker deployment#
# Pull the ROCm PyTorch image
docker pull rocm/pytorch:rocm7.0.2_ubuntu24.04_py3.12_pytorch_release_2.8.0
# Launch container
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
# Inside the container
pip install amd-aiter
git clone https://github.com/ROCm/ATOM.git && cd ATOM && pip install .
# Start serving
python -m atom.entrypoints.openai_server \
--model deepseek-ai/DeepSeek-R1 \
--kv_cache_dtype fp8 -tp 8
Engine CLI arguments (EngineArgs)#
These arguments are available for all entrypoints (server, examples, and any
script using EngineArgs.add_cli_args):
Argument |
Default |
Description |
|---|---|---|
|
|
Model name or path |
|
|
Trust remote code from HuggingFace |
|
|
Tensor parallel size |
|
|
Data parallel size |
|
|
Disable CUDA graph capture; use eager execution |
|
|
Enable prefix caching |
|
|
Emit the periodic engine-status line (throughput, running/waiting reqs, KV usage, prefix-cache hit rate) |
|
|
Seconds between engine-status lines |
|
|
Internal engine communication port |
|
|
KV cache dtype: |
|
|
KV cache block size |
|
|
Maximum context length (defaults to HF config) |
|
|
Maximum tokens per batch |
|
|
Maximum sequences per batch |
|
|
GPU memory utilization (0.0 to 1.0) |
|
|
Delay factor before scheduling next prompt |
|
|
Batch sizes for CUDA graph capture |
|
|
Compilation level (0-3); 3 = torch.compile |
|
|
Dummy weights (no checkpoint read). Bare flag / |
|
|
Enable expert parallelism for MoE |
|
|
Enable data-parallel attention |
|
|
Directory for torch profiler traces |
|
|
Speculative decoding method ( |
|
|
Number of speculative tokens per step |
Accuracy validation#
ATOM supports accuracy validation through the lm-eval framework via the OpenAI-compatible API.
Setup#
pip install lm-eval[api]
Run evaluation#
Start an ATOM server, then run lm-eval against it:
# Start server
python -m atom.entrypoints.openai_server \
--model meta-llama/Meta-Llama-3-8B \
--kv_cache_dtype fp8
# 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
Any lm-eval task can be used. The local-completions model type sends
requests to the /v1/completions endpoint, making it compatible with the ATOM
server without modification.
Source files#
File |
Description |
|---|---|
|
OpenAI-compatible API server (FastAPI + Uvicorn) |
|
|
|
SSE frame encoding ( |
|
|
|
Splits the reasoning channel from the answer; |
|
The dialects, and |
|
Kimi-K3’s channel tokens, split by owner: what the reasoning stage strips and what only the tool parser may |
|
The one reader: the engine both delivery modes run through |
|
Renders the chat template, and the two startup probes of it: |
|
Which format a model emits, resolved once at startup from its chat template |
|
Per-format tool-call syntax; each format declares its markers and a |
|
|
|
|
|
|
|
Simple batch inference example |
|
Offline profiling tool |
|
Online serving benchmark ( |
|
Async HTTP request functions for each backend ( |
|
|
|
|
|
|
|
|
|
|