fix(kb): guard warm-start replay on the shape the canonical_id cannot express

[rpoornac]

The gap

The recipe canonical_id is a seven-tuple:

inference:{model}:{hardware}:{framework_name}:{model_type}:{architectures}:{framework_version}:{precision}

Expert parallelism is not a dimension of it, and neither is the compute-partition mode. Before this change kb_hardware_slug appended topology only when nodes >= 2, so on a single node it returned the bare GPU type. A run in SPX and a run in CPX produced the same id — and recipe_canonical_id has no partition parameter, so the difference was not merely unrecorded, it was inexpressible.

Nothing downstream caught it either. The warm-start cascade relaxes only conc/isl/osl, and RecipeScope.matches_workload_shape compares a hardcoded (tp, conc, isl, osl). So an exact tier hit at confidence 1.0 could hand the PRELUDE auto-replay a config recorded with eight times the partitions. --warm-replay-min-reproduce-pct is a real backstop and the run does recover, but only after paying for the verify round — and status=drift reports it as a recipe that failed to reproduce rather than as a machine that never matched.

What this changes

The approach is to make the collision unrepresentable rather than detectable. An earlier revision of this PR compared shapes at read time and demoted disagreeing rows; that was blocked on mechanism and removed — see Why identity below.

The shape is part of the hardware identity. kb_hardware_slug takes a partition_mode and encodes it, plus expert parallelism, at any node count:

MI300X_cpx                     <- CPX pod
MI300X_dpx                     <- DPX pod
MI300X_ep8                     <- single node, EP=8
MI300X_ws16_tp8_ep8_cpx        <- 2 nodes, EP=8, CPX

The default shape is no shape. SPX and ep <= 1 append nothing, so every key written before this change is byte-identical after it:

MI300X                         <- whole card, dense (the historical key)
MI300X                         <- no mode published at all
MI300X                         <- EP=1 is dense, not a shape

One resolver feeds both sides. resolve_kb_topology now returns partition_mode, read from the state's compute_partition["mode"] and falling back to published_shape(). The KB reader and the writer derive the key from that one call, so they cannot drift into agreeing only by coincidence.

The fallback tiers carry the shape. _hardware_fallback_values preserves the topology suffix, so a partitioned pod reaches same-ISA siblings at its own shape and never the whole-card row at any tier:

>>> _hardware_fallback_values("MI300X_cpx")
['mi300x_cpx', 'mi308x_cpx', 'mi325x_cpx']
>>> _hardware_is_compatible("MI300X_cpx", "MI300X")
False

_TOPOLOGY_SUFFIX_RE is extended so a shaped single-node slug still parses, without admitting arbitrary suffixes:

MI300X_cpx                     -> '_cpx'
MI300X_ep8_cpx                 -> '_ep8_cpx'
MI300X_ws16_tp8_ep8_cpx_sglang -> '_ws16_tp8_ep8_cpx_sglang'
MI300X_garbage                 -> ''

The read-side guard is deleted, along with _shape_mismatch, _GUARDED_SHAPE_KEYS, the strict_shape parameter, and --recipe-kb-strict-fingerprint — a flag that was declared in the parser and read nowhere.

Why identity, not a read-side check

@ZhengGong-amd blocked the earlier revision, and all three defects reproduced:

1. The comparison read a dict that cannot hold the keys. _with_exact_history builds a fixed ten-key mapping with no tp, ep or partitions in it, so the check reported a mismatch for dimensions the dict is structurally incapable of carrying: 'tp: row=None pod=8, ep: row=None pod=8, partitions: row=None pod=8'.
2. Demotion did not deny the replay. Setting warm_tier/warm_conf to seed_only/0.0 leaves the config-donor block reachable, because it is entered on warm_point being truthy. _donor_is_trustworthy then compares only conc/isl/osl and never looks at the demoted dimensions, so the config replayed anyway — a guard that logged a refusal and handed over the config.
3. DEFAULT_EP = 1 demoted the whole corpus. Every dense CLI run has state.ep == 1, so workload_shape published ep: 1 on every run and the check reported 'ep: row=None pod=1' against every historical row.

The first and third are consequences of comparing at read time, not implementation slips. Putting the distinction in the key removes the comparison altogether.

It also closes @meinali-566's point without a separate change: RecipeScope had no way to filter on ep or partitions, and distinct canonical_ids mean there is nothing left to filter.

Blast radius

