Linear Attention

Attention Gym provides functional linear-attention operators with eager reference implementations.

This page is the API and implementation reference. Start with the operation you need:

• Train or prefill: Gated Delta Rule or Kimi Delta Attention, using the chunk_* operations.
• Carry state between calls or decode: the recurrent_* operations and the paged/decode variants in those sections.
• Split a packed stream across ranks: Context Parallelism.
• Understand changing sequence lengths under capture: the Ragged CUDA Graphs guide.

Gated Delta Rule

chunk_gdn and recurrent_gdn use the token-major layout [batch, sequence, heads, dimension] and return an output/state tuple. For each token, the scalar natural-log gate decays the previous state before the delta update, and the query reads the updated state:

decayed_state = exp(gate) * state
residual = beta * (value - decayed_state @ key)
state = decayed_state + outer(residual, key)
output = scale * state @ query

Public and persistent state uses axis order [N, H, V, K]; paged pools replace N with the slot count. chunk_gdn uses a chunk-parallel decomposition for training and prefill. recurrent_gdn consumes tokens in order for decoding, inference prefill, and state-carrying correctness checks. Like chunk_kda, chunk_gdn defaults to the repo-local fused chunk pipeline (impl="fused"); impl="reference" selects eager PyTorch. Pass kernel_options={"backend": "cudnn"} to select the optional cuDNN CuTeDSL backend; it accepts the same approximate split_forward and split_backward forgetting-horizon options as cuDNN KDA (see Kimi Delta Attention), with the same no-state restriction. They help only when few sequences leave the GPU idle and every head forgets quickly; heads that never forget stay uncut. recurrent_gdn(..., impl="fused") selects the inference-only Triton scan.

from attn_gym.linear import chunk_gdn

output, final_state = chunk_gdn(query, key, value, gate, beta, output_final_state=True)

Supported capabilities

• Fixed-length inputs in [batch, sequence, heads, dimension] layout and packed batch-one inputs with cu_seqlens.
• Separate recurrent and chunked operations with explicit initial and final state.
• CPU and CUDA execution through eager PyTorch operations.
• A repo-local fused chunk pipeline on CUDA capability 8.0+ with dense and packed inputs, grouped Q/K heads, tails, empty sequences, initial/final state, strict torch.compile, and CUDA Graph support. On CUDA capability 8.0+, paged_chunk_gdn advances selected [num_slots, H, V, K] cache rows in place for inference prefill without caller-side gather/scatter copies.
• An opt-in cuDNN fused chunk implementation with the same public training/state contract and direct paged-cache prefill on SM100/SM103.
• An inference-only fused recurrent implementation on CUDA, including mutable paged state caches shaped [num_slots, H, V, K] selected by state_indices.
• Autograd for the reference and fused chunk implementations, including initial-state gradients.
• Q/K/V share one dtype. FP16 and BF16 inputs use FP32 recurrence math and state while returning output in the Q dtype.
• Gate and beta may use independent floating dtypes and are converted to the recurrence compute dtype. A provided initial state uses FP32 for low-precision QKV.
• FP64 reference inputs retain FP64 recurrence math and state.

What the fused paths keep in FP32

Every optimized GDN and KDA path (repo-local fused and cuDNN) carries the recurrent state from initial_state to final_state, and the state cotangent from d_final_state to d_initial_state, in FP32 accumulators with FP32 decay. The chunk-entry state that each chunk reads, the per-chunk state checkpoints, and the per-chunk state-cotangent tape are staged in the Q/K/V dtype because they are tensor-core operands. BF16 shares the FP32 exponent range, so this only costs mantissa precision. FP16 also limits the range: state values that round beyond the FP16 finite range (about 65504) become non-finite when a chunk reads them, and cotangent values in the FP16 subnormal range (below 2^-14) lose precision or round to zero before they reach dgate, dk, and dv. Keep FP16 states and state cotangents within the FP16 normal range, for example through loss scaling, or use BF16 when state magnitudes can leave it. Use L2-normalized Q/K with FP16, as the training example does: unnormalized Q/K or unusually large V can likewise push attention, solve, and value intermediates outside the FP16 range between FP32-accumulating GEMMs. The cuDNN KDA backend additionally applies the per-chunk decay to the carried state, and to the state cotangent of its no-state local backward, through a diagonal MMA in the Q/K/V dtype rather than in FP32.

The repo-local fused chunk backend requires CUDA capability 8.0+, matching FP16 or BF16 Q/K/V, and equal K = V of 64 or 128, but has no cuDNN runtime dependency. Its scalar natural-log gate is not lower-bounded: kernels contract raw QK/KK before applying masked nonpositive causal decay differences. The public chunk size is BT64; internal 16x16 blocks are only the hierarchical triangular-solve representation. Layouts requiring int64 tensor offsets are rejected until every repo-local kernel has a wide-address path. For K = V = 128, backward uses the tuned CuTe kernels on SM100/SM103 and portable Triton kernels on other supported architectures. The K = V = 64 backward uses portable Triton on every architecture.

The cuDNN chunk backend requires the optional cudnn dependencies and SM100/SM103, FP16/BF16 Q/K/V, FP32 state, and K = V = 128 contract. It also consumes the scalar natural-log gate without a lower bound and uses exact execution by default.

The fused recurrent implementation requires CUDA with Triton, Q/K/V in FP16, BF16, or FP32, and K <= 256. Packed offsets must begin at zero, be nondecreasing, and end within the physical token capacity. Output rows beyond the terminal offset are inactive capacity and unspecified. Unsupported implementation choices fail rather than falling back to the reference.

Migration from the prototype API

• naive_recurrent_gated_delta_rule(...) and gated_delta_rule(..., mode="recurrent") become recurrent_gdn(...).
• naive_chunk_gated_delta_rule(...) and gated_delta_rule(..., mode="chunked") become chunk_gdn(...).
• Inputs use [batch, sequence, heads, dimension], matching the KDA operations; keyword callers use q, k, and v.
• The public chunk size is fixed at 64, matching KDA's fused decomposition.
• initial_state may be positional, and output_final_state controls the optional state output.
• Both operations return (output, final_state), matching the KDA operations.

recurrent_gdn_decode is the serving-specific one-token path, mirroring recurrent_kda_decode: it consumes the packed post-convolution QKV buffer plus raw gate and beta projections, computes the gate transform (-exp(A_log) * softplus(raw + dt_bias)), the beta sigmoid, and the query/key L2 normalization in-kernel, supports grouped q/k heads, and advances the paged FP32 state pool in place through state_indices, so no separate elementwise kernels run per decode step.

attn_gym.linear.chunk_gdn(q, k, v, gate, beta, initial_state=None, *, cu_seqlens=None, scale=None, output_final_state=False, impl=Impl.FUSED, kernel_options=None)

Apply chunk-parallel gated delta rule attention for training and prefill.

Inputs use the token-major layout [batch, sequence, heads, dimension]. The scalar natural-log gate decays the previous state before each beta-scaled delta update, and the query reads the updated state. Chunking changes only the decomposition and floating-point order of that recurrence. FP16 and BF16 inputs use FP32 recurrence math and state.

Parameters:

Name	Type	Description	Default
q	Tensor	Queries shaped [B, T, HK, K]. HK may divide the value head count H for grouped-head attention: each block of H // HK consecutive value heads shares one query/key head.	required
k	Tensor	Keys shaped like q and using the same dtype.	required
v	Tensor	Values shaped [B, T, H, V] and using the same dtype as q.	required
gate	Tensor	Floating per-token scalar natural-log decay shaped [B, T, H].	required
beta	Tensor	Floating per-token write gate shaped [B, T, H].	required
initial_state	Tensor | None	Initial recurrent state shaped [N, H, V, K] in the recurrence compute dtype or BF16, where N is the number of logical sequences. BF16 state is promoted to the compute dtype for recurrence math.	None
cu_seqlens	Tensor | None	Optional packed offsets shaped [N + 1] for batch-one inputs. They start at zero, never decrease, and may end before T; output beyond the terminal offset is unspecified.	None
scale	float | None	Query scale. Defaults to 1 / sqrt(K).	None
output_final_state	bool	Return the final recurrent state with the output.	False
impl	Impl | str	"fused" (default) uses the repo-local scalar chunk pipeline on CUDA capability 8.0+ with FP16/BF16 QKV and equal K = V of 64 or 128; "reference" uses eager PyTorch.	FUSED
kernel_options	SplitOptions | None	Backend-specific options for fused execution. The repo-local path is the default; {"backend": "cudnn"} selects the optional CuTeDSL 4.7 cuDNN backend. split_forward and split_backward opt into its approximate forgetting-horizon split schedules (see :class:attn_gym.linear.types.SplitOptions) for calls without initial_state or output_final_state.	None

Returns:

Type	Description
tuple[Tensor, Tensor | None]	The output in q.dtype and either the final recurrent state or None.

attn_gym.linear.paged_chunk_gdn(q, k, v, gate, beta, state_cache, state_indices, *, cu_seqlens=None, has_initial_state=None, scale=None, kernel_options=None)

Apply inference-only chunk GDN while advancing a paged state cache in place.

Requires CUDA capability 8.0 or newer.

Parameters:

