[jinyouzhi]

Stacked on #9010 (only the last two commits are new here; I will rebase after #9010 merges).

_validate_holders compared every piece with every other piece, _sources compared every source region with every other region, and _shards_tile compared each rank's pieces pairwise by shard address. Planning time therefore grew quadratically with the number of pieces, and a map with 1024 pieces already exceeded the default max_pair_checks.

When no rank repeats a box and distinct boxes do not overlap, a rank covers a box exactly when it holds that very box, so the holders of each box can be read from a table. This PR adds boxes_disjoint to affine_geometry. It sweeps along the axis whose box origins vary most and compares each box only with the boxes still open where it starts. The same sweep proves that a rank's pieces tile its shard when they are disjoint by address and fill it (suggested by 0z5a in review). Maps that fail the fast checks fall back to the original pairwise walk, which names the offending piece, so error messages are unchanged, and maps that pass no longer hit the per-rank len(held)**2 budget.

The sweep counts its comparisons against max_pair_checks too (suggested by Achyuthan-S in review). A disjoint layout can keep every box open for the whole sweep: for example, h full-height one-column boxes plus 2h one-element boxes beside them. When the sweep runs out of budget it reports "not shown disjoint". It compares at most n·(n−1)/2 pairs, so running out implies n² > max_pair_checks, and the pairwise fallback refuses at once. With 4,095 such pieces, planning went from 19.2 s to a budget refusal in 0.6 s.

Planning time (8 ranks with interleaved rows, re-sharded to 8 row shards, 64 columns; ms, single CPU process):

pieces	before	holders/sources only	+ shard addresses (this PR)
256	211	15	7
512	827	46	15
1024	3278 (refused by default budget)	156	33
2048	12926	569	76
4096	50469	2227	235

Tests:

• Randomized oracle over 200 seeds: plans and errors match the pairwise walk, including replicas split at different points, overlapping pieces and wrong holder lists.
• boxes_disjoint is checked against an all-pairs reference over 100 seeds.
• Mutation check: the tests catch dropped intersections and wrong holder sets.
• A rank holding 1024 pieces plans under the default max_pair_checks (refused before this change).
• The always-open layout above plans correctly under the default budget and is refused with max_pair_checks just below its sweep cost. Without the sweep budget, that test plans instead of refusing.
• tests/unit/checkpoint -k affine passes on 8× NVIDIA L20 (645 passed). At 12663c3 on CPU (DS_ACCELERATOR=cpu, 4 ranks): 642 passed, 5 skipped.

The remaining near-quadratic steps, the source × target join and _covers_union, overlap with #8936 and will follow after it lands; segment merging in _coalesce has its own max_merge_checks budget and will follow separately.

[0z5a]

#9010 and #9011 carry the same code as the fork branches I read on Nov 1; range-diff shows only the trailer dropped. So my reading there stands: #9010 is a pure move, and #9011's table check accepts exactly what the pairwise walk accepts. On CPU, the affine tests pass with each PR merged into current master (290 and 591 tests). They also pass with #9011 merged with #8936 (656), where the only conflict is the typing import line in affine_transfer.py.

One addition to the remaining-work list: _shards_tile still compares each rank's pieces pairwise. On the interleaved-rows layout it is now most of the planning time. boxes_disjoint applies to it directly, with the same fallback, so error messages stay the same:

    for rank, held in sorted(by_rank.items()):
        if boxes_disjoint([frame.local for frame in held]) and sum(frame.local.volume
                                                                   for frame in held) == held[0].shard_numel:
            continue
        # ... existing budget check and pairwise walk ...

I tested 8 ranks with interleaved rows, re-sharded to 8 row shards with 64 columns. Planning drops from 1920 ms to 190 ms at 4096 pieces, and from 148 ms to 29 ms at 1024. The affine tests still pass (581). Maps that pass also no longer hit the per-rank len(held)**2 budget. It could go in here or with the follow-up, whichever you prefer.

[jinyouzhi]

