Skip to content

CAP-1: trace the published model at TP1 and TP2, and pin where it specialises #142

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
jgong5 merged 3 commits into from
Sep 23, 2026
Merged

CAP-1: trace the published model at TP1 and TP2, and pin where it specialises #142

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ae8b43a
compass: trace the published model at TP1 and TP2, and pin where it s…
Sep 22, 2026
9fcd6c7
compass: make the pin fail for a repair inside the buffer's construct…
Sep 22, 2026
24d7574
compass: say where the third site is, and drop the doc labels from th…
Sep 22, 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
14 changes: 13 additions & 1 deletion atom/compass/design/04_model_capture_and_cost_ir.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -198,9 +198,21 @@ fake-less aiter ops sit *inside* those leaves and are never reached.
> this section was taken with the fake-tensor capture module that PR #10 added under
> `atom/compass/capture/`, driven by scripts that were never committed and writing JSON
> records that no longer exist. That module and its tests have been withdrawn from the
> tree, so **none of the numbers below is reproducible here**. They are kept because they
> tree, so **none of the numbers below was reproducible here**. They are kept because they
> are the specification a replacement capture is written from, not because they can be
> re-run; anything that builds on them re-takes them first.
>
> **Partly re-taken since, by `tests/compass/test_capture_real_model.py`.** That test
> traces the 27B at both widths through a group of the honest width, and agrees with the
> 27B row on what it is a claim about rather than on its totals: at TP2 it records
> `aiter.all_reduce_` **129** — 128 row-parallel at `communication_op.py:58` plus the
> vocab-parallel one at `embed_head.py:175`, the 128 predicted from the config's layer
> types and not read off the inventory — one functional all-gather at
> `embed_head.py:257` carrying `[2, 124160]` to `[4, 124160]`, and one broadcast. The raw

@jgong5 jgong5 Sep 22, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Finding 5 (principle 8) — a shape nothing recorded, and a missing configuration.

one functional all-gather at embed_head.py:257 carrying [2, 124160] to [4, 124160]

Two problems:

  1. _Recorder stores "shapes": shapes_in only, and the test asserts gathered == [["2", "124160"]] and stops. [4, 124160] is in no record and no assertion — it is inference from the width, not a measurement. I confirmed the input shape reproduces exactly; the output shape is simply not captured.
  2. [4, 124160] is the functional substitute's arrangement, not ATOM's. ATOM hands all_gather_into_tensor an output buffer of (world_size,) + input_size = [2, 2, 124160], which _functional_collectives' shim gathers into [4, 124160] and then reshapes. Writing the substitute's shape into 04 as what the 27B does at TP2 is a small instance of the class the apply_simulated_tp correction row exists to warn about.

Separately: this paragraph does not record that the TP2 inventory was taken with ATOM_USE_CUSTOM_ALL_GATHER=0, which selects a non-default ATOM path — a real TP2 deployment takes the ca_comm custom gather, which is not what was traced. The test docstring and the PR body both say this; the design record, which is what gets read later, does not. Every number in this paragraph is conditional on it.

@jgong5 jgong5 Sep 22, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Closed — both arrangements are measured now, and the configuration is in the paragraph.

_functional_collectives' shim has ATOM's own output buffer in hand, so it records it: the record's gather_buffers carries input, atom_output and functional_output per call, read off the live call. Measured at TP2:

input              ['2', '124160']
atom_output        ['2', '2', '124160']     <- ATOM's (world_size,) + input_size buffer
functional_output  ['4', '124160']          <- the substitute's concatenation

test_the_two_arrangements_of_the_vocab_gather_are_both_measured asserts that triple, and asserts TP1 records none.

04's paragraph now says: one all-gather at embed_head.py:257 (no inferred destination shape), plus a separate short paragraph giving both arrangements with [4, 124160] labelled as the substitute's, and a statement that an earlier revision gave it as the 27B's and gave it from the width rather than from a record. A third paragraph states the two declared substitutions every number in the section is conditional on — ATOM_USE_CUSTOM_ALL_GATHER=0 as a non-default path (the default takes ca_comm, which is not what was traced), and the legacy-to-functional routing as the source of every _c10d_functional.* entry including both wait_tensors, with only the 129 aiter.all_reduce_ being ATOM's own dispatch.

> Triton traffic agrees exactly: **33 launches across 3 kernels**. The operator totals do
> **not** match and are not expected to: they are 2,521 at TP1 and 2,662 at TP2 against
> 2,471 and 2,611 here, on a decode step of two sequences at a block size and a batch
> budget this file never recorded, which is the reason a total is not the assertion.