Against a historical row that recorded no shape at all:

Pod	Outcome
dense, no partition env	kept — key byte-identical
dense, SPX published	kept — byte-identical
dense or MoE with ep=1	kept — ep=1 is not a shape
DPX published	not reached (pod key MI300X_dpx)
CPX published	not reached (MI300X_cpx)
MoE, EP=8	not reached (MI300X_ep8)

A partitioned pod now misses rather than mis-hits. It falls through the normal cascade and, absent a row recorded at its shape, seeds cold — the honest outcome for a corpus that predates the distinction, and it self-heals as sessions republish under the shaped key. The previous revision's table promised demotions here; that was the mechanism that got removed.

Scope

No value{} schema change and no allowlist change: workload_shape is an existing published key, sanitize_shared_knowledge filters secrets and host paths rather than structural keys, and RecipeScope.as_dict() — the Store query payload — is untouched.

workload_shape still publishes ep and partitions for disclosure, now omitting a dense ep and preferring the partition count the run actually published over one re-derived from the mode. Nothing compares those fields; they are there to be read by a human and by the prompt.

Verification

• 14 new tests: slug encoding for divided vs. whole cards, single-node EP, SPX treated as no shape; resolve_kb_topology carrying the mode, preferring a persisted mode on resume, and leaving it unset when nobody published one; a partitioned pod not reaching a whole-card row, a whole-card pod still reaching the row it recorded, SPX staying identical to an unrecorded mode; workload_shape publication, SPX and dense-ep omission, published partition count winning; and projection completeness against the publisher.
• 264 passed across test_recipe_kb_t0_anchor, test_canonical_id_5tuple, test_multi_node_scripts and test_remote_recipe_v2.
• CI green: 30 checks passing, 0 failing.
• ruff check and ruff format --check clean.
• +412/−27 across 10 files: +260 tests, +49 CHANGELOG, and +103/−24 of business code.

Follow-up

rocm and aiter are still written into every row's stack_fingerprint and compared nowhere at read time. That is a real gap and it is tracked in #1507, with the prerequisite recording fix in #1508. It is deliberately out of scope here: ROCm and AITER cannot go into the identity on these terms, because every patch bump would mint a fresh canonical_id and no row would ever be reusable. Defect 2 above is the specific trap that design has to avoid.

[ZhengGong-amd]

Blocking: the shape belongs in the identity, not in a read-side guard

The diagnosis is right and worth fixing — SPX and CPX really do share one canonical_id. The mechanism is what I'm blocking on: three of the guard's behaviours are wrong on the main paths today (all green under the new tests), and the rule it implements already exists in this repo.

Required change

kb_hardware_slug already states this PR's argument verbatim:

_tp{tp} / _ep{ep} — the parallel formation. These are fixed at launch (not explored), so a best_config tuned at one split (e.g. TP4 vs TP2/DP2 over the same world size) is invalid at another […] ep==1 (no expert parallelism) is omitted so dense keys stay clean.

It is gated to nodes >= 2 only to keep single-node keys byte-identical — which is the same "omit the default so historical rows don't move" rule you applied to SPX. Apply it in that one place:

1. In kb_hardware_slug, let single-node participate in suffixing, but only for a non-default shape: SPX or unpublished mode → no suffix; ep <= 1 → no suffix; DPX/QPX/CPX → append the mode. Have resolve_kb_topology read the mode from state.compute_partition / HYPERLOOM_PARTITION_*.
2. Delete _shape_mismatch, _GUARDED_SHAPE_KEYS, the strict_shape parameter, the cli/kb.py passthrough, and --recipe-kb-strict-fingerprint itself (record the removal in the CHANGELOG).
3. Keep the workload_shape publication and the SHAPE_KEYS derivation in knowledge_to_warm_recipe — that part is right. Two fixes: omit ep <= 1 for the same reason you omit SPX, and prefer state.compute_partition["partitions"] (already populated by published_shape() from HYPERLOOM_PARTITION_COUNT) over re-deriving the count from the mode table.

What that buys, and why it is the only shape that closes the issue: a CPX pod cannot read an SPX row because it is a different canonical_id — no flag, no demotion, no "a dimension neither side recorded" special case. Historical keys stay byte-identical, so nothing demotes corpus-wide. Read and write derive the key from one call, so the row a guard checked and the row that gets replayed cannot be different rows. And _hardware_is_compatible already requires the topology suffix to match exactly, so _find_config_donor is covered by construction rather than by a second check.

