perf(kda): use FlashInfer fused decode on Blackwell

docs/recipes/models.md

@@ -214,8 +214,9 @@ bounded draft attention for long-context deployments.
 
 Kimi-K3 combines a MoonViT vision encoder with a hybrid KDA
 (linear-attention) / NoPE-MLA (full-attention) decoder and a
-DeepSeek-V3-style latent MoE. The KDA layers currently use
-flash-linear-attention kernels on NVIDIA, so install it first:
+DeepSeek-V3-style latent MoE. KDA prefill can use the optional
+flash-linear-attention kernels on NVIDIA, so install it when selecting that
+backend:
 
 ```bash
 pip install flash-linear-attention
@@ -226,8 +227,16 @@ Notes:
 - K3 uses the cache-group scheduler and KDA state groups.
 - KDA dispatch is vendor-neutral at the runtime boundary. The kernel registry
   selects the existing FLA-derived NVIDIA implementation or the native AMD
-  implementation, including each backend's preferred recurrent-state layout.
-  The runtime does not transpose or reinterpret that state.
+  implementation for prefill and speculative verify. On B200/B300, ordinary
+  single-token decode uses FlashInfer's fused KDA operator when its backend is
+  available; TokenSpeed bulk-copies prefix-cache COW rows before invoking the
+  public single-index API. DSpark T=8 verify remains on the FP32-state
+  ReplaySSM path. Layerwise L2 restore also retains the native per-layer path
+  so its load fences are not bypassed.
+- The persistent KDA convolution layout is sequence-major on Blackwell and
+  feature-major elsewhere. Cache-transfer P/D peers must therefore use the
+  same GPU-generation layout; mixed H100/B200-or-B300 K3 PD is rejected by the
+  cache contract.
 - NVIDIA auto-selects `--attention-backend tokenspeed_mla` for K3
   (fp8 KV required). AMD uses the `mla` backend.
 - `tokenspeed serve` auto-selects the `kimi_k3` reasoning and tool-call

python/tokenspeed/runtime/layers/attention/backends/hybrid_kda.py

@@ -30,6 +30,7 @@
 from tokenspeed_kernel.ops.activation.triton import rmsnorm_gated_sigmoid
 from tokenspeed_kernel.ops.attention import (
     kda_batched_replay_uses_raw_gate,
+    kda_conv_state_layout,
     kda_fused_paged_decode,
     kda_fused_paged_verify,
     kda_paged_decode,
@@ -48,11 +49,13 @@
 from tokenspeed_kernel.ops.attention.triton.verify_state_blocks import (
     commit_state_pages,
 )
+from tokenspeed_kernel.ops.kvcache.triton import KdaGroupedStateCopyDescriptor
 from typing_extensions import override
 
 from tokenspeed.runtime.layers.attention.backends.hybrid_linear_attn import (
     HybridLinearAttnBackend,
     MambaAttnBackend,
+    allocate_conv_state_rows_like,
     logger,
 )
 
@@ -121,6 +124,10 @@ def __init__(
         self._batched_replay_launch = None
         self._batched_replay_ready = False
         self._replay_descriptor_bound: set[int] = set()
+        self._decode_cow_descriptors: tuple[KdaGroupedStateCopyDescriptor, ...] = ()
+        self._decode_cow_first_layer: int | None = None
+        self._decode_cow_l2_notice_emitted = False
+        self._sequence_major_conv = kda_conv_state_layout() == "sequence_major"
         self.kda_backend = (kda_backend or "auto").strip().lower()
         if self.kda_backend not in KDA_PREFILL_BACKENDS:
             raise ValueError(
@@ -136,6 +143,33 @@ def __init__(
     @override
     def set_kv_pool(self, kv_pool) -> None:
         super().set_kv_pool(kv_pool)
+        self._decode_cow_descriptors = ()
+        self._decode_cow_first_layer = None
+        if self._sequence_major_conv:
+            layer_ids = tuple(self._state_layer_ids())
+            group_ids = tuple(self._state_groups())
+            # Kimi-K3 currently publishes three state groups. Keep a fail-closed
+            # fallback if a future recipe changes that contract: prepared FI
+            # decode weights will be ignored rather than running without COW.
+            if layer_ids and len(group_ids) == 3:
+                group_rows = {group_id: row for row, group_id in enumerate(group_ids)}
+                group_sel = tuple(
+                    group_rows[self._state_group_for(layer_id)]
+                    for layer_id in layer_ids
+                )
+                components = tuple(
+                    self._state_components(layer_id) for layer_id in layer_ids
+                )
+                if components and components[0][0].is_cuda:
+                    self._decode_cow_descriptors = (
+                        KdaGroupedStateCopyDescriptor.build(
+                            tuple(component[0] for component in components), group_sel
+                        ),
+                        KdaGroupedStateCopyDescriptor.build(
+                            tuple(component[1] for component in components), group_sel
+                        ),
+                    )
+                    self._decode_cow_first_layer = layer_ids[0]
         if self._replay_active and self.speculative_num_draft_tokens > 1:
             rows = self.max_bs * self.speculative_num_draft_tokens
             layer_ids = tuple(self._state_layer_ids())
@@ -244,6 +278,8 @@ def _bind_replay_descriptor(self, layer_id: int, weights: tuple) -> None:
             strides = (
                 first_qkv.stride(0),
                 first_conv.stride(0),
+                first_conv.stride(1),
+                first_conv.stride(2),
                 first_fa.stride(0),
                 first_beta.stride(0),
                 first_state.stride(0),
@@ -266,6 +302,8 @@ def _bind_replay_descriptor(self, layer_id: int, weights: tuple) -> None:
                 if (
                     layer_qkv.stride(0),
                     layer_conv.stride(0),
+                    layer_conv.stride(1),
+                    layer_conv.stride(2),
                     layer_fa.stride(0),
                     layer_beta.stride(0),
                     layer_state.stride(0),
@@ -308,10 +346,12 @@ def launch(read_indices, write_indices, accepted_length):
                         f_a_dim=geometry[2],
                         qkv_stride=strides[0],
                         conv_stride=strides[1],
-                        f_a_stride=strides[2],
-                        beta_stride=strides[3],
-                        state_stride=strides[4],
-                        gate_stride=strides[5],
+                        conv_feature_stride=strides[2],
+                        conv_history_stride=strides[3],
+                        f_a_stride=strides[4],
+                        beta_stride=strides[5],
+                        state_stride=strides[6],
+                        gate_stride=strides[7],
                         conv_width=conv_width,
                         layers_per_group=layers_per_group,
                         lower_bound=next(iter(lower_bounds)),
@@ -343,9 +383,7 @@ def _ensure_verify_scratch(self, bs: int, draft_token_num: int) -> None:
                 (
                     conv
                     if self._replay_uses_raw_gate
-                    else torch.empty(
-                        (rows, *conv.shape[1:]), dtype=conv.dtype, device=conv.device
-                    )
+                    else allocate_conv_state_rows_like(conv, rows, zero=False)
                 ),
                 None,
             )
@@ -387,6 +425,47 @@ def _kda_gate(
         else:
             return g_raw
 
+    def _stage_flashinfer_decode_cow(self, layer_id: int, batch_size: int) -> bool:
+        """Prepare single-index FlashInfer state once at the first KDA layer.
+
+        Layerwise L2 restore owns the first access to each layer's cache fields.
+        A cross-layer bulk copy cannot safely jump those fences, so such batches
+        stay on the native dual-index Triton path.
+        """
+        descriptors = self._decode_cow_descriptors
+        if not descriptors or self._decode_cow_first_layer is None:
+            return False
+        load_tracker = getattr(self.kv_pool, "layerwise_load_tracker", None)
+        # This branch is fixed while a CUDA graph is captured. Merely checking
+        # the current consumer set would bake cross-layer copies into a graph
+        # that a later L2-restored batch replays. Disable this backend whenever
+        # the pool participates in layerwise restore, preserving its fences.
+        if load_tracker is not None:
+            if not self._decode_cow_l2_notice_emitted:
+                logger.info(
+                    "FlashInfer KDA decode uses the Triton fallback while "
+                    "layerwise cache restore is active"
+                )
+                self._decode_cow_l2_notice_emitted = True
+            return False
+        metadata = self.forward_metadata
+        if (
+            metadata.state_in_blocks_by_group is None
+            or metadata.state_out_blocks_by_group is None
+        ):
+            return False
+        if layer_id == self._decode_cow_first_layer:
+            group_ids = self._state_groups()
+            read_groups = tuple(
+                metadata.state_in_blocks_by_group[group_id] for group_id in group_ids
+            )
+            write_groups = tuple(
+                metadata.state_out_blocks_by_group[group_id] for group_id in group_ids
+            )
+            for descriptor in descriptors:
+                descriptor.copy(read_groups, write_groups, batch_size=batch_size)
+        return True
+
     @override
     def _decode(
         self,
@@ -409,13 +488,23 @@ def _decode(
         output_gate: torch.Tensor | None,
         norm_weight: torch.Tensor | None,
         norm_eps: float | None,
+        prepared_weights: object | None = None,
+        layer_id: int | None = None,
     ) -> torch.Tensor | None:
         if output_gate is not None and (norm_weight is None or norm_eps is None):
             raise ValueError(
                 "norm_weight and norm_eps are required with a KDA output gate"
             )
         if f_a_out is None:
             return None
+        use_prepared_decode = prepared_weights is not None
+        if use_prepared_decode and (
+            layer_id is None
+            or not self._stage_flashinfer_decode_cow(layer_id, mixed_qkv.shape[0])
+        ):
+            use_prepared_decode = False
+        active_prepared_weights = prepared_weights if use_prepared_decode else None
+        dispatch_read_indices = write_indices if use_prepared_decode else read_indices
 
         num_value_heads = value_dim // attn_tp_size // head_v_dim
         result = kda_fused_paged_decode(
@@ -428,7 +517,7 @@ def _decode(
             A_log,
             dt_bias,
             state_pool=ssm_states,
-            read_indices=read_indices,
+            read_indices=dispatch_read_indices,
             write_indices=write_indices,
             num_heads=num_value_heads,
             head_dim=head_v_dim,
@@ -437,6 +526,7 @@ def _decode(
             output_gate=output_gate,
             norm_weight=norm_weight,
             norm_eps=norm_eps,
+            prepared_weights=active_prepared_weights,
             recurrent_layout=self.kda_recurrent_layout,
         )
         if result is None:

python/tokenspeed/runtime/layers/attention/backends/hybrid_linear_attn.py

@@ -108,6 +108,31 @@ def _mask_fresh_initial_state(
     return torch.where(mask, recurrent_state, torch.zeros_like(recurrent_state))
 
 
+def allocate_conv_state_rows_like(
+    conv_state: torch.Tensor,
+    rows: int,
+    *,
+    zero: bool,
+) -> torch.Tensor:
+    """Allocate dense rows while preserving the convolution cache layout."""
+    if conv_state.ndim != 3:
+        raise ValueError("convolution state must be [rows, channels, history]")
+    channels, history = conv_state.shape[1:]
+    feature_stride, history_stride = conv_state.stride()[1:]
+    if (feature_stride, history_stride) not in {(history, 1), (1, channels)}:
+        raise ValueError(
+            "convolution state must use a dense feature-major or "
+            f"sequence-major layout, got stride={tuple(conv_state.stride())}"
+        )
+    result = torch.empty_strided(
+        (rows, channels, history),
+        (channels * history, feature_stride, history_stride),
+        dtype=conv_state.dtype,
+        device=conv_state.device,
+    )
+    return result.zero_() if zero else result
+
+
 @dataclass(frozen=True)
 class _StateBlockIndexPlan:
     checkpoint_granularity: int
@@ -535,11 +560,7 @@ def _ensure_verify_scratch(self, bs: int, draft_token_num: int) -> None:
         for layer_id in layer_ids:
             conv, ssm = self._state_components(layer_id)
             self._verify_scratch[layer_id] = (
-                torch.zeros(
-                    (rows_needed, *conv.shape[1:]),
-                    dtype=conv.dtype,
-                    device=conv.device,
-                ),
+                allocate_conv_state_rows_like(conv, rows_needed, zero=True),
                 (
                     None
                     if self.replay_ssm
@@ -624,10 +645,14 @@ def _verify_copy_tables_get(self) -> dict:
 
         def _row_stride_i32(t: torch.Tensor) -> int:
             # Slab components are page-interleaved as_strided views: row
-            # payload contiguous, row-to-row stride the physical page.
-            if t[0].numel() and not t[0].is_contiguous():
+            # payload dense, row-to-row stride the physical block.
+            row = t[0]
+            row_is_dense = row.is_contiguous() or (
+                row.ndim == 2 and row.transpose(0, 1).is_contiguous()
+            )
+            if row.numel() and not row_is_dense:
                 raise RuntimeError(
-                    "batched verify state copy requires contiguous row payloads"
+                    "batched verify state copy requires dense row payloads"
                 )
             stride_bytes = t.stride(0) * t.element_size()
             if stride_bytes % 4:
@@ -1448,6 +1473,7 @@ def forward_decode(
         output_gate = kwargs.get("output_gate")
         norm_weight = kwargs.get("norm_weight")
         norm_eps = kwargs.get("norm_eps")
+        prepared_weights = kwargs.get("flashinfer_kda_decode_weights")
         gate_lower_bound = kwargs.get("lower_bound")
         A_log = kwargs["A_log"]
         dt_bias = kwargs["dt_bias"]
@@ -1479,6 +1505,8 @@ def forward_decode(
             output_gate=output_gate,
             norm_weight=norm_weight,
             norm_eps=norm_eps,
+            prepared_weights=prepared_weights,
+            layer_id=layer_id,
         )
         if fused_out is not None:
             return fused_out
@@ -1556,6 +1584,8 @@ def _decode(
         output_gate: torch.Tensor | None,
         norm_weight: torch.Tensor | None,
         norm_eps: float | None,
+        prepared_weights: object | None = None,
+        layer_id: int | None = None,
     ) -> torch.Tensor | None:
         """Whole-step decode attempt; ``None`` falls through to the shared flow.
 
@@ -1586,6 +1616,8 @@ def _decode(
             output_gate: Optional KDA gated-norm logits.
             norm_weight: Optional KDA output RMSNorm weight.
             norm_eps: Optional KDA output RMSNorm epsilon.
+            prepared_weights: Opaque backend plan retained by the model.
+            layer_id: Layer owning the state components, when applicable.
 
         Returns:
             The layer output when a fused kernel ran, else None.

python/tokenspeed/runtime/layers/attention/kda_geometry.py

@@ -0,0 +1,53 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""KDA convolution-state geometry shared by cache producers and consumers."""
+
+from __future__ import annotations
+
+
+def kda_conv_state_channel_axis(
+    shape: tuple[int, ...],
+    *,
+    channels: int | None = None,
+    history: int | None = None,
+) -> int:
+    """Return the channel axis of a physical KDA convolution-state row."""
+    if (channels is None) == (history is None):
+        raise ValueError("provide exactly one of channels or history")
+    expected = channels if channels is not None else history
+    if (
+        expected is None
+        or expected <= 0
+        or len(shape) != 2
+        or any(dim <= 0 for dim in shape)
+    ):
+        raise ValueError(
+            "invalid KDA convolution state geometry: "
+            f"shape={shape}, channels={channels}, history={history}"
+        )
+    matches = tuple(axis for axis, dim in enumerate(shape) if dim == expected)
+    if len(matches) != 1:
+        raise ValueError(
+            "KDA convolution state must be [channels, history] or "
+            "[history, channels], got "
+            f"shape={shape}, channels={channels}, history={history}"
+        )
+    return matches[0] if channels is not None else 1 - matches[0]