Module verify_key

Module verify_key 

Source
Expand description

Batched-verify graph key: the canonical depth→slot assignment and the key bytes derived from it. ONE ordering rule, shared by the scheduler that dispatches the batch (mtp_dcut::plan, mtp_step) and the model that builds the CUDA-graph cache key (verify_e2::verify_batched_graph_key).

§The measured defect (nsys + A/B on dgx2, binary b508679e4)

The batched-verify graph cache keys on the per-row (ssm slot, depth k) pairs in batch order, because a capture bakes each row’s pool state addresses AND the depth-run launch structure. D-Cut re-ranks WHICH sequence gets WHICH depth every step, so the key space was the set of ARRANGEMENTS of the step’s depth multiset over the batch’s slots:

n=8, the three multisets actually observed:
  8!/(5!·2!·1!) + 8!/(4!·4!) + 8!/(6!·2!) = 168 + 70 + 28 = 266 keys

against VERIFY_BATCHED_GRAPH_CAP = 32. Measured key counts: n=2 → 2, n=4 → 10, n=8 → 160-253, n=16 → 1 (D-Cut is off above width 8). nsys at C=8: 149 captures in 167 steps (89% of steps), cuGraphInstantiate + cuGraphExecDestroy + cuGraphDestroy = 23.2 ms/step ≈ 20% of the step; GPU busy 96.3% → 77.2%. A/B at C=8: control 78.89 tok/s vs 84.35 with D-Cut off (+6.9%, key count 253 → 1) — and that leg also LOSES the row pruning, so the thrash alone costs more than 6.9%.

§The fix: canonical depth→slot assignment

D-Cut’s ranking chooses HOW MANY drafts survive at each depth (the multiset) — that is where its row saving comes from. It also chooses WHO gets them, and that half is what multiplies the key space. So the multiset stays confidence-chosen and the ARRANGEMENT becomes a pure function of the batch: depths descending are paired with slots ascending. The key is then determined by (slot set, depth multiset) alone — at n=8 the 266 observed arrangements collapse to the 3 multisets that produced them (worst case over all reachable shapes: multisets of size 8 over depths {2,3,4} = C(10,2) = 45, versus 3^8 = 6561 arrangements).

★ The two orderings RECONCILE instead of fighting. The dispatch needs depths descending (equal depths must form contiguous runs — the batched conv+WY fast path launches once per run, trait_decode_batched_conv_gdn_multi.rs) and the SSM batched arms need slots ascending in batch order (ssm_batched_recurrent.rs, decode_step.rs, mtp_step.rs). Under the confidence-chosen arrangement those two demands are in direct conflict: a ragged batch sorted deepest-first scrambles the slot order, so each depth run gets an arbitrary SUBSET of the pool slots and the consecutive-slot precondition fails. Pairing depths-descending with slots-ascending makes the two orders THE SAME order. A depth run owns a consecutive slot block only when the selected pool slots are themselves consecutive; the model checks actual pointers and declines the batched fast path when fragmentation leaves gaps.

Correctness: which sequence gets which depth is a pure PERFORMANCE choice. Every batchable sequence enters the step with exactly ladder_nd drafts (mtp_step truncates the surplus), each assigned depth is in 1..=ladder_nd drafts, and a verify of a shorter draft prefix is the same math on fewer rows. Σ rows is unchanged, so the row budget and chunking are unchanged. What is NOT free to change is the pairing between a batch POSITION and the slot whose pointers the graph baked there — hence one ordering rule, used by both the dispatch and the key.

Kill switch ATLAS_NO_CANONICAL_VERIFY_KEY (PRESENCE — house convention, =0 is NOT off) restores the pre-canonical behaviour: each sequence keeps its own confidence-chosen depth and the batch is sorted deepest-first, ssm-slot second.

§The width gate: it only pays where the key space explodes

Collapsing the key space is not free — forcing the assignment overrides D-Cut’s confidence pairing and re-shapes the depth runs — and the key space only explodes at the TOP of D-Cut’s width range. Measured key counts against VERIFY_BATCHED_GRAPH_CAP = 32: n=2 → 2, n=4 → 10, n=8 → 160-253, n=16 → 1. At n=2 and n=4 there is essentially nothing to collapse, and the A/B says so — see CANONICAL_KEY_MIN_WIDTH, which is the ONE threshold and carries the table. Below it the pre-canonical assignment is restored BYTE-IDENTICALLY; at/above it the canonical assignment applies. canonical_assignment is the single gate; call sites never re-derive it.

Constants§

CANONICAL_KEY_MIN_WIDTH
Batch WIDTH (sequences) at or above which the canonical depth→slot assignment is applied. Below it verify_batch_order / verify_batch_permutation take their canonical = false arm, which is the pre-canonical (pre-PR-#552) behaviour byte for byte: each sequence keeps its own confidence-chosen depth and the batch sorts deepest-first, ssm-slot second, ties on input index (a stable sort_by_key(|(a, k)| (Reverse(k), slot)), exactly what both call sites used before).

Functions§

canonical_assignment
THE GATE — the one decision “does this batch get the canonical depth→slot assignment?”. Both seams ask this and nothing else: mtp_dcut::plan (which decides order AND assignment) and mtp_step (permutation only), each passing the FULL batch width so the two can never disagree — plan gates on batchable.len(), and a chunked dispatch must use that same width, not the chunk’s.
canonical_key_min_width
Sweep the threshold without a rebuild: ATLAS_CANONICAL_KEY_MIN_WIDTH=<n> (VALUE-parsed; 0 = canonical at every width, a value above the widest batch = never). Unset or unparseable ⇒ CANONICAL_KEY_MIN_WIDTH. Parsed once per process, like dcut_width_cap.
canonical_verify_key_enabled
Canonical assignment ON unless ATLAS_NO_CANONICAL_VERIFY_KEY is present. Read once per process.
verify_batch_order
Order one verify batch AND assign its depths — the planner’s entry point (mtp_dcut::plan), the one place a sequence’s verify depth is decided.
verify_batch_permutation
Dispatch ORDER for one verify batch — the permutation only.
verify_graph_key
The batched-verify CUDA-graph cache key for one batch: the (ssm slot, row count) pairs in DISPATCH order, then a wy-tables-present sentinel.