[None][perf] Integrate FlashInfer KDA prefill with fused convolution

.pre-commit-config.yaml

@@ -1469,7 +1469,7 @@ legacy-files: &legacy_files |
 # list; the hook's own `files:` pattern only gates *when* the hook triggers.
 
 # Global exclude: vendored code + trtllm-gen FMHA artifacts (cubin pointers, export headers, cuda_ptx)
-exclude: '(^cpp/tensorrt_llm/common/sha256/|^triton_kernels/|^tensorrt_llm/_torch/attention/backends/prims_ts/|^tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/|trtllmGenKernels/fmha/cubin/kernelMetaInfo\.h$|cubin\.cpp$|cubin\.h$|trtllmGenKernels/fmha/trtllmGen_fmha_export/|trtllmGenKernels/fmha/cuda_ptx/|trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelMetaInfo\.h$|trtllmGenKernels/gemm/trtllmGen_gemm_export/KernelMetaInfo\.h$|\.cubin\.tar\.zst$)'
+exclude: '(^tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_chunked_bt16(_small_bh)?\.py$|^cpp/tensorrt_llm/common/sha256/|^triton_kernels/|^tensorrt_llm/_torch/attention/backends/prims_ts/|^tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/|trtllmGenKernels/fmha/cubin/kernelMetaInfo\.h$|cubin\.cpp$|cubin\.h$|trtllmGenKernels/fmha/trtllmGen_fmha_export/|trtllmGenKernels/fmha/cuda_ptx/|trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelMetaInfo\.h$|trtllmGenKernels/gemm/trtllmGen_gemm_export/KernelMetaInfo\.h$|\.cubin\.tar\.zst$)'
 
 default_install_hook_types: [pre-commit, commit-msg]
 repos:

3rdparty/vendor_patches/flashinfer-kda-bt16.patch

