KDA

For framework integrations that already target FlashKDA Python APIs, prefer the flash_kda wrapper package first. Use the MATE API below when wrapper coverage is not enough or when direct operator-level control is required.

KDA (Kimi Delta Attention) covers MATE’s direct operator APIs for chunked KDA and KDA decode on MUSA. Use this page when the flash_kda wrapper is not enough, or when you need direct control over mate.kda entrypoints such as chunk_kda and gated_delta_rule_decode.

chunk_kda is MATE’s fused chunked KDA operator on MUSA, while gated_delta_rule_decode provides the direct KDA decode API.

Minimal KDA example:

import torch
from mate.kda import chunk_kda

q = torch.randn((1, 4, 2, 128), device="musa", dtype=torch.bfloat16)
k = torch.randn((1, 4, 2, 128), device="musa", dtype=torch.bfloat16)
v = torch.randn((1, 4, 4, 128), device="musa", dtype=torch.bfloat16)
g = torch.randn((1, 4, 4, 128), device="musa", dtype=torch.bfloat16)
beta = torch.randn((1, 4, 4), device="musa", dtype=torch.bfloat16)

output = chunk_kda(q=q, k=k, v=v, g=g, beta=beta)

Chunk KDA at a glance

  • Public API: mate.chunk_kda / mate.kda.chunk_kda

  • Device: MUSA

  • Input dtypes: torch.float16 and torch.bfloat16

  • Head dimension: currently fixed to 128

  • Sequence modes: - Dense: [B, T, H, 128] - Varlen: [S, H, 128] or [1, S, H, 128] with cu_seqlens

  • Optional recurrent state input/output

  • Optional preallocated output and final_state

Chunk KDA toolchain requirements

  • The repository-wide install baseline still applies, but build the current fused chunk KDA path with MUSA SDK / MTCC 5.1.0 or newer.

  • The 4.3.6 toolchain may fail to compile KDA kernels.

Chunk KDA shape contract

For dense mode:

  • q / k: [B, T, Hqk, 128]

  • v / g: [B, T, Hv, 128]

  • beta: [B, T, Hv]

For varlen mode:

  • q / k: [S, Hqk, 128] or [1, S, Hqk, 128]

  • v / g: [S, Hv, 128] or [1, S, Hv, 128]

  • beta: [S, Hv] or [1, S, Hv]

  • cu_seqlens: cumulative sequence lengths with shape [num_seqs + 1]

Additional constraints:

  • k.shape == q.shape

  • g.shape == v.shape

  • beta.shape == v.shape[:3]

  • Hv must be divisible by Hqk

  • the last dimension of every tensor passed to the kernel must be contiguous

Chunk KDA state tensors

  • initial_state is optional

  • final_state is optional unless output_final_state=True

  • state shape: [num_seqs, Hv, 128, 128]

  • state dtype: same as value dtype or torch.float32

  • if both initial_state and final_state are provided, their dtypes must match

Chunk KDA gate parameters

When gate parameters are enabled:

  • A_log and dt_bias must be provided together

  • A_log shape: [Hv]

  • dt_bias shape: [Hv, 128]

  • lower_bound defaults to -5.0

Chunk KDA outputs

  • default return: output

  • if output_final_state=True: returns (output, final_state)

Decode support

The decode entry point is mate.kda.gated_delta_rule_decode. It remains module-scoped because mate.gated_delta_rule_decode is the separate GDN API.

Capability

Supported KDA decode scope

Device

MUSA through the TileLang backend

Sequence layout

Dense [B, T, H, D] or varlen [1, total_tokens, H, D] with cu_seqlens

Q/K/V dtypes

Matching torch.float16 or torch.bfloat16 tensors

Head dimensions

K == V and the dimension must be divisible by 32

Head grouping

The number of value/state heads must be divisible by the number of Q/K heads

State

Optional torch.float32 or torch.bfloat16 state in V-first [pool, HV, V, K] or K-first [pool, HV, K, V] layout

Output dtype

torch.float16, torch.bfloat16, or torch.float32

Serving modes

Fixed or variable-length decode, continuous batching through state_indices, and speculative decode through num_accepted_tokens

State output

Separate final-state output or in-place state update

Decode inputs use a with shape [B, T, HV, K] and b with shape [B, T, HV] or [B, T, HV, V]. A_log is optional float32 metadata with shape [HV]; dt_bias is optional float32 or bfloat16 data with shape [HV, K].

API reference

mate.kda.chunk_kda(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, g: torch.Tensor, beta: torch.Tensor, scale: float | None = None, initial_state: torch.Tensor | None = None, output_final_state: bool = False, cu_seqlens: torch.Tensor | None = None, A_log: torch.Tensor | None = None, dt_bias: torch.Tensor | None = None, lower_bound: float = -5.0, use_qk_l2norm_in_kernel: bool = True, output: torch.Tensor | None = None, final_state: torch.Tensor | None = None) → torch.Tensor | tuple[torch.Tensor, torch.Tensor][source]

Run the fused chunk KDA kernel.

Parameters:

  • q – Query tensor with shape [B, T, Hqk, 128] for dense mode or [S, Hqk, 128] / [1, S, Hqk, 128] for varlen mode.

  • k – Key tensor with the same shape and dtype as q.

  • v – Value tensor with shape [B, T, Hv, 128] or varlen equivalent.

  • g – Gate input tensor with the same shape as v.

  • beta – Beta logits tensor with shape [B, T, Hv] or varlen equivalent.

  • scale – Optional QK scaling factor. Defaults to 128**-0.5.

  • initial_state – Optional recurrent state tensor.

  • output_final_state – Whether to return the final recurrent state.

  • cu_seqlens – Optional cumulative sequence lengths for varlen mode.

  • A_log – Optional per-head gate parameter tensor.

  • dt_bias – Optional per-head, per-channel gate bias tensor.

  • lower_bound – Gate lower bound used when gate parameters A_log and dt_bias are enabled. Defaults to -5.0.

  • use_qk_l2norm_in_kernel – Whether to normalize Q/K in the kernel.

  • output – Optional preallocated output tensor.

  • final_state – Optional preallocated final-state tensor.

mate.kda.gated_delta_rule_decode(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, state: torch.Tensor | None, A_log: torch.Tensor | None, a: torch.Tensor, dt_bias: torch.Tensor | None, num_accepted_tokens: torch.Tensor | None, cu_seqlens: torch.Tensor | None, b: torch.Tensor, state_indices: torch.Tensor | None = None, scale: float | None = None, output: torch.Tensor | None = None, use_qk_l2norm: bool | None = None, lower_bound: float | None = None, use_gate_in_kernel: bool | None = None, use_lower_bound: bool | None = None, apply_beta_sigmoid: bool | None = None, allow_neg_eigval: bool | None = None, use_initial_state: bool | None = None, state_v_first: bool | None = None, is_varlen: bool | None = None, store_final_state: bool = False, inplace_final_state: bool = False, state_dtype: torch.dtype | None = None) → tuple[torch.Tensor, torch.Tensor | None][source]

Unified Gated Delta Rule Decode API.

Parameters:

  • q (torch.Tensor) – Query of shape [B, T, H, K] or [1, total_tokens, H, K].

  • k (torch.Tensor) – Key of shape [B, T, H, K] or [1, total_tokens, H, K]. Must match q in shape and dtype.

  • v (torch.Tensor) – Value of shape [B, T, HV, V] or [1, total_tokens, HV, K]. Must match q in batch/time dims.

  • state (Optional[torch.Tensor]) – State buffer in VK layout. Shape is [B_or_pool, HV, V, K] when state_v_first is True, otherwise [B_or_pool, HV, K, V]. Supported dtypes are torch.float32 and torch.bfloat16. It may be None when initial-state reads and in-place writes are disabled. The returned final state uses the same dtype. BF16 state is widened for arithmetic and rounded back to BF16 after each decoded token, matching repeated single-token decode.

  • A_log (Optional[torch.Tensor]) – Optional log decay of shape [HV]. When None, the backend uses zeros.

  • a (torch.Tensor) – Input-dependent decay of shape [B, T, HV, K].

  • dt_bias (Optional[torch.Tensor]) – Optional decay bias of shape [HV, K]. When None, the backend runs with has_dt_bias=False and supplies a dummy tensor internally. Supported dtypes are torch.float32 and torch.bfloat16.

  • num_accepted_tokens (Optional[torch.Tensor]) – Optional speculative-decoding accepted-token counts of shape [B].

  • cu_seqlens (Optional[torch.Tensor]) – Optional cumulative sequence lengths for varlen mode.

  • b (torch.Tensor) – Update gate of shape [B, T, HV] or [B, T, HV, V].

  • state_indices (Optional[torch.Tensor]) – Optional [B] or [B, T] int32/int64 mapping batch entries and decode steps to a state pool. Every non-negative index must be smaller than the state pool size.

  • scale (Optional[float]) – Query scale. If None, defaults to 1 / sqrt(K).

  • output (Optional[torch.Tensor]) – Optional pre-allocated output tensor of shape [B, T, HV, V].

  • use_qk_l2norm (Optional[bool]) – Whether to L2-normalize q and k in-kernel. None uses the current default.

  • lower_bound (Optional[float]) – Optional backend lower-bound value. None uses the current default.

  • use_gate_in_kernel (Optional[bool]) – Optional explicit control of in-kernel gate application.

  • use_lower_bound (Optional[bool]) – Optional explicit control of backend lower-bound logic.

  • apply_beta_sigmoid (Optional[bool]) – Optional explicit control of in-kernel sigmoid application on b.

  • allow_neg_eigval (Optional[bool]) – Optional explicit control of backend negative-eigenvalue handling.

  • use_initial_state (Optional[bool]) – Optional explicit control of whether the backend reads the provided state.

  • state_v_first (Optional[bool]) – State matrix layout selector. None defaults to True.

  • is_varlen (Optional[bool]) – Whether to run the varlen path. None defaults to whether cu_seqlens is provided.

  • store_final_state (bool) – Whether to request final-state output from the backend.

  • inplace_final_state (bool) – Whether final-state writes should reuse state.

  • state_dtype (Optional[torch.dtype]) – State storage dtype. When omitted, it is inferred from state, or defaults to torch.float32 when state is None. Supported dtypes are torch.float32 and torch.bfloat16.

Returns:

  • output: backend output tensor

  • final_state: semantic final-state result, or None when not requested

Return type:

Tuple[torch.Tensor, Optional[torch.Tensor]]