For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).
Mojo struct
KdaChunk
struct KdaChunk
KDA chunk-parallel prefill forward pass.
Same algorithm and same numerics-within-tolerance contract as
kda_decode (see Kernels/lib/kda/chunk_fwd.mojo's M2c docstring and
test_kda_chunk_parallel.mojo), but restructures the recurrence so a
prefill's sequential depth is O(total_T / 16) instead of O(total_T):
L1 (kda_chunk_prepare_gpu) and L3 (kda_chunk_output_gpu) run one CTA
per (chunk, value-head) in parallel; only L2 (kda_chunk_scan_gpu)
carries a sequential dependency, and it carries it per-CHUNK rather than
per-token. Intended for prefill (multi-token sequences); kda_decode
remains the single-token decode path.
Takes the SAME tensors as kda_decode -- CHUNK_SIZE (16) and the chunk
map are internal to this op, not part of its call contract, so a caller
can swap between the two ops without restructuring its inputs.
Default parameters match kda_decode's (backward-compatible with the S1
graph op):
gate_mode = "original" (stable-softplus gate)
beta_mode = "logits" (sigmoid applied to beta_logits)
state_layout = "K_FIRST" (pool shape [N, HV, K, V])
Tensor shapes: identical to kda_decode (see that op's docstring).
Implemented traits
Methods
execute
static def execute[qkv_dtype: DType, gate_dtype: DType, state_dtype: DType, output_dtype: DType, target: StringSpan[ImmStaticOrigin], gate_mode: StringSpan[ImmStaticOrigin] = StringSpan("original"), beta_mode: StringSpan[ImmStaticOrigin] = StringSpan("logits"), state_layout: StringSpan[ImmStaticOrigin] = StringSpan("K_FIRST")](output: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=output.static_spec], q: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=q.static_spec], k: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=k.static_spec], v: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=v.static_spec], raw_gate: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=raw_gate.static_spec], beta_logits: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=beta_logits.static_spec], a_log: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=a_log.static_spec], dt_bias: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=dt_bias.static_spec], cu_seqlens: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=cu_seqlens.static_spec], state_pool: ManagedTensorSlice[IOSpec[_, _].MutableInput, static_spec=state_pool.static_spec], state_indices: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=state_indices.static_spec], ctx: DeviceContext)