@@ -0,0 +1,207 @@
+diff --git a/kda_chunked_bt16_small_bh.py b/kda_chunked_bt16_small_bh.py
+index e8bc9feb5ece3dbfe3156c90bb27e68fa5a4521a..7478dc054a5a9afb9de5d416588a4b074971f38c 100644
+--- a/kda_chunked_bt16_small_bh.py
++++ b/kda_chunked_bt16_small_bh.py
+@@ -1,3 +1,4 @@
++# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ # Copyright (c) 2026 by FlashInfer team.
+ #
+ # Licensed under the Apache License, Version 2.0 (the "License");
+@@ -12,6 +13,16 @@
+ # See the License for the specific language governing permissions and
+ # limitations under the License.
+
++"""FlashInfer Small-BH KDA, with indexed recurrent-state pool access.
++
++Vendored from flashinfer-ai/flashinfer commit
++b8e81a7dae6c2aa5c9aa19c06e81e9a729271493,
++flashinfer/kda_kernels/kda_chunked_bt16_small_bh.py (Apache-2.0).
++Local changes add state-slot indirection and padded pool-row strides, reject
++unwarmed graph specializations, and use the active device SM count for scheduling.
++The K1/K2 arithmetic and layouts are retained from upstream.
++"""
++
+ import functools
+ from typing import Optional
+
+@@ -24,7 +35,6 @@ from cutlass._mlir.dialects import llvm
+ from cutlass.cute.core import _pack_shape
+ from cutlass.cute.nvgpu import cpasync, warp
+ from cutlass.cutlass_dsl import dsl_user_op
+-from flashinfer.cute_dsl.utils import get_max_active_clusters
+
+
+ RCP_LN2 = 1.4426950408889634
+@@ -1803,6 +1813,7 @@ class ChunkKdaFwdK2:
+         mS_in: Optional[cute.Tensor] = None,  # (B, H, D, D)
+         mS_out: Optional[cute.Tensor] = None,  # (B, H, D, D)
+         mCuSeqlens: Optional[cute.Tensor] = None,  # (B + 1,)
++        mStateIndices: Optional[cute.Tensor] = None,  # (B,)
+         stream=None,
+     ):
+         if cutlass.const_expr(mCuSeqlens is not None):
+@@ -2076,6 +2087,7 @@ class ChunkKdaFwdK2:
+             mS_in,
+             mS_out,
+             mCuSeqlens,
++            mStateIndices,
+             sQK_layout_staged,
+             sQKt_layout_staged,
+             sV_layout_staged,
+@@ -2123,6 +2135,7 @@ class ChunkKdaFwdK2:
+         mS_in: Optional[cute.Tensor],
+         mS_out: Optional[cute.Tensor],
+         mCuSeqlens: Optional[cute.Tensor],
++        mStateIndices: Optional[cute.Tensor],
+         sQK_layout_staged: cute.ComposedLayout,
+         sQKt_layout_staged: cute.ComposedLayout,
+         sV_layout_staged: cute.ComposedLayout,
+@@ -2148,6 +2161,9 @@ class ChunkKdaFwdK2:
+         warpgroup_idx = cute.arch.make_warp_uniform(tidx // self.warpgroup_threads)
+
+         v_idx, h_idx, b_idx = cute.arch.block_idx()
++        state_idx = b_idx
++        if cutlass.const_expr(mStateIndices is not None):
++            state_idx = mStateIndices[b_idx]
+
+         seq_len = cute.size(tma_tensor_Q, mode=[0])
+         seq_offset = 0
+@@ -2475,11 +2491,11 @@ class ChunkKdaFwdK2:
+             if cutlass.const_expr(mS_in is not None):
+                 if cutlass.const_expr(self.transpose_S):
+                     gS_in = cute.local_tile(
+-                        mS_in[None, None, h_idx, b_idx], (self.BV, self.D), (v_idx, 0)
++                        mS_in[None, None, h_idx, state_idx], (self.BV, self.D), (v_idx, 0)
+                     )
+                 else:
+                     gS_in = cute.local_tile(
+-                        mS_in[None, None, h_idx, b_idx], (self.D, self.BV), (0, v_idx)
++                        mS_in[None, None, h_idx, state_idx], (self.D, self.BV), (0, v_idx)
+                     )
+                     gS_in = cute.make_tensor(
+                         gS_in.iterator, cute.select(gS_in.layout, mode=[1, 0])
+@@ -2763,11 +2779,11 @@ class ChunkKdaFwdK2:
+             if cutlass.const_expr(mS_out is not None):
+                 if cutlass.const_expr(self.transpose_S):
+                     gS_out = cute.local_tile(
+-                        mS_out[None, None, h_idx, b_idx], (self.BV, self.D), (v_idx, 0)
++                        mS_out[None, None, h_idx, state_idx], (self.BV, self.D), (v_idx, 0)
+                     )
+                 else:
+                     gS_out = cute.local_tile(
+-                        mS_out[None, None, h_idx, b_idx], (self.D, self.BV), (0, v_idx)
++                        mS_out[None, None, h_idx, state_idx], (self.D, self.BV), (0, v_idx)
+                     )
+                     gS_out = cute.make_tensor(
+                         gS_out.iterator, cute.select(gS_out.layout, mode=[1, 0])
+@@ -2808,12 +2824,12 @@ def _cute_dtype(dtype: torch.dtype):
+     raise ValueError(f"unsupported dtype: {dtype}")
+
+
+-def _fake(dtype, shape):
++def _fake(dtype, shape, assumed_align=16):
+     return cute.runtime.make_fake_compact_tensor(
+         dtype,
+         shape,
+         stride_order=tuple(reversed(range(len(shape)))),
+-        assumed_align=16,
++        assumed_align=assumed_align,
+     )
+
+
+@@ -2838,6 +2854,8 @@ def _compile_k1(
+     is_varlen: bool,
+     max_active_clusters: int,
+ ):
++    if torch.cuda.is_current_stream_capturing():
++        raise RuntimeError("Warm Small-BH KDA before CUDA graph capture")
+     cu_dtype = _cute_dtype(dtype)
+     b, t, tp, nt = (cute.sym_int() for _ in range(4))
+     feature: tuple[object, ...]
+@@ -2904,7 +2922,11 @@ def _compile_k2(
+     has_state_output: bool,
+     transpose_state: bool,
+     is_varlen: bool,
++    state_row_stride: int | None,
++    has_state_indices: bool,
+ ):
++    if torch.cuda.is_current_stream_capturing():
++        raise RuntimeError("Warm Small-BH KDA before CUDA graph capture")
+     cu_dtype = _cute_dtype(dtype)
+     b, bs, t, tp, h, nt = (cute.sym_int() for _ in range(6))
+     feature: tuple[object, ...]
+@@ -2925,6 +2947,14 @@ def _compile_k2(
+         cu_seqlens = None
+     state = (b, h, d, d)
+
++    def fake_state():
++        if state_row_stride is None:
++            return _fake(_cute_dtype(state_dtype), state)
++        return cute.runtime.make_fake_tensor(
++            _cute_dtype(state_dtype), state,
++            stride=(state_row_stride, d * d, d, 1), assumed_align=16,
++        )
++
+     q, k, kr = [_fake(cu_dtype, feature) for _ in range(3)]
+     kernel = ChunkKdaFwdK2(D=d, BT=_BT, BV=bv, transpose_S=transpose_state)
+     stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
+@@ -2938,9 +2968,10 @@ def _compile_k2(
+         _fake(cu_dtype, matrix),
+         _fake(cutlass.Float32, gate),
+         _fake(cu_dtype, output),
+-        _fake(_cute_dtype(state_dtype), state) if has_state_input else None,
+-        _fake(_cute_dtype(state_dtype), state) if has_state_output else None,
++        fake_state() if has_state_input else None,
++        fake_state() if has_state_output else None,
+         cu_seqlens,
++        _fake(cutlass.Int32, (cute.sym_int(),), assumed_align=4) if has_state_indices else None,
+         stream,
+         options="--enable-tvm-ffi",
+     )
+@@ -3040,6 +3071,7 @@ def chunk_kda_fwd_k2(
+     output: torch.Tensor | None = None,
+     state_input: torch.Tensor | None = None,
+     state_output: torch.Tensor | None = None,
++    state_indices: torch.Tensor | None = None,
+     max_active_clusters: int = 152,
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+     is_varlen = cu_seqlens is not None
+@@ -3072,6 +3104,8 @@ def chunk_kda_fwd_k2(
+         has_output,
+         transpose_state,
+         is_varlen,
++        state_input.stride(0) if state_indices is not None else None,
++        state_indices is not None,
+     )
+     compiled(
+         q,
+@@ -3085,6 +3119,7 @@ def chunk_kda_fwd_k2(
+         state_input,
+         state_output,
+         cu_seqlens,
++        state_indices,
+     )
+     return output, state_output
+
+@@ -3108,11 +3143,12 @@ def chunk_kda_fwd(
+     output: torch.Tensor | None = None,
+     state_input: torch.Tensor | None = None,
+     state_output: torch.Tensor | None = None,
++    state_indices: torch.Tensor | None = None,
+     workspace_tensors: tuple[torch.Tensor, ...] | None = None,
+     max_active_clusters: int = 0,
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+     if max_active_clusters <= 0:
+-        max_active_clusters = get_max_active_clusters(1)
++        max_active_clusters = torch.cuda.get_device_properties(q.device).multi_processor_count
+     qd, kd, kr, mqk, mkk, gk = chunk_kda_fwd_k1(
+         q,
+         k,
+@@ -3142,6 +3178,7 @@ def chunk_kda_fwd(
+         output=output,
+         state_input=state_input,
+         state_output=state_output,
++        state_indices=state_indices,
+         max_active_clusters=max_active_clusters,
+     )
+

3rdparty/vendor_sources.lock.yaml

@@ -1,5 +1,17 @@
 schema_version: 1
 vendors:
+  flashinfer-kda-bt16:
+    url: https://github.com/flashinfer-ai/flashinfer.git
+    branch: main
+    commit: b8e81a7dae6c2aa5c9aa19c06e81e9a729271493
+    source: flashinfer/kda_kernels
+    destination: tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda
+    include:
+    - kda_chunked_bt16.py
+    - kda_chunked_bt16_small_bh.py
+    patch: 3rdparty/vendor_patches/flashinfer-kda-bt16.patch
+    patch_digest: sha256:bc9b96a69873e0ddfcba66ab899c666cf52766976f24d850048d2a0bf8d967ff
+    digest: sha256-tree-v1:d2d382335ccb68b55078523e63835c28b699607227afb81308cb7bb1d85845ac
   flashinfer-prims-ts:
     url: https://github.com/yuxianq/flashinfer.git
     branch: trtllm-prims-ts-dev-20261012

pyproject.toml

@@ -40,6 +40,9 @@ fix = true
 force-exclude = true
 exclude = [
     "**3rdparty/**",
+    # Preserve vendored FlashInfer KDA sources as recorded in the vendor lock.
+    "tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_chunked_bt16.py",
+    "tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_chunked_bt16_small_bh.py",
     "tensorrt_llm/_torch/attention/backends/prims_ts/**",
     "tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/**",
     "triton_kernels/**",