[Bugfix] Fuse GLM-5.3 KDA q|k|v convs into one conv1d parameter

tests/models/glm5next/test_kda_conv_weight_refit.py

@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""GLM-5.3-Flash KDA q|k|v short convs are one fused parameter.
+
+The conv kernels read ``conv1d.weight`` directly, so an RL refit must reach it
+through every write path (#55087), and FULL cudagraphs capture its address, so
+it must never move.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm.model_executor.model_loader.reload.layerwise import (
+    finalize_layerwise_reload,
+    initialize_layerwise_reload,
+    record_metadata_for_reloading,
+)
+from vllm.models.glm5next.common import kda
+from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
+
+CHANNELS, KERNEL = 8, 4
+
+
+def _checkpoint_weights(seed: int) -> list[torch.Tensor]:
+    gen = torch.Generator().manual_seed(seed)
+    return [torch.randn(CHANNELS, 1, KERNEL, generator=gen) for _ in range(3)]
+
+
+def _load(conv1d: nn.Module, weights: list[torch.Tensor]) -> None:
+    for shard_id, weight in enumerate(weights):
+        conv1d.weight.weight_loader(conv1d.weight, weight, shard_id)
+
+
+def _kernel_view(conv1d: nn.Module) -> torch.Tensor:
+    return conv1d.weight.view(conv1d.weight.size(0), conv1d.weight.size(-1))
+
+
+@pytest.fixture
+def conv1d(dist_init) -> nn.Module:
+    return kda._fused_qkv_conv1d(KERNEL, CHANNELS, 1, 0, prefix="conv1d")
+
+
+def test_fused_conv_loads_checkpoint_shards_in_qkv_order(conv1d):
+    weights = _checkpoint_weights(seed=0)
+    _load(conv1d, [w.view(CHANNELS, KERNEL) for w in weights[:1]] + weights[1:])
+    assert conv1d.weight.shape == (3 * CHANNELS, 1, KERNEL)
+    assert torch.equal(_kernel_view(conv1d), torch.cat(weights).squeeze(1))
+
+
+def test_refit_updates_fused_conv_without_moving_it(conv1d):
+    """Refits through load_weights (RL weight sync) and through the layerwise
+    reload (``reload_weights``) both land at the original address."""
+    data_ptr = conv1d.weight.data_ptr()
+    _load(conv1d, _checkpoint_weights(seed=0))
+
+    refit = _checkpoint_weights(seed=1)
+    _load(conv1d, refit)
+    assert torch.equal(_kernel_view(conv1d), torch.cat(refit).squeeze(1))
+
+    reload = _checkpoint_weights(seed=2)
+    model = nn.Sequential(conv1d)
+    record_metadata_for_reloading(model)
+    initialize_layerwise_reload(model)
+    _load(conv1d, reload)
+    finalize_layerwise_reload(model, model_config=None)
+    assert torch.equal(_kernel_view(conv1d), torch.cat(reload).squeeze(1))
+    assert conv1d.weight.data_ptr() == data_ptr
+
+
+def _decode_layer(conv1d: nn.Module) -> kda.Glm5NextLinearAttention:
+    layer = object.__new__(kda.Glm5NextLinearAttention)
+    nn.Module.__init__(layer)
+    layer.prefix = "layer.0"
+    layer.conv1d = conv1d
+    layer.local_num_heads, layer.head_dim = 1, CHANNELS
+    layer.local_projection_size = CHANNELS
+    layer.kda_safe_gate, layer.kda_lower_bound = True, -5.0
+    layer.A_log = layer.dt_bias = None
+    layer._conv_state_dim_first = True
+    layer.kv_cache = (torch.zeros(1, 3 * CHANNELS, KERNEL - 1), None)
+    return layer
+
+
+def test_decode_reads_refitted_fused_conv(conv1d, monkeypatch):
+    """Decode must hand the conv kernel ``conv1d.weight`` itself, not a copy
+    cached by an earlier forward, so refits reach the address FULL cudagraphs
+    capture (#55087)."""
+    layer = _decode_layer(conv1d)
+    metadata = GDNAttentionMetadata(
+        num_prefills=0,
+        num_prefill_tokens=0,
+        num_decodes=1,
+        num_decode_tokens=1,
+        num_spec_decodes=0,
+        num_spec_decode_tokens=0,
+        num_actual_tokens=1,
+        non_spec_query_start_loc=torch.tensor([0, 1], dtype=torch.int32),
+        non_spec_state_indices_tensor=torch.zeros(1, dtype=torch.int32),
+    )
+    monkeypatch.setattr(
+        kda,
+        "get_forward_context",
+        lambda: SimpleNamespace(attn_metadata={layer.prefix: metadata}),
+    )
+    seen = []
+
+    def conv_update(x, conv_state, weight, *args, **kwargs):
+        seen.append(weight)
+        return x
+
+    monkeypatch.setattr(kda, "causal_conv1d_update", conv_update)
+    monkeypatch.setattr(
+        kda, "fused_recurrent_kda", lambda **kwargs: (kwargs["out"], None)
+    )
+
+    def decode() -> torch.Tensor:
+        layer._forward(
+            torch.zeros(1, 3 * CHANNELS),
+            torch.zeros(1, 1, 1, CHANNELS),
+            torch.zeros(1, 1, 1),
+            torch.zeros(1, 1, 1, CHANNELS),
+        )
+        return seen.pop()
+
+    for seed in (0, 1):
+        weights = _checkpoint_weights(seed)
+        _load(conv1d, weights)
+        weight = decode()
+        assert weight.data_ptr() == conv1d.weight.data_ptr()
+        assert torch.equal(weight, torch.cat(weights).squeeze(1))

tests/models/glm5next/test_kda_recurrent.py

@@ -303,8 +303,9 @@ def test_prefill_checkpoint_resumes_suffix(monkeypatch, dim_first, num_spec, off
     layer.kda_lower_bound = -5.0
     layer.kda_prefill_backend = "flashkda"
     layer._conv_state_dim_first = dim_first
-    layer._merged_conv_weight = torch.randn(width, 4, device=device) * 0.1
-    layer.q_conv1d = SimpleNamespace(bias=None)
+    layer.conv1d = SimpleNamespace(
+        weight=torch.randn(width, 1, 4, device=device) * 0.1, bias=None
+    )
     layer.A_log = torch.randn(1, 1, heads, 1, device=device) * 0.1
     layer.dt_bias = torch.randn(heads * D, device=device) * 0.1
     state_shape = (2, heads, D, D)

vllm/models/glm5next/common/kda.py

@@ -21,6 +21,9 @@
     MambaPrefillCheckpointMetadata,
 )
 from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention
+from vllm.model_executor.layers.mamba.gdn.kimi_gdn_linear_attn import (
+    _make_fused_conv1d_weight_loader,
+)
 from vllm.model_executor.layers.mamba.kda_checkpoint import (
     FlashKDAPrefillCheckpointExporter,
     kda_prefill_checkpoint_alignment,
@@ -160,6 +163,36 @@ def _resolve_kda_prefill_backend(
     return "flashkda" if supported and backend != "triton" else "triton"
 
 
+def _fused_qkv_conv1d(
+    conv_size: int, projection_size: int, tp_size: int, tp_rank: int, prefix: str
+) -> ColumnParallelLinear:
+    """The q|k|v short convs as one ``[3 * C_local, 1, K]`` parameter.
+
+    Checkpoint ``q/k/v_conv1d`` weights load as shards 0/1/2 (see
+    ``stacked_params_mapping``). The conv kernels read this parameter directly,
+    so every weight write (load, refit, offload, IPC cache) lands in it and its
+    address stays fixed for cudagraphs (#55087).
+    """
+    conv1d = ColumnParallelLinear(
+        input_size=conv_size,
+        output_size=3 * projection_size,
+        bias=False,
+        params_dtype=torch.float32,
+        prefix=prefix,
+    )
+    conv1d.weight.data = conv1d.weight.data.unsqueeze(1)
+    delattr(conv1d.weight, "weight_loader")
+    set_weight_attrs(
+        conv1d.weight,
+        {
+            "weight_loader": _make_fused_conv1d_weight_loader(
+                [projection_size] * 3, tp_size, tp_rank
+            )
+        },
+    )
+    return conv1d
+
+
 class Glm5NextLinearAttention(GatedDeltaNetAttention):
     head_dim: int
     num_heads: int
@@ -266,37 +299,13 @@ def __init__(
 
         set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
 
-        self.q_conv1d = ColumnParallelLinear(
-            input_size=self.conv_size,
-            output_size=projection_size,
-            bias=False,
-            params_dtype=torch.float32,
-            prefix=f"{prefix}.q_conv1d",
-        )
-        self.k_conv1d = ColumnParallelLinear(
-            input_size=self.conv_size,
-            output_size=projection_size,
-            bias=False,
-            params_dtype=torch.float32,
-            prefix=f"{prefix}.k_conv1d",
-        )
-        self.v_conv1d = ColumnParallelLinear(
-            input_size=self.conv_size,
-            output_size=projection_size,
-            bias=False,
-            params_dtype=torch.float32,
-            prefix=f"{prefix}.v_conv1d",
+        self.conv1d = _fused_qkv_conv1d(
+            self.conv_size,
+            projection_size,
+            self.tp_size,
+            self.tp_rank,
+            prefix=f"{prefix}.conv1d",
         )
-        # unsqueeze to fit conv1d weights shape into the linear weights shape.
-        # Can't do this in `weight_loader` since it already exists in
-        # `ColumnParallelLinear` and `set_weight_attrs`
-        # doesn't allow to override it
-        self.q_conv1d.weight.data = self.q_conv1d.weight.data.unsqueeze(1)
-        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.
-        self._merged_conv_weight: torch.Tensor | None = None
 
         self.A_log = nn.Parameter(
             torch.empty(1, 1, self.local_num_heads, 1, dtype=torch.float32)
@@ -552,20 +561,11 @@ def _forward(
         # One merged short-conv over q|k|v instead of three separate calls. The
         # 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.
-        if self._merged_conv_weight is None:
-
-            def _w(m):
-                return m.weight.view(m.weight.size(0), m.weight.size(2))
-
-            self._merged_conv_weight = torch.cat(
-                [_w(self.q_conv1d), _w(self.k_conv1d), _w(self.v_conv1d)],
-                dim=0,
-            ).contiguous()
-        conv_weights = self._merged_conv_weight
-        conv_bias = self.q_conv1d.bias
+        # three calls. conv_state is already stored as the merged q|k|v state.
+        conv_weights = self.conv1d.weight.view(
+            self.conv1d.weight.size(0), self.conv1d.weight.size(-1)
+        )
+        conv_bias = self.conv1d.bias
 
         # Split projections / gating into spec (draft-verify) and non-spec token
         # groups when speculative decoding is active. Spec tokens carry

vllm/models/glm5next/common/model.py

@@ -874,6 +874,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             (".in_proj_qkvbfg_a", ".b_proj", 3),
             (".in_proj_qkvbfg_a", ".f_a_proj", 4),
             (".in_proj_qkvbfg_a", ".g_a_proj", 5),
+            # KDA: fuse the q, k, v short convs into one conv1d
+            (".conv1d", ".q_conv1d", 0),
+            (".conv1d", ".k_conv1d", 1),
+            (".conv1d", ".v_conv1d", 2),
         ]
         if _is_moe(self.config):
             # Params for weights, fp8 weight scales, fp8 activation scales