From 5ced06c0a3412816cc9bd52f2f0551f8311891e4 Mon Sep 17 00:00:00 2001 From: kliuae Date: Tue, 11 Aug 2026 17:50:23 +0000 Subject: [PATCH 1/4] remove first layers prefill bubbles Signed-off-by: kliuae --- vllm/models/kimi_k3/amd/kda.py | 7 ++ vllm/models/kimi_k3/amd/kda_metadata.py | 109 ++++++++++++++++++ .../kimi_k3/amd/ops/third_party/kda/chunk.py | 19 ++- vllm/v1/attention/backends/gdn_attn.py | 84 ++++++++------ 4 files changed, 178 insertions(+), 41 deletions(-) create mode 100644 vllm/models/kimi_k3/amd/kda_metadata.py diff --git a/vllm/models/kimi_k3/amd/kda.py b/vllm/models/kimi_k3/amd/kda.py index 868a1df171f5..dd35bfdcab38 100644 --- a/vllm/models/kimi_k3/amd/kda.py +++ b/vllm/models/kimi_k3/amd/kda.py @@ -38,6 +38,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.third_party.kda import ( chunk_kda_with_fused_gate, fused_recurrent_kda, @@ -45,10 +46,14 @@ ) 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 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") @@ -466,6 +471,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 diff --git a/vllm/models/kimi_k3/amd/kda_metadata.py b/vllm/models/kimi_k3/amd/kda_metadata.py new file mode 100644 index 000000000000..9681aca1d878 --- /dev/null +++ b/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 diff --git a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py index 9d0709ce4bf3..0a43249d6e3f 100644 --- a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py +++ b/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 diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 0b32dfcf9444..dd92f3bd8c30 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -165,6 +165,49 @@ 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 +374,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 +394,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 From d16dc9d65a271bb48cf1edea6fa60f2029bd0ebb Mon Sep 17 00:00:00 2001 From: kliuae Date: Wed, 12 Aug 2026 06:42:18 +0000 Subject: [PATCH 2/4] mypy Signed-off-by: kliuae --- .../layers/mamba/gdn/kimi_gdn_linear_attn.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py index bc49226c75b2..42e717f9b44e 100644 --- a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py +++ b/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 Any import torch from einops import rearrange @@ -393,19 +394,30 @@ 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. + # These copies may have different signatures for Kimi-K3, but they agree + # on the arguments used here. + chunk_kda_with_fused_gate: Callable[..., Any] if current_platform.is_rocm(): from vllm.models.kimi_k3.amd.ops.third_party.kda import ( - chunk_kda_with_fused_gate, + chunk_kda_with_fused_gate as _amd_chunk_kda_with_fused_gate, + ) + from vllm.models.kimi_k3.amd.ops.third_party.kda import ( fused_recurrent_kda, fused_recurrent_kda_packed_decode, ) + + chunk_kda_with_fused_gate = _amd_chunk_kda_with_fused_gate else: from vllm.models.kimi_k3.nvidia.ops.third_party.kda import ( - chunk_kda_with_fused_gate, + chunk_kda_with_fused_gate as _nvidia_chunk_kda_with_fused_gate, + ) + from vllm.models.kimi_k3.nvidia.ops.third_party.kda import ( fused_recurrent_kda, fused_recurrent_kda_packed_decode, ) + chunk_kda_with_fused_gate = _nvidia_chunk_kda_with_fused_gate + assert isinstance(attn_metadata_raw, dict) attn_metadata_narrowed = attn_metadata_raw[self.prefix] assert isinstance(attn_metadata_narrowed, GDNAttentionMetadata) From 991091cb266b72f4fc3c6fae6aeabede169098d6 Mon Sep 17 00:00:00 2001 From: kliuae Date: Wed, 12 Aug 2026 06:44:16 +0000 Subject: [PATCH 3/4] precommit Signed-off-by: kliuae --- vllm/v1/attention/backends/gdn_attn.py | 1 - 1 file changed, 1 deletion(-) diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index dd92f3bd8c30..30758455f4e5 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -171,7 +171,6 @@ def _build_chunk_metadata( 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 d30b99699bd8757501d0208cd2e2c0eeee46c6f2 Mon Sep 17 00:00:00 2001 From: kliuae Date: Wed, 12 Aug 2026 07:10:00 +0000 Subject: [PATCH 4/4] mypy Signed-off-by: kliuae --- .../layers/mamba/gdn/kimi_gdn_linear_attn.py | 23 ++++++++----------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py index 42e717f9b44e..cf0a17720f46 100644 --- a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable -from typing import Any +from typing import TYPE_CHECKING import torch from einops import rearrange @@ -396,28 +396,25 @@ def _forward( # under kimi_k3/{amd,nvidia}/ops so each can diverge independently. # These copies may have different signatures for Kimi-K3, but they agree # on the arguments used here. - chunk_kda_with_fused_gate: Callable[..., Any] - if current_platform.is_rocm(): - from vllm.models.kimi_k3.amd.ops.third_party.kda import ( - chunk_kda_with_fused_gate as _amd_chunk_kda_with_fused_gate, + 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, ) - from vllm.models.kimi_k3.amd.ops.third_party.kda import ( + 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, ) - - chunk_kda_with_fused_gate = _amd_chunk_kda_with_fused_gate else: from vllm.models.kimi_k3.nvidia.ops.third_party.kda import ( - chunk_kda_with_fused_gate as _nvidia_chunk_kda_with_fused_gate, - ) - from vllm.models.kimi_k3.nvidia.ops.third_party.kda import ( + chunk_kda_with_fused_gate, fused_recurrent_kda, fused_recurrent_kda_packed_decode, ) - chunk_kda_with_fused_gate = _nvidia_chunk_kda_with_fused_gate - assert isinstance(attn_metadata_raw, dict) attn_metadata_narrowed = attn_metadata_raw[self.prefix] assert isinstance(attn_metadata_narrowed, GDNAttentionMetadata)