A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 13, 08:15 UTC
[Bugfix] fix kimi-linear crash (#28445)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Jiangyun Zhu committed
Nov 13, 2025 at 15:59 UTC
fa183e92713456dec682088a362dd9908100cc03
1 file changed
+12
-9
vllm/model_executor/layers/kda.py
+12
-9
@@ -44,7 +44,6 @@ def kda_attention(
44
k_proj_states: torch.Tensor,
45
v_proj_states: torch.Tensor,
46
g1: torch.Tensor,
47
- g2: torch.Tensor,
47
beta: torch.Tensor,
48
core_attn_out: torch.Tensor,
49
layer_name: str,
@@ -56,7 +55,6 @@ def kda_attention(
55
k_proj_states=k_proj_states,
56
v_proj_states=v_proj_states,
57
g1=g1,
59
- g2=g2,
58
beta=beta,
59
core_attn_out=core_attn_out,
60
)
@@ -67,7 +65,6 @@ def kda_attention_fake(
65
k_proj_states: torch.Tensor,
66
v_proj_states: torch.Tensor,
67
g1: torch.Tensor,
70
- g2: torch.Tensor,
68
beta: torch.Tensor,
69
core_attn_out: torch.Tensor,
70
layer_name: str,
@@ -284,7 +281,6 @@ class KimiDeltaAttention(nn.Module, MambaBase):
281
k,
282
v,
283
g1,
287
- g2,
284
beta,
285
core_attn_out,
286
self.prefix,
@@ -299,7 +295,6 @@ class KimiDeltaAttention(nn.Module, MambaBase):
295
k_proj_states: torch.Tensor,
296
v_proj_states: torch.Tensor,
297
g1: torch.Tensor,
302
- g2: torch.Tensor,
298
beta: torch.Tensor,
299
core_attn_out: torch.Tensor,
300
) -> None:
@@ -316,8 +311,15 @@ class KimiDeltaAttention(nn.Module, MambaBase):
311
has_initial_state = attn_metadata.has_initial_state
312
non_spec_query_start_loc = attn_metadata.non_spec_query_start_loc
313
non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501
314
+ num_actual_tokens = attn_metadata.num_actual_tokens
315
constant_caches = self.kv_cache[forward_context.virtual_engine]
316
317
+ q_proj_states = q_proj_states[:num_actual_tokens]
318
+ k_proj_states = k_proj_states[:num_actual_tokens]
319
+ v_proj_states = v_proj_states[:num_actual_tokens]
320
+ g1 = g1[:num_actual_tokens]
321
+ beta = beta[:num_actual_tokens]
322
+
323
(conv_state_q, conv_state_k, conv_state_v, recurrent_state) = constant_caches
324
# deal with strides
325
conv_state_q = conv_state_q.transpose(-1, -2)
@@ -372,7 +374,7 @@ class KimiDeltaAttention(nn.Module, MambaBase):
374
).transpose(0, 1)
375
else:
376
decode_conv_indices = non_spec_state_indices_tensor[
375
- : attn_metadata.num_decodes
377
+ : attn_metadata.num_actual_tokens
378
]
379
q = causal_conv1d_update(
380
q_proj_states,
@@ -438,8 +440,9 @@ class KimiDeltaAttention(nn.Module, MambaBase):
440
beta=beta,
441
initial_state=recurrent_state,
442
use_qk_l2norm_in_kernel=True,
441
- cu_seqlens=non_spec_query_start_loc,
443
+ cu_seqlens=non_spec_query_start_loc[: attn_metadata.num_decodes + 1],
444
ssm_state_indices=non_spec_state_indices_tensor,
445
)
444
- assert core_attn_out_non_spec.shape == core_attn_out.shape
445
- core_attn_out[:] = core_attn_out_non_spec
446
+ core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[
447
+ 0, :num_actual_tokens
448
+ ]