diff --git a/tests/models/glm5next/test_kda_conv_weight_refit.py b/tests/models/glm5next/test_kda_conv_weight_refit.py new file mode 100644 index 000000000000..0cd6c11203df --- /dev/null +++ b/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)) diff --git a/tests/models/glm5next/test_kda_recurrent.py b/tests/models/glm5next/test_kda_recurrent.py index 06b2eafac6a2..e904ce8b6ac1 100644 --- a/tests/models/glm5next/test_kda_recurrent.py +++ b/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) diff --git a/vllm/models/glm5next/common/kda.py b/vllm/models/glm5next/common/kda.py index 68d405b40cfc..f7103e75fd45 100644 --- a/vllm/models/glm5next/common/kda.py +++ b/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 diff --git a/vllm/models/glm5next/common/model.py b/vllm/models/glm5next/common/model.py index 5305a36793cb..6b0cbe4ad61b 100644 --- a/vllm/models/glm5next/common/model.py +++ b/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