Skip to content

fix: replay one prefill token-graph bucket across DP ranks; prune CI - #2696

Merged
valarLip merged 2 commits into
mainfrom
fix/dp-prefill-bucket-and-ci-cleanup
Oct 24, 2026
Merged

valarLip merged 2 commits into
mainfrom
fix/dp-prefill-bucket-and-ci-cleanup

Conversation

@valarLip

Copy link
Copy Markdown
Collaborator

DP attention: one prefill token-graph bucket for the whole group

Under DP attention, a rank with no prefill of its own runs a dummy prefill while a peer prefills. Before this change, each prefilling rank padded to its own bucket, and the dummy sat on the smallest one. A recording's collectives are sized by its bucket, so a dummy replaying 256 beside a peer replaying 2048 hung the group.

ForwardMode._pad_to_token_graph now picks the bucket that holds the group's largest token count and raises every prefilling rank to it, dummy prefill included. The function no longer takes this rank's scheduled_tokens. test_every_prefilling_rank_replays_the_groups_bucket covers the dummy case.

Docs updated to match:

  • the dummy_execution docstring;
  • ATOM_TOKEN_GRAPH_SIZES in environment_variables.md;
  • the token-graph section of compilation_cudagraph_guide.md;
  • the dummy-batch section of distributed_guide.md, which named a _execute_dummy_prefill that no longer exists.

Known limitation (unchanged by this PR): without a prefill delayer (ATOM_ENABLE_PREFILL_DELAYER=0, or RapidServe), no rank learns before scheduling that a peer will prefill, so a decoding rank keeps decoding. It then runs eagerly beside a peer replaying a prefill token graph, and the MoE collectives' sizes can disagree. Making decode rows ride the prefill path when any rank prefills would remove the dummy prefill and this case together; that is follow-up work.

CI

  • Fix main: tests/test_otel_scheduler.py::test_engine_initializes_tracing_before_ready compared the whole EngineCore startup call list. refactor: one graph mechanism per region, settled DP step shapes, cheaper V4.1 decode prep #2688 added a prefill_token_graphs call, and feat(observability): add opt-in standalone OpenTelemetry tracing #2644 merged on an older base, which turned main red (job). The test now asserts only allocate_kv_cache → initialize → freeze → ready.
  • Accuracy: remove DeepSeek-V4-Pro DEP RCCL, DeepSeek-V4-Pro MTP, MiniMax-M2.7, GLM-5-FP8, GLM-5.2-FP8, GLM-5.2-MXFP4 and GLM-5.2-MXFP4 MTP. The atomesh accuracy whitelist drops DeepSeek-V4-Pro MTP too.
  • Benchmark: remove GLM-5.2-FP8 and GLM-5.2-MXFP4, along with their atom-benchmark.yaml dispatch inputs, and DeepSeek-V4-Pro's MTP3 / DPA MTP3 variants. test_benchmark_catalog.py no longer names them.
  • Agentic nightly: add DeepSeek-V4.1-Flash DSpark5 DTP2 (-tp 2 --enable-dp-attention) and DEP2 (-tp 2 --enable-dp-attention --enable-expert-parallel), both at concurrency 128 on GPUs 0,1.

Test plan

  • 18 affected test files (forward mode, scheduler, DP metadata, DSpark, runner padding, benchmark / agentic catalogs, atomesh workflow, OTel): 672 passed, 45 skipped
  • The old OTel assertion fails on this base; the new one passes
  • validate_catalog.py; the agentic builder expands both new variants to the expected server args
  • Agentic nightly DTP2 / DEP2 on GPU

DP attention: every prefilling rank, an idle rank's dummy prefill included,
now replays the bucket holding the group's largest token count. A
recording's collectives are sized by its bucket, so ranks replaying
different buckets (a dummy prefill on the smallest one beside a peer on a
larger one) hung the group.

CI:
- test_otel_scheduler: assert only the order around tracing.initialize
  instead of the whole startup RPC list, which #2688 extended with
  prefill_token_graphs and turned main red after #2644 merged.
- Drop accuracy entries DeepSeek-V4-Pro DEP RCCL, DeepSeek-V4-Pro MTP,
  MiniMax-M2.7 and every GLM before 5.3; drop GLM-5.2 and the V4-Pro MTP3 /
  DPA MTP3 benchmark variants with their dispatch inputs.
