From cd28ec295b933898e8208b0694eed004e0aec9c3 Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Sat, 10 Oct 2026 07:21:55 +0000 Subject: [PATCH 1/3] [Bugfix] Alias GLM-5.3 KDA q|k|v conv weights into the merged conv weight Glm5NextLinearAttention runs one causal conv over q|k|v and cached the concatenated weight on first forward, so an RL refit that rewrote the three ColumnParallelLinear conv weights left the kernels reading the old values (#55087). Invalidating that cache is not enough: with the default FULL_AND_PIECEWISE cudagraphs the KDA forward is captured inside the FULL decode graphs, which keep the merged weight's address, so a rebuilt tensor never reaches decode replays. Allocate the merged [3*C, 1, K] buffer up front and make the q/k/v conv weights views into it. Every write path (initial load, load_weights from a weight-transfer engine, the layerwise reload_weights, kernel-format copy_) then lands in the weight the kernels read, at an address that never changes. The buffer is a non-persistent parameter-alias buffer, which the layerwise reload already skips and restores. Co-authored-by: Jyotirmoy Roy Co-authored-by: Mustafa YILDIRIM <2341474+mustafayildirim@users.noreply.github.com> Co-authored-by: Claude Fable 5.1 Signed-off-by: Jared Wen --- .../glm5next/test_kda_conv_weight_refit.py | 199 ++++++++++++++++++ vllm/models/glm5next/common/kda.py | 59 ++++-- 2 files changed, 236 insertions(+), 22 deletions(-) create mode 100644 tests/models/glm5next/test_kda_conv_weight_refit.py 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..76853d391834 --- /dev/null +++ b/tests/models/glm5next/test_kda_conv_weight_refit.py @@ -0,0 +1,199 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GLM-5.3-Flash KDA q|k|v conv weights share one merged buffer. + +The conv kernels read ``_merged_conv_weight`` while weight loading writes the +three ``ColumnParallelLinear`` weights, so an RL refit must reach the merged +weight through every write path (#55087). FULL cudagraphs capture the merged +weight's address, so it must also never move. +""" + +import os + +import pytest +import torch +from torch import nn + +from vllm import LLM, SamplingParams +from vllm.model_executor.layers.linear import ColumnParallelLinear +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.platforms import current_platform + +CHANNELS, KERNEL = 8, 4 +CONV_WEIGHTS = ("q_conv1d.weight", "k_conv1d.weight", "v_conv1d.weight") +# Layers 0-3 of the real checkpoint at full width; needs SM90/SM100 (indexer). +PRUNED_MODEL = os.getenv("GLM5_PRUNED_MODEL", "JaredforReal/GLM-5.3-Flash-4L") + + +class _ConvLayer(nn.Module): + """The conv part of ``Glm5NextLinearAttention``, wired like the layer.""" + + def __init__(self) -> None: + super().__init__() + self.q_conv1d, self.k_conv1d, self.v_conv1d = ( + ColumnParallelLinear( + KERNEL, + CHANNELS, + bias=False, + params_dtype=torch.float32, + prefix=f"{name}_conv1d", + ) + for name in "qkv" + ) + self.register_buffer( + "_merged_conv_weight", + kda._merge_conv_weights(self.q_conv1d, self.k_conv1d, self.v_conv1d), + persistent=False, + ) + + @property + def convs(self) -> tuple[ColumnParallelLinear, ...]: + return (self.q_conv1d, self.k_conv1d, self.v_conv1d) + + +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 _merged(weights: list[torch.Tensor]) -> torch.Tensor: + return torch.cat([w.view(CHANNELS, KERNEL) for w in weights]) + + +@pytest.fixture +def layer(dist_init) -> _ConvLayer: + return _ConvLayer() + + +def test_conv_weights_alias_the_merged_buffer(layer): + """Loading each conv weight through its stock loader fills the merged + weight in place; the kernels' view keeps the checkpoint channel order.""" + merged = layer._merged_conv_weight + assert merged.shape == (3 * CHANNELS, KERNEL) and merged.is_contiguous() + for conv in layer.convs: + assert conv.weight.shape == (CHANNELS, 1, KERNEL) + assert conv.weight.untyped_storage().data_ptr() == ( + merged.untyped_storage().data_ptr() + ) + assert "_merged_conv_weight" not in layer.state_dict() + + weights = _checkpoint_weights(seed=0) + for conv, weight in zip(layer.convs, weights): + conv.weight.weight_loader(conv.weight, weight) + assert torch.equal(merged, _merged(weights)) + + +def test_refit_updates_merged_weight_without_moving_it(layer): + """A refit rewrites the conv weights through load_weights (RL weight sync) + and through the layerwise reload (``reload_weights``); both must show up + in the merged weight at its original address.""" + merged = layer._merged_conv_weight + data_ptr = merged.data_ptr() + for conv, weight in zip(layer.convs, _checkpoint_weights(seed=0)): + conv.weight.weight_loader(conv.weight, weight) + + refit = _checkpoint_weights(seed=1) + for conv, weight in zip(layer.convs, refit): + conv.weight.weight_loader(conv.weight, weight) + assert torch.equal(merged, _merged(refit)) + assert merged.data_ptr() == data_ptr + + reload = _checkpoint_weights(seed=2) + model = nn.Sequential(layer) + record_metadata_for_reloading(model) + initialize_layerwise_reload(model) + for conv, weight in zip(layer.convs, reload): + conv.weight.weight_loader(conv.weight, weight) + finalize_layerwise_reload(model, model_config=None) + assert layer._merged_conv_weight is merged + assert torch.equal(merged, _merged(reload)) + assert merged.data_ptr() == data_ptr + assert "_merged_conv_weight" in layer._non_persistent_buffers_set + for conv in layer.convs: + assert conv.weight.untyped_storage().data_ptr() == ( + merged.untyped_storage().data_ptr() + ) + + +def _merged_conv_weight_ptrs(model: nn.Module) -> list[int]: + return [ + module._merged_conv_weight.data_ptr() + for module in model.modules() + if isinstance(module, kda.Glm5NextLinearAttention) + ] + + +def _refit_from_checkpoint(worker, conv_scale: float) -> list[int]: + """RL refit shape: stream the checkpoint through ``reload_weights`` (the + layerwise reload), with the KDA conv weights scaled.""" + from vllm.model_executor.model_loader import get_model_loader + + runner = worker.model_runner + model = runner.get_model() + loader = get_model_loader(runner.load_config) + + def weights(): + for name, tensor in loader.get_all_weights(runner.model_config, model): + if name.endswith(CONV_WEIGHTS): + tensor = tensor * conv_scale + yield name, tensor + + worker.reload_weights(weights_iterator=weights()) + return _merged_conv_weight_ptrs(model) + + +def _refit_kernel_format(worker, conv_scale: float) -> list[int]: + """Kernel-format refit: ``reload_weights(is_checkpoint_format=False)`` + copies into the parameters without any weight loader.""" + model = worker.model_runner.get_model() + weights = [ + (name, param.data * conv_scale) + for name, param in model.named_parameters() + if name.endswith(CONV_WEIGHTS) + ] + worker.reload_weights(weights_iterator=iter(weights), is_checkpoint_format=False) + return _merged_conv_weight_ptrs(model) + + +@pytest.mark.skipif( + not ( + current_platform.is_cuda() + and torch.cuda.is_available() + and torch.cuda.get_device_capability()[0] in (9, 10) + ), + reason="GLM-5.3-Flash sparse indexer needs SM90/SM100", +) +def test_refit_reaches_cudagraph_decode(monkeypatch: pytest.MonkeyPatch): + """With CUDA graphs on, decode must follow the refitted conv weights and + return to the original output once they are restored (#55087).""" + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") + llm = LLM( + model=PRUNED_MODEL, + max_model_len=2048, + max_num_seqs=4, + gpu_memory_utilization=0.8, + limit_mm_per_prompt={"image": 0, "video": 0}, + ) + params = SamplingParams(temperature=0.0, max_tokens=8) + + def generate() -> list[int]: + return list( + llm.generate(["The capital of France is"], params)[0].outputs[0].token_ids + ) + + reference = generate() + ptrs = [] + for refit, scale, same in ( + (_refit_from_checkpoint, 2.0, False), + (_refit_from_checkpoint, 1.0, True), + (_refit_kernel_format, 2.0, False), + (_refit_kernel_format, 0.5, True), + ): + ptrs.append(llm.collective_rpc(refit, kwargs={"conv_scale": scale})[0]) + assert (generate() == reference) == same, (refit.__name__, scale) + assert all(p == ptrs[0] for p in ptrs) diff --git a/vllm/models/glm5next/common/kda.py b/vllm/models/glm5next/common/kda.py index 141a8cc7cb82..216a2fdc1ad2 100644 --- a/vllm/models/glm5next/common/kda.py +++ b/vllm/models/glm5next/common/kda.py @@ -160,6 +160,29 @@ def _resolve_kda_prefill_backend( return "flashkda" if supported and backend != "triton" else "triton" +def _merge_conv_weights(*convs: ColumnParallelLinear) -> torch.Tensor: + """Re-point the conv layers' weights at slices of one merged buffer. + + The buffer is laid out like the checkpoint's conv1d weights (``[C, 1, K]`` + per layer, concatenated along the channel dim), so the stock + ``ColumnParallelLinear`` loader fills it in place. Returns the ``[3*C, K]`` + view the causal-conv kernels take. + """ + weights = [conv.weight for conv in convs] + merged = torch.empty( + sum(weight.size(0) for weight in weights), + 1, + weights[0].size(-1), + dtype=weights[0].dtype, + device=weights[0].device, + ) + offset = 0 + for weight in weights: + weight.data = merged[offset : offset + weight.size(0)] + offset += weight.size(0) + return merged.view(merged.size(0), merged.size(-1)) + + class Glm5NextLinearAttention(GatedDeltaNetAttention): head_dim: int num_heads: int @@ -285,16 +308,17 @@ def __init__( params_dtype=torch.float32, prefix=f"{prefix}.v_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 + # The q|k|v conv weights are views into one merged buffer, so every + # path that rewrites them (initial load, RL refit through load_weights + # or the layerwise reload, kernel-format copy_) lands directly in the + # weight the conv kernels read: there is no cache to invalidate, and the + # buffer keeps its address so FULL cudagraphs captured with it stay + # valid across refits (#55087). + self.register_buffer( + "_merged_conv_weight", + _merge_conv_weights(self.q_conv1d, self.k_conv1d, self.v_conv1d), + persistent=False, + ) self.A_log = nn.Parameter( torch.empty(1, 1, self.local_num_heads, 1, dtype=torch.float32) @@ -550,18 +574,9 @@ 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() + # three calls. The q/k/v conv weights are views into _merged_conv_weight + # (see __init__) and conv_state is already stored as the merged q|k|v + # state, so both are used directly. conv_weights = self._merged_conv_weight conv_bias = self.q_conv1d.bias From 92ba789bc59973d3ebb9d8e43e8f65d8666d07e3 Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Mon, 12 Oct 2026 06:06:08 +0000 Subject: [PATCH 2/3] [Bugfix] Fuse GLM-5.3 KDA q|k|v convs into one conv1d parameter Replace the q/k/v conv weights aliased into a non-persistent merged buffer with a single [3 * C_local, 1, K] conv1d parameter. The checkpoint's q_conv1d/k_conv1d/v_conv1d weights load into it as shards 0/1/2 through stacked_params_mapping and the fused conv1d weight loader already used by the Kimi and OLMo KDA/GDN layers. Parameter views did not survive every weight path. The UVA and prefetch CPU offloaders replace each parameter's .data after the layer is built, and the IPC weight cache clones every entry on its own, so the q/k/v weights stopped being views of the buffer the kernels read and cpu_offload_gb produced wrong tokens. A single parameter is the tensor all of these paths write, so loads, refits and offload reach the kernels at a fixed address and FULL cudagraph decode stays correct after a refit (#55087). Co-Authored-By: Claude Opus 5.5 Signed-off-by: Jared Wen --- .../glm5next/test_kda_conv_weight_refit.py | 150 +++++++----------- tests/models/glm5next/test_kda_recurrent.py | 5 +- vllm/models/glm5next/common/kda.py | 93 +++++------ vllm/models/glm5next/common/model.py | 4 + 4 files changed, 103 insertions(+), 149 deletions(-) diff --git a/tests/models/glm5next/test_kda_conv_weight_refit.py b/tests/models/glm5next/test_kda_conv_weight_refit.py index 76853d391834..22352ae398a3 100644 --- a/tests/models/glm5next/test_kda_conv_weight_refit.py +++ b/tests/models/glm5next/test_kda_conv_weight_refit.py @@ -1,11 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""GLM-5.3-Flash KDA q|k|v conv weights share one merged buffer. +"""GLM-5.3-Flash KDA q|k|v short convs are one fused parameter. -The conv kernels read ``_merged_conv_weight`` while weight loading writes the -three ``ColumnParallelLinear`` weights, so an RL refit must reach the merged -weight through every write path (#55087). FULL cudagraphs capture the merged -weight's address, so it must also never move. +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. Weight offloading replaces parameter storage after +construction, which the fused parameter must also survive. """ import os @@ -15,7 +15,7 @@ from torch import nn from vllm import LLM, SamplingParams -from vllm.model_executor.layers.linear import ColumnParallelLinear +from vllm.distributed import cleanup_dist_env_and_memory from vllm.model_executor.model_loader.reload.layerwise import ( finalize_layerwise_reload, initialize_layerwise_reload, @@ -25,104 +25,60 @@ from vllm.platforms import current_platform CHANNELS, KERNEL = 8, 4 -CONV_WEIGHTS = ("q_conv1d.weight", "k_conv1d.weight", "v_conv1d.weight") +CKPT_CONV_WEIGHTS = ("q_conv1d.weight", "k_conv1d.weight", "v_conv1d.weight") # Layers 0-3 of the real checkpoint at full width; needs SM90/SM100 (indexer). PRUNED_MODEL = os.getenv("GLM5_PRUNED_MODEL", "JaredforReal/GLM-5.3-Flash-4L") -class _ConvLayer(nn.Module): - """The conv part of ``Glm5NextLinearAttention``, wired like the layer.""" - - def __init__(self) -> None: - super().__init__() - self.q_conv1d, self.k_conv1d, self.v_conv1d = ( - ColumnParallelLinear( - KERNEL, - CHANNELS, - bias=False, - params_dtype=torch.float32, - prefix=f"{name}_conv1d", - ) - for name in "qkv" - ) - self.register_buffer( - "_merged_conv_weight", - kda._merge_conv_weights(self.q_conv1d, self.k_conv1d, self.v_conv1d), - persistent=False, - ) - - @property - def convs(self) -> tuple[ColumnParallelLinear, ...]: - return (self.q_conv1d, self.k_conv1d, self.v_conv1d) - - 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 _merged(weights: list[torch.Tensor]) -> torch.Tensor: - return torch.cat([w.view(CHANNELS, KERNEL) for w in weights]) +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 layer(dist_init) -> _ConvLayer: - return _ConvLayer() - - -def test_conv_weights_alias_the_merged_buffer(layer): - """Loading each conv weight through its stock loader fills the merged - weight in place; the kernels' view keeps the checkpoint channel order.""" - merged = layer._merged_conv_weight - assert merged.shape == (3 * CHANNELS, KERNEL) and merged.is_contiguous() - for conv in layer.convs: - assert conv.weight.shape == (CHANNELS, 1, KERNEL) - assert conv.weight.untyped_storage().data_ptr() == ( - merged.untyped_storage().data_ptr() - ) - assert "_merged_conv_weight" not in layer.state_dict() +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) - for conv, weight in zip(layer.convs, weights): - conv.weight.weight_loader(conv.weight, weight) - assert torch.equal(merged, _merged(weights)) + _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_merged_weight_without_moving_it(layer): - """A refit rewrites the conv weights through load_weights (RL weight sync) - and through the layerwise reload (``reload_weights``); both must show up - in the merged weight at its original address.""" - merged = layer._merged_conv_weight - data_ptr = merged.data_ptr() - for conv, weight in zip(layer.convs, _checkpoint_weights(seed=0)): - conv.weight.weight_loader(conv.weight, weight) +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) - for conv, weight in zip(layer.convs, refit): - conv.weight.weight_loader(conv.weight, weight) - assert torch.equal(merged, _merged(refit)) - assert merged.data_ptr() == data_ptr + _load(conv1d, refit) + assert torch.equal(_kernel_view(conv1d), torch.cat(refit).squeeze(1)) reload = _checkpoint_weights(seed=2) - model = nn.Sequential(layer) + model = nn.Sequential(conv1d) record_metadata_for_reloading(model) initialize_layerwise_reload(model) - for conv, weight in zip(layer.convs, reload): - conv.weight.weight_loader(conv.weight, weight) + _load(conv1d, reload) finalize_layerwise_reload(model, model_config=None) - assert layer._merged_conv_weight is merged - assert torch.equal(merged, _merged(reload)) - assert merged.data_ptr() == data_ptr - assert "_merged_conv_weight" in layer._non_persistent_buffers_set - for conv in layer.convs: - assert conv.weight.untyped_storage().data_ptr() == ( - merged.untyped_storage().data_ptr() - ) + assert torch.equal(_kernel_view(conv1d), torch.cat(reload).squeeze(1)) + assert conv1d.weight.data_ptr() == data_ptr -def _merged_conv_weight_ptrs(model: nn.Module) -> list[int]: +def _conv_weight_ptrs(model: nn.Module) -> list[int]: return [ - module._merged_conv_weight.data_ptr() + module.conv1d.weight.data_ptr() for module in model.modules() if isinstance(module, kda.Glm5NextLinearAttention) ] @@ -139,12 +95,12 @@ def _refit_from_checkpoint(worker, conv_scale: float) -> list[int]: def weights(): for name, tensor in loader.get_all_weights(runner.model_config, model): - if name.endswith(CONV_WEIGHTS): + if name.endswith(CKPT_CONV_WEIGHTS): tensor = tensor * conv_scale yield name, tensor worker.reload_weights(weights_iterator=weights()) - return _merged_conv_weight_ptrs(model) + return _conv_weight_ptrs(model) def _refit_kernel_format(worker, conv_scale: float) -> list[int]: @@ -154,10 +110,17 @@ def _refit_kernel_format(worker, conv_scale: float) -> list[int]: weights = [ (name, param.data * conv_scale) for name, param in model.named_parameters() - if name.endswith(CONV_WEIGHTS) + if name.endswith("self_attn.conv1d.weight") ] + assert weights worker.reload_weights(weights_iterator=iter(weights), is_checkpoint_format=False) - return _merged_conv_weight_ptrs(model) + return _conv_weight_ptrs(model) + + +def _greedy(llm: LLM) -> list[int]: + params = SamplingParams(temperature=0.0, max_tokens=8) + output = llm.generate(["The capital of France is"], params)[0] + return list(output.outputs[0].token_ids) @pytest.mark.skipif( @@ -168,25 +131,23 @@ def _refit_kernel_format(worker, conv_scale: float) -> list[int]: ), reason="GLM-5.3-Flash sparse indexer needs SM90/SM100", ) -def test_refit_reaches_cudagraph_decode(monkeypatch: pytest.MonkeyPatch): +@pytest.mark.parametrize("cpu_offload_gb", [0, 10]) +def test_refit_reaches_cudagraph_decode( + monkeypatch: pytest.MonkeyPatch, cpu_offload_gb: int +): """With CUDA graphs on, decode must follow the refitted conv weights and - return to the original output once they are restored (#55087).""" + return to the original output once they are restored (#55087), including + when the KDA layers are offloaded to CPU.""" monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") llm = LLM( model=PRUNED_MODEL, max_model_len=2048, max_num_seqs=4, gpu_memory_utilization=0.8, + cpu_offload_gb=cpu_offload_gb, limit_mm_per_prompt={"image": 0, "video": 0}, ) - params = SamplingParams(temperature=0.0, max_tokens=8) - - def generate() -> list[int]: - return list( - llm.generate(["The capital of France is"], params)[0].outputs[0].token_ids - ) - - reference = generate() + reference = _greedy(llm) ptrs = [] for refit, scale, same in ( (_refit_from_checkpoint, 2.0, False), @@ -195,5 +156,8 @@ def generate() -> list[int]: (_refit_kernel_format, 0.5, True), ): ptrs.append(llm.collective_rpc(refit, kwargs={"conv_scale": scale})[0]) - assert (generate() == reference) == same, (refit.__name__, scale) + assert (_greedy(llm) == reference) == same, (refit.__name__, scale) assert all(p == ptrs[0] for p in ptrs) + + del llm + cleanup_dist_env_and_memory() 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 216a2fdc1ad2..9c13f19e8cd3 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,27 +163,34 @@ def _resolve_kda_prefill_backend( return "flashkda" if supported and backend != "triton" else "triton" -def _merge_conv_weights(*convs: ColumnParallelLinear) -> torch.Tensor: - """Re-point the conv layers' weights at slices of one merged buffer. +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. - The buffer is laid out like the checkpoint's conv1d weights (``[C, 1, K]`` - per layer, concatenated along the channel dim), so the stock - ``ColumnParallelLinear`` loader fills it in place. Returns the ``[3*C, K]`` - view the causal-conv kernels take. + 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). """ - weights = [conv.weight for conv in convs] - merged = torch.empty( - sum(weight.size(0) for weight in weights), - 1, - weights[0].size(-1), - dtype=weights[0].dtype, - device=weights[0].device, + 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 + ) + }, ) - offset = 0 - for weight in weights: - weight.data = merged[offset : offset + weight.size(0)] - offset += weight.size(0) - return merged.view(merged.size(0), merged.size(-1)) + return conv1d class Glm5NextLinearAttention(GatedDeltaNetAttention): @@ -287,37 +297,12 @@ 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", - ) - # The q|k|v conv weights are views into one merged buffer, so every - # path that rewrites them (initial load, RL refit through load_weights - # or the layerwise reload, kernel-format copy_) lands directly in the - # weight the conv kernels read: there is no cache to invalidate, and the - # buffer keeps its address so FULL cudagraphs captured with it stay - # valid across refits (#55087). - self.register_buffer( - "_merged_conv_weight", - _merge_conv_weights(self.q_conv1d, self.k_conv1d, self.v_conv1d), - persistent=False, + self.conv1d = _fused_qkv_conv1d( + self.conv_size, + projection_size, + self.tp_size, + self.tp_rank, + prefix=f"{prefix}.conv1d", ) self.A_log = nn.Parameter( @@ -574,11 +559,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 q/k/v conv weights are views into _merged_conv_weight - # (see __init__) and conv_state is already stored as the merged q|k|v - # state, so both are used directly. - 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 e35887ca7c54..92b877b179b5 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 From 8c75d806acb0bd76e799c2b56bf641bbab19c0f7 Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Sat, 24 Oct 2026 11:22:27 +0000 Subject: [PATCH 3/3] [Test] Replace GLM-5.3 KDA refit e2e test with a _forward guard The pruned-checkpoint e2e test downloaded 19 GB and started four engines in the kernel-level glm5next unit step. Check the regression directly instead: a decode step through `_forward` with the kernels stubbed must hand the conv kernel `conv1d.weight` itself before and after a refit. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Jared Wen --- .../glm5next/test_kda_conv_weight_refit.py | 156 +++++++----------- 1 file changed, 63 insertions(+), 93 deletions(-) diff --git a/tests/models/glm5next/test_kda_conv_weight_refit.py b/tests/models/glm5next/test_kda_conv_weight_refit.py index 22352ae398a3..0cd6c11203df 100644 --- a/tests/models/glm5next/test_kda_conv_weight_refit.py +++ b/tests/models/glm5next/test_kda_conv_weight_refit.py @@ -4,30 +4,24 @@ 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. Weight offloading replaces parameter storage after -construction, which the fused parameter must also survive. +it must never move. """ -import os +from types import SimpleNamespace import pytest import torch from torch import nn -from vllm import LLM, SamplingParams -from vllm.distributed import cleanup_dist_env_and_memory 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.platforms import current_platform +from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata CHANNELS, KERNEL = 8, 4 -CKPT_CONV_WEIGHTS = ("q_conv1d.weight", "k_conv1d.weight", "v_conv1d.weight") -# Layers 0-3 of the real checkpoint at full width; needs SM90/SM100 (indexer). -PRUNED_MODEL = os.getenv("GLM5_PRUNED_MODEL", "JaredforReal/GLM-5.3-Flash-4L") def _checkpoint_weights(seed: int) -> list[torch.Tensor]: @@ -76,88 +70,64 @@ def test_refit_updates_fused_conv_without_moving_it(conv1d): assert conv1d.weight.data_ptr() == data_ptr -def _conv_weight_ptrs(model: nn.Module) -> list[int]: - return [ - module.conv1d.weight.data_ptr() - for module in model.modules() - if isinstance(module, kda.Glm5NextLinearAttention) - ] - - -def _refit_from_checkpoint(worker, conv_scale: float) -> list[int]: - """RL refit shape: stream the checkpoint through ``reload_weights`` (the - layerwise reload), with the KDA conv weights scaled.""" - from vllm.model_executor.model_loader import get_model_loader - - runner = worker.model_runner - model = runner.get_model() - loader = get_model_loader(runner.load_config) - - def weights(): - for name, tensor in loader.get_all_weights(runner.model_config, model): - if name.endswith(CKPT_CONV_WEIGHTS): - tensor = tensor * conv_scale - yield name, tensor - - worker.reload_weights(weights_iterator=weights()) - return _conv_weight_ptrs(model) - - -def _refit_kernel_format(worker, conv_scale: float) -> list[int]: - """Kernel-format refit: ``reload_weights(is_checkpoint_format=False)`` - copies into the parameters without any weight loader.""" - model = worker.model_runner.get_model() - weights = [ - (name, param.data * conv_scale) - for name, param in model.named_parameters() - if name.endswith("self_attn.conv1d.weight") - ] - assert weights - worker.reload_weights(weights_iterator=iter(weights), is_checkpoint_format=False) - return _conv_weight_ptrs(model) - - -def _greedy(llm: LLM) -> list[int]: - params = SamplingParams(temperature=0.0, max_tokens=8) - output = llm.generate(["The capital of France is"], params)[0] - return list(output.outputs[0].token_ids) - - -@pytest.mark.skipif( - not ( - current_platform.is_cuda() - and torch.cuda.is_available() - and torch.cuda.get_device_capability()[0] in (9, 10) - ), - reason="GLM-5.3-Flash sparse indexer needs SM90/SM100", -) -@pytest.mark.parametrize("cpu_offload_gb", [0, 10]) -def test_refit_reaches_cudagraph_decode( - monkeypatch: pytest.MonkeyPatch, cpu_offload_gb: int -): - """With CUDA graphs on, decode must follow the refitted conv weights and - return to the original output once they are restored (#55087), including - when the KDA layers are offloaded to CPU.""" - monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") - llm = LLM( - model=PRUNED_MODEL, - max_model_len=2048, - max_num_seqs=4, - gpu_memory_utilization=0.8, - cpu_offload_gb=cpu_offload_gb, - limit_mm_per_prompt={"image": 0, "video": 0}, +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) ) - reference = _greedy(llm) - ptrs = [] - for refit, scale, same in ( - (_refit_from_checkpoint, 2.0, False), - (_refit_from_checkpoint, 1.0, True), - (_refit_kernel_format, 2.0, False), - (_refit_kernel_format, 0.5, True), - ): - ptrs.append(llm.collective_rpc(refit, kwargs={"conv_scale": scale})[0]) - assert (_greedy(llm) == reference) == same, (refit.__name__, scale) - assert all(p == ptrs[0] for p in ptrs) - - del llm - cleanup_dist_env_and_memory() + + 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))