iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm][Perf] Kimi-K3 Remove prefill pipeline stall in chunk KDA (#51862)

Browse Source 
Signed-off-by: kliuae <kuanfu.liu@embeddedllm.com>
This commit is contained in:
kliuae
2026-08-13 09:11:45 +00:00
committed by GitHub
parent 50ba4bc6b2
commit f96261637c
5 changed files with 188 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py
+11 -2
View File
@@ -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 @@ class KimiGatedDeltaNetAttention(GatedDeltaNetAttention):
# 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
+7
View File
@@ -40,6 +40,7 @@ from vllm.model_executor.layers.mamba.ops.gather_initial_states import (
)
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.models.kimi_k3.amd.ops.third_party.kda import (
)
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 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
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
+109
View File
@@ -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
Vendored
+14 -5
View File
@@ -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
+47 -36
View File
@@ -165,6 +165,48 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
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,
@@ -330,10 +372,6 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
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.
@@ -354,38 +392,11 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
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:
context_lens_tensor = m.compute_num_computed_tokens()