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, *, backend: Literal['cute-dsl', 'cake', 'auto'] = 'cute-dsl', state_indices_mode: Literal['positive_unique', 'unique_or_null', 'repeated_positive'] | 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 8, 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].

  • backend – Implementation backend. "cute-dsl" preserves the existing FlashInfer implementation. "cake" strictly selects an exported Cake kernel and raises when no registered route matches. "auto" selects Cake only when state_indices_mode is supplied and a route matches on SM100a or SM103a, otherwise preserving CuTe DSL. Each architecture uses its own compiled module. Default: "cute-dsl".

  • state_indices_mode – Host-known assertion about state_indices used only by the Cake dispatcher. "positive_unique" means every index is positive and unique; "unique_or_null" means positive indices are unique and non-positive null rows may be present; "repeated_positive" means at least one positive slot repeats (null rows may also be present). Required for backend="cake". The dispatcher never reads the CUDA tensor to infer this property.

Returns:

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