- Agentic nightly: add DeepSeek-V4.1-Flash DSpark5 DTP2 and DEP2 at
  concurrency 128.
@github-actions

Copy link
Copy Markdown
Contributor

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

  • ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

  • ✅ Run after the PR is approved and Pre Checkin passes
  • ✅ Run immediately when an approval review is submitted
  • ✅ Can be requested before approval with labels
Label Tests
ci:full Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom Run native ATOM model accuracy tests
ci:vllm Run ATOM vLLM OOT model accuracy tests
ci:sglang Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 2696 --add-label <label>

@valarLip
valarLip merged commit f554c19 into main Oct 24, 2026
26 of 27 checks passed
@valarLip
valarLip deleted the fix/dp-prefill-bucket-and-ci-cleanup branch October 24, 2026 17:26
yhl-amd added a commit that referenced this pull request Oct 24, 2026
#2696 makes every prefilling DP rank replay one token-graph bucket, but a
step can still pair a prefill with a decode: without the prefill delayer
(ATOM_ENABLE_PREFILL_DELAYER=0, or RapidServe) no rank learns before
scheduling that a peer will prefill, so an idle rank runs a dummy decode
and a busy one keeps decoding. The prefilling rank then replays its
bucket's recording while the decoding rank runs eagerly, and the MoE's
DP exchange differs on the two sides: a recorded all_gather at the
bucket's height against an eager variable-length gather. With the
delayer off, the first request already hangs on TP2 DP attention
(DECIDE prefill_dp=(True, False): 526 tokens replaying the 1024 bucket
beside a 6-row dummy decode).

`_pad_to_token_graph` now replays a bucket only when every rank
prefills; otherwise every rank runs eagerly, where both sides take the
same variable-length gather from the group's table. This closes the
"known limitation" #2696 left for follow-up without changing the
delayer path, where every rank already prefills together.

Validated on MI355X, DeepSeek-V4.1-Flash TP2 --enable-dp-attention with
ATOM_ENABLE_PREFILL_DELAYER=0: main hangs on the first request; with
this change a single request, eight concurrent prompts of different
lengths, and new prompts arriving during a 1500-token decode all
complete. With the delayer on (default) main and this change both pass.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
yhl-amd added a commit that referenced this pull request Oct 25, 2026
#2696 makes every prefilling DP rank replay one token-graph bucket, but a
step can still pair a prefill with a decode: without the prefill delayer
(ATOM_ENABLE_PREFILL_DELAYER=0, or RapidServe) no rank learns before
scheduling that a peer will prefill, so an idle rank runs a dummy decode
and a busy one keeps decoding. The decoding rank then ran its early
layers eagerly beside a peer replaying the prefill recording. The MoE's
DP exchange differs on the two sides -- a recorded custom all_gather at
the bucket's height, on the peer buffers the group registered for that
recording, against an eager variable-length gather -- and the group
hangs: with the delayer off, the first request already does on TP2 DP
attention (526 tokens replaying the 1024 bucket beside a 6-row dummy
decode).

On a step some rank prefills, every rank now replays the prefill bucket
holding the group's largest row count, a decoding rank on its decode
rows padded to it:

- `ForwardMode._pad_to_token_graph` settles one bucket for the whole
  group whenever any rank prefills, and `ForwardMode.prefill_token_graph`
  marks the step;
- `region_graph._phase` keys such a step's lookup under "prefill", so
  the decoding rank finds the recording its peers replay. Its attention
  is an eager region and still runs the decode it scheduled; the
  recording's `rope_quant_window` hands it the window row, which the
  decode branch of `_attend` writes before it attends;
- `_record_tr_graphs` records each bucket with every rank at it.

The decoder after the KV source is unchanged: on such a step the
prefilling ranks run `DecoderGraph` eagerly (`req_graphs` stays off) and
a decoding rank runs every row, under the stages `decoder_stages` sizes
from the group's table.

Validated on MI355X, DeepSeek-V4.1-Flash TP2 --enable-dp-attention:
with the delayer off, main hangs on the first request; with this change
a single request, eight concurrent prompts of different lengths and new
prompts during a 1500-token decode complete, and GSM8K (full, 3-shot,
c64) scores 0.9287 / 0.9174 in two runs, where 459 steps had a rank
decode on the prefill bucket. With the delayer on (default): 0.9158
against main's 0.9227.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
zufayu added a commit that referenced this pull request Oct 26, 2026
… that exist