Name	Type	Description	Default
q	Tensor	Queries shaped [B, T, HK, K]. HK may divide the value-head count.	required
k	Tensor	Keys shaped like q and using the same dtype.	required
v	Tensor	Values shaped [B, T, H, V] and using the same dtype as q.	required
gate	Tensor	Floating per-token scalar natural-log decay shaped [B, T, H].	required
beta	Tensor	Floating per-token write gate shaped [B, T, H].	required
state_cache	Tensor	Mutable FP32 or BF16 state pool shaped [num_slots, H, V, K]. Recurrence math remains FP32; BF16 affects persistent storage only. The optional cuDNN backend requires FP32.	required
state_indices	Tensor	Contiguous int32 slot indices, one per logical sequence. Positive, unique indices select slots to advance; non-positive indices produce zero output and leave the cache untouched.	required
cu_seqlens	Tensor | None	Optional packed offsets shaped [N + 1] for batch-one inputs. Values must start at zero, be nondecreasing, and end at or before the physical token capacity T; the tensor defines one interval per state_indices entry. Output rows past cu_seqlens[-1] are not written and hold unspecified values.	None
has_initial_state	Tensor | None	Optional contiguous boolean mask, one per logical sequence. False entries start from zero and overwrite the selected slot.	None
scale	float | None	Query scale. Defaults to 1 / sqrt(K).	None
kernel_options	SplitOptions | None	Backend options. The repo-local path is the default; {"backend": "cudnn"} selects the optional CuTeDSL 4.7 cuDNN backend, which requires 16-byte-aligned pool bases and slot origins. Split schedules are not supported.	None

Returns:

Type	Description
Tensor	The output in q.dtype. state_cache is advanced in place.

attn_gym.linear.recurrent_gdn(q, k, v, gate, beta, initial_state=None, *, cu_seqlens=None, scale=None, output_final_state=False, state_indices=None, has_initial_state=None, autotune=True, impl=Impl.FUSED)

Apply recurrent gated delta rule attention for decoding and inference prefill.

The recurrence consumes tokens in order, carrying an explicit [N, H, V, K] state. Inputs and outputs use the token-major layout [batch, sequence, heads, dimension]. FP16 and BF16 inputs use FP32 recurrence math. State may use BF16 storage and is promoted for recurrence math.

Parameters:

Name	Type	Description	Default
q	Tensor	Queries shaped [B, T, HK, K]. HK may divide the value head count H for grouped-head attention: each block of H // HK consecutive value heads shares one query/key head.	required
k	Tensor	Keys shaped like q and using the same dtype.	required
v	Tensor	Values shaped [B, T, H, V] and using the same dtype as q.	required
gate	Tensor	Floating per-token scalar natural-log decay shaped [B, T, H].	required
beta	Tensor	Floating per-token write gate shaped [B, T, H].	required
initial_state	Tensor | None	Initial recurrent state shaped [N, H, V, K] in the recurrence compute dtype or BF16, where N is the number of logical sequences. BF16 state is promoted to the compute dtype for recurrence math.	None
cu_seqlens	Tensor | None	Optional packed offsets shaped [N + 1] for batch-one inputs. They start at zero, never decrease, and may end before T; output beyond the terminal offset is unspecified.	None
scale	float | None	Query scale. Defaults to 1 / sqrt(K).	None
output_final_state	bool	Return the final recurrent state with the output. Rejected with state_indices, which advances the state pool in place instead.	False
state_indices	Tensor | None	Optional contiguous int32 slot indices selecting rows of a mutable initial_state pool shaped [num_slots, H, V, K]. Positive indices must be unique and in [1, num_slots); nonpositive indices produce zero output and leave the pool untouched. These value constraints are caller responsibilities.	None
has_initial_state	Tensor | None	Optional contiguous boolean mask indicating whether each selected slot should be loaded. False entries mark freshly assigned slots whose contents are garbage: they start from zero and overwrite the slot, even for empty sequences.	None
autotune	bool	Benchmark candidate value-tile sizes for non-paged execution when true; paged execution always uses deterministic heuristics because it mutates the state cache.	True
impl	Impl | str	"fused" uses the inference-only Triton scan; "reference" uses eager PyTorch with autograd support.	FUSED

Returns:

Type	Description
tuple[Tensor, Tensor | None]	The output in q.dtype and either the final recurrent state or None.

attn_gym.linear.recurrent_gdn_decode(packed_qkv, raw_gate, raw_beta, A_log, dt_bias, state_cache, state_indices, *, has_initial_state=None, scale=None, out=None)

Run one-token paged GDN decode with preprocessing fused into the recurrence.

One Triton kernel slices the post-convolution QKV buffer, computes the gate as -exp(A_log) * softplus(raw_gate + dt_bias) and the write gate as sigmoid(raw_beta), L2-normalizes q and k, and advances the selected cache slots in place, so serving callers launch no separate elementwise kernels.

Parameters:

Name	Type	Description	Default
packed_qkv	Tensor	Post-convolution QKV shaped [B, HK*K + HK*K + H*V]. Each token stores [Q for HK heads | K for HK heads | V for H heads]; within each section head rows are contiguous. HK may divide H for grouped-head attention: each block of H // HK consecutive value heads shares one query/key head.	required
raw_gate	Tensor	Unactivated per-head gate projection shaped [1, B, H], matching the vLLM-style single-token decode convention used by recurrent_kda_decode.	required
raw_beta	Tensor	Unactivated write gate shaped [1, B, H].	required
A_log	Tensor	FP32 per-head log decay parameter shaped [H].	required
dt_bias	Tensor	FP32 per-head gate bias shaped [H].	required
state_cache	Tensor	FP32 or BF16 paged state pool shaped [num_slots, H, V, K]. Slots may have padding between them but each [H, V, K] row must be dense. Recurrence math remains FP32 and the updated state is cast to the pool dtype. K must be at most 256.	required
state_indices	Tensor	Contiguous int32 slot indices shaped [B]. Non-positive indices are padding/null entries: they produce zero output and leave the cache untouched. Each positive index must be in [1, num_slots) and unique among active rows because duplicate in-place updates race. These value constraints are caller responsibilities and are not host-validated.	required
has_initial_state	Tensor | None	Optional contiguous boolean mask, one per sequence. False entries mark freshly assigned slots whose contents are garbage: the step starts from the zero state and overwrites the slot.	None
scale	float | None	Query scale. Defaults to 1 / sqrt(K).	None
out	Tensor | None	Optional caller-owned contiguous output buffer shaped [1, B, H, V] in packed_qkv.dtype on the same device. When supplied, the kernel writes into and returns this exact tensor. It must not alias any input.	None

Returns:

Type	Description
Tensor	Decode output shaped [1, B, H, V] in packed_qkv.dtype. This is out
Tensor	itself when a buffer is supplied. The operation is inference-only and advances
Tensor	state_cache in place.

Gate Transforms

gate_transform turns a raw gate projection into the FP32 natural-log decay consumed by the chunked and recurrent operations. The transform kind and the gate shape are independent:

kind	formula	range
"bounded"	lower_bound * sigmoid(exp(A_log) * (raw + dt_bias))	(lower_bound, 0)
"softplus"	-exp(A_log) * softplus(raw + dt_bias)	(-inf, 0)

Either kind applies to per-head gates (raw [B, T, H], dt_bias [H], the GDN convention) or per-channel gates (raw [B, T, H, D], dt_bias [H, D], the KDA convention). kind is required: a default could silently select the wrong recurrence convention. bound_gate is the KDA-specific spelling of kind="bounded" for per-channel gates with lower_bound=-5.0. Both producers default to fastmath=True; see Fast math and precision.

impl="fused" has a portable and a fast route. Per-channel D=128 gates on CUDA capability 9.0+ use the CuTeDSL kernels for both kinds (near the bf16->fp32 cast roofline with fastmath=True); every other "softplus" configuration uses a portable Triton kernel, and per-head "bounded" gates are reference-only. Both backwards reduce the A_log and dt_bias gradients in FP32 from per-block partials, so they are deterministic. The decode kernels (recurrent_kda_decode, recurrent_gdn_decode) evaluate the same transforms in-kernel from the raw projections and do not call this operation. Unbounded softplus gates exceed the finite range some fused chunk-KDA schedules assume; keep per-channel softplus gates for backends that document support for it.

from attn_gym.linear import chunk_gdn, gate_transform

gate = gate_transform(raw_gate, A_log, dt_bias, kind="softplus")
output, final_state = chunk_gdn(q, k, v, gate, beta)

attn_gym.linear.gate_transform(raw_gate, A_log, dt_bias, *, kind, lower_bound=None, fastmath=True, impl=Impl.FUSED)

Map raw gate projections to per-token natural-log decays.

kind="bounded" computes lower_bound * sigmoid(exp(A_log) * (raw_gate + dt_bias)) (Kimi/KDA convention, range (lower_bound, 0)); kind="softplus" computes -exp(A_log) * softplus(raw_gate + dt_bias) (Mamba2/GDN convention, unbounded below). Arithmetic and output use FP32; a fused raw-gate gradient matches the input dtype. The decode kernels evaluate the same transforms in-kernel from the raw projections.

Parameters:

Name	Type	Description	Default
raw_gate	Tensor	Floating-point projection output shaped [B, T, H] (per-head gate) or [B, T, H, D] (per-channel gate).	required
A_log	Tensor	FP32 per-head log scale shaped [H].	required
dt_bias	Tensor	FP32 bias shaped [H] or [H, D] to match raw_gate.	required
kind	GateTransform | str	Transform kind; required so a default cannot silently select the wrong recurrence convention.	required
lower_bound	float | None	Finite nonpositive floor, required by "bounded" and rejected by "softplus".	None
fastmath	bool	Allow approximate amplitude, sigmoid, and softplus math (default True). False selects non-fast primitives for forward and backward. The reference implementation uses normal PyTorch math for either value.	True
impl	Impl | str	"reference" uses ordinary PyTorch. "fused" uses private CuTeDSL kernels for per-channel D=128 gates on CUDA capability 9.0+ (both kinds) and a portable Triton kernel for every other "softplus" configuration. Fused paths require FP16, BF16, or FP32 raw_gate; per-head "bounded" gates are reference-only.	FUSED

