Repository navigation
Exact constrained inference for PCs (pyjuice.constraints) - #67
Merged
Merged
Conversation
…lazy composition
New package pyjuice.constraints (exposed as juice.constraints) for
constrained inference on PCs. A Constraint is a formal language over
token ids and knows nothing about any PC: compiled against a PC with n
variables, it denotes "x_0..x_{n-1}, read in variable-index order, is in
the language".
Backends need different things from a constraint, so the base class is
capability-based rather than requiring one canonical form. Membership
(`accepts`) is the only requirement; `matcher` (incremental prefix
tracking, the hard support mask for SMC and decoding), `wmc` / `sample`
under factorised weights, `automaton` (exact lifting) and `relax` (a
finite-state over-approximation for coarsened proposals) are optional.
`capabilities()` reports what a constraint supports, and an unsupported
capability raises NotImplementedError. The Matcher contract requires
`allowed_next` to be sound: it may never exclude a token from which an
accepted sequence is reachable.
`&`, `|`, `~` and `concat` build a lazy expression tree (And, Or, Not,
Concat); And/Or are flattened and commutative. For now the composition
nodes claim only membership. Every constraint has a stable sha256
fingerprint of its structure, which defines equality, so identical
constraints can share compiled tables.
The only edit outside the new package is one import line in
pyjuice/__init__.py.
…omata
DFA stores transitions over token classes: token_class[V] and
delta[K, C], where tokens share a class exactly when they lead to the
same successor from every state (the constructor merges identical
columns and numbers classes by their smallest token). Most states treat
most tokens alike, so this is small: Ctrl-G's tutorial DFAs need 9 and 6
classes over the 50257-token GPT-2 vocabulary. The DFA is complete, with
an explicit dead state.
DFA has every capability: an exact DFAMatcher (allowed_next with an
optional exact remaining length), wmc and sample under factorised
weights (one backward pass over per-class summed weights), and
automaton/relax (itself). minimize() is canonical (Moore refinement,
then BFS renumbering), so two DFAs accept the same language iff their
minimised forms are equal. Builders: anything, nothing, from_dense,
from_transitions (a partial {state: {token: next}} map with arbitrary
ids, the shape outlines-core produces), contains (Aho-Corasick with
absorbing acceptance, Ctrl-G's keyword semantics) and from_ctrlg.
Composition nodes now build an automaton when every child has one
(product over joint token classes for And/Or, complement for Not,
subset construction for Concat), cache it, and get every capability
from it. Joint classes use a mixed-radix key and a counting pass rather
than torch.unique over the vocabulary, which on CPU cost ~180 ms per call
inside a build.
Tests check every builder and composition against its definition on all
short strings, minimisation against automata-lib's state counts, matcher
masks against brute force, and wmc / sample against enumeration.
…aints as tests DFA.word_count(lo, hi, token_kind) accepts sequences with lo..hi words, reading word boundaries off a per-token kind (starts with a separator, contains a letter or digit). This is the definition of Ctrl-G's WordCountBuilder; the sequence is treated as preceded by a separator. constraints/text.py holds the tokenizer-dependent part: word_token_kinds(tokenizer) classifies every token as Ctrl-G does (decoding each token after a special token so that leading spaces survive), and word_token_kinds_from_strings does the same from token texts. ctrlg_examples_test.py builds both Ctrl-G tutorial constraints (keyword sets AND exactly 10 words; ordered keywords AND 7-12 words) over a 36-token synthetic vocabulary, so no tokenizer is downloaded, and checks them against independent references: hand-written sentences, 4000 random strings for both the composite and its automaton, wmc / sample against enumeration, and matcher masks. On the real GPT-2 tokenizer the same builds give canonical DFAs identical to Ctrl-G's (checked in the constraints-in-pyjuice repository).
…anguage/ The package is organised by concern: `language/` holds what a constraint is (independent of any PC), and later siblings hold compilation and the inference backends. base.py, dfa.py and text.py move unchanged into `language/`, and their tests into `tests/constraints/language/`. The public API is unchanged (`juice.constraints.DFA` etc.); only direct submodule paths move, e.g. `pyjuice.constraints.text` is now `pyjuice.constraints.language.text`.
analyze_structure(pc) reports, for every node group, how its scope sits in the variables' index order -- the order a constraint reads them in -- since that decides how a constraint combines with the PC exactly: whole / suffix / prefix (a vector over one boundary state), interval (a matrix over the entry and exit states) or fragmented (several runs, an entry and exit state per run). It also reports whether the PC is right-linear (every inner scope a suffix, as in an HMM), contiguous and its largest number of runs, plus the features the first constrained backends do not handle (multi-variable leaves, non-Categorical leaves, external-parameter sum nodes). It only reports; refusing is left to the compiler. The result carries a structural signature: a hash of the node graph (node kinds, block sizes, scopes, children, edges, distribution and external-parameter signatures) that ignores parameter values and tying, so a compiled constraint can later be rebound to a PC with the same structure. Results are cached per compiled PC in a weak-key dictionary; nothing is attached to the circuit.
…t lifting The constraint side of the lifted backend's plan. Boundary t is the point before the PC reads x_t; a state is active at t if it is reachable from the initial state in exactly t tokens and can still reach acceptance in exactly n - t. Only active states get columns, numbered in increasing state id, and next_col[t, k, c] gives the column at t+1 reached from column k by a token of class c (-1 if that successor is inactive). Boundary 0 holds only the initial state and every active state at n is accepting, so neither needs its own table. An unsatisfiable constraint gets empty boundaries rather than an error. The layout depends only on the automaton and n, never on the PC, its parameters or evidence; build_layout caches it by the automaton's fingerprint. On a keyword AND exactly-8-words constraint over 12 positions the automaton has 35 states but the active width per boundary ranges from 1 to 23.
juice.constraints.compile(constraint, pc, backend=None) checks that the lifted backend can serve the pair and builds its plan: the PC's structural analysis and the boundary layout of the constraint's automaton over the PC's variables. When it cannot, one ConstraintCompileError lists every reason at once: the constraint has no automaton, unsupported nodes (grouped by reason, with an example variable set), a vocabulary mismatch with the Categorical leaves, or fragmented scopes. The fragmented-scope message names the worst node's variables and run count, how many node groups are fragmented, and the K^(2r) state size such a node would need. The result, CompiledConstraint, holds the PC, the constraint, the structure, the automaton and the layout, and reports satisfiability, the active width per boundary, state and class counts, scope-shape counts, an estimate of the lifted node-value memory per sample and the compile time. Nothing in it depends on parameter values or evidence; with_pc rebinds the plan to another PC whose structural signature matches. The query methods (marginal, conditional, sample, decoder) raise NotImplementedError until the reference backend exists. The module is compiler.py rather than compile.py: exporting a function under its own module's name hides the module (`import pyjuice.constraints.compile as m` would return the function).
The six DFAs of Ctrl-G's own correctness tests (two random DFAs, a keyword, alternative keywords, ordered keywords, a keyword intersected with a random DFA) over a 4-token vocabulary, each placed in five "prefix . x . suffix" windows. Ctrl-G's DFA reads only x; here the constraint reads every PC variable, so a window (m, L, s) becomes n = m + L + s and the constraint "any m tokens . M . any s tokens", with |x| = L exactly. Every string of length n (at most 4^7) is enumerated, and admissibility comes from the definitions (substring tests, every split for the concatenation, the random DFAs' own tables), never from our automata. Each part of compilation is checked for its own job: the automaton accepts exactly the admissible strings; the HMM is right-linear; the columns at each boundary are exactly the automaton states admissible strings pass through; and walking next_col keeps a string alive exactly while its prefix can still be completed, on the automaton's own state at every step. A guard test keeps the matrix meaningful (no case accepts every string, some pairs are unsatisfiable, the layout prunes reachable states), and a rebinding test shares one plan between two HMMs with different parameters.
The package's conventional short name, next to `import pyjuice as juice`. The package docstring shows it in a short example, and the constraint tests import the package as `jc` and refer to its public names as `jc.DFA`, `jc.compile`, etc. Internal modules that the package does not export (structure, the lifted plan, language.text) are still imported by path. No behaviour change.
…(Categorical leaves) A plain-PyTorch reference for log p(C, e) on the lifted plan, used as the ground truth for the library's GPU implementation. It lives in tests/constraints/backends/lifted/conftest.py and is exposed as the `reference` fixture: the suite runs with --import-mode=importlib, so test modules cannot import a helper module directly. Every node over the contiguous scope [a, b] gets a block of log-values over (column at boundary a) x (column at boundary b + 1): a leaf turns its Categorical token probabilities into per-token-class masses moved by next_col, a product chains its children's blocks in scope order, a sum takes the weighted sum of its children's blocks, and the root sums over the accepting columns. This covers every (PC, constraint) pair the compiler accepts. Leaves must be Categorical, like the compiler's first version requires; any other leaf raises NotImplementedError. It runs in float32, max-shifted in log space, and reads the parameters from the compiled PC's live tensors, so in-place updates are seen. reference_test.py validates it: the extracted parameters equal what pc.update_parameters() writes back; with the one-state constraint it reproduces juice.queries.marginal on an HMM (tied and untied), a 1-D PD and a hand-built circuit with three-child products listed out of scope order, a sum over leaves, node-level product edges and a sum over two products; under six constraints it equals brute force over every string on tiny PCs, also after an in-place parameter update; and its log-space block operations stay exact far outside exp's range, including a 300-token HMM.
pyjuice calls a PC's leaves input nodes (input layers, input distributions), so the constraints code and its tests now do too: the compiler's and the structure analysis' messages, the reference's `leaf_block` (now `input_block`) and its docstrings, and the test helpers' `leaf_dist` arguments. "Leaf" is kept where it means a non-composite constraint (as opposed to And / Or / Concat), which is not a PC node.
What `jc.compile(constraint, pc)` returns holds the PC as well as the constraint (and everything compiled from the pair), which the old name did not say. `ConstrainedCircuit` names both halves, parallels pyjuice's `TensorCircuit`, does not depend on the backend, and keeps `cc` as the natural variable name. The module moves with it, compiled.py -> constrained_circuit.py; `jc.compile` and `ConstraintCompileError` keep their names, as they describe the compile step.
…constructions, accumulate repeated sum edges in the reference The PCs the lifted-backend tests use now come from one `build_pc` fixture in tests/constraints/backends/lifted/conftest.py (test modules cannot import each other under --import-mode=importlib), and cover nine kinds: HMMs (tied, untied, and with 4 node blocks per position and random block-sparse transitions), 1-D PDs (plain, prod_dominated with shared sums and block-sparse sum edges, and blockified from block size 1), and three hand-built circuits with three- and four-child products listed out of scope order, a sum over an input node, an input node next to an interval node, node-level and permuted block-level product edges, children with different block counts, explicit block-sparse sum edges and block size 1 throughout. Covering the block-sparse HMM exposed a bug in the reference: its sum weights were built with `dense[idx] = blocks`, as pyjuice's own `get_params(as_matrix = True)` does, and so kept only one copy of an edge block that appears twice. `block_sparse_rnd_blk_edge_constructor` does produce repeated edge blocks, and the forward adds both copies, so the reference's distribution lost mass (log total -1.02 instead of 0, 1.1 nats off pyjuice's marginal). It now accumulates the copies. The parameter test checks that every node's weights are normalized, and compares with `get_params(as_matrix = True)` only where the edge blocks are distinct.
build_pc_tables(structure, layout, pc) in backends/lifted/plan.py reads
everything the lifted forward needs about the circuit, as int32 tensors
on the PC's device, and compile sets them on the ConstrainedCircuit:
* columns_per_sample -- S, the most columns any sum or product node
keeps per sample. A node over [a, b] keeps (columns at boundary a) x
(columns at boundary b + 1), a sequence end counting as one: the
widest boundary for a right-linear PC, a few times more with interval
nodes.
* product_rows -- per product layer and pattern ("input_suffix": an
input node followed by an inner node over the rest of the sequence,
every product of an HMM; "chain": any other product), each product
node's element row, its children's node rows in scope order and the
boundaries between them (-1 padded).
* input_range and root_rows.
Input nodes need no table of their own: in pyjuice a sum's children are
always products (SumNodes puts a one-child product above an input
child). Element rows are only unique within a product layer, since
every product layer reuses element_mars from its start.
Rows are numbered from the PC's own layers (each numbers its node
groups consecutively in the order of its `nodes` list) and checked
against the node groups' record of them, which describes the circuit
their graph was compiled into last: a graph compiled again into another
circuit is refused instead of being read with that circuit's rows.
with_pc reuses the structure analysis, the automaton and its layout,
and reads the new PC's tables (one pass over its node groups);
compiling with other options gives the same rows, but nothing
guarantees it. bytes_per_sample is now the memory of the PC's two
buffers at S columns plus the class masses.
The tests check the child rows of every product against the
reference's edge formula on all nine test PCs, the patterns, S on hand-
counted cases, the device and dtype, with_pc onto a CPU copy compiled
with other options, and the refusal of a recompiled graph.
A ConstrainedCircuit assumes its PC does not change (its parameters may, in place): `pc` is now a read-only property, so switching PCs goes through `with_pc`, and every query first checks that the PC is still on the device its tables were built on. Moved to another device, a query raises a RuntimeError asking to compile again or to use `with_pc`, instead of reading tables that live elsewhere. The check is one attribute comparison, about 145 ns per query.
…he constructor fix block_sparse_rnd_blk_edge_constructor no longer repeats an edge block, so the block-sparse HMM stopped exercising the reference's accumulation of repeated edge blocks. hand_permuted now lists one of its sum edge blocks twice, which brings that back (with a plain assignment instead, seven of its tests fail). get_params(as_matrix = True) now accumulates too, so the parameter test compares with it for every PC. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Releases the buffers a forward or backward pass fills -- node_mars, element_mars, node_mars_tempered, node_flows and element_flows -- with their backing storage, and forgets the CUDA graphs recorded against them, leaving the circuit as it was before its first pass: the next pass allocates the buffers again, and whatever a query reads from the last pass (e.g. conditional sampling) needs a new forward pass first. The parameter flows (param_flows, denom_param_flows) are kept, since they accumulate statistics across batches for the next update. A buffer still referenced elsewhere is released when that reference goes; in particular a forward pass's output carries autograd history, so a copy of it kept without .detach() keeps node_mars alive. Tests: after a forward and a backward pass on a 256-latent HMM at batch 512, freeing returns the buffers' storage (18.9 MB), keeps the parameter flows bit-identical, and the next passes give the same log-likelihoods and keep accumulating flows; recorded CUDA graphs are dropped; freeing a circuit that never ran is a no-op.
…tion buffers on creation Constrained queries will run in buffers of their own, laid out for the lifted plan, so creating a ConstrainedCircuit (by compile or with_pc) now calls pc.free_activation_buffers(): the PC's activation buffers are released, its parameter flows kept, and it allocates them again on its next own pass. cc.to(device) moves the PC and everything compiled from the pair (the product tables) to the device, in place like TensorCircuit.to, and returns the constrained circuit. A PC moved on its own still makes every query raise; the message now points to cc.to. Tests: compiling releases a PC's buffers but keeps its parameter flows and its own results, and with_pc releases them again; cc.to between the GPU and the CPU moves the PC and the tables together and keeps queries valid.
…he length The 300-position test compared the float32 reference with juice.queries.marginal within 1e-5 x max(|log-likelihood|, 1). For the sample with every position missing, pyjuice returns exactly 0 while the reference sums every token's probability at each position, with a non-deterministic index_add_, and landed anywhere between 5.7e-6 and 1.6e-5 -- so the test failed now and then. Float32 rounding there grows with the number of positions, so the tolerance is now 1e-5 x |log-likelihood| + n x 1e-6. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Constrained queries get buffers of their own, named as in TensorCircuit and laid out for the lifted plan. build_pc_tables also returns the regions, which depend only on the structure: one per sum node group (its node_mars rows and the slots its block takes per sample -- per group rather than per layer, since a sum and its children share their scope) and one per product layer group (its element_mars rows and its largest block). buffer_layout(..., batch_size) places them. node_mars holds, each region on a 16-float boundary: the input region (rows 0 .. end of the inputs, batch_size columns exactly, so pyjuice's own InputLayer.forward writes it as it writes a TensorCircuit's node_mars), the class region (the input rows x the token classes, the class masses of missing positions) and one region per sum node group, align16(batch_size x slots) columns wide, slot-major (column s * batch_size + b). element_mars is scratch for one product layer group at a time. Rows padded to 16 floats let the kernels take the column count as a run-time argument at full speed (one compile for every batch size and width). ConstrainedCircuit._buffers(batch_size) cuts the buffers from storage the circuit owns, reused while large enough and at most 4 times too large, filled with -inf only when allocated; .to() drops it. buffer_bytes(batch_size) is the exact size, and bytes_per_sample now counts what the buffers hold per sample: every sum node group's blocks, the largest product layer group and the input rows. Per sample under Ctrl-G's ex2 constraint at n = 16: an HMM 18.8 KB (61.2 KB with one width for everything), a 1-D PD 8.75 MB (17.9 MB with one width per layer, 36.2 MB with one width for everything). Tests: regions follow the PC's layers on all nine test PCs; the layout is aligned, ordered and disjoint; an input layer writes the input region bit-identically to a TensorCircuit's node_mars and touches nothing else; the views sit where the layout says; storage is reused, grown and shrunk as intended; per-group widths are smaller than per-layer widths on a PD; .to() reallocates on the new device.
cc.marginal(data, missing_mask = None, precision = "fp32") returns log p(C, e) for every sample, shaped [B, number of roots] as in pyjuice.queries.marginal. That is the log-probability that the PC generates a string which satisfies the constraint and agrees with the observed tokens. It runs the lifted forward on the GPU (CUDA only), over the buffers the circuit owns: * input layers: pyjuice's own input forward. Marginalized positions also need each node's mass per token class; one Triton kernel computes it over the distinct parameter rows, so tied HMM emissions are summed once; * product layers: a Triton kernel for an input times a suffix block (the HMM pattern), and batched PyTorch block products for the other chains (a first version); * sum layers: dense node blocks (children contiguous, parameters strided by the block size, block size >= 64) exponentiate their children once and multiply on cuBLAS using views of the PC's parameters. Other blocks take a fused Triton online-logsumexp kernel with tensor-core dots. "fp32" gives fp32-level products (tf32x3, exact fp32 cuBLAS) and "tf32" round-to-nearest TF32. The static tables (Program) are built on the first query and dropped by cc.to(). On Ctrl-G's 4096-latent HMM (50257 tokens, 32 positions) with its two example constraints, all positions marginalized, fp32 / tf32: batch ex1 ex2 1 6.3 / 5.6 ms 5.8 / 5.8 ms 16 25.9 / 18.0 ms 19.4 / 12.3 ms 128 156 / 89 ms 107 / 60 ms For comparison, pyjuice's own forward at the same number of columns (batch x mean live DFA states, 128 sequences) takes 175 / 124 ms. Tests: * the marginal against the reference implementation on every PC kind in the suite, at batch 1, 3 and 16, with two constraints and both sum paths; * a one-state constraint reproduces pyjuice's marginal; * the forms of missing_mask; * an unsatisfiable constraint gives -inf; * tf32 stays close to fp32. compile_test now runs the query where it used to expect NotImplementedError.
The lifted product layer no longer picks a code path by circuit shape
(the HMM's input x suffix pattern in Triton, every other chain in
batched PyTorch). build_pc_tables now gives every product layer one
table: each product node's element row, its children's node rows in
scope order, and the boundaries before, between and after them. The
Program folds every product right to left into binary steps, intermediate
blocks in scratch rows, and groups the steps into one launch per
contraction and operand kinds. Each step is one of four contractions,
chosen by its operands, never by the circuit's shape:
* input @ block: an input at t, then a block from t + 1 -- a gather over
the token classes through next_col. The block may be the identity: a
product with a single input child, or an input next to an input;
* block @ input: the mirror, a gather over each column's predecessors
(predecessor_tables);
* block @ block: a per-sample log-space matrix product over the shared
boundary on tensor cores, each operand's row or column maxima taken
once per launch over its distinct operands;
* copy: a product over a single sum.
Products are always fp32-level (tf32x3 for block @ block, exact fp32
log-sum-exps for the gathers); the query's precision applies to the sum
layers.
Against the previous implementation, product layers only (Ctrl-G's ex1,
32 positions, all marginalized):
1-D PD, 16 latents batch 1: 1425 -> 17.6 ms batch 16: 1454 -> 172 ms
HMM, 4096 latents batch 1: x1.23-1.39 batch 16: x1.03-1.09
batch 128: x1.00-1.05, bit-identical marginals
Tests: a hand-built left-linear PC (hand_left: a prefix chain and a
three-child product) joins the suite; plan_test reads the new tables.
…s for block @ block Most entries of a lifted block are structurally -inf: the pair (column at a, column at b) is finite only if the automaton joins the two states in exactly b - a tokens. Under Ctrl-G's ex1 on a 32-position PD, only 21% of the block @ block operand entries and 5% of its multiply-adds are feasible. In state-id order they are scattered, though: 92% of the 16 x 16 x 16 tile combinations are non-empty. * build_layout orders each boundary's columns by the automaton's fewest tokens to acceptance (within n), farthest first, ties by state id. Any order is valid; this one gathers the infeasible pairs into whole tiles (39% of (16 x 16 tile, 16-state chunk) pairs left on that workload, 42% in the reverse order). The distance comes from the accept_within table _build already computes. * Reachability(layout, tile): which pairs every interval joins, kept per interval only as tile occupancy and computed on demand, one forward sweep per entry boundary (a scatter through next_col, no stored matrices). * skip_masks(reach, triples, tm, tn, tk) -> SkipMasks: per (t0, t1, t2) triple and output tile, one bit per chunk of the shared boundary, in several 32-bit words when the boundary has more than 32 chunks. A kernel with larger tiles ORs its sub-tiles' masks. * The Program builds the masks of every block @ block step once and attaches each step's triple index to its launch. Nothing reads them yet: the product kernel that does comes next. On that workload the masks take 61 ms of the 1.8 s Program build and 1.07 MB (the block @ block step tables: 3.49 MB). Marginals are unchanged. Tests: the column order on random automata; Reachability against a brute-force search over token strings, sequence end included; the masks against their definition at four tile shapes, with several words per tile on a 120-state automaton; larger tiles equal the OR of their sub-tiles; every block @ block step indexes its own triple. ctrlg_style_test checks the set of active states rather than their order.
…le, skipping empty tiles Lifted buffers now keep each sample's block contiguous: slot s of sample b sits in column b * slots + s (was s * batch_size + b). Sums, copies and the root's readout act on whole rows or single columns and are unchanged; the product kernels are rewritten around the new order. * block @ block: one program per (step, sample) walks its own 16 x 32 output tiles, so the sample's two blocks stay in L1 across them. Each tile takes its rows' and columns' maxima over the shared chunks it needs, then the product of the shifted exponentials over the same chunks, with exact fp32 dots (PRODUCT_PRECISION; "tf32x3" is the alternative, same speed here). Chunks the Program's skip masks rule out are never loaded (SKIP_EMPTY_TILES). The maxima pre-pass over distinct operands, its buffers and the operand tables are gone: in the old slot-major order the batch dimension took the register tile, capping output tiles at 16 x 16 and re-reading every operand about 7 times. * input @ block: one program per step and tile of its whole output, flattened over (sample, entry column, exit column) -- one contiguous run in the new order, and full lanes whatever the widths. One program per (step, sample) had left an HMM's 1-column suffix blocks with per-lane lookups and idle lanes (products up to 2.9x slower). * block @ input: one program per (step, sample, exit column), its lanes over entry columns. PD with 16 latents (Ctrl-G's ex1 / ex2, 32 positions, all marginalized), the whole marginal against the previous commit: batch 1 18.4 -> 8.4 ms (x2.18) 8.8 -> 4.8 ms (x1.85) batch 16 200 -> 118 ms (x1.69) 93 -> 63 ms (x1.48) 4096-latent HMM: unchanged (same GPU, alternating runs, batch 1 and 16 within 3.6%; batch 128 within 3%), identical marginals. Tests: every sample's block of every sum node is where a batch of one puts it (four PC kinds); skipping empty tiles gives what computing them gives on a 48-state automaton whose blocks span several tiles (a mutation that inverts the masks fails it).
A product with a single child that is a sum (a copy) was materialized: the sum's block copied into element_mars, for the next sum layer to read. On a 16-latent PD those copies took 22% of the marginal at batch 16. * The Program no longer launches a copy whose element row only fused node blocks read. Those node blocks read the sum's own row instead: each fused part's edges are split in two lists, the children in element_mars and the aliased ones as (row within its sum region, region, pid). The fused sum kernel runs a second loop over the aliased list into the same online log-sum-exp. (Selecting between the two sources inside one loop doubled the loads of every chunk and cost 25% of the sum kernel, even with nothing aliased.) * A copy that a dense node block reads is still made: the dense path needs its children as one run of rows. * ALIAS_COPIES turns it off. The aliased children are summed after the others, so results agree up to the accumulation order (identical in fp32 on the workloads below). PD with 16 latents (Ctrl-G's ex1 / ex2, 32 positions, all marginalized), the whole marginal against the previous commit (same GPU, alternating runs): batch 1 8.8 -> 7.3 ms (x1.20) 5.0 -> 4.4 ms (x1.15) batch 16 118 -> 93 ms (x1.28) 63 -> 51 ms (x1.23) 4096-latent HMM: no copies, unchanged, identical marginals. Tests: aliased and materialized copies give the same marginal on every PC kind through both sum paths, and no row a dense node block reads is aliased; on a PD every copy is aliased and none is launched. Mutations that skip the aliased loop, read the wrong row or drop row 0 fail.
The dense sum path (an HMM's transitions) made seven elementwise passes around every cuBLAS product: the children's column maxima, a finite check, the shift (on a strided view, so not vectorized), the exponential, then the log and the shift added back on the product. On a 4096-latent HMM at batch 128 they took 12% of the marginal. Now the column maxima (torch.amax), then one Triton pass that shifts and exponentiates into the product's operand, and after the product one in-place pass that takes the log and adds the maxima back. A column of -inf children is shifted by 0 and stays all 0; its product is then exactly 0 and log 0 - inf = -inf, so the second pass needs no guard. 4096-latent HMM (Ctrl-G's ex1 / ex2, 32 positions), the whole marginal against the previous commit (same GPU, alternating runs): batch 1 x1.15-1.17 batch 16 fp32 x1.03-1.07, tf32 x1.12-1.16 batch 128 fp32 x1.06 (ex1 158 -> 150 ms), tf32 x1.11-1.13 (89 -> 80) Marginals identical or within 1.5e-5. The PD's sum layers do not take the dense path and are unchanged. Tests: the existing dense-path tests; mutations that drop the -inf shift guard or the maxima added back fail them.
…ct list -- JSON hard_1 diffusion 3158 -> 2709 ms Lever D let a team stage the masses of the codes its terms read alone, but its staging loop still visited every (child, code) slot of the boundary to check the code's mark: 615 per child on JSON hard_1, 979 on regex 400words. That is a few times more work than the few marked codes need. After marking, a warp now takes the prefix of the marked codes over the mask's words (`wpre`, in shared memory after the mask). The loop runs over (child, marked code) alone. The i-th marked code is found by a binary search of the words' prefix, then its bit within the word. Both loops compute a mass with the same `mass_of`, so the values are the same bits. The Python and C++ shared-memory rules both count the prefix. Codpula 4096, batch 16, full decode (toast; base 146b54a): - JSON hard_1 diffusion: 3158 -> 2709 ms (23.8 -> 20.3 ms a call; `ibtd_kernel` 1029 -> 578 ms) - regex diffusion: 279 -> 259 (`ibtd_kernel` 71 -> 48 ms) - the l2r cases: unchanged Eight recorded sessions bit-identical; four mutants are killed (the next marked code; an inclusive word prefix; and lever D's marking mutants, nothing marked and the first live column's terms alone). tests/constraints and strong_check --slow pass.
…long the rows -- HMM ex108 B8 diffusion 5766 -> 5468 ms `_live_shifted_exp_kernel` runs a program per live column and reads that column at the row stride: one float per memory sector, twice (its maximum, then its exponentials). An HMM's decode keeps its dense sums compact, so every column of a product is live, and the pass reads every [4096, live] block of a session that way. Passes of 256 live columns or more now run `_live_shifted_exp_tiles_kernel`: 16 columns a program (32 from 4096), 512 rows of them at a time, read along the rows. The values are the same bits: a maximum is exact in any order, and every exponential is the same. Narrower passes keep the per-column kernel, whose many programs are faster there. HMM ex108, full decode (toast; two passes with the GPUs swapped): - diffusion B8: 5766 / 5787 -> 5468 / 5464 ms (the exp passes 561 -> 232 ms) - diffusion B4: 3295 / 3295 -> 3150 / 3151 (263 -> 110) - l2r B16: GPU time unchanged (92.6 ms; its passes are narrow) Nine recorded sessions are bit-identical (HMM ex108 B8 diffusion among them). The test covers both kernels' tiles and asserts the tiled kernel's bits; six mutants are killed. tests/constraints, strong_check --slow and the Ctrl-G fp64 check pass.
… out transposed -- HMM ex108 B8 diffusion 5454 -> 5260 ms The dense path's live-column products ran `W @ ex_t^T`, [rows, live]. On tensor cores cuBLAS picks faster kernels for the transposed product `ex_t @ W^T`, [live, rows]. Replaying every product of an HMM ex108 B8 diffusion session (M = K = 4096, 1 to 83k live columns), TF32 takes 1806 -> 1469 ms; the transposed form is faster in every band of widths, with the weights stacked or not. At fp32 (SIMT kernels) it measured 4423 -> 4480 ms, so fp32 queries keep their products as they were (`_TRANSPOSED_PRODUCTS`: TF32 only, which tf32 and bf16 queries run). The log scatters (`_log_scatter_kernel`, `_log_scatter_accumulate_kernel`) read the product transposed under a new `TRANS` flag. The results are not bit-identical at TF32 (other cuBLAS kernels): the recorded HMM sessions move by at most 1.9e-5 in their conditionals and 6.1e-5 in their log marginals, and their difference from fp32 is unchanged. fp32 queries and every Codpula session are bit-identical. HMM ex108, full decode (toast; two passes with the GPUs swapped; base 08fbcb0): - diffusion B8: 5454 / 5466 -> 5260 / 5263 ms - diffusion B4: 3147 / 3148 -> 3058 / 3056 - l2r B16: within noise (124 ms) Tests: no strong-check case reached the dense products. The strong checker gets an HMM (64-node dense blocks) in `test_a_session_against_float64`, and `test_unstacked_dense_products_against_float64`, an HMM of two 64-node blocks a transition multiplied block by block, so a run's later block log-adds its top-down values (the decoder stacks every run up to STACK_MAX_COLUMNS columns, so nothing else reached that scatter at TF32). Four mutants are killed. tests/constraints, strong_check --slow and the Ctrl-G fp64 check pass.
… on bf16 operands -- HMM ex108 B8 diffusion 5253 -> 4553 ms At precision "bf16" the CUDA fused sums multiplied bf16 operands, but the dense cuBLAS path still ran TF32. Its live branches now round the exponentials to bf16 (the exp pass writes them so) and multiply bf16 copies of the weights, summing and returning fp32 (`torch.mm(..., out_dtype = torch.float32)`). The copies, of the stacked weights or of each block's, are kept in the query's weight stacks: built once inside a pyjuice.fast_inference scope. Where torch.mm cannot (`out_dtype` is recent PyTorch), and on every column (no live columns), the products run as at "tf32". HMM ex108, full decode (croissant; two passes with the GPUs swapped; base 23a5d46): - diffusion B8: 5253 / 5267 -> 4553 / 4548 ms - diffusion B4: 3057 / 3059 -> 2708 / 2710 - l2r B16: within noise Conditionals against fp32 (B8 diffusion): at most 2.7-3.8e-4 off (TF32: 1.9-2.2e-4), no -inf mismatch. The recorded HMM sessions move by up to 2.8e-4 (conditionals) and 9.2e-4 (log marginals); every Codpula session is bit-identical, and fp32 queries are unchanged. Tests: `test_bf16_dense_products_multiply_bf16_weight_copies` (stacked or block by block, tied or not: one copy per distinct weights, built once in a scope, and what is multiplied; within bf16's error of the reference); the strong checker's HMM cases now run real bf16. Five mutants are killed. tests/constraints (croissant: 8563 passed; the one failure needs automata-lib, absent there), strong_check --slow and the Ctrl-G fp64 check pass.
…ernal parameters sample by sample -- Codpula ex108 unpruned B8 diffusion 7162 -> 6564 ms A sum layer under external parameters (BlockScale) ran every node block fused, each chunk of children scaled by its gate. When its node blocks are one dense run -- every block over the same contiguous children, the blocks consecutive rows of one region (`block_scale.dense_run`, recorded with the part's gate tables) -- the top-down pass now multiplies its compact rows sample by sample instead: each sample's weights `theta * phi_b / Z_b` made into one matrix (`sample_weights`, one buffer reused across samples), then the dense path's transposed product over that sample's lanes (bf16 operands at a bf16 query, as in `dense_sum_top_down`) and its log scatter. A layer takes it on compact rows, on tensor cores (tf32 and bf16 queries), its children unwritten by an earlier launch, and from `PER_SAMPLE_MIN_LANES` (256) lanes a row per sample: replaying every fused launch of a Codpula ex108 B8 diffusion session, the fused kernel won every launch below 128 lanes a sample, the products by sample 99% from 256. Narrower rows, dense storage, fp32 queries and every layer without external parameters run as before. The forward stays fused: by sample, its bf16 products left the conditionals of that session's improbable tokens with 10x the fused kernels' error against fp32 (1.6e-2 vs 1.6e-3), while the top-down pass's matches them (1.6e-3). Codpula ex108, full decode (croissant; two passes with the GPUs swapped; base 4573354): - unpruned diffusion B8: 7167 / 7158 -> 6572 / 6557 ms (peak 61.2 -> 62.6 GB) - unpruned diffusion B4: 4062 / 4066 -> 3736 / 3738 - pruned diffusion B8, l2r B16: unchanged Tests: `test_sums_by_sample_against_float64` (Codpula's top-down pass by sample at any width: a gate per pair of blocks, node and child gates finer than a block, every sample's gates the same; diffusion and l2r; tf32 and bf16) against the float64 reference. Five mutants are killed. The test files that build layers under external parameters pass, with strong_check --slow; with the products by sample forced at every width, all but one tolerance check at bf16.
…values only for the rows with new evidence in its scope -- Codpula ex108 B8 diffusion 6559 -> 6283 ms, HMM JSON hard_1 B8 4029 -> 3887 ms A region some rows of a call hold the inside values of (nothing newly observed in its scope for them) and others not is PARTIAL: its held rows' values are re-gathered into the call's layout, and the forward pass computes the other rows' lanes alone -- the dense sums multiply those columns and leave the others as they are, the fused sums take them as their columns (fresh_order), the input@block / block@input walks skip the tiles of held rows (fresh_counts). The move into the new layout goes through the previous call's top-down storage where it holds nothing the call keeps, so the peak does not grow (HMM ex108 B8 51.78 -> 51.87 GB); the fresh lanes are built in one launch, the lane offsets copied with the layout's one wait. 31-call decode, two passes on croissant: HMM ex108 B8 4566 -> 4499 ms; left-to-right decoding unchanged (test_decoding_speed). Tests: decoder_test (rows holding values get them back; sums compute the fresh rows' lanes); the dfa_exact suites that reach these paths, --slow; Ctrl-G fp64 cross-check.
…algebra of byte automata, and reads oneOf, allOf, not and if/then/else -- 1809 JSONSchemaBench schemas accepted (was 1678), hard_0 50 -> 0.35 s language/backend/automata.py: the sets of texts byte automata accept -- literals, characters by code point, regular expressions (Regex's checks) -- combined by concatenation, union, repetition (one subset construction with empty moves, numba), intersection and difference (products), every result minimal in a canonical form (equal tables iff equal languages). byte_level._trimmed finds the co-reachable states by a reverse breadth-first search (numba; it was quadratic on chains). Regex and JSONSchema share a text-constraint base (regex._Text). language/backend/json_schema.py writes outlines-core's canonical texts piece for piece, each schema object compiled once: of the 1621 JSONSchemaBench schemas the checker accepted, 1615 have exactly outlines-core's language; the others hold property names it does not escape, or references nested past its third, after which it drops properties (required ones too). outlines-core is now a test dependency only; two JSONSchemas are equal when their languages are. The corpus compiles in 144 s, against 32361 s through outlines-core's regex and interegular. oneOf, allOf (any number of schemas), not and if/then/else are read exactly up to the canonical form. Subtraction is sound because what is subtracted is a superset: below a subtracting combinator every schema object is compiled twice, its written texts (A) and every text written there that is an instance (B), in one canonical form of its own (one order of names, numbers and strings one way). B's any value counts the nesting alone (exact brackets cost 2 ** depth states). A combinator beside other keywords is still refused, with where. Tests: automata_test (every operation against Python's re and its operands on every byte string up to 5 bytes; canonical forms under renumbering); json_schema_test (outlines-core parity on every form and two whitespace patterns, and -- slow -- on the cached JSONSchemaBench; combinators: cases, every text of a small universe of values in any order of names and spelling of numbers accepted only if an instance, exact draws valid). 20 mutants killed. On the corpus, the 137 combinator schemas with texts: every exact draw valid (jsonschema, numbers exact).
…l, trim on class tables and join in one pass -- JSON Schema compile 3.4-5.2x on the slowest JSONSchemaBench schemas, hard_0 0.35 -> 0.11 s numpy.unique(axis = 0) over the [states, 256] tables, to find the bytes no column tells apart, was two thirds of a compile. _column_classes hashes every column in one pass over the rows and compares each byte's column with its hash's classes, numbering them by their smallest byte (the canonical form unchanged); the dead state is found on the minimal table over classes, before it is expanded to bytes; the parts of a concatenation, union or repetition are gathered in one vectorized pass. Github_trivial 129 27.2 -> 7.0 s, Github_hard 308 188 -> 55 s, Github_medium 46 99 -> 19 s, 230 67 -> 14 s, 340 37 -> 9 s. Same automata: the canonical-form tests (renumbered automata, equal languages, equal tables) pass unchanged.
…s -- JSON Schema compile 1.2-1.6x on the largest JSONSchemaBench schemas, tables 4-5x smaller A ByteAutomaton holds a table over its byte classes (the bytes no state tells apart, numbered by their smallest byte) and the class of every byte; the table over all 256 bytes is made on first use (made from one, an automaton finds its classes on first use). Operations refine their operands' classes from the 256-entry maps instead of scanning [states, 256] tables, minimization merges the classes the minimal table does not tell apart (the canonical form stays unique), and the token lift walks the vocabulary trie through classes. Github_hard 308 55.1 -> 44.3 s (701605 states x 68 classes), Github_medium 46 18.9 -> 11.7 s, 230 13.8 -> 8.8 s, 340 8.8 -> 6.1 s, Github_trivial 129 7.0 -> 6.1 s, hard_0 0.11 -> 0.08 s.
…y byte lifts to tokens in parallel, minimal without Hopcroft -- 312k-state JSON Schema automaton 100 -> 8.1 s, hard_0 1.2 -> 0.3 s Every byte a token (GPT-2's, byte-fallback vocabularies): every state a byte text reaches is reached at a token boundary, so the states are a breadth-first search away and their trie walks are apart -- walkers take every P-th state, each summing into its own hashes. A minimal byte automaton's states are then told apart by tokens as by bytes, so the token DFA is minimal once the padding state is merged into the state accepting the empty text alone (where there is one): it is renumbered, not minimized. Other vocabularies, and byte automata not known minimal (a Regex's), lift as before. A DFA merges its identical class columns in one pass (backend.dfa._column_classes, moved there) instead of torch.unique(dim = 1), for every DFA made. Github_medium 340 (312099 states, GPT-2): walk 58 s sequential -> parallel, Hopcroft 27 s -> none, column merging 13 s -> one pass. Test: the lift equals the walk from the initial state minimized, with and without padding, with 1, 2, 3 and all threads.
…ta are evaluated lazily, once per expression -- JSON Schema compile up to 4.4x on the largest JSONSchemaBench schemas concat, union, repeat, star and optional return an expression (automata.Expr), evaluated on first use all at once: its automata wired into one automaton with empty moves (_Joined.wire), determinized and minimized once -- instead of after every operation, which built a bounded array's body three or four times over (the repetition, the first item, the optional, the brackets). A repeated part is evaluated first; the JSON Schema compiler evaluates every schema object's automaton once (its memo), and what it reuses (free-form values, B's any value, a combinator's result), so that no subexpression is wired again at each use. Every one of the 1809 JSONSchemaBench schemas the checker accepts compiles to the same canonical automaton as before; compile time 1187 -> 876 s over them, Github_easy 23 29.0 -> 6.6 s (2.05M states), Github_hard 308 43.7 -> 24.1 s, Github_medium 549 18.6 -> 10.2 s, 46 12.7 -> 7.2 s, 230 9.3 -> 5.0 s (40 compiles at once). Test: random nested expressions -- subexpressions shared among them -- evaluated at once equal every operation evaluated in turn.
…-bit -- an illegal memory access past 2^31 entries (JSON hard_0 from n = 768, the 10x schema at 1024) Eleven Triton reads computed `next_col + (t * W + i) * C + class` in int32, wrapping once n * W * C exceeds 2^31 (the CUDA liveness sweep and the transition kernel already used 64-bit offsets): the class weights' sample walk faulted first, the forward and top-down walks' and the Triton sweeps' reads were the same. On JSON hard_0 (W 3098, C 981) n = 640 ran and n = 768 faulted, as the bound says; fixed, HMM hard_0 decodes left to right at n = 1024 (compile 61 s, tables 30 GB, 25 ms a call, peak 80 GB). Test (slow, about 20 GB of GPU memory): a 129-state automaton of 46715 classes over 400 positions, 2.4e9 entries, the last 41 positions -- whose offsets overflow -- asked of a two-latent HMM: the rows are the float64 reference's within 3 fp32 ulps of the log values near -4000 they are differences of (the same at positions without overflow); on the kernels before, an illegal memory access. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…rough the boundaries' column maps past 4 GiB -- JSON hard_0 at n = 1024 compiles in 38 s on 4.7 GB (was 59 s on 29.6 GB), the 10x schema's 61 GB of tables go A boundary's transitions are the automaton's through its column maps: next_col[t, i, c] = col_of_state[t + 1][delta[state_id[t, i], c]]. The layout keeps delta [K, C] and the two maps and derives a boundary's table on demand (BoundaryLayout.next_col_at, Transitions); the program's transition tables -- entry and predecessor lists, class lists, enumerated classes, statistics (transition_stats) -- are built from them a boundary at a time on the device, and pair_of [n, W, C] is gone (the transition masses compare successor columns). The table over every boundary stays, built per boundary with no host copy, while it has at most TRANSITION_TABLE_MAX_ENTRIES (4 GiB) entries; past that the kernels read a transition through the maps (common._next_col, HOMOGENEOUS; the CUDA liveness sweep the same). The maps cost the walks over every class of a column that read the table per lane -- those of automata of few classes (Ctrl-G ex108, 14 classes: class weights +11%, input @ block group +17%), whose tables fit the bound at any practical size (ex110 at n = 1024) and stay. Automata of many classes walk entry lists: JSON hard_0 left to right at n = 1024 10.1 -> 9.5 ms a call through the maps (they stay in L2), hard_1 B8 the same. Reachability advances through every state's distinct successors. Bit-identical conditionals (ex108 HMM and Codpula B8, hard_1 B8 diffusion and left to right), -0.1 .. +0.6% in time; Kubernetes schemas (24k states, 1.6-1.9k classes) decode at n = 128 on 45-54 GB where their tables alone were 20 GB. Test: the maps forced on small automata (threshold 0) give the tables' rows bit for bit -- dense and compact storage, marginals, every liveness sweep, a soft-evidence decoder session.
…ool under a memory budget, the ones a call reads held and the rest computed again as decoding reaches them -- JSON Kubernetes_171 (24k states) decodes left to right at n = 256 on one GPU (54.8 GB peak; the 82 GiB of suffix values failed to allocate), n = 512 at 74 GB (its 190 GB of values never fit) A decoder session keeps the inside values of the scopes no sample observes once, for every sample and every call (prior_mars): decoding left to right, the suffixes after the prefix -- 4096 x n x K floats for an automaton whose every state is possible at most boundaries, 82 GiB for Kubernetes_171 at n = 256. They are now a PriorPool: a buffer under a budget, an extent per sum region while its values are held, the values an interval no query needs evicted when the pool is full (their validity withdrawn, so that a later query computes them again from the ones still held below), the one read longest ago first and the checkpoints (depth a multiple of the stride) last. The stride is the smallest for which the checkpoints and a segment of values fit the pool together, so no value is recomputed from more than a segment below: left to right, every segment of the suffixes is computed once more as decoding enters it, about one forward pass more a session (Kubernetes_171 n = 256 B = 4: 255 calls 67.8 s, 38.5 s of them call 0, ordinary calls 10 ms, a segment's entry 350-460 ms; n = 512: 511 calls 168 s, call 0 82 s, every fifth call computes 4 values again, 386 of 512 over the session). The stride of the least estimated recomputation when no stride's checkpoints and segment fit together. The planner computes, of the values kept once, those the launches READ -- the children of the products computed per sample or top down, the root's for a marginal -- and, so that each can be computed, their children not held, down to the held ones (QueryPlan.prior_read: what the query reads, held or computed, pinned while it runs; held children of the computed ones among them). When they do not fit the pool beside its checkpoints, the lower depths are computed first in pieces, each a plan of its own (plans.kept_plan, forward.kept_values), the pieces keeping what the query reads, so that every piece leaves fewer to compute. A query whose values do not fit at all is refused with their size and the pool's: a PD's first call reads nearly every value kept once (its top-down values kept once read their siblings' inside values), so no budget below their total serves a PD. The budget: cc.decoder(memory_budget = <bytes>); None (the default) keeps every value when their total fits the memory the process can take (the device's free memory and the blocks the caching allocator holds unused), as before, bit for bit, else PRIOR_MEMORY_SHARE (60%) of it. Kubernetes_171 n = 256 under 16 GB: 32.3 GB peak, 75.4 s. The kernels are unchanged: they take the offset tables they took. A/B against 5bc68f9 at n = 128 (hard_0 left to right and diffusion, ex110 left to right, Kubernetes_171 HMM and Codpula): equal within noise. Program.iv_children / iv_depth; State.prior_buffer(plan) allocates the extents of what the plan computes; prior_pool(), set_prior_budget(). Tests: kept_values_test (a half-size pool bit-identical to the full one over a left-to-right HMM session, as beams, with a marginal; every value read has an extent; the refusal and the option); plans_test's flags follow the read closure.
…s the positions where it rules classes out as observed -- their values per sample, pruned, instead of kept once over every column: Codpula on JSON Kubernetes_171 at n = 256, B = 1: call 0 50.8 s -> 87 ms, peak 76.9 -> 34.5 GB, 255 calls 84 -> 5.4 s A decoder session's per-sample inputs of a batch mark their positions as observed: their values differ per sample, so nothing over them is kept once. Inputs of one row for all did not: their values are every row's, so the input-independent intervals were kept once -- and a block kept once is computed without evidence pruning, over every column (plans.py). Under soft evidence that rules classes out, the per-sample values of the same intervals are pruned by the candidate classes (kernels.live.class_support): on Kubernetes_171 at n = 256 the suffixes take 0.07 GB per sample against 88 GB kept once, and a session of two rows ran its first call in 0.1 s where one row took 51 s. A query now plans as observed the positions ExternalInputs.observed_under(prog, masses) gives: the session's inputs' own (observed: evidence of a batch, gate scopes of a batch) and, under evidence of one row for all, the variables where its class support rules classes out -- read from the support the session's first query computes anyway for its pruning (live_successors keeps it as ClassMasses.support), once per masses. Evidence that rules nothing out, gates alone, and inputs of a batch plan as before. The decision is one per session: a rule deciding per call (keep once only when the block is no larger than the samples' live slots) was tried and makes intervals flip between the two modes as liveness changes with the prefix -- HMM n = 128, B = 4 went from 11.5 to 288 ms a call. Kubernetes_171 Codpula B = 1 l2r: n = 128 call 0 21 s -> ~0.1 s; n = 256 (pool of 43 GB before) as in the title, the same as B = 2 and 4 (36-39 GB, 5.1-5.2 s). ex110 n = 128 B = 1 l2r: 5.9 -> 1.9 s, 69 -> 16 GB. JSON hard_0 n = 128 B = 1, warmed: l2r 825 -> 255 ms (call 0 592 -> 42 ms, calls 1.8 -> 1.7 ms), diffusion 1372 -> 649 ms (call 0 826 -> 56 ms, calls 17.6 -> 19.1 ms). HMM n = 128 B = 4: unchanged (15.6 s call 0, 17.1 s, 45.3 GB). Tests: softevi_categorical_test -- one row's evidence ruling nothing out is kept once, per-row evidence never, one row's that rules classes out is kept once only over positions where it rules nothing out, all the brute force's values; a one-row session with top-k evidence, left to right, equals the brute force and holds its values per sample.
…ir transition tables -- JSON Kubernetes_171 (24k states): 81 distinct boundaries of 1024, tables 27 -> 2.8 GB at n = 512, 3.1 GB at n = 1024 where compile ran out of memory; the HMM decodes left to right at n = 1024 on one GPU A boundary's transition tables -- its (column, class) entries with their successors, by column, by successor and by class, ungrouped for the product kernels and grouped for liveness and the class weights -- depend on its possible states and the next boundary's alone: the columns, their successors' columns, the classes live there (an enumerated token's code counts from the boundary's first token), and whether the boundary is the sequence end. Decoding at a length well past the shortest accepted string, nearly every boundary has every state possible, so the tables were n copies of one: 5.3 GB at n = 128, 27 GB at n = 512, and the builder's int64 lists over every boundary ran out of memory at n = 1024 (HMM) and n = 512 (Codpula, 94 GB). transition_tables now groups the boundaries by key (a hash of the states at t, of the states at t + 1, the end flag), builds one segment per distinct key, and points every boundary's rows of entry_ptr, pred_ptr, cls_ptr and enum_ptr into its segment. The kernels read the tables as before. One structure held a value per boundary and entry -- a sample's successor lists (live.SuccessorLists, SampleSuccessors), a per-sample copy of the entry space addressed by the entry numbers --, where alike boundaries would overwrite each other's lists (the suite caught it): the tables carry boundary_base [n + 1], the boundaries' entries counted one after the other, the lists are laid out by it, and a list's start is counted from the sample's first (prod._list_range reads it without the boundary's first entry). Kubernetes_171 HMM: n = 128 tables 5.3 -> 2.6 GB, compile peak 7.5 GB; n = 512 27 -> 2.8 GB, peak 10.8 GB; n = 1024 3.1 GB, peak 11.6 GB, compile 116 s (the layout's reachability on the host, next). Decoding at n = 1024 B = 4 left to right, under the pool (56 GB of the 393 GB of values, stride 8): 1023 calls 415 s, call 0 169 s, ordinary calls 42 ms, 142 segment entries 208 s, peak 65.9 GB -- 24k states at 1024 positions on one GPU. A/B against d80bcd0, warmed, both GPU orders (hard_0 n = 128 HMM l2r 722/738 -> 731/720 ms, diffusion 9838/9950 -> 9845/9956, Codpula l2r 271/274 -> 275/275, diffusion 961/968 -> 972/962; ex110 HMM l2r 2414/2416 -> 2411/2407; Kubernetes_171 HMM n = 128 l2r 16152/16143 -> 16127/16133, peak 45.4 -> 42.7 GB; Codpula n = 256 l2r 2509/2543 -> 2554/2584, peak 39.3 -> 28.0 GB): equal within noise, the peaks lower by the tables. Tests: layout_test -- alike boundaries share one segment (the entries stored are the distinct boundaries', pointer rows of alike boundaries equal) while every boundary's entries read through its pointers are its own (the existing test against next_col); liveness, forward, conditional and enumerated tests cover the kernels and the per-sample lists.
…PC's device and stop at their fixed point -- JSON Kubernetes_171 (24k states) compiles in 2.1 s at n = 1024 (was 116 s), 1.6 s at n = 256 (was 28 s); jc.compile(..., verbose = True) shows the compilation's stages A constraint's layout -- the states possible at every boundary: reachable from the start in exactly t tokens and able to accept in exactly n - t more -- was two sweeps over the automaton's transition table on the host, one step per boundary, each a gather over every (state, class) edge: 39M on Kubernetes_171, 99% of the compile time (115 of 116 s at n = 1024). Both sweeps now run on the PC's device (the automaton's tables already live there for its other uses) and stop at their FIXED POINT: a row equal to the one before repeats for ever, the step being a function of the row, so the rest are copies -- decoding at a length well past the shortest accepted string, the sweeps settle after a few dozen boundaries of a thousand (Kubernetes_171: 81 distinct). The column numbering reuses the row before where the states are the same. The layout is the host's, bit for bit, and still cached by fingerprint and n. jc.compile takes verbose (False): a tqdm bar per stage, as pyjuice.compile shows its own -- the layout's two sweeps and its columns, the PC's tables (compiler.Progress; Backend.build takes it as report). The tables a backend builds at its first query are not part of it (Kubernetes_171 n = 1024: 5 s, per-boundary loops that could share alike boundaries too). Tests: layout_test -- the sweeps against the brute force on an automaton that settles, a parity counter that alternates for ever (every boundary swept), and random ones, at lengths past where they settle; the device's layout equals the host's. compile_test -- verbose shows its stages and gives the same circuit.
…s work, not to the sequence -- Kubernetes_171 (24k states) HMM left to right at n = 1024: an ordinary call 47 -> 8 ms, 2080 -> 33 kernel launches; planning 28 -> 2 ms, the liveness sweep 19 -> 3 ms; Codpula at n = 256: a session 3.1 -> 0.9 s Left to right, a call reads two suffix blocks and the prefix's chain, yet its time grew with n (4 ms at n = 128, 47 at 1024) through three costs per boundary rather than per boundary touched. The top-down planner's mode pass looped over every input step and class weight destination of the program (1024 positions x 4096 nodes = 4.2M of each) testing whether its position is a target: _top_down_tables now groups a stage's input steps and a launch's destinations by position (_position_groups), and the pass visits the targets' groups alone. Its work arrays indexed by launch child, destination or product row were allocated and zeroed every call (46 MB): they live on the top-down program (_mode_workspace) and a call resets the entries the call before wrote. Both step builders then looped in Python over every step of the program (2048 top-down, 1024 forward) though all but the target's had nothing to do: the modes' counts give the active steps by prefix sums (_top_down_active, _forward_active) and the builders visit those, the idle steps a template built once. 28 -> 2 ms a call at n = 1024. The liveness sweep of an automaton wider than 4096 columns ran a Triton launch per boundary and direction (2048 a call at n = 1024, 15 us of host time each; the CUDA sweep in shared memory takes narrower automata). kernels/c/ liveness.cu gets the wide sweep, the rows in global memory, launched cooperatively with up to 16 CTAs a sample (a grid barrier a boundary; single-CTA kernels where the grid cannot be resident, the Triton launches without the extension). Forward: the observed prefix is the leader CTA's alone, through a frontier of live columns in shared memory (one column a boundary); from the first missing token every CTA takes a slice of the row -- where one CTA spent its time: 49k columns' successor lists a boundary over ~100 boundaries on an automaton of 16 classes (ex110), which had cost its first calls 12 ms against 4.7. Backward, after the forward: the forward-live columns alone (a live column's successors under the classes allowed are live, and the AND keeps live columns alone), rows scanned 16 bytes at a time, the dense boundaries with every CTA and the sparse prefix with the leader. Saturation, with every class allowed: a full forward row stays full through missing tokens (a state active at t + 2 is reached from one active at t + 1), so once the row is full and nothing is observed ahead the rest is written at once; backward, the rows after the last observed position are full and the sweep starts there. A sample's own lists (its inputs rule classes out) take no shortcut. 19 -> 3 ms a call at n = 1024, 2.2 -> 0.7 at n = 128. Sessions, warmed, against f592e13 (toast GPU 7 / 6): Kubernetes_171 Codpula n = 256 B = 4 l2r 3144 -> 897 ms (a call 6.6 -> 2.8 ms); HMM n = 128 B = 4 16156 -> 15889 (4.3 -> 2.4 a call, call 0 the rest); HMM n = 1024 B = 4 415 -> 327 s (42 -> 10 a call, the pool's segments and call 0 the rest); ex110 HMM B = 8 2456 -> 2235 (4.2 -> 2.4); hard_0 HMM 750 -> 705 (2.3 -> 1.9), Codpula 314 -> 248 (2.0 -> 1.4); hard_0 diffusion B = 8 (an automaton the shared-memory sweep takes) 9837 -> 9941, noise. Tests: liveness_test's sweeps fixture and the per-boundary equality test take the wide CUDA sweep as a fourth arm (brute force, per-sample classes, every evidence pattern); the planner tests cover the plans' equality to their references.
…ssor when the boundaries' live classes outnumber their successor pairs eightfold, the query's masses within a share of the device's memory and its walks not pruned per sample -- JSON hard_0 x HMM 4096: a diffusion decode of 32 calls 9985 -> 3049 ms, the input @ block kernel 107 -> 2.4 ms a call The product kernels sum a missing token's input over its live classes (hard_0: 30 a column, Kubernetes_171: 47), or over the column's distinct successors (2-3) through transition masses built once per computation of the class masses: log-sum-exp of the classes leading to each successor. Grouping was chosen when a column's live classes, on average, were GROUP_MIN_RATIO times the MOST successors of any single column: a per-column maximum (85 on hard_0, 67 on Kubernetes_171) that no average meets, so JSON Schema automata were never grouped though their columns branch into 2-3 successors. grouped_worth compares totals: the boundaries' live (column, class) pairs against their (column, successor) pairs, the work the two loops do; the ratio stays 8, calibrated on automata whose every column branches alike (where the two readings agree). The masses take a float per (input step, successor pair): node blocks x positions x pairs, 1.2G floats (4.7 GB) for hard_0 x HMM 4096 at n = 128, against a fixed budget of 2^26. transition_budget allows TRANSITION_MEMORY_SHARE (a quarter) of the device's free memory at program build when that is more, and a sample's masses stay under 2^31 floats (the steps' offsets into them are int32). The choice is the query's (Program.grouped_for, through Program.transitions, which the drivers call before the liveness, whose successor lists follow the tables): grouped when its class masses are one set for every sample and fit, else the per-class kernels through the ungrouped tables, built when first needed. A query whose inputs prune the walks per sample (a session's soft evidence that rules classes out: its successor lists are each sample's live ones) takes the per-class kernels whatever the budget: the walks touch each sample's live columns, while the masses would be built for every column's pairs, a sample at a time (hard_0 x Codpula, B = 4: a diffusion step 5x slower grouped) -- GROUP_PRUNED_WALKS, off, lets such a query be grouped when its samples' masses fit (the tests of the per-sample masses turn it on). Both table sets are the program's; the launches carry the masses' offsets either way; the masses buffer is a float under the per-class tables (Kubernetes_171 at n = 128 would otherwise ask 119 GiB for masses no kernel reads); the step groups of the per-class kernels are built for every program and used by them alone; the top-down kernels read the query's tables (trans_top_down a property), the CUDA input @ block walk its own ungrouped ones. Fixed on the way: the walk kernels reached a sample's block of the masses by an int32 product of the sample and the stride (prod.py, prod_top_down.py), wrong past 2^31 floats of per-sample masses (hard_0 x Codpula at B >= 2 under grouping); the masses kernel already used int64. Measured (croissant GPU 7, warmed; against d47b3b6): hard_0 x HMM 4096 diffusion B = 8 9985 -> 3049 ms; l2r B = 4 705 -> 570 (call 0 467 -> 343); hard_0 x Codpula diffusion B = 4 967 (the per-class kernels, as before), l2r 248 -> 242; Kubernetes_171 HMM n = 128 l2r 15889 -> 15888, Codpula n = 256 897 -> 892 (the per-class kernels: 32G floats of masses would not fit). Outputs of the grouped and per-class kernels on hard_0 agree where finite to 3.7e-4 under the bf16 default (the log-sum-exp reordered; 1.3e-5 at fp32 on Codpula B = 4 with per-sample masses, forced), with the same -inf set and no NaN. Tests: enumerated_test -- an automaton of 40 classes and 2 successors a state groups, one of 8 classes and ~5 successors does not; its tests of the grouped kernels under per-sample evidence turn GROUP_PRUNED_WALKS on; softevi_categorical_test -- a marginal without inputs is grouped (the masses buffer its size), a session of 3 rows' evidence that rules classes out takes the per-class kernels (the ungrouped tables built only then, a one-float buffer), the same values, the marginal again grouped and not under a budget below its masses, the session grouped once pruned walks may be and its 3 samples' masses fit, the same values; forward_test's budget test zeroes the memory share; the grouped / classes fixtures of the existing tests cover the kernels. The constraint tests (8283) pass on croissant.
… masses are each sample's own over the entries its inputs leave possible, built once per session -- JSON hard_0 x Codpula 4096: a diffusion decode of 32 calls at B = 4 966 -> 827 ms, at B = 8 1259 -> 974; a call 29.3 -> 24.1 / 37.7 -> 27.5 Grouped transitions (d9f8df2) sum a missing token's input over a column's distinct successors through transition masses: one per (input step, pair) when the class masses are one set for every sample. A decoder session's soft evidence makes them each sample's own, and the grouped kernels then needed a sample's masses over EVERY pair of the boundary: hard_0 x Codpula at B = 4, 4 x 1.17G floats, 19 GB built at call 0 (83 -> 1303 ms, the classes' masses multiplied against the pairs for all 981 classes), for walks pruned to the sample's live columns -- so such a query took the per-class kernels (31 ms a call), though grouped its walks were 7 ms a call cheaper (the input @ block walk 4.0 -> 1.5 ms, the top-down one 7.7 -> 1.4). The masses are now laid out over the sample's successor entries alone (kernels/transitions.py SampleMasses): the (column, successor) pairs some class the sample's inputs allow leads along -- the per-sample successor lists a session already builds (live.SampleSuccessors, which now carry every program entry's rank among the sample's, -1 for none, and the inverse) -- few against the boundary's pairs when the inputs rule most classes out. Sample b's masses of a step of rank r among its position's are the count[b, t] floats at base[b, t] + r count[b, t]; one set of masses lies the same way over the boundary's pairs (pair_off[t] + r pairs[t]; the launch rows' seventh field is the step's rank now, the same for both). sample_transition_masses builds them from per-sample lists of the (entry, class) pairs the sample allows (the program's, live.entry_classes, against its allowed classes: 3.3 a kept entry on hard_0, where hundreds of classes lead along an entry and a sample's inputs allow 24 of 981 at a position): for each entry the log-sum of the row's masses of those classes, a stored class's from the row's table, an enumerated class's from its token's probability and weight -- 12 ms for 4 samples, 23 for 8, once per session (a loop over every class leading along an entry, masked, took 400 / 800). The four grouped walk kernels take the layout through Program.mass_args: the step's block from base and count by the lane's sample, the entry's place its rank there (one dependent load more than the pair's number). The query's choice (Program.transitions) is by size: a sample's layout's floats against the budget, or trans_size for one set; the layout alone decides (0.2 ms), the class lists come once a query takes it; GROUP_PRUNED_WALKS (the rule that pruned walks take the per-class kernels) goes. Kubernetes_171 x Codpula at n = 256 stays with the per-class kernels: a sample keeps 5.3k entries of a boundary's 82k pairs, 22G floats over 4 samples and 4096 nodes, past the budget. Measured (croissant GPU 7, warmed; against d9f8df2 on the same GPU): hard_0 x Codpula diffusion B = 4 966 -> 827 ms (a call 29.3 -> 24.1, call 0 82 -> 93: the masses built), B = 8 1259 -> 974 (a call 37.7 -> 27.5, call 0 118 -> 136); l2r B = 4 242 -> 258 (call 0 63 -> 78, a call 1.4 either way); Kubernetes_171 Codpula n = 256 l2r 892 -> 908 (call 0 159 -> 170: the lists' rank map); hard_0 x HMM diffusion B = 8 3050 -> 3051 (one set of masses, as before). Outputs of the grouped and per-class kernels on hard_0 x Codpula agree where finite to 7.6e-6 (B = 4) / 1.1e-5 (B = 8) at fp32, with the same -inf set and no NaN. Tests: liveness_test -- the rank and entry maps against a brute force over random automata; enumerated_test -- a session's per-sample masses equal a float64 brute force per (sample, step, entry), the enumerated classes' included, an entry no allowed class leads along having none, one set of masses as before; softevi_categorical_test -- a session whose evidence rules classes out takes the per-sample layout (fewer floats than three sets), the per-class tables under a budget below it, the same values either way; forward_test's masses check reads the rank-based offsets. The constraint tests (8284) pass on croissant and brezel.
…extension, its pre-pass a CUDA kernel launched with it -- JSON hard_0 x Codpula 4096, a diffusion decode of 32 calls: B = 4 827 -> 729 ms, B = 8 974 -> 898; a call 24.1 -> 21.1 / 27.5 -> 25.5 A Codpula diffusion step was host-bound: the host took 22-23 ms to issue a call's ~1160 kernels (B = 4 and B = 8 alike, host_probe.py) against 17-22 ms of GPU work; the wall was the host's time and a tail. Nine kernels a position, a Triton launch ~15 us of host time through the fast launcher and a driver call 20-35 us with its Python: forward the fused sum's pre-pass (Triton), the input @ block walk (Triton), the sum and its split merge (CUDA, one call); top-down the pre-pass, the walk, the sum and merge, the class weights. The pre-pass -- the children's chunks (top-down: the parents') shifted-exponentiated and laid out in the fragment order the tensor cores read -- is now a CUDA kernel in kernels/c/fused_sum.cu (prepass_kernel: a CTA per chunk and 64 columns, the values staged in shared memory, their column maxima, e written contiguously by the inverse of the fragment mapping, on the TF32 grid or in bf16 as the products read it; top-down less the column's sample's log Z, the sample found as the sum kernel finds it), launched by fused_sum_cuda ahead of the sums on the same stream: one Python call per fused sum, the Triton pre-pass kernels left to the launches without a CUDA plan. sum.py's _cuda_forward takes the operand and its tables for it; the per-sum numpy isin over the held intervals is a set lookup. Tried and dropped: merging a split's shares in the sum kernel by the last share in -- a tile's 16k outputs on one CTA where the merge kernel spreads them over 64, the sums 4.7 -> 10 ms a call. Measured (croissant GPU 7, warmed; against 041de0f on the same GPU): hard_0 x Codpula diffusion B = 4 827 -> 729 ms (a call 24.1 -> 21.1), B = 8 974 -> 898 (27.5 -> 25.5; the host's issue time of a call 22.9 -> 18.6 ms, against ~23 ms of GPU work now: the step is no longer host-bound at B = 8, and balanced at B = 4); l2r B = 4 258 -> 250; hard_0 x HMM diffusion B = 8 3050 -> 3050 (its sums are cuBLAS's). Launches a call 1158 -> 906. Tests: cuda_sum_test -- a CUDA launch is one call: the CUDA pre-pass's events and no Triton pre-pass beyond the launches without a plan, the merge kernel's from the same call, the values the Triton kernels' (the existing equality tests cover the pre-pass in every precision, split and pruning). The constraint tests reaching the sums (6520: cuda_sum, blockscale, forward, top_down, softevi_categorical, conditional, decoder) pass on croissant.
…s -- a pruned launch of few live columns runs two CTAs an SM, its chunks in twice the shares; JSON hard_0 x Codpula 4096, a diffusion decode of 32 calls: a call at B = 8 27.5 -> 23.5 ms (its GPU work 24.5 -> 20.0, the sums 10.3 -> 6.0), at B = 4 23 -> 21 (19.4 -> 15.2) A Codpula diffusion step's fused sums (kernels/c/fused_sum.cu) are 252 launches a call, each streaming the tied 33.6 MB bf16 weight copy against 8 live columns (one a sample at B = 8). The kernel took tiles of 64 columns at the narrowest: 32 tiles of 4 node blocks, split 5 ways -- 160 CTAs on 188 SMs, one CTA an SM, a single wave of 26 chunks each, 38 us a launch. Timestamps in the kernel (a scratch build) put a chunk at 1.1-1.2 us of which the wait for its copies is 0.06: the copies are hidden, deeper pipelines bought nothing (measured: 3 and 4 stages equal), and the time is the chunk's instruction issue -- some 820 SASS instructions a warp a chunk for 15 async copies and 4 tensor-core products (64-bit address math and predicates recomputed every chunk, runtime divisions, argument fields reloaded), nearly the same at 16 columns as at 64 -- plus a prologue of 6.4 us (where_of's chain of dependent loads in one thread). An SM running one CTA's chain is half idle: a second CTA on the SM runs its chain at nearly the same speed. A tile of 16 columns takes 47 KB of shared memory and 64 registers a thread, so two CTAs share an SM, and _cuda_split (whose residency is the extension's ctas_per_sm) puts a tile's 128 chunks in 11 shares: 352 CTAs in one wave of 12 chunks. Replayed on the captured hard_0 launch: the sum kernel 37.7 -> 20.7 us (merge 5.5 -> 4.3, the pre-pass 4.7 unchanged); at 16 columns with one CTA an SM (more stages, or 2 waves) it stays at 36, and 16 or more shares lose to a second wave. The kernel's pass 1 (a thread's one or two (block, column)s) lets threads past the tile's idle; the dispatch and ctas_per_sm take the two widths. Tiles of 64 and wider are as before: _cuda_shape picks the narrowest width covering a launch's columns in as few tiles as 128 would, so a launch of more than 32 columns a tile is unchanged. Measured (croissant GPU 7, warmed, diffusion_profile over 32 calls, against the same tree at 64-column tiles): hard_0 x Codpula B = 8 a steady call 27.5 -> 23.5 ms, call 10's GPU work 24.5 -> 20.0 (fused_sum_kernel 10.3 -> 6.0 ms, merge 1.4 -> 1.1); B = 4 a call 23 -> 21, GPU 19.4 -> 15.2 (the sums 9.5 -> 5.5; the call is host-bound there at ~18.6 ms). hard_0 x HMM unchanged (cuBLAS sums). Outputs against eaa6c1f (digests of every call): HMM hard_0 diffusion bit-identical; Codpula hard_0 diffusion B = 8 / B = 4, l2r, and Kubernetes_171 l2r at n = 256: the same -inf sets, max |diff| 7.6e-6 / 1.1e-5 (11 shares instead of 5 fold the fp32 partial sums in another order). Tests: the constraint tests reaching the sums (6520: cuda_sum, blockscale, forward, top_down, softevi_categorical, conditional, decoder) pass on croissant, the small launches of the suite now on 16- and 32-column tiles.
…copies from descriptors settled once -- a chunk 1.1 -> 0.73 us at 16 columns, 1.2 -> 0.79 at 64; JSON hard_0 x Codpula 4096, a diffusion decode's sum kernels: B = 8 6.0 -> 5.0 ms a call, B = 64 (tiles of 80 and 64) 5.6 -> 4.7 A chunk of the fused sum kernel (kernels/c/fused_sum.cu) waits for its copies, issues the next chunk's, and runs its products and fold. Timestamps in the kernel (a scratch build) on the Codpula hard_0 launch put the wait at 0.06 us of a 1.1-1.2 us chunk, the products and fold at 0.3-0.4, and the ISSUE of the next chunk's copies at 0.75: load() recomputed every copy's address, bounds and validity from scratch each chunk -- the block's liveness from shared memory, the weight's source behind a runtime test of the affine layout, the e copy's column and its bounds test, the t copy's byte count, the scale copy's block and lane by a division by the runtime q4, and the kernel's arguments reloaded -- some 820 SASS instructions a warp a chunk for 15 async copies and 4 tensor-core products, the same at 16 columns as at 64. The invariants of a thread's copies are now settled before the loop (a `settle` lambda): a mask of the live blocks, the thread's place in a weight tile folded into every block's first weight (affine: advanced by wstride a chunk, else the share's table), its e copies -- E_COPIES of 16 bytes THREADS copies apart, a bit each for "within the tile and its column one of the launch's" --, its t copy's bytes and columns, its scale copy's block and samples (or, by column, the columns' samples) and mode. A chunk's load adds the chunk's offsets to them and issues; a factor lane's two shifts T are read once into registers past pass 1 (one more barrier a CTA). The copies land where they did with the bytes they did, the products and fold run in the same order: the outputs are bit-identical. The three-pass (precision fp32) products past 64 columns and the fp32 operands at 128 sit at the register cap -- settled once, their descriptors spilled (up to 136 bytes a thread; tiles of 128 columns 15% slower) -- so those variants (DESC false) settle per chunk, as before, and keep their registers. Measured (croissant GPU 7 vs 9a9d5f8 on the same GPU; the per-chunk split from the timestamped build): a chunk at 16 columns 1.07 -> 0.73 us (the issue 0.78 -> 0.52, products 0.25 -> 0.16), at 64 1.17 -> 0.79 (0.75 -> 0.38); the hard_0 B = 8 launch's sum kernel 20.7 -> 16.3 us at 16 columns (11 shares, two CTAs an SM), 37.7 -> 26.2 at 64 (5 shares); a 16-column CTA's loop now moves its ~9 KB a chunk at the L2 -> SM copy ceiling with two CTAs an SM, and its prologue (where_of's dependent loads, pass 1: 5.9 us) is 40% of it. Codpula hard_0 diffusion, 32 calls: B = 8 fused_sum kernels 6.0 -> 5.0 ms a call (GPU 20.0 -> 18.9), B = 4 5.5 -> 4.5 (15.2 -> 14.2), B = 64 (tiles of 80 and 64) 5.6 -> 4.7 (137 -> 133.5); precision fp32 (three-pass) B = 8 8.3 -> 7.7, B = 128 (tiles of 80 and 128, settled per chunk) 38.0 -> 36.5, B = 96 (tiles of 112 and 96) 27.6 -> 26.7, B = 112 (128 and 112) 29.1 -> 28.4; precision tf32 B = 128 23.0 -> 20.4. The B = 4 and B = 8 walls (21 / 23.5 ms) are the host's now: it issues a call's 906 launches in ~18.7 ms. Outputs against 9a9d5f8 (digests of every call): bit-identical on hard_0 x Codpula diffusion B = 8 and B = 4, l2r, hard_0 x HMM diffusion, Kubernetes_171 x Codpula l2r at n = 256. Tests: the constraint tests reaching the sums (6520: cuda_sum, blockscale, forward, top_down, softevi_categorical, conditional, decoder) pass on croissant.
…under a JSON Schema and under a regular expression tests/constraints/examples/json_schema_regex_decoding_test.py: a small vocabulary made in-test (every printable ASCII character, a few longer pieces, an end-of-sequence id), an HMM over it, and token-by-token decoding under jc.JSONSchema (an object with a bounded string, an optional integer and an enum) and under jc.Regex (a YYYY-MM-DD date); every drawn text is checked with json.loads and the schema's terms, or re.fullmatch, and by the constraint's accepts. Skipped without interegular, as the Regex and JSON Schema tests are. Both pass on croissant (10 s and 3 s).
…s its walk over the rows -- JSON hard_0 x Codpula 4096, a diffusion decode of 32 calls: B = 8 900 -> 808 ms, B = 64 4538 -> 3685 (the scoring kernels 36.4 -> 11.8 ms a call) The conditionals' scoring (kernels/score.py) runs a program per column and tile of lanes -- stored classes in _score_weights_kernel, candidate tokens in _score_candidates_kernel, enumerated tokens in _score_tokens_kernel, a class's tokens for a tile of columns in _score_kernel --, each walking the K rows of the column's class weights. Most programs have no live lane: a column has few live stored classes (Codpula hard_0 B = 8: 3.7 of 70 on average, none for 736 of the 1024 target columns -- every sample's column of a position some sample observes is a target), and a candidate token's class is often dead at the position. Their walks loaded nothing (every load masked off) but ran the loop all the same: 64 tiles of rows, two passes in the shifts' kernel. Such a program now walks no rows: _score_weights_kernel's row range is empty when no class of its tile is live, the others pass K = 0 to their dot (_log_dot) or loop when no lane is on. What the walk would have left is what the program starts from -- -inf shifts and parts, zero sums -- so it writes the same values; a live program runs the same code as before: the outputs are bit-identical. Measured (croissant GPU 7 against 8712aad on GPU 6, warmed, diffusion over 32 calls; a call's scoring kernels replayed alone for the per-kernel times): Codpula hard_0 B = 8 the shifts' kernel 3.40 -> 0.78 ms a call, the candidates' 1.13 -> 0.89, a decode 900 -> 808 ms (a steady call 24 -> 21); B = 4 767 -> 730; B = 64 the kernels 26.9 -> 5.3 and 9.5 -> 6.5 ms, a decode 4538 -> 3685 (a call 150 -> 124); HMM hard_0 B = 8 the shifts' kernel 1.29 -> 0.85, the class-tile scores (_score_kernel) 3.16 -> 2.53, its enumerated tokens' unchanged (every lane live). Compacting the shifts' lanes to a column's live classes would take B = 8 to 0.47 ms, but changes the reduction's layout and so the classes' parts of p(e, C) in their last bits: not done. Outputs against 8712aad (digests of every call): bit-identical on hard_0 x Codpula diffusion B = 8 and l2r, hard_0 x HMM diffusion and l2r, Kubernetes_171 x Codpula l2r at n = 256. Tests: tests/constraints and strong_check_test --slow pass on croissant.
…live pairs share one row of class weights -- JSON hard_0, a diffusion decode of 32 calls at B = 8: HMM 4096 3036 -> 2178 ms (peak memory 36.6 -> 28.0 GB), Codpula's GeneralizedHMM without per-sample inputs 3066 -> 2145 (42.0 -> 33.4) A conditional query's class weights (kernels/score.py) are a row of K floats per target column (a sample at a target position) and live class: the class's weight for every input node of the position, the sum over the class's (column, successor) pairs whose ends are both live (_class_weight_kernel). Under evidence most classes of a boundary keep the same few live pairs, and each was walked and stored on its own: HMM 4096 x JSON hard_0, B = 8, 912535 live classes of the 1024 target columns had 267218 distinct live pair lists at a decode's call 1, 296742 and 6824 at call 10, 258961 and 1741 at call 20; Codpula 4096 6643 and 1256 at call 1. The class weights were 228 ms of call 1's 527 ms of GPU work, 34 ms of call 10's 54. Classes with the same live pairs now share one row, that of the first of them in the boundary's class order (their representative). _live_classes_kernel keys every live class: the sum of its live pairs' hashes (splitmix64 of the pair's id), mixed with their number and with whether the class is stored or enumerated. class_weight_rows sorts each column's keys in class order and proposes the first class of a key as the representative of the others; _class_reps_kernel checks every proposal exactly -- as many live pairs, the same kind, and every live pair (i, q) of the class taken by the candidate too (its successor from column i is q) -- and a class failing it keeps a row of its own, so a collision of keys loses a merge and never changes a value. _class_rows_kernel gives the representatives consecutive rows and points every other live class at its representative's. CLASS_ROWS_SHARED turns it off. The scoring exponentiates a stored class's row in place: with shared rows _score_weights_kernel runs twice (SPLIT), every class reading its row for its shift and its part of p(e, C), then every row exponentiated once, by its owner. A shared row holds what each of its classes' rows held: the walk sums a class's live pairs in the same order whichever class it walks. Measured (croissant, against 8712aad with the scoring skips, warmed, diffusion over 32 calls): HMM hard_0 B = 8 a decode 3036 -> 2178 ms, its class weights call 1 228 -> 156 ms, call 3 98 -> 34, call 10 34 -> 4.6, call 20 25 -> 2.2, a call's GPU work 527 -> 441 ms (call 1), 54 -> 21 (call 10), the steady calls 50 -> 28 ms; Codpula's GeneralizedHMM 4096 with no per-sample inputs (its base weights, V = 126464) hard_0 B = 8 3066 -> 2145 ms; the peak memory 36.6 -> 28.0 GB and 42.0 -> 33.4, the rows' buffer. Codpula hard_0 B = 8 771 -> 767 ms (its class weights 2.1 -> 1.3 ms a call; the call is the host's), B = 64 3728 -> 3532 (9.4 -> 3.5 ms). Outputs against 8712aad with the scoring skips (digests of every call): bit-identical on hard_0 x Codpula diffusion B = 8 and l2r, hard_0 x HMM diffusion and l2r, Kubernetes_171 x Codpula l2r at n = 256. The rows shared and not in one process, a diffusion decode of 8 calls at B = 8 on a 60-state automaton whose 40 classes take 8 transition columns (3% of the entries changed): PD 16 (n = 32) 72280 live classes in 56285 rows, Codpula's GeneralizedHMM 256 72183 in 56205, every call's conditionals bit-identical; on a random automaton (75988 in 75051) PD, HMM and Codpula too. Tests: conditional_test checks the rows against a brute force with the rows shared and not, and the weights with every key the same (each class a candidate for every other) over HMM, PD and hand-built kinds, dense and compact; softevi_categorical_test's score-weights paths run with the rows shared and not. The full suite passes on brezel (10242 tests, with the scoring skips), strong_check_test --slow on croissant (82).
…n generated for its call pattern and hands Triton each tensor of the current device as its address -- a launch 22 -> 15 us of host in a dfa_exact decode call; JSON hard_0 diffusion decodes of 32 calls: Codpula 4096 B = 4 681 -> 619 ms, B = 8 763 -> 688, HMM 4096 B = 8 2187 -> 2127; plain PC steps at batch 16 -9% FastJITFunction3x (utils/kernel_launcher.py) relaunches a cached compiled kernel: _gather binds the call's arguments to the parameters in a Python loop over every keyword, _signature builds the cache key over every parameter, and CompiledKernel.run gets the tensors -- whose launcher in Triton 3.7 (extractPointer, backends/nvidia/driver.c) calls each tensor's data_ptr() and asks the driver for its device pointer (cuPointerGetAttribute), for every tensor argument of every launch. The dfa_exact walks pass some 70 arguments, 44 of them tensors, most by keyword: a launch took 22 us of host in a Codpula hard_0 decode call (393 launches, 8.6 ms of the call's 19.3 ms issue). A launch now binds through a function generated once per call pattern -- how many arguments it passes by position, and the names it passes by keyword, in order --, which unpacks the arguments and the keyword values into locals named as the parameters, fills in the defaults, and builds the very signature _signature builds: the cache keys are unchanged. A pattern the generic binding does not bind plainly (a name it does not know, one given twice, a required parameter missing) keeps _gather, as before. The values CompiledKernel.run gets carry each tensor of the current CUDA device as its address -- the int its signature read --, which Triton's launcher takes as it is: for memory of the current device the driver's device pointer is that address, so the kernel gets the same arguments. Any other tensor -- a CPU tensor (PyTorch's pinned memory is one), or a tensor of another GPU, which the driver refuses in the current context when that GPU's memory is not mapped there -- is passed as it is, and Triton's launcher translates or refuses it as before. Nothing else changes: the cache, the kernels, the patterns left to the generic binding, callable grids and PYJUICE_FAST_LAUNCH=0 (the stock launch). Measured (croissant: the decode arms alternated on one GPU with nothing else on it; brezel: the plain PCs): a launch of 44 pointers, 14 ints and 15 constexprs, hot, 17.0 -> 13.1 us (Triton's run 6.4 -> 3.6); in a decode call a launch 22 -> 15 us (Codpula hard_0 B = 8), 13 -> 10 (HMM hard_0 B = 8). Decodes of 32 calls, two rounds each: Codpula B = 4 682 / 680 -> 620 / 618 ms (a steady call 19.8 -> 17.6), B = 8 754 / 772 -> 687 / 689 (20.7 -> 18.5), HMM B = 8 2193 / 2181 -> 2119 / 2134 (28.5 -> 26.0). A plain PC's eager step (forward, backward, mini-batch EM): HMM of 1024 latents B = 16 5.71 / 5.61 -> 5.17 / 5.09 ms, B = 256 8.62 -> 8.62, HCLT of 256 latents B = 16 1.81 / 1.84 -> 1.67 / 1.66, B = 256 1.99 / 2.01 -> 1.87 / 1.88. Checked: a scratch build of the launcher checked every fast launch of the full suite against the generic binding -- 2.6 million launches of 4125 kernels: the same signature, the same values, and each of the 28.4 million addresses it handed over equal to the driver's device pointer for it (what Triton's launcher would have passed) --; none differed, and no launch fell to the generic binding. Outputs: five decoder sessions' digests bit-identical (hard_0 x Codpula diffusion and l2r, x HMM diffusion and l2r, Kubernetes_171 x Codpula l2r); with PYJUICE_AUTOTUNE=0, plain PCs' log-likelihoods (an HMM and a PD at batches 16 and 256, an HCLT at 16) and a constrained decode's conditionals bit-identical -- the HCLT's at batch 256 differ from run to run on either tree alike. On two GPUs, a PC's and a constrained decode's outputs on cuda:1 equal cuda:0's, and a PC on cuda:1 queried with cuda:0 current is refused as before. Tests: kernel_launcher_test checks the binders against the generic binding over 400 call patterns (positional and keyword mixes, defaults, meta-parameters, CPU and CUDA tensors, unaligned views, the tensors' device current or not) and the patterns left to it, the results against triton.jit for every parameter kind, a CPU tensor refused as on the stock path, and -- in a child process with two GPUs -- launches with the other device's tensors, which must end as the stock launch's: the test fails when every CUDA tensor is handed over as an address, which ran such a launch and returned wrong values. The full suite passes on brezel (10244 tests), strong_check_test --slow on croissant.
Brings in #62: the sum layers' precision option, the cuBLAS dense path, the tf32x3 backward, and the shared choice of the tlmm fast path. #63, the z-correction default, is on this branch already as 1ed9d48. Conflicts: tensorcircuit.py and block_scale.py, both split into modules on this branch. Both take the branch's side. tensorcircuit.py adds #62's import of PRECISIONS and _pop_removed_precision_flags, and #62's change to the element-flow helper goes to model/inspection.py, where that helper now lives. block_scale.py's split modules already carry #63.
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
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.
pyjuice.constraints(conventionallyimport pyjuice.constraints as jc) answers a PC's queries exactly under a constraint on its variables: a regular language over token ids.What it adds
jc.DFA.contains;jc.Regexandjc.JSONSchemaover a tokenizer's vocabulary (jc.Vocabulary;pip install pyjuice[constraints]bringsinteregular);jc.LTLfformulas over sets of token ids; and their combinations with&,|,~and concatenation.marginalgives log p(e, C), andconditionalgives log p(x_t = v | e, C) over the vocabulary or over given candidate tokens. A decoder session (cc.decoder()) keeps what the queries of an autoregressive or masked-diffusion decode share, and takes per-sample inputs for the whole session: soft evidence and BlockScale gates, as CoDD uses them. Sampling is not served yet.constraints/inference/dfa_exact): the PC lifted over the automaton's states, with structural sparsity, evidence pruning, compact per-sample storage, and fused Triton and CUDA kernels. What it compiles depends only on the PC's structure and the constraint.tests/constraints, about 9,000 tests, checks the queries against brute force and reference implementations, and against float64 references (strong_check_test.py --slow).tests/constraints/examplesshows autoregressive decoding, including under a JSON Schema and under a regular expression.Outside the constraints package
model/tensorcircuit.pyinto mixin modules (build,common,cuda_graphs,external_params,inspection,learning),SoftEvidenceCategoricalintonodes/distributions/softevi/, andBlockScaleSumParamsintonodes/external_params/blockscale/.utils/kernel_launcher.py: a fast launch binds its arguments through a function generated per call pattern, and hands Triton each tensor of the current device as its address. That skips adata_ptr()call and a driver call per tensor argument; tensors on the CPU or on another GPU still go to Triton as tensors and are checked as before. A launch takes about 15 µs of host time instead of 22 µs in a decode, and plain PC training steps at batch 16 run about 9% faster. A shadow build compared all 2.6 million fast launches of the full suite against the old binding, and every address against the driver's device pointer, without a mismatch.tests/conftest.pypins each test process to its GPU when the conftest is imported, before a nested conftest can initialize CUDA. Several long tests are now marked@pytest.mark.slow.import pyjuicealso importspyjuice.constraints. There is no new required dependency, and import time is unchanged.pyproject.toml: theconstraintsextra, test-only dev dependencies (jsonschema,outlines-core), and the backend's CUDA sources as package data.Verification
After merging main, on RTX PRO 6000 GPUs with Python 3.12:
tests/constraints/language/dfa_test.pyneedsfrozendictinstalled, and passes with it.strong_check_test.py --slow: 82 passed.Known limitations