Rebasing onto main surfaced it: #2696 pruned every GLM before 5.3 and the
V4-Pro MTP3 / DPA MTP3 variants from .github/benchmark/models.json. Two of
the four judged entries no longer existed, and build_perf_check_matrix.py
did exactly what it is for -- refused to build a short matrix and named
both by hand:

  PERF_MODELS names 'glm-5-2-fp8', which is not a catalog prefix
  'deepseek-v4-pro' has no variant with suffix '-mtp3' in the catalog

Unrebased, this merges as a workflow that fails at its first step.

GLM-5.2-FP8 -> Qwen3.5-397B-A17B-FP8, the nearest analogue still in this
catalog: an FP8 MoE at tp=4. GLM-5.3 is not a candidate -- it carries
accuracy entries but no ATOM-native benchmark one, and GLM's remaining
perf coverage lives on the SGLang, vLLM and OOT lines, which run a server
atom_test.sh does not launch. Adding a GLM-5.3 benchmark entry is its own
change to shared config, not a line in this PR.

DeepSeek-V4-Pro MTP3 -> MiniMax-M3-MXFP8 EAGLE3 for the pair. V4-Pro's
only remaining speculative variant, DPA DSpark, starts at conc_min=64 and
so cannot reach the c=32 level this line judges at; M3-MXFP8 EAGLE3 spans
1..256 and its base is already an entry here, which is what the pair needs
-- only the speculative half moves, or both do. DeepSeek-V4.1-Flash was
tried first and rejected: the catalog gives it only a DSpark variant, so
it is itself a speculative entry with no base to pair against, and taking
it would have put two speculative entries in a four-entry matrix.

The shape is unchanged: four entries, two levels, eight cells, three
non-speculative plus one pair. PERF_EXPECTED_ENTRIES stays 4.

The WARMUP_MULT and PERF_PROMPT_MULTS comments keep GLM-5.2-FP8's measured
bias, which is the only evidence about what WARMUP_MULT 1 costs, but now
say the entry is gone and that none of the four here has been measured for
the same thing either way.

