Repository navigation
refactor: one graph mechanism per region, settled DP step shapes, cheaper V4.1 decode prep - #2688
Merged
Merged
Conversation
This reverts commit 2c772a5.
…aper V4.1 decode prep Graph regions - TokenGraph (token buckets) and ReqGraph (request buckets) in atom/utils/region_graph.py are the runner's two graph mechanisms; decode, the V4.1 prefill decoder (DecoderGraph) and the DSpark draft record and replay through them. Bucket lookups live on the graphs (TokenGraph.bucket, __contains__), not in the forward context. - Prefill graph recording is one pass (_record_prefill_graphs); capture boilerplate is shared (_build_capture); dummy prefill and decode share dummy_execution. DP step shape - ForwardMode.decide runs the step's single DP exchange and settles the running tokens per DP rank into a runner-owned numpy buffer in place. DPMetadata holds only that table; gather_running_tokens is the one all-reduce path (skipped at dp=1). set_forward_context takes the table as a required argument; num_tokens names become scheduled/running tokens. V4.1 decode preparation (64k ISL, conc 4, TP4: prepare_decode 310 -> 176 us, decode step median 3105 -> 2919 us, mean GPU idle per step 169 -> 34 us) - Step indptrs are filled by their first reader (PagedAttentionCache.indptrs) and dropped by begin_forward, like unit tiles: a mono step fills none, a graphed forward records the fill. - The FP4 score stage's walk plan is laid out on the device in the mono step-meta launch (step_meta._walk_plan); index_plan.fill_walk_plan stays as its host reference, checked against it on random steps. WalkPlanner and the installers' metadata_builder argument are gone. - prepare_batch_step finds group members by name instead of scanning every member twice; CompressPlan drops cu_compress_cpu / compress_plan_cpu, which nothing read. - The previous step's speculative status is taken after the input ids are staged, published and their fill launched: that wait is the step's one CPU block on the GPU, and only the attention metadata after it reads the status. Cleanup - Patch-style branches, duplicate plumbing and dead code removed across the runner, attention builders (shared Fp4LaunchStorages, V4/MLA indexer rows), MoE, plugins and tests; docs updated to match. Validation: V4.1-Flash TP4 GSM8K 3-shot 0.9242 (CI baseline 0.9224); unit tests in the affected areas pass apart from the known environment failures.
Contributor
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
valarLip
added a commit
that referenced
this pull request
Oct 24, 2026
#2688 made `begin_forward`, `end_forward` and `engram_forward` TR-graph breaks by decorating the methods with `eager_region()`. Those methods are called from `_EarlyLayers` / `_LateLayers`, which are torch.compiled, and `eager_region` calls straight through while Dynamo traces. The compiled stage then calls `torch.ops.aiter.v41_begin_forward` / `v41_engram` directly, so no break happens at replay or record time. Recording the prefill token graphs, the Engram staging fork (`EngramStaging.start`: `stream.wait_stream(compute)`) lands inside the first segment, its join is in `v41_end_forward` at the end of the forward, and the first real break (layer 0's `_compress`) ends a capture with an unjoined stream: `hipErrorStreamCaptureUnjoined` at server start (DeepSeek-V4.1-Flash DSpark accuracy CI on main). Put `eager_region()` under `torch_compile_guard` on the three ops, as `dspark_block_attention`, `kda_attention_with_output` and `unified_attention_with_output_base` already do, and drop the method-level decorators that compilation discards. The test fake forward context gains `main_stream` / `in_hipgraph`, which an unrecorded break reads. Validation (CI image rocm/atom-dev:nightly_202610231620 + CI aiter wheel, MI355X TP4, --kv_cache_dtype bf16 --index_cache_dtype fp4 --method dspark --num-speculative-tokens 5): - before: same hipErrorStreamCaptureUnjoined as CI - after: prefill token graphs (256..3072) and request graphs recorded, server ready; GSM8K 3-shot flexible 0.9257 / strict 0.9265 - black, ruff clean; tests/test_region_graph.py, tests/attentions/deepseek_v41, tests/models/deepseek_v41 and the V4.1/Engram tests: 1504 passed Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
valarLip
added a commit
that referenced
this pull request
Oct 24, 2026
…d TBO decode as request graphs (#2692) * fix(tbo): lay decode ubatch metadata out for the requests each micro-batch runs #2688 dropped the TBO decode CUDA graph (`capture_tbo_graph`), so a decode step the DP group splits into micro-batches now always runs eagerly. Its micro-batches are cut from the scheduled requests (`maybe_create_ubatch_slices`: 25 scheduled -> [0, 12) / [12, 25)), but the attention backends still laid each ubatch's metadata out for half the padded batch, the split the removed graph had baked in (32 running -> [0, 16) / [16, 32)). Micro-batch 1's rows then attended with, and wrote KV slots for, other requests' context. gpt-oss-120b, TP2 + DP attention + EP + `--enable-tbo all`, GSM8K 3-shot flexible: 0.8901 (CI, 2c772a5) -> 0.5951 (CI, 9a0a520). Locally on the CI image: 2c772a5 0.8992, 9a0a520 0.6080; on 9a0a520, `--enable-tbo prefill` 0.8901 and TBO off 0.8931, while disabling prefill token graphs (0.5929) and `--enforce-eager` (0.7619) do not recover it. - `decode_ubatch_request_ranges` is the one decode split: the slices (`maybe_create_ubatch_slices`), the per-ubatch token counts the DP sync agrees on (`_precompute_decode`) and every backend's per-ubatch metadata (`_prepare_ubatch_decode` in aiter_attention, aiter_mla and deepseek_v4_attn) all take it. - Each backend keeps its padded per-ubatch width, which always covers the ubatch's requests; only the request ranges change. DeepSeek-V4's non-eager branch cut ranges at the DP-agreed width and is aligned too. - aiter_mla: `build_ubatch_metadata` sliced `dcp_local_context_lens` from half the padded batch as well; it now reads the ubatch's own start (`_tbo_full_running_bs` -> `_tbo_ubatch_request_ranges`). - tbo/__init__: sort the imports and `__all__` the new export touches. Tests: tests/test_tbo_decode_ubatch_split.py (CPU) checks slices, DP sync and backend widths agree; the H2D GPU tests that hard-coded the padded split now expect the scheduled one and check each ubatch's context_lens; the MLA dcp_local_context_lens test covers 3 scheduled requests. Validation (rocm/atom-dev:nightly_202610231620, MI355X): - gpt-oss-120b TP2 DP-attn EP `--enable-tbo all`, GSM8K 3-shot flexible 0.8825 / strict 0.1766 (threshold 0.87) - black, ruff clean on the changed files; CPU + RUN_H2D_GPU_TESTS=1 attention tests: 1217 passed. tests/test_h2d_v4_{,indexer_}publication.py fail the same 28 cases on main and on this branch (stale fakes, e.g. a config without max_num_seqs), so DeepSeek-V4's non-eager branch has no GPU test coverage here. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * feat(tbo): record TBO decode buckets as request graphs #2688 dropped `capture_tbo_graph` and left every decode the DP group splits into micro-batches to run eagerly. Record them on `ReqGraph` instead: with TBO decode on, a bucket that splits (`decode_bucket_splits`: running_bs > 2) is recorded split at the bucket, and a decode step replays it. - `ModelRunner.capture_cudagraph`: such a bucket's capture context carries its micro-batch slices (`decode_ubatch_request_ranges(bs)`) and per-rank rows -- every rank records the same bucket, so each micro-batch's rows are its own, repeated -- and `req_graph.record` records the UBatchWrapper forward as is: its persistent workers issue onto streams forked from the capture stream and `_run_ubatches` joins the comm stream before returning. TR_HYBRID has no split token graph; there a split decode stays eager. - `ForwardMode.decide(tbo_decode_graphs=...)`: a split decode no longer forces eager when the buckets were recorded split; prefill TBO still does. `ForwardMode.ubatch_split_bs` is the batch the step's micro-batches are cut from: the bucket on a replay (whatever each rank scheduled into it), the scheduled requests on an eager split. - Attention builders take it through `build(ubatch_split_bs=)` and every `_prepare_ubatch_decode` cuts its per-ubatch metadata at it (`split_bs`). DeepSeek-V4's read its widths from the forward context, which at build time is the previous forward's; it now takes the split explicitly, and its token offsets are indexed from the scheduled requests a padded ubatch can start past. - `decode_bucket_splits` is the one bs > 2 gate (`_precompute_decode`, the recording, V4's prepare and capture). - docs/compilation_cudagraph_guide.md: how TBO decode records and replays. Validation, gpt-oss-120b TP2 + DP attention + EP (rocm/atom-dev nightly_202610231620, MI355X): - `--enable-tbo all`: every decode bucket records; GSM8K 3-shot flexible 0.8984 / strict 0.1804 (2c772a5: 0.8992). - Replay vs 2c772a5's `capture_tbo_graph`, ISL/OSL 1024, mean TPOT / output tok/s: conc 16 8.85 ms / 1711 vs 8.71 / 1563; conc 64 11.59 / 4999 vs 14.57 / 3940; conc 128 14.43 / 8028 vs 14.35 / 7875. The eager split this replaces did not finish its conc-16 point in 30 minutes. - On this model TBO decode is slower than none with either recording: `--enable-tbo prefill` gives 5.78 / 8.09 / 10.33 ms TPOT (2c772a5 5.49 / 8.02 / 10.25). - tests/test_forward_mode.py: which decodes replay split and what they cut. CPU tests 1364 passed; RUN_H2D_GPU_TESTS=1 attention tests pass, and the V4 H2D tests fail the same 28 cases as main. black, ruff clean. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * ci: fail fast on GPU hosts that cannot run a job Add `GPU host health check` as the first step of the atom-test accuracy job on the DO and predownload runners, before any Docker cleanup or image pull. It fails the job, with an `Unhealthy GPU host` annotation, when: - the host has no rocm-smi. Runners registered under linux-atom-do-mi350x-8 without ROCm tools on the host (spur-*, crsuse2-m2m-*) failed 51 jobs at `rocm-smi --showmemuse`, 2026-10-23; - rocm-smi fails, times out (60 s), or prints no GPU rows; - any GPU reports N/A temperature or power. Of the failed accuracy jobs since 2026-10-21, all 30 that hung at server start with `HW Exception ... GPU Hang` (about 10 minutes in) showed this in the rocm-smi table the job printed before starting, on five hosts. None of the 83 other failures and none of 80 successful jobs (32 hosts) did. The GPUs that hung are exactly the ones reporting no power; one host had a single such GPU, and four hosts reported edge temperature N/A on every GPU with one or two lacking power. At least 9 of 16 GPU preflight timeouts started within 11 s of a hang ending, on a runner that replaced the hung one, matching a wedged GPU on the same host. For the GPUs reporting no power it prints `rocm-smi --showuniqueid`, which is rocminfo's GPU UUID, so the host can be found and drained. tests/test_gpu_host_health_check.py runs the step through a stand-in rocm-smi on the three tables CI hosts printed (healthy, one dead GPU, edge-only host with two dead GPUs) and on a missing, failing and row-less rocm-smi. Replaying the step on the rocm-smi tables of 121 historical CI jobs: all 30 that hung fail it, all 91 others pass. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
3 of 4 tasks
valarLip
added a commit
that referenced
this pull request
Oct 24, 2026
…2696) * fix: replay one prefill token-graph bucket across DP ranks; prune CI 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. * test(benchmark): key the agentic nightly grid by parallel layout
This was referenced Oct 26, 2026
ZLkanyo009
added a commit
that referenced
this pull request
Oct 26, 2026
Conflict in atom/utils/envs.py only: #2688 removed ATOM_DUAL_STREAM_PIECEWISE, which sat next to this branch's ATOM_TRITON_TOPK_SOFTMAX. Keep the removal and the new ATOM_TRITON_TOPK_SOFTMAX entry. topK.py merges cleanly; the Triton path still gates on torch.cuda.is_current_stream_capturing(), which holds under the Co-authored-by: Cursor <cursoragent@cursor.com> #2688 graph mechanism.
28 tasks
jiayyu
added a commit
that referenced
this pull request
Oct 26, 2026
…ill traces Review (Copilot, 10-26): - `prepare_mixed` called `prepare_prefill(batch, n_p_seqs)`, but main's #2688 made `running_tokens` a required argument, so every mixed step raised TypeError before attention metadata was built. The prefill half is never padded, so it now passes `running_bs=n_p_seqs, running_tokens=n_p_tokens`. The existing tests checked `prepare_mixed`'s source only; the new `test_prepare_mixed_calls_both_halves_with_their_real_signatures` runs it with both halves stubbed and their arguments bound against the real methods, and fails without this fix. - The schedule-time-advance loop in `postprocess` covers prefill rows only, yet still gated `tracing.prefill_complete` on `i >= total_seqs_num_decode` (from #2644). Prefill rows come first in a mixed batch, so a final chunk could fail that check and leave its prefill span open. Any final chunk now completes it; a pure prefill batch has no decode rows, so flag-off is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 task
yitingw1
added a commit
that referenced
this pull request
Oct 26, 2026
…split-tp main (#2688) moved the indexer scoring into `atom.model_ops.indexer` (`score_topk` over `IndexRows`, with a row split and an in-place gather), which supersedes this branch's own chunk loop and gather. The conflicted code takes main's side; the split is re-applied on top of it in the next commit. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
yitingw1
added a commit
that referenced
this pull request
Oct 26, 2026
…he shared scorer Re-apply the TP query-row split on top of the shared scorer from #2688. The V4 builder hands the prefill rows a `DcpSplit` over the TP group (`over_tp`), so each rank scores and top-ks only its slice of the rows against the full key set; the bands, the FP4 launch plan and the in-place gather are the scorer's, as for GLM under DCP. `deepseek_v4.py` only drops the up to tp-1 tail rows the gathered selection carries. The split's gather now runs on the group's pynccl communicator, for DCP as well. On the torch ProcessGroup the NCCL communicator and its stream are created at the first collective; nothing else in the step uses that group and warmup skips the indexer, so they appeared at the first real prefill, after the KV budget was fixed. With a tight budget that ended in an HSA out-of-resources abort while launching the RCCL kernel. pynccl is built with the group, counted in the memory profile, and runs on the current stream. The DCP path falls back to the ProcessGroup only without a usable pynccl; the TP split is not turned on without one. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
1 task
This was referenced Oct 26, 2026
1 task
valarLip
pushed a commit
that referenced
this pull request
Oct 27, 2026
…2125) * [PCP] Use TP for shared_experts * [PCP] Sync sampled ids outside the speculative branch, and reject PCP+TBO+spec * [Fix] Fix PCP + DSpark part I * [PCP] Step the shared-expert fold down to a block-aligned grid `pcp_merged_tp_grid` cut unconditionally on the full pcp*tp grid, which faults in the weight loader whenever `moe_inter_dim / (pcp*tp)` is not a 128 multiple -- any tp*pcp=16 deployment, where 3072/16 = 192 splits a block-quant scale block in half. It now takes the largest divisor that stays aligned and reports a `prescale` of g/(pcp*tp). The two cases that were hand-written fall out of that formula -- no scale on the full grid, the familiar 1/pcp_size when the module stays replicated -- so the pre-scale branch that had become unreachable carries weight again. 8-card configs are unchanged: still cut 8 ways, still prescale 1.0. Also from review: - Refuse ATOM_PCP_MOE_MERGE together with MoE's flatten-TP-across-DP. That fold shards routed across pcp*dp*tp, while a module summed only by the tp all_reduce and the pcp collective never sees a DP reduce. - Refuse online re-quantisation when the effective TP grid is wider than the real TP group: the gather runs on the real group and the shard-back on the override, so it would narrow 1/pcp of a weight as if it were whole. DCP's override narrows instead and is unaffected. - One `_needs_id_sync` predicate behind both the PCP and the TP id broadcast, which have to stay the same condition. - Tests: the two folds asserted against each other on every rank (a reversed flatten goes red on tp4pcp2/tp2pcp4 and, as expected, not on pcp8 where tp=1 cannot tell them apart), the aligned step-down, and that shards tile the ranks evenly -- which is what makes prescale a scalar rather than a per-rank factor. - The import guard uses skip_if_dependency_missing rather than `except Exception`, which would read a SyntaxError as a skip. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * [PCP] Fold the padding rows into the last request's FP4 indexer group Since #2688 the FP4 prefill scorer is the row-group kernel, which reads group b's keys off block_tables[b] and asserts the table covers every group. pcp_owned_rows_by_request still gave the PCP padding rows a group of their own, one more than the requests the table has rows for, so any PCP prefill whose token count is not a multiple of pcp asserted (block_tables.shape[0] >= batch). The old prefill kernel looked rows up by batch_id_per_q_token and skipped bid < 0, so the extra group was harmless there. Padding rows sit at position 0, so their bound is 0 and they read no key in whichever group they land in. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * [PCP] Sum the DSpark draft FFN across PCP like Block.forward Under PCP the MoE weights are folded into the TP grid for every model, the draft included, so each rank's FFN output is a partial that needs the PCP all_reduce. Block.forward (target and MTP) gets it from moe_pcp_merge_forward; DSparkLayer.forward_block rewrote the FFN step as a plain self.ffn() call and never summed, which halved the DSpark acceptance rate under PCP. The draft graph is recorded with is_prefill=False and is_dummy_run=False, so the decode branch's all_reduce is what goes into it. V4-Pro-DSpark, gsm8k 200 x 5-shot, k=5: pcp8 32.96% -> 63.14% acceptance against a same-source tp8 at 63.48%. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * style: sort the pcp_utils imports in deepseek_v2 Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * [PCP] Address review: drop DSpark's second aux gather, settle pcp_split_rows in every forward - DSparkProposer.propose called aux_for() only to check it was not None, which under a PCP prefill all_gathered the aux buffers a second time (compute_draft_kv already gathered them) and threw the result away. Check _captures_aux instead. - pcp_split_rows was set only in ModelRunner.forward; RapidServe's decode forward overrides it and never did, so aux was not gathered there. Both forwards now call _settle_pcp_split_rows() after run_model. - Tests: the FP4 owned-rows case now has a real padding row, plus a last-request-owns-only-padding case and a brute-force comparison; a structural test pins DSpark's forward_block to the same moe_pcp_merge_forward branch as Block.forward. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * [PCP] Keep an online-requantized shared expert on the TP grid Online re-quantization gathers the full weight on the real TP group before quantizing (per_Token on the row-parallel w2 always needs it). Folded onto the wider pcp*tp grid, that gather rebuilds only 1/pcp of the weight, which linear.py now refuses at load -- turning a config that ran on main (the shared expert unsharded under PCP) into a load failure. When any of the shared expert's linears will be re-quantized online, do not fold it: it stays on TP exactly as on main. Offline-quantized checkpoints are unaffected. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * [PCP] Put forward's decorators back on forward f324bd9 inserted _settle_pcp_split_rows between ModelRunner.forward and its decorators, so the helper picked up @_profile_runner_stage, @torch.inference_mode() and @with_eplb_forward_monitor -- the last one binds a batch argument, so every forward raised TypeError -- and forward lost all three. A test now calls the helper bare and checks that forward is still wrapped. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * [PCP] Warn when a PCP model lacks pcp_split_active; trim comments - _settle_pcp_split_rows warns once when PCP is on, a drafter is present and the model has no pcp_split_active(): it then assumes no row-split, and a future PCP model that forgot the method would silently leave the drafter's aux ungathered. - Shorten the comments and docstrings this branch added to what/why; the derivations live in the PR description. No code change beyond the warning (AST-compared with docstrings stripped). Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * [PCP] Drop shared-expert sharding and the DSpark FFN branch ahead of #2725 main's #2725 moves the PCP gather/reduction into FusedMoE._run_experts, around the routed experts only, and computes the shared expert on each rank's own 1/pcp rows. That already takes the prefill gain the shared-expert sharding was after, and a shared expert sharded across PCP can no longer be summed in prefill (the ranks hold different rows). It also deletes moe_pcp_merge_forward, which the DSpark draft FFN fix called; FusedMoE now all-reduces the draft's MoE itself. Removed: pcp_merged_tp_grid / PcpSharedGrid, Expert(tp_grid), the shared prescale, the online re-quant check, linear.py's override forwarding and gather guard, DSparkLayer.forward_block's merge branch, and their tests. Kept: sampled-id sync and the PCP+TBO+spec guard, pcp_split_active and the DSpark aux gather, the FP4 indexer padding fix. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * test(pcp): give the draft-split tests a forward context #2725 makes _pcp_split_draft_inputs flag the pass as split on the forward context for FusedMoE's PCP gather, so the CPU tests have to supply one. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * test(pcp): pin the DSpark aux gather and the draft-split decision aux_for must gather exactly when the target forward row-split, and the draft prefill split must follow that answer (pcp_split_rows), not the current pass's pcp_split, which drafting clears. Reverting either fix now fails a test instead of silently halving DSpark's acceptance. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * [PCP] Rename pcp_split_rows to target_pcp_split Next to #2725's per-pass pcp_split the old name read as a synonym. It is the target forward's answer, fixed right after it and kept through drafting, which clears pcp_split per pass; it is now also a Context __init__ argument like its sibling. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * [PCP] Snapshot pcp_split for the drafter instead of re-deriving it Since #2725 both deepseek_v2 and deepseek_v4 write the target forward's split decision to context.pcp_split, and FusedMoE gathers on it. target_pcp_split now copies that value right after run_model, before drafting clears it, instead of re-deriving it through a separate pcp_split_active(). The aux gather and FusedMoE thus read one decision. Drops pcp_split_active() from pcp_utils and both models, the _pcp_active delegation, and the missing-method warning; this PR no longer touches the model files. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
1 task
whx-sjtu
added a commit
that referenced
this pull request
Oct 28, 2026
Add `--sequence-parallel-size`, riding the PCP process group as an axis independent of TP: non-attention layers run on T/W tokens, and an all-to-all around attention trades the sequence for a head slice. MoE gathers the SP shards, runs this rank's expert slice and reduce-scatters the partial sums; large MXFP4 batches route and quantize before the gather. ATOM_SP_MOE_MEGA swaps the M3 prefill MoE block for AITER's MegaMoE, with the native norm, gate and top-k in front of it. Ported onto the request/token graph runner (#2688) and the route-before-gather MoE (#2725): SP graphs capture this rank's token shard at a fixed address, gathered SP rows place their experts after the gather and skip the per-rank row count, and SP keeps its own reduce-scatter instead of the PCP-merge all-reduce. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
whx-sjtu
added a commit
that referenced
this pull request
Oct 28, 2026
Add `--sequence-parallel-size`, riding the PCP process group as an axis independent of TP: non-attention layers run on T/W tokens, and an exchange around attention trades the sequence for a head slice (GQA KV heads and an indexer tail are replicated to the ranks whose queries use them). MoE gathers the SP shards, runs this rank's expert slice and reduce-scatters the partial sums; large MXFP4 batches route and quantize before the gather. ATOM_SP_MOE_MEGA (default on) runs eligible SP MXFP4 prefill MoE blocks through AITER's MegaMoeTP, with the native norm, gate and top-k in front of it. Ineligible configs, layers or a failed runtime setup keep the native MoE. ATOM_TP_REPLICATED_O_PROJ (experimental, default off) keeps o_proj replicated under TP and exchanges heads instead of all-reducing. Both are generic and opted into per model; MiniMax-M3 is the first user. Ported onto the request/token graph runner (#2688) and the route-before-gather MoE (#2725): SP graphs capture this rank's token shard at a fixed address, gathered SP rows place their experts after the gather and skip the per-rank row count, and SP keeps its own reduce-scatter instead of the PCP-merge all-reduce. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
whx-sjtu
added a commit
that referenced
this pull request
Oct 28, 2026
Add `--sequence-parallel-size`, riding the PCP process group as an axis independent of TP: non-attention layers run on T/W tokens, and an exchange around attention trades the sequence for a head slice (GQA KV heads and an indexer tail are replicated to the ranks whose queries use them). MoE gathers the SP shards, runs this rank's expert slice and reduce-scatters the partial sums; large MXFP4 batches route and quantize before the gather. ATOM_SP_MOE_MEGA (default on) runs eligible SP MXFP4 prefill MoE blocks through AITER's MegaMoeTP, with the native norm, gate and top-k in front of it. Ineligible configs, layers or a failed runtime setup keep the native MoE. ATOM_TP_REPLICATED_O_PROJ (experimental, default off) keeps o_proj replicated under TP and exchanges heads instead of all-reducing. Both are generic and opted into per model; MiniMax-M3 is the first user. Ported onto the request/token graph runner (#2688) and the route-before-gather MoE (#2725): SP graphs capture this rank's token shard at a fixed address, gathered SP rows place their experts after the gather and skip the per-rank row count, and SP keeps its own reduce-scatter instead of the PCP-merge all-reduce. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
demonsan
added a commit
that referenced
this pull request
Oct 29, 2026
Brings the stacked Ulysses-SP + gemm_fusion series (70 commits; 68 files vs main) onto current main, resolving its collisions with #2688 "one graph mechanism per region", #2725 "MoE route before rows move", and the other commits merged since 4438322. Conflicts resolved (7 files): - atom/model_engine/model_runner.py (9 hunks): SP hooks ported onto main's req_graph / ForwardMode.token_graph mechanism (padded-input staging, local width capture, SP gather on replay, tr-graph record path); the gemm_fusion launch-gate check stays immediately before the single req_graph.replay(); the SP graph-capture context manager and the gemm_fusion before/after_graph_capture hooks are kept. - atom/model_engine/engine_core.py (1 hunk): kept the KV-scrub RPC next to main's has_pending_kv_work gate. - atom/model_ops/attention_mha.py (1 hunk): kept the M3 sparse fused-prologue methods and the dense-prologue eligibility change. - atom/model_ops/moe.py (10 hunks): SP gather/prequantize/local-topk hooks re-expressed on main's select_experts + _run_experts flow; the x_scale and scheduled_tokens_gpu handoff is kept; sp_moe_reduce_scatter restored before the reduce_results all-reduce. Cross-hunk adaptations: `_run_experts` gained keyword-only `pcp_handled=False` (default False is main's behavior exactly; SP/token-gather callers pass True so the PCP fold does not move the rows a second time), and `FusedMoE.select_experts` keeps the union's `shared_experts_fused=None` argument (its default reproduces main's fuse_shared logic). - atom/utils/forward_context.py (2 hunks): imports unioned; main's DPMetadata(running_tokens_per_dp_rank) is kept, with the SP expansion of context.running_tokens_across_dp when Ulysses SP is enabled. - atom/utils/backends.py (1 hunk): took main's typing import; this branch's change there was a stale rebase artifact, so the file is now identical to main. - tests/test_model_runner_decode_padding.py (2 hunks): main's token_graph fixture plus the SP assertions (split input, gathered output, padded height). Re-baselined onto the new APIs (test-only): - tests/model_ops/test_gemm_fusion_rebase_compat.py: the compute_hash stand-in gained main's new env/geometry factors; the prepare_inputs pin follows the running_tokens_per_dp_rank API. - tests/model_ops/test_gemm_fusion_r12_abort_path.py: the replay pin follows self.req_graph.replay() and still requires the launch-gate check right before it. - tests/model_ops/test_sp_moe_dispatch.py: ported to main's route-before-gather API (26 tests). - tests/model_ops/test_shared_expert_dispatch.py: the branch's select_experts_with_record test re-expressed on FusedMoE.select_experts' shared_experts_fused argument. The 5 shipped files (atom/model_ops/gemm_fusion/resources.py, hooks.py, site.py, tests/model_ops/test_gemm_fusion.py, tests/model_ops/test_gemm_fusion_aiter_contract.py) are blob-identical to efb0252. Verified without a GPU: 474 CPU tests pass (8 PR-scoped files + 2 ported MoE test files; test_shared_expert_dispatch.py skips locally without aiter), Black clean, Ruff has no findings on lines this merge adds. GPU e2e for SP graph replay, gemm_fusion gate capture/replay placement, and the r18 band is still pending.
demonsan
added a commit
that referenced
this pull request
Oct 29, 2026
CI non-GPU unit tests failed on four test-side drifts introduced when main's #2688/#2725 refactors merged into the SP series; product code is unchanged: - test_moe_layouts_cannot_reuse_incompatible_sp_compilation_artifacts: Config.compute_hash now reads max_model_len / max_num_seqs directly. - test_ulysses_cpu_metadata prepare_inputs test: supply running_tokens_per_dp_rank, forward_vars['batch_info'], the sync ub fields, ubatch_split_bs and token_graph; assert the SP table the merged set_forward_context publishes. - test_forward_context_without_aiter_or_triton: subprocess snippet now uses the running_tokens_per_dp_rank signature and SP table semantics. Local: 34 passed across the two files; E0-equivalent selection 562 passed / 3 skipped; Black and Ruff clean.
This was referenced Oct 29, 2026
jiayyu
added a commit
that referenced
this pull request
Oct 30, 2026
…ill traces Review (Copilot, 10-26): - `prepare_mixed` called `prepare_prefill(batch, n_p_seqs)`, but main's #2688 made `running_tokens` a required argument, so every mixed step raised TypeError before attention metadata was built. The prefill half is never padded, so it now passes `running_bs=n_p_seqs, running_tokens=n_p_tokens`. The existing tests checked `prepare_mixed`'s source only; the new `test_prepare_mixed_calls_both_halves_with_their_real_signatures` runs it with both halves stubbed and their arguments bound against the real methods, and fails without this fix. - The schedule-time-advance loop in `postprocess` covers prefill rows only, yet still gated `tracing.prefill_complete` on `i >= total_seqs_num_decode` (from #2644). Prefill rows come first in a mixed batch, so a final chunk could fail that check and leave its prefill span open. Any final chunk now completes it; a pure prefill batch has no decode rows, so flag-off is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
jiayyu
added a commit
that referenced
this pull request
Oct 30, 2026
…ill traces Review (Copilot, 10-26): - `prepare_mixed` called `prepare_prefill(batch, n_p_seqs)`, but main's #2688 made `running_tokens` a required argument, so every mixed step raised TypeError before attention metadata was built. The prefill half is never padded, so it now passes `running_bs=n_p_seqs, running_tokens=n_p_tokens`. The existing tests checked `prepare_mixed`'s source only; the new `test_prepare_mixed_calls_both_halves_with_their_real_signatures` runs it with both halves stubbed and their arguments bound against the real methods, and fails without this fix. - The schedule-time-advance loop in `postprocess` covers prefill rows only, yet still gated `tracing.prefill_complete` on `i >= total_seqs_num_decode` (from #2644). Prefill rows come first in a mixed batch, so a final chunk could fail that check and leave its prefill span open. Any final chunk now completes it; a pure prefill batch has no decode rows, so flag-off is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
PerryZhang01
added a commit
to jgong5/ATOM
that referenced
this pull request
Oct 30, 2026
…2495) `_graph_pool` asks the runner's `_tr_hybrid_active`, the name upstream 9a0a520 (ROCm#2688) gave `_piecewise_cg_active`, so a simulated start that captures graphs gets past the graph-pool reservation again. A test compiles ATOM's `_estimate_cudagraph_overhead` from source and checks `_graph_pool` reserves the same bytes under each graph mode, so the next rename fails the CPU gate. Closes #2489. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
valarLip
pushed a commit
that referenced
this pull request
Oct 30, 2026
) * perf(v4): split the prefill indexer by query row across plain TP The DeepSeek-V4 indexer is replicated across TP (ReplicatedLinear projections, unsharded index heads), so under plain TP every rank scored the same prefill rows against the same compressed keys. Each rank now scores and top-ks only its slice of the rows, against the full key set, and the [rows, topk] selections are all-gathered back in place over the TP group. The FP4 prefill schedule is built for the rank's rows only. Applies to the FP8 and gfx950 FP4 prefill paths under plain TP; DP attention, TBO, PCP, DCP, simulated TP, plugin mode and the gfx1250 OPUS scorer keep the full per-rank scoring. ATOM_V4_INDEXER_QSPLIT=0 turns it off. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * perf(v4): split the prefill indexer by query row across plain TP on the shared scorer Re-apply the TP query-row split on top of the shared scorer from #2688. The V4 builder hands the prefill rows a `DcpSplit` over the TP group (`over_tp`), so each rank scores and top-ks only its slice of the rows against the full key set; the bands, the FP4 launch plan and the in-place gather are the scorer's, as for GLM under DCP. `deepseek_v4.py` only drops the up to tp-1 tail rows the gathered selection carries. The split's gather now runs on the group's pynccl communicator, for DCP as well. On the torch ProcessGroup the NCCL communicator and its stream are created at the first collective; nothing else in the step uses that group and warmup skips the indexer, so they appeared at the first real prefill, after the KV budget was fixed. With a tight budget that ended in an HSA out-of-resources abort while launching the RCCL kernel. pynccl is built with the group, counted in the memory profile, and runs on the current stream. The DCP path falls back to the ProcessGroup only without a usable pynccl; the TP split is not turned on without one. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * test(v4): stub the TP-split names in the transfer-regions builder test; note the GLM TP split The test loads the real V4 builder with only the names it imports stubbed; the builder now also names `get_tp_group`, `DcpSplit` and `atom.model_ops.indexer.score.split_gather_comm`, so all 41 cases failed at setup on the non-GPU runner. Stub them (the test runs at TP 1, so the split gate returns before touching any of them). Also leave a note at the GLM builder's DCP split: the GLM indexer is replicated across TP too, so the same split could cover dcp=1. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * refactor(indexer): rename DcpSplit to RowSplit, over_tp to group, dcp_qsplit_rows to qsplit_rows The row split serves the TP-replicated V4 indexer as well as GLM under DCP, so neither the type nor the helper is DCP's any more. `RowSplit.group` names the ranks the rows are split over and gathered on ("dcp" or "tp"), in place of the `over_tp` flag. No behavior change. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * fix(indexer): check ranks agree on the split gather's transport; harden the split; cover the builder Review follow-ups for the TP row split: - `split_gather_comm_on_every_rank`: aiter builds a group's pynccl per rank and only warns when that fails, so ranks could disagree on whether to gather over pynccl or the ProcessGroup (or, under TP, whether to split at all) and hang. The V4 gate and the GLM DCP builder check it once at startup over the CPU group and raise on disagreement. - The pynccl gather asserts contiguous rows (it reads raw pointers, where the ProcessGroup would refuse a strided view). - `RowSplit.group` is `Literal["dcp", "tp"]` and `RowSplit.of` rejects anything else instead of falling through to the DCP group. - `_index_rows` refuses a split on the FP8 paged path, which never gets one. - Tests: the builder hands each prefill path this TP rank's split and the gate follows the TP pynccl (via the transfer-regions builder harness); the startup agreement; a multi-GPU test of the in-place gather on pynccl, empty ranks included. Wording left over from DCP-only in docstrings, the model-ops guide and a test name; the env doc lists the pynccl condition. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * fix(indexer): vote on the host, and only when the split would run - The startup pynccl agreement builds its vote under the runner's CUDA default device; put it on the host explicitly for the gloo CPU group. - The V4 gate checks every rank-invariant condition first and only asks the ranks about pynccl when it would actually split, so a run with the flag off (or DCP/PCP/TBO/...) never reaches the collective. - Multi-GPU test: also runs the startup agreement under the CUDA default device, and kills ranks still alive at its deadline rather than leaving them for interpreter exit to join. - Comment why the GLM builder imports the check locally; gate test covers the flag off (and that it does not ask). Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * perf(v4): split the prefill indexer only on steps with enough index logits The row split saves index scoring, which grows with a step's index logits -- each request's new rows times its committed keys -- and, on FP4, the launches of logits bands a rank no longer needs. It pays a per-layer top-k gather, which grows with the rows and waits on the slowest rank. A step under `QSPLIT_MIN_LOGITS` (2^26) is scored whole: 8192 rows of short cold prompts (0.5M logits) gain nothing from the split, while 4096 rows on a 516,096-token prefix (533M) run a third faster. The bound is a compromise measured on one setup (TP8 MI355X, FP4 index, 32 seqs, 6 speculative tokens, 1M max_model_len, prefix hits). The two nearest points go the other way, step time split vs whole, though both drop from two FP4 bands to one: 66.1M (512 rows, host-bound) -3.6 %, 68.2M (2048 rows, GPU-bound, 1 sample) +3.0 %. FP8, TP4 and cold prompts near the bound are uncalibrated. The logits come from the scheduled batch's host metadata (`extend_lens_np`, `n_committed`), the same on every rank, summed in int64 (16K rows of a 524,288-token request reach 2^31). PCP and TBO, which never split, pass no per-request rows. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * fix(plugin): the sglang GLM-5.2 DSA bridge splits rows with RowSplit `DcpSplit` was renamed `RowSplit` in this PR; the bridge imports it inside a function, so the old name only failed when a DCP prefill reached it. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * perf(v4): stripe the TP indexer split's rows; gather on the ProcessGroup Addresses the #2722 review: - Rows are dealt round-robin: rank r scores rows r, r + W, ... (RowSplit.striped). A row scores the keys it sees, (pos + 1) // 4, so contiguous shares left the last rank the most: 1.88x the mean in a cold 16K chunk, and nearly all of a long prefix hit's rows on one rank beside short cold prompts. The scorer takes this rank's rows of the query and weights, scores them into its slot, gathers, and transposes the rank-major slots back to row order. FP4 rows are copied with clone(): the FP4 kernels assert dense rows, and .contiguous() hands a single row back with the stripe's stride. FP8 reads them strided. A striped split returns a new tensor, so the scorer refuses it with out, on_band, keys or request_local=False. The builder slices the visible ends, request ids, FP8 row starts and FP4 query starts per request to this rank's rows. - The gate counts the keys a step's rows see, sum(n * n_committed) - sum(n * n) / 8, in int64, against ATOM_V4_INDEXER_QSPLIT_MIN_LOGITS (new, default 2^26). - The gather is on the TP group's ProcessGroup again, not pynccl, so the pynccl gate and the startup vote are gone. RCCL allocates on a group's first all-gather (448 MiB a GPU at TP8), and warmup's prefill stays under the bound, so a run that splits gathers once at builder construction, before the memory profile, which then counts it. - Tests: striped rows back in row order over threaded ranks (CPU) and 2/4/8 ranks, one-row and empty ranks included (GPU); the gate's count; the startup gather; the V4 op entry handing back the step's rows alone, in row order, as its fake promises. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * docs(indexer): say where the gather allocates; label the model-entry test - _gather_split: "beyond the communicator's own first-use buffers" could read as nothing being allocated once the communicator exists (after warmup's broadcast, say); RCCL also sets memory aside on a group's first all-gather, so say that. - qsplit_rows: the last ranks' buffers can have tails; they have none when the rows divide evenly. - The model-entry test is renamed for what it checks (the step's rows, in order), and its cases carry ids naming the tail each one leaves. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(v4): keep a P/D decode instance's prefills whole A decode instance (mooncake / moriio, kv_role kv_consumer) prefills only to recompute, too rarely to pay for the split's startup gather (448 MiB a GPU at TP8). The gate turns the split off there, so the gather never runs. A load-only LMCache offload node sets the same role and still splits. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [V4] qsplit: multi-aware P/D decode gate, startup gather out of the getter Address review: - P/D decode detection reads the kv_role of the P/D sub-connector (factory.pd_transfer_config), so a decode instance under `multi` no longer splits. A `multi` with a top-level kv_consumer and a P/D sub without a role now splits: that sub runs as a producer. - _indexer_qsplit_world_size is a pure query; the startup gather moves to _init_indexer_qsplit, called from __init__. - ATOM_V4_INDEXER_QSPLIT is parsed with _flag_env ("true" is on). - RowSplit validates `group` in both constructors. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [V4] qsplit: test that __init__ inits the split; pin three more P/D shapes - The builder's __init__ calls _init_indexer_qsplit (checked on its source: the builder is too heavy to construct in a unit test). - moriio decode, multi with a load-only offload alone, and a role on multi alone (the sub runs as a producer, so the run splits). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [V4] qsplit: drop the plugin input from the gate The V4 builder is only built by ATOM's own ModelRunner; plugin bridges use their own builders, so is_plugin_mode() was always False there. Drops the atom.plugin import from deepseek_v4_attn.py. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [V4] qsplit: leave striped top-k rank-major, remap in csa_translate_pack Drop the per-layer transpose copy after the striped gather. score_topk returns the gathered [W*S, K] buffer as is; csa_translate_pack reads row t at (t % W) * S + t // W via constexpr ROW_WORLD and runtime row_stride. ROW_WORLD=1 compiles the same kernel as before. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
yitingw1
added a commit
that referenced
this pull request
Oct 30, 2026
The file failed 20 of 22 cases on main with RUN_H2D_GPU_TESTS=1, all from fixtures that fell behind the code: - set_pcp patched pcp_round_robin_query_indices (removed in #2688); the builder reads get_pcp_rank. - make_runner lacked forward_vars["batch_info"], config.max_num_seqs and the builder's csa_layers, max_model_len_idx, _local_heads and index_logits_workspace. - The padded decode case modelled the bucket width through UBatchWrapper._decode_ub_running_bs; since #2692 the split comes from _prepare_ubatch_decode's split_bs, so the test passes the bucket. - prepare_prefill was called with a total_tokens_num the fake batch lacks. 22/22 pass on this branch and on its merge base. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
1 task
whx-sjtu
added a commit
that referenced
this pull request
Oct 30, 2026
Add `--sequence-parallel-size`, riding the PCP process group as an axis independent of TP: non-attention layers run on T/W tokens, and an exchange around attention trades the sequence for a head slice (GQA KV heads and an indexer tail are replicated to the ranks whose queries use them). MoE gathers the SP shards, runs this rank's expert slice and reduce-scatters the partial sums; large MXFP4 batches route and quantize before the gather. ATOM_SP_MOE_MEGA (default on) runs eligible SP MXFP4 prefill MoE blocks through AITER's MegaMoeTP, with the native norm, gate and top-k in front of it. Ineligible configs, layers or a failed runtime setup keep the native MoE. ATOM_TP_REPLICATED_O_PROJ (experimental, default off) keeps o_proj replicated under TP and exchanges heads instead of all-reducing. Both are generic and opted into per model; MiniMax-M3 is the first user. Ported onto the request/token graph runner (#2688) and the route-before-gather MoE (#2725): SP graphs capture this rank's token shard at a fixed address, gathered SP rows place their experts after the gather and skip the per-rank row count, and SP keeps its own reduce-scatter instead of the PCP-merge all-reduce. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
jiayyu
added a commit
that referenced
this pull request
Nov 2, 2026
…ill traces Review (Copilot, 10-26): - `prepare_mixed` called `prepare_prefill(batch, n_p_seqs)`, but main's #2688 made `running_tokens` a required argument, so every mixed step raised TypeError before attention metadata was built. The prefill half is never padded, so it now passes `running_bs=n_p_seqs, running_tokens=n_p_tokens`. The existing tests checked `prepare_mixed`'s source only; the new `test_prepare_mixed_calls_both_halves_with_their_real_signatures` runs it with both halves stubbed and their arguments bound against the real methods, and fails without this fix. - The schedule-time-advance loop in `postprocess` covers prefill rows only, yet still gated `tracing.prefill_complete` on `i >= total_seqs_num_decode` (from #2644). Prefill rows come first in a mixed batch, so a final chunk could fail that check and leave its prefill span open. Any final chunk now completes it; a pure prefill batch has no decode rows, so flag-off is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 task done
whx-sjtu
added a commit
that referenced
this pull request
Nov 3, 2026
Add `--sequence-parallel-size`, riding the PCP process group as an axis independent of TP: non-attention layers run on T/W tokens, and an exchange around attention trades the sequence for a head slice (GQA KV heads and an indexer tail are replicated to the ranks whose queries use them). MoE gathers the SP shards, runs this rank's expert slice and reduce-scatters the partial sums; large MXFP4 batches route and quantize before the gather. ATOM_SP_MOE_MEGA (default on) runs eligible SP MXFP4 prefill MoE blocks through AITER's MegaMoeTP, with the native norm, gate and top-k in front of it. Ineligible configs, layers or a failed runtime setup keep the native MoE. ATOM_TP_REPLICATED_O_PROJ (experimental, default off) keeps o_proj replicated under TP and exchanges heads instead of all-reducing. Both are generic and opted into per model; MiniMax-M3 is the first user. Ported onto the request/token graph runner (#2688) and the route-before-gather MoE (#2725): SP graphs capture this rank's token shard at a fixed address, gathered SP rows place their experts after the gather and skip the per-rank row count, and SP keeps its own reduce-scatter instead of the PCP-merge all-reduce. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
whx-sjtu
added a commit
that referenced
this pull request
Nov 3, 2026
Add `--sequence-parallel-size`, riding the PCP process group as an axis independent of TP: non-attention layers run on T/W tokens, and an exchange around attention trades the sequence for a head slice (GQA KV heads and an indexer tail are replicated to the ranks whose queries use them). MoE gathers the SP shards, runs this rank's expert slice and reduce-scatters the partial sums; large MXFP4 batches route and quantize before the gather. ATOM_SP_MOE_MEGA (default on) runs eligible SP MXFP4 prefill MoE blocks through AITER's MegaMoeTP, with the native norm, gate and top-k in front of it. Ineligible configs, layers or a failed runtime setup keep the native MoE. ATOM_TP_REPLICATED_O_PROJ (experimental, default off) keeps o_proj replicated under TP and exchanges heads instead of all-reducing. Both are generic and opted into per model; MiniMax-M3 is the first user. Ported onto the request/token graph runner (#2688) and the route-before-gather MoE (#2725): SP graphs capture this rank's token shard at a fixed address, gathered SP rows place their experts after the gather and skip the per-rank row count, and SP keeps its own reduce-scatter instead of the PCP-merge all-reduce. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
whx-sjtu
added a commit
that referenced
this pull request
Nov 3, 2026
Add `--sequence-parallel-size`, riding the PCP process group as an axis independent of TP: non-attention layers run on T/W tokens, and an exchange around attention trades the sequence for a head slice (GQA KV heads and an indexer tail are replicated to the ranks whose queries use them). MoE gathers the SP shards, runs this rank's expert slice and reduce-scatters the partial sums; large MXFP4 batches route and quantize before the gather. ATOM_SP_MOE_MEGA (default on) runs eligible SP MXFP4 prefill MoE blocks through AITER's MegaMoeTP, with the native norm, gate and top-k in front of it. Ineligible configs, layers or a failed runtime setup keep the native MoE. ATOM_TP_REPLICATED_O_PROJ (experimental, default off) keeps o_proj replicated under TP and exchanges heads instead of all-reducing. Both are generic and opted into per model; MiniMax-M3 is the first user. Ported onto the request/token graph runner (#2688) and the route-before-gather MoE (#2725): SP graphs capture this rank's token shard at a fixed address, gathered SP rows place their experts after the gather and skip the per-rank row count, and SP keeps its own reduce-scatter instead of the PCP-merge all-reduce. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 task done
davzhuAMD
added a commit
to davzhuAMD/ATOM
that referenced
this pull request
Nov 4, 2026
Let the Kimi-K3 attention builder take a prefill padded to a recorded token-graph bucket (supports_token_graph), so its dense layers replay as region graphs between the KDA and MLA eager regions. - kimi_mla_gdn_attn.py: supports_token_graph; a padded prefill keeps the real rows' metadata, marks them (prefill_real_rows) and gives the pad rows slot -1; image prefills run eagerly (step_needs_eager); a capture builder on the request-less dummy prefill. - base_attention.py, kimi_k3/model.py: on a marked step the MLA region cuts its inputs to the real rows, rebuilding the group-major FP8 q_scale, and both regions hand their output back at bucket height. - atom/utils/token_graph_rows.py (new, no AITER): the cut and the hand-back. - region_graph.py, envs.py, docs: ATOM_TOKEN_GRAPH_EAGER_PAD (default off), which runs a padded step eagerly instead of replaying it. - tests/test_kimi_k3_prefill_token_graphs.py: CPU tests. Off with ATOM_TOKEN_GRAPH_SIZES= (empty). Derived from our earlier piecewise prefill-graph patch, rebuilt on ROCm#2688's region graphs. Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: David Zhu <david.zhu@amd.com>
jiayyu
added a commit
that referenced
this pull request
Nov 4, 2026
…ill traces Review (Copilot, 10-26): - `prepare_mixed` called `prepare_prefill(batch, n_p_seqs)`, but main's #2688 made `running_tokens` a required argument, so every mixed step raised TypeError before attention metadata was built. The prefill half is never padded, so it now passes `running_bs=n_p_seqs, running_tokens=n_p_tokens`. The existing tests checked `prepare_mixed`'s source only; the new `test_prepare_mixed_calls_both_halves_with_their_real_signatures` runs it with both halves stubbed and their arguments bound against the real methods, and fails without this fix. - The schedule-time-advance loop in `postprocess` covers prefill rows only, yet still gated `tracing.prefill_complete` on `i >= total_seqs_num_decode` (from #2644). Prefill rows come first in a mixed batch, so a final chunk could fail that check and leave its prefill span open. Any final chunk now completes it; a pure prefill batch has no decode rows, so flag-off is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
whx-sjtu
added a commit
that referenced
this pull request
Nov 5, 2026
Add `--sequence-parallel-size`, riding the PCP process group as an axis independent of TP: non-attention layers run on T/W tokens, and an exchange around attention trades the sequence for a head slice (GQA KV heads and an indexer tail are replicated to the ranks whose queries use them). MoE gathers the SP shards, runs this rank's expert slice and reduce-scatters the partial sums; large MXFP4 batches route and quantize before the gather. ATOM_SP_MOE_MEGA (default on) runs eligible SP MXFP4 prefill MoE blocks through AITER's MegaMoeTP, with the native norm, gate and top-k in front of it. Ineligible configs, layers or a failed runtime setup keep the native MoE. ATOM_TP_REPLICATED_O_PROJ (experimental, default off) keeps o_proj replicated under TP and exchanges heads instead of all-reducing. Both are generic and opted into per model; MiniMax-M3 is the first user. Ported onto the request/token graph runner (#2688) and the route-before-gather MoE (#2725): SP graphs capture this rank's token shard at a fixed address, gathered SP rows place their experts after the gather and skip the per-rank row count, and SP keeps its own reduce-scatter instead of the PCP-merge all-reduce. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
jiayyu
added a commit
that referenced
this pull request
Nov 5, 2026
…ill traces Review (Copilot, 10-26): - `prepare_mixed` called `prepare_prefill(batch, n_p_seqs)`, but main's #2688 made `running_tokens` a required argument, so every mixed step raised TypeError before attention metadata was built. The prefill half is never padded, so it now passes `running_bs=n_p_seqs, running_tokens=n_p_tokens`. The existing tests checked `prepare_mixed`'s source only; the new `test_prepare_mixed_calls_both_halves_with_their_real_signatures` runs it with both halves stubbed and their arguments bound against the real methods, and fails without this fix. - The schedule-time-advance loop in `postprocess` covers prefill rows only, yet still gated `tracing.prefill_complete` on `i >= total_seqs_num_decode` (from #2644). Prefill rows come first in a mixed batch, so a final chunk could fail that check and leave its prefill span open. Any final chunk now completes it; a pure prefill batch has no decode rows, so flag-off is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
jiayyu
added a commit
that referenced
this pull request
Nov 5, 2026
…ill traces Review (Copilot, 10-26): - `prepare_mixed` called `prepare_prefill(batch, n_p_seqs)`, but main's #2688 made `running_tokens` a required argument, so every mixed step raised TypeError before attention metadata was built. The prefill half is never padded, so it now passes `running_bs=n_p_seqs, running_tokens=n_p_tokens`. The existing tests checked `prepare_mixed`'s source only; the new `test_prepare_mixed_calls_both_halves_with_their_real_signatures` runs it with both halves stubbed and their arguments bound against the real methods, and fails without this fix. - The schedule-time-advance loop in `postprocess` covers prefill rows only, yet still gated `tracing.prefill_complete` on `i >= total_seqs_num_decode` (from #2644). Prefill rows come first in a mixed batch, so a final chunk could fail that check and leave its prefill span open. Any final chunk now completes it; a pure prefill batch has no decode rows, so flag-off is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Two commits on top of main:
atom/models/deepseek_v4/mono/engine.py) is written againstrunner.graphs,runner.graph_poolandrunner.graph_logits, which commit 2 replaces with request-bucket graphs (ReqGraph). It needs to be re-landed on the region graphs, where it can be validated at TP8.Graph regions
TokenGraph(token buckets) andReqGraph(request buckets) inatom/utils/region_graph.pyare the runner's two graph mechanisms; decode, the V4.1 prefill decoder (DecoderGraph) and the DSpark draft record and replay through them. Bucket lookups live on the graphs (TokenGraph.bucket,__contains__)._record_prefill_graphs); capture staging is shared (_build_capture); dummy prefill and decode sharedummy_execution.DP step shape
ForwardMode.decideruns the step's single DP exchange and settles the running tokens per DP rank into a runner-owned numpy buffer in place.DPMetadataholds only that table;DPMetadata.gather_running_tokensis the one all-reduce path (skipped at dp=1).num_tokensnames become scheduled / running tokens.V4.1 decode preparation
PagedAttentionCache.indptrs) and dropped bybegin_forward, like unit tiles: a mono step fills none; a graphed forward records the fill.step_meta._walk_plan);index_plan.fill_walk_planstays as its host reference and a GPU test checks them equal on random steps.WalkPlannerand the installers'metadata_builderargument are removed.prepare_batch_stepfinds group members by name instead of scanning every member twice;CompressPlandropscu_compress_cpu/compress_plan_cpu, which nothing read.Measured on V4.1-Flash, TP4, DSpark 5, 64k ISL / 512 OSL, concurrency 4 (rank 0 trace):
prepare_decodehost timeCleanup
Patch-style branches, duplicate plumbing and dead code removed across the runner, attention builders (shared
Fp4LaunchStorages, V4 / MLA indexer rows), MoE, plugins and tests; docs updated to match.Test plan
--kv_cache_dtype bf16 --index_cache_dtype fp4 --method dspark --num-speculative-tokens 5 --level 0) GSM8K 3-shot: 0.9242 (CI baseline 0.9224, threshold 0.91)ReqGraphand validated at TP8