Skip to content

[ROCm][Kimi-K3] Don't read W_UK_T before weight loading in the MLA wrapper - #61654

Merged
AndreasKaratzas merged 3 commits into
vllm-project:mainfrom
Rohan138:fix-kimi-mla-w-uk-t
Oct 14, 2026
Merged

AndreasKaratzas merged 3 commits into
vllm-project:mainfrom
Rohan138:fix-kimi-mla-w-uk-t

Conversation

@Rohan138

@Rohan138 Rohan138 commented Oct 14, 2026 •

Copy link
Copy Markdown
Collaborator

Overview

Fixes a Kimi-K3 startup crash on ROCm ('MLAAttention' object has no attribute 'W_UK_T') that appears once #53123 turns AITER FP8 BMM off by default.

Claims

  • KimiK3MultiHeadLatentAttentionWrapper._fused_qk_prep_supported() no longer reads attn.W_UK_T in __init__.
  • The fused decode Q-prep path stays enabled when neither AITER BMM is on, instead of being silently disabled.

Validation

FP4 BMM FP8 BMM main this PR
off off AttributeError: W_UK_T True
off on False False
on off True True
  • Not run end to end locally: the 475 GB checkpoint doesn't fit on my box. CI's MI355 Spec Decode job covers it.

Details

_fused_qk_prep_supported() runs in __init__, but MLAAttention.W_UK_T is only created in process_weights_after_loading. Before #53123, FP8 BMM was on by default, so the expression short-circuited and never read it. process_weights_after_loading always creates W_UK_T when FP4 and FP8 BMM are both off, so the check isn't needed. Using getattr(..., None) instead would have turned the fused path off for every such config. No open PR fixes this; #60110 edits _fused_decode but not this gate.


Pull Request Checklist
  • I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

  • AI assistance was used during the creation of this PR.

  • Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

  • Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

  • Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

  • Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

🤖 Generated with Claude Code

…apper

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

@mergify mergify Bot added kimi k3 rocm Related to AMD ROCm labels Oct 14, 2026
@github-project-automation github-project-automation Bot moved this to Todo in AMD Oct 14, 2026
@Rohan138 Rohan138 added the ready ONLY add when PR is ready to merge/full CI is needed label Oct 14, 2026
@Rohan138

Copy link
Copy Markdown
Collaborator Author

/ci run

@github-actions

Copy link
Copy Markdown

✅ Triggered Buildkite CI #94873 for commit 599cd5d955f5.

@Rohan138

Copy link
Copy Markdown
Collaborator Author

/ci run

@Rohan138

Copy link
Copy Markdown
Collaborator Author

merged main and rerunning CI to pick up fix from #61650

@github-actions

Copy link
Copy Markdown

✅ Triggered Buildkite CI #94893 for commit 4a498a86deca.

@AndreasKaratzas
AndreasKaratzas enabled auto-merge (squash) October 14, 2026 18:48
@AndreasKaratzas
AndreasKaratzas merged commit 9a648b4 into vllm-project:main Oct 14, 2026
95 checks passed
@github-project-automation github-project-automation Bot moved this from Todo to Done in AMD Oct 14, 2026
AndreasKaratzas pushed a commit to AndreasKaratzas/vllm that referenced this pull request Oct 16, 2026
…apper (vllm-project#61654)

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Signed-off-by: Andreas Karatzas <AndreasKaratzas@users.noreply.github.com>
michael-gpu pushed a commit to practical-inference/vllm-ampere-extended that referenced this pull request Oct 20, 2026
…apper (vllm-project#61654)

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

k3 kimi ready ONLY add when PR is ready to merge/full CI is needed rocm Related to AMD ROCm

Projects

Status: Done

Development

Successfully merging this pull request may close these issues.

2 participants