Raises:

Type	Description
ValueError	On shape, lower_bound, or unsupported fused configurations.

attn_gym.linear.GateTransform

Bases: str, Enum

Pointwise transform from a raw gate projection to a natural-log decay.

Both kinds apply to per-head ([B, T, H]) and per-channel ([B, T, H, D]) gates; the kind is independent of the gate shape.

Stateful Short Convolution

causal_conv1d is the differentiable dense or packed operation. It accepts compact per-sequence history and can return a newly allocated final history. Serving callers can keep history in a shared [num_slots, W - 1, C] pool instead: paged_causal_conv1d reads selected slots during multi-token prefill and advances them in place, while causal_conv1d_decode does the same for one token per sequence. Both are inference-only and support padding between slots. Positive state_indices must be unique and within the pool; non-positive routes return zero without touching it. has_initial_state=False gives paged prefill a zero history before it overwrites a fresh slot.

attn_gym.linear.causal_conv1d(x, weight, *, activation=None, cu_seqlens=None, initial_state=None, return_final_state=False, forward_config=None, input_grad_config=None, weight_grad_config=None, persistent_tma_input_gradient=False)

Apply causal depthwise convolution with an optional fused activation.

Width is inferred from weight.shape[1] and compile-time specialized. Schedule fields control register shapes, vector partitioning, unrolled loops, and block mapping, so changing one compiles and caches a distinct kernel.

Parameters:

Name	Type	Description	Default
x	Tensor	CUDA FP16, BF16, or FP32 input with shape [B, T, C] and contiguous channels. Batch and token strides are arbitrary, so slices, permutations, and expanded views are read in place. Each batch row is convolved as an independent sequence.	required
weight	Tensor	Contiguous depthwise weights with shape [C, W] matching x dtype and device.	required
activation	str | None	Optional activation fused into the convolution epilogue and recomputed by the backward kernels. None applies no activation; "silu" is built in and :func:register_activation adds custom names. The name is an ordinary string argument, so compiled graphs specialize on it like any other static scalar.	None
cu_seqlens	Tensor | None	Optional contiguous CUDA int32 offsets delimiting independent sequences in a packed [1, T, C] input. Offsets must be nondecreasing, begin at zero, and end at an active endpoint L <= T; repeated offsets represent empty padding slots for static-shape CUDA Graph replay. Only output positions before L are defined.	None
initial_state	Tensor | None	Optional causal history with shape [N, W - 1, C], where N is the dense batch size or the number of packed sequences. Absence is equivalent to an all-zero history.	None
return_final_state	bool	Return the final W - 1 input positions with the output.	False
forward_config	ShortConvConfig | None	Optional forward schedule specialization.	None
input_grad_config	ShortConvConfig | None	Optional input-gradient schedule specialization.	None
weight_grad_config	ShortConvConfig | None	Optional weight-gradient schedule specialization.	None
persistent_tma_input_gradient	bool	Use an explicit machine-bounded persistent worker grid for packed TMA input gradients. Disabled by default because it favors overcaptured replays and can regress full-capacity backward. Requires x's base and outer strides to be 16-byte aligned; backward raises otherwise.	False

Returns:

Type	Description
Tensor | tuple[Tensor, Tensor]	A contiguous tensor with the same shape, dtype, and device as x. The input
Tensor | tuple[Tensor, Tensor]	gradient is also computed contiguous and flows back through any view that produced
Tensor | tuple[Tensor, Tensor]	x. For packed input, only output and input-gradient positions in [0, L) are
Tensor | tuple[Tensor, Tensor]	defined; parameter, state, and final-state values exclude the inactive suffix. When
Tensor | tuple[Tensor, Tensor]	return_final_state is true, also returns [N, W - 1, C] history.

attn_gym.linear.paged_causal_conv1d(x, weight, state_cache, state_indices, *, activation=None, cu_seqlens=None, has_initial_state=None)

Apply multi-token convolution while advancing a paged history in place.

Parameters:

Name	Type	Description	Default
x	Tensor	[B, T, C] input with contiguous channels. Batch and token strides are arbitrary, as in a column slice of a fused projection, and are read in place. With cu_seqlens, B must be one and the physical token axis contains packed logical sequences.	required
weight	Tensor	Contiguous depthwise weights [C, W] with W >= 2.	required
state_cache	Tensor	Mutable history pool [num_slots, W - 1, C]. Each slot must be compact, while padding between slots is supported. The storage base must be 16-byte aligned.	required
state_indices	Tensor	Contiguous int32 slot indices, one per logical sequence. Positive indices must be unique and less than state_cache.shape[0]. Non-positive entries produce zero output and leave the pool untouched.	required
activation	str | None	Any registered activation name, or None.	None
cu_seqlens	Tensor | None	Optional contiguous int32 packed boundaries [N + 1].	None
has_initial_state	Tensor | None	Optional contiguous bool mask with one entry per sequence. False entries read zero history before overwriting the slot.	None

Returns:

Type	Description
Tensor	The activated contiguous output with the same shape and dtype as x. For packed
Tensor	input, only rows before cu_seqlens[-1] are defined. state_cache is advanced in place.
Tensor	Empty fresh sequences clear their selected history; empty resumed sequences
Tensor	preserve it.

Batch size, token capacity, and sequence count bind at launch and reuse a compiled kernel. Packed/dense mode, optional-mask presence, channels, width, dtype, activation, schedule, and address width remain specialized.

This operation is inference-only. Use :func:causal_conv1d for autograd.

attn_gym.linear.causal_conv1d_decode(x, weight, state, *, activation=None, state_indices=None, has_initial_state=None, forward_config=None)

Advance the convolution by one token per sequence over a paged history.

The decode counterpart to :func:causal_conv1d, fused: one program reads a slot's history, convolves the new token in, and writes the shifted history back, so no gather or scatter surrounds the call.

Parameters:

Name	Type	Description	Default
x	Tensor	[num_sequences, C] input, one token per sequence, matching weight and state in dtype. Channels must be contiguous; the row stride is arbitrary, as in a column slice of a fused projection, and rows are read in place.	required
weight	Tensor	Contiguous depthwise weights [C, W]. W >= 2; a width-1 convolution carries no history and needs no update.	required
state	Tensor	Causal history [num_slots, W - 1, C] holding each slot's trailing W - 1 inputs, oldest first -- the layout :func:causal_conv1d takes as initial_state and returns as its final state. Advanced in place: the rows shift left and x becomes the newest. Each slot must be contiguous, but padding between slots is supported. The storage must be 16-byte aligned, since the launcher cannot realign it without losing the update.	required
activation	str | None	Any name registered with :func:register_activation, or None.	None
state_indices	Tensor | None	Optional contiguous int32 slot indices, one per sequence, selecting rows of a paged state pool. Without them, sequence i uses slot i. Positive slots must be distinct. Non-positive entries produce zero output and leave the pool untouched.	None
has_initial_state	Tensor | None	Optional contiguous CUDA bool mask, one entry per sequence. True reads the selected history; False starts from zeros without reading that slot, even if it contains NaNs. None resumes every live slot. Non-positive explicit indices remain inactive regardless of this mask; without state_indices, slot zero is a valid sequence slot.	None
forward_config	ShortConvConfig | None	Optional schedule specialization.	None

Returns:

Type	Description
Tensor	The activated contiguous output [num_sequences, C], same dtype as x.

Inference-only: the single-token step has no backward, so gradient-tracking inputs are rejected when autograd is enabled.

Kimi Delta Attention

The KDA references use token-major tensors: query, key, and per-channel gate are [batch, sequence, heads, key_dimension]; value is [batch, sequence, heads, value_dimension]; beta is [batch, sequence, heads]. Both recurrent and chunked forms support ordinary PyTorch autograd and an optional recurrent state shaped [N, H, V, K]. Paged prefill and decode use the same axis order with the leading dimension interpreted as cache slots.

examples/linear/delta_rule_training.py builds these operations into a small trainable [B, T, hidden_size] -> [B, T, hidden_size] attention module (--variant=kda; the same module runs the GDN recipe with --variant=gdn, swapping only the gate and the core). To mirror the main Kimi block structure, projected Q/K/V pass through a causal depthwise SiLU convolution, the forget and output gates use two-stage factorized projections, and the per-head output uses learned RMS normalization before sigmoid gating. --backend=reference uses the PyTorch reference throughout. The optimized chunk_kda core used by --backend=fused supports Ampere or newer and owns the BT64 scan behind the same public boundary. The complete fused KDA variant currently requires Hopper because its separate bound_gate producer uses TMA. Implementations may inline the scan, but cumulative gates and chunk boundaries are not caller-visible representations.

