A training batch has a lot of structure: many completions per prompt in a GRPO or CISPO group, a system prompt repeated across otherwise-unrelated tasks, several SFT completions distilled from one source. Whenever two sequences in the batch share a token prefix, encoding that prefix twice is wasted compute for any token mixing layers that are causal. ART packs the batch as a radix tree instead of a list of sequences, then solves a bin-packing problem to fit that tree into fixed-length training rows (the fixed length has a maximum set by the available VRAM, which is usually what you're configuring it to for highest MFU, but you may choose to keep it smaller for sample efficiency).
Treat each tokenized sequence as a root-to-leaf path. Two sequences that agree on their first k tokens can share one node covering those k tokens; where they disagree, the tree branches. Building this shape is a small recursive procedure: find the longest run every active sequence agrees on, emit it once, then partition the rest by their next token and recurse on each partition.
def build_segments(sequences, start=0, parent=None):
active = [s for s in sequences if len(s) > start]
if not active:
return
end = start
while end < min(len(s) for s in active) and len({s[end] for s in active}) == 1:
end += 1
if end > start: # a shared run: emit it once, recurse past it
segment = emit(active, start, end, parent)
yield from build_segments(active, start=end, parent=segment)
return
for _, branch in groupby(sorted(active, key=lambda s: s[start]),
key=lambda s: s[start]):
yield from build_segments(list(branch), start, parent) # split by next token
That's the shape of it, simplified from ART's actual walk in
prefix_tree_packing.py. Two refinements matter in practice.
First, a
configurable minimum shared-segment length (64 tokens by default) keeps the tree
from paying attention-kernel and bookkeeping overhead to share a handful of
tokens.
Second, sharing currently excludes trainable tokens for simplicity of loss computation. Trainable tokens likely have different advantages and the advantage tensor is the same length as the token ids tensor (matching to their respective tokens), so you can't overlap trainable tokens without changing how this is handled. I have an idea for generalizing the losses/algorithms in ART and how to support prefix tree packing on that which will replace this and probably allow sharing trainable tokens, which I'll get to in the next month or so.
Finally, an off-by-one note, each branch keeps one duplicated, non-trainable context
token right at its join point, so it still gets its own next-token query
position when the model is asked to predict the first token that's actually
unique to the branch, since it will have a unique loss.
A simple boundary to draw is one tree per RL group — that's where shared
prefixes are supposed to live. ART draws no such boundary. Every tokenized
completion queued for a training step, regardless of which group or task it came
from, walks into one tree together. A system prompt repeated across ten unrelated
tasks becomes one shared root the same way a GRPO group's prompt does; the walk
doesn't know or care that they're “different groups” — it only
knows token sequences. The same goes for a handful of SFT completions distilled
from one source, or a batch of eval rollouts replaying the same tool schema: none
of it needs a group in the RL sense to benefit from the tree.
You can also just fill in the bins as you go, checking the size of the tree at each step, starting a new one when you can't fit the next trajectory. But it's a little more fun and there's a little more potential gain if you think about and solve it more generally.
What comes with that generality is that deciding which leaves land in which row becomes bin packing: partition a set of items with individual costs into the fewest fixed-capacity bins. Drawing an arbitrary tree-per-group boundary would have sidestepped that problem entirely, at the price of every shared prefix that happens to fall on the wrong side of the boundary. Alternatively, in-order is just a simple bin-packing heuristic. ART takes the trade the other way: pay for a real (NP-hard) bin-packing pass, get back every sharing opportunity actually present in the batch.
That trade is measurable. On a synthetic batch of 48 RL groups spread across three task families — math, tool-use, code review — all sharing one 350-token system prompt and packed into 4096-token rows, building the tree independently per group and only then bin-packing needs 66 rows; pooling everything into one global tree first needs 61 — about eight percent fewer rows purely from where the sharing boundary was drawn, with the bin-packing algorithm itself held fixed. The full historical comparison, including what ART's packing looked like before any of this existed, is at the end.
The same global-tree-then-bin-pack computation gets reused for planning, too. ART's pipeline autotuner estimates how many packed rows a candidate set of groups would need by running this exact tree-and-bin-pack pass over sampled subsets of past data, before ever committing to a training step's batch composition — the same machinery that packs a step also gets asked “what if” questions about how to build the next one.
Bin packing — fixed-size items, fixed-capacity bins, minimize bins used — is one of the oldest NP-hard problems around, so which heuristic to reach for is a real decision, not an afterthought (Claude's words, honestly probably not that important to most inputs, but fun to think about and nice to choose something pretty general). First-fit and best-fit decreasing are the standard practical answers: sort items largest-first, then drop each one into whichever open bin leaves the least leftover space (opening a new bin only if nothing fits). ART uses a best-fit-decreasing variant.
But the tree changes what “item size” even means. A leaf's total token count is fixed, but its marginal cost of joining a particular bin is not: any segment that bin already contains — because a sibling landed there first — is free. Classical bin packing has no notion of one item making another cheaper. Here it's the whole point.
class _PrefixTreeBin:
def __init__(self):
self.leaves = []
self.occupied_segments = set()
self.token_count = 0
def insertion_delta(self, leaf):
return sum(
length for segment_id, length in leaf.segment_path
if segment_id not in self.occupied_segments
)
def add(self, leaf):
self.token_count += self.insertion_delta(leaf)
self.occupied_segments.update(sid for sid, _ in leaf.segment_path)
self.leaves.append(leaf)
def _place_prefix_tree_leaves(groups, *, seq_len):
# empty_bin_cost is just the full, unshared length
ordered = sorted(groups, key=lambda g: max(leaf.empty_bin_cost for leaf in g),
reverse=True)
bins = []
for leaf in (leaf for group in ordered for leaf in group):
best_bin, best_remaining = None, seq_len + 1
for candidate in bins:
count = candidate.token_count + candidate.insertion_delta(leaf)
# this fill most heuristic (inherited from best-fit-decreasing) is a little odd, say you have two equally full bins, it
# will pick the one that shares fewer tokens. But, practically I could not
# show one generally performed better on semi-realistic workloads than the other, so we stick with this
if count <= seq_len and seq_len - count < best_remaining:
best_bin, best_remaining = candidate, seq_len - count
if best_bin is None:
best_bin = _PrefixTreeBin()
bins.append(best_bin)
best_bin.add(leaf)
return bins
Trimmed directly from ART's preprocessing/pack.py — a leaf's
segment_path is the list of (segment_id, length) pairs
on its root-to-leaf path in the global tree, and empty_bin_cost is
its full, unshared length. Everything else is ordinary best-fit-decreasing;
insertion_delta is the only place the tree shows up at all.
Worth being honest about: once item cost depends on what's already in the bin, this isn't bin packing in the strict textbook sense any more — the classical proofs behind first-fit and best-fit decreasing assume fixed item sizes. It's still NP-hard, and best-fit-decreasing over the real marginal cost is a principled thing to do, but it doesn't inherit the same tight approximation guarantee as its fixed-size ancestor. If you want the precise framing (and where a more principled algorithm might come from), see the appendix.
Step through it on the tree from the last section, plus a second, unrelated group:
A small hand-picked example (row capacity 1000 tok) — the real measured numbers are at the end. Every completion here happens to cost more than half a row on its own, so no arrangement of raw, unshared completions could ever fit two per row; every row saved below comes from the tree, not from cleverer bin packing.
Packing is useful if token mixing layers are causal, i.e. a token may attend only to its own ancestor chain, never a sibling branch or an unrelated tree packed into the same row. This is prefix caching and radix attention in vLLM and SGLang respectively, just for training.
The pattern that holds for every layer below except GDN: the attention kernel itself never changes, only the input we build for it does. For ordinary masked attention, flex/flash kernels already accept an arbitrary mask specification — a token's tree-correct visible set just decomposes into one fully-visible block per ancestor segment plus a causal triangle over its own segment, exactly the two primitive shapes those kernels already special-case. Building that mask correctly is the whole job; most blocks in a packed row end up empty and get skipped entirely, and the kernel is none the wiser.
Rows/columns 0–2 are R, 3–5 are branch A, 6–8 are branch B (sibling of A). The outlined row is query position 7 in branch B: it sees all of R (a full block, since R fully precedes it) and its own branch causally up to itself — never branch A.
Same pattern, different input shape. ART's sparse top-k indexer attention
variants (GLM-5.2- and DeepSeek-V4-style) don't take a mask — their
attention kernels (sparse_mla, sparse_attn_tilelang)
already just gather and attend over whatever key/value indices they're handed,
because that's what top-k sparse attention over long context looks like whether
or not anything is packed. So the kernel doesn't actually have to change here (nice not to have to work on this). All the
tree-awareness lives one level up, in the indexer: it only ever scores or
constructs candidates from a query's actual ancestor chain, using the same
pre-order/exit-order (these words relate to how the mask is constructed, from a DFS, an increasing int is assigned, you can check your ancestors by an [enter, exit] interval) interval check as the mask above, so a sibling branch's
keys are never even proposed, let alone selected.
GDN (Gated DeltaNet) is the layer that needs a more significant change. It's recurrent
— there's no mask to compute, only a running state that has to be threaded
token by token, forked once per branch point, and reset at every tree root.
ART's GDN planner walks the tree depth by depth; at each depth, every segment
across every tree in the packed batch — regardless of which row or group it
came from — gets coalesced into two kernel calls: one varlen causal-conv
call and one varlen delta-rule call, using the same flash-style
cu_seqlens layout an unpacked batch would use.
Codex can handle this kind of work, writing custom kernels and efficient work between kernels pretty well when it has built a solid set of input shapes (workloads) and precise profiling, otherwise I probably wouldn't be building this.
# the shape of it: one pass per tree depth, two kernel calls per pass
state_cache = {}
for depth, bucket in enumerate(tree_segment_buckets_by_depth):
init_conv, init_state = fork_parent_states(bucket, state_cache) # zero-copy expand
conv_out = causal_conv1d_varlen(bucket.qkv, init_conv, cu_seqlens=bucket.cu_seqlens)
rec_out = chunk_gated_delta_rule(conv_out, init_state, cu_seqlens=bucket.cu_seqlens)
state_cache[bucket.family_indices] = (conv_out.final_state, rec_out.final_state)
scatter(rec_out, rows=bucket.row_indices, positions=bucket.position_indices)
A child's initial state is a zero-copy broadcast view of its parent's final
state — a plain .expand(), so autograd sums gradients from
multiple children back into it for free; root segments start from a zeroed
state. One wrinkle: the delta-rule kernel can only fork state cleanly at a fixed
chunk boundary (64 tokens), so a branch point whose length isn't a multiple of
64 gets its ragged tail replayed once per child — cheap, at most 63 extra
tokens per child. (This gets substantially hairier once context parallelism is
involved — cross-rank state exchange, chunk-boundary alignment across
ranks — but we'll discuss some of that in a future CP post.)
None of the planning above touches a real tensor. Segment boundaries, group and parent ids, bin assignments — the entire packing plan is computed over plain integers: sequence indices and lengths. Only once a bin's final leaf set is fixed does ART do a single materialization pass that slices real token ids, logprobs, advantages, and (when present) MoE routing-replay data into the packed tensors.
segment ids, lengths, group/parent ids, bin assignment — integers only
tokens, logprobs, advantages, routing replay — sliced into place exactly once
Large per-token payloads, like MoE routing-replay bundles, are kept as lazy segment lists until that pass and sliced exactly once each.
Numbers below replay the original prefix-tree placement rule from the first
version of the PR against ART's current global-occupancy placement
(_place_prefix_tree_leaves, prefix_tree_pack_segments).
The input is 32 agent rollouts in four groups of eight, packed into
131,072-token training rows. Every rollout starts with the same 4,096-token
system and tool prefix, then a 12,000-token scenario prefix unique to its
group. The groups arrive with completion lengths of 2,000, 2,000, 12,000,
and 7,000 tokens respectively: short A, short B, long, then medium.
Worth knowing what came before, too. The very first version (we called it “shared prefix,” not “prefix tree,” since sharing was exactly one level deep and scoped to a single RL group) deduplicated the largest prompt shared by a trajectory group — that goes pretty far for GRPO, but doesn't cover any of the tree structured cases above. The version that replaced it kept the same in-order, one-row-at-a-time filling until the row was full, but rebuilt a real multi-segment tree fresh for whatever already happened to be in the current row — measured below as “original in-order tree.” I note the very first version simply to give some vision on alternatives you might build.
| arrival order | group-specific prefix | completion length | rollouts |
|---|---|---|---|
| 1. short A | 12,000 | 2,000 | 8 |
| 2. short B | 12,000 | 2,000 | 8 |
| 3. long | 12,000 | 12,000 | 8 |
| 4. medium | 12,000 | 7,000 | 8 |
Each row also pays for the 4,096-token system/tool root once. Prefix and completion lengths exclude that common root. For next-token loss alignment, ART leaves the final prompt token on each rollout leaf: the exact physical accounting is therefore 11,999 shared group-prefix tokens plus completion-length + 1 tokens per rollout.
| placement | row contents | used tokens |
|---|---|---|
| in-order row 1 | short A + short B + half of long | 120,113 / 131,072 |
| in-order row 2 | remaining long + 7 medium | 125,105 / 131,072 |
| in-order row 3 | final medium rollout | 23,096 / 131,072 |
| global row 1 | long + 3 short A | 130,105 / 131,072 |
| global row 2 | medium + 5 short A + short B | 122,114 / 131,072 |
The original packer permanently moved on whenever the next rollout did not fit, so the last medium rollout stranded almost an entire row. Global placement sees every leaf first, puts the long group down before the short groups, and fills the remaining space deliberately. The same 32 rollouts use two rows instead of three: allocated training tokens fall 33% (393,216 → 262,144). The tree and training semantics are unchanged; only row placement differs. The example is not super unrealistic, though I think often the in-order version does perfectly fine (equivalent-ish to the current version).
Packing quality is one question; planning cost is a separate one. Using the same multi-task-curriculum batch shape as the scope comparison earlier in the post (not the 32-rollout example above), here's how the tree walk and the bin-packing search each scale as the number of leaves grows:
| groups | completions | tree walk | bin packing | rows |
|---|---|---|---|---|
| 12 | 105 | 0.9 ms | 1.4 ms | 15 |
| 48 | 450 | 9.6 ms | 4.2 ms | 61 |
| 192 | 1,668 | 42.0 ms | 48.7 ms | 227 |
| 768 | 6,686 | 291.2 ms | 743.5 ms | 918 |
The tree walk stays close to linear. The bin-packing pass doesn't — its best-fit search checks every open bin for every leaf, so cost grows closer to quadratic in leaf count. At the batch sizes ART actually trains with this is well under a second; it's the part to watch if batch sizes grow by another order of magnitude.
The overall idea for prefix tree packing is pretty simple, just some neat connections to classical problems here in a cool place. And the benefits can be very substantial. Despite that, only a few other open source RL frameworks implement it, VeRL with a plugin and AReaL that codex could find and I don't believe those support it for layers such as GDN. So, ART is probably the most advanced open version of this.
Classical bin packing — fixed-size items, fixed-capacity bins, minimize bins
used — assumes an item's cost doesn't depend on what else shares its bin.
ART's version doesn't have that property: a leaf's marginal cost of joining a bin
is insertion_delta, which depends on which segments that bin already
contains. That changes what “item size” means, so it's worth being
precise about what's preserved and what isn't.
Any classical bin-packing instance embeds trivially: give every item its own
private segment, shared with nothing. insertion_delta then reduces
to plain fixed item size, and tree-packing specializes exactly to ordinary bin
packing. The reduction is immediate; the decision problem stays NP-complete.
A bin's occupied-token count is the size of the union of every segment
referenced by its assigned leaves. Union-size is a monotone submodular
(coverage) function of the assigned set: it never decreases as leaves are
added, and the marginal gain of adding one more leaf — exactly
insertion_delta — can only shrink as more leaves that already
cover its segments get added first. So the precise statement of the problem is
bin packing where each bin's load is a monotone submodular function of its
assigned set, not an additive sum of fixed sizes.
The worst-case bounds for first-fit and best-fit decreasing (∼11/9·OPT + O(1)) are proved with arguments that lean on additivity — pairing up under-full bins, volume-counting against total item size. Those proofs don't carry over to a submodular load function. Submodularity does give subadditivity (f(S∪T) ≤ f(S) + f(T)), enough to recover some constant-factor bound via a similar volume argument, but a looser one than the tight classical result. Best-fit-decreasing over true marginal cost is a principled generalization of the standard heuristic, not one with a matching proof of near-optimality.
Sorting leaves by empty_bin_cost — standalone, worst-case size
— is the natural analogue of “decreasing,” but it's blind to
shareability. Two leaves can have identical standalone cost while one shares
heavily with whatever's already placed and the other shares nothing;
cost-sorting alone can't tell them apart, only the online best-fit search
discovers the difference after the fact. Keeping sibling leaves adjacent in the
placement order (grouping by originating tree-family, ordering families by
their most expensive member) is a partial mitigation, not a fix for the
ordering problem itself.
No proven superior ratio for this exact problem shape is being claimed here
— these are just plausible next places to look: a placement score that
weighs segment overlap alongside remaining capacity, not best-fit alone;
walking the tree in a locality-preserving order (e.g. depth-first) so segment
reuse is discovered by construction instead of by luck of the best-fit search;
or borrowing from the submodular-cover literature (Wolsey's greedy
submodular-cover result gives a 1 + ln(max cost)-type guarantee for
a related, though not identical, covering formulation). For a cost structure
this non-standard, the credible way to know whether any of these actually beats
what's here is the same way the numbers earlier in this post were produced
— measured on real or realistic batches, not reasoned out from bin-packing
theory alone.