From 8015fca1bbb6e8ba3febff956a134c824524388f Mon Sep 17 00:00:00 2001 From: Mustafa YILDIRIM Date: Fri, 18 Sep 2026 15:04:10 -0700 Subject: [PATCH] [Bugfix] Invalidate KDA merged conv-weight cache on weight reload GLM-5.3 KDA lazily merges q/k/v conv weights into a single cached tensor. Online-RL refits update q_conv1d/k_conv1d/v_conv1d through their weight loaders but never invalidate the cache, so inference keeps serving the pre-refit (e.g. dummy-warmup) merge. Wrap each conv weight's stamped weight_loader so any write clears _merged_conv_weight; it is rebuilt from the new params on the next forward. Covers checkpoint-format refits (weight-loader path, including layerwise reload which re-dispatches through the stamped loader). Fixes #55087 Signed-off-by: Mustafa YILDIRIM --- vllm/models/glm5next/common/kda.py | 31 +++++++++++++++++++++++++++--- 1 file changed, 28 insertions(+), 3 deletions(-) diff --git a/vllm/models/glm5next/common/kda.py b/vllm/models/glm5next/common/kda.py index 436b9685cb12..d8793d527a73 100644 --- a/vllm/models/glm5next/common/kda.py +++ b/vllm/models/glm5next/common/kda.py @@ -51,6 +51,26 @@ ) +def _invalidate_merged_conv_weight_cache( + kda_layer: "Glm5NextLinearAttention", conv: nn.Module +) -> None: + """Wrap ``conv.weight.weight_loader`` so any weight write clears the + layer's merged q|k|v conv weight cache (rebuilt lazily on the next + forward). Weight loading dispatches on the param's stamped loader, not + the module attribute.""" + original_loader = conv.weight.weight_loader + + def weight_loader(param, *args, **kwargs): + original_loader(param, *args, **kwargs) + kda_layer._merged_conv_weight = None + + set_weight_attrs(conv.weight, {"kda_owner_layer": kda_layer}) + # set_weight_attrs refuses to overwrite an existing attr; delete the + # stamped loader and re-stamp the wrapper. + delattr(conv.weight, "weight_loader") + set_weight_attrs(conv.weight, {"weight_loader": weight_loader}) + + class _Glm5NextMergedColumnParallelLinear(MergedColumnParallelLinear): """Merged projection with multiple replicated output shards. @@ -273,8 +293,12 @@ def __init__( self.k_conv1d.weight.data = self.k_conv1d.weight.data.unsqueeze(1) self.v_conv1d.weight.data = self.v_conv1d.weight.data.unsqueeze(1) # Lazily-built merged q|k|v conv weight (built on first forward, after - # weights are loaded). See _forward. + # weights are loaded). See _forward. Any write through a conv weight + # loader invalidates the cache, so a weight refit (online-RL reload) + # never serves the stale pre-refit merge (#55087). self._merged_conv_weight: torch.Tensor | None = None + for conv in (self.q_conv1d, self.k_conv1d, self.v_conv1d): + _invalidate_merged_conv_weight_cache(self, conv) self.A_log = nn.Parameter( torch.empty(1, 1, self.local_num_heads, 1, dtype=torch.float32) @@ -500,8 +524,9 @@ def _forward( # 1D conv is independent per channel, so concatenating q/k/v along the # channel dim and running a single causal_conv1d is bit-identical to # three calls. The merged weight is q|k|v conv weights concatenated; - # built once and cached (params are fixed after load). conv_state is - # already stored as the merged q|k|v state, so it is used directly. + # cached until any conv weight is reloaded (see __init__: weight + # loaders invalidate the cache). conv_state is already stored as the + # merged q|k|v state, so it is used directly. if self._merged_conv_weight is None: def _w(m):