The example explicitly runs projections in BF16 while retaining FP32 parameters and gate math; no ambient autocast context is required. Distributed mixed-precision policies must preserve A_log and dt_bias as FP32 when parameters are materialized, rather than only casting activations inside forward. A module-wide BF16 FSDP policy violates that contract. A correctness-first integration can keep the KDA unit under an FP32 policy while its projections explicitly compute in BF16; isolating only the strict-FP32 decay state is a future bandwidth optimization. The optimized chunk_kda core requires Ampere or newer and head_dim=128; its public boundary accepts FP16, BF16, or FP32 inputs. SM100/SM103 select the specialized CuTe path; SM120 and earlier supported architectures use the portable Triton stages. Homogeneous FP16 and BF16 Q/K/V stay in their input dtype, while FP32 or mixed-dtype inputs retain the existing BF16 normalization. The core chunks internally at 64 tokens. Complete B=1 inputs whose length is a multiple of the chunk size run on the direct dense route; other dense [B, T, H, D] inputs are lowered internally to equal-length packed sequences, while chunk_kda(..., cu_seqlens=offsets) accepts explicitly packed [1, T, H, D] inputs. All forms carry sequence boundaries through the forward, backward, and recurrent states; logical sequences may have tails or be empty. For fixed-capacity execution, the terminal offset may be smaller than physical T; primitive forward values outside [0, cu_seqlens[-1]) are unspecified. paged_chunk_kda uses the same Ampere-or-newer forward route while updating selected FP32 state-cache slots in place. The internal reverse scan returns zero cotangents for inactive gate rows. This does not sanitize arbitrary parameterized gate producers: callers still need the masking rules below because 0 * NaN can poison their reductions.

FP16 intermediate range

Use L2-normalized Q/K with FP16, as the training example does, and keep states and state cotangents within the FP16 normal range. See "What the fused paths keep in FP32" in the Gated Delta Rule section for the boundary between the FP32 carry and the Q/K/V-dtype staging that this shares with GDN.

A captured graph with sequence capacity N keeps cu_seqlens.shape == (N + 1,). If a replay has M <= N nonempty sequences and L <= T active tokens, repeat the terminal endpoint through the unused tail:

[0, sequence_start_1, ..., L, L, ..., L]

The repeated ranges are ordinary empty sequences. Stateful APIs therefore retain N state rows even when only M sequences are nonempty. Both L and M may change on replay, but physical token capacity T and metadata capacity N may not.

Scheduling for this over-capture regime is automatic and selected independently for each eligible chunk-parallel ragged kernel: a kernel switches to a bounded persistent worker grid only when its capacity task count exceeds a few waves of that grid. Persistent workers stride over the active chunk count built on device from cu_seqlens, so chunk-compute launch overhead tracks active work rather than captured capacity. Capacity-sized initialization and reduction work may remain. Exact or mildly padded shapes keep capacity-sized STATIC grids whose padding CTAs return immediately; there is no user-facing scheduling knob.

Callers that use dynamic active lengths within fixed-capacity tensors must opt in to masking. Ragged primitives read only [0, cu_seqlens[-1]) from token-shaped inputs, including output cotangents, and leave the suffix of token-shaped outputs and input gradients unspecified. Four edge rules make those primitives safe to compose:

• Caller buffer → ordinary operation: value-mask the buffer. A zero cotangent does not neutralize a NaN activation in a weight reduction: 0 * NaN is still NaN.
• Parameterized producer → ragged primitive: add a gradient barrier so the primitive's unspecified input-gradient suffix cannot enter the producer's reduction.
• Ragged primitive → ordinary operation: value-mask the primitive output before the ordinary operation saves it for backward.
• Ragged primitive → ragged primitive: do nothing; neither primitive reads the inactive suffix.

Construct one device-resident predicate inside the captured graph and reuse it at every boundary. This keeps all masks consistent and lets replay recompute them when cu_seqlens[-1] changes without a host read or recapture.

from attn_gym.linear.kda import (
    active_token_mask,
    mask_inactive_token_gradients,
    mask_inactive_tokens,
)

active_mask = active_token_mask(hidden, cu_seqlens)
hidden = mask_inactive_tokens(hidden, active_mask)  # Caller → ordinary projection.
projected = input_projection(hidden)
projected = mask_inactive_token_gradients(projected, active_mask)
stage = ragged_primitive_one(projected, cu_seqlens=cu_seqlens)
output = ragged_primitive_two(stage, cu_seqlens=cu_seqlens)  # No mask between primitives.
output = mask_inactive_tokens(output, active_mask)
output = output_projection(output)
output = mask_inactive_token_gradients(output, active_mask)  # Model boundary.

mask_inactive_token_gradients(x, active_mask) preserves forward values while zeroing inactive automatic-differentiation paths. For contiguous packed CUDA tensors in eager mode its autograd path aliases x in the forward and row-masks tangents and cotangents; compiled graphs and unsupported layouts keep the torch.where(mask, x, x.detach()) form (which materializes an elementwise result) so Inductor can own fusion. Automatic-differentiation paths are zero on inactive rows, and subsequent derivatives inherit the same mask. Recurrent and convolution states have one row per logical sequence rather than one row per physical token, so token masks must not be applied to them.

attn_gym.linear.kda.active_token_mask(x, cu_seqlens)

Construct a reusable predicate for the active packed token prefix.

Construct the predicate inside CUDA Graph capture so replay recomputes it from the device-resident cu_seqlens[-1] endpoint.

attn_gym.linear.kda.mask_inactive_tokens(x, active_mask)

Zero inactive token values and derivatives using a reusable mask.

Passing None returns x unchanged, which keeps the default dense and exact-packed paths free of masking operations.

attn_gym.linear.kda.mask_inactive_token_gradients(x, active_mask)

Preserve values while blocking inactive automatic-differentiation paths.

Place this barrier between a parameterized producer and a primitive that ignores inactive forward rows but leaves their input-gradient suffix undefined. Forward values pass through unchanged while cotangents and tangents are row-masked, so automatic-differentiation paths are zero on inactive tokens and subsequent derivatives inherit the same mask. On the eligible eager CUDA path the result aliases x storage; compiled and fallback paths materialize torch.where. Passing None returns x unchanged.

DeltaRuleAttention.forward passes explicit offsets to its short convolution and chunk_kda (or chunk_gdn); the selected implementation owns its sequence-local scan and scheduling. Set mask_inactive_capacity=True only when the packed tensor reserves physical rows beyond cu_seqlens[-1]; dense and exact-packed callers leave it disabled and pay no masking cost. The optimized boundaries are first-order and do not support higher-order autograd. Run python examples/linear/delta_rule_training.py --backend=fused --packed --batch-size=4 --tokens=256 to sample token-level lengths from a truncated Zipf distribution, pack them exactly into one physical batch, print their cu_seqlens, and pass those offsets through the complete training step. The complete composed core forward and backward use private custom operators with fake-tensor registrations and first-order autograd wrappers, so fused chunk_kda supports strict torch.compile(fullgraph=True) and CUDA Graph capture for fixed physical token capacity and sequence count. Packed reference execution is eager-only: it reads cu_seqlens on the host to run each logical sequence independently. Boundary values and the active token count may change on replay; changing the physical token capacity or sequence count requires recompilation or recapture. Pass --compile to compile the complete example as one full graph. This keeps the custom KDA core behind its registered operator boundary while allowing Inductor to fuse the surrounding PyTorch normalization and remaining pointwise work. The bounded gate itself uses private CuTeDSL forward and backward operators. It can be combined with --profile; compilation warmups run before the trace starts. Profiling requires transformer-nuggets and writes a native Perfetto .pftrace through the shared attn_gym.testing.profile_trace helper. Like FLA's default training path, its backward recomputes the W/U, gated Q/K, recurrent-state, and corrected-value intermediates instead of retaining them across the forward/backward boundary.

Graph-safe active-token replay does not by itself make complete model time proportional to L. The ragged short convolution, gate scan, and KDA core avoid reading inactive token values, but the example's projections, output normalization, output gate, and output projection still process physical capacity T. An end-to-end integration needs active-prefix-aware surrounding operations to turn smaller L into a comparable step-time reduction. Million-token training also requires model-level activation checkpointing and context parallelism; this single-device example implements neither distributed policy.

The module can sit behind a transformer layer's attention slot while state is threaded explicitly:

from examples.linear.delta_rule_training import DeltaRuleAttention

attention = DeltaRuleAttention(hidden_size=512, num_heads=4, head_dim=128).cuda()
first = attention(hidden_states[:, :128], return_final_state=True)
second = attention(
    hidden_states[:, 128:],
    first.final_state,
    initial_conv_state=first.final_conv_state,
    return_final_state=True,
)

The example intentionally is not checkpoint-compatible with Kimi K3. A model adapter must still provide exact checkpoint parameter names and initialization, the model's packed-sequence metadata, cache layout, and distributed execution policy. The short convolution, factorized gates, and learned gated RMS normalization match the production structure but remain ordinary PyTorch teaching implementations.

The two public KDA cores share one gate contract: gate is the per-token natural-log decay before any prefix sum. chunk_kda owns the natural-log-to-log2 conversion and its sequence-local BT64 cumulative sum; recurrent_kda performs only the conversion because recurrence consumes one token decay at a time. This keeps chunking out of model code and lets callers switch execution modes without changing gate representation. Custom producers should return finite, nonpositive values. The repo-local fused backend requires approximately [-5.914, 0]. cuDNN BF16 instead requires every aligned 16-token per-channel sum to exceed -126 * ln(2); a uniform lower_bound >= -5.45 is safe, including the common -5 bound. cuDNN stages these MMA operands in the Q/K dtype rather than TF32, so cuDNN FP16 does not support usual model-range gates. These limits are not checked at runtime. The training example uses the Kimi-style FP32 transform lower_bound * sigmoid(exp(A_log) * (raw_gate.float() + dt_bias)), but that model policy is not part of the public KDA API.

