Skip to content

vllm.models.kimi_k3.amd.kda

_PaddedKimiGDNMergedColumnParallelLinear

Bases: _KimiGDNMergedColumnParallelLinear

Merged projection with padding_size zero rows appended per rank.

Source code in vllm/models/kimi_k3/amd/kda.py 
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
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_()