[RFC] Graph capture support for prefill and training: GDN/KDA chunk and convolution paths #1155

Proposal

Enable graph capture for the full prefill/training compute path of GDN and KDA models — the chunk attention forward/backward and the short-convolution forward/backward — so that a bounded input regime can be captured once and replayed many times.

"Graph capture" is backend-neutral in this proposal: CUDA Graphs on NVIDIA, and the equivalent graph mechanism on Ascend NPUs (ACL Graph, as already used by vLLM-Ascend), both rely on the same prerequisites — static grids, static addresses, capture-safe kernels. The contract proposed here is written once, against those prerequisites, and both backends benefit from the same work.

The dynamic parts of the varlen path today:

1. Grids are sized by the number of chunks / tokens, which depends on the input lengths.
2. The chunk index list (which sequence each chunk belongs to, its offset inside the sequence) is built on the host from cu_seqlens; its output length is data-dependent and the build involves host synchronization. Neither is allowed during capture.
3. Tensor addresses change from call to call.

Batch size enters through (2): a static formulation must pin a maximum sequence count N_max alongside the token budget T_max. For GDN/KDA training, N_max is the micro-batch size — fixed per training configuration, so a run captures once and keeps replaying. For prefill, N_max is the serving engine's max sequences per batch. The pinned (T_max, N_max) pair also prices the memory overhead (below), so batch size is a first-class input to this design.

Principles

• Scope: GDN and KDA, chunk attention fwd+bwd plus short-convolution fwd+bwd. Decode-specific paths (fused_recurrent, conv update-with-cache) are out of scope — serving stacks already handle those.
• Switch off: behavior is bit-identical to today. Switch on: kernel math is unchanged; only scheduling metadata (grids, index tables) changes.
• The graph path pins the maximum batch: inputs must satisfy tokens ≤ T_max and sequences ≤ N_max. Larger inputs transparently take the existing eager path with identical results.
• One contract, both backends: the kernel-level changes (static grid, sentinel early-exit, device-side indices) are written once and are valid for both CUDA Graph and Ascend graph capture.
• Success metric: host launches per layer drop to one replay within the bounded regime; end-to-end gain is measured by same-hardware benchmarks.

Core design