Both cores select their implementation with impl: "fused" runs the optimized kernels and enforces their constraints (the chunked core requires head_dim=128 and Ampere or newer, preserves homogeneous FP16/BF16 Q/K/V, normalizes FP32 or mixed inputs to BF16, and chunks at 64 tokens; the fused recurrent scan is inference-only), while "reference" runs the eager FP32 oracle behind the identical packed contract on any hardware and head dimension, and stays differentiable. There is no automatic fallback between implementations, and the chunk-versus-recurrent switch is caller policy (on B200 the scan wins below roughly 32 tokens per sequence).

recurrent_kda_decode is the serving-specific one-token path. It consumes channel-major post-convolution QKV ([Q for all heads | K for all heads | V for all heads]), raw gate and beta projections, and a paged state cache. Q/K normalization, gate activation, beta sigmoid, recurrence, output, and state-cache update run in one Triton kernel. Callers may provide a stable output buffer for allocation-free CUDA Graph replay.

Token-at-a-time decode groups floating-point additions differently than 64-row chunk_kda training, so decode drifts from training numerics. examples/linear/kda_zero_drift_decode.py replays the current partial chunk with chunk_kda from the last FP32 boundary state, which makes decode output bitwise-identical to training.

chunk_kda(..., kernel_options={"backend": "cudnn"}) selects an opt-in SM100/SM103 FP16/BF16 training backend for applications that already hold per-token natural-log gate increments. Q/K/V share one dtype. It uses public CuTeDSL 4.7 primitives and has no cuDNN Frontend runtime dependency. Like the other implementations, it returns (output, final_state); request the latter with output_final_state=True. Dense no-state calls with complete 64-token chunks use the dense kernels and partial tails take the packed path with synthesized boundaries; packed/stateful calls pass explicit contiguous int32 boundaries and FP32 [N, H, V, K] states. Q/K/V/gate/state accept TMA-compatible innermost modes with aligned dynamic outer strides; beta requires an element-aligned contiguous inner mode. By default the forward uses the unsplit recurrence: one persistent work item per sequence and head, which leaves the GPU underused for a few long sequences at low head counts. FP16 and BF16 use unsplit backward execution by default; eligible packed no-state long contexts may use the cuDNN backward with one unsplit work item per sequence and head. Callers that guarantee normalized keys, post-sigmoid beta, and nonpositive decay increments may set split_backward and/or split_forward to True in a cuDNN kernel_options mapping to opt into approximate forgetting-horizon splitting: a sequence is cut only where its cumulative decay has crossed the kernel's threshold, and each cut item rebuilds its entry state from zero over a short warmup window. Gates that never forget produce no cuts and reproduce the unsplit result exactly; when cuts are accepted the result stays within the low-precision error budget while exposing several times the parallel work. The threshold is a margin of bits past the output dtype's half-ulp, not an underflow; see the module docstring and DEFAULT_LOG2_THRESHOLD in attn_gym/linear/_delta_rule/cudnn_fe/common/split_k.py. split_forward affects only the forward recurrence: unless split_backward is also enabled, backward computes the exact unsplit recurrence's gradient rather than the derivative of the split forward. The implementation chooses the split count from the input geometry; no public split-size knob is exposed. Split schedules currently require a no-state call, so context parallelism never uses them. paged_chunk_kda(..., kernel_options={"backend": "cudnn"}) uses the same unsplit forward while updating selected cache slots directly; paged execution never uses forgetting-horizon splitting. Install the cudnn extra to use this backend.

kernel_options={"schedule": "persistent"} is for CUDA graphs whose replays can carry far fewer tokens than the capacity they were captured for. The default static grid launches one CTA per capacity chunk, so a replay at 1/16 of capacity still pays for the empty CTAs; the persistent grid strides over only the active work. Below roughly a quarter of capacity active, persistent wins (measured 0.1-0.7x the static time on GB200); at full capacity it loses (1.1-2.4x). Outputs are bitwise identical. Requires the fused backend and, for packed inputs, the TMA output path (Hopper or newer).

Fast math and precision

fastmath controls gating exponentials in KDA's forward, backward, recomputation, and CP summaries.

Backend	True (default)	False
Repo-local fused KDA	Fast, approximate exponentials	Non-fast exponentials that preserve FP32 subnormal results
impl="reference"	Normal PyTorch math	Normal PyTorch math
kernel_options={"backend": "cudnn"}	Native backend math	ValueError: unsupported

Use the same value for staged forward and backward preparation. Choose it before warmup and CUDA Graph capture. FP16/BF16 rounding and the documented gate-range limits apply in both modes.

bound_gate and gate_transform have separate flags, also defaulting to True. For non-fast math in both gate production and attention:

from attn_gym.linear import chunk_kda
from attn_gym.linear.kda import bound_gate

gate = bound_gate(raw_gate, A_log, dt_bias, fastmath=False)
output, final_state = chunk_kda(q, k, v, gate, beta, fastmath=False, output_final_state=True)

The training and CP examples use --fastmath by default; --no-fastmath selects False. Paged chunk and replay execution use True.

attn_gym.linear.chunk_kda(q, k, v, gate, beta, initial_state=None, *, cu_seqlens=None, scale=None, output_final_state=False, fastmath=True, autotune=True, impl=Impl.FUSED, kernel_options=None)

Apply chunk-parallel KDA for training and prefill.

Parameters:

Name	Type	Description	Default
q	Tensor	Queries shaped [B, T, H, K], scaled by scale internally. Use L2-normalized Q/K with fused FP16: unnormalized values can overflow the FP16 intermediates passed between FP32-accumulating GEMMs.	required
k	Tensor	Keys shaped like q and subject to the same fused FP16 range limitation.	required
v	Tensor	Values shaped [B, T, H, V].	required
gate	Tensor	Finite, nonpositive per-token natural-log decay shaped like q. At each token the previous state is multiplied channelwise by exp(gate). Pass per-token values, not cumulative gates; chunking and log-base conversion are internal. The repo-local fused implementation requires approximately [-5.914, 0]. cuDNN BF16 instead requires every aligned 16-token per-channel sum to exceed -126 * ln(2); a uniform lower_bound >= -5.45 is safe, including the common -5 bound. cuDNN stages these MMA operands in the Q/K dtype rather than TF32, so cuDNN FP16 does not support usual model-range gates. These limits are not checked at runtime and do not apply to reference or recurrent execution.	required
beta	Tensor	Per-token write gate shaped [B, T, H].	required
initial_state	Tensor | None	Starting recurrent state, with one [H, V, K] entry per logical sequence.	None
cu_seqlens	Tensor | None	Packed offsets shaped [N + 1] for batch-one inputs, as contiguous int32 on q.device; they start at zero, never decrease, may repeat for empty sequences whose states pass through, and may end before T. When the endpoint is below T, output values and token-input gradients beyond it are undefined. Fixed-capacity callers must sanitize outputs with mask_inactive_tokens and protect parameterized input producers with mask_inactive_token_gradients.	None
scale	float | None	Query scale, applied inside the kernels in FP32. Defaults to 1 / sqrt(K).	None
output_final_state	bool	Return the final recurrent state with the output.	False
fastmath	bool	True is the default because it’s faster, and turning it off usually didn’t improve accuracy in our tests. If you’re seeing unexpected outputs or gradients, especially with strong negative gates, try False and compare against the reference. It can preserve small contributions that fast math drops when intermediate values fall into the subnormal range. The reference backend uses normal PyTorch math for either value; native cuDNN requires True. Configure gate production separately with bound_gate or gate_transform.	True
autotune	bool	Benchmark candidate fused-kernel configurations when true. cuDNN accepts this argument for API compatibility but uses its fixed schedule.	True
impl	Impl | str	"fused" uses optimized kernels and "reference" uses differentiable eager PyTorch in FP32. There is no automatic fallback.	FUSED
kernel_options	KernelOptions | None	Backend-specific options. This is a sneaky BC trick. It is annoying to have a bunch of kwargs that its hard to know when they apply but, this trick allows us to add new options without breaking the API. I will do my best to add new options by providing a nice typed dict and hopefully your agent will now how to use this and be able to find the right option to wiggle.	None

Returns:

Type	Description
Tensor	The output in q.dtype and either an FP32 final state with one entry
Tensor | None	per logical sequence or None.

attn_gym.linear.paged_chunk_kda(q, k, v, gate, beta, state_cache, state_indices, *, cu_seqlens=None, has_initial_state=None, autotune=True, kernel_options=None, replay_state=None)

Apply inference-only chunk KDA while advancing a paged state cache in place.

Paged and replay-backed calls use fastmath=True.

Parameters:

Name	Type	Description	Default
q	Tensor	Queries shaped [B, T, H, K], scaled by 1/sqrt(K) internally.	required
k	Tensor	Keys shaped like q.	required
v	Tensor	Values shaped [B, T, H, V].	required
gate	Tensor	Finite, nonpositive per-token natural-log decay shaped like q. Pass per-token values, not cumulative gates; chunking and log-base conversion are internal, matching :func:chunk_kda.	required
beta	Tensor	Per-token write gate shaped [B, T, H].	required
state_cache	Tensor	Mutable FP32 or BF16 state pool shaped [num_slots, H, V, K]. Recurrence math remains FP32; BF16 affects persistent storage only. The optional cuDNN backend requires FP32.	required
state_indices	Tensor	Contiguous int32 slot indices, one per logical sequence. Positive, unique indices select cache slots to read and advance; non-positive indices produce zero output and leave the cache untouched.	required
cu_seqlens	Tensor | None	Packed offsets shaped [N + 1] for batch-one inputs, as contiguous int32 on q.device. They start at zero, never decrease, may repeat for empty sequences, and may end before T; output rows past cu_seqlens[-1] are not written and hold unspecified values.	None
has_initial_state	Tensor | None	Optional contiguous boolean mask, one per logical sequence. False entries ignore the selected cache contents and start from zero before advancing that slot. This is useful when a slot has just been assigned.	None
autotune	bool	For the repo-local fused backend, benchmark candidate kernel configurations when true; cuDNN uses its fixed schedule.	True
kernel_options	KernelOptions | None	Backend-specific options. {"backend": "cudnn"} selects the optional CuTeDSL 4.7 cuDNN backend. split_backward and split_forward are not supported by this paged operation; schedule applies to the fused backend as in chunk_kda.	None
replay_state	ReplayState | None	Optional :class:ReplayState persistent cache. Q, K, and V use BF16, gate and beta use FP32, and count uses int32. The five token caches have shape [num_slots, 64, ...]; every count must remain in [0, 63], and no replay-state tensors may alias. Together, state_cache and replay_state own the logical recurrent state: state_cache is the FP32 state after the last complete 64-token chunk, while replay_state retains the unfinished suffix. For example, after 150 tokens they hold the state at token 128 and the final 22 tokens, respectively. Callers must preserve and pass both across invocations; switching to state-only decode would discard that suffix. When supplied, calls with one token per dense batch row use replay-backed decode; other calls use streaming prefill. Both paths support torch.compile and CUDA graph capture. Persistent storage consists of the five per-slot token caches plus their counts. Replay prefill additionally allocates five prefix buffers with token capacity total_input_capacity + 63 * num_requests rounded up to a multiple of 64 and five tail buffers shaped [num_requests, 64, ...], plus routing and compute intermediates. CUDA graph users such as vLLM must budget this workspace for every captured bucket. Graph replay may change packed offsets, slot routing, initialization flags, and replay counts while tensor shapes remain fixed to that bucket.	None

Returns:

Type	Description
Tensor	The output in q.dtype. state_cache is advanced in place.

attn_gym.linear.recurrent_kda(q, k, v, gate, beta, initial_state=None, *, cu_seqlens=None, scale=None, output_final_state=False, state_indices=None, has_initial_state=None, autotune=True, impl=Impl.FUSED)

Apply recurrent KDA for decoding and inference prefill.

Parameters:

Name	Type	Description	Default
q	Tensor	Queries shaped [B, T, HK, K], scaled by scale internally. HK may divide the value head count H for multi-value attention (MVA): each block of H // HK consecutive value heads shares one query/key head, while the gate, beta, and state stay per value head.	required
k	Tensor	Keys shaped like q.	required
v	Tensor	Values shaped [B, T, H, V].	required
gate	Tensor	Finite, nonpositive per-token natural-log decay shaped [B, T, H, K], one decay vector per value head. At each token the previous state is multiplied channelwise by exp(gate). Use the same non-cumulative representation for chunked and recurrent execution; recurrent execution has no chunk-rebase lower limit.	required
beta	Tensor	Per-token write gate shaped [B, T, H].	required
initial_state	Tensor | None	Starting recurrent state, with one [H, V, K] entry per logical sequence.	None
cu_seqlens	Tensor | None	Packed offsets shaped [N + 1] for batch-one inputs, as contiguous int32 on q.device; they start at zero, never decrease, may repeat for empty sequences whose states pass through, and may end before T.	None
scale	float | None	Query scale. Defaults to 1 / sqrt(K).	None
output_final_state	bool	Return the final recurrent state with the output. Rejected together with state_indices, which advances the pool in place instead.	False
state_indices	Tensor | None	Contiguous int32 slot indices, one per logical sequence, selecting rows of a paged initial_state pool shaped [num_slots, H, V, K]. Each sequence reads and advances initial_state[state_indices[i]] in place. An index not in [1, num_slots) implies padding and are ignored by the kernel. The active indices also have to be unique to prevent two sequences from writing the same slot concurrently. Example: num_slots = 6 state_indices = [3, -1, 5, 0] seq 0 reads and updates initial_state[3] seq 1 is padding (index is -1) so this is ignored seq 2 reads and updates initial_state[5] seq 3 is padding bc 0	None
has_initial_state	Tensor | None	Optional contiguous boolean mask, one per logical sequence. False entries ignore stale contents in the selected slot and start from zero.	None
autotune	bool	Benchmark candidate value-tile sizes for non-paged execution when true; winners are cached and reused. Paged execution and false use a deterministic sequence-length heuristic.	True
impl	Impl | str	"fused" uses the inference-only optimized scan; "reference" uses differentiable eager PyTorch in FP32, with no automatic fallback.	FUSED

Returns:

Type	Description
Tensor	The output in q.dtype and either an FP32 final state with one entry
Tensor | None	per logical sequence or None.

Serving limitations: without state_indices state rows map directly to logical sequences and final states are written out of place, decode preprocessing and scan are separate launches, and speculative-decoding rollback is unsupported.

attn_gym.linear.recurrent_kda_decode(packed_qkv, raw_gate, raw_beta, A_log, dt_bias, state_cache, state_indices, *, has_initial_state=None, gate_transform='bounded', lower_bound=-5.0, scale=None, out=None)

Run one-token paged KDA decode with preprocessing fused into the recurrence.

Parameters:

Name	Type	Description	Default
packed_qkv	Tensor	Post-convolution QKV shaped [B, H * (2 * K + V)]. Each token stores [Q for all heads | K for all heads | V for all heads]; within each section, head rows are contiguous.	required
raw_gate	Tensor	Unactivated gate shaped [1, B, H, K].	required
raw_beta	Tensor	Unactivated write gate shaped [1, B, H].	required
A_log	Tensor	FP32 per-head log decay parameter shaped [H].	required
dt_bias	Tensor	FP32 per-head/channel gate bias shaped [H, K].	required
state_cache	Tensor	FP32 or BF16 paged state pool shaped [num_slots, H, V, K]. Slots may have padding between them but each [H, V, K] row must be dense. Recurrence math remains FP32 and the updated state is cast to the pool dtype. K must be at most 256. Paged chunk and recurrent prefill use the same [H, V, K] slot layout, so the cache can transition directly from prefill to decode without a layout conversion.	required
state_indices	Tensor	Contiguous int32 slot indices shaped [B]. Non-positive indices are padding/null entries: they produce zero output and leave the cache untouched. Each positive index must be in [1, num_slots) and unique among active rows because duplicate in-place updates race. These value constraints are caller responsibilities and are not host-validated.	required
has_initial_state	Tensor | None	Optional contiguous boolean mask, one per sequence. False entries mark freshly assigned slots whose contents are garbage: the step starts from the zero state and overwrites the slot.	None
gate_transform	Literal['bounded', 'softplus']	Pointwise gate transform. "bounded" computes lower_bound * sigmoid(exp(A_log) * (raw_gate + dt_bias)); "softplus" computes -exp(A_log) * softplus(raw_gate + dt_bias).	'bounded'
lower_bound	float	Finite nonpositive bound used only by the "bounded" transform.	-5.0
scale	float | None	Query scale. Defaults to 1 / sqrt(K).	None
out	Tensor | None	Optional caller-owned contiguous output buffer shaped [1, B, H, V] in packed_qkv.dtype on the same device. When supplied, the kernel writes into and returns this exact tensor. It must not alias any input.	None

Returns:

Type	Description
Tensor	Decode output shaped [1, B, H, V] in packed_qkv.dtype. This is
Tensor	out itself when a buffer is supplied.

Q/K L2 normalization, gate activation, beta sigmoid, the recurrent update, and the output projection from recurrent state are performed in one Triton kernel. The operation is inference-only and advances state_cache in place.

attn_gym.linear.Impl

Bases: str, Enum

Select a fused or reference implementation without automatic fallback.

Context Parallelism

Every delta-rule token step is affine in the V-major recurrent state, so any token range collapses to one FP32 map H_out = H_in @ A + B, packed as [HV, V + K, K] = [bias; transition] (one map per value head; GQA key heads are expanded by the factor kernels). Context parallelism is therefore a prefix scan over these summaries, and the same machinery serves KDA and GDN (whose scalar per-head gate broadcasts onto the per-channel summary kernels). The public surface has three tiers so that new partitionings or communication topologies never add options to an op.

Terminology and index spaces

Anything a plan is built from is global; anything that touches a tensor on a rank is local. attn_gym.linear.context_parallel carries the canonical NOTE [Terminology].

term	meaning	space
global stream	the whole packed token stream, cu_seqlens_global	global
sequence	one document in the global stream	global
fragment	one contiguous global token range a rank owns; may cover several sequences. A rank owns a list of them: one for a contiguous shard, two for zig-zag	global
subsequence	the tokens of one fragment that belong to one sequence; a fragment covering several sequences has several subsequences, and each becomes one cu_seqlens segment of the owner's span	global
span	the concatenation of a rank's fragments, in the order it listed them: the packed q/k/v/output tensors on that rank, with cu_seqlens marking its subsequence boundaries	local
chunk	the attn_gym.linear.kda.stages.CHUNK_SIZE (64) token block the fused kernels work in (WY factors and one state step per block); every subsequence is chunked on its own, starting at its first token, so its last chunk may be partial and no fragment cut needs to be chunk-aligned	local
summary	the [bias; transition] map of a token range of the span	local range
slot	gathered[cp_rank][f]: one per fragment. Forward, the summary of fragment f's last subsequence; backward, the reverse map of its first; the identity when nobody needs it	(cp_rank, fragment)
predecessors / successors	the slots holding a fragment's sequence's earlier / later pieces, in token order	(cp_rank, fragment)
terminal	subsequences that end their sequence; true final states live there	local index