The open issue *"whether ATOM's real model classes trace cleanly under this mode at
TP>1"* is **answered yes**, on two models at both widths, with the collectives in the
Expand Down
3 changes: 2 additions & 1 deletion atom/compass/design/12_open_items.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,7 @@ M-f `14`; M-g `01` D3.5.
| T9 | Declare a row-ordering treatment for decode attention |
| **T51** | Enumerate the layer-pattern shapes for Qwen3.8-27B and Kimi-K3; confirm the nested-`Repeat` detector reaches the hierarchical form on both |
| **T52** | Root-cause the `TorchDispatchMode` 8-rank hang at `dspark_scheduler.py:264` — gates T5 |
| **T81** | Make the D18 capture *symbolic* on ATOM's real forward, or record that it cannot be. **Two independent specialisation sites are measured, and only one of them is reachable from the caller.** The capture is concrete — 0 non-numeric shape entries of 13,047 at TP2 — and the tracing mechanism is not the limitation: under the mode a GEMM and a softmax keep their free symbol with `shape_env.replacements` empty. **Site one, the bound.** `prepare_decode` derives one count per staged buffer and uses it to fill the buffer's numpy view before handing it to `copy_to_gpu`; anything that needs an `int` takes `__index__` of a `SymInt` and gets its hint, recording the symbol as a constant with no error and no warning. Measured on ATOM's path as `s56 -> 2` through `aiter_attention.py:1115 in prepare_decode`, and it is not a numpy behaviour — a bare `__index__()` and a plain list slice do the same. Passing the symbol as the bound closes this one. **Site two, the copy.** `copy_to_gpu` is `self.gpu[:n].copy_(self.cpu[:n])`, and `self.cpu` is a real numpy-backed tensor with constant dimensions, so `copy_` solves every symbolic dimension of the destination that the slice does not cover. Measured as `s64 -> 16384` through `aiter_attention.py:1142` → `atom/utils/__init__.py:725 in copy_to_gpu`, where `s64` is `block_tables: ['512', 's64']` — a dimension no bound controls. `A1_tp1.json` reaches `{s64: 16384}` with **no injected SymInt at all**, so this site is not an artefact of the injection attempts. **Consequences.** `CpuGpuBuffer` is *not* unchanged by a repair: it contains site two. Supplying a symbolic bound from `prepare_decode` is therefore **not shown sufficient**, and nothing here shows it is. Of the two repair routes originally recorded, "a symbolic `CpuGpuBuffer`" is the one site two points at, and "a capture entry point below `prepare_inputs`" is **untested** — neither closed nor shown unreachable. No claim is made here that a symbolic capture requires changing ATOM's serving path; that would be inference, and the experiment that would settle it has not been run. **No longer pinned.** The test that held both sites — `tests/compass/test_capture_symbolic_shapes.py`, which also checked that ATOM still shares the one bound — was withdrawn from the tree together with the capture module it exercised. The measurements above stand exactly as taken; what is gone is their reproduction, so **nothing on this tree holds these two sites in place today**, and a change to `prepare_decode` or `copy_to_gpu` would pass unnoticed here. Gates T5 alongside T52. |
| **T81** | Make the D18 capture *symbolic* on ATOM's real forward, or record that it cannot be. **Two independent specialisation sites are measured, and only one of them is reachable from the caller.** The capture is concrete — 0 non-numeric shape entries of 13,047 at TP2 — and the tracing mechanism is not the limitation: under the mode a GEMM and a softmax keep their free symbol with `shape_env.replacements` empty. **Site one, the bound.** `prepare_decode` derives one count per staged buffer and uses it to fill the buffer's numpy view before handing it to `copy_to_gpu`; anything that needs an `int` takes `__index__` of a `SymInt` and gets its hint, recording the symbol as a constant with no error and no warning. Measured on ATOM's path as `s56 -> 2` through `aiter_attention.py:1115 in prepare_decode`, and it is not a numpy behaviour — a bare `__index__()` and a plain list slice do the same. Passing the symbol as the bound closes this one. **Site two, the copy.** `copy_to_gpu` is `self.gpu[:n].copy_(self.cpu[:n])`, and `self.cpu` is a real numpy-backed tensor with constant dimensions, so `copy_` solves every symbolic dimension of the destination that the slice does not cover. Measured as `s64 -> 16384` through `aiter_attention.py:1142` → `atom/utils/__init__.py:725 in copy_to_gpu`, where `s64` is `block_tables: ['512', 's64']` — a dimension no bound controls. `A1_tp1.json` reaches `{s64: 16384}` with **no injected SymInt at all**, so this site is not an artefact of the injection attempts. **Consequences.** `CpuGpuBuffer` is *not* unchanged by a repair: it contains site two. Supplying a symbolic bound from `prepare_decode` is therefore **not shown sufficient**, and nothing here shows it is. Of the two repair routes originally recorded, "a symbolic `CpuGpuBuffer`" is the one site two points at, and "a capture entry point below `prepare_inputs`" is **untested** — neither closed nor shown unreachable. No claim is made here that a symbolic capture requires changing ATOM's serving path; that would be inference, and the experiment that would settle it has not been run. **Pinned again, and re-measured rather than restored.** `tests/compass/test_capture_real_model.py` builds the published Qwen3.8-27B under `FakeTensorMode` through ATOM's own `ModelRunner`, traces one two-sequence decode step, and holds **both sites by value and by the innermost frames they happen through** — `-> 2` at `aiter_attention.py:1115 in prepare_decode`, and `-> 16384` at `aiter_attention.py:1142` into `atom/utils/__init__.py:725 in copy_to_gpu`. The symbol names differ from the ones recorded above, because a symbol is numbered by the order its ShapeEnv created it and that ordering is not a property of either site; everything else is identical, including the 16,384, which is the published config's 262,144 positions over ATOM's default 16-token blocks. The test also asserts the two symbols are distinct and that the first is the bound the caller was handed, which is the independence claim: **a symbolic bound closes site one and leaves site two exactly as it is**, so a repair to one of them fails this test on the other. A change to `prepare_decode` or `copy_to_gpu` no longer passes unnoticed. Gates T5 alongside T52. |

