Factor in AMD compute-partition modes (SPX/DPX/QPX/CPX) as an optimizer lever - #1287
Closed
rpoornac wants to merge 10 commits into
Closed
Factor in AMD compute-partition modes (SPX/DPX/QPX/CPX) as an optimizer lever#1287rpoornac wants to merge 10 commits into
rpoornac wants to merge 10 commits into
Conversation
rpoornac
requested review from
a team,
Ahmedhasssan-aig,
devalshahamd and
tsrikris
as code owners
August 25, 2026 15:44
Adds SPX/DPX/QPX/CPX as a first-class lever so a throughput-sensitive
session can trade single-stream latency for aggregate throughput, and
adds the constraint that keeps that trade honest. Off by default: with no
``--compute-partition-modes`` the hardware is never touched and the KEEP
ladder is unchanged.
Measured on one MI355X over HunyuanWorld-Mirror's 30 example scenes, all
five configurations run through the production ``run_scriptable`` path:
SPX x1 @2 12.24 fwd/s 163 ms (anchor)
DPX x2 @2 14.14 fwd/s 283 ms +15.5%
QPX x4 @2 14.64 fwd/s 543 ms +19.6%
CPX x8 @1 14.66 fwd/s 540 ms +19.8%
CPX x8 @2 OOM, predicted from the per-stream footprint
Throughput saturates long before the hardware does: partitioning past DPX
buys 3.7% for 1.9x the latency. That shape is why the lever ships with a
budget rather than alone -- unconstrained, the search picks CPX and makes
every request three times slower to get there.
* ``common/gpu_partition.py`` is the hardware boundary, and enforces the
two invariants that make a partition measurement trustworthy. A set is
not a set until it reads back: ``amd-smi set`` reports success for a
change that was only staged, and exits 0 on a permission failure, so
every mutation re-reads and compares. And partitions are not identified
by device index: HIP enumerates whole cards first, so on one card of
eight split into DPX the partitions are devices 7 and 8, and selecting
by index measures a full card while labelling it a partition. Callers
get a CU-count predicate instead.
* A card refuses to repartition while a process is resident, and teardown
is not synchronous -- ``docker rm -f`` returns before the runtime has
released the device. The set waits that refusal out, matching on
``AMDSMI_STATUS_BUSY``. Treating it as fatal cost a sweep two
configurations: the restore failed, left a shared card in QPX, and the
next run inherited that as the state to restore to.
* ``_latency_budget.py`` gates KEEP on ``--max-latency-ms`` between the
throughput and accuracy checks. Absolute, because an SLA is; fail-closed
on an unmeasured latency, because a throughput win with no latency
evidence is exactly the case the budget exists to stop.
* ``_partition_lever.py`` carries the mode as ordinary variant env, so
fingerprinting, dedup and the journal need no changes, and holds the
hardware only around the benchmark via a context manager that restores
on both paths. Device enumeration stays with the benchmark script,
which is the only layer that can see the partitions.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
(cherry picked from commit 2900bdb)
… path The rebase brought in ``ea45bf687``, which gave operator-supplied session knobs a three-step shape: restore from archive when the resume omits them, apply (validating) so a re-passed flag wins, then persist the live env back onto state for the resume after. The lever had hand-rolled its own half of that in one expression and got two things wrong that the convention gets right. * The env tier was missing. Resume resolved CLI flag > archive, so a ``HYPERLOOM_MAX_LATENCY_MS`` exported for this resume was overwritten by the older archived value. ``_restore_partition_lever_from_state`` now applies flag > env > archive, writing onto ``args`` so ``_export_partition_lever`` stays the only writer of that env and the flag still gets validated rather than trusted. * Nothing was persisted. A resume that re-passed ``--max-latency-ms`` published the new budget to the executors while the manifest kept the old one, so the session was measured under one contract and recorded under another -- and the next resume restored the stale number. ``_persist_partition_lever`` mirrors it back. * ``--streams-per-partition`` defaulted to 2, which made "not passed" indistinguishable from "passed 2" and would overwrite a persisted 4 on every resume. Now defaults to None, with the 2 applied where the flag is resolved. A parser-level test locks this, because re-adding the 2 reads as harmless -- 2 *is* the documented default -- and silently changes the experiment. Also drops a duplicated pair of env-name constants: ``cli/__init__`` had its own copy of ``PARTITION_MODES_ENV`` / ``STREAMS_PER_PARTITION_ENV`` alongside ``_partition_lever``'s. The names now come from the module that owns them, and ``read_session_lever`` is the one parser of those three variables, shared by the manifest seed and the resume persist. Conflicts in the rebase were all additive and taken whole from both sides, except ``run_scriptable``, where main's candidate-list diagnostic for a missing entrypoint replaced the one-line error the lever's branch still had above its partition planning. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: Cursor <cursoragent@cursor.com> (cherry picked from commit 259c6d3)
…able The mode list was validated only against the four known names, and the partition ladder came from a static table. Neither was checked against the board. ``amd-smi partition -a`` states both, so it is now asked. * ``read_partition_profiles`` reads the card's own profiles: mode, index, ``num_partitions``, XCC instances per partition, and the memory-partition caps each compute mode can pair with. The report is sparse -- a profile's first row names it and carries its XCC count, the rows after it continue the same profile with its DECODER/DMA/JPEG resources under blank identity fields -- so only the named rows are taken. A parser treating every row as a profile would invent four per mode. * The scope is checked at launch. A mode this board does not offer now exits 2 with the card's actual list, instead of reaching ``set_partition_mode`` and failing there: partway through a session, on a privileged mutation of shared hardware. That is the same argument ``parse_modes`` already made for refusing a typo at parse time; the board's real capabilities were simply not available to it before. * The query needs the same elevation as the set -- unprivileged, amd-smi fills every field with "N/A" -- so it degrades to *no answer*, never to *supports nothing*, and an unelevated session is told its request went unvalidated rather than being blocked. "Not validated" and "validated as fine" do not look alike in the log. * ``partition_count_conflicts`` turns ``MODE_PARTITION_COUNTS`` from an assumption into a checked one. That table drives every CU calculation, and partition devices are then found by matching the CU count exactly, so a board whose ladder is not 1/2/4/8 would not disagree loudly -- the benchmark would just find no device of the expected width. * ``layout_for`` refuses a CU count that does not divide evenly instead of flooring it, for the same reason: a floored width matches no device, and the eventual error names the wrong cause. All four boards in the identity table divide evenly today; this is what catches the one that does not. Memory partitioning is still not a lever. The per-profile NPS caps are captured because the card reports them and the pairing is a real constraint -- on MI355X, SPX is NPS1-only while the split modes accept NPS2 -- but nothing acts on them yet, and switching NPS needs a driver reload. Also corrects the DEFAULT_MODE comment, which had that relationship backwards. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: Cursor <cursoragent@cursor.com> (cherry picked from commit f2b58a2)
The lever validated and recorded --compute-partition-modes but nothing turned that list into variants, so the modes were only ever reachable by hand. partition_lever_grid expands the validated list into one env-only variant per mode and seeds them ahead of the grid. Ordering them first is what makes the rest of the round useful: explore already stacks a KEEP'd variant's envs onto everything after it, so the winning mode becomes the topology the remaining knobs are explored inside, and each later mode has to beat the best mode so far rather than the original baseline. The prepend-and-dedupe the framework levers already did is factored out and shared, keeping its payload in its own name because the attribution pass needs it. Scriptable frameworks only. plan_partition_run is called by the scriptable runner and not the serving path, so a serving framework would take the env, change nothing, and file the number under a mode the card was never in -- refused at launch, with the generator declining as a second line. Co-authored-by: Cursor <cursoragent@cursor.com> (cherry picked from commit 69c4731)
…load A narrow mode divides the card's memory while every stream on it keeps a full copy of the weights, so CPX at two streams is the first to run out. Measuring that costs a run and returns an OOM instead of a number: the variant is scored as a failure and the operator learns from a stack trace what arithmetic could have said for free. fits_in_partition already encoded the predicate and was called from nowhere; prune_infeasible_modes now calls it per requested mode and the generator emits what survives. Two inputs had to exist. read_hbm_gib asks the card its capacity, since boards sharing an ISA do not share one -- and answers only in SPX, because amd-smi reports VRAM per device and under a split mode a device is a partition, with nothing in the payload saying which the number is. Dividing an already-divided figure again would understate capacity eightfold and prune every mode that fits, so an ambiguous reading is reported as unknown. The footprint has two sources. peak_gib_per_stream, read from a report that carries it and carried onto current_best from the baseline, is the real one. Absent that, the checkpoint's own weight bytes serve as a lower bound and are used strictly as one: each stream holds its own copy, so "does not fit by the weights alone" is a proof while "fits" is no evidence. Pruning acts only on the former, so it cannot drop a mode that would have won -- and an unknown capacity or unknown footprint drops nothing at all rather than guessing. Every drop is warned with its arithmetic, because a list that quietly comes back shorter than the operator's is how this gets rediscovered as a bug. Co-authored-by: Cursor <cursoragent@cursor.com> (cherry picked from commit a6cee1c)
Every other lever is reachable by the optimizer on its own, so an operator meets it in the results. This one stays off unless a flag names the modes, which means the only people who can use it are the ones who already knew about it. The report is the artefact read after every run, so it is where the lever gets introduced. Advertising is not recommending. Partitioning only ever gives a single stream fewer CUs, so the copy carries the cost as plainly as the 20%: the same measurement took per-request latency from 183 ms to 1211 ms, it needs a privileged amd-smi, and repartitioning evicts every process on the card. An operator who turns this on because a report suggested it and then finds latency quadrupled was misled by that report. Silent where the advice would not apply -- serving frameworks, which refuse the lever at launch, and multi-node sessions, where it manages one card. When the lever did run the section reports what happened instead, including the two states that read as gaps if left blank: no mode kept, which is a measured loss rather than a lever that never ran, and no latency budget, which is how a session ends up on the narrowest mode on offer. Co-authored-by: Cursor <cursoragent@cursor.com> (cherry picked from commit 93b0616)
…mpted The session report already tells an operator which capabilities never ran. The partition lever belongs on that list more than anything else on it: every other capability is reachable by the optimizer on its own, so not_attempted there means the search declined to go somewhere. This one is off unless a flag names the modes, so not_attempted means nobody offered it -- and an operator who has not heard of the lever cannot discover it from results that never mention it. The reason carries the flag, so the row is actionable rather than a reproach. Absent, not not_attempted, on a framework that cannot apply it. Listing it there would be false in the way that matters: it reads as a missed opportunity when the launch would in fact have been refused, sending the reader to a flag that exits 2. Modes that were offered and lost read as tried, because a measured loss is evidence and filing it under not_attempted would claim the lever never ran. reason was in the documented contract and rendered nowhere, which left a row able to explain itself with no way to say it; the capability table now prints it. No existing row sets it, so nothing else changes. Co-authored-by: Cursor <cursoragent@cursor.com> (cherry picked from commit 406c332)
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
Collaborator
|
@rpoornac can you pls check to see why major changes have been added to the TraceLens files? I'm not sure if aspects related to TraceLens need to be handled in this PR |
tsrikris
reviewed
Aug 25, 2026
Each is a way the lever could mislabel a measurement or mutate something it was never scoped to, reached from a direction that had no check. * A failed restore made every later baseline a silent lie. ``partitioned`` logs a restore failure rather than raising, which is right -- it must not mask the exception that caused the exit -- but the session then carries on with the card left split, and the runs requesting *no* mode are precisely the ones with nothing to notice. They would measure a split card and be recorded as the unpartitioned baseline: the mislabelling this module refuses everywhere else, arrived at from the one side that never looked. ``plan_partition_run`` now reads the mode before a mode-less run and refuses a split card. Silent while the lever is off, because then nothing here has touched the hardware and a split card is the operator's own arrangement; silent too on an unreadable mode, since "cannot tell" is not evidence of a problem. * Launch validation checked card 0 while the session mutated another. ``supported_modes``, ``unsupported_modes`` and ``partition_count_conflicts`` were all asked about GPU 0, but the apply path manages ``partition_gpu_id()``. On a heterogeneous node that validated a different board than the one that gets repartitioned -- defeating the whole reason this check happens at launch rather than at the privileged mode change. * Multi-node was silent in the report but ungated in the lever. The report skips ``multi_node_session``, on the grounds that the lever manages one card; the launch path refused serving frameworks for the same class of reason but never refused a cluster. So a multi-node scriptable session could repartition one node's GPU, keep the mode, and have nothing in the report mention that a privileged change happened at all. Refused at launch now, beside the serving-framework check. Also hardens the test file's own lever-env fixture. ``_export_partition_lever`` writes ``os.environ`` directly, which ``monkeypatch`` cannot undo on its behalf, so the first test to complete an export leaked a mode list and a latency budget into every test that ran after it -- which reaches the KEEP gate, and turned 15 unrelated explore cases red. The fixture now saves and restores the five variables outright. Co-authored-by: Cursor <cursoragent@cursor.com>
tsrikris
requested changes
Aug 25, 2026
Collaborator
There was a problem hiding this comment.
Recommend reverting these changes
rpoornac
force-pushed
the
feat/gpu-partition-modes
branch
from
August 25, 2026 17:17
406c332 to
9f06ad1
Compare
``ruff format --check .`` is a hard gate in lint.yml alongside ``ruff check``, and three files on this branch fail it. All three are line joins the configured line length allows, so nothing about the code changes. Two of them predate the review pass and would have failed CI on the original branch too: the ``layout_for`` odd-CU case in ``test_gpu_partition.py`` and the over-budget reason string in ``_latency_budget.py``. The third is the ``unsupported_modes`` error string this branch just rewrote to name the GPU. Worth recording why they were missed: the branch was verified with ``pytest`` and with ``ruff check`` on the changed files, and neither sees a formatting difference. Only ``ruff format`` does, and only the repo-wide invocation CI actually runs. Co-authored-by: Cursor <cursoragent@cursor.com>
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Three findings, all on files this branch adds or edits. The earlier ones the bots reported were against ``prompt_builder.py``, ``codex_session.py`` and the TraceLens tools -- files this PR never meant to touch, which came from the stale lineage the branch was rebased off. These are the ones that are actually ours. * **Empty except** in ``per_stream_footprint_gib``. The ``pass`` was deliberate but said so nowhere, which is the same finding the repo already answered with a comment in ``_steal_stale_claim``. A report carrying ``peak_gib_per_stream`` as a non-number is treated as not carrying it, which is what the caller's "prunes nothing when the footprint is unknown" contract already promises. * **Implicit string concatenation in a list**, five times over, in the not-exercised branch of the partition report section. Adjacent literals in a list display are ambiguous with a forgotten comma -- a real class of bug, and here it would have split one markdown bullet into two. Each paragraph is now bound to a name before the list. The rendered section is byte-for-byte identical, checked by diffing the output against the previous commit's. * **Unreachable code**, twice, in the ``partitioned`` restore tests. A false positive with a real cost: CodeQL does not model ``pytest.raises`` as suppressing, so an inline ``raise`` as the block's last statement makes everything after it read as dead -- and what follows is the entire point of those two cases, which is what the context manager did on its way out. A ``_raise`` helper keeps the flow analysable without weakening either assertion. Not changed: the same concatenation pattern at ``explore.py:593`` is main's code, not this branch's, and the one in ``_latency_budget.py`` is already parenthesized, which is why the rule does not fire on it. Co-authored-by: Cursor <cursoragent@cursor.com>
CI E2E report — ❌ Failed
|
Collaborator
Author
|
Closing this in favour of #1288 - same work, but based on the live @tsrikris your question was the right one to ask, and the answer is that the TraceLens files were never part of this change. This branch was cut from a local ref for Rebasing the seven feature commits onto the live The bot findings split the same way. The ones against The rebase also surfaced that Opening fresh rather than force-pushing this one so the review starts from the actual diff instead of 20 comment threads anchored to commits that no longer exist. |
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
6 participants
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
Makes AMD compute partitioning a first-class optimizer lever instead of something an operator does by hand around the loop.
--compute-partition-modes spx,dpx,qpx,cpxexpands into one env-only explore variant per mode, seeded ahead of the grid. Explore already stacks a KEEP'd variant's envs onto everything after it, so the winning mode becomes the topology the remaining knobs are explored inside, and each later mode has to beat the best mode so far rather than the original baseline.read_hbm_gib()asks the card its capacity andprune_infeasible_modes()drops what cannot fit, reporting the arithmetic for each drop.reports/final.md(with the flag, the privilege requirement, and the latency cost) and as acompute_partitionrow in the session report's Capabilities not attempted.Where it deliberately refuses
plan_partition_runis called by the scriptable runner and not the serving path, so a serving framework would take the env, change nothing, and file the number under a mode the card was never in. Refused at launch, with the generator declining as a second line.amd-smi setreports success for a change that has only been staged, so every mutation re-reads the mode and compares.amd-smireports VRAM per device, and under a split mode a device is a partition, with nothing in the payload saying which. Dividing an already-divided figure again would understate capacity eightfold and prune every mode that fits.Known limitation
No in-tree harness reports peak memory yet. The measured path is plumbed end to end (
peak_gib_per_stream-> measurement ->current_best) and prunes correctly the moment a report carries the number, but on a stock report the axis prunes on weight bytes alone. That bound is sound but loose for activation-heavy models: a 1.26B diffusion model carries ~2.5 GiB of weights against a ~20.7 GiB real footprint. Closing it is a harness-side change.Test plan
test_external_multi_node.pyare pre-existing and unrelated - confirmed identical on a cleanHEADworktree.read_hbm_gib()verified against a live MI355X: reads 287.98 GiB, andlayout_fornow reports per-partition memory (CPX (8 x 32 CU, 36 GiB)).