[ROCm][Perf] Kimi-K3 Remove prefill pipeline stall in chunk KDA

vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from collections.abc import Callable
+from typing import TYPE_CHECKING
 
 import torch
 from einops import rearrange
@@ -393,8 +394,16 @@ def _forward(
 
         # Vendor-specific KDA kernels: AMD/ROCm and NVIDIA keep their own copies
         # under kimi_k3/{amd,nvidia}/ops so each can diverge independently.
-        if current_platform.is_rocm():
-            from vllm.models.kimi_k3.amd.ops.third_party.kda import (
+        # These copies may have different signatures for Kimi-K3, but they agree
+        # on the arguments used here.
+        if TYPE_CHECKING:
+            from vllm.models.kimi_k3.nvidia.ops.third_party.kda import (
+                chunk_kda_with_fused_gate,
+                fused_recurrent_kda,
+                fused_recurrent_kda_packed_decode,
+            )
+        elif current_platform.is_rocm():
+            from vllm.models.kimi_k3.amd.ops.third_party.kda import (  # type: ignore[assignment]
                 chunk_kda_with_fused_gate,
                 fused_recurrent_kda,
                 fused_recurrent_kda_packed_decode,

vllm/models/kimi_k3/amd/kda.py

@@ -40,6 +40,7 @@
 )
 from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader
 from vllm.model_executor.utils import set_weight_attrs
+from vllm.models.kimi_k3.amd.kda_metadata import KimiK3ROCmKDABackend
 from vllm.models.kimi_k3.amd.ops.kda_decode import (
     is_fused_kda_decode_supported,
     make_decode_conv1d_weight_loader,
@@ -52,12 +53,16 @@
 )
 from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated
 from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
+from vllm.v1.attention.backend import AttentionBackend
 from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
 
 logger = init_logger(__name__)
 
 
 class KimiK3DeltaAttention(GatedDeltaNetAttention):
+    def get_attn_backend(self) -> type[AttentionBackend]:
+        return KimiK3ROCmKDABackend
+
     def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]:
         if self.model_config is None or self.cache_config is None:
             raise ValueError("model_config and cache_config must be set")
@@ -546,6 +551,8 @@ def _prefill_conv(
                     output_final_state=True,
                     use_qk_l2norm_in_kernel=True,
                     cu_seqlens=prefill_query_start_loc,
+                    chunk_indices=m.chunk_indices,
+                    chunk_offsets=m.chunk_offsets,
                 )
                 # Init cache
                 recurrent_state[prefill_state_indices] = last_recurrent_state

vllm/models/kimi_k3/amd/kda_metadata.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""ROCm Kimi-K3 specialization of GDN attention metadata.
+
+The request classification and cudagraph staging intentionally mirror
+``GDNAttentionMetadataBuilder``. Only the FLA chunk metadata is built
+differently on device rather than on the host.
+"""
+
+import torch
+
+from vllm.logger import init_logger
+from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE
+from vllm.triton_utils import tl, triton
+from vllm.utils.math_utils import next_power_of_2
+from vllm.v1.attention.backends.gdn_attn import (
+    GDNAttentionBackend,
+    GDNAttentionMetadataBuilder,
+)
+
+logger = init_logger(__name__)
+
+_BLOCK_T = 256
+_MIN_BLOCK_N = 128
+
+
+@triton.jit(do_not_specialize=["N"])
+def _chunk_metadata_kernel(
+    cu_seqlens,
+    chunk_indices,
+    chunk_offsets,
+    N,
+    BT: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_T: tl.constexpr,
+):
+    i_n = tl.program_id(0)
+    offs_n = tl.arange(0, BLOCK_N)
+    is_seq = offs_n < N
+    bos = tl.load(cu_seqlens + offs_n, mask=is_seq, other=0).to(tl.int32)
+    eos = tl.load(cu_seqlens + offs_n + 1, mask=is_seq, other=0).to(tl.int32)
+    nt = tl.where(is_seq, tl.cdiv(eos - bos, BT), 0)
+
+    base = tl.sum(tl.where(offs_n < i_n, nt, 0))
+    num_chunks = tl.sum(tl.where(offs_n == i_n, nt, 0))
+
+    tl.store(chunk_offsets + i_n, base)
+    if i_n == 0:
+        tl.store(chunk_offsets + N, tl.sum(nt))
+
+    for t0 in range(0, num_chunks, BLOCK_T):
+        offs_t = t0 + tl.arange(0, BLOCK_T)
+        mask_t = offs_t < num_chunks
+        row = (base + offs_t) * 2
+        tl.store(chunk_indices + row, tl.full([BLOCK_T], i_n, tl.int32), mask=mask_t)
+        tl.store(chunk_indices + row + 1, offs_t.to(tl.int32), mask=mask_t)
+
+
+def prepare_chunk_metadata_device(
+    cu_seqlens: torch.Tensor,
+    cu_seqlens_cpu: torch.Tensor,
+    chunk_size: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Build FLA chunk metadata on device, with no host<->device transfer."""
+    num_seqs = cu_seqlens_cpu.numel() - 1
+    seq_lens = cu_seqlens_cpu[1:] - cu_seqlens_cpu[:-1]
+    num_chunks = int(((seq_lens + chunk_size - 1) // chunk_size).sum())
+
+    chunk_indices = torch.empty(
+        num_chunks, 2, dtype=cu_seqlens.dtype, device=cu_seqlens.device
+    )
+    chunk_offsets = torch.empty(
+        num_seqs + 1, dtype=torch.int64, device=cu_seqlens.device
+    )
+    _chunk_metadata_kernel[(num_seqs,)](
+        cu_seqlens,
+        chunk_indices,
+        chunk_offsets,
+        num_seqs,
+        BT=chunk_size,
+        BLOCK_N=max(_MIN_BLOCK_N, next_power_of_2(num_seqs + 1)),
+        BLOCK_T=_BLOCK_T,
+        num_warps=4,
+    )
+    return chunk_indices, chunk_offsets
+
+
+class KimiK3ROCmKDAMetadataBuilder(GDNAttentionMetadataBuilder):
+    def _build_chunk_metadata(
+        self,
+        prefill_query_start_loc: torch.Tensor,
+        prefill_query_start_loc_cpu: torch.Tensor,
+        device: torch.device,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return prepare_chunk_metadata_device(
+            prefill_query_start_loc,
+            prefill_query_start_loc_cpu,
+            FLA_CHUNK_SIZE,
+        )
+
+
+class KimiK3ROCmKDABackend(GDNAttentionBackend):
+    @staticmethod
+    def get_name() -> str:
+        return "KIMI_K3_KDA_ROCM"
+
+    @staticmethod
+    def get_builder_cls() -> type[KimiK3ROCmKDAMetadataBuilder]:
+        return KimiK3ROCmKDAMetadataBuilder

vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py

@@ -601,6 +601,7 @@ def _chunk_kda_fwd_with_cumulative_g(
     output_final_state: bool,
     cu_seqlens: torch.Tensor | None = None,
     chunk_indices: torch.Tensor | None = None,
+    chunk_offsets: torch.Tensor | None = None,
     chunk_size: int = FLA_CHUNK_SIZE,
     safe_gate: bool = False,
 ):
@@ -638,6 +639,7 @@ def _chunk_kda_fwd_with_cumulative_g(
         output_final_state=output_final_state,
         cu_seqlens=cu_seqlens,
         chunk_indices=chunk_indices,
+        chunk_offsets=chunk_offsets,
         use_exp2=True,
     )
     del w, u, kg
@@ -711,13 +713,15 @@ def chunk_kda_with_fused_gate_fwd(
     output_final_state: bool,
     lower_bound: float | None = None,
     cu_seqlens: torch.Tensor | None = None,
+    chunk_indices: torch.Tensor | None = None,
+    chunk_offsets: torch.Tensor | None = None,
 ):
     chunk_size = FLA_CHUNK_SIZE
-    chunk_indices = (
-        prepare_chunk_indices(cu_seqlens, chunk_size)
-        if cu_seqlens is not None
-        else None
-    )
+    # Building chunk metadata from a device `cu_seqlens` costs a blocking
+    # device->host copy. Callers that already hold it, i.e. the attention
+    # metadata builder, should pass it in.
+    if chunk_indices is None and cu_seqlens is not None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
     g, beta = fused_kda_gate_chunk_cumsum(
         raw_g,
         raw_beta=raw_beta,
@@ -739,6 +743,7 @@ def chunk_kda_with_fused_gate_fwd(
         output_final_state=output_final_state,
         cu_seqlens=cu_seqlens,
         chunk_indices=chunk_indices,
+        chunk_offsets=chunk_offsets,
         chunk_size=chunk_size,
         safe_gate=lower_bound is not None,
     )
@@ -792,6 +797,8 @@ def chunk_kda_with_fused_gate(
     lower_bound: float | None = None,
     use_qk_l2norm_in_kernel: bool = False,
     cu_seqlens: torch.Tensor | None = None,
+    chunk_indices: torch.Tensor | None = None,
+    chunk_offsets: torch.Tensor | None = None,
     **kwargs,
 ):
     """Run chunk KDA from raw gate and beta projections."""
@@ -815,6 +822,8 @@ def chunk_kda_with_fused_gate(
         output_final_state=output_final_state,
         lower_bound=lower_bound,
         cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        chunk_offsets=chunk_offsets,
     )
     return o, final_state
 

vllm/v1/attention/backends/gdn_attn.py

@@ -165,6 +165,48 @@ def __init__(
             device=device,
         )
 
+    def _build_chunk_metadata(
+        self,
+        prefill_query_start_loc: torch.Tensor,
+        prefill_query_start_loc_cpu: torch.Tensor,
+        device: torch.device,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE
+
+        if self.gdn_prefill_backend == "cutedsl":
+            from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import (
+                prepare_metadata_cutedsl,
+            )
+
+            assert prefill_query_start_loc is not None
+            assert prefill_query_start_loc_cpu is not None
+            total_tokens = int(prefill_query_start_loc_cpu[-1].item())
+            return prepare_metadata_cutedsl(
+                prefill_query_start_loc,
+                total_tokens,
+                FLA_CHUNK_SIZE,
+            )
+
+        # Only prefill batches use FLA chunk ops.
+        # Pre-compute on CPU and async-copy to GPU to avoid
+        # GPU→CPU sync (.tolist()) in prepare_chunk_indices.
+        from vllm.third_party.flash_linear_attention.ops.index import (
+            prepare_chunk_indices,
+            prepare_chunk_offsets,
+        )
+
+        assert prefill_query_start_loc_cpu is not None
+        return (
+            async_tensor_h2d(
+                prepare_chunk_indices(prefill_query_start_loc_cpu, FLA_CHUNK_SIZE),
+                device=device,
+            ),
+            async_tensor_h2d(
+                prepare_chunk_offsets(prefill_query_start_loc_cpu, FLA_CHUNK_SIZE),
+                device=device,
+            ),
+        )
+
     def build(  # type: ignore[override]
         self,
         common_prefix_len: int,
@@ -331,10 +373,6 @@ def build(  # type: ignore[override]
         prefill_state_indices: torch.Tensor | None = None
         prefill_has_initial_state: torch.Tensor | None = None
         if num_prefills > 0:
-            from vllm.third_party.flash_linear_attention.ops.utils import (
-                FLA_CHUNK_SIZE,
-            )
-
             # In a mixed non-spec batch, decodes are peeled off to the recurrent
             # kernel (decode-first front slice), so build chunk metadata from the
             # rebased prefill-only cu_seqlens; otherwise use the full non-spec one.
@@ -355,38 +393,11 @@ def build(  # type: ignore[override]
                 prefill_query_start_loc_cpu = non_spec_query_start_loc_cpu
                 prefill_state_indices = non_spec_state_indices_tensor
 
-            if self.gdn_prefill_backend == "cutedsl":
-                from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import (
-                    prepare_metadata_cutedsl,
-                )
-
-                assert prefill_query_start_loc is not None
-                assert prefill_query_start_loc_cpu is not None
-                total_tokens = int(prefill_query_start_loc_cpu[-1].item())
-                chunk_indices, chunk_offsets = prepare_metadata_cutedsl(
-                    prefill_query_start_loc,
-                    total_tokens,
-                    FLA_CHUNK_SIZE,
-                )
-            else:
-                gpu_device = query_start_loc.device
-                # Only prefill batches use FLA chunk ops.
-                # Pre-compute on CPU and async-copy to GPU to avoid
-                # GPU→CPU sync (.tolist()) in prepare_chunk_indices.
-                from vllm.third_party.flash_linear_attention.ops.index import (
-                    prepare_chunk_indices,
-                    prepare_chunk_offsets,
-                )
-
-                assert prefill_query_start_loc_cpu is not None
-                chunk_indices = async_tensor_h2d(
-                    prepare_chunk_indices(prefill_query_start_loc_cpu, FLA_CHUNK_SIZE),
-                    device=gpu_device,
-                )
-                chunk_offsets = async_tensor_h2d(
-                    prepare_chunk_offsets(prefill_query_start_loc_cpu, FLA_CHUNK_SIZE),
-                    device=gpu_device,
-                )
+            chunk_indices, chunk_offsets = self._build_chunk_metadata(
+                prefill_query_start_loc,
+                prefill_query_start_loc_cpu,
+                query_start_loc.device,
+            )
 
         if num_prefills > 0:
             has_initial_state = context_lens_tensor > 0