@jgong5 jgong5 Sep 22, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Finding 4 (principle 8) — this sentence states a counterfactual that the PR's own instrument contradicts, and hides a third site.

a symbolic bound closes site one and leaves site two exactly as it is

I ran that counterfactual for the first time, on node 18, by repairing site one alone — aiter_attention.py:1115 → np[: int(batch.total_tokens_num_decode)], line-count preserving so site two's line numbers do not move. Half the sentence holds; half does not:

s27 -> 16384  aiter_attention.py:1142 in prepare_decode  ->  atom/utils/__init__.py:725 in copy_to_gpu
s13 -> 2      atom/utils/forward_context.py:437 in assert_shape_contract  ->  :424 in _rows

Site two is indeed untouched. But the bound is not closed — it specialises fourteen lines later in a module this PR does not name. forward_context.py:424 is return None if t is None else int(t.shape[0]), the _rows helper of assert_shape_contract, reached from slot_rows = _rows("slot_mapping") at line 437: an ATOM assertion helper taking int() of a symbolic dimension.

So "two independent specialisation sites" is an artefact of ordering. There are at least three, and this test pins the two that happen first — which is fine as a pin, but the row should say so rather than assert a closure it has not demonstrated. This is the most useful thing the new instrument has produced and CAP-2 will meet forward_context.py:437 the moment it lands its site-one repair; it belongs in T81, not in a review comment.

@jgong5 jgong5 Sep 22, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Closed, and the counterfactual is now a test rather than a review experiment.

You are right on both counts, and the sentence is gone. T81 now says the sites are an order, not a set, names the third — atom/utils/forward_context.py:437 in assert_shape_contract into :424 in _rows — and records that closing site one relocates the bound rather than resolving it.

Rather than cite your measurement, the test reproduces it. A third capture pass (--repair-site-one) gives each staging buffer's numpy view a numpy.ndarray subclass that reads a SymInt bound's node.hint instead of taking __index__ of it. That is the site-one repair applied from outside ATOM, so no ATOM source is edited to obtain the result. What it records on node 18:

s27 -> 16384  aiter_attention.py:1142 in prepare_decode  ->  atom/utils/__init__.py:725 in copy_to_gpu
s13 -> 2      atom/utils/forward_context.py:437 in assert_shape_contract  ->  :424 in _rows

Identical to yours, including the order. test_closing_site_one_moves_the_bound_to_a_third_site asserts site two is untouched, asserts site three by value and frames, and asserts that the symbol solved at site three is the same bound that was injected — so it is the relocation that is pinned, not just a third location.

