[Bugfix][Kimi K3] Skip absent metadata during CUDA graph profiling (#53581)
Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: Codex <codex@openai.com>
Kevin H. Luu committed
Aug 24, 2026 at 08:16 UTC
4c56e62c85cea8fc2251efc25159836c214402aa
3 files changed
+20
-2
tests/models/kimi_k3/test_kda.py
+14
@@ -20,6 +20,7 @@ from vllm.model_executor.layers.mamba.ops.gather_initial_states import (
20
from vllm.models.kimi_k3.amd.ops.third_party.kda import (
21
fused_recurrent_kda_packed_decode as fused_recurrent_kda_packed_decode_amd,
22
)
23
+from vllm.models.kimi_k3.nvidia import kda as nvidia_kda
24
from vllm.models.kimi_k3.nvidia.kda import (
25
_flashkda_prefill,
26
_store_cache_checkpoints_kernel,
@@ -59,6 +60,19 @@ PACKED_DECODE_IMPLS = {
60
}
61
62
63
+def test_kda_warmup_skips_missing_metadata(monkeypatch):
64
+ monkeypatch.setattr(
65
+ nvidia_kda,
66
+ "get_forward_context",
67
+ lambda: SimpleNamespace(attn_metadata={}),
68
+ )
69
+ layer = object.__new__(nvidia_kda.KimiK3DeltaAttention)
70
+ object.__setattr__(layer, "prefix", "language_model.model.layers.0.self_attn")
71
+ empty = torch.empty(0, device=DEVICE)
72
+
73
+ assert layer._forward(empty, empty, empty, empty, empty) is None
74
+
75
+
76
def test_kda_recoverssm_config_state_layout():
77
vllm_config = SimpleNamespace(
78
model_config=SimpleNamespace(
vllm/models/kimi_k3/amd/kda.py
+3
-1
@@ -313,7 +313,9 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
313
return
314
315
assert isinstance(attn_metadata_raw, dict)
316
- attn_metadata_narrowed = attn_metadata_raw[self.prefix]
316
+ attn_metadata_narrowed = attn_metadata_raw.get(self.prefix)
317
+ if attn_metadata_narrowed is None:
318
+ return
319
assert isinstance(attn_metadata_narrowed, GDNAttentionMetadata)
320
m = attn_metadata_narrowed
321
has_initial_state = m.has_initial_state
vllm/models/kimi_k3/nvidia/kda.py
+3
-1
@@ -658,7 +658,9 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
658
)
659
660
assert isinstance(attn_metadata_raw, dict)
661
- attn_metadata_narrowed = attn_metadata_raw[self.prefix]
661
+ attn_metadata_narrowed = attn_metadata_raw.get(self.prefix)
662
+ if attn_metadata_narrowed is None:
663
+ return
664
assert isinstance(attn_metadata_narrowed, KimiK3KDAMetadata)
665
m = attn_metadata_narrowed
666
has_initial_state = m.has_initial_state