Please leave TP out of this PR: encoding tp single-node would re-key the existing corpus. Partition mode plus ep > 1 already covers the SPX/CPX case your description is built on, at zero migration cost. If TP needs guarding, add a key to the _shape_matches tuple already inside _donor_is_trustworthy in a separate PR, so the self and borrowed paths share one gate.

Why the guard cannot be patched into correctness

Three defects, each on a main path, each invisible to the six new tests.

1. _shape_mismatch reads a dict that structurally cannot hold shape keys. When a row carries exact_history, the comparison reads exact_history — but _with_exact_history builds a fixed 10-key dict (canonical_id, view, the five experiential lists, sessions, the two gain fields). No tp, no ep, no partitions, ever:

class S:
    tp = 8; ep = 8; conc = 64; isl = 1024; osl = 256
    baseline_workload_extra = {}
    compute_partition = {"mode": "CPX"}

row = {"tp": 8, "ep": 8, "partitions": 8,
       "exact_history": {"canonical_id": "x", "lessons": [], "pitfalls": []}}
_shape_mismatch(S(), row)
# 'tp: row=None pod=8, ep: row=None pod=8, partitions: row=None pod=8'

_shape_mismatch(S(), {k: v for k, v in row.items() if k != "exact_history"})
# ''

Same row, same pod, exact match — reported as three mismatches purely because history was attached. This is the common path, not an edge case: in local mode T0 writes its own anchor row before reading, that row is non-actionable on a first run, so _cascade_warm_start_search sets exact_history and every borrowed hit goes through this branch. The semantics are also inverted — exact_history supplies the priors, the top level supplies the config you intend to deny.

2. The demotion does not deny the replay. It only sets warm_tier/warm_conf; the config-donor block right below still runs (warm_point is truthy), and _find_config_donor returns a row gated only by _donor_is_trustworthy, which compares conc/isl/osl and never looks at tp/ep/partitions. Its confidence becomes recommended_replay.config_confidence directly:

warm-start row shape does not match this pod (partitions: row=8 pod=None); demoting exact to seed_only
tier      : seed_only
confidence: 0.0
replay    : {'extra_server_args': '--also-from-a-cpx-pod', ...,
             'config_tier': 'model_family', 'config_confidence': 0.9, ...}

0.9 clears the 0.7 default, so PRELUDE replays a config from a CPX pod anyway. The guard fires, logs, and hands the replay from a row it checked to a row it did not.

3. DEFAULT_EP = 1 invalidates the blast-radius table. --ep defaults to DEFAULT_EP = 1 via _int_arg("ep", DEFAULT_EP), so state.ep == 1 on every CLI run, and _positive_int(1) publishes it:

workload_shape(dense_run_state)
# {'tp': 8, 'ep': 1, 'conc': 64, 'isl': 1024, 'osl': 256}

_shape_mismatch(dense_run_state, {"tp": 8, "conc": 64, "isl": 1024, "osl": 256})
# 'ep: row=None pod=1'

So "ep is absent unless the run actually sets it" and the table's first row (dense, no partition env → kept) don't hold: under the flag, every row written before this PR demotes on every run. It's the false disagreement you carefully avoided for SPX, and kb_hardware_slug already documents the rule that prevents it.

All three are the same failure mode. The guard compares rows across publishers that disagree by construction: workload_shape, _build_t0_trace_extras and _collect_workload_tags each restate the shape key list, and they differ on partitions (only the first writes it, so it never self-heals on a local row), on pp (only the second writes it, and it is as replay-sensitive as tp/ep), and on the EP env fallback (the last two have it, the first doesn't — another false demotion). Add a dimension and all of it has to be revised again, with the next exact_history-shaped mistake equally invisible. Making the mismatch unrepresentable removes the class instead of the instance.

One note on the flag's original promise

It isn't vacuous. framework_version and precision are identity dimensions, but rocm_version and aiter_commit — both written into every row's stack_fingerprint — are compared nowhere at read time. So the M5 consumer still has a real job; it is just a different one. That argues for deleting this flag and filing that separately, not for keeping the name and swapping the semantics: a flag whose name says fingerprint and whose behaviour is workload shape turns an unimplemented promise into a misnamed one.

[meinali-566]

@rpoornac Thanks for adding the ep / partitions publication and the replay-side safety guard. I think there is one KB Store selection gap worth clarifying.

RecipeScope already partitions the current KB Store write/read path by:

kernel_optimizer, tp, conc, isl, osl

The remote client sends scope.as_dict() to KB Store, validates the returned View’s scope, and separately validates knowledge.workload_shape through RecipeScope.matches_workload_shape().

Therefore, tp is already an exact Store-scope dimension on the current remote path. conc, isl, and osl are also represented in the scope and handled by the T0 relaxation cascade where appropriate.

This PR adds ep and partitions to workload_shape and checks:

("tp", "ep", "partitions")

after a warm row has already been returned. However, RecipeScope.as_dict() and the KB Store scope remain unchanged.

That is a useful safety backstop: it prevents an incompatible config from being replayed by demoting the row to seed_only. However, it does not allow KB Store to select the correct compatible Recipe in the first place.

For example, these two sessions can still share the same current Store scope:

A: tp=8, ep=1, partitions=1
B: tp=8, ep=8, partitions=8

If A is the scoped champion while the current pod matches B, the new guard safely demotes A. But the reader cannot continue to B because KB Store never partitions or queries candidates by ep and partitions.

Therefore, the current implementation solves:

Do not replay an incompatible Recipe.

but not yet:

Select the compatible Recipe from KB Store.

Would you agree that ep and partitions should also become RecipeScope / KB Store scope dimensions, while retaining the strict shape check as defense in depth?

If so, please confirm the intended compatibility behavior for historical rows, especially:

• missing ep
• missing partitions
• omitted SPX / partitions=1

After confirmation, I can make the corresponding KB Store service, schema, and query changes.

I would also prefer to avoid introducing another parallel set of variables or mappings. Ideally, both the scope projection and replay guard should derive from the existing workload_shape() publisher/shared state as a single source of truth. conc, isl, and osl should continue using the existing RecipeScope and cascade behavior rather than being introduced again elsewhere.

[rpoornac]

@ZhengGong-amd Agreed on all three, and I reproduced each one before changing anything rather than taking them on faith. Pushed as 1734757.

The three defects, confirmed. _shape_mismatch on a row carrying exact_history reported tp: row=None pod=8, ep: row=None pod=8, partitions: row=None pod=8 for a row that matched the pod exactly, and '' for the same row with exact_history stripped — _with_exact_history builds a fixed 10-key dict, so the shape keys could never be in the dict I was reading. The demotion also did not deny the replay: it sets warm_tier/warm_conf and returns, and the config-donor block below is still entered because warm_point is truthy, with _donor_is_trustworthy comparing only conc/isl/osl. And DEFAULT_EP = 1 means state.ep == 1 on every CLI run, so workload_shape published {'tp': 8, 'ep': 1, ...} and every pre-PR row produced ep: row=None pod=1 — the mass demotion I had carefully avoided for SPX, reintroduced one dimension over.

What the PR does now. kb_hardware_slug suffixes the partition mode and ep at any node count, so the collision is unrepresentable rather than detected: MI300X vs MI300X_ep8_cpx are different identities and a CPX pod never reaches an SPX row. _shape_mismatch, _GUARDED_SHAPE_KEYS, strict_shape, the cli/kb.py passthrough and --recipe-kb-strict-fingerprint are all gone, with the removal recorded in the CHANGELOG under ### Removed. workload_shape and the SHAPE_KEYS derivation stay, now omitting ep <= 1 and preferring compute_partition["partitions"] over re-deriving from the mode table. TP is left out, for the reason you gave. Verified that the default shape is byte-identical — nodes=1 with tp=8, ep=1, partition_mode="SPX" still returns exactly MI300X, as does an unrecognised mode.

One companion change you did not ask for, because it would otherwise have regressed silently. _TOPOLOGY_SUFFIX_RE required a mandatory leading _ws[1-9]\d*, so a single-node _cpx or _ep8 suffix failed to parse, _parse_hardware_topology returned None, and _hardware_fallback_values stopped offering the same-ISA SKUs for exactly the rows these suffixes were added for. I extended it with an alternation that keeps _ws mandatory for the cluster form, so MI300X_cpx now yields ['mi300x_cpx', 'mi308x_cpx', 'mi325x_cpx'] while MI300X_garbage still fails to parse rather than being mistaken for a backend name. _hardware_is_compatible("MI300X_cpx", "MI300X") is False, which is the original bug closed at the identity level.

On the flag's original promise: you are right that it is not vacuous, and rocm_version / aiter_commit being written into every stack_fingerprint and compared nowhere at read time is a real gap. I have noted it in the CHANGELOG as tracked separately and will file it, rather than keeping a flag whose name says fingerprint and whose behaviour is workload shape.

Full suite on the new base: 22113 passed, 9 failed — all nine reproduce on a clean main worktree (test_external_multi_node, test_agentx_repair, test_preflight_auth_override, and two load-flaky test_supervisor cases), none in the KB or recipe paths. ruff check and ruff format --check clean.

[rpoornac]

@meinali-566 Thanks — the selection gap you describe was real, and your example is exactly the right one. The answer changed with the rework I just pushed (1734757), so let me answer against the new shape rather than the one you reviewed.

Short answer: no new RecipeScope dimensions are needed for this, because the two sessions in your example no longer share an identity. The mode and ep now suffix the hardware slug in kb_hardware_slug, so your A and B are:

A: tp=8, ep=1, partitions=1  ->  inference:{model}:mi300x:...
B: tp=8, ep=8, partitions=8  ->  inference:{model}:mi300x_ep8_cpx:...

A pod matching B queries B's canonical_id directly, so KB Store selects the compatible Recipe by identity — there is no "cannot continue to B", because B was never behind A. This also means the strict shape check is deleted rather than retained as defense in depth: with the mismatch unrepresentable, a read-side comparison has nothing left to refuse, and keeping one would mean a second place that has to agree about the shape. That was the substance of ZhengGong's block, and I think it dissolves your gap as a side effect rather than deferring it.

Your three compatibility questions, explicitly. All three cases are deliberately indistinguishable from each other and all three keep their existing key:

historical row	treated as	key
missing ep	dense (ep <= 1)	unchanged
missing partitions	whole card	unchanged
omitted SPX / partitions=1	whole card	unchanged

A mode this build does not recognise also falls in that bucket, so an unknown string cannot invent a key. I verified the byte-identical property directly: kb_hardware_slug("MI300X", nodes=1, tp=8, ep=1, partition_mode="SPX") is exactly "MI300X".

One honest limitation, since you asked specifically about historical rows. This prevents future collisions; it does not retroactively re-file rows already written under a colliding key. A single-node ep=8 row recorded before this change still sits under the bare MI300X key, so a dense pod can still reach it. Re-keying it would require a migration pass over the corpus, which I have deliberately kept out of this PR.

On your last point, we agree, and that is how it is built: workload_shape() remains the single publisher, SHAPE_KEYS is derived from it so the warm-row projection cannot drop a dimension the publisher emits, and the key derives from one resolve_kb_topology() call that both the reader and the writer go through. conc/isl/osl keep the existing RecipeScope and cascade behaviour untouched — nothing about them moved.

So I do not think there is KB Store service, schema or query work needed for this issue. If you would still like ep/partitions as first-class Store scope dimensions for a separate reason — querying or reporting across formations, say — that seems worth doing on its own terms, and workload_shape() is the right source to project from.

[rpoornac]

Housekeeping: I've rewritten the PR description, which was still describing the previous revision's read-side guard — the --recipe-kb-strict-fingerprint flag, the seed_only/0.0 demotion, and a blast-radius table promising demotions. None of that is in the branch anymore. The description now matches what the code does (shape in the hardware identity, so a partitioned pod misses rather than mis-hits), records the three defects and why read-time comparison was the wrong mechanism, and corrects the verification numbers: 14 new tests and +103/−24 of business code, not 6 and ~106.

No code changed — head is still 1734757 and CI is unchanged at 30 green. Apologies to anyone who read the old text from an email notification.

The out-of-scope rocm/aiter gap now has a number: #1507, with the recording prerequisite in #1508.

[meinali-566]

@rpoornac I re-reviewed the latest revision (1734757) and agree with this approach.

And no other block issues are found.

LGTM cc @ZhengGong-amd

[rpoornac]

Re-landed as #1542, at @ZhengGong-amd's recommendation, now that 1.1.1 is cut.

Restored by reverting the revert onto current main rather than rebuilding by hand, so this is the reviewed content: eight of the nine source files are byte-identical to 31143ab. parser.py differs only because main has since refactored --max-hours into DEFAULT_MAX_HOURS, and against current main the only change to that file is still the --recipe-kb-strict-fingerprint removal. The changelog entries moved into the [Unreleased] section the release opened, and the "tracked separately" line now names #1507.