Static upper bound with sentinels. Given (T_max, N_max), the chunk count has a closed-form upper bound NT_max = cdiv(T_max, 64) + N_max − 1 (the cdiv term covers the first sequence's partial chunk; every further sequence contributes at most one more). Every chunk-count-sized kernel launches a static NT_max grid; unused index-table slots carry a sentinel, and the kernel returns immediately upon reading one.

Device-side chunk indices. The index list is generated on the device: a searchsorted-based routine fills a fixed-shape [NT_max, 2] buffer with (sequence id, intra-chunk offset) pairs and writes the sentinel into unused slots. cu_seqlens is padded to N_max + 1 entries with zero-length tail sequences so its shape is fixed. The operator can build this table internally, or the caller can precompute and pass it in — GDN already accepts external chunk_indices since #1132; KDA gains the same option here.

Static buffer contract. Replay reissues the captured addresses, so every call copies inputs into fixed-shape static buffers and reads outputs from them. This holds for prefill and training alike, independent of recompute mode, and defines how callers drive the operator.

Enable rule and fallback. An input is graph-eligible iff tokens ≤ T_max and sequences ≤ N_max (dense input is the N = 1 special case). Larger inputs take the eager path. Bucket choices (how many (T_max, N_max) tiers) are caller policy; the operator layer guarantees that any fixed tier is capturable.

Memory model (worst-case pricing). All graph-held tensors are sized at the static upper bound regardless of the actual input, and replay always touches the full static shapes. Cost follows the bound, in three parts: static input/output buffers at [T_max, …]; intermediates allocated at NT_max (for KDA-class states, K = V = 128 in fp32, that is 64 KB per head per chunk — T_max = 4096, N_max = 32, H = 32 gives ≈ 190 MB resident); and one graph private pool per tier. Convolution adds only its W-window states, negligible next to the chunk states. The first version ships a single tier; benchmarks decide whether tiering pays.

What needs to change

Four steps, each independently mergeable:

1. Utils (fla/ops/utils/): device-side static chunk-index generation (searchsorted + sentinel fill), the NT_max helper, zero-length sequence padding, and static-grid variants of the varlen cumsum kernels. Pure-torch parity tests included.
2. Kernel contract. Every varlen kernel in the two paths gains a compile-time GRAPH_CAPTURE flag and a sentinel early-exit. The flag defaults to off, in which case the new code is dead-code-eliminated and behavior is bit-identical. Wave-one checklist:
   • KDA fwd: kda_gate_chunk_cumsum, chunk_local_cumsum, chunk_kda_fwd_intra, chunk_gated_delta_rule_fwd_h, chunk_gla_fwd_o_gk.
   • KDA bwd: chunk_kda_bwd_intra, chunk_kda_bwd_wy_dqkg_fused, recompute_w_u_fwd, chunk_gated_delta_rule_bwd_dhu, the recomputed chunk_gated_delta_rule_fwd_h, kda_gate_bwd.
   • GDN fwd: chunk_local_cumsum, chunk_gated_delta_rule_fwd_intra, recompute_w_u_fwd, chunk_gated_delta_rule_fwd_h, chunk_fwd_o.
   • GDN bwd: prepare_wy_repr_bwd, chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_gated_delta_rule_bwd_dhu.
   • Conv fwd/bwd (fla/modules/conv/triton/): causal_conv1d_fwd_kernel, causal_conv1d_states_fwd_kernel, causal_conv1d_bwd_kernel, compute_dh0_kernel. The update/cache decode kernels stay untouched.
3. Operator entry: KDA accepts external chunk_indices (mirroring GDN's [GDN] Allow passing pre-computed chunk_indices to chunk_gated_delta_rule #1132); the conv module and both operators gain the graph switch — an explicit argument, or automatic detection of an active capture (see Open questions).
4. Tests: capture/replay bitwise parity against eager on identical inputs; repeated replay under varying cu_seqlens in the same static buffers; one coverage pass per gate branch; conv covered in the same capture/replay tests since it sits inside the same layer path. The context-parallel path needs no dedicated kernel change (its pre-process kernels are already static-grid); a two-rank capture/replay test guards it.

Limitations

• Out-of-budget inputs fall back to eager: no gain, no loss. The expected benefit is short-sequence regimes; at long sequences the launch overhead is amortized away.
• Capture requires a warmup: the kernels are autotuned, so the same static configuration must run eagerly first to fix the autotune choice; capturing cold triggers benchmarking during capture or records the wrong config. Warmup is a one-time caller cost.
• Each (T_max, N_max) tier holds resident memory; with large training batches the NT_max-sized intermediates raise the memory floor. Size tiers from the model's actual shapes before enabling.
• Do not combine with torch.compile's own cudagraph mode; pick one.

Open questions

• Switch shape: an explicit graph_capture=True argument versus automatically taking the static path whenever a capture is active. I lean explicit.
• Bucket ownership: operator layer exposes only the contract and callers (vLLM, training frameworks) own tiering and static buffers, versus FLA shipping a reference wrapper. I lean caller-owned, with an in-tree example.
• Rollout order: after GDN/KDA stabilize, extend the same contract to the remaining chunk operators, or keep the scope closed.

Rationale

A chunked forward issues a dozen-plus host launches per layer, and a backward more; the short convolution adds its own launches on the same path. For short inputs (typically ≤ 4096 tokens) this per-launch overhead is a visible fraction of end-to-end time — short prefill requests in serving, and short-sequence training steps alike. Graph capture collapses the whole sequence into one replay. Serving stacks already rely on capture for decode; this proposal extends the same benefit to prefill and training, on both NVIDIA and Ascend, with the eager path preserved bit-identically for everything outside the bounded regime.