[JaredforReal]

Overview

Glm5NextLinearAttention cached the merged q|k|v conv weight on the first forward, so an RL refit left the conv kernels reading stale values (#55087). This PR replaces the three q/k/v_conv1d modules with a single conv1d parameter of shape [3 * C_local, 1, K], which the kernels read directly. Checkpoint q_conv1d / k_conv1d / v_conv1d weights load into it as shards 0/1/2 through stacked_params_mapping, using the _make_fused_conv1d_weight_loader already used by Kimi-Linear KDA and OLMo-Hybrid GDN. There is no second tensor to keep in sync, so every write path lands in the tensor the kernels read, at a fixed address.

Alternative fixes for the same issue: #56298 and #57638 (see Details; their authors are co-authors on the commit). The fused-parameter design was suggested by hao-aaron in review and replaces this PR's first approach (parameter views into a merged buffer), which broke under weight offloading and ipc_cache copy mode.

Claims

• After load_weights / reload_weights (checkpoint or kernel format), GLM-5.3 KDA decode uses the refitted conv weights, in eager mode and with the default FULL_AND_PIECEWISE cudagraphs.
• The conv weight keeps its address across refits, so FULL decode graphs captured before a refit stay valid.
• Works with --cpu-offload-gb (with and without refit) and with --load-format ipc_cache in both copy and zero_copy modes.
• No numerical change: the per-rank layout is [q_r | k_r | v_r], the same as the old torch.cat. Direct-load outputs are identical to main. One copy of the conv weights is dropped.

Validation

Hardware: GB300 (SM100), TP1, 4-layer pruned GLM-5.3-Flash at full width (JaredforReal/GLM-5.3-Flash-4L).

pytest tests/models/glm5next/test_kda_conv_weight_refit.py tests/models/glm5next/test_kda_recurrent.py -v

• Unit tests (CPU, conv1d sub-module): checkpoint q/k/v shards (2D and 3D) land in q|k|v order; load_weights-style refit and a real initialize_layerwise_reload / finalize_layerwise_reload cycle update the weight without moving it. 22 passed with test_kda_recurrent.py.
• test_refit_reaches_cudagraph_decode[cpu_offload_gb=0|10] (pruned model, default cudagraphs, one engine per case): direct load -> reference tokens; reload_weights of the checkpoint with the conv weights x2 -> output changes; checkpoint as-is -> == reference; kernel-format reload_weights(is_checkpoint_format=False) on conv1d.weight x2 / x0.5 -> changes / == reference; the conv weight address is identical after all four refits. 2 passed (~5 min each, mostly startup). On main the first refit step fails ([Bug]: GLM-5.3 KDA keeps stale merged convolution weights after refit #55087).
• pre-commit run on the changed files: passed (ruff, mypy, SPDX, test tethering).

CI note: the e2e test is the first in tests/models/glm5next to use the pruned checkpoint (19 GB download, needs SM90/SM100, skipped elsewhere). It runs in the existing glm5next-unit-tests step; the two offload cases take about 10 min together on GB300, mostly engine startup. #57537 adds a dedicated step for pruned-checkpoint e2e tests that this test could move to once it lands.

Full refit flow, 3 prompts x 24 greedy tokens, compared with main direct load:

tree	config	direct load	dummy -> reload_weights()	load_weights partial update / kernel-format refit
main	cudagraphs	reference	wrong (#55087)	wrong
previous revision (buffer views)	cudagraphs	== main	== main	== main
previous revision (buffer views)	cpu_offload_gb=10	wrong	–	–
this PR	cudagraphs	== main	== main	== main
this PR	eager	== main	== main	== main
this PR	cpu_offload_gb=10	== main	–	–

After the reload, every language-model parameter and buffer checksum equals a fresh load.

--load-format ipc_cache (a vllm preload daemon per GPU, default cudagraphs). After each refit the check compares the conv weight the kernel reads against the checkpoint × scale (max abs diff), then compares the output:

tree	mode	load	kernel-format refit x2 / x1
main	copy	== main	weight stale (diff 0.72), output unchanged (#55087)
previous revision (buffer views)	copy	== main	weight stale (diff 0.72), output unchanged
previous revision (buffer views)	zero_copy	== main	diff 0, output changes / == main
this PR	copy	== main	diff 0, output changes / == main, address unchanged
this PR	zero_copy	== main	diff 0, output changes / == main, address unchanged

Checkpoint-format refits are not supported on an ipc_cache-loaded model on main either, so they're not compared here. _apply_entries re-registers every parameter as a plain nn.Parameter (399/399 lose weight_loader), so load_weights of any stacked weight raises AttributeError (verified on main with self_attn.q_proj), and the layerwise reload_weights fails with Cannot copy out of meta tensor. On main the three conv weights happened to fall through to default_weight_loader; with this PR they're stacked like the other projections. Making checkpoint-format refit work after ipc_cache is a separate issue.

Details

Root cause. The KDA layer runs one causal conv over the concatenated q|k|v channels. The merged weight was built lazily on the first forward with torch.cat and cached, so later writes to q/k/v_conv1d.weight (RL refit through the weight-transfer engines, reload_weights, or kernel-format copy_) never reached it.

Why not invalidate the cache (#56298, #57638). Both PRs drop the cached tensor on refit and rebuild it on the next forward. That works with enforce_eager=True, but with the default FULL_AND_PIECEWISE cudagraphs the KDA forward is captured inside the FULL decode graphs, which keep the old tensor's address, so decode replays never see the rebuilt weight.

Why one parameter instead of views into a merged buffer. Views only hold while nothing replaces a parameter's storage. The UVA and prefetch CPU offloaders set each parameter's .data right after the layer is built, and the ipc_cache loader in copy mode clones every entry separately; either way the three weights stop aliasing the tensor the kernels read. With a single parameter, those paths, the layerwise reload, and kernel-format copy_ all write the tensor the kernels use. This matches how Kimi-Linear and OLMo-Hybrid already store their fused short convs.

Trade-offs.

• The vLLM parameter name changes from q/k/v_conv1d.weight to conv1d.weight. Checkpoint-format updates keep the HF names and go through the stacked mapping. Kernel-format writers need the fused name, and GLM-5.3 artifacts saved with vLLM parameter names (sharded_state, vLLM-serialized tensorizer) before this change will not load.
• Open PRs [glm] [pcp] [2/n] add KCP for hybrid GLM KDA #59304 (KCP) and [CPU] Add KDA backend for GLM5Next #57496 (CPU KDA backend) reference q_conv1d / _merged_conv_weight and will need a small rebase.

Not covered here. The TP2 cudagraph reload drift from the FlashInfer TRT-LLM FP8 MoE gemm1_clamp_limit tensor and the DSA indexer's lazily cached _wp_fp32 are separate issues, left for follow-ups.

AI assistance (Claude Code) was used for the investigation and the draft of the change; every line was reviewed and the runs above were executed by the submitter.

---

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.

[Copilot]

🟡 Changes recommended

IPC cache copy mode independently replaces the aliased tensors, reintroducing stale convolution weights after refits.

1 open finding

• IPC copy mode breaks storage aliases after refit · New

What changed in this PR

Fixes stale GLM-5.3 KDA convolution weights during refits by sharing storage between q/k/v parameters and the merged kernel buffer.

Changes:

• Allocates and aliases the merged convolution buffer during initialization.
• Adds unit and CUDA-graph refit regression tests.

File	Description
vllm/​models/​glm5next/​common/​kda.py	Introduces shared q/k/v convolution storage.
tests/​models/​glm5next/​test_kda_conv_weight_refit.py	Tests aliasing, reloads, and graph-visible storage stability.

🧠 Review effort: Balanced

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[hao-aaron]

Nice find on the FULL-graph address problem. Making the merged weight live at a fixed address is the right goal. I think the aliasing approach can still go stale, though, in two load paths where the q/k/v params stop sharing storage with _merged_conv_weight:

1. Weight offloading (--cpu-offload-gb, UVA or the prefetch offloader).
   • make_layers runs the offloader right after each layer's __init__. It moves each param's .data to new CPU/UVA memory (offloader/uva.py _maybe_offload_to_cpu, offloader/prefetch.py:677). Buffers aren't offloaded.
   • The checkpoint then loads into the offloaded params, while _merged_conv_weight keeps its torch.empty contents from the first request.
   • On main this case works today, because the lazy torch.cat reads the loaded params.
2. --load-format ipc_cache with mode: copy.
   • _apply_entries clones each exported tensor separately, so the params and the buffer come back as independent copies.
   • Outputs are correct at startup and go stale after the next refit.

Suggestion: fuse the three convs into one param, the way Kimi-Linear's KDA (kimi_gdn_linear_attn.py) and OLMo-Hybrid's GDN (olmo_gdn_linear_attn.py) already do. Then there's no second tensor to keep in sync:

# kda.py __init__: replaces q/k/v_conv1d, the unsqueezes, and the cached merge
self.conv1d = ColumnParallelLinear(
    input_size=self.conv_size,
    output_size=3 * projection_size,
    bias=False,
    params_dtype=torch.float32,
    prefix=f"{prefix}.conv1d",
)
self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
delattr(self.conv1d.weight, "weight_loader")
set_weight_attrs(
    self.conv1d.weight,
    {"weight_loader": _make_fused_conv1d_weight_loader(
        [projection_size] * 3, self.tp_size, self.tp_rank)},
)

# _forward
conv_weights = self.conv1d.weight.view(
    self.conv1d.weight.size(0), self.conv1d.weight.size(-1)
)
conv_bias = self.conv1d.bias

# model.py stacked_params_mapping
(".conv1d", ".q_conv1d", 0),
(".conv1d", ".k_conv1d", 1),
(".conv1d", ".v_conv1d", 2),

_make_fused_conv1d_weight_loader can be imported from kimi_gdn_linear_attn, as kimi_k3/amd/kda.py already does. It accepts both 2D and 3D checkpoint conv tensors.

From reading the code, this should:

• keep outputs identical: the per-rank layout is [q_r | k_r | v_r], the same as today's torch.cat;
• cover every conv call: the spec-decode, prefill and decode paths already all read the one merged weight;
• keep a fixed address for FULL cudagraphs, and make every write path land in the tensor the kernel reads: layerwise reload, kernel-format copy_, offloading, and ipc_cache copy mode;
• drop one copy of the conv weights.

Trade-offs:

• The param is renamed to conv1d.weight.
  • Kernel-format writers would need the fused name.
  • GLM-5.3 artifacts saved with vLLM param names (sharded_state, vLLM-serialized tensorizer) wouldn't load.
  • Checkpoint-format updates keep the HF names and go through the stacked mapping.
• Tests and other PRs reference the old names: tests/models/glm5next/test_kda_recurrent.py mocks q_conv1d / _merged_conv_weight, and [glm] [pcp] [2/n] add KCP for hybrid GLM KDA #59304 and [CPU] Add KDA backend for GLM5Next #57496 reference them too.

Your refit test should carry over, with conv1d.weight for the kernel-format step. A --cpu-offload-gb run would cover case 1.

Caveat: I haven't tested any of this. It's from reading the code, so please treat it as a suggestion, not a verified fix.

This review/comment was written by an AI agent (Claude Code) on my behalf.

[hao-aaron]

The fused param change looks good to me. One concern about the new e2e test: test_refit_reaches_cudagraph_decode is quite heavy for glm5next-unit-tests. Everything else in that step is kernel/unit-level and the step has a 20-min timeout, while this test pulls a 19 GB checkpoint and starts four engines (two per cpu_offload_gb case). Two lighter alternatives:

1. A _forward-level guard, no engine. test_prefill_checkpoint_resumes_suffix in test_kda_recurrent.py already drives layer._forward on a mocked layer. A sibling test could stub the conv kernel to record the weight it receives and assert it shares storage with layer.conv1d.weight, before and after an in-place refit. That targets the actual regression (a cached copy of the merged weight) and runs in milliseconds. Together with your CPU tests for the fused loader and the layerwise reload, I think that covers the change.

2. Keep the e2e shape, but with dummy weights. test_initialization already builds GLM-5.3 with load_format="dummy" + dummy_hf_overrides, so only the config (and the tokenizer, or use skip_tokenizer_init with token IDs) gets downloaded. For the checkpoint-format step, refit only the three q/k/v_conv1d tensors (split the current fused weight and scale it). finalize_layerwise_processing puts back the original tensors for layers that receive no weights. Since greedy tokens from random weights may not move, compare logprobs: they should differ after ×2 and match exactly after restoring. The address-stability check stays as is. I haven't verified that a shrunk dummy GLM-5.3 generates cleanly with cudagraphs, so this needs a quick try.

If you'd rather keep the pruned-checkpoint e2e, I'd move it into the glm5next-pruned-e2e step from #57537 instead of the unit step.

This review/comment was written by an AI agent (Claude Code) on my behalf.

[hao-aaron]

LGTM, thanks for the fused param and the _forward guard.

For the record, the same bug pattern exists in the GLM-5.3 indexer: Indexer._wp_fp32 (vllm/models/glm5next/common/attention.py) is an fp32 copy of the weights_proj rows of the fused wk_weights_proj, built lazily on the first forward and never refreshed, so after a refit the head gate keeps using the old weights (and FULL cudagraphs hold its address). It exists because wk_weights_proj is one bf16 param but the gate is computed in fp32. A structural fix along the lines of this PR would be to unfuse it and make weights_proj its own fp32 parameter loaded directly by load_weights. Not for this PR, just noting it here so it's tracked.

This review/comment was written by an AI agent (Claude Code) on my behalf.

[github-actions]

✅ @JaredforReal, 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.

[JaredforReal]

/ci run

[github-actions]

✅ Triggered Buildkite CI #97525 for commit 52fd007fe381.

ℹ️ This PR is 10 commits behind upstream main. Merged upstream main into this branch as 52fd007fe381 before starting CI. Pull the branch before pushing more changes.