Also drops .coverage, a 53KB binary pytest-cov leaves in the repo root and
this branch committed by accident, and adds the .gitignore line that stops
the next one -- the rule the file's own comment says to prefer over care.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
valarLip pushed a commit that referenced this pull request Oct 26, 2026
…et (#2697)

* fix(dp): a decoding rank replays the group's prefill token-graph bucket

#2696 makes every prefilling DP rank replay one token-graph bucket, but a
step can still pair a prefill with a decode: without the prefill delayer
(ATOM_ENABLE_PREFILL_DELAYER=0, or RapidServe) no rank learns before
scheduling that a peer will prefill, so an idle rank runs a dummy decode
and a busy one keeps decoding. The decoding rank then ran its early
layers eagerly beside a peer replaying the prefill recording. The MoE's
DP exchange differs on the two sides -- a recorded custom all_gather at
the bucket's height, on the peer buffers the group registered for that
recording, against an eager variable-length gather -- and the group
hangs: with the delayer off, the first request already does on TP2 DP
attention (526 tokens replaying the 1024 bucket beside a 6-row dummy
decode).

On a step some rank prefills, every rank now replays the prefill bucket
holding the group's largest row count, a decoding rank on its decode
rows padded to it:

- `ForwardMode._pad_to_token_graph` settles one bucket for the whole
  group whenever any rank prefills, and `ForwardMode.prefill_token_graph`
  marks the step;
- `region_graph._phase` keys such a step's lookup under "prefill", so
  the decoding rank finds the recording its peers replay. Its attention
  is an eager region and still runs the decode it scheduled; the
  recording's `rope_quant_window` hands it the window row, which the
  decode branch of `_attend` writes before it attends;
- `_record_tr_graphs` records each bucket with every rank at it.

The decoder after the KV source is unchanged: on such a step the
prefilling ranks run `DecoderGraph` eagerly (`req_graphs` stays off) and
a decoding rank runs every row, under the stages `decoder_stages` sizes
from the group's table.

Validated on MI355X, DeepSeek-V4.1-Flash TP2 --enable-dp-attention:
with the delayer off, main hangs on the first request; with this change
a single request, eight concurrent prompts of different lengths and new
prompts during a 1500-token decode complete, and GSM8K (full, 3-shot,
c64) scores 0.9287 / 0.9174 in two runs, where 459 steps had a rank
decode on the prefill bucket. With the delayer on (default): 0.9158
against main's 0.9227.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* docs(dp): every rank of a prefill step replays the group's token-graph bucket

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
yhl-amd added a commit to SemiAnalysisAI/InferenceX that referenced this pull request Oct 26, 2026
…OM AgentX on nightly_202610261558

Image rocm/atom-dev:nightly_202610261558 (ATOM 713eb06, digest
sha256:57477e7a...) adds ROCm/ATOM#2696/#2697, which fix the hang every
DP-attention point hit on nightly_202610240802, and #2719 (trace labels).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
ZhangLirong-amd added a commit that referenced this pull request Oct 28, 2026
…A recipe (#2784)

* ci(agentic): run the DSv4.1 DTP2/DEP2 c128 cells with InferenceX's DPA recipe

#2696 added the DTP2 and DEP2 c128 cells on the model-wide flags only. The
InferenceX DSv4.1 ATOM AgentX DPA run (SemiAnalysisAI/InferenceX run
40362093275, TP2/DPA c128) also sets the DPA scheduling knobs. Without them a
prefill can land on almost every step: our 10-26 nightly DTP2 c128 reached
1334 output tok/s/GPU at 37.3 ms mean ITL, InferenceX 2123 at 15.8 ms.

Both DPA variants now match that recipe:
- env: ATOM_PREFILL_DECODE_INTERVAL=10, ATOM_DP_SESSION_AFFINITY=1,
  ATOM_DP_LB_REQ_EQUIV=512, ATOM_ENABLE_PREFILL_DELAYER=1, GPU_MAX_HW_QUEUES=5
- args: --dp-load-balance least_tokens, --kv_cache_dtype bf16,
  --gpu-memory-utilization 0.9, capture sizes up to 128

Variant args follow the model-wide ones, so the TP2/TP4 cells are unchanged.

* ci(agentic): keep the model-wide KV cache dtype for the DSv4.1 DPA cells

* ci(agentic): drop GPU_MAX_HW_QUEUES and the default prefill-delayer flag from the DSv4.1 DPA cells
zufayu added a commit that referenced this pull request Oct 29, 2026
… that exist

Rebasing onto main surfaced it: #2696 pruned every GLM before 5.3 and the
V4-Pro MTP3 / DPA MTP3 variants from .github/benchmark/models.json. Two of
the four judged entries no longer existed, and build_perf_check_matrix.py
did exactly what it is for -- refused to build a short matrix and named
both by hand:

  PERF_MODELS names 'glm-5-2-fp8', which is not a catalog prefix
  'deepseek-v4-pro' has no variant with suffix '-mtp3' in the catalog

Unrebased, this merges as a workflow that fails at its first step.

GLM-5.2-FP8 -> Qwen3.5-397B-A17B-FP8, the nearest analogue still in this
catalog: an FP8 MoE at tp=4. GLM-5.3 is not a candidate -- it carries
accuracy entries but no ATOM-native benchmark one, and GLM's remaining
perf coverage lives on the SGLang, vLLM and OOT lines, which run a server
atom_test.sh does not launch. Adding a GLM-5.3 benchmark entry is its own
change to shared config, not a line in this PR.

DeepSeek-V4-Pro MTP3 -> MiniMax-M3-MXFP8 EAGLE3 for the pair. V4-Pro's
only remaining speculative variant, DPA DSpark, starts at conc_min=64 and
so cannot reach the c=32 level this line judges at; M3-MXFP8 EAGLE3 spans
1..256 and its base is already an entry here, which is what the pair needs
-- only the speculative half moves, or both do. DeepSeek-V4.1-Flash was
tried first and rejected: the catalog gives it only a DSpark variant, so
it is itself a speculative entry with no base to pair against, and taking
it would have put two speculative entries in a four-entry matrix.

The shape is unchanged: four entries, two levels, eight cells, three
non-speculative plus one pair. PERF_EXPECTED_ENTRIES stays 4.

The WARMUP_MULT and PERF_PROMPT_MULTS comments keep GLM-5.2-FP8's measured
bias, which is the only evidence about what WARMUP_MULT 1 costs, but now
say the entry is gone and that none of the four here has been measured for
the same thing either way.

Also drops .coverage, a 53KB binary pytest-cov leaves in the repo root and
this branch committed by accident, and adds the .gitignore line that stops
the next one -- the rule the file's own comment says to prefer over care.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant