Context Parallelism in ART

01 / The workload

What does a longer packed row buy us?

It buys room for more unique work. With a shared prompt, that can mean substantially more completed trajectories per pass through the model.

A trajectory here is one prompt followed by one sampled completion. Several trajectories can share a prompt, or share a longer conversation before branching. The model must produce the same outputs as if each trajectory had run separately. It must not let a completion read another completion. But computation on an identical, non-trainable prefix can be shared.

It helps to distinguish two lengths. Logical length is the context seen by a trajectory along its root-to-leaf path. Physical length is the number of stored token positions in the packed row. A tree with a 24K prompt and five 8K completions has 64K physical tokens, while every trajectory has a 32K logical context. Throughout the capacity examples, K means 1,024. A longer physical row is not necessarily a longer context for any individual query.

one prompt of length P, completions of lengths C_1 ... C_m

flattened storage = m * P + sum(C_i)
tree storage      =     P + sum(C_i)

The saving is not limited to storage. Tokenwise projections and MLPs process the prompt once. Dense causal attention computes its prompt-to-prompt triangle once. The completion-to-prompt rectangles still exist: every completion needs to read the prompt. Sharing removes duplicated computation, not legitimate dependencies.

Context parallelism (CP) partitions tokens among a group of GPUs. A rank is one participating process, normally one GPU. CP4 means four ranks collaborate on this packed row. The model parameters need not be partitioned by CP itself: each rank can apply the same linear layer to a different set of tokens. Other parallelisms can partition parameters at the same time.

This differs from tensor parallelism (TP), which splits a layer's channels or matrix dimensions and combines partial results. Pipeline parallelism (PP) splits layers across stages. Expert parallelism (EP) splits the experts of a mixture-of-experts model and routes token activations to them. These are composable axes, not four competing names for the same operation. In many of our large-model runs, CP handles sequence-dependent work and EP provides the necessary expert-weight sharding.

The scaling question must match the workload

In strong scaling, we keep a problem fixed and ask more GPUs to finish it faster. In weak scaling, we give N times as much work to N times as many GPUs and ask them to finish in the same time, with similar memory per GPU. The second question is especially natural for training a larger collection of trajectories.

There is a crucial qualification for attention. Making one causal sequence N times longer creates roughly N squared times as many query-key pairs. That is not constant work per GPU. Our ordinary weak-scaling workloads instead add independently sized families of prompts and completions until the larger packed row fills. Family shapes have the same distribution, so token count and visible attention-pair count grow roughly together. A fixed global long-chain experiment is a different, explicitly labeled strong-scaling test.

Prefix reuse and CP therefore help in two distinct ways. Reuse reduces the work represented by a collection of trajectories. CP distributes the remaining work and grows the capacity of the row. The capacity figure explains the first opportunity; the runtime measurements later test the second. We should not infer the runtime of 29 trajectories merely from the fact that they fit.

m_flat = floor(S / (P + C))
m_tree = floor((S - P) / C), for S >= P

visible pairs in the tree:
A_tree = P(P+1)/2 + sum_i [C_i P + C_i(C_i+1)/2]

A_flat - A_tree = (m-1) P(P+1)/2

02 / Start on one GPU

What information defines one attention output?

A weighted sum of values and its normalization. This small interface is what lets us split attention without approximating it.

For one attention head, a linear projection produces a query q for the token whose output we want, a key k_j for every token it may read, and a value v_j to read from that token. A dot product measures query-key compatibility. Softmax turns the scores into positive weights summing to one. A mask excludes keys that this query must not see.

s_j = dot(q, k_j) / sqrt(d) + any allowed score bias
p_j = exp(s_j) / sum_visible_keys exp(s_k)
o   = sum_visible_keys p_j * v_j

Here d is the query/key head dimension; o has the value dimension. We will omit the head index when the same reasoning applies independently to each head. In grouped-query attention, several query heads share one KV head. This reduces the amount of KV data stored and moved, even though each query head still computes its own scores.

For a causal sequence, query i can read positions up through i. For a tree, it can read earlier tokens in its own segment and all tokens in its ancestors. Notice what is absent from the formula: no term requires visible keys to be adjacent in memory, on the same GPU, or processed in a particular tile order. Those are implementation choices. Only membership in the visible set and correct normalization matter.

Efficient attention already avoids the giant matrix

The naive implementation forms a score matrix with one row per query and one column per key. At long context this intermediate is enormous. FlashAttention instead visits tiles of keys, keeps small running statistics, and never writes the complete probability matrix to GPU memory. Its major contribution is reducing traffic to high-bandwidth memory (HBM), the GPU's large off-chip memory, while retaining exact attention in real arithmetic. FlashAttention paper.

ART wants to retain that efficient inner computation. The local attention kernel should still perform its dot products, softmax, and value accumulation using its tuned GPU schedule. Distributed execution should provide the kernel with an appropriate query block, a KV block, and a mask, then combine its answers. Replacing a mature attention kernel with a Python-level distributed formula would lose the benefit we are trying to preserve.

The extra result we need from each kernel call is log-sum-exp, abbreviated LSE. For this query and this subset of keys, LSE is the logarithm of the softmax denominator. Keeping it in log space avoids overflowing exp(s) when scores are large.

L_A = log(sum_{j in A} exp(s_j))
o_A = sum_{j in A} exp(s_j - L_A) * v_j

The pair (o_A, L_A) summarizes everything subset A contributes to this query's output. It is not a reusable compression of those keys for every possible query: change q and both quantities change. It is a sufficient summary for merging this query's already-computed attention. That distinction will matter when we compare attention with recurrent state later.

In ART's dense path, PyTorch FlexAttention supplies the local operator and mask interface, with the qualified Flash backend for the H200 measurements in this article. ART supplies the work decomposition, movement of ranges, and merge around it. A kernel launch is not synonymous with a complete layer: one distributed layer can invoke several local attention stages, as well as small layout and merge kernels.

03 / Partition the same computation

Keep the query; bring it the keys it needs

A query's owner computes its complete output, using local KV immediately and remote KV in scheduled stages.

Let us use a tiny tree for the rest of the attention discussion. Its root has four tokens, physical positions 0 through 3. Completion A occupies 4 through 6. Completion B occupies 7 and 8. These are addresses in the packed row, not one causal sequence of length nine. Query 6 can see 0 through 6. Query 8 can see 0 through 3, then 7 and 8; it cannot see A.

For an illustrative two-rank assignment, rank 0 owns tokens {0,1,4,5,6}; rank 1 owns {2,3,7,8}. This assignment is not claimed to be optimal. It makes an important point visible: each rank's tokens can consist of several disjoint ranges. The owner stores the residual activations for those tokens and performs their token-local work. Initially it also holds their projected Q, K, and V.

Rank 0 can start query 6's attention against its local visible keys {0,1,4,5,6}. In parallel it requests the needed remote KV range {2,3}. Once that range is ready, another kernel computes query 6 against those keys. A stage is one such local attention computation over a planned query/KV subset. A wave groups remote ranges into a scheduled exchange and the computation that consumes them.

We do not need to broadcast Q to every KV owner in this design. Q remains at its owner, and the final O stays there too. That keeps the output projection, residual addition, normalization, and subsequent tokenwise work local to the same owner. Remote KVs are temporary inputs, not a change of ownership. Multiple local queries can reuse a fetched KV range; it should not be sent again for every query.

Two softmax outputs cannot simply be averaged

Each stage normalizes over a different set of keys. If the local stage accounts for 30% of the total exponential score mass and the remote stage accounts for 70%, their outputs must receive those weights. Counting keys does not give the weights; two very high-scoring keys may dominate two hundred low-scoring keys.

L = logaddexp(L_A, L_B)
o = exp(L_A - L) * o_A + exp(L_B - L) * o_B

The expression logaddexp(a,b) means log(exp(a)+exp(b)), evaluated stably. The two multipliers sum to one. The result is precisely the output that a single attention call over A union B would produce, up to floating-point evaluation order. It is also composable: merge two stages, then merge the result with a third.

This is why CP can build on a fast attention kernel. Its inner dot products and softmax do not become distributed operations. The distributed layer asks for the local O and LSE, then uses a comparatively small merge over query rows. ART fuses these gathers and merges where possible, rather than leaving chains of indexing, casts, and temporary tensors between large kernels.

Communication has a deadline, not just a byte count

A CUDA stream is an ordered queue of GPU operations. Work on different streams can overlap when the hardware and data dependencies permit it. ART uses that separation to begin a remote KV fetch while the compute stream runs local attention. It then waits for the fetch only where a remote stage actually needs the data. Later fetches can be issued while earlier remote stages compute.

Suppose local compute lasts 5 ms and a KV fetch lasts 3 ms. The fetch consumes network resources, but it need not add 3 ms to the layer's wall time. If the same fetch takes 7 ms, the compute stream waits about 2 ms after local compute. We call that delay exposed communication. More generally it is communication on the layer's critical path, the chain of dependencies that determines when the layer can finish.

The schedule cannot hide everything just by assigning more streams. Packing a send buffer must finish before its transfer reads it; attention must wait for its receive buffer; the buffer must remain alive while either stream uses it. GPU memory bandwidth and compute resources may also be shared between communication, packing, and attention. These constraints are part of the execution we measure, not overhead assumed away by an asynchronous API.

o = (N_A + N_B) / (Z_A + Z_B)
  = [Z_A / (Z_A + Z_B)] o_A
    + [Z_B / (Z_A + Z_B)] o_B.

a = exp(L_A - m)
b = exp(L_B - m)
L = m + log(a + b)
o = (a * o_A + b * o_B) / (a + b)

04 / Close the ownership loop

Where do the gradients go?

dQ stays with the query owner. Contributions to remote dK and dV must return to the token owner and be added there.

In a normal attention backward pass, each query contributes to its own q gradient and to the k and v gradients of all keys it read. A shared root may be read by hundreds of completion queries. Those contributions add. CP changes where they are produced, not how many there are or how they should be summed.

Return to query 6 on rank 0. Its local stage produces gradient contributions for local keys {0,1,4,5,6}. Its remote stage produces contributions for keys {2,3}. Those two remote tokens are owned by rank 1, which may also produce dKV for them from its own queries. Rank 0 must send its contributions back, and rank 1 must add them to the others before running the relevant projection backward.

The transpose of a gather that replicates values is an accumulation, not a plain copy. If three ranks fetch one root KV, backward has three contributions to that KV. If a rank has no useful tokens in a particular stage, it still has to obey the collective schedule when other ranks need its participation. A zero-sized tensor does not grant permission to skip an ordered collective.

There is a second gradient sum outside the attention operator. If CP ranks use replicated projection or MLP weights on different tokens, each rank computes only its token subset's contribution to those parameter gradients. The training system must reduce them with the model's intended loss normalization. Returning dKV to its token owner does not replace parameter-gradient reduction. EP and TP change which parameters are shared and which groups perform those sums; a correct topology must account for both kinds of ownership.

Use the global softmax derivative

Backward must also respect the normalization we just merged. For a query with incoming output gradient g, define D=dot(g,o), where o is the complete attention output. The derivative of a score is p_j times (dot(g,v_j)-D). The D term couples all visible keys, including keys handled by other stages.

There are two equivalent interfaces for an efficient staged backward. A kernel can use the saved global O and LSE directly. Or the merge backward can produce the appropriate adjoints for each local O and LSE. What is not correct is to run every stage's ordinary softmax backward with the original g and pretend the local denominators were global. ART's stage and merge machinery supplies this missing information around the local backward kernels.

After a remote stage's backward produces dKV, its return can overlap backward compute for another stage. Near the end, however, all outstanding contributions must be complete before the owner consumes its final gradients. This gives backward both a useful overlap opportunity and a possible communication tail.

This explains why a profiler's summed kernel times are not the layer time. Concurrent communication and compute overlap. CPU launch ranges can contain queueing or compilation rather than GPU execution. To reconcile a profile, we need a wall-clock span, stream dependencies, and the union of intervals on the critical path. Adding all visible bars can count the same millisecond twice.

dJ / do_r = w_r * g
dJ / dL_r = w_r * dot(g, o_r - o)

dJ/ds_j = p_global_j * dot(g, v_j - o)
p_global_j = w_r * p_local_j

dV_j += p_global_j * g
dQ   += (dJ/ds_j) * K_j / sqrt(d)
dK_j += (dJ/ds_j) * Q   / sqrt(d)

05 / Turn algebra into a schedule

Equal tokens are not equal work

The planner chooses ownership and communication waves to reduce the slowest rank's predicted completion time, not simply to make token counts match.

On an ordinary causal sequence, early queries read few keys and late queries read many. Split the sequence into equal contiguous intervals and the last rank gets much more attention work. A ring can circulate KV shards correctly, yet still leave early-query ranks underused. The communication algorithm and the load-balancing problem are different problems.

In ring attention, each rank starts with its local Q/K/V, computes against a KV shard, then passes that shard to the next rank while receiving another. Repeating the rotation visits the required remote KVs. The output/LSE merge keeps a running answer without retaining every shard at once. It is a useful distributed analogue of tiled attention. The limitation relevant here is not the merge math; it is committing to a regular circulation and ownership pattern when the useful work is irregular.

A familiar improvement is a zigzag or folded assignment. Give each rank one early piece and one late piece. Their causal workloads partly compensate: a cheap early piece accompanies an expensive late piece. For one regular causal triangle, this is a good geometrical shortcut. It can still reuse the same local attention kernel and the same merge algebra.

For eight causal tokens, query workloads are 1,2,3,4,5,6,7,8 visible keys. A contiguous two-rank split gives 10 versus 26 pairs. Split into four two-token blocks instead: their work is 3,7,11,15. Assigning the first and last to one rank and the middle two to the other gives 18 pairs each. That is the intuition behind the fold; it is not a change to which keys a query may read.

Now replace the triangle with our tree. Late physical positions are not necessarily expensive. They may belong to short independent branches. A shared root is read by many branches, so placing it affects KV traffic as well as its own query work. With sliding windows, a late query may read only a short path suffix. No fixed fold of the packed row captures all these cases.

The idea we took from MagiAttention

MagiAttention was the central inspiration: decompose attention into distributable work and let a planner decide its placement. The essential step is to stop treating a rank's contiguous sequence interval as a law of attention. A plan may assign different collections of query work to different ranks, then schedule the required data movement.

ART did not import MagiAttention's entire runtime or its attention kernel. Nor did we impose a restriction that all supported masks must be only full rectangles or causal triangles. Those shapes often occur in trees and are useful for explaining them, but the executor's masking interface preserves more general visibility. The inheritance is the decomposition-and-planning approach; the local dense arithmetic remains in the FlexAttention stack.

Three objects the planner must keep separate

First is visibility: which query-key pairs are valid. Second is ownership: where the token activations live. Third is the execution schedule: which query/KV subsets run in which stages and when their data arrives. Visibility is fixed by model and tree semantics. Ownership and scheduling may change to improve performance without changing visibility.

Searching all assignments of individual tokens is hopeless. ART groups nearby physical tokens into contiguous chunks for planning. A rank may own any collection of these chunks; its whole assignment need not be contiguous. Contiguity limits the search space and makes movement representable as a compact set of ranges. It is not a mathematical constraint that neighboring tokens must execute together.

The planner counts visible pairs between query chunk i and key chunk j. A matrix W[i,j] records these counts. Summing a query chunk's row estimates its arithmetic. The nonzero columns identify the KV it needs. After assigning owners, those same columns determine local work, remote dependencies, and opportunities to share one fetch across several queries.

Token-local work must not disappear from the objective either. A root's early queries may be cheap in attention but still run projections and an MLP. Giving one rank most tokens to equalize only attention pairs can make that rank the bottleneck in the surrounding model. Geometry, model-derived work profiles, and memory constraints complement the attention-time prediction.

Consider a root chunk used by ten query chunks on one rank. Its arithmetic contributes to ten query-key interactions, but its KV need only be fetched once for a stage containing those queries. A model that charges one transfer per interaction overestimates traffic. A model that balances only pair counts misses this reuse and may create an unnecessarily fragmented plan.

Predict milliseconds on the actual execution shape

Visible pair count is a useful starting point, not a complete runtime model. Tiny attention calls underfill the GPU. More ranges require gathering and assembly. Each stage adds launch and merge work. A remote transfer costs latency as well as bytes divided by bandwidth. Forward and backward have different throughput and different overlap opportunities.

The generic dense planner uses calibrated coefficients for these effects and simulates the stage schedule. Local and remote attention have distinct forward/backward pair-time coefficients. A measured underfill term accounts for small remote stages. Range counts approximate layout overhead. Communication coefficients must reflect the hardware; explicit rank-to-node link profiles in later planner work further distinguish cross-node traffic from local traffic. The measured snapshot is not a claim that one bandwidth coefficient predicts every placement.

stage time ~= launch cost
            + visible pairs / achieved pair throughput
            + merge work + range handling + underfill

transfer time ~= latency + bytes / achieved bandwidth

start next compute = max(compute available, required KV ready)
finish backward    = max(compute finished, all dKV returned)

These are empirical predictions, not proofs or peak-FLOP estimates. A coefficient should explain observed runtime across execution shapes; it should not be chosen merely to make a named benchmark select a desired plan. The simulator is also an approximation: contention and CPU scheduling can make actual overlap worse. Comparing predicted and measured times on alternative plans is the test of whether its decisions deserve trust.

The bounded search starts from reasonable chunk assignments, tries a limited set of wave counts, and evaluates candidate ownership moves around costly ranks. It stops when it cannot find a worthwhile improvement or reaches its search budget. This avoids exponential enumeration. More waves can improve overlap and limit temporary memory, but they also create more small calls. Fewer waves reduce launches, but can delay useful compute or enlarge working buffers.

Architecture-specific planners need architecture-specific work. GLM's profile includes indexer tile pairs, sparse-attention and projection work, bytes, memory, peers, and fragmentation; its resource-based comparisons are not identical to the generic dense scalar-time model. GDN has a separate runtime model for local versus chained recurrence. Calling all of these simply “balanced tokens” would hide the most important decision they make.

The plan is per input, not per layer

A changing packed tree requires a new visibility layout and potentially a new ownership plan. But dozens of compatible model layers can use that plan. CPU lookahead prepares future inputs while the GPU executes the current one. GPU block metadata is materialized for a prepared plan and reused where the attention geometry matches. Tensor values, such as actual Q and K, are still computed afresh in every layer.

Block metadata tells the local attention kernel which query/key tiles contain valid work, which are entirely visible, and which require an elementwise mask check. It is not a stored dense probability matrix. Compact construction is important at long context: materializing a temporary token-by-token Boolean mask can consume more memory than the attention operation itself. ART's preparation constructs the kernel's required representation without changing the visibility the kernel receives.

That separation is why a 40 ms planning time need not add 40 ms to every layer. It is also why we must report it. If lookahead cannot keep up, the model waits. If a new input repeatedly triggers compilation, the nominally fast layer is not fast in training. Caching an exactly repeated tree is useful in a fixed-input benchmark, but cannot establish performance on varying completions.

compute_done = local_attention_ms
fetch_ready = first_fetch_ms

for stage in remote_stages:
    start = max(compute_done, fetch_ready)
    fetch_ready = start + stage.next_fetch_ms
    compute_done = start + stage.attention_and_merge_ms

forward_done = compute_done

compute_done = 0
returns_done = []
for stage in backward_stages:
    compute_done += stage.backward_ms
    returns_done.append(compute_done + stage.return_dkv_ms)
compute_done += local_backward_ms
backward_done = max(compute_done, *returns_done)

Implementation anchor: ART context-parallel planner and runtime. Coefficients and search budgets describe this revision, not a claim of universal optimality.

06 / Change visibility, preserve the algebra

Trees, sliding windows, and attention sinks

The same staged attention algebra supports these cases, but storage addresses must not masquerade as positions or probabilities.

A tree mask is not an ordinary causal mask applied to the packed row. Physical position 8 follows completion A in storage, but A is not an ancestor of B. Position IDs also follow the root-to-leaf path. In our example, B's tokens at physical positions 7 and 8 have logical positions 4 and 5. That distinction affects rotary position embeddings as well as visibility.

For sliding-window attention (SWA), a query sees only a recent part of its logical history. If the window contains three tokens including the current token, query B8 at logical position 5 can see logical positions 3,4,5: physical tokens 3,7,8. Taking the last three physical addresses instead would return 6,7,8, leaking a sibling and dropping the correct root token.

The planner can intersect tree ancestry with the layer's window rule when counting work and requesting ranges. The local kernel receives the resulting mask. Smaller windows reduce arithmetic and often KV traffic, but they also reduce how much compute is available to hide communication. A model alternating global and windowed layers may therefore have very different bottlenecks in those two layer classes.

A learned sink is probability mass, not another token

Some models, including GPT-OSS, add a learned attention-sink logit per head. In this use of the term, the sink contributes to the softmax denominator but has zero value. It allows the head to assign some probability to “read nothing.” This is distinct from keeping the first few real tokens globally visible, another technique sometimes called attention sinks.

Suppose ordinary attention over visible tokens has output o and LSE L, and the sink logit is a. Including the sink rescales the token output by exp(L)/(exp(L)+exp(a)). The sink belongs to the whole query, not to every remote stage. Applying one copy in each of four stage denominators would silently multiply its mass by four and make the model depend on CP size.

ART can apply the sink once when composing the complete output, while retaining the ordinary local kernel and its LSE interface. Backward must propagate through that rescaling, including into the learned sink parameter. This small example illustrates a broader principle: preserve model semantics at the correct algebraic boundary, rather than independently applying a whole-model correction inside every partition.

L_total = logaddexp(L, a)
w = exp(L - L_total)
y = w * o

for g = dJ/dy:
dJ/do = w * g
dJ/dL = w * (1-w) * dot(g, o)
dJ/da = -w * (1-w) * dot(g, o)

07 / Change how keys are selected

GLM-5.2: parallelize the search as well as the attention

A learned indexer first chooses the keys. ART distributes that choice, merges canonical token IDs, and gives a combined KV buffer to an efficient sparse-attention kernel.

Dense attention computes scores for every visible key in every attention head. GLM-5.2 uses a separate, cheaper learned indexer to choose a limited number of visible tokens, then performs its main attention only on those choices. Multi-head latent attention (MLA) stores a compressed representation of KV, reducing KV storage and transport relative to storing a separate expanded KV vector for every query head.

The model configuration behind these experiments has 78 layers and hidden width 6,144. Attention has 64 heads, a 2,048-dimensional compressed query projection, a 512-dimensional KV latent, and value dimension 256 per head. Its indexer has 32 heads of dimension 128 and selects up to 2,048 tokens per query. The first three MLP layers are dense; later layers have 256 routed experts plus a shared expert. The MLP and attention designs are independent reasons this model differs from an ordinary dense Transformer.

The indexer's score is itself a small multi-head computation. Each indexer head compares its projected query with a shared projected key, applies ReLU to that dot product, and contributes through an input-dependent head weight. Summing these contributions gives one scalar candidate score, with the model's scale factors included in the projected quantities and weights below. This is not the main attention head's softmax probability. Tree visibility first determines which candidates are eligible; the indexer then chooses among them.

index_score(i,j) = sum_h weight(i,h)
                         * max(dot(index_q(i,h), index_k(j)), 0)

selected_ids(i) = top_k visible j by index_score(i,j)

The expensive part is not only the final sparse operation. The indexer must inspect the query's visible candidates, score them, and select top-k. Its scoring and selection are separate from the softmax in the main attention. The architecture amortizes this work with IndexShare: a full-index layer computes routes and following shared-index layers reuse them. In this configuration the first three layers compute their own indices, followed by repeating groups of one full-index layer and three shared-index layers. Shared layers reuse token choices, not another layer's KV values.

Two attention-like operations, two different planning costs

The indexer is essentially a smaller attention scoring operation: queries compare against visible keys. Its distributed execution therefore follows the same pattern as dense attention. Keep queries with their owners, bring in candidate keys in stages, compute local query–key scores, and combine stage results. What is combined differs: the indexer merges top-k candidates rather than softmax outputs and their log-normalizers. It does not compute an attention-weighted sum of values.

The subsequent sparse MLA has a different work shape. Once a query has at least k visible candidates, it attends to k selected keys whether its tree path contains 5,000 or 50,000 tokens. Its main attention work is approximately fixed per query, rather than proportional to the visible area of the prefix tree. Queries with fewer than k eligible keys are the exception. The indexer still pays to score the eligible candidates; selecting a fixed number does not make that search fixed-cost.

A planner must account for both terms: candidate-scoring work on full-index layers, and selected-key attention work on every layer. Balancing only visible query–key area can leave the sparse operation unbalanced, because equal areas can contain very different numbers of queries. Nor can we reuse dense attention's stage splits unchanged for the sparse operation: the selected keys can be distributed very unevenly among those stages. ART retains the staged search, but assembles the selected address space for the main sparse kernel, as described below.

hidden tokens
  -> Q / compressed-KV projections
  -> full layer: score candidates and choose token IDs
     shared layer: reuse preceding full layer's token IDs
  -> sparse MLA over the selected KV rows
  -> output projection and residual
  -> dense or mixture-of-experts MLP

The sparse kernel already takes a list of KV rows

This interface is the key to reuse. A sparse MLA kernel accepts queries, a KV buffer, and a per-query list of indices into that buffer. It does not need to know that the values came from different GPUs or that visibility originated in a prefix tree. ART can reorganize and remap the input without replacing all of the kernel's highly tuned arithmetic.

MLA's compression survives this boundary. ART folds the key expansion into the query projection, so the kernel compares a transformed query directly with the compressed KV latent plus its positional component. It accumulates a weighted sum in latent space; the value expansion happens afterward. It does not send an expanded 64-head KV tensor merely to run the distributed operation. That is both a memory saving and a reason to preserve this kernel interface carefully.

On a full-index layer, each scheduled candidate stage produces scores paired with canonical token IDs: identities in the packed tree, independent of the stage's temporary storage. The query owner merges these candidates into its global top-k. Then it assembles the required stage KVs into a combined buffer and maps each winning ID to that buffer's row. One sparse MLA call consumes the resulting list.

Why not call sparse MLA separately for every candidate stage? After selecting a fixed total of k keys, each stage may hold only a small and uneven part of them. Separate calls repeat launch, bookkeeping, and intermediate-output work. Combining the selected address space allows the kernel to see one normal sparse problem. Backward splits dKV using the saved buffer boundaries and returns contributions to their owners, just as in dense CP.

IndexShare makes canonical IDs especially important. Reusing temporary offsets from another layer or plan would confuse “row 17 in this receive buffer” with “token 17 in the input.” The former changes with placement; the latter is a stable model choice. Any routing tie rule must also remain independent of message arrival or rank order.

Where ART actually changed the kernels

The sparse kernels were adapted from Miles GLM and TileLang's DeepSeek-V3.2 work. Their GPU schedules are valuable starting points, not untouchable black boxes. ART added variable-length CP integration, canonical route handling, combined KV execution, supported TP-local head counts, and kernel changes where the original geometry was inefficient or incorrect under the pinned toolchain.

One important backward issue is atomic contention. Many queries may select the same shared-prefix key. They all contribute to its dK/dV, and parallel programs must add to common locations safely. Atomics provide that addition, but contending updates can serialize. ART used four FP32 accumulation shards for dKV followed by a reduction. This spends temporary memory and an extra reduction to reduce contention in the dominant kernel. It is a general response to shared-key fan-in, not a special rule for one prompt length.

Indexer work also suffered from launch shape. Grouping rectangles that share query ranges reduced 1,715 rectangle programs to 286 grouped-query programs in the recorded case. Making query/key extents runtime arguments removed repeated specialization as sequences changed. These changes affect different costs: grouping improves GPU work organization; dynamic extents remove compilation from changing-input execution. They should not be conflated into a single arithmetic speedup.

Other work repaired a pinned TileLang backward problem that could produce NaN dQ, handled invalid route sentinels, preserved the LSE log-base convention, and tuned SM100 dQ using its tensor-core instructions with explicit shared-buffer synchronization. Correctness at these interfaces is essential: a kernel with a faster but incompatible LSE convention breaks the same merge derivative we derived earlier.

That last distinction matters. A valid, matched-shape stress comparison can still be the wrong headline for realistic training. The quarter-off backward result and the combined-KV operation result are more useful estimates of production impact. None of these percentages should be multiplied together: they measure different scopes and, in some cases, different baselines.

q_i^T k_j = q_i^T W_K c_j
          = (W_K^T q_i)^T c_j

sum_j p_ij v_j = sum_j p_ij W_V c_j
              = W_V (sum_j p_ij c_j)

Implementation evidence: combined-KV execution, dynamic grouped indexer and dKV work; architecture and runtime at the GLM source snapshot.

08 / Change what represents the past

Gated DeltaNet context parallelism

A token shard can compute its effect on an unknown incoming state. These effects compose, letting ranks do most of their local work concurrently.

So far, every distributed operation has let a query inspect individual keys, either all visible keys or a learned subset. GDN changes that interface. There is no list of past KV rows to fetch and no softmax output to merge. Instead, each token reads and updates a fixed-size state. To distribute this layer, we need to understand what information crosses a token-shard boundary and how to compute it without serializing the ranks.

Qwen3.5 mixes ordinary attention with Gated DeltaNet (GDN) layers. An attention layer keeps access to individual past keys and values. GDN compresses the past into a finite matrix state. That makes its token-dependent work linear in sequence length for fixed model dimensions, but a literal reading of its recurrence appears to introduce a sequential dependency at every token.

We will separate three problems: how the single-GPU kernel avoids a token-by-token loop, how its state transformations compose across GPUs, and how the prefix tree lets independent branches share an initial state without sharing their subsequent updates. The short convolution around the recurrence has a separate, much smaller boundary dependency. Together these give us a distributed GDN layer, not an attention algorithm applied to a different kernel.

What a GDN token does to the state

For one head, write the state S as a d_k by d_v matrix. A key k selects a direction in its key dimension, and k-transpose times S reads the value currently associated with that direction. Each new token has a query q, a key k, a value v, a decay alpha, and a write strength beta. First decay the old state, then correct its prediction of v in the direction k:

S_pre = alpha * S_previous
delta = beta * (v - S_pre^T k)
S_new = S_pre + k delta^T
y     = S_new^T q

All q, k, and v in this notation are column vectors. The superscript T means transpose, not sequence length. If the state already predicts v accurately, the correction is small. If it does not, the outer product k delta-transpose changes the relevant key direction. Alpha controls forgetting; beta controls the correction. This is not merely adding k v-transpose at every step, and not an interpolation that replaces the entire state.

The surrounding layer projects its inputs, applies a short causal depthwise convolution to the Q/K/V streams, normalizes Q and K, constructs the gates, and after the recurrence applies gated normalization and an output projection. A depthwise convolution filters each channel independently over nearby preceding tokens. These token layouts and short filters matter to runtime even though they are not the recurrence itself.

First parallelize within a GPU

We should not run the equations above as a Python loop over thousands of tokens. Flash Linear Attention (FLA) realizes the same recurrence using small chunks. Within a chunk, it expresses the interactions as triangular matrix operations suitable for GPU kernels; between chunks it carries the finite state. ART builds on these efficient gated-delta kernels. FLA implementation.

The important property is that a chunk's input-dependent preparation can be computed before its incoming state is known. It produces matrices that express both its outputs and its final state as affine functions of that incoming state. “Affine” means a linear transformation plus an additive term. This gives us a much better distributed algorithm than waiting for rank 0 to finish all its tokens, then rank 1, then rank 2.

Then parallelize across GPUs

For a whole rank-local shard of a chain, combine its chunk transformations into:

S_out = M * S_in + E

M: d_k by d_k       E: d_k by d_v

M describes how this shard transforms any prior state. E describes what it writes starting from zero. These depend on the shard's tokens, not on the actual state arriving from earlier shards. Every rank can therefore prepare its local token work and its summary (M,E) concurrently. It does not fetch the preceding ranks' Q/K/V to replay their recurrence.

Two summaries compose in order. If A runs before B, substitute A's output into B:

M_BA = M_B * M_A
E_BA = M_B * E_A + E_B

This composition is associative, although not commutative. A prefix scan computes the composition of all earlier summaries for every rank. It is the same algorithmic idea as cumulative sums, with affine composition replacing addition. After the scan, each rank knows the true initial state for its shard and applies it to the work it prepared locally.

The source snapshot uses doubling rounds: at distance 1, then 2, then 4, each rank incorporates an earlier aggregate. These are collective all-to-all exchanges with zero-count entries for uninvolved peers, not a rank-conditional chain of ad hoc sends. An extra shift makes the result exclusive, meaning “all shards before me,” and the final chain state is distributed where the tree needs it. Backward uses the corresponding suffix information in reverse order.

That detail explains both the benefit and the cost. A chain of N ranks does not require N full shard computations in series. But there are logarithmically many rounds of summary communication and matrix composition, plus boundary work. Summary size is independent of token count but grows with head count and squared key dimension. On short shards, it can be a substantial fraction of the whole operation.

The tree supplies a second kind of parallelism

Not every segment should be chained. If an entire independent family fits efficiently on one rank, a local execution avoids distributed summary costs. If one long root dominates the row, distributing that root can reduce the maximum rank time. When a root finishes, all its children inherit the same state and can run independently; at the next tree level, their children do the same.

ART represents these dependencies as a directed acyclic graph (DAG). For a dependency depth it batches compatible local work together and chained work together, rather than launching a separate recurrence for every family. Packed metadata uses one token dimension and segment offsets, so varying lengths do not require every segment to occupy the longest segment's padded width. The local and chain cases use the same mathematical operator, with different ownership of its state boundary.

Backward follows the DAG in reverse. Two children using one parent state contribute two gradients to it; those are summed before propagating through the parent. Within a chained segment, a rank's state gradient includes effects from later ranks as well as its own token outputs. Merely backpropagating each shard from a zero final-state gradient would lose the cross-rank dependency.

Convolution needs a halo, not the whole prefix

A causal filter of width K needs at most K-1 preceding input tokens to begin a shard. Those boundary tokens are called a halo. If a shard is extremely short, its halo may span more than one earlier rank. At a tree fork, the children begin from the parent's tail, not the tail of a sibling adjacent in storage.

ART's true variable-length convolution handles the packed segments, activation, and state boundary. Communication supplies the needed tails; backward returns and accumulates gradients for their duplicated uses. This avoids converting a ragged family into a large padded batch merely to call a rectangular convolution. The custom layout and convolution kernels are useful because they remove real memory work and launches around the already-efficient recurrent kernel.

The remaining planner choice is economic: does saved token compute exceed layout exchange, summaries, scans, and parent-state traffic? GDN's runtime model includes these terms and scales state sizes with actual head/key/value dimensions. Its fit uses more than the 35B shape, including 9B and 397B. It is still an empirical model with bounded search, not a guarantee of globally optimal plans. A fixed “chain above this many tokens” rule would miss both model-size differences and the other work competing for each rank.

Consecutive GDN layers can retain a GDN-friendly token ownership, an island, rather than exchange back after every layer. The residual stream and tokenwise modules must agree on that layout. Attention-to-GDN boundaries perform the required redistribution; TP sequence-parallel split/gather operations must preserve the same token order through the island. Layout is part of the layer contract, including recomputation in backward.

S_pre_i = alpha_i S_(i-1)
delta_i = beta_i (v_i - S_pre_i^T k_i)
S_i = S_pre_i + k_i delta_i^T

S_pre_i = exp(G_i) S_start
        + sum_{j<i} exp(G_i-G_j) k_j delta_j^T

delta_i^T = beta_i [v_i^T
  - exp(G_i) k_i^T S_start
  - sum_{j<i} exp(G_i-G_j) (k_i^T k_j) delta_j^T]

P[i,j] = exp(G_i-G_j) * dot(k_i,k_j), for j<i
P[i,j] = 0, otherwise
K_start[i,:] = exp(G_i) * k_i^T

D = B (V - K_start S_start - P D)
(I + B P) D = B (V - K_start S_start)

U = A B V           [c,d_v]
W = A B K_start     [c,d_k]
D = U - W S_start   [c,d_v]

S_end = exp(G_last) S_start + R^T D
      = [exp(G_last) I - R^T W] S_start + R^T U

M = exp(G_last) I - R^T W
E = R^T U

C[i,j] = exp(G_i-G_j) * dot(q_i,k_j), for j<=i
Y = Q_start S_start + C D
  = (Q_start - C W) S_start + C U

M_c M_b M_a S
  + M_c M_b E_a + M_c E_b + E_c

dS_in = F^T dY + M^T dS_out

Source anchors: FLA CP summary/scan implementation, shape-aware runtime model. The derivation expands the project's GDN layer guide.

09 / Move the partition instead

Mamba-2: each rank gets the whole tree for fewer heads

The qualified Nemotron path uses an all-to-all transpose from token shards to head shards, preserving the native SSD kernel instead of using GDN's distributed summary implementation.

Mamba-2 is also recurrent, but it is not GDN with another name. Its selective state-space model constructs an input-dependent write and read around a decaying state. Within a head, the decay has a simpler structure than GDN's key-dependent matrix correction. Mamba-2's state-space duality (SSD) connects this recurrence to structured matrix operations, enabling efficient chunked GPU execution. Mamba-2 paper.

In a simplified per-head notation, x is a vector of channels; B and C are state write/read vectors; dt determines the discretized time step. The head's learned decay parameter and dt produce a scalar decay a. A state H maps state coordinates to channels:

H_t = a_t H_(t-1) + x_t b_t^T
y_t = H_t C_t + D_skip * x_t

b_t includes the input/time-step scaling in this notation. The full layer projects x, B, C, dt, and a gate z, applies a short convolution to the appropriate streams, runs SSD, and applies gated normalization and output projection. Several heads can share B/C through groups, so those projected widths need not equal the full number of heads times state dimension.

The Nemotron Nano configuration studied here contains 52 blocks: 23 Mamba, 23 MoE, and six attention blocks. Its Mamba layers have 64 heads, eight B/C groups, state dimension 128, and convolution width four. These independent heads offer another distributed decomposition: instead of splitting a head's recurrence across ranks, give a rank fewer heads but all physical tokens.

What crosses the network

The input projection runs on the original token owners. ART then partitions projected x/B/C/dt by recurrent head assignment and uses one all-to-all to assemble each destination's full-token, partial-head streams. An all-to-all lets every rank send a different slice to every other rank. This exchange is a transpose of the token-by-head ownership grid, not an all-gather that replicates every head's tokens everywhere.

The gate z is not needed by SSD itself. It remains on token owners. After the recurrent computation, the reverse all-to-all brings each token's output head slices home, and gated normalization/output projection proceed in that original layout. Backward reverses the two transposes and accumulates any replicated grouped B/C contributions.

This differs deliberately from our dense attention strategy. There, queries remain token-sharded while needed KV visits them. Here, the independent head axis lets the recurrence run locally across all tokens. The cost is moving the projected recurrent streams and the output, and the benefit is leaving the native SSD kernel in control of the full sequence for its heads.

What the tree wrapper changes

After redistribution, each rank still sees a tree, not one ordinary causal array. ART computes the physical tree's convolution once for its channels, gathers recurrent inputs into tree-aware scan buckets, and supplies the correct parent states. The numerical core is the official mamba_ssm 2.3.2.post1 SSD scan. Custom Triton kernels handle permutation, gathering, assembly, and the corresponding gradients around it.

The native chunk interface also shapes state reuse. Complete shared chunks can hand their states to descendants; an unaligned shared suffix is replayed in the child's scan so the native kernel sees its expected chunk realization. That replay is bounded by the chunk size, not the entire prompt. It is a real overhead and must be counted, but it is not the central explanation of scaling. The dominant design choice is the token/head transpose and reuse of native SSD.

There are limits. A head shard must have a supported width, so CP size cannot grow independently of the number of available heads. If several rank shards need the same B/C group, those narrower streams are replicated and their gradients summed. Full-token metadata is present per recurrent rank, and poorly chosen bucket shapes can introduce padding or extra layout work. “Every rank gets fewer heads” is not proof that every allocation is exactly divided by CP.

This is also why we do not call Mamba's implementation “the GDN scan path.” Both exploit recurrent structure, and both preserve shared-prefix states, but their network interfaces are different. The measured Nemotron implementation has a qualified CP2 training result and pipeline comparison; we do not yet have a matching CP1/2/4/8 curve from which to claim the same scaling as dense attention.

rank r before: tokens owned by r, heads 0 ... H-1
send to rank j: those tokens, head slice j
rank j after:  all tokens, head slice j

reverse exchange:
send each token owner's rows back, then concatenate head slices

Qualified source: Nemotron tree operator and exchanges, implementation and validation PR.

10 / Account for the whole rank

Which memory scales, and which does not?

CP divides token-dependent storage. It does not, by itself, divide model weights, eliminate communication buffers, or make every architecture's temporary state scale identically.

Suppose a residual tensor has T tokens and hidden width H. With balanced token ownership across N ranks, each rank stores about T H/N elements. Many intermediate activations, Q/K/V projections, and tokenwise gradients have the same dependence. At fixed total T, CP can reduce them. If T grows with N in weak scaling, they can stay approximately constant per rank.

Parameters behave differently. A rank that owns fewer tokens still needs the weights used on those tokens. CP alone generally replicates dense weights across its ranks. TP, EP, PP, optimizer sharding, or weight offload may address that memory. In a CP8/EP8 MoE run, lower expert-weight memory comes from EP, not from an unexplained property of CP. Adding both axes together is useful, but attribution matters.

Activation checkpointing, also called recomputation, saves fewer intermediate activations and recomputes them during backward. It can combine with CP. The recomputed layer must use the same tree, ownership, and communication contract as the original forward. A cached mutable “current layout” that changes before recomputation can corrupt results even if the first forward looked correct.

Why CP can be an attractive alternative to more TP or PP

Increasing TP splits each matrix multiplication further. That can reduce local matrix efficiency and requires collective results at linear-layer boundaries. Depending on the model, moving compressed KV for attention can be cheaper than moving hidden-width partial linear results repeatedly. CP also leaves larger local channel dimensions for tokenwise matrix multiplication. These are reasons to try it, not a claim that CP always sends fewer bytes on every architecture.

PP keeps each GPU busy only when its pipeline stage has a microbatch available. Long packed sequences consume substantial memory, limiting how many microbatches we can keep in flight. Hybrid models also make stages hard to balance: an attention block, a GDN block, and an MoE block need not take equal time. Bubbles and uneven stage time can then outweigh small pipeline transport costs.

CP works within one packed microbatch and can use all ranks on its token-dependent work. That makes it particularly valuable when there are too few independent packed microbatches to fill a deep pipeline. Data parallelism remains valuable once enough separate work and enough memory are available; it avoids the within-sequence communication entirely. The right conclusion is to choose axes according to the limiting resource, not to treat one topology as universally best.

Finally, distinguish allocated memory, live tensors known to the allocator, from reserved memory, the allocator's retained blocks. Report the peak across ranks. Different completion lengths and concurrent buffers can change fragmentation and peak reservation without changing the mathematical tensor count. A layer-memory estimate is useful for planning, but a real training run is the proof that the job fits.

11 / Put the mechanisms against measurements

How close do we get to the scaling target?

Dense attention is very close after its first distributed step. Sparse MLA has a larger fixed CP cost. GDN scales usefully but has a more exposed scan cost. The distinctions are visible in the measurements.

The plots below show time relative to each series' CP1 time. Ideal weak scaling is a horizontal line at 1: N times the GPUs process roughly N times the physical work in the original time. A ratio of 1.2 at CP8 means 8/1.2, or about 6.7 times, the aggregate throughput of CP1, not a 20% throughput gain. Scaling efficiency is the reciprocal of the time ratio for matched weak-scaling work.

What the workload names actually mean

A name describes one family that is repeatedly added to fill a packed row. A row can end with a clipped final segment. Roots remain fixed in these ordinary presets. Except for the explicitly fixed control, completion lengths are independently sampled for each new sequence using rounded, clipped normal draws. This changes tree geometry and ownership, not merely tensor values.

For these weak-scaling dense rows, global packed length is base times CP size. Thus fixed/varied 5K use 40,960, 81,920, 163,840, and 327,680 tokens at CP1/2/4/8. The long-chain row uses 98,304 through 786,432 tokens. The “K” in historical preset names is shorthand, not a substitute for the exact integers in this table.

The short control is useful for launch and communication overhead. Dominant families test skew and shared-prefix demand. Medium-long completions change the arithmetic-to-traffic ratio. Long segments make within-trajectory dependencies much more expensive. The historical name “true completion chain” describes its long branches; under family-stacking weak scaling, it does not guarantee that the GDN planner chooses a distributed chain at every CP size. That is a property of the actual plan, not the name. Deeper tree tests separately exercise nested shared roots and ragged branching. No one default prompt can establish all of these properties.

Dense attention: understand the first jump

CP1 has local attention and no remote KV stage or distributed merge. CP2 activates fetches, additional attention stages, normalization merging, and inverse dKV return. On fixed 5K this costs about 2.2 ms in total. Increasing from CP2 to CP8 then quadruples global tokens for only 0.93 ms more operation time. Varied 5K adds 1.13 ms over the same interval. Those are about 95% weak-scaling efficiency relative to CP2.

This is a mechanistic explanation of the onset, not a claim that every microsecond of the 2.2 ms has been isolated into a universal constant. Our accepted Nsight runs estimated exposed communication at about 0.53, 0.54, and 0.70 ms for fixed CP2/4/8. Most transfer duration overlapped useful work. Additional calls, layout work, and merges still occupy time even when the network is hidden.

A later source check reproduced fixed CP2/CP4 at 19.405/19.853 ms and varied at 19.536/20.498 ms. The table preserves the accepted CP2/4/8 matrix rather than silently replacing individual entries with whichever run was fastest. The local shape is 32 query heads, four KV heads, and head dimension 128; these are the model's attention tensors, not its GDN state dimensions.

Preparation is measurable even when it is amortized

Fixed trees benefit from plan caching; changing trees pay the search and rank-materialization cost. GPU block metadata is a significant setup operation in either case. These costs occur per prepared input/layout, not for every compatible layer. They cannot simply be omitted from end-to-end throughput, and they cannot simply be added to every layer time either. Lookahead determines how much is exposed.

For the changing CP2 run, the separately summarized forward and backward medians were 5.416 and 14.070 ms, while the end-to-end median was 19.536 ms. Separately selected medians do not necessarily add, and the enclosing measurement includes its own boundaries. Nested span totals across all ranks are even less suitable for additive attribution.

GDN: useful scaling, but a less hidden boundary cost

The corresponding CP8/CP1 time ratios are 1.47, 1.59, and 1.72. This is substantial aggregate throughput growth, but not dense attention's nearly flat curve. More of the summary scan, reverse-state propagation, parent-state movement, and layout work remains exposed. Chaining is a tool to reduce local imbalance, not an operation with the same constant cost as a local bucket.

