diff --git a/vllm/models/kimi_k3/amd/kda.py b/vllm/models/kimi_k3/amd/kda.py index 48ac78338571..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,17 +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_() self.f_b_proj = ColumnParallelLinear( self.head_dim,