Skip to content

(recipe) Kimi-K3 AgentX: FP8 prefill attention, PrefillDelayer from C16, capture from bs=1 #2382

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zhuyuhua-v merged 2 commits into from
Sep 24, 2026
Merged

(recipe) Kimi-K3 AgentX: FP8 prefill attention, PrefillDelayer from C16, capture from bs=1 #2382

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3c1880a
(recipe) Kimi-K3 AgentX: FP8 prefill attention, PrefillDelayer from C…
gbyu-amd Sep 24, 2026
41dc7b8
(recipe) Kimi-K3 AgentX: keep the change to env vars and launch args
gbyu-amd Sep 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 13 additions & 2 deletions recipes/Agentic-Kimi-K3.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ with ATOM on 8×AMD MI355X GPUs. The validated configuration uses:
- FP8 KV cache
- native GPU prefix caching; LMCache CPU offload from CONC 8 up
- DSpark speculative decoding on CONC ≤ 16, with synthetic acceptance
- FlyDSL FP8 prefill attention on every band; the PrefillDelayer from CONC 16 up
- the SemiAnalysis Weka AgentX workload

The workload uses the AIPerf scenario `inferencex-agentx-mvp` and public
Expand Down Expand Up @@ -42,14 +43,16 @@ changes.
| Prefix cache | Enabled (required for AgentX multi-turn prefix hits) |
| Speculative decoding | DSpark (`--method dspark`), not native MTP |
| Synthetic acceptance | `--spec-decode-acceptance-length` (performance-only; see below) |
| Prefill attention | FlyDSL FP8 (`ATOM_USE_FLYDSL_FP8_PREFILL_ATTN=1`), every band |
| PrefillDelayer | CONC ≥ 16: `ATOM_PREFILL_DECODE_INTERVAL=4` (4 decode passes after each prefill), `ATOM_PREFILL_DELAYER_MAX_QUEUE_MS=5000` |
| Profiling duration | 3,600 seconds |
| Warmup | 10 additional one-token requests per lane |
| AIPerf | `0.12.0` (`agentx-v1.0.4`) |

### Per-concurrency server table

`graph_max = 2 * CONC * (1 + spec)`. Capture sizes are the dense range
`[2, 3, …, graph_max]`. `max-num-seqs` is **not** always `2 * CONC`.
`[1, 2, …, graph_max]`. `max-num-seqs` is **not** always `2 * CONC`.

| CONC | DCP | spec | AL | LMCache | ReplaySSM | `max-num-seqs` | batched tokens | GPU util | `graph_max` |
|---:|---:|---:|---:|---|---:|---:|---:|---:|---:|
Expand All @@ -74,6 +77,9 @@ C8–C48 use 128 GiB LMCache; C56–C80 use **192 GiB**. LMCache CPU size is
`AITER_REUSE_IDENTICAL_COMM_GROUPS=1` on **C56/C64/C72/C80**, and `0` on every
other band (C1…C48).

`ATOM_PREFILL_DECODE_INTERVAL=4` and `ATOM_PREFILL_DELAYER_MAX_QUEUE_MS=5000` on
**C16 and up**; unset on C1…C14.

## 0. Container prerequisites

Two things about the container decide whether the numbers above are
Expand Down Expand Up @@ -131,6 +137,7 @@ export AITER_SITUV2_A4W4=1
export AITER_FLYDSL_STAGE2_FP8=1
export ATOM_STATE_CHECKPOINT_DEMAND=0
export ATOM_GDN_SSM_DTYPE="${ATOM_GDN_SSM_DTYPE:-fp16}"
export ATOM_USE_FLYDSL_FP8_PREFILL_ATTN=1
# Set explicitly to 1 for reproducibility; this is also the current default.
export ATOM_USE_FLYDSL_GATHER_KV_B_PROJ=1
export PYTHONNOUSERSITE=1
Expand Down Expand Up @@ -190,14 +197,18 @@ if [[ "${CONC}" -ge 56 ]]; then
fi
AITER_REUSE_IDENTICAL_COMM_GROUPS="${AITER_REUSE_IDENTICAL_COMM_GROUPS:-0}"
export AITER_REUSE_IDENTICAL_COMM_GROUPS
if [[ "${CONC}" -ge 16 ]]; then
export ATOM_PREFILL_DECODE_INTERVAL="${ATOM_PREFILL_DECODE_INTERVAL:-4}"
export ATOM_PREFILL_DELAYER_MAX_QUEUE_MS="${ATOM_PREFILL_DELAYER_MAX_QUEUE_MS:-5000}"
fi
export ATOM_ENABLE_REPLAYSSM
SPEC_TOKENS_FOR_GRAPH=0
if [[ "${NUM_SPECULATIVE_TOKENS}" != "0" ]]; then
SPEC_TOKENS_FOR_GRAPH="${NUM_SPECULATIVE_TOKENS}"
fi
CUDAGRAPH_MAX_NUM_SEQS="${CUDAGRAPH_MAX_NUM_SEQS:-$((2 * CONC))}"
GRAPH_MAX=$((CUDAGRAPH_MAX_NUM_SEQS * (1 + SPEC_TOKENS_FOR_GRAPH)))
CUDAGRAPH_CAPTURE_SIZES="[$(seq -s, 2 "${GRAPH_MAX}")]"
CUDAGRAPH_CAPTURE_SIZES="[$(seq -s, 1 "${GRAPH_MAX}")]"

ATOM_CMD=(
python3 -m atom.entrypoints.openai_server
Expand Down