feat(kda): fused KDA decode for 16/64 heads and the DS conv-state layout

flashinfer/kda_decode.py

@@ -521,8 +521,8 @@ def fused_kda_decode(
 
     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.
+    head dimension 128 and 8, 12, 16, 24, 32, 48, 64, 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.
@@ -536,9 +536,12 @@ def fused_kda_decode(
             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]``.
+            ``[num_slots, 3 * H * 128, 3]`` and dtype bfloat16. Each slot uses
+            either the sequence-dimension (SD) cache layout with strides
+            ``[slot_stride, 1, 3 * H * 128]`` or the dimension-first (DS)
+            layout with strides ``[slot_stride, 3, 1]``. Slots must start
+            8-byte aligned (``slot_stride`` a multiple of 4). The Cake backend
+            supports SD only.
         raw_gate:
             Raw per-channel recurrence gate with shape
             ``[1, num_rows, H, 128]`` and dtype bfloat16.
@@ -678,9 +681,10 @@ def packed_fused_kda_decode(
     T>1 fails closed. T>1 uses the SM10x CuTe DSL packed backend, which
     accepts float32 or bfloat16 recurrent state and requires a finite negative
     ``lower_bound``. Both paths require head dimension 128 and convolution
-    width four; T=1 accepts 8, 12, 24, 32, 48, or 96 heads and T>1 accepts
-    12, 24, 32, 48, or 96 heads. Other tensors follow :func:`fused_kda_decode`,
-    with the packed row count supplied by ``x``.
+    width four; T=1 accepts 8, 12, 16, 24, 32, 48, 64, or 96 heads and T>1
+    accepts 12, 16, 24, 32, 48, 64, or 96 heads. T>1 requires the SD
+    convolution-cache layout; the DS layout is T=1 only. Other tensors follow
+    :func:`fused_kda_decode`, with the packed row count supplied by ``x``.
 
     ``query_start_loc`` must be contiguous int32 of shape ``[N+1]``, start at
     zero, and contain nondecreasing offsets within ``x`` with active lengths no

flashinfer/kda_kernels/fused_kda_decode.py

@@ -56,7 +56,7 @@
 _CONV_THREADS = 3 * _HEAD_DIM // 4
 _Q_SCALE = _HEAD_DIM**-0.5
 _L2_EPS = 1.0e-6
-_SUPPORTED_HEADS = (8, 12, 24, 32, 48, 96)
+_SUPPORTED_HEADS = (8, 12, 16, 24, 32, 48, 64, 96)
 _CUTE_DSL_MODULE = "fused_kda_decode"
 _SOURCE_FILES = (str(Path(__file__).resolve()),)
 
@@ -75,6 +75,17 @@ def _aligned_tensor(tensor, alignment):
     )
 
 
+def _ds_conv_chunks(conv_state, channel_in_warp, vector_idx, head_idx, slot):
+    """View one thread's DS conv state as three contiguous bf16x4 chunks.
+
+    In DS a channel's three history taps are adjacent, so the 4 channels x 3
+    taps a thread owns are 12 contiguous values; chunk c holds flat values
+    4c..4c+3, where flat value 3 * channel + tap.
+    """
+    taps = conv_state[(None, channel_in_warp, None, vector_idx, head_idx, slot)]
+    return cute.make_tensor(taps.iterator, cute.make_layout((4, 3), stride=(1, 4)))
+
+
 def _sigmoid(value):
     """Evaluate a cancellation-safe approximate sigmoid."""
     return cute.rcp(cute.exp(-value, fastmath=True) + 1.0, approx=True, ftz=True)
@@ -101,6 +112,7 @@ def _fused_kda_decode_kernel(
     use_lower_bound: cutlass.Constexpr,
     lower_bound: cutlass.Constexpr,
     norm_eps: cutlass.Constexpr,
+    conv_dim_first: cutlass.Constexpr,
 ):
     thread_idx, _, _ = cute.arch.thread_idx()
     head_idx, row_idx, _ = cute.arch.block_idx()
@@ -195,24 +207,44 @@ def _fused_kda_decode_kernel(
                 ),
                 weight_registers[(None, width_idx)],
             )
-        for history_idx in cutlass.range_constexpr(3):
-            cute.copy(
-                copy_bf16x4,
-                _aligned_tensor(
-                    conv_state[
-                        (
-                            None,
-                            channel_in_warp,
-                            history_idx,
-                            vector_idx,
-                            head_idx,
-                            slot,
-                        )
-                    ],
-                    8,
-                ),
-                conv_registers[(None, history_idx)],
+        if cutlass.const_expr(conv_dim_first):
+            # DS layout: load the thread's 12 contiguous values as three
+            # coalesced vectors, then transpose to (channel, tap) in registers.
+            conv_chunks = _ds_conv_chunks(
+                conv_state, channel_in_warp, vector_idx, head_idx, slot
             )
+            conv_flat = cute.make_rmem_tensor((4, 3), BF16)
+            for chunk_idx in cutlass.range_constexpr(3):
+                cute.copy(
+                    copy_bf16x4,
+                    _aligned_tensor(conv_chunks[(None, chunk_idx)], 8),
+                    conv_flat[(None, chunk_idx)],
+                )
+            for element_idx in cutlass.range_constexpr(4):
+                for history_idx in cutlass.range_constexpr(3):
+                    flat_idx = 3 * element_idx + history_idx
+                    conv_registers[(element_idx, history_idx)] = conv_flat[
+                        (flat_idx % 4, flat_idx // 4)
+                    ]
+        else:
+            for history_idx in cutlass.range_constexpr(3):
+                cute.copy(
+                    copy_bf16x4,
+                    _aligned_tensor(
+                        conv_state[
+                            (
+                                None,
+                                channel_in_warp,
+                                history_idx,
+                                vector_idx,
+                                head_idx,
+                                slot,
+                            )
+                        ],
+                        8,
+                    ),
+                    conv_registers[(None, history_idx)],
+                )
 
         mixed_registers = cute.make_rmem_tensor((4,), F32)
         for qkv_idx in cutlass.range_constexpr(4):
@@ -236,24 +268,42 @@ def _fused_kda_decode_kernel(
             _aligned_tensor(mixed_warp[(None, thread_idx)], 16),
         )
         if is_live:
-            for history_idx in cutlass.range_constexpr(3):
-                cute.copy(
-                    copy_bf16x4,
-                    conv_registers[(None, history_idx)],
-                    _aligned_tensor(
-                        conv_state[
-                            (
-                                None,
-                                channel_in_warp,
-                                history_idx,
-                                vector_idx,
-                                head_idx,
-                                slot,
-                            )
-                        ],
-                        8,
-                    ),
+            if cutlass.const_expr(conv_dim_first):
+                conv_chunks = _ds_conv_chunks(
+                    conv_state, channel_in_warp, vector_idx, head_idx, slot
                 )
+                conv_flat = cute.make_rmem_tensor((4, 3), BF16)
+                for element_idx in cutlass.range_constexpr(4):
+                    for history_idx in cutlass.range_constexpr(3):
+                        flat_idx = 3 * element_idx + history_idx
+                        conv_flat[(flat_idx % 4, flat_idx // 4)] = conv_registers[
+                            (element_idx, history_idx)
+                        ]
+                for chunk_idx in cutlass.range_constexpr(3):
+                    cute.copy(
+                        copy_bf16x4,
+                        conv_flat[(None, chunk_idx)],
+                        _aligned_tensor(conv_chunks[(None, chunk_idx)], 8),
+                    )
+            else:
+                for history_idx in cutlass.range_constexpr(3):
+                    cute.copy(
+                        copy_bf16x4,
+                        conv_registers[(None, history_idx)],
+                        _aligned_tensor(
+                            conv_state[
+                                (
+                                    None,
+                                    channel_in_warp,
+                                    history_idx,
+                                    vector_idx,
+                                    head_idx,
+                                    slot,
+                                )
+                            ],
+                            8,
+                        ),
+                    )
     elif thread_idx >= _NUM_THREADS - _HEAD_DIM:
         # Four otherwise-idle warps build the per-channel recurrence decay and
         # gated RMSNorm scale once, rather than recomputing them in every warp.
@@ -443,6 +493,7 @@ def _fused_kda_decode_launch(
     use_lower_bound: cutlass.Constexpr,
     lower_bound: cutlass.Constexpr,
     norm_eps: cutlass.Constexpr,
+    conv_dim_first: cutlass.Constexpr,
 ):
     num_heads = A_log.shape[0]
     num_rows = x.shape[0]
@@ -463,18 +514,17 @@ def _fused_kda_decode_launch(
             stride=(1, 4, hidden_size, 4 * hidden_size, _HEAD_DIM),
         ),
     )
+    # (channel in vector, vector, history, q/k/v, head, slot). SD stores the
+    # three history columns as [3, 3 * H * 128] per slot, DS as [3 * H * 128, 3].
+    if cutlass.const_expr(conv_dim_first):
+        conv_strides = (3, 12, 1, 3 * hidden_size, 3 * _HEAD_DIM)
+    else:
+        conv_strides = (1, 4, qkv_size, hidden_size, _HEAD_DIM)
     conv_state_layout = cute.make_tensor(
         conv_state.iterator,
         cute.make_layout(
             (4, 32, 3, 3, num_heads, conv_state.shape[0]),
-            stride=(
-                1,
-                4,
-                qkv_size,
-                hidden_size,
-                _HEAD_DIM,
-                conv_state.stride[0],
-            ),
+            stride=(*conv_strides, conv_state.stride[0]),
         ),
     )
     gate_layout = cute.make_tensor(
@@ -543,6 +593,7 @@ def _fused_kda_decode_launch(
         use_lower_bound,
         lower_bound,
         norm_eps,
+        conv_dim_first,
     ).launch(
         grid=[num_heads, num_rows, 1],
         block=[_NUM_THREADS, 1, 1],
@@ -551,7 +602,7 @@ def _fused_kda_decode_launch(
     )
 
 
-def _make_compile_inputs(state_dtype):
+def _make_compile_inputs(state_dtype, conv_dim_first=False):
     """Build shape- and stride-dynamic inputs for CuTe compilation."""
     num_rows = cute.sym_int()
     num_heads = cute.sym_int()
@@ -574,11 +625,19 @@ def compact(shape, dtype):
         stride=(cute.sym_int64(), 1),
         assumed_align=16,
     )
+    # DS keeps a channel's three history taps contiguous; SD keeps channels
+    # contiguous within each tap.
+    if conv_dim_first:
+        conv_stride = (cute.sym_int64(), 3, 1)
+    else:
+        conv_stride = (cute.sym_int64(), 1, qkv_size)
+    # Both layouts move conv state as bf16x4 vectors, so the launcher only
+    # requires 8-byte aligned slots; do not assume more here.
     conv_state = cute.runtime.make_fake_tensor(
         BF16,
         shape=(conv_slots, qkv_size, 3),
-        stride=(cute.sym_int64(), 1, qkv_size),
-        assumed_align=16,
+        stride=conv_stride,
+        assumed_align=8,
     )
     raw_beta = cute.runtime.make_fake_tensor(
         BF16,
@@ -632,7 +691,9 @@ def compact(shape, dtype):
 
 
 @functools.cache
-def _get_compiled_kernel(state_dtype, lower_bound, norm_eps, packed_t1=False):
+def _get_compiled_kernel(
+    state_dtype, lower_bound, norm_eps, packed_t1=False, conv_dim_first=False
+):
     """Get a shape-dynamic specialization from the two-level CuTe DSL cache."""
     state_is_bf16 = state_dtype == torch.bfloat16
     compile_state_dtype = BF16 if state_is_bf16 else F32
@@ -648,18 +709,20 @@ def _get_compiled_kernel(state_dtype, lower_bound, norm_eps, packed_t1=False):
         f"d128_w4_{gate_name}_state{state_name}"
         f"_eps{str(float(norm_eps)).replace('.', '_').replace('-', 'm')}"
         f"_packed{int(packed_t1)}"
+        f"{'_ds' if conv_dim_first else ''}"
     )
     kernel = build_and_load_cute_dsl_kernel(
         _CUTE_DSL_MODULE,
         kernel_name,
         lambda: cute.compile(
             _fused_kda_decode_launch,
-            *_make_compile_inputs(compile_state_dtype),
+            *_make_compile_inputs(compile_state_dtype, conv_dim_first),
             state_is_bf16,
             packed_t1,
             use_lower_bound,
             compile_lower_bound,
             float(norm_eps),
+            conv_dim_first,
             options="--enable-tvm-ffi --generate-line-info",
         ),
         extra_key_files=_SOURCE_FILES,
@@ -738,6 +801,12 @@ def _select_cake_variant(
 
     num_rows = int(x.shape[0])
     num_heads = int(x.shape[1]) // (3 * _HEAD_DIM)
+    # Cake routes cover a subset of _SUPPORTED_HEADS (16 and 64 have none); the
+    # selector raises for other head counts, so fall back to CuTe DSL here.
+    if not any(
+        num_heads in rule.heads for variant in variants for rule in variant.eligibility
+    ):
+        return None
     state_dtype = "bfloat16" if state.dtype == torch.bfloat16 else "float32"
     return select_cake_fused_kda_decode_variant(
         target=target,
@@ -919,13 +988,19 @@ def run_fused_kda_decode(
         conv_state.ndim != 3
         or conv_state.shape[1:] != (3 * hidden_size, 3)
         or conv_state.stride(0) < 9 * hidden_size
-        or conv_state.stride(1) != 1
-        or conv_state.stride(2) != 3 * hidden_size
+        or conv_state.stride()[1:] not in ((1, 3 * hidden_size), (3, 1))
     ):
         raise ValueError(
-            "conv_state must have shape [slots, 3 * H * 128, 3] "
-            "and use the SD cache layout"
+            "conv_state must have shape [slots, 3 * H * 128, 3] and use the SD "
+            "(strides [slot, 1, 3 * H * 128]) or DS (strides [slot, 3, 1]) "
+            "cache layout"
         )
+    if conv_state.stride(0) % 4 or conv_state.data_ptr() % 8:
+        # Both layouts move conv state as bf16x4 vectors.
+        raise ValueError(
+            "conv_state slots must start 8-byte aligned (slot stride a multiple of 4)"
+        )
+    conv_dim_first = conv_state.stride(2) == 1
     if raw_gate.shape != (1, num_rows, num_heads, _HEAD_DIM):
         raise ValueError("raw_gate must have shape [1, num_rows, H, 128]")
     if not raw_gate.is_contiguous():
@@ -1018,7 +1093,7 @@ def run_fused_kda_decode(
         raise ValueError("packed T=1 metadata requires backend='cute-dsl'")
 
     cake_variant = None
-    if backend != "cute-dsl" and state_indices_mode is not None:
+    if backend != "cute-dsl" and state_indices_mode is not None and not conv_dim_first:
         cake_variant = _select_cake_variant(
             x=x,
             conv_state=conv_state,
@@ -1056,7 +1131,11 @@ def run_fused_kda_decode(
         )
 
     kernel = _get_compiled_kernel(
-        state.dtype, lower_bound, float(norm_eps), packed_t1=packed_t1
+        state.dtype,
+        lower_bound,
+        float(norm_eps),
+        packed_t1=packed_t1,
+        conv_dim_first=conv_dim_first,
     )
     args = (
         x,

flashinfer/kda_kernels/fused_kda_decode_multitoken.py

@@ -1021,8 +1021,8 @@ def run_fused_kda_decode_multitoken(
         raise ValueError("x must have a contiguous channel dimension of 3 * H * 128")
     num_heads = x.shape[1] // (3 * _D)
     hidden_size = num_heads * _D
-    if num_heads not in (12, 24, 32, 48, 96):
-        raise ValueError("H must be one of 12, 24, 32, 48, or 96")
+    if num_heads not in (12, 16, 24, 32, 48, 64, 96):
+        raise ValueError("H must be one of 12, 16, 24, 32, 48, 64, or 96")
     if weight.shape != (3, _W, hidden_size) or not weight.is_contiguous():
         raise ValueError("weight must be contiguous with shape [3, 4, H * 128]")
     if (

flashinfer/trace/templates/kda.py

@@ -277,7 +277,7 @@ def _recurrent_kda_template(
         "singleton == 1",
         "projection_groups == 3",
         "head_dim == 128",
-        "num_heads in (8, 12, 24, 32, 48, 96)",
+        "num_heads in (8, 12, 16, 24, 32, 48, 64, 96)",
         "conv_width == 4",
         "conv_history == 3",
     ],
@@ -343,7 +343,7 @@ def _recurrent_kda_template(
         "singleton == 1",
         "projection_groups == 3",
         "head_dim == 128",
-        "num_heads in (12, 24, 32, 48, 96)",
+        "num_heads in (12, 16, 24, 32, 48, 64, 96)",
         "conv_width == 4",
         "conv_history == num_tokens + 2",
     ],