flashinfer.kda_decode.fused_kda_decode

flashinfer.kda_decode.fused_kda_decode(x: Tensor, weight: Tensor, conv_state: Tensor, raw_gate: Tensor, raw_beta: Tensor, A_log: Tensor, dt_bias: Tensor, state_indices: Tensor, state: Tensor, output_gate: Tensor, norm_weight: Tensor, lower_bound: float | None = -5.0, norm_eps: float = 1e-05, output: Tensor | None = None) → Tensor

Run the fused Kimi KDA decode pipeline.

This operator fuses a width-four depthwise causal convolution with SiLU, one recurrent KDA update, and gated RMSNorm. It is specialized for head dimension 128 and 12, 24, 32, 48, or 96 heads. conv_state and state are updated in-place.

Slot zero is reserved as a null slot. Rows whose state_indices value is non-positive produce zeros and do not update either cache.

Parameters:

  • x – Packed QKV projection with shape [num_rows, 3 * H * 128] and dtype bfloat16. The channel dimension must be contiguous.

  • weight – Depthwise convolution weights with shape [3, 4, H * 128] and dtype float32.

  • conv_state – Paged convolution cache with shape [num_slots, 3 * H * 128, 3] and dtype bfloat16. Each slot must use the sequence-dimension cache layout with strides [slot_stride, 1, 3 * H * 128].

  • raw_gate – Raw per-channel recurrence gate with shape [1, num_rows, H, 128] and dtype bfloat16.

  • raw_beta – Raw delta-rule learning-rate logits with shape [1, num_rows, H] and dtype bfloat16.

  • A_log – Log decay parameter with H elements and dtype float32.

  • dt_bias – Per-channel decay bias with H * 128 elements and dtype float32.

  • state_indices – Cache slot selected by each decode row. Must be a contiguous int32 tensor with num_rows elements. Live indices must be in [1, num_slots); non-positive indices select the null path.

  • state – Paged recurrent state with shape [num_slots, H, 128, 128] and dtype float32 or bfloat16. The recurrence is evaluated in float32; a bfloat16 state is rounded when written back. Each slot’s [H, 128, 128] contents must be contiguous. state and conv_state must have the same num_slots.

  • output_gate – Gated RMSNorm logits with shape [num_rows, H, 128] or [1, num_rows, H, 128] and dtype bfloat16.

  • norm_weight – RMSNorm weight with 128 elements and dtype float32.

  • lower_bound – Negative lower bound used by the recurrence gate. Defaults to -5.0 for Kimi K3. Pass None to use the original Kimi-Linear softplus gate.

  • norm_eps – Non-negative RMSNorm epsilon. Defaults to 1e-5.

  • output – Optional preallocated contiguous bfloat16 output with shape [1, num_rows, H, 128].

Returns:

The bfloat16 output tensor with shape [1, num_rows, H, 128].