Blockscale sampling - #55
Merged
Merged
Conversation
Triton fails to compile the sampling kernels whenever a tile has a size-1
dimension taking part in a reduction ("PassManager::run failed"), which is a
hard crash rather than a slow path. Two configurations were reachable:
* BLOCK_M == 1 (one node block per layer) with BLOCK_S == 32
(num_samples >= 4096) -- every HMM-shaped chain, at large batch;
* BLOCK_S == 1 (num_samples < 256) -- PD-structured circuits at small batch,
including every conditional draw on them, since those inherit the
evidence's batch size.
Floor every tile dimension at 2 in all four launchers, which removes the class.
Both cases predate the recent sampling work (verified against a3259ca).
tests/queries/sample_structures_test.py covers HMM / HCLT / PD / RAT-SPN and a
hand-built ragged circuit at num_samples in {1, 16, 512, 4096} plus conditional
draws. Its PD is sized 64 vars / 32 latents deliberately: at 16/8 the layer
shapes never instantiate the failing tile and the file passes with the fix
removed.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
For any two product nodes over the same scope, the partition into their
children's scopes must be identical -- the condition under which a single vtree
exists. Computed once in `TensorCircuit.__init__` (0.1-0.5 ms).
Its purpose is the top-down sampler: when it holds, the frontier's shape is a
function of the scopes alone and never of which node a draw selected, so the
whole index plan repeats identically on every call and can be computed once.
Measured, `SD => plan invariant` on HMM, HCLT, PD, RAT-SPN, a hand-built ragged
circuit and an adversarial mixed-arity one. It is sufficient, not necessary:
RAT-SPN is not structured decomposable yet is invariant anyway, since its
several random splits share an arity -- so the flag errs toward disabling the
cache, which is the safe direction.
Two subtleties, both regression-tested:
* a ONE-CHILD product splits nothing and must be skipped. pyjuice caps every
PC with `summate(multiply(ns), ...)`, so a unary product over the full scope
always exists and conflicts with that scope's real split; counting it made
even an HMM come out non-decomposable;
* skipping unary products opens one hole -- a sum node offering both a flat
component and a decomposed one over the same scope -- closed with a local
per-sum-node check that all its child products share one partition.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Most of a top-down pass derives WHERE things go rather than drawing them: per layer a torch.where, a device-to-host copy, a serial slot allocation and a copy back. Measured, that bookkeeping is ~80% of the pass -- wall time scales with depth (x7.1 for x8 layers) and barely with batch (x1.9 for x256 samples), and GPU work is 15-20% of it. Those indices are the same on every call exactly when the circuit respects one vtree, so `pc.is_structured_decomposable` gates recording them once and replaying them afterwards, LRU-bounded per (num_samples, conditional). HMM 3.12x | HCLT 3.05x | ragged 2.47x (PD / RAT-SPN unchanged, no cache) The gate is load-bearing, not conservative dressing: forcing the cache on a circuit whose plan varies shifts per-variable marginals by 48 sigma. Circuits without the property take the original path unchanged. Also makes `num_samples` a runtime argument rather than a Triton constexpr in the three shared-parameter kernels. It is the size of the frontier, so on a circuit that is NOT structured decomposable it changes every call and Triton recompiled per call: PD sampling went 768 ms -> 3.65 ms, mixed-arity 274 ms -> 0.95 ms. It is only compared against, never used to index. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
With the index plan cached, every shape in the top-down pass is static and its
buffers are reused, which is what graph capture needs. `sample(use_cudagraph =
True)` captures the pass once per (num_samples, conditional) and replays it,
collapsing ~80 host launches into one.
HMM 0.72 -> 0.23 ms (3.19x) | HCLT 0.64 -> 0.24 (2.73x) | ragged 0.38 -> 0.16
(2.39x), on top of the plan cache -- ~11x against the uncached path, and close
to the ~0.15 ms of actual GPU work in the pass.
OPT-IN, because a graph owns a private memory pool and pins the frontier buffers
it was captured with for the circuit's lifetime: worth it inside a sampling
loop, not for a one-off draw. Requires `pc.is_structured_decomposable` -- a
captured pass replays one specific index plan -- and says so rather than
capturing something wrong.
Two things this needed:
* a graph freezes scalar kernel arguments, so seeding the RNG inside the sum
kernels would make every replay redraw the IDENTICAL sample. Uniforms now
come from a buffer refilled outside the graph, with a per-launch slice so
layers do not share a stream. Only the graph path uses it; an ordinary draw
still seeds in-kernel and is unchanged;
* `node_samples[ind_n, ind_b] = -1` is not capturable (advanced-index
assignment invalidates the capture). Replaced with an equivalent flat
`scatter_` over cached indices -- the positions are unique, so it writes the
same entries.
Correctness: 7 tests in tests/queries/sample_cudagraph_test.py, including that
consecutive replays draw DIFFERENT samples (the frozen-seed trap) and that
graphed draws match ungraphed ones per variable within 5 sigma, conditional and
unconditional. Full suite 569 passed.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`on: push` had no branch filter, so every push to every branch rebuilt the docs and force-pushed them to `gh-pages` with `CLEAN: true` -- the public site was being republished from whichever feature branch was pushed last. It currently documents `sum_external_params`, which exists only on an unmerged branch. Pull requests still build (that catches Sphinx errors before merge) but no longer publish. :note: for `push` events GitHub reads the workflow file from the pushed commit, so branches that predate this change keep publishing until they pick it up. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Both messages stated the mismatch and left the reader to work out what could possibly have staged tensors at another batch size. The answer is not obvious: a forward given no `sum_external_params` runs the layer as a plain sum layer and leaves the PREVIOUS staging in place rather than replacing it, so an intervening ungated forward at a different batch is what usually does it -- and the calls that look external (`categorical_evidence_logp`, `soft_evidence_value_mask`) are input-layer kwargs that stage no gate at all, which makes the omission easy to miss. Both now name the cause, sketch the call sequence that produces it, and say what to do. The sampling one also states why it cannot simply take the gate from its own arguments: it must use what the forward staged, or it would be sampling one distribution conditioned on another. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`pc._staged_external_params` was only ever overwritten, never cleared, so it did not describe the most recent forward -- it described the most recent forward that happened to pass `sum_external_params`. Both consumers read it unconditionally, and must: `pc.backward()` and a conditional `queries.sample()` run against a `node_mars` that was computed with those values. So an ungated forward in between made them apply a gate the forward they follow never used -- silently at a matching batch size, and as a confusing shape error from inside the sampler at a differing one. Reported from a decode loop that interleaves a 1-row gated forward with a batched ungated one; the batch difference is what made it visible, not what caused it. Verified: with the gate supplied on every forward, interleaving batch sizes was always fine. A forward given no external parameters now clears them, so what is staged always describes the latest forward. The `StagedExternalParams` branch is deliberately kept separate: `forward` mutates `kwargs` in place and registers the backward hook with `**kwargs`, so a staged dict really does come back round, and clearing there would make `lls.backward()` lose its own forward's gate. The cost, accepted deliberately: a FORGOTTEN gate is now silent rather than loud. The backward or draw simply runs ungated, consistent with the forward it follows -- which is the honest reading of "no external parameters were supplied". Alternating gated and ungated forwards is legitimate (an ungated baseline likelihood, a marginal query), so it is not warned about. Five tests in external_params_staging_test.py cover the lifecycle, including the autograd hook keeping its gate. The sampler's batch-mismatch assertion stays as defence in depth, but is now unreachable through the public API, and the test that exercised it asserts the new behaviour instead. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Triton compiles a separate kernel per divisibility class (`% 16 == 0`, `== 1`) of every integer argument. `num_samples` here is the number of SELECTED nodes, which changes from call to call on any circuit whose index plan is not fixed, and `seed` is random -- so the sampler kept hitting new specializations forever, at ~75 ms of compilation each. Measured on a `PD` circuit (not structured decomposable, so no plan cache), 60 sampling calls at batch 64: before 95 compiles, on 24 of 60 calls last-30 median 3.41 ms, MAX 365.3 ms after 52 compiles, all up front last-30 median 3.30 ms, MAX 3.3 ms The steady-state cost barely moves; what goes away is the tail, which is what a decode loop actually feels -- a third of its sampling calls were paying for a compile. Attribution was measured rather than guessed: of 21 compiled variants of the sum kernel, the layer shapes alone need 12, `BLOCK_S` accounted for none of the excess (it takes one value per layer shape), and the integer specialization of `num_samples` accounted for all of it. Neither value is used to index or to size a tile, so `do_not_specialize` costs nothing. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Step one of removing the top-down pass's bookkeeping. Today the frontier layout is discovered at run time -- a `torch.where` per layer to find its entries, a compaction to keep the buffer dense, a per-column cursor to hand out slots. MEASURED on a `PD` circuit at batch 512, that is 92% of the pass's GPU time and ~80% of its wall time; the sampling kernels themselves are 0.064 ms of 4.15 ms. None of it depends on the draw. A frontier entry stands for a SCOPE, and which scopes a layer owns is a property of the circuit, so the layout can be computed once: a layer's rows are known, rows never move, and a child's destination is its own scope's row. `prod_crows` is the piece that matters for circuits that are not structured decomposable -- it encodes "this child's scope owns this row" rather than "the Nth slot allocated", which is exactly what a varying decomposition breaks. Liveness becomes a mask rather than a shape, which the kernels already handle: an inactive row holds `-1`, matches no compiled node, and costs a masked-off lane. Measured lanes processed against lanes live: structured-decomposable circuits pay 1.00x, `PD` 3.49x, `RAT-SPN` 3.80x -- multiplying the 8% to remove the 92%. The derived layout needs exactly the buffers the driver already allocates (`_num_nscopes` / `_num_escopes`), verified on HMM, HCLT, PD, RAT-SPN and a hand-built ragged circuit. Nothing consumes it yet; the kernels and the driver follow. Tests pin the derivation, since a wrong row does not crash -- it writes one node's child into another node's slot. Four seeded faults, all caught. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Consumes the scope plan: `_use_scope_plan = True` replaces the discovered
frontier layout with the derived one. No `torch.where`, no compaction, no slot
cursor, no numba, no device-to-host round trip -- every index comes from the
circuit.
Two things fall out of the row being structural rather than allocated:
* the row is a SCALAR per program, so the grid becomes (row, sample-tile) and
`node_samples[row, b]` is contiguous along `b`, where the pair-list form
gathers a different row per lane;
* `count_prod_nch` and `sample_prod_layer` merge into ONE kernel. The count
existed only to feed the cursor, so with each child writing to its own
scope's row it, its three intermediate buffers and a launch per layer all go.
At batch 512, against the pass it replaces and against a backward:
bwd old new speedup new/bwd launches
HMM 3.06 2.78 1.48 1.88x 0.48 410 -> 155
HCLT 0.72 0.89 0.55 1.62x 0.77 84 -> 53
PD 1.14 4.09 0.71 5.75x 0.62 418 -> 65
RAT-SPN 0.48 2.14 0.43 4.99x 0.90 219 -> 43
Every circuit now samples faster than it backpropagates, which was the goal, and
the two that are NOT structured decomposable -- the ones no plan cache can help --
gain the most. It also beats the plan cache on the circuits that had one, so it
supersedes rather than complements it.
Verified against the existing path: per-variable marginals agree within 5 sigma
over 40k draws on all four structures, conditional and unconditional, and every
sample assigns every variable exactly once.
OPT-IN for now. It does not yet dispatch to `ExternalParamsSumLayer.sample_layer`,
so a gated circuit would be drawn from its shared parameters; that is refused
rather than sampled wrongly. Wiring that, then making it the default and deleting
the pair-list path, is the remaining work.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The scoped path called the shared-parameter kernels directly and refused gated
circuits from a guard in the driver. That inverted the arrangement the rest of
the external-parameter code follows -- the descriptor owns its kernels, and the
driver asks the layer rather than knowing which layer types exist -- and it put a
`NotImplementedError` about `ExternalParamsSumLayer` inside a module that is
supposed to be parameterization-agnostic.
Restored, with one convention per name rather than a second dispatch:
* `sample_layer` now means the STRUCTURAL frontier layout (rows / erows) at
every level -- layer, descriptor, `BlockScaleSumParams`;
* the pair-list convention it replaces is `sample_layer_pairs`, and goes when
the pair-list driver does.
The gated kernel gains a scoped form, `_bs_scoped_sample_kernel`. It shares
`_gate_weights` with the pair-list one, so the two cannot disagree about what an
edge weighs; only the addressing differs.
Verified against the EXACT gated distribution rather than against the other path:
drawn frequencies vs `phi * theta / Z` over 100k samples, max |z| = 2.59, and
one-hot gates still pin the drawn child deterministically.
Also factors the staged-batch check into `_check_staged_batch`, since both
conventions need it and duplicating that message was how it would drift.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Passing both selected the scoped pass and silently discarded the graph, so a caller who asked for capture got none and no indication of it. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Two things were wrong with the previous state, both raised by Anji:
* the fast path was opt-in behind an underscore kwarg, so in practice nobody
got it -- `sample()` was still the slow pair-list pass by default;
* `_use_scope_plan` + `use_cudagraph` refused each other, when combining them
is exactly what should work.
Both fixed. The scoped pass is now the default and is capturable, including on
circuits that are NOT structured decomposable -- that requirement belonged to the
pair-list pass, which replays a recorded index plan and so is only correct where
the plan repeats. The scoped pass derives its layout from the circuit, so its
shapes are static for any circuit at all. The requirement is still enforced for
the legacy path.
At batch 512, against a backward:
bwd default +graph graph/bwd
HMM 3.10 1.48 0.46 0.15
HCLT 0.71 0.49 0.22 0.31
PD 1.14 0.63 0.26 0.23
RAT-SPN 0.45 0.37 0.20 0.44
`PD` went 4.09 -> 0.26 ms, a 16x improvement on a circuit no plan cache can help.
Capture needs the RNG out of the kernels' scalar arguments, since a graph bakes
those in and every replay would redraw the identical sample. The seed is now
LOADED from a one-element tensor written before each replay -- simpler than the
uniforms buffer the pair-list pass uses, and verified: consecutive replays differ,
and captured draws match uncaptured ones within 5 sigma on every structure.
Tests that asserted the old behaviour are updated rather than deleted, and the
non-SD refusal is now pinned in BOTH directions -- allowed for the scoped pass,
still required for the pair-list one -- since getting that backwards is a wrong
answer rather than a crash. The plan-cache tests now ask for the legacy path
explicitly; they go when it does.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…. captured graph holds stale buffer pointers 3. graph cache key ignored gated-vs-ungated 4. caches ignore device changes 5. _do_calibration ignored on the scoped path
…er paths
An unconditional draw staged its external parameters into the SAME buffer a forward
pass uses, and so rewrote state the following backward depends on. `pc.backward()`
deliberately takes its gates from what the forward staged -- that is what stops a
backward running against gates its `node_mars` was not built with -- so a draw placed
between the two corrupted it: an ungated draw cleared the staging and the backward ran
ungated against gated marginals, a gated draw substituted its own. Neither raised;
measured as param-flow differences of 17-21 against a truth of 0. Draws now stage into
a private buffer. Saving and restoring the field would not have been enough, because
the buffer is reallocated whenever the batch size differs and that invalidates the
forward's views in place.
A gated draw could not use the shared CUDA-graph workspace at a smaller batch than the
one captured: the gated kernels take the frontier's width as their sample count, using
it both to validate the staged tensors and to stride the gate table, so a width-4096
replay of an 8-sample draw was rejected -- and had it not been, would have read past the
gate buffer. Gated draws now keep one workspace per batch size; ungated ones still share,
which is where that optimisation was aimed.
`_sample_input_ns = False` handed back the live frontier under `use_cudagraph`, so the
caller's result was overwritten by their next draw. It is copied now. (`.contiguous()`
does not do this: on an already-contiguous slice it returns the same object.)
Three more defects surfaced while verifying those, each caught only because the fix was
checked rather than assumed:
* `_init_buffer` set its "unknown name" flag and then indexed `__dict__[name]` anyway,
raising `KeyError` before reaching the allocation it had just decided to make -- so
the branch was dead and no new buffer name was usable;
* staging into a private buffer left `_graph_bindings` blind to the draw's own gates,
since it reads `_staged_external_params`. It is passed what the call staged now.
This also restored a regression test's ability to fail: without it the shrinking-batch
test passed on broken code;
* `_group_fast_stage` hardcoded `self.external_params`, which is `None` on a circuit
that has only ever been sampled from. This broke gated HMM sampling consistently and
was caught by the whole suite rather than by the touched files.
Tests: four regressions, each verified to fail when its own fix is reverted. Also seed
the `random` MODULE in the gated builder -- `sample()` draws its kernel seed from there,
which `torch.manual_seed` does not control, so every draw depended on how many times
other tests had called it and `test_conditional_multi_tile_edge_axis` failed about 2 runs
in 14 when it followed another file.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`test_hclt_zero_preserving` and its two HMM siblings check that parameters masked to zero are still zero after training, by reading them back off the nodes. But nothing copied the TRAINED parameters back first, so the loop was re-reading the values the test itself had just written with `set_params` -- including the very zeros it was about to assert were zero. The check could not fail. Seeding a real zero-leak (an `em_par_update` that ignores `keep_zero_params`) left all three green. `pc.update_parameters()` before the check. All three now fail with that same defect seeded and pass without it, so the property these tests exist for is guarded for the first time. It does hold: nothing was actually broken, only unwatched. `test_external_categorical_dist_fw_w_mask_speed` had no assertion at all -- it timed a forward pass and printed the result. A regression that made the forward skip its work reported a 29x speed-up and still passed. It now checks the masked half of the batch against a closed form, which is cheap there because the answer is a single category per row; the unmasked half is a log-sum-exp over 329800 categories whose reference would materialise about 2.7 GiB, which is why the neighbouring correctness tests use a small alphabet. A second assertion pins that the two branches have not collapsed into one. Verified by halving the external log-probabilities at dispatch: the test now fails by 8.0 where it used to pass. The timing output is kept, as FYI rather than as the test. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
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
1 participant
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.
No description provided.