Fragment cut points are arbitrary global tokens: chunking restarts at every subsequence, so they need no alignment to sequences or chunks. With cu_seqlens_global = (0, 40, 232, 384) and the zig-zag table fragments_global = [[(0, 96), (288, 384)], [(96, 192), (192, 288)]]:

global tokens   0     40             96       192   232      288      384
sequences       |- s0 -|------------ s1 --------------|-------- s2 --------|
rank 0           [====== frag A ======]                        [= frag B =]
rank 1                                 [ frag C ][= frag D ==]

rank 0 span:  A∩s0 | A∩s1 | B∩s2     local cu_seqlens (0, 40, 96, 192)
rank 1 span:  C∩s1 | D∩s1 | D∩s2     local cu_seqlens (0, 96, 136, 192)
slot order:   s1 = A∩s1 -> C∩s1 -> D∩s1        s2 = D∩s2 -> B∩s2

state_summaries(bounds) takes an int32 [R, 2] device tensor of local span ranges and is exact over whole chunks of one subsequence: start = sub_start + 64·i, stop = sub_start + 64·j or the subsequence end, never crossing a cu_seqlens boundary; start == stop is the identity. The recipe always passes whole subsequences, (cu_seqlens[i], cu_seqlens[i + 1]).

Local computation: prepare, summarize, run

Primitives (attn_gym.linear.kda.chunk_kda_prepare / chunk_kda_prepare_backward, and attn_gym.linear.gdn.chunk_gdn_prepare / chunk_gdn_prepare_backward) split the fused core around the communication point without exposing its WY factors:

from attn_gym.linear.kda import chunk_kda_prepare, chunk_kda_prepare_backward

prepared = chunk_kda_prepare(q, k, v, gate, beta, cu_seqlens=cu_seqlens)
summaries = prepared.state_summaries(bounds)  # [bias; transition] per range
# ...exchange summaries and compose each subsequence's entry state...
output, final_state = prepared.run(initial_state, output_final_state=True)

grads = chunk_kda_prepare_backward(
    prepared.saved, d_output, initial_state, scale=prepared.scale, schedule=prepared.schedule
)
grad_summaries = grads.state_grad_summaries(bounds)  # reverse maps [bias; transition]
# ...exchange and compose each subsequence's exit cotangent...
dq, dk, dv, dgate, dbeta, _ = grads.run(d_final_state)

Each bounds row names one local cu_seqlens segment (or a whole-chunk run inside one), and one call returns every row's map; the context-parallel routing supplies the rows. Both the native SM100 summaries and the portable Triton summaries split long ranges into work items when their static budget permits, then compose the partial maps on the device. The budget depends on the span length, head count, backend tile size, and SM count; budget 1 scans whole ranges in one kernel without a planner or compose launch. Work planning reads only device values, preserving CUDA Graph replay across different bounds values of the same shape. attn_gym.linear.context_parallel holds the pure-PyTorch algebra on summaries (merge_state, compose_summaries, neutral_summary).

Ownership: from global fragments to local routing

Ownership plans (attn_gym.linear.context_parallel.ContextParallelPlan.from_fragments) take every rank's fragments. You choose them in plain Python; the plan cuts each fragment at sequence boundaries into Subsequences, one span cu_seqlens segment each, and derives the routing: which fragments need a forward or reverse summary (one gather slot per fragment), which slots to fold for each entry state or exit cotangent, and which subsequences end their sequence and therefore hold true final states. Two pieces of one document on the same rank are simply two span segments whose entry states the routing supplies, so contiguous shards, zig-zag load balancing, and document-aligned partitions differ only in the fragment lists. The fragments must tile [0, total_tokens) exactly once. plan.routing(device) materializes the same routing as fixed-shape device tensors; summary_slots / grad_summary_slots fill a rank's [slots, ...] buffer from a prepared handle and a routing, and compose_entry_states / compose_exit_cotangents fold the gathered [world, slots, ...] buffers into each subsequence's entry state or exit cotangent; the collective in between is the caller's. examples/linear/delta_rule_context_parallel.py builds contiguous and zig-zag fragment lists in a dozen lines.

Communication: the all-gather recipe

Reference recipe (attn_gym.linear.context_parallel.context_parallel_chunk, bound as attn_gym.linear.kda.context_parallel_kda and attn_gym.linear.gdn.context_parallel_gdn) moves summaries with one padded all-gather per direction and owns the autograd function; it is generic over the delta-rule variant through a StagedOp pair of the two prepare entry points, but not over the communication topology. context_parallel_conv_history builds the short convolution's initial_state from the same routing (plan.routing(device, conv_history=W - 1)). For a point-to-point pipeline, a recursive-doubling scan over compose_summaries, DTensor, or communication overlap, compose the staged primitives and the routing helpers (summary_slots, compose_entry_states, ...) around your own collective; the primitives and plans are unchanged. Sharding does not cost accuracy: against an FP32 reference, sharded gradients match the unsharded fused op's error to within noise for both ops. When a rank's span is one whole subsequence (a routing built with max_subsequences=1, the usual case for one long document over contiguous shards), the recipe passes cu_seqlens=None to the stages, so a chunk-aligned span runs the same dense kernels as an unpacked chunk_kda call; the choice follows the routing's shape, so it is fixed per captured graph.

from attn_gym.linear.context_parallel import ContextParallelPlan
from attn_gym.linear.kda import context_parallel_kda

# Two ranks, one fragment each: the global stream split down the middle.
total_tokens = cu_seqlens_global[-1]
half = total_tokens // 2
plan = ContextParallelPlan.from_fragments(
    cu_seqlens_global, [[(0, half)], [(half, total_tokens)]], cp_rank=dist.get_rank(group)
)
token_ids = plan.global_token_ids(device)  # gather global tensors into this rank's span
routing = plan.routing(device)  # the layout as fixed-shape device tensors
output, final_state = context_parallel_kda(q, k, v, gate, beta, routing=routing, group=group)
true_final_states = final_state[routing.terminal]  # boolean indexing syncs; see below

CUDA Graph replay: fixed storage, changing layout

The recipe starts each sequence from zero, always returns every subsequence's exit state, and does not accept initial_state or output_final_state=False. Every index it uses comes from the routing tensors, so a CUDA Graph captured around it replays for any document layout and fragment table within the routing caps: capture with plan.routing(device, slots=S, max_subsequences=N) (at most S fragments per rank, N segments per span, the same span length), then copy each new layout's routing tensors into the captured ones (test/test_context_parallel_distributed.py checks three layout/table pairs). Eager callers pass no caps.

The host plan is not rerun by graph replay. Before each replay, build the new plan outside capture, arrange the new inputs in its span order, and copy its routing values into the existing device tensors. Do not replace those tensors: the captured kernels use their original addresses. The process group, per-rank span length, routing caps, and convolution history width remain fixed.

The recipe returns N exit-state rows; only those where routing.terminal is set are true final states. Boolean indexing produces a variable-size result and synchronizes, so keep all N rows inside capture and use a fixed-shape device mask or a per-row cotangent buffer. The replay test updates that cotangent buffer before each replay, setting nonterminal rows to zero. Host-side plan.terminal indices are useful outside capture, or when terminal rows never change; capturing an index selection from them does not make it follow later layouts. Update routing only between complete forward/backward steps, not while backward still needs the previous layout.

Document-aligned fragments: CP-invariant bits

If every fragment boundary is a document boundary, no rank needs another rank's state: each document runs from the zero state to its end on the rank that owns it, exactly as the unsharded op runs it. The recipe above then composes nothing (every entry state is zero), and its output, gradients, and final states are bitwise identical for any CP degree and any document-aligned fragment table, including CP=1; for KDA they are also bitwise the public chunk_kda. Nothing else is required: no flag, no tiling, no change to the kernels. document_aligned_fragments in examples/linear/delta_rule_context_parallel.py (--partition documents) assigns whole documents to ranks, longest first onto the least-loaded rank, so ranks stay within one document's length of each other. What it cannot do is split a document that does not fit on one rank; that is what the summaries above are for.

cuDNN local execution

kernel_options={"backend": "cudnn"} runs each rank's local passes with cuDNN's kernels. The recipe is the same as for the fused backend: summarize every local range as an affine map of its entry state, exchange and compose the maps, then run the local pass from the composed state.

