@tduncan / Nvidia-Optimization-Stack-VLLM / commits / f5cd862dbd

[ROCm][Bugfix] Kimi-K3 Fix KDA NaN on mixed batches and racy autotune config (#50649)

Signed-off-by: kliuae <kuanfu.liu@embeddedllm.com>

kliuae committed Aug 5, 2026 at 16:07 UTC f5cd862dbdafb9d9870945a4f61454f90d87e2ae
3 files changed +556 -7
vllm/models/kimi_k3/amd/kda.py new
+530
@@ -0,0 +1,530 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +from einops import rearrange
6 +from torch import nn
7 +
8 +from vllm.compilation.breakable_cudagraph import eager_break_during_capture
9 +from vllm.config import VllmConfig
10 +from vllm.distributed import divide
11 +from vllm.forward_context import get_forward_context
12 +from vllm.model_executor.layers.linear import (
13 + ColumnParallelLinear,
14 + RowParallelLinear,
15 +)
16 +from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention
17 +
18 +# Generic KDA helpers, shared with Kimi-Linear. They are neither ROCm- nor
19 +# K3-specific, so they are imported rather than duplicated. (nvidia/kda.py keeps
20 +# its own copies; duplicating here instead would be a one-line change.)
21 +from vllm.model_executor.layers.mamba.gdn.kimi_gdn_linear_attn import (
22 + _KDA_GATE_LOGBOUND_MIN,
23 + _KimiGDNMergedColumnParallelLinear,
24 + _make_fused_conv1d_weight_loader,
25 + a_log_weight_loader,
26 +)
27 +from vllm.model_executor.layers.mamba.mamba_utils import (
28 + MambaStateDtypeCalculator,
29 + MambaStateShapeCalculator,
30 + is_conv_state_dim_first,
31 +)
32 +from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
33 + causal_conv1d_fn,
34 + causal_conv1d_update,
35 +)
36 +from vllm.model_executor.layers.mamba.ops.gather_initial_states import (
37 + gather_initial_states,
38 +)
39 +from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader
40 +from vllm.model_executor.utils import set_weight_attrs
41 +from vllm.models.kimi_k3.amd.ops.third_party.kda import (
42 + chunk_kda_with_fused_gate,
43 + fused_recurrent_kda,
44 + fused_recurrent_kda_packed_decode,
45 +)
46 +from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated
47 +from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
48 +from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
49 +
50 +
51 +class KimiK3DeltaAttention(GatedDeltaNetAttention):
52 + def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]:
53 + if self.model_config is None or self.cache_config is None:
54 + raise ValueError("model_config and cache_config must be set")
55 + return MambaStateDtypeCalculator.kda_state_dtype(
56 + self.model_config.dtype, self.cache_config.mamba_cache_dtype
57 + )
58 +
59 + def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]:
60 + return MambaStateShapeCalculator.kda_state_shape(
61 + self.tp_size,
62 + self.num_heads,
63 + self.head_dim,
64 + conv_kernel_size=self.conv_size,
65 + num_spec=self.num_spec,
66 + )
67 +
68 + def __init__(
69 + self,
70 + config: KimiLinearConfig,
71 + vllm_config: VllmConfig,
72 + prefix: str = "",
73 + ) -> None:
74 + super().__init__(config, vllm_config, prefix)
75 +
76 + kda_config = config.linear_attn_config # type: ignore[attr-defined]
77 + assert kda_config is not None, "linear_attn_config must be set"
78 + assert kda_config.get("use_full_rank_gate", False), (
79 + "KimiK3DeltaAttention requires use_full_rank_gate; the low-rank "
80 + "gate path belongs to the shared Kimi-Linear layer."
81 + )
82 +
83 + self.head_dim = kda_config["head_dim"]
84 + self.num_heads = kda_config["num_heads"]
85 + assert self.num_heads % self.tp_size == 0
86 + self.local_num_heads = divide(self.num_heads, self.tp_size)
87 +
88 + self.projection_size = self.head_dim * self.num_heads
89 + self.local_projection_size = divide(self.projection_size, self.tp_size)
90 + self.conv_size = kda_config["short_conv_kernel_size"]
91 + self.use_full_rank_gate = True
92 +
93 + # Keep f_a before the narrow beta shard, then pad each TP-local row to
94 + # select the aligned BF16 GEMM path. The padding also avoids an Inductor
95 + # correctness issue seen with the row-strided G view.
96 + qkvg_output_sizes = [self.projection_size] * 4
97 + in_proj_output_sizes = qkvg_output_sizes + [self.head_dim, self.num_heads]
98 + local_output_size = (
99 + 4 * self.local_projection_size + self.head_dim + self.local_num_heads
100 + )
101 + self.in_proj_padding = -local_output_size % 16
102 + if self.in_proj_padding:
103 + in_proj_output_sizes.append(self.in_proj_padding * self.tp_size)
104 +
105 + self.in_proj_qkvgfab = _KimiGDNMergedColumnParallelLinear(
106 + self.hidden_size,
107 + in_proj_output_sizes,
108 + replicated_shard_id=4,
109 + tp_size=self.tp_size,
110 + bias=False,
111 + quant_config=self.quant_config,
112 + prefix=f"{prefix}.in_proj_qkvgfab",
113 + )
114 + if self.in_proj_padding:
115 + self.in_proj_qkvgfab.weight.data[-self.in_proj_padding :].zero_()
116 +
117 + self.f_b_proj = ColumnParallelLinear(
118 + self.head_dim,
119 + self.projection_size,
120 + bias=False,
121 + quant_config=self.quant_config,
122 + prefix=f"{prefix}.f_b_proj",
123 + )
124 + self.dt_bias = nn.Parameter(
125 + torch.empty(self.local_projection_size, dtype=torch.float32)
126 + )
127 + set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
128 +
129 + # One packed parameter and cache let decode run a single conv update.
130 + # Prefill slices them back into Q/K/V to obtain dense outputs cheaply.
131 + self.conv1d = ColumnParallelLinear(
132 + input_size=self.conv_size,
133 + output_size=3 * self.projection_size,
134 + bias=False,
135 + params_dtype=torch.float32,
136 + prefix=f"{prefix}.conv1d",
137 + )
138 + self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
139 + delattr(self.conv1d.weight, "weight_loader")
140 + set_weight_attrs(
141 + self.conv1d.weight,
142 + {
143 + "weight_loader": _make_fused_conv1d_weight_loader(
144 + [self.projection_size] * 3,
145 + self.tp_size,
146 + self.tp_rank,
147 + )
148 + },
149 + )
150 +
151 + self.A_log = nn.Parameter(
152 + torch.empty(self.local_num_heads, dtype=torch.float32)
153 + )
154 + set_weight_attrs(self.A_log, {"weight_loader": a_log_weight_loader(0)})
155 +
156 + self.gate_lower_bound: float | None = kda_config.get("gate_lower_bound", None)
157 + if self.gate_lower_bound is not None:
158 + assert _KDA_GATE_LOGBOUND_MIN <= self.gate_lower_bound < 0, (
159 + "KDA gate lower bound must be in "
160 + f"[{_KDA_GATE_LOGBOUND_MIN}, 0). "
161 + f"Got {self.gate_lower_bound}."
162 + )
163 + self.use_safe_gate = self.gate_lower_bound is not None
164 +
165 + additional_config = vllm_config.additional_config
166 + backend = (
167 + additional_config.get("kda_prefill_backend", "auto")
168 + if isinstance(additional_config, dict)
169 + else "auto"
170 + )
171 + backend = "triton" if backend == "auto" else backend
172 + assert backend == "triton", (
173 + "The ROCm Kimi-K3 KDA layer only supports the Triton KDA prefill "
174 + f"backend, got {backend!r}."
175 + )
176 +
177 + self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid")
178 + self.o_proj = RowParallelLinear(
179 + self.projection_size,
180 + self.hidden_size,
181 + bias=False,
182 + quant_config=self.quant_config,
183 + prefix=f"{prefix}.o_proj",
184 + )
185 +
186 + compilation_config = vllm_config.compilation_config
187 + if prefix in compilation_config.static_forward_context:
188 + raise ValueError(f"Duplicate layer name: {prefix}")
189 + compilation_config.static_forward_context[prefix] = self
190 +
191 + def rearrange_mixed_qkv(
192 + self, mixed_qkv: torch.Tensor
193 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
194 + seq_len = mixed_qkv.shape[0]
195 + qkv = mixed_qkv.view(seq_len, 3, self.local_num_heads, self.head_dim)
196 + # Materialize all three row-strided inputs with one token-major to
197 + # QKV-major permutation. Each unbound tensor is then contiguous.
198 + qkv = qkv.permute(1, 0, 2, 3).contiguous().unsqueeze(1)
199 + return qkv.unbind(0)
200 +
201 + def forward(
202 + self,
203 + hidden_states: torch.Tensor,
204 + positions: torch.Tensor,
205 + output: torch.Tensor,
206 + ) -> None:
207 + num_tokens = hidden_states.size(0)
208 + projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0]
209 +
210 + split_sizes = [
211 + 3 * self.local_projection_size,
212 + self.local_projection_size,
213 + self.head_dim,
214 + self.local_num_heads,
215 + ]
216 + if self.in_proj_padding:
217 + split_sizes.append(self.in_proj_padding)
218 + projected = projected_qkvgfab.split(split_sizes, dim=-1)
219 + mixed_qkv, g_proj_states, f_a, beta = projected[:4]
220 +
221 + g1 = self.f_b_proj(f_a)[0]
222 + beta = beta.unsqueeze(0)
223 + g1 = rearrange(g1, "n (h d) -> 1 n h d", d=self.head_dim)
224 + g2 = rearrange(g_proj_states, "... (h d) -> ... h d", d=self.head_dim)
225 +
226 + core_attn_out = torch.empty(
227 + (1, num_tokens, self.local_num_heads, self.head_dim),
228 + dtype=hidden_states.dtype,
229 + device=hidden_states.device,
230 + )
231 +
232 + self._forward(
233 + mixed_qkv=mixed_qkv,
234 + g1=g1,
235 + g2=g2,
236 + beta=beta,
237 + core_attn_out=core_attn_out,
238 + )
239 + core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)")
240 + output[:] = self.o_proj(core_attn_out)[0]
241 +
242 + @eager_break_during_capture
243 + def _forward(
244 + self,
245 + mixed_qkv: torch.Tensor,
246 + g1: torch.Tensor,
247 + g2: torch.Tensor,
248 + beta: torch.Tensor,
249 + core_attn_out: torch.Tensor,
250 + ) -> None:
251 + forward_context = get_forward_context()
252 + attn_metadata_raw = forward_context.attn_metadata
253 +
254 + if attn_metadata_raw is None:
255 + return
256 +
257 + assert isinstance(attn_metadata_raw, dict)
258 + attn_metadata_narrowed = attn_metadata_raw[self.prefix]
259 + assert isinstance(attn_metadata_narrowed, GDNAttentionMetadata)
260 + m = attn_metadata_narrowed
261 + has_initial_state = m.has_initial_state
262 + non_spec_query_start_loc = m.non_spec_query_start_loc
263 + non_spec_state_indices_tensor = m.non_spec_state_indices_tensor
264 + spec_sequence_masks = m.spec_sequence_masks
265 + spec_token_indx = m.spec_token_indx
266 + non_spec_token_indx = m.non_spec_token_indx
267 + spec_state_indices_tensor = m.spec_state_indices_tensor
268 + spec_query_start_loc = m.spec_query_start_loc
269 + num_accepted_tokens = m.num_accepted_tokens
270 + num_actual_tokens = m.num_actual_tokens
271 + mixed_qkv = mixed_qkv[:num_actual_tokens]
272 + g1 = g1[:, :num_actual_tokens]
273 + beta = beta[:, :num_actual_tokens]
274 +
275 + constant_caches = self.kv_cache
276 +
277 + conv_state, recurrent_state = constant_caches
278 + # conv_state must be (..., dim, width-1) for the conv kernels.
279 + # DS layout stores it that way directly; SD layout needs a transpose.
280 + if not is_conv_state_dim_first():
281 + conv_state = conv_state.transpose(-1, -2)
282 +
283 + conv_weights = self.conv1d.weight.view(
284 + self.conv1d.weight.size(0), self.conv1d.weight.size(2)
285 + )
286 + q_conv_weight, k_conv_weight, v_conv_weight = conv_weights.split(
287 + self.local_projection_size, dim=0
288 + )
289 + q_conv_state, k_conv_state, v_conv_state = conv_state.split(
290 + self.local_projection_size, dim=-2
291 + )
292 +
293 + # Split tokens into the multi-query spec-decode part and the remaining
294 + # (prefill / plain decode) part.
295 + if spec_sequence_masks is not None:
296 + if m.num_prefills == 0 and m.num_decodes == 0:
297 + mixed_qkv_spec = mixed_qkv
298 + g1_spec, beta_spec = g1, beta
299 + mixed_qkv_ns = g1_ns = beta_ns = None
300 + else:
301 + mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx)
302 + g1_spec = g1.index_select(1, spec_token_indx)
303 + beta_spec = beta.index_select(1, spec_token_indx)
304 + mixed_qkv_ns = mixed_qkv.index_select(0, non_spec_token_indx)
305 + g1_ns = g1.index_select(1, non_spec_token_indx)
306 + beta_ns = beta.index_select(1, non_spec_token_indx)
307 + else:
308 + mixed_qkv_spec = g1_spec = beta_spec = None
309 + mixed_qkv_ns, g1_ns, beta_ns = mixed_qkv, g1, beta
310 +
311 + # ---------- spec-decode multi-query path ----------
312 + core_attn_out_spec = None
313 + if spec_sequence_masks is not None:
314 + assert spec_state_indices_tensor is not None
315 + assert spec_query_start_loc is not None
316 + spec_conv_indices = spec_state_indices_tensor[:, 0][: m.num_spec_decodes]
317 + spec_max_query_len = spec_state_indices_tensor.size(-1)
318 +
319 + # Sibling beta and, for full-rank gates, output-gate views remain
320 + # live, so write the convolution output separately.
321 + spec_conv_out = torch.empty(
322 + mixed_qkv_spec.shape,
323 + dtype=mixed_qkv_spec.dtype,
324 + device=mixed_qkv_spec.device,
325 + )
326 + mixed_qkv_spec = causal_conv1d_update(
327 + mixed_qkv_spec,
328 + conv_state,
329 + conv_weights,
330 + self.conv1d.bias,
331 + activation="silu",
332 + conv_state_indices=spec_conv_indices,
333 + num_accepted_tokens=num_accepted_tokens,
334 + query_start_loc=spec_query_start_loc,
335 + max_query_len=spec_max_query_len,
336 + validate_data=False,
337 + out=spec_conv_out,
338 + )
339 + q_spec, k_spec, v_spec = (
340 + rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim)
341 + for x in mixed_qkv_spec.split(self.local_projection_size, dim=-1)
342 + )
343 + spec_cu_seqlens = spec_query_start_loc[: m.num_spec_decodes + 1]
344 + # Spec-only batches write directly into core_attn_out.
345 + spec_out = (
346 + core_attn_out[:, : q_spec.shape[1]]
347 + if m.num_prefills == 0 and m.num_decodes == 0
348 + else None
349 + )
350 + core_attn_out_spec, _ = fused_recurrent_kda(
351 + q=q_spec,
352 + k=k_spec,
353 + v=v_spec,
354 + raw_g=g1_spec,
355 + raw_beta=beta_spec,
356 + A_log=self.A_log,
357 + dt_bias=self.dt_bias,
358 + lower_bound=self.gate_lower_bound,
359 + initial_state=recurrent_state,
360 + cu_seqlens=spec_cu_seqlens,
361 + ssm_state_indices=spec_state_indices_tensor,
362 + num_accepted_tokens=num_accepted_tokens,
363 + out=spec_out,
364 + )
365 +
366 + # ---------- non-spec path (prefill or plain decode) ----------
367 + core_attn_out_non_spec = None
368 + if mixed_qkv_ns is not None:
369 + assert g1_ns is not None and beta_ns is not None
370 + if m.num_prefills > 0:
371 + q_ns, k_ns, v_ns = mixed_qkv_ns.split(
372 + self.local_projection_size, dim=-1
373 + )
374 +
375 + # Packed prefill conv would require copying V solely to make
376 + # it dense for KDA. Separate calls accept the strided inputs
377 + # and produce dense Q/K/V without that extra traffic.
378 + # TODO: Use packed conv once every KDA prefill backend accepts
379 + # row-strided Q/K/V directly.
380 + def _prefill_conv(
381 + x: torch.Tensor,
382 + state: torch.Tensor,
383 + weight: torch.Tensor,
384 + ) -> torch.Tensor:
385 + return causal_conv1d_fn(
386 + x.transpose(0, 1),
387 + weight,
388 + None,
389 + activation="silu",
390 + conv_states=state,
391 + has_initial_state=has_initial_state,
392 + cache_indices=non_spec_state_indices_tensor,
393 + query_start_loc=non_spec_query_start_loc,
394 + metadata=m,
395 + ).transpose(0, 1)
396 +
397 + q_ns = _prefill_conv(q_ns, q_conv_state, q_conv_weight)
398 + k_ns = _prefill_conv(k_ns, k_conv_state, k_conv_weight)
399 + v_ns = _prefill_conv(v_ns, v_conv_state, v_conv_weight)
400 + q_ns, k_ns, v_ns = (
401 + rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim)
402 + for x in (q_ns, k_ns, v_ns)
403 + )
404 +
405 + assert non_spec_state_indices_tensor is not None
406 + assert has_initial_state is not None
407 +
408 + # A mixed non-spec batch is decode-first: the decodes are
409 + # length-1 sequences, and the chunk kernel returns NaN for those.
410 + # Send them to the recurrent kernel and give the chunk kernel the
411 + # prefill tail only.
412 + core_attn_out_decode = None
413 + split_non_spec = spec_sequence_masks is None and m.num_decodes > 0
414 + if split_non_spec:
415 + assert non_spec_query_start_loc is not None
416 + nd_tok = m.num_decode_tokens
417 + core_attn_out_decode, _ = fused_recurrent_kda(
418 + q=q_ns[:, :nd_tok],
419 + k=k_ns[:, :nd_tok],
420 + v=v_ns[:, :nd_tok],
421 + raw_g=g1_ns[:, :nd_tok],
422 + raw_beta=beta_ns[:, :nd_tok],
423 + A_log=self.A_log,
424 + dt_bias=self.dt_bias,
425 + lower_bound=self.gate_lower_bound,
426 + initial_state=recurrent_state,
427 + cu_seqlens=non_spec_query_start_loc[: m.num_decodes + 1],
428 + ssm_state_indices=non_spec_state_indices_tensor[
429 + : m.num_decodes
430 + ],
431 + )
432 + q_ns = q_ns[:, nd_tok:]
433 + k_ns = k_ns[:, nd_tok:]
434 + v_ns = v_ns[:, nd_tok:]
435 + g1_ns = g1_ns[:, nd_tok:]
436 + beta_ns = beta_ns[:, nd_tok:]
437 + prefill_query_start_loc = m.prefill_query_start_loc
438 + prefill_state_indices = m.prefill_state_indices
439 + prefill_has_initial_state = m.prefill_has_initial_state
440 + assert prefill_query_start_loc is not None
441 + assert prefill_state_indices is not None
442 + assert prefill_has_initial_state is not None
443 + else:
444 + prefill_query_start_loc = non_spec_query_start_loc
445 + prefill_state_indices = non_spec_state_indices_tensor
446 + prefill_has_initial_state = has_initial_state
447 +
448 + initial_state = gather_initial_states(
449 + recurrent_state,
450 + prefill_state_indices,
451 + prefill_has_initial_state,
452 + )
453 + (
454 + core_attn_out_non_spec,
455 + last_recurrent_state,
456 + ) = chunk_kda_with_fused_gate(
457 + q=q_ns,
458 + k=k_ns,
459 + v=v_ns,
460 + raw_g=g1_ns,
461 + raw_beta=beta_ns,
462 + A_log=self.A_log,
463 + g_bias=self.dt_bias,
464 + lower_bound=self.gate_lower_bound,
465 + initial_state=initial_state,
466 + output_final_state=True,
467 + use_qk_l2norm_in_kernel=True,
468 + cu_seqlens=prefill_query_start_loc,
469 + )
470 + # Init cache
471 + recurrent_state[prefill_state_indices] = last_recurrent_state
472 +
473 + if split_non_spec:
474 + # Restore decode-first token order for the merge below.
475 + core_attn_out_non_spec = torch.cat(
476 + [core_attn_out_decode, core_attn_out_non_spec], dim=1
477 + )
478 +
479 + else:
480 + # pure-decode non-spec batch
481 + assert non_spec_state_indices_tensor is not None
482 + decode_conv_indices = non_spec_state_indices_tensor[
483 + : mixed_qkv_ns.size(0)
484 + ]
485 + # Sibling beta and, for full-rank gates, output-gate views
486 + # remain live, so write the conv output separately.
487 + packed_conv_out = torch.empty(
488 + mixed_qkv_ns.shape,
489 + dtype=mixed_qkv_ns.dtype,
490 + device=mixed_qkv_ns.device,
491 + )
492 + mixed_qkv_ns = causal_conv1d_update(
493 + mixed_qkv_ns,
494 + conv_state,
495 + conv_weights,
496 + self.conv1d.bias,
497 + activation="silu",
498 + conv_state_indices=decode_conv_indices,
499 + validate_data=True,
500 + out=packed_conv_out,
501 + )
502 + core_attn_out_non_spec, _ = fused_recurrent_kda_packed_decode(
503 + mixed_qkv=mixed_qkv_ns,
504 + raw_g=g1_ns,
505 + raw_beta=beta_ns,
506 + A_log=self.A_log,
507 + dt_bias=self.dt_bias,
508 + lower_bound=self.gate_lower_bound,
509 + initial_state=recurrent_state,
510 + state_indices=decode_conv_indices,
511 + )
512 +
513 + # ---------- merge spec and non-spec outputs ----------
514 + if core_attn_out_spec is not None and core_attn_out_non_spec is not None:
515 + # Mixed batches require indexed placement in the original order.
516 + merged = torch.empty(
517 + (1, num_actual_tokens, *core_attn_out_spec.shape[2:]),
518 + dtype=core_attn_out_spec.dtype,
519 + device=core_attn_out_spec.device,
520 + )
521 + merged.index_copy_(1, spec_token_indx, core_attn_out_spec)
522 + merged.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)
523 + core_attn_out[0, :num_actual_tokens] = merged[0, :num_actual_tokens]
524 + elif core_attn_out_non_spec is not None:
525 + core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[
526 + 0, :num_actual_tokens
527 + ]
528 + else:
529 + assert core_attn_out_spec is not None
530 + core_attn_out.copy_(self.o_norm(core_attn_out, g2))
vllm/models/kimi_k3/amd/linear.py
+18 -6
@@ -28,7 +28,7 @@ from vllm.model_executor.layers.linear import (
28 )
29 from vllm.model_executor.layers.logits_processor import LogitsProcessor
30 from vllm.model_executor.layers.mamba.gdn.kimi_gdn_linear_attn import (
31 - KimiGatedDeltaNetAttention,
31 + KimiGatedDeltaNetAttention as KimiLinearGatedDeltaNetAttention,
32 )
33 from vllm.model_executor.layers.mamba.mamba_utils import (
34 MambaStateCopyFunc,
@@ -61,6 +61,7 @@ from vllm.model_executor.models.utils import (
61 make_layers,
62 maybe_prefix,
63 )
64 +from vllm.models.kimi_k3.amd.kda import KimiK3DeltaAttention
65 from vllm.models.kimi_k3.amd.ops.attn_res import attn_res
66 from vllm.sequence import IntermediateTensors
67 from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
@@ -477,11 +478,22 @@ class KimiDecoderLayer(nn.Module):
478 quant_config = vllm_config.quant_config
479
480 if config.is_kda_layer(layer_idx):
480 - self.self_attn = KimiGatedDeltaNetAttention(
481 - config,
482 - vllm_config,
483 - prefix=f"{prefix}.self_attn",
484 - )
481 + # Kimi-K3 sets use_full_rank_gate and uses the ROCm-specific K3 KDA
482 + # layer; Kimi-Linear keeps the shared low-rank-gate implementation.
483 + kda_config = config.linear_attn_config
484 + assert kda_config is not None
485 + if kda_config.get("use_full_rank_gate", False):
486 + self.self_attn = KimiK3DeltaAttention(
487 + config,
488 + vllm_config,
489 + prefix=f"{prefix}.self_attn",
490 + )
491 + else:
492 + self.self_attn = KimiLinearGatedDeltaNetAttention(
493 + config,
494 + vllm_config,
495 + prefix=f"{prefix}.self_attn",
496 + )
497 else:
498 qk_nope_head_dim = config.qk_nope_head_dim
499 qk_rope_head_dim = config.qk_rope_head_dim
vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py
+8 -1
@@ -28,6 +28,13 @@ BT_LIST_AUTOTUNE = [32, 64, 128]
28 NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if is_amd else [4, 8, 16, 32]
29
30
31 +# num_stages=4 is excluded. The `u` loop has a trip count of 2 (V / BV), and a
32 +# 4-stage pipeline over it emits a third, consumer-less async copy into LDS that
33 +# the `w` loop then reuses; at num_warps=4 that races on gfx950 and `u` comes
34 +# back with non-deterministic O(1e38) garbage once a batch reaches 4096 tokens.
35 +_RECOMPUTE_W_U_NUM_STAGES = [2, 3]
36 +
37 +
38 @triton.heuristics(
39 {
40 "STORE_QG": lambda args: args["qg"] is not None,
@@ -39,7 +46,7 @@ NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if is_amd else [4, 8, 16, 32]
46 configs=[
47 triton.Config({}, num_warps=num_warps, num_stages=num_stages)
48 for num_warps in [2, 4, 8]
42 - for num_stages in [2, 3, 4]
49 + for num_stages in _RECOMPUTE_W_U_NUM_STAGES
50 ],
51 key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"],
52 )