From a2920277af6ddbb2fecae394c09fab5ac44bb1ea Mon Sep 17 00:00:00 2001 From: whx-sjtu Date: Fri, 16 Oct 2026 03:33:52 +0000 Subject: [PATCH 1/3] [Bugfix][ROCm] Fix Kimi K3 online loading of padded KDA projections Declare checkpoint-backed elements so each merged projection completes after its real shards arrive, and initialize declared padding on materialization. Exercise real KDA construction and loading to verify source release and padding preservation across TP ranks and shard orders. Co-authored-by: Codex Signed-off-by: whx-sjtu --- .../model_loader/test_reload.py | 64 +++++++++++++++++++ .../model_loader/reload/meta.py | 3 + vllm/models/kimi_k3/amd/kda.py | 4 ++ 3 files changed, 71 insertions(+) diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index b58c2c350795..41d47fb7e0a2 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -973,6 +973,70 @@ def test_online_processing_waits_for_late_registered_bias(): assert torch.equal(quant_method.bias_at_process, loaded_bias) +@pytest.mark.parametrize("num_heads", [96, 128]) +@pytest.mark.parametrize("tp_rank", [0, 3]) +@pytest.mark.parametrize("reverse", [False, True]) +def test_kda_online_loading_releases_sources_and_zeros_padding( + monkeypatch, default_vllm_config, num_heads, tp_rank, reverse +): + """Finish KDA's merged projection after six inputs, before model finalization.""" + from vllm.models.kimi_k3.amd.kda import KimiK3DeltaAttention + from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig + + monkeypatch.setattr( + "vllm.distributed.parallel_state.get_tp_group", + lambda: SimpleNamespace(world_size=8, rank_in_group=tp_rank), + ) + default_vllm_config.model_config = SimpleNamespace(dtype=torch.float32) + config = KimiLinearConfig( + hidden_size=16, + linear_attn_config={ + "kda_layers": [0], + "full_attn_layers": [], + "head_dim": 128, + "num_heads": num_heads, + "short_conv_kernel_size": 4, + "use_full_rank_gate": True, + }, + ) + attention = KimiK3DeltaAttention( + config, default_vllm_config, prefix="model.layers.0.self_attn" + ) + projection = attention.in_proj_qkvgfab + # Exercise online loading on CPU without selecting a GPU quantization kernel. + projection.weight = to_meta_tensor(projection.weight) + projection.quant_method = None + initialize_online_processing(projection) + + empty_strided = torch.empty_strided + + def poisoned_empty(*args, **kwargs): + return empty_strided(*args, **kwargs).fill_(float("nan")) + + monkeypatch.setattr(torch, "empty_strided", poisoned_empty) + widths = [128 * num_heads] * 4 + [128, num_heads] + expected = {} + sources = [] + for shard in reversed(range(6)) if reverse else range(6): + source = torch.arange(widths[shard] * 16, dtype=torch.float32).reshape(-1, 16) + sources.append(ref(source)) + local_width = widths[shard] if shard == 4 else widths[shard] // 8 + start = 0 if shard == 4 else tp_rank * local_width + expected[shard] = source[start : start + local_width].clone() + assert projection.weight.is_meta + projection.weight.weight_loader(projection.weight, source, shard) + del source + + assert not projection.weight.is_meta + assert all(source() is None for source in sources) + expected_weight = torch.cat([expected[shard] for shard in range(6)]) + torch.testing.assert_close( + projection.weight[: len(expected_weight)], expected_weight + ) + padding = projection.weight[len(expected_weight) :] + assert torch.equal(padding, torch.zeros_like(padding)) + + def test_layerwise_reload_skips_non_persistent_parameter_alias_buffers(monkeypatch): layer = _AliasedBufferLayer() model = torch.nn.Sequential(layer) diff --git a/vllm/model_executor/model_loader/reload/meta.py b/vllm/model_executor/model_loader/reload/meta.py index 9874ebd67be9..3898de844898 100644 --- a/vllm/model_executor/model_loader/reload/meta.py +++ b/vllm/model_executor/model_loader/reload/meta.py @@ -56,6 +56,9 @@ def materialize_meta_tensor(meta_tensor: torch.Tensor) -> torch.Tensor: dtype=meta_tensor.dtype, requires_grad=False, ) + # Initialization on meta does not preserve storage for unloaded padding. + if get_tensor_load_numel(meta_tensor) < meta_tensor.numel(): + tensor.zero_() tensor.__class__ = meta_tensor.__class__ tensor.__dict__ = meta_tensor.__dict__.copy() return tensor diff --git a/vllm/models/kimi_k3/amd/kda.py b/vllm/models/kimi_k3/amd/kda.py index 611a0c5d96b5..a4adeb6babcf 100644 --- a/vllm/models/kimi_k3/amd/kda.py +++ b/vllm/models/kimi_k3/amd/kda.py @@ -155,6 +155,10 @@ def __init__( ) if self.in_proj_padding: self.in_proj_qkvgfab.weight.data[-self.in_proj_padding :].zero_() + # Alignment rows are not part of the checkpoint. + self.in_proj_qkvgfab.weight.weight_loader_numel = ( + local_output_size * self.hidden_size + ) self.f_b_proj = ColumnParallelLinear( self.head_dim, From 17c7e2ee6e68a6b63590fe60a26c281c5f8159a0 Mon Sep 17 00:00:00 2001 From: whx-sjtu Date: Wed, 28 Oct 2026 08:22:37 -0500 Subject: [PATCH 2/3] Remove KDA online loading test The test imports the Kimi K3 KDA layer, which pulls numba into the shared model_executor pytest process and breaks the Triton dispatcher tests that run after it. Signed-off-by: whx-sjtu Co-Authored-By: Claude Opus 5.5 --- .../model_loader/test_reload.py | 64 ------------------- 1 file changed, 64 deletions(-) diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index 95186e4d1390..eac5245bd988 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -977,70 +977,6 @@ def test_online_processing_waits_for_late_registered_bias(): assert torch.equal(quant_method.bias_at_process, loaded_bias) -@pytest.mark.parametrize("num_heads", [96, 128]) -@pytest.mark.parametrize("tp_rank", [0, 3]) -@pytest.mark.parametrize("reverse", [False, True]) -def test_kda_online_loading_releases_sources_and_zeros_padding( - monkeypatch, default_vllm_config, num_heads, tp_rank, reverse -): - """Finish KDA's merged projection after six inputs, before model finalization.""" - from vllm.models.kimi_k3.amd.kda import KimiK3DeltaAttention - from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig - - monkeypatch.setattr( - "vllm.distributed.parallel_state.get_tp_group", - lambda: SimpleNamespace(world_size=8, rank_in_group=tp_rank), - ) - default_vllm_config.model_config = SimpleNamespace(dtype=torch.float32) - config = KimiLinearConfig( - hidden_size=16, - linear_attn_config={ - "kda_layers": [0], - "full_attn_layers": [], - "head_dim": 128, - "num_heads": num_heads, - "short_conv_kernel_size": 4, - "use_full_rank_gate": True, - }, - ) - attention = KimiK3DeltaAttention( - config, default_vllm_config, prefix="model.layers.0.self_attn" - ) - projection = attention.in_proj_qkvgfab - # Exercise online loading on CPU without selecting a GPU quantization kernel. - projection.weight = to_meta_tensor(projection.weight) - projection.quant_method = None - initialize_online_processing(projection) - - empty_strided = torch.empty_strided - - def poisoned_empty(*args, **kwargs): - return empty_strided(*args, **kwargs).fill_(float("nan")) - - monkeypatch.setattr(torch, "empty_strided", poisoned_empty) - widths = [128 * num_heads] * 4 + [128, num_heads] - expected = {} - sources = [] - for shard in reversed(range(6)) if reverse else range(6): - source = torch.arange(widths[shard] * 16, dtype=torch.float32).reshape(-1, 16) - sources.append(ref(source)) - local_width = widths[shard] if shard == 4 else widths[shard] // 8 - start = 0 if shard == 4 else tp_rank * local_width - expected[shard] = source[start : start + local_width].clone() - assert projection.weight.is_meta - projection.weight.weight_loader(projection.weight, source, shard) - del source - - assert not projection.weight.is_meta - assert all(source() is None for source in sources) - expected_weight = torch.cat([expected[shard] for shard in range(6)]) - torch.testing.assert_close( - projection.weight[: len(expected_weight)], expected_weight - ) - padding = projection.weight[len(expected_weight) :] - assert torch.equal(padding, torch.zeros_like(padding)) - - def test_layerwise_reload_skips_non_persistent_parameter_alias_buffers(monkeypatch): layer = _AliasedBufferLayer() model = torch.nn.Sequential(layer) From 15668c30976eb63469ef76bcabefd53996d34eb4 Mon Sep 17 00:00:00 2001 From: whx-sjtu Date: Thu, 29 Oct 2026 00:31:34 -0500 Subject: [PATCH 3/3] Move KDA padding handling into an AMD-local linear subclass Mirror the NVIDIA implementation: the merged projection declares its checkpoint-backed size and re-zeroes padding rows after each weight load, so online processing no longer needs the shared meta.py change. Signed-off-by: whx-sjtu Co-Authored-By: Claude Opus 5.5 --- .../model_loader/reload/meta.py | 3 -- vllm/models/kimi_k3/amd/kda.py | 38 +++++++++++++++---- 2 files changed, 31 insertions(+), 10 deletions(-) diff --git a/vllm/model_executor/model_loader/reload/meta.py b/vllm/model_executor/model_loader/reload/meta.py index 3898de844898..9874ebd67be9 100644 --- a/vllm/model_executor/model_loader/reload/meta.py +++ b/vllm/model_executor/model_loader/reload/meta.py @@ -56,9 +56,6 @@ def materialize_meta_tensor(meta_tensor: torch.Tensor) -> torch.Tensor: dtype=meta_tensor.dtype, requires_grad=False, ) - # Initialization on meta does not preserve storage for unloaded padding. - if get_tensor_load_numel(meta_tensor) < meta_tensor.numel(): - tensor.zero_() tensor.__class__ = meta_tensor.__class__ tensor.__dict__ = meta_tensor.__dict__.copy() return tensor diff --git a/vllm/models/kimi_k3/amd/kda.py b/vllm/models/kimi_k3/amd/kda.py index 8b75e65d1e09..7e193054c5c7 100644 --- a/vllm/models/kimi_k3/amd/kda.py +++ b/vllm/models/kimi_k3/amd/kda.py @@ -73,6 +73,35 @@ logger = init_logger(__name__) +class _PaddedKimiGDNMergedColumnParallelLinear(_KimiGDNMergedColumnParallelLinear): + """Merged projection with ``padding_size`` zero rows appended per rank.""" + + def __init__(self, *args, padding_size: int = 0, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.padding_size = padding_size + if padding_size: + assert isinstance(self.weight, torch.Tensor) + # Exclude padding rows from the load-completion count. + set_weight_attrs( + self.weight, + {"weight_loader_numel": self.weight[:-padding_size].numel()}, + ) + self.weight.data[-padding_size:].zero_() + + def weight_loader(self, param, loaded_weight, loaded_shard_id=None) -> None: + super().weight_loader(param, loaded_weight, loaded_shard_id) + self._zero_padding(param) + + def weight_loader_v2(self, param, loaded_weight, loaded_shard_id=None) -> None: + super().weight_loader_v2(param, loaded_weight, loaded_shard_id) + self._zero_padding(param) + + def _zero_padding(self, param: torch.Tensor) -> None: + # Online processing materializes the weight after __init__ zeroing. + if self.padding_size and param is self.weight: + param.data[-self.padding_size :].zero_() + + class KimiK3DeltaAttention(GatedDeltaNetAttention): def get_attn_backend(self) -> type[AttentionBackend]: return KimiK3ROCmKDABackend @@ -159,21 +188,16 @@ def __init__( if self.in_proj_padding: in_proj_output_sizes.append(self.in_proj_padding * self.tp_size) - self.in_proj_qkvgfab = _KimiGDNMergedColumnParallelLinear( + self.in_proj_qkvgfab = _PaddedKimiGDNMergedColumnParallelLinear( self.hidden_size, in_proj_output_sizes, replicated_shard_id=4, tp_size=self.tp_size, + padding_size=self.in_proj_padding, bias=False, quant_config=self.quant_config, prefix=f"{prefix}.in_proj_qkvgfab", ) - if self.in_proj_padding: - self.in_proj_qkvgfab.weight.data[-self.in_proj_padding :].zero_() - # Alignment rows are not part of the checkpoint. - self.in_proj_qkvgfab.weight.weight_loader_numel = ( - local_output_size * self.hidden_size - ) self.f_b_proj = ColumnParallelLinear( self.head_dim,