Skip to content

fix(ci-mesh): K3 AgentX cells — FP8 prefill attention and capture from bs=1 #2417

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
gbyu-amd merged 2 commits into from
Sep 28, 2026
Merged

fix(ci-mesh): K3 AgentX cells — FP8 prefill attention and capture from bs=1 #2417

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
434a4d7
fix(ci-mesh): align Kimi-K3 AgentX cells with the recipe (#2382)
ganyi1996ppo Sep 28, 2026
7cd958f
fix(ci-mesh): drop PrefillDelayer env from K3 P/D cells
ganyi1996ppo Sep 28, 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
4 changes: 3 additions & 1 deletion .github/benchmark/models_atomesh.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -479,6 +479,8 @@ models:
# move the published numbers.
ATOM_USE_FLYDSL_GATHER_KV_B_PROJ: "1"
ATOM_STATE_CHECKPOINT_DEMAND: "0"
# AgentX runs FlyDSL FP8 prefill attention on every band.
ATOM_USE_FLYDSL_FP8_PREFILL_ATTN: "1"
# AgentX band default: ReplaySSM off on CONC 1/2/4, on for the mid band
# (CONC 8-16), off again from CONC 32. Cases that declare their own env
# block must restate the key, because a YAML merge key loses to an
Expand Down Expand Up @@ -526,7 +528,7 @@ models:
workers: 1
tp: 8
# AgentX pins the compile path and captures the dense query-token
# range 2..2*CONC*(1+spec) that `auto` derives per concurrency.
# range 1..2*CONC*(1+spec) that `auto` derives per concurrency.
cudagraph: auto
cudagraph_mode: FULL
compilation_level: 3
Expand Down
10 changes: 5 additions & 5 deletions .github/scripts/atomesh/pd_server_atom.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -407,7 +407,7 @@ if [[ "${DECODE_ENABLE_DP}" == "true" ]]; then
fi

# AgentX captures every query-token count the engine can produce, i.e. the dense
# range [2, graph_max] with graph_max = seqs * (1 + spec_tokens), where seqs
# range [1, graph_max] with graph_max = seqs * (1 + spec_tokens), where seqs
# defaults to 2 * CONC. Concurrencies whose in-flight window is wider than
# 2 * CONC pin seqs explicitly via cudagraph_max_num_seqs.
auto_cudagraph_capture_sizes() {
Expand All @@ -422,11 +422,11 @@ auto_cudagraph_capture_sizes() {
seqs=$(( 2 * conc ))
fi
graph_max=$(( seqs * (1 + spec) ))
if (( graph_max < 2 )); then
graph_max=2
if (( graph_max < 1 )); then
graph_max=1
fi
echo "[${role}] cudagraph auto range 2..${graph_max} (seqs=${seqs} spec=${spec})" >&2
echo "[$(seq -s, 2 "${graph_max}")]"
echo "[${role}] cudagraph auto range 1..${graph_max} (seqs=${seqs} spec=${spec})" >&2
echo "[$(seq -s, 1 "${graph_max}")]"
}

build_cudagraph_args() {
Expand Down