The comparison also illustrates why asymptotic complexity is not a runtime prediction. The short dense-attention control can be faster than GDN despite its quadratic per-trajectory arithmetic. Efficient attention kernels have high arithmetic intensity and well-amortized tiling; GDN carries state, performs multiple preparation and backward kernels, and pays boundaries. At long completion lengths the difference reverses: the dense long-chain operation is much more expensive. A full-model prediction must use the actual layer mix and dimensions, not an O(T) label.

GLM: the learned indexer remains part of the bill

Eight times the tokens takes 1.289 times the operation time, with similar memory per rank. CP2 introduces route merging, remote candidate stages, combined-KV assembly, and gradient return; the time then stays in a much narrower band through CP8. Sparse backward remains a major optimization target. We do not give an additive indexer/forward/backward split here because the archived component summaries use a different aggregation from the total and do not reconcile as a sum.

A separate four-layer, production-dimension GLM training test took 4.994/5.644/5.803/5.880 seconds at CP1/2/4/8 with independently varying completions. The CP8 time ratio was 1.177. Peak allocation went from 94.90 to 73.49 GiB per rank, and useful model-FLOP utilization stayed in the 19.1–22.4% range. This includes surrounding projections, experts, recomputation, loss and backward; it is not the 78-layer model's step time.

Across the node boundary

The middle row isolates placement: same width and work, but some traffic crosses nodes. It retains 95.96% of the one-node throughput. The last row then doubles both GPUs and tokens. Relative to the split CP8 row, throughput rises about 1.89 times, or 94.60% of ideal twofold weak scaling. This is not 94.60% of the old absolute throughput; it is efficiency relative to doubling it.

The result includes EP as well as CP. It says the combined training topology scales well on the measured fabric, not that a network link has a universal 5% penalty. A fixed-work comparison also favored the tested CP-only placements over PP: on eight GPUs, 14,011 versus 11,024 tokens/s for PP2; on sixteen, 27,964 versus 21,610 for PP2 and 20,134 for PP4. Pipeline bubbles and stage balance matter here, not only transport bytes.

Mamba: keep the evidence boundary visible

The exact 52-layer Nemotron run at CP2/EP2, 128K, BF16 and compilation enabled on B300 measured a steady 3.455-second step and 10.13% useful MFU. Five steady samples had about 1.36% coefficient of variation, without measured post-warmup compilation. This establishes an actual full-schedule training result; it does not establish a CP8 scaling curve.

A separate reduced 39-layer fixture measured 43,624 physical tokens/s in isolation, 43,469 for matched core training on live-pipeline inputs, and 41,273 end to end. The live core retained 99.64% of isolated throughput, and end to end retained 94.61%. Those comparisons test integration overhead at the same reduced fixture. Mixing them with the exact 52-layer time would create a misleading “speedup.”

12 / What makes a result transferable

Measure the computation real training will execute

A good local kernel result is necessary. It is not sufficient until changing trees, gradients, preparation, memory, and the surrounding model agree with it.

The experiments use several levels of measurement for different questions. Operation labs make ownership, stages, kernel time and communication observable without full-model setup. Stacked-layer benches expose per-layer repetition, layout transitions, and specialization on changing inputs. Isolated backend training adds real parameters, experts, optimizer/adapter work, recomputation and model memory. Live-pipeline checks finally measure how much of that rate survives input supply and job boundaries.

The throughput tables report physical packed tokens per second, not generated tokens or trajectories per second. A shared prefix contributes physical tokens even when it has no direct loss. When comparing different packing schemes, count completed trajectories or the same training objective as well: removing redundant prefix copies can improve useful throughput while reducing the number of physical tokens executed.

We keep instrumentation on the lab or harness side. Nsight/NVTX ranges are useful for explaining a trace, but heavy per-module timing and memory hooks can themselves change throughput. A clean throughput run uses the intended compilation setting and avoids those hooks. A diagnostic run supplies attribution, with its overhead and measurement scope stated separately.

Changing inputs are part of the performance contract

A kernel can be fast on a warmed input and unusable if it recompiles whenever the next completion has a different length. Token counts, segment counts, and variable total extents should therefore be runtime inputs where possible. Static model dimensions and a finite set of execution roles can still specialize. The relevant question is whether warmup covers those roles in a reasonable number of independent examples, not whether an exact replay hits its cache.

In the dense source checks here, sixteen measured trees were distinct from the warmup trees. The fixed control deliberately repeated geometry, while varied workloads sampled lengths afresh. The cold-role example above is why four warmup inputs were not sufficient evidence on their own. A benchmark should record the per-sequence time series and new compilations, not discard an arbitrary number of slow samples until a preferred average appears.

Correctness needs more than one output comparison

Packed-versus-flattened tests establish that branches cannot read siblings and shared computation matches independent trajectories. Topology oracles compare canonical outputs and gradients across parallel layouts. Tracing must reconstruct original token identity before comparing; otherwise a storage permutation looks like a model error. Exact routing checks catch sparse-index changes that a mean output tolerance could miss.

For recurrent layers, tests also need parent-state gradients, empty participants, short segments, long chains, and recomputation. A rank skipping a collective is a protocol error even if it has no local arithmetic. Conversely, a large relative gradient error should not be dismissed because the tensor's absolute values are small. Correctness is the algebra and its complete distributed realization.

Low-precision kernel outputs need dtype-appropriate tolerances, but tuning those tolerances is not a performance optimization. The useful optimization target is preserving the intended math while eliminating needless movement, launches, replication, and serial dependencies. This is why the kernel-reuse boundary is important: we keep mature arithmetic where it works and modify the surrounding organization, or the kernel itself, only where evidence identifies a real limitation.

The common idea, after the differences

There is no one “CP primitive” that explains all four implementations. Dense attention brings KV to query owners and merges normalized results. GLM first merges a learned choice of token identities, then runs sparse attention over remapped KV. GDN composes transformations of the recurrent state. Mamba redistributes the independent head dimension so native SSD can process each head's tree locally.

What they share is a method: start from the information dependency, find a compact interface to the efficient local computation, preserve token and gradient ownership, and schedule the remaining communication against actual compute. The planner must model that execution, not a convenient surrogate such as token balance alone. The benchmark must measure that execution on changing inputs, not just its most cache-friendly snapshot.

The original motivation was a longer useful row. Prefix trees remove duplicate prompt work; CP makes room to process more of the remaining trajectories together. The evidence shows this can be a strong throughput strategy as well as a memory strategy, including across nodes. Understanding why requires both halves: the algebra that permits partitioning, and the execution details that keep partitioning from becoming expensive.

useful MFU = useful model FLOPs
           / (elapsed seconds * total GPU peak FLOPs/second)

weak-scaling efficiency = time_at_base / time_at_N
  when useful work and GPU count both grow by N