[Kimi K3 Bug] Fix kimi k3 KeyError: 'language_model.model.layers.0.self_attn'

[yewentao256]

Purpose

This bug is introduced by #53306

It skips Mamba-family attention metadata during profile/warmup dummy runs. This PR fixes the issue

Test

vllm serve moonshotai/Kimi-K3 \
  --trust-remote-code \
  -tp 8 \
  --load-format fastsafetensors \
  --enable-prefix-caching \
  --reasoning-parser kimi_k3 \
  --enable-auto-tool-choice \
  --tool-call-parser kimi_k3 \
  --host 0.0.0.0 \
  --port 30000

Main

(Worker_TP5 pid=3131903) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]     attn_metadata_narrowed = attn_metadata_raw[self.prefix]
(Worker_TP5 pid=3131903) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]                              ~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^
(Worker_TP3 pid=3131901) ERROR 08-24 13:18:52 [multiproc_executor.py:1050] KeyError: 'language_model.model.layers.0.self_attn'
(Worker_TP5 pid=3131903) ERROR 08-24 13:18:52 [multiproc_executor.py:1050] KeyError: 'language_model.model.layers.0.self_attn'
(Worker_TP3 pid=3131901) ERROR 08-24 13:18:52 [multiproc_executor.py:1050] 
(Worker_TP5 pid=3131903) ERROR 08-24 13:18:52 [multiproc_executor.py:1050] 
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]   File "/home/yewentao256/vllm-source/vllm/models/kimi_k3/nvidia/model.py", line 1367, in forward
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]     hidden_states, prefix_sum, residual = layer(
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]                                           ^^^^^^
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]   File "/home/yewentao256/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]     return self._call_impl(*args, **kwargs)
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]   File "/home/yewentao256/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]     return forward_call(*args, **kwargs)
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]   File "/home/yewentao256/vllm-source/vllm/models/kimi_k3/nvidia/model.py", line 1101, in forward
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]     hidden_states = self._run_self_attn(positions, hidden_states)
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]   File "/home/yewentao256/vllm-source/vllm/models/kimi_k3/nvidia/model.py", line 1011, in _run_self_attn
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]     return self.self_attn(
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]            ^^^^^^^^^^^^^^^
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]   File "/home/yewentao256/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]     return self._call_impl(*args, **kwargs)
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]   File "/home/yewentao256/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]     return forward_call(*args, **kwargs)
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]   File "/home/yewentao256/vllm-source/vllm/models/kimi_k3/nvidia/kda.py", line 628, in forward
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]     self._forward(
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]   File "/home/yewentao256/vllm-source/vllm/compilation/breakable_cudagraph.py", line 97, in wrapper
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]     return fn(*args, **kwargs)
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]            ^^^^^^^^^^^^^^^^^^^
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]   File "/home/yewentao256/vllm-source/vllm/models/kimi_k3/nvidia/kda.py", line 661, in _forward
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]     attn_metadata_narrowed = attn_metadata_raw[self.prefix]
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]                              ~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050] KeyError: 'language_model.model.layers.0.self_attn'
(Worker_TP1 pid=3131899) ERROR 08-24 13:18:52 [multiproc_executor.py:1050]

Now

(APIServer pid=3218033) INFO:     Started server process [3218033]
(APIServer pid=3218033) INFO:     Waiting for application startup.
(APIServer pid=3218033) INFO:     Application startup complete.

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

[yewentao256]

/ci run

[yewentao256]

Also CC @njhill

[github-actions]

✅ Triggered Buildkite CI #85324 for commit c45f90d0e060.

[njhill]

Thanks @yewentao256! yes my mistake 🙈

There is already a PR to fix though #53581

[yewentao256]

Close this PR as duplicated