[ROCm][Perf] Kimi-K3 Remove prefill pipeline stall in chunk KDA (#51862)
Signed-off-by: kliuae <kuanfu.liu@embeddedllm.com>
kliuae committed
Aug 13, 2026 at 17:11 UTC
f96261637c05d372e077a8aeb6df03cda7f5b354
5 files changed
+188
-43
vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py
+11
-2
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from collections.abc import Callable
5
+from typing import TYPE_CHECKING
6
7
import torch
8
from einops import rearrange
@@ -393,8 +394,16 @@ class KimiGatedDeltaNetAttention(GatedDeltaNetAttention):
394
395
# Vendor-specific KDA kernels: AMD/ROCm and NVIDIA keep their own copies
396
# under kimi_k3/{amd,nvidia}/ops so each can diverge independently.
396
- if current_platform.is_rocm():
397
- from vllm.models.kimi_k3.amd.ops.third_party.kda import (
397
+ # These copies may have different signatures for Kimi-K3, but they agree
398
+ # on the arguments used here.
399
+ if TYPE_CHECKING:
400
+ from vllm.models.kimi_k3.nvidia.ops.third_party.kda import (
401
+ chunk_kda_with_fused_gate,
402
+ fused_recurrent_kda,
403
+ fused_recurrent_kda_packed_decode,
404
+ )
405
+ elif current_platform.is_rocm():
406
+ from vllm.models.kimi_k3.amd.ops.third_party.kda import ( # type: ignore[assignment]
407
chunk_kda_with_fused_gate,
408
fused_recurrent_kda,
409
fused_recurrent_kda_packed_decode,
vllm/models/kimi_k3/amd/kda.py
+7
@@ -40,6 +40,7 @@ from vllm.model_executor.layers.mamba.ops.gather_initial_states import (
40
)
41
from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader
42
from vllm.model_executor.utils import set_weight_attrs
43
+from vllm.models.kimi_k3.amd.kda_metadata import KimiK3ROCmKDABackend
44
from vllm.models.kimi_k3.amd.ops.kda_decode import (
45
is_fused_kda_decode_supported,
46
make_decode_conv1d_weight_loader,
@@ -52,12 +53,16 @@ from vllm.models.kimi_k3.amd.ops.third_party.kda import (
53
)
54
from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated
55
from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
56
+from vllm.v1.attention.backend import AttentionBackend
57
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
58
59
logger = init_logger(__name__)
60
61
62
class KimiK3DeltaAttention(GatedDeltaNetAttention):
63
+ def get_attn_backend(self) -> type[AttentionBackend]:
64
+ return KimiK3ROCmKDABackend
65
+
66
def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]:
67
if self.model_config is None or self.cache_config is None:
68
raise ValueError("model_config and cache_config must be set")
@@ -546,6 +551,8 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
551
output_final_state=True,
552
use_qk_l2norm_in_kernel=True,
553
cu_seqlens=prefill_query_start_loc,
554
+ chunk_indices=m.chunk_indices,
555
+ chunk_offsets=m.chunk_offsets,
556
)
557
# Init cache
558
recurrent_state[prefill_state_indices] = last_recurrent_state
vllm/models/kimi_k3/amd/kda_metadata.py
new
+109
@@ -0,0 +1,109 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""ROCm Kimi-K3 specialization of GDN attention metadata.
4
+
5
+The request classification and cudagraph staging intentionally mirror
6
+``GDNAttentionMetadataBuilder``. Only the FLA chunk metadata is built
7
+differently on device rather than on the host.
8
+"""
9
+
10
+import torch
11
+
12
+from vllm.logger import init_logger
13
+from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE
14
+from vllm.triton_utils import tl, triton
15
+from vllm.utils.math_utils import next_power_of_2
16
+from vllm.v1.attention.backends.gdn_attn import (
17
+ GDNAttentionBackend,
18
+ GDNAttentionMetadataBuilder,
19
+)
20
+
21
+logger = init_logger(__name__)
22
+
23
+_BLOCK_T = 256
24
+_MIN_BLOCK_N = 128
25
+
26
+
27
+@triton.jit(do_not_specialize=["N"])
28
+def _chunk_metadata_kernel(
29
+ cu_seqlens,
30
+ chunk_indices,
31
+ chunk_offsets,
32
+ N,
33
+ BT: tl.constexpr,
34
+ BLOCK_N: tl.constexpr,
35
+ BLOCK_T: tl.constexpr,
36
+):
37
+ i_n = tl.program_id(0)
38
+ offs_n = tl.arange(0, BLOCK_N)
39
+ is_seq = offs_n < N
40
+ bos = tl.load(cu_seqlens + offs_n, mask=is_seq, other=0).to(tl.int32)
41
+ eos = tl.load(cu_seqlens + offs_n + 1, mask=is_seq, other=0).to(tl.int32)
42
+ nt = tl.where(is_seq, tl.cdiv(eos - bos, BT), 0)
43
+
44
+ base = tl.sum(tl.where(offs_n < i_n, nt, 0))
45
+ num_chunks = tl.sum(tl.where(offs_n == i_n, nt, 0))
46
+
47
+ tl.store(chunk_offsets + i_n, base)
48
+ if i_n == 0:
49
+ tl.store(chunk_offsets + N, tl.sum(nt))
50
+
51
+ for t0 in range(0, num_chunks, BLOCK_T):
52
+ offs_t = t0 + tl.arange(0, BLOCK_T)
53
+ mask_t = offs_t < num_chunks
54
+ row = (base + offs_t) * 2
55
+ tl.store(chunk_indices + row, tl.full([BLOCK_T], i_n, tl.int32), mask=mask_t)
56
+ tl.store(chunk_indices + row + 1, offs_t.to(tl.int32), mask=mask_t)
57
+
58
+
59
+def prepare_chunk_metadata_device(
60
+ cu_seqlens: torch.Tensor,
61
+ cu_seqlens_cpu: torch.Tensor,
62
+ chunk_size: int,
63
+) -> tuple[torch.Tensor, torch.Tensor]:
64
+ """Build FLA chunk metadata on device, with no host<->device transfer."""
65
+ num_seqs = cu_seqlens_cpu.numel() - 1
66
+ seq_lens = cu_seqlens_cpu[1:] - cu_seqlens_cpu[:-1]
67
+ num_chunks = int(((seq_lens + chunk_size - 1) // chunk_size).sum())
68
+
69
+ chunk_indices = torch.empty(
70
+ num_chunks, 2, dtype=cu_seqlens.dtype, device=cu_seqlens.device
71
+ )
72
+ chunk_offsets = torch.empty(
73
+ num_seqs + 1, dtype=torch.int64, device=cu_seqlens.device
74
+ )
75
+ _chunk_metadata_kernel[(num_seqs,)](
76
+ cu_seqlens,
77
+ chunk_indices,
78
+ chunk_offsets,
79
+ num_seqs,
80
+ BT=chunk_size,
81
+ BLOCK_N=max(_MIN_BLOCK_N, next_power_of_2(num_seqs + 1)),
82
+ BLOCK_T=_BLOCK_T,
83
+ num_warps=4,
84
+ )
85
+ return chunk_indices, chunk_offsets
86
+
87
+
88
+class KimiK3ROCmKDAMetadataBuilder(GDNAttentionMetadataBuilder):
89
+ def _build_chunk_metadata(
90
+ self,
91
+ prefill_query_start_loc: torch.Tensor,
92
+ prefill_query_start_loc_cpu: torch.Tensor,
93
+ device: torch.device,
94
+ ) -> tuple[torch.Tensor, torch.Tensor]:
95
+ return prepare_chunk_metadata_device(
96
+ prefill_query_start_loc,
97
+ prefill_query_start_loc_cpu,
98
+ FLA_CHUNK_SIZE,
99
+ )
100
+
101
+
102
+class KimiK3ROCmKDABackend(GDNAttentionBackend):
103
+ @staticmethod
104
+ def get_name() -> str:
105
+ return "KIMI_K3_KDA_ROCM"
106
+
107
+ @staticmethod
108
+ def get_builder_cls() -> type[KimiK3ROCmKDAMetadataBuilder]:
109
+ return KimiK3ROCmKDAMetadataBuilder
vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py
+14
-5
@@ -601,6 +601,7 @@ def _chunk_kda_fwd_with_cumulative_g(
601
output_final_state: bool,
602
cu_seqlens: torch.Tensor | None = None,
603
chunk_indices: torch.Tensor | None = None,
604
+ chunk_offsets: torch.Tensor | None = None,
605
chunk_size: int = FLA_CHUNK_SIZE,
606
safe_gate: bool = False,
607
):
@@ -638,6 +639,7 @@ def _chunk_kda_fwd_with_cumulative_g(
639
output_final_state=output_final_state,
640
cu_seqlens=cu_seqlens,
641
chunk_indices=chunk_indices,
642
+ chunk_offsets=chunk_offsets,
643
use_exp2=True,
644
)
645
del w, u, kg
@@ -711,13 +713,15 @@ def chunk_kda_with_fused_gate_fwd(
713
output_final_state: bool,
714
lower_bound: float | None = None,
715
cu_seqlens: torch.Tensor | None = None,
716
+ chunk_indices: torch.Tensor | None = None,
717
+ chunk_offsets: torch.Tensor | None = None,
718
):
719
chunk_size = FLA_CHUNK_SIZE
716
- chunk_indices = (
717
- prepare_chunk_indices(cu_seqlens, chunk_size)
718
- if cu_seqlens is not None
719
- else None
720
- )
720
+ # Building chunk metadata from a device `cu_seqlens` costs a blocking
721
+ # device->host copy. Callers that already hold it, i.e. the attention
722
+ # metadata builder, should pass it in.
723
+ if chunk_indices is None and cu_seqlens is not None:
724
+ chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
725
g, beta = fused_kda_gate_chunk_cumsum(
726
raw_g,
727
raw_beta=raw_beta,
@@ -739,6 +743,7 @@ def chunk_kda_with_fused_gate_fwd(
743
output_final_state=output_final_state,
744
cu_seqlens=cu_seqlens,
745
chunk_indices=chunk_indices,
746
+ chunk_offsets=chunk_offsets,
747
chunk_size=chunk_size,
748
safe_gate=lower_bound is not None,
749
)
@@ -792,6 +797,8 @@ def chunk_kda_with_fused_gate(
797
lower_bound: float | None = None,
798
use_qk_l2norm_in_kernel: bool = False,
799
cu_seqlens: torch.Tensor | None = None,
800
+ chunk_indices: torch.Tensor | None = None,
801
+ chunk_offsets: torch.Tensor | None = None,
802
**kwargs,
803
):
804
"""Run chunk KDA from raw gate and beta projections."""
@@ -815,6 +822,8 @@ def chunk_kda_with_fused_gate(
822
output_final_state=output_final_state,
823
lower_bound=lower_bound,
824
cu_seqlens=cu_seqlens,
825
+ chunk_indices=chunk_indices,
826
+ chunk_offsets=chunk_offsets,
827
)
828
return o, final_state
829
vllm/v1/attention/backends/gdn_attn.py
+47
-36
@@ -165,6 +165,48 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
165
device=device,
166
)
167
168
+ def _build_chunk_metadata(
169
+ self,
170
+ prefill_query_start_loc: torch.Tensor,
171
+ prefill_query_start_loc_cpu: torch.Tensor,
172
+ device: torch.device,
173
+ ) -> tuple[torch.Tensor, torch.Tensor]:
174
+ from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE
175
+
176
+ if self.gdn_prefill_backend == "cutedsl":
177
+ from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import (
178
+ prepare_metadata_cutedsl,
179
+ )
180
+
181
+ assert prefill_query_start_loc is not None
182
+ assert prefill_query_start_loc_cpu is not None
183
+ total_tokens = int(prefill_query_start_loc_cpu[-1].item())
184
+ return prepare_metadata_cutedsl(
185
+ prefill_query_start_loc,
186
+ total_tokens,
187
+ FLA_CHUNK_SIZE,
188
+ )
189
+
190
+ # Only prefill batches use FLA chunk ops.
191
+ # Pre-compute on CPU and async-copy to GPU to avoid
192
+ # GPU→CPU sync (.tolist()) in prepare_chunk_indices.
193
+ from vllm.third_party.flash_linear_attention.ops.index import (
194
+ prepare_chunk_indices,
195
+ prepare_chunk_offsets,
196
+ )
197
+
198
+ assert prefill_query_start_loc_cpu is not None
199
+ return (
200
+ async_tensor_h2d(
201
+ prepare_chunk_indices(prefill_query_start_loc_cpu, FLA_CHUNK_SIZE),
202
+ device=device,
203
+ ),
204
+ async_tensor_h2d(
205
+ prepare_chunk_offsets(prefill_query_start_loc_cpu, FLA_CHUNK_SIZE),
206
+ device=device,
207
+ ),
208
+ )
209
+
210
def build( # type: ignore[override]
211
self,
212
common_prefix_len: int,
@@ -330,10 +372,6 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
372
prefill_state_indices: torch.Tensor | None = None
373
prefill_has_initial_state: torch.Tensor | None = None
374
if num_prefills > 0:
333
- from vllm.third_party.flash_linear_attention.ops.utils import (
334
- FLA_CHUNK_SIZE,
335
- )
336
-
375
# In a mixed non-spec batch, decodes are peeled off to the recurrent
376
# kernel (decode-first front slice), so build chunk metadata from the
377
# rebased prefill-only cu_seqlens; otherwise use the full non-spec one.
@@ -354,38 +392,11 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
392
prefill_query_start_loc_cpu = non_spec_query_start_loc_cpu
393
prefill_state_indices = non_spec_state_indices_tensor
394
357
- if self.gdn_prefill_backend == "cutedsl":
358
- from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import (
359
- prepare_metadata_cutedsl,
360
- )
361
-
362
- assert prefill_query_start_loc is not None
363
- assert prefill_query_start_loc_cpu is not None
364
- total_tokens = int(prefill_query_start_loc_cpu[-1].item())
365
- chunk_indices, chunk_offsets = prepare_metadata_cutedsl(
366
- prefill_query_start_loc,
367
- total_tokens,
368
- FLA_CHUNK_SIZE,
369
- )
370
- else:
371
- gpu_device = query_start_loc.device
372
- # Only prefill batches use FLA chunk ops.
373
- # Pre-compute on CPU and async-copy to GPU to avoid
374
- # GPU→CPU sync (.tolist()) in prepare_chunk_indices.
375
- from vllm.third_party.flash_linear_attention.ops.index import (
376
- prepare_chunk_indices,
377
- prepare_chunk_offsets,
378
- )
379
-
380
- assert prefill_query_start_loc_cpu is not None
381
- chunk_indices = async_tensor_h2d(
382
- prepare_chunk_indices(prefill_query_start_loc_cpu, FLA_CHUNK_SIZE),
383
- device=gpu_device,
384
- )
385
- chunk_offsets = async_tensor_h2d(
386
- prepare_chunk_offsets(prefill_query_start_loc_cpu, FLA_CHUNK_SIZE),
387
- device=gpu_device,
388
- )
395
+ chunk_indices, chunk_offsets = self._build_chunk_metadata(
396
+ prefill_query_start_loc,
397
+ prefill_query_start_loc_cpu,
398
+ query_start_loc.device,
399
+ )
400
401
if num_prefills > 0:
402
context_lens_tensor = m.compute_num_computed_tokens()