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..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,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, diff --git a/vllm/models/kimi_k3/amd/kda.py b/vllm/models/kimi_k3/amd/kda.py index b519c6a367ac..870c13a3ed16 100644 --- a/vllm/models/kimi_k3/amd/kda.py +++ b/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 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..30758455f4e5 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/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