Forward. prepare stores only the inputs; cuDNN keeps its WY factors on chip and never writes them out. state_summaries produces each range's [B; A] map (S_exit = S_entry @ A + B) by running cuDNN's state-only pass over the range twice: from a zero entry state, whose final state is B, and from an identity entry state with the value term disabled, whose final state is A. Which ranges are probed and how their rows are ordered is decided on the device, so the call is CUDA Graph replayable and an empty range yields the identity map. Every range must be a subsequence (one rank's piece of a document, one segment of the local cu_seqlens), which is all the CP recipes ever request; a nonempty range that is not one is filled with NaN. run is cuDNN's forward from the composed entry state.

Backward. run is cuDNN's own stateful backward (kda_chunk_cudnn_packed_bwd_with_state): the state pass from the saved entry state writes checkpoints, then the BT16 backward consumes them together with the exit cotangent, returning the token gradients and the entry-state cotangent. state_grad_summaries produces each subsequence's [C; R] map with the same kernels: C is the entry-state cotangent of cuDNN's backward run with a zero exit cotangent, and R is the forward transition transposed. No fused factors are computed anywhere on the cuDNN path.

Numerics. A document cut on 16-token boundaries whose fragments exchange cuDNN's own state and cotangent reproduces the unsharded cuDNN gradients bit for bit. The composed entry states do not: the maps carry cuDNN's 16-token-chunk rounding, so cuDNN under CP is its own numerical baseline, neither the fused CP baseline nor unsharded cuDNN. The staged handles and this recipe are eager-only; torch.compile is not supported through them.

attn_gym.linear.context_parallel.context_parallel_chunk(stages, q, k, v, gate, beta, *, routing, group)

Run a staged delta-rule op over this rank's span with state exchanged by all-gather.

Parameters:

Name	Type	Description	Default
stages	StagedOp	The op's entry points; attn_gym.linear.kda.context_parallel_kda and attn_gym.linear.gdn.context_parallel_gdn bind them with the op's options.	required
q	Tensor	This rank's span of queries, laid out as plan.subsequences; k, v, gate, and beta follow the same layout and the op's chunk_* contract.	required
routing	ContextParallelRouting	This layout's tensors from ContextParallelPlan.routing.	required
group	ProcessGroup	Process group containing exactly the plan's ranks, in order.	required

Returns:

Type	Description
Tensor	The span's output and one FP32 [max_subsequences, HV, V, K] exit state per span
Tensor	segment. Only the entries where routing.terminal is set are the sequences' true final
tuple[Tensor, Tensor]	states; the others are intermediate states that the owner of the next subsequence
tuple[Tensor, Tensor]	continued from, or unused rows of padded segments.

initial_state is not accepted: every sequence starts from zero and the plan supplies the entry states. Every rank in group must call this function the same number of times. A CUDA Graph captured around a call replays for any layout whose routing has the same shapes: copy the new layout's tensors into the captured ones.

attn_gym.linear.kda.context_parallel.context_parallel_kda(q, k, v, gate, beta, *, routing, group, scale=None, autotune=True, fastmath=True, kernel_options=None)

Run chunk_kda over this rank's span with state exchanged by all-gather.

See attn_gym.linear.context_parallel.context_parallel_chunk for the argument contract; scale, autotune, fastmath, and kernel_options follow chunk_kda. Forward, backward, and their summaries share the same math policy. kernel_options={"backend": "cudnn"} selects cuDNN for the local pass and summaries and requires fastmath=True.

attn_gym.linear.gdn.context_parallel.context_parallel_gdn(q, k, v, gate, beta, *, routing, group, scale=None)

Run fused chunk_gdn over this rank's span with state exchanged by all-gather.

See attn_gym.linear.context_parallel.context_parallel_chunk for the argument contract; scale follows chunk_gdn.

attn_gym.linear.context_parallel.ContextParallelPlan dataclass

Host ownership and routing metadata for one rank, derived from every rank's fragments.

Attributes:

Name	Type	Description
table	tuple[tuple[Fragment, ...], ...]	table[cp_rank][slot] is that rank's fragment behind gather slot slot.
cp_rank	int	This rank's row in table: its rank within the context-parallel group.
predecessors	tuple[tuple[Slot, ...], ...]	Per local fragment, the slots of the same sequence's pieces before its first subsequence, in increasing token order.
successors	tuple[tuple[Slot, ...], ...]	Per local fragment, the slots of the same sequence's pieces after its last subsequence, in decreasing token order, i.e. the order a reverse fold visits them.
terminal	tuple[int, ...]	Local indices of subsequences that end their sequence, whose exit states are the sequence's true final states.

cu_seqlens property

Span offsets, one segment per subsequence.

fragments property

This rank's fragments in span order.

slots property

Gather slots per rank: one per fragment of the widest row.

subsequences property

This rank's subsequences in span order.

from_fragments(cu_seqlens_global, fragments, cp_rank) classmethod

Build cp_rank's plan from every rank's fragments, half-open global token ranges.

fragments[cp_rank] lists that rank's fragments in span order; the plan cuts each one at sequence boundaries. Cut points need no alignment to sequences or chunks. Together the fragments must tile [0, total_tokens) exactly once.

global_token_ids(device)

Global token id of every span token: the bridge from global tensors to the span.

routing(device, *, slots=None, max_subsequences=None, conv_history=0)

Materialize this layout as device tensors; the caps make the shapes layout-independent.

By default every tensor is sized for this layout alone. Under CUDA Graph replay pass the caps the loader promises never to exceed: slots (fragments per rank; default this table's maximum) and max_subsequences (segments per span; default this layout's count). Every layout within the caps and with the same span length then yields identically shaped tensors, whatever its fragment table. conv_history is the short convolution's W - 1 (0: no halo tensors). Predecessor and successor rows are world * slots - 1 wide.

attn_gym.linear.context_parallel.Subsequence dataclass

The part of one fragment inside one sequence: global tokens [start, stop).

attn_gym.linear.context_parallel.ContextParallelRouting dataclass

One layout of a plan as fixed-shape device tensors; see :meth:ContextParallelPlan.routing.

Every shape depends only on the caps (slots, max_subsequences, conv_history) and the span length, so a CUDA Graph captured around context_parallel_chunk replays with the tensors of another layout copied in place. Predecessors and successors are stored per local fragment, since only a fragment's first subsequence has predecessors and only its last has successors; entry_sources and exit_sources map every span segment to its fragment's folded state or to zero.

Attributes:

Name	Type	Description
cu_seqlens	Tensor	int32 [max_subsequences + 1] span offsets, padded with the span length.
forward_bounds	Tensor	int32 [slots, 2] range of each local fragment's last subsequence when it continues elsewhere, else empty; feeds state_summaries.
reverse_bounds	Tensor	int32 [slots, 2] range of each local fragment's first subsequence when it continues from elsewhere, else empty; feeds state_grad_summaries.
reverse_sources	Tensor	int64 [slots] local index of each fragment's first subsequence, whose exit-state cotangent is folded into the reverse slot's bias.
predecessors	Tensor	int64 [slots, world * slots - 1] flat cp_rank * slots + fragment indices of the slots before each fragment in its sequence, in token order, padded with world * slots (the identity).
successors	Tensor	Same for each fragment's successors, in reverse token order.
entry_sources	Tensor	int64 [max_subsequences] local fragment whose folded predecessors are a segment's entry state, or slots for a segment that starts its sequence.
exit_sources	Tensor	Same for exit cotangents: the fragment whose folded successors apply, or slots for a segment that ends its sequence (or is padding).
terminal	Tensor	bool [max_subsequences] marks subsequences that end their sequence.
tail_sources	Tensor	int64 [slots, conv_history] span token behind each position of the fragment's tail row, or tokens (a zero row) where the tail is shorter or the slot is unused; feeds context_parallel_conv_history.
conv_sources	Tensor	int64 [max_subsequences, conv_history] flat (rank * slots + fragment) * conv_history + position index into the gathered tails for each token preceding a segment, or world * slots * conv_history (zero row).

from_fragments(cu_seqlens_global, fragments, cp_rank, device, *, slots=None, max_subsequences=None, conv_history=0) classmethod

Build device routing directly from global sequence offsets and rank fragments.

validate(span, group)

Reject a routing built for another group, rank, span length, or device (host-only).

attn_gym.linear.kda.stages.chunk_kda_prepare(q, k, v, gate, beta, *, cu_seqlens=None, scale=None, autotune=True, fastmath=True, kernel_options=None)

Run the factor half of chunk_kda and return a handle for summaries and output.

Arguments follow chunk_kda with two restrictions: q/k/v must already share a float16 or bfloat16 dtype (no silent cast, because the caller owns autograd), and the batch dimension must be one so token offsets index one packed span (NOTE [Terminology] in attn_gym.linear.context_parallel). fastmath follows chunk_kda and applies to preparation, summaries, and run. Pass the same value to chunk_kda_prepare_backward. Native cuDNN requires True. kernel_options={"backend": "cudnn"} runs the local pass and both summaries with cuDNN's kernels. Split schedules are not available with entry states.

attn_gym.linear.kda.stages.chunk_kda_prepare_backward(saved, d_output, initial_state, *, scale, autotune=True, fastmath=True, schedule=ScheduleRequest.AUTO)

Recompute the local backward tensors before any reverse-summary exchange.

initial_state is the entry state the forward run consumed and scale is the forward handle's resolved prepared.scale; there is no default because a silently re-derived scale would corrupt every gradient. Both live outside saved because the caller's autograd function owns them. Use the forward's fastmath and prepared.schedule for matching factors and launch geometry. A cuDNN tape (:class:ChunkKDACudnnSaved) returns :class:ChunkKDACudnnBackward and requires fastmath=True.

attn_gym.linear.gdn.stages.chunk_gdn_prepare(q, k, v, gate, beta, *, cu_seqlens=None, scale=None)

Run the factor half of fused chunk_gdn and return a handle for summaries and output.

Arguments follow chunk_gdn with the batch dimension fixed to one so token offsets index one packed stream.

attn_gym.linear.gdn.stages.chunk_gdn_prepare_backward(saved, d_output, initial_state, *, scale)

Recompute the local backward tensors before any reverse-summary exchange.

initial_state is the entry state the forward run consumed and scale is the forward handle's resolved prepared.scale; see chunk_kda_prepare_backward.