[whx-sjtu]

Overview

Fix ROCm Kimi K3 online weight loading so its padded merged KDA projection finishes as soon as all checkpoint shards arrive. Following review, the fix lives in an AMD-local _PaddedKimiGDNMergedColumnParallelLinear, mirroring the NVIDIA implementation; no shared loader code is changed.

Summary

six KDA checkpoint shards arrive
  → complete using real elements, excluding alignment rows (weight_loader_numel)
  → layerwise processing materializes the weight and replays the loads
  → loader re-zeroes the padding rows
  → process the layer and release its input tensors

Claims

• Avoid accumulating full GPU checkpoint tensors across incomplete KDA layers during online quantization.
• Keep the alignment padding zero after online materialization, which allocates with empty_strided and drops the constructor's zero-init.
• Only vllm/models/kimi_k3/amd/kda.py changes; the shared Kimi-Linear projection and reload/meta.py are untouched.

Validation

Evidence

Local regression on ROCm (vllm/vllm-openai-rocm:nightly container), using the KDA online-loading test from this PR's first revision as an out-of-tree harness: real KimiK3DeltaAttention construction at TP8, two TP ranks × two shard orders × padded/unpadded head counts, empty_strided poisoned with NaN.

Configuration	Result
main	4 padded cases fail: projection stays on meta and holds all six sources
This PR	8 pass
This PR with loader re-zeroing disabled	4 padded cases fail: padding rows are NaN

pytest tests/models/kimi_k3/test_amd_kda_fused_conv1d.py \
  tests/models/kimi_k3/test_amd_kda_direct_return.py \
  tests/models/kimi_k3/test_amd_kda_checkpoint.py \
  tests/models/kimi_k3/test_kda.py tests/models/kimi_k3/test_kda_metadata.py \
  tests/models/test_alice_ai.py -q
# 181 passed, 31 skipped (MI355X)
pytest tests/models/kimi_k3/test_amd_kda_decode.py -k test_decode_weights_staged_after_loading -q
# 4 passed

ruff passes; the mypy hook reports no new errors in amd/kda.py (28 vs. 29 pre-existing on main).

No unit test is added. The first revision's test imported the Kimi K3 KDA layer into the shared model_executor pytest process; that pulls in numba, which exposes a latent triton_utils/dispatcher issue and fails the Triton dispatcher tests that run afterwards. That dispatcher issue is independent of this fix and is left to a separate PR.

Full-model GPU validation used vLLM 81198e97ba7eee2a22540caaa756b7fdddcb4d93 plus the first revision of this fix (same weight_loader_numel accounting; padding zeroed during materialization instead of in the loader), on 2 × 8 MI355X GPUs, TP8/DCP8, fastsafetensors, target-model FP8 PTPC, DSpark4 with synthetic AL=3.36, C48/3600s. Both runs had the same P/D hosts and started with ~287.707 GiB free per GPU. The current revision has not been rerun at full-model scale.

Measurement	Original run	Fixed run
Loading	HSA out-of-resources at 10/12 checkpoint groups	All 16 workers loaded; 513 target online layers per worker
Maximum buffered GPU source references per worker	~42.19 GiB	0.451 GiB
Lowest sampled free VRAM, P / D	22 / 88 MiB	25.818 / 24.869 GiB
Benchmark	Did not start	Full 3600s coverage; 4224 successful requests

Source-reference bytes overlap allocated/reserved memory and must not be added to those counters. This verifies loading and serving completion, not decoding accuracy or a throughput gain. The fixed run had 7 empty-content request errors and a 300s cancel-drain timeout, retained in the run logs.

Details

At TP8, KDA's local merged projection has 6284 real rows and 4 alignment rows. Without weight_loader_numel, online processing waits for 6288 rows although the checkpoint only provides the six real projections, so it defers processing and retains every source until model finalization. With GPU-backed fastsafetensors inputs this exhausts VRAM before loading completes.

During online loading the constructor's padding zero_() runs on meta, and materialization allocates with empty_strided. Layerwise processing then replays the buffered calls through the layer's weight loader, so re-zeroing the padding rows there restores them. CopyCounter only counts copy_, so the zeroing does not affect completion accounting. It adds no steady-state inference work.

The subclass extends the shared Kimi-Linear projection instead of importing the NVIDIA class, which would pull NVIDIA backend modules into the ROCm path and change AMD loading through its block-scale replication.

Related work / duplicate check: #50523 fixes general zero initialization for rematerialized weights; it does not declare KDA's loadable size and does not resolve this source-retention failure. If it lands, the loader re-zeroing becomes redundant but harmless. Open KDA quantization PRs and #61287 were also checked; no matching KDA completion fix was found.

Merge Danger

Door: two-way. Blast Radius: Kimi K3 ROCm KDA weight loading.

AI assistance was used: Codex prepared the first revision; Claude Code prepared the review revision and the CI analysis. Results above describe agent-run validation.

---

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

[claude]

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.

[github-actions]

✅ @whx-sjtu, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Run commands test your branch with every commit currently on its upstream target branch. If it is behind, the bot merges the target branch into it first (pull before pushing again). Append --allow-stale to a run command to test an outdated branch as is, at your own risk.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[whx-sjtu]

/ci run

[github-actions]

✅ Triggered Buildkite CI #97873 for commit 5853e6f0b5bc.

[whx-sjtu]

/ci run

[github-actions]

✅ Triggered Buildkite CI #97964 for commit 05b19461750e.

[simondanielsson]

Thanks for the fix!

[simondanielsson]

Suggestion: can we instead fix this in the same way as nvidia, i.e. add a new _KimiGDNMergedColumnParallelLinear and add the padding/zeroing fix there? Alternatively just import their implementation directly

That way we don't have to touch any shared code

[whx-sjtu]

Done — moved the padding handling into an AMD-local _PaddedKimiGDNMergedColumnParallelLinear (subclassing the shared Kimi-Linear one), mirroring the NVIDIA class; meta.py is no longer touched. I didn't import the NVIDIA class directly since nvidia/kda.py pulls the NVIDIA backend's ops into the ROCm path, and its block-scale replication would change AMD loading behavior.

[whx-sjtu]

/ci run

[github-actions]

✅ Triggered Buildkite CI #98256 for commit f97c0866a2c1.