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
KdaDecode
struct KdaDecode
KDA decode recurrence forward pass.
Implements the Kimi Delta Attention decode recurrence with per-key alpha decay and in-kernel gate+beta activation fusion.
Default parameters (backward-compatible with 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: - output : [1, total_T, num_value_heads, value_head_dim] (OUT) - q : [1, total_T, num_key_heads, key_head_dim] (bf16) - k : [1, total_T, num_key_heads, key_head_dim] (bf16) - v : [1, total_T, num_value_heads, value_head_dim] (bf16) - raw_gate : [1, total_T, num_value_heads, key_head_dim] (fp32) - beta_logits : [1, total_T, num_value_heads] (fp32) - a_log : [num_value_heads] (fp32) - dt_bias : [num_value_heads, key_head_dim] (fp32) - cu_seqlens : [batch_size + 1] (int32) - state_pool : [max_slots, num_value_heads, key_head_dim, value_head_dim] (MUT) - state_indices : [batch_size] (int32)
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)