Thanks for the review and for the _shards_tile suggestion. I added it here (bea2e09) with the same fallback, so error messages are unchanged and maps that pass no longer hit the per-rank len(held)**2 budget. A new test plans a map with 1024 pieces on one rank under the default max_pair_checks; it was refused before. Re-measured on the same layout, 64 columns: 1024 pieces 160 → 33 ms, 4096 pieces 2313 → 234 ms. The table in the description now separates the holders/sources step from this one.

Both #9010 and this PR are rebased onto current master. tests/unit/checkpoint -k affine passes on 8× NVIDIA L20 (343 and 645). #9010 is now ready for review; this one stays a draft until #9010 lands.

While looking for other pairwise checks, I found one in ParamAffineMap._validate_destinations from #8922. It compares every strided box of a rank with every other one, with no budget, and validate() runs on every map build and conversion. With interleaved columns, 2048 pieces per rank took 6.7 s. #9039 (stacked on this PR) applies the same sweep there: 68 ms, with errors unchanged. The source × target join and _covers_union remain for after #8936.

[Achyuthan-S]

Thanks @jinyouzhi. I checked #9011 the same way as #9039, running plan_transfer under master (b9100c5) and under this PR on 20,000 random (target, source) pairs: row splits, column splits, replicated and ZeRO-3 layouts of 2-D parameters, where each map has an even chance of one changed piece (a dropped or extra holder, a source or destination off by one, or a duplicated piece).

• With the default budget, every pair gave the same plan or the same error.
• With max_pair_checks=40, the only differences were 265 pairs that master refused for the holder-check budget. This PR planned 176 of them, refused 72 with the specific defect, and hit a later budget check on 17. All 1,695 plans it made at that budget reproduce extract exactly.
• The affine test files pass on CPU with 4 ranks (640 passed).

One observation, not blocking: boxes_disjoint has no budget of its own, so its worst case is still quadratic. A disjoint layout of h one-column boxes spanning every row, plus h one-row boxes beside them, keeps every column box open for the whole sweep: 4,096 boxes take about 4.6 s on my laptop, against about 2 ms for a column layout of the same size. No builder emits anything like it, so real maps are fine, and max_pieces_per_map bounds it. But PlanBudget says a map whose piece count is itself the bug cannot force an unbounded enumeration, and master refused a map that size at once. Counting the sweep's comparisons against max_pair_checks would keep that guarantee.

LGTM otherwise.

[jinyouzhi]

Thanks @Achyuthan-S, good catch. 12663c3 gives boxes_disjoint an optional max_checks, and the planner passes max_pair_checks at all three call sites. Once the sweep exceeds the budget, it reports "not shown disjoint". Because it compares at most n·(n−1)/2 pairs, that implies n² > max_pair_checks, so the existing pairwise fallback refuses immediately with the usual budget error. Your layout (h one-column boxes plus 2h one-row boxes beside them, 4,095 pieces) now refuses in 0.6 s instead of 19 s here. A new test covers it: it plans correctly under the default budget and is refused just below its sweep cost. #9039 is rebased on top. ParamAffineMap.validate() takes no PlanBudget and keeps its unbudgeted pairwise fallback, the same as on master.

[chatgpt-codex-connector]

Codex Review Summary: ✅ Completed

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-11-12T17:11:39.048047Z	c91a7ba	Draft marked ready

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: c91a7ba441

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[Achyuthan-S]

Thanks @jinyouzhi, 82b30c7 does it: boxes_disjoint now takes max_checks, all three planner call sites pass max_pair_checks, and a layout that keeps every box open runs into the budget instead of an unbounded sweep. LGTM.

On the Codex note: I'd keep the pairwise walk as the oracle, since it is what master checks, but compare the plans and the exception types rather than the message text, and check that every accepted plan reproduces extract. Then a later change to a message doesn't fail the test, while the equivalence with master stays pinned (the same point as the monkeypatch on #9039).

[jinyouzhi]

@delock ready to review, Thanks!