(An ndarray subclass and not a delegating wrapper: ATOM's pack_rows takes memoryview() of the staging view, and a wrapper is a bytes-like object is required there.)

The cost is a fourth subprocess capture, ~11.6 s. I judged that worth paying for a claim that otherwise has no source in the tree; say if you disagree.

Neither T81 nor the test claims three is all there are — only that these are the three this instrument reaches.


### Topic 06 — workload harness contract

Expand Down Expand Up @@ -276,5 +276,6 @@ Corrections identified while writing later documents, not yet applied to earlier
| `04` | D18's five `torch.cuda` stubs are not enough to import ATOM on this stack. Construction needs three more (`get_device_properties`, `current_device`, `get_device_capability` — the first is read at *import* by aiter's Triton attention configs) and running `ModelRunner.__init__` needs fourteen more, tagged `needed_for: "model_runner"` by the capture module's `install_runner_stubs` (the import-path set carried `needed_for: "import"`) and pinned by name in that module's tests — **both the module and those tests have since been withdrawn from the tree**, so the counts stand as measured and nothing here re-takes them. `get_device_properties` and `mem_get_info` are *declared readings* in the `03` D14 sense and must be passed in, never read from a host |
| `04` | D18 does not cover raw `@triton.jit` launches. They bypass the dispatcher entirely, so `FakeTensorMode` cannot fake them and the first one reached kills the trace in `triton/backends/amd/driver.py:369`. Any inventory taken with them skipped is a **diagnostic**, not a capture, and must be labelled so wherever it is reported |
| `04` | A TP>1 capture taken through `apply_simulated_tp` at one physical rank both **erases** and **fabricates**: 129 `all_reduce` per forward become the identity and appear nowhere in the inventory, while one `all_gather` becomes six real dispatched ops over a half-zeros tensor. Neither direction is visible in the operator list itself. **Measured since: the substitution is not needed.** At an honest width the same forward completes and records all 129 (`04` D18, *Collectives at TP>1*), where the substituted one refused at 2,477 ops with none. Removing it from the capture path was left as a separate task; the module has since been withdrawn from the tree, so no capture here takes the substituted route |
| `04` | D18 requires `torch.cuda.is_available()` to report **True**, measured on a host whose driver was *wedged*: False hangs inside `_ensureCUDADeviceGuardSet` and True completes. On a host with **no driver at all** the dependency runs the other way, and the same flag is read by two parties that need opposite answers. ATOM needs True, as D18 says. `FakeTensorMode` needs False: that one flag gates `_only_lift_cpu_tensors`, which keeps `torch.tensor` on the host and moves it afterwards — without it ATOM's own `torch.tensor([])` under a CUDA default device is `No HIP GPUs are available`, below anything the mode can intercept — and it gates `_ensureCUDADeviceGuardSet`, which makes CUDA kernels traceable, and the skipping of constant propagation across a device conversion, which otherwise runs the next operator on a small fake **for real** on the destination device. `tests/compass/test_capture_real_model.py` answers both by overriding `FakeTensorMode.avoid_device_init` rather than by choosing one. Two further readings are not `torch.cuda` at all and D18 does not name them: aiter shells out to `rocminfo` at import through `get_gfx_runtime`, which ignores `GPU_ARCHS` and needs `/dev/kfd`, and Triton's active driver asks the live device for its target, after which aiter falls back to a jax import that is not installed. Both are the architecture, and it is configured. |

@jgong5 jgong5 Sep 22, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Finding 6 — the CpuGpuBuffer straddle finding is not in any design document. This is the right table for it and it is not here: the new row covers is_available and the two unnamed device readings and stops.

The PR body calls the straddle CAP-2's starting point, and it survives the squash only in _stage_buffers's docstring. That is better than the PR body, but CAP-2 starts from T81 and from these correction rows, not from a test's internals. Both traps are the kind that cost a day twice: a .numpy() taken while the mode is active leaves the real storage marked not resizable, and converting self.cpu rather than a discarded template memoises the concrete side as symbolic. T81 currently says only that CpuGpuBuffer contains site two.

Add it here or extend T81 — either survives the squash where a reader will find it.

@jgong5 jgong5 Sep 22, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Closed — it is a correction row in this table now.

The new row under 04 states that D18 does not say how CpuGpuBuffer is built under the mode and that it has to straddle it; that .numpy() raising is why no ModelRunner constructs at all until it is dealt with; and both traps by name — a .numpy() taken while the mode is active leaving the real storage not resizable, and converting self.cpu rather than a discarded template memoising the concrete side as symbolic (Trying to resize storage that is not resizable). It also records that this is where T81's site two lives, and that the constructor is executed rather than substituted precisely so a repair here cannot land unnoticed — which is finding 1's closure, stated where a reader of this table will meet it.

T81 keeps its own sentence on what the pin covers, so either entry point reaches it.

| `15` | **D92's decision row still reads *"inherits the TP group"* and *"remainder included"***, and P0.6 measured both false — the group is built in aiter and spans DP, and an indivisible expert count is refused rather than rounded (`15` D92, T65). D92's **conclusion** and D93's formula survive unchanged, so this is a correction to the decision's stated reasons, not to the decision. Rewriting a decision is the owner's call; the body of D92 carries the measurement in the meantime. |
| `04`, `07` | The MoE all-to-all's block-count cap is stated as a single number in `04`'s `exclusive` join-policy text and in `07`'s calibration table; it is `min(128, CU)` for prefill and `min(64, CU)` for decode (**T87**) |
Loading