vllm.models.kimi_k3.amd.linear

Classes:

• KimiMLAAttention –

  Main reference: DeepseekV2 vllm Implementation.

• KimiRoutedInputLinear –

  Distribute large prefill projections over identical TP/DCP groups.

KimiMLAAttention

Bases: Module

Main reference: DeepseekV2 vllm Implementation.

Source code in vllm/models/kimi_k3/amd/linear.py

class KimiMLAAttention(nn.Module):
    """Main reference: DeepseekV2 vllm Implementation."""

    def __init__(
        self,
        config: KimiLinearConfig,
        hidden_size: int,
        num_heads: int,
        qk_nope_head_dim: int,
        qk_rope_head_dim: int,
        v_head_dim: int,
        q_lora_rank: int | None,
        kv_lora_rank: int,
        use_nope: bool = False,
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
        **kwargs,
    ) -> None:
        super().__init__()
        self.hidden_size = hidden_size
        self.qk_nope_head_dim = qk_nope_head_dim
        self.qk_rope_head_dim = qk_rope_head_dim
        self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
        self.v_head_dim = v_head_dim
        self.q_lora_rank = q_lora_rank
        self.kv_lora_rank = kv_lora_rank
        self.num_heads = num_heads
        tp_size = get_tensor_model_parallel_world_size()
        self.num_local_heads = num_heads // tp_size
        self.scaling = self.qk_head_dim**-0.5
        self.use_nope = use_nope
        assert self.use_nope is True
        assert num_heads % tp_size == 0
        self.use_output_gate = config.mla_use_output_gate
        if self.q_lora_rank is not None and self.use_output_gate:
            self.fused_qkv_a_proj = KimiK3MergedQKVGateLinear(
                hidden_size=self.hidden_size,
                q_lora_rank=self.q_lora_rank,
                kv_lora_rank=self.kv_lora_rank,
                qk_rope_head_dim=self.qk_rope_head_dim,
                total_num_heads=num_heads,
                v_head_dim=self.v_head_dim,
                bias=False,
                quant_config=quant_config,
                prefix=f"{prefix}.fused_qkv_a_proj",
            )
        elif self.q_lora_rank is not None:
            self.fused_qkv_a_proj = MergedColumnParallelLinear(
                self.hidden_size,
                [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
                bias=False,
                quant_config=quant_config,
                prefix=f"{prefix}.fused_qkv_a_proj",
                disable_tp=True,
            )
        else:
            self.kv_a_proj_with_mqa = ReplicatedLinear(
                self.hidden_size,
                self.kv_lora_rank + self.qk_rope_head_dim,
                bias=False,
                quant_config=quant_config,
                prefix=f"{prefix}.kv_a_proj_with_mqa",
            )
        if self.q_lora_rank is not None:
            self.q_a_layernorm = RMSNorm(
                self.q_lora_rank,
                eps=config.rms_norm_eps,
            )
            self.q_b_proj = ColumnParallelLinear(
                self.q_lora_rank,
                self.num_heads * self.qk_head_dim,
                bias=False,
                quant_config=quant_config,
                prefix=f"{prefix}.q_b_proj",
            )
        else:
            self.q_proj = ColumnParallelLinear(
                self.hidden_size,
                self.num_heads * self.qk_head_dim,
                bias=False,
                quant_config=quant_config,
                prefix=f"{prefix}.q_proj",
            )
        self.kv_a_layernorm = RMSNorm(
            self.kv_lora_rank,
            eps=config.rms_norm_eps,
        )
        self.kv_b_proj = ColumnParallelLinear(
            self.kv_lora_rank,
            self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.kv_b_proj",
        )
        self.o_proj = RowParallelLinear(
            self.num_heads * self.v_head_dim,
            self.hidden_size,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.o_proj",
        )

        if self.use_output_gate and self.q_lora_rank is None:
            projection_size = self.num_heads * self.v_head_dim
            self.g_proj = ColumnParallelLinear(
                self.hidden_size,
                projection_size,
                bias=False,
                quant_config=quant_config,
                prefix=f"{prefix}.g_proj",
            )

        # TODO: Remove this mypy workaround once the K3 PR is fully merged.
        mla_modules = MLAModules(  # type: ignore[call-arg]
            kv_a_layernorm=self.kv_a_layernorm,
            kv_b_proj=self.kv_b_proj,
            rotary_emb=None,
            o_proj=self.o_proj,
            fused_qkv_a_proj=self.fused_qkv_a_proj
            if self.q_lora_rank is not None
            else None,
            kv_a_proj_with_mqa=self.kv_a_proj_with_mqa
            if self.q_lora_rank is None
            else None,
            q_a_layernorm=self.q_a_layernorm if self.q_lora_rank is not None else None,
            q_b_proj=self.q_b_proj if self.q_lora_rank is not None else None,
            q_proj=self.q_proj if self.q_lora_rank is None else None,
            indexer=None,
            is_sparse=False,
            topk_indices_buffer=None,
            g_proj=getattr(self, "g_proj", None),
        )
        self.mla_attn = KimiK3MultiHeadLatentAttentionWrapper(
            self.hidden_size,
            self.num_local_heads,
            self.scaling,
            self.qk_nope_head_dim,
            self.qk_rope_head_dim,
            self.v_head_dim,
            self.q_lora_rank,
            self.kv_lora_rank,
            mla_modules,
            cache_config,
            quant_config,
            prefix,
        )

    def get_input_quant_key(self) -> QuantKey | None:
        return self.mla_attn.get_input_quant_key()

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor | QuantizedActivation,
    ) -> torch.Tensor:
        return self.mla_attn(positions, hidden_states)

KimiRoutedInputLinear

Bases: ReplicatedLinear

Distribute large prefill projections over identical TP/DCP groups.

Methods:

• forward_quantized –

  Exchange the same MXFP4 activations consumed by the routed experts.

Source code in vllm/models/kimi_k3/amd/linear.py

class KimiRoutedInputLinear(ReplicatedLinear):
    """Distribute large prefill projections over identical TP/DCP groups."""

    # Measured crossover on MI355X: smaller batches lose to gather overhead.
    min_sharded_tokens = 8192
    _token_group = None

    def enable_token_sharding(self) -> None:
        from vllm.platforms.rocm import on_gfx950

        config = get_current_vllm_config()
        parallel = config.parallel_config
        if (
            not on_gfx950()
            or config.lora_config is not None
            or parallel.use_ubatching
            or parallel.tensor_parallel_size != 8
            or parallel.decode_context_parallel_size != 8
            or self.input_size != 7168
            or self.output_size != 3584
        ):
            return
        tp, dcp = get_tp_group(), get_dcp_group()
        if tp.ranks == dcp.ranks and tp.rank_in_group == dcp.rank_in_group:
            self._token_group = dcp

    def forward(self, x: torch.Tensor):
        group = self._token_group
        # Small batches keep replicated GEMM: gathering costs more than it saves.
        if (
            group is None
            or x.ndim != 2
            or x.dtype != torch.bfloat16
            or x.shape[0] < self.min_sharded_tokens
        ):
            return super().forward(x)

        projected, bias = self._project_local(x)
        return group.all_gather(projected, dim=0)[: x.shape[0]], bias

    def _project_local(self, x: torch.Tensor):
        group = self._token_group
        assert group is not None
        tokens = x.shape[0]
        chunk = (tokens + group.world_size - 1) // group.world_size
        start = group.rank_in_group * chunk
        local = x[start : min(start + chunk, tokens)]
        if local.shape[0] != chunk:
            local = F.pad(local, (0, 0, 0, chunk - local.shape[0]))
        return super().forward(local)

    def forward_quantized(self, x: torch.Tensor):
        """Exchange the same MXFP4 activations consumed by the routed experts."""
        from aiter import dtypes
        from aiter.ops.quant import per_1x32_mx_quant_hip

        group = self._token_group
        assert group is not None
        projected, bias = self._project_local(x)
        assert bias is None
        values, scales = per_1x32_mx_quant_hip(
            projected, quant_dtype=dtypes.fp4x2, shuffle=False
        )
        packed = torch.cat((values.view(torch.uint8), scales.view(torch.uint8)), dim=1)
        gathered = group.all_gather(packed.view(torch.int32), dim=0)
        gathered = gathered.view(torch.uint8)[: x.shape[0]]
        split = self.output_size // 2
        values = gathered[:, :split].contiguous().view(dtypes.fp4x2)
        scales = gathered[:, split:].contiguous().view(dtypes.fp8_e8m0)
        return values, scales

forward_quantized(x)

Exchange the same MXFP4 activations consumed by the routed experts.

Source code in vllm/models/kimi_k3/amd/linear.py

def forward_quantized(self, x: torch.Tensor):
    """Exchange the same MXFP4 activations consumed by the routed experts."""
    from aiter import dtypes
    from aiter.ops.quant import per_1x32_mx_quant_hip

    group = self._token_group
    assert group is not None
    projected, bias = self._project_local(x)
    assert bias is None
    values, scales = per_1x32_mx_quant_hip(
        projected, quant_dtype=dtypes.fp4x2, shuffle=False
    )
    packed = torch.cat((values.view(torch.uint8), scales.view(torch.uint8)), dim=1)
    gathered = group.all_gather(packed.view(torch.int32), dim=0)
    gathered = gathered.view(torch.uint8)[: x.shape[0]]
    split = self.output_size // 2
    values = gathered[:, :split].contiguous().view(dtypes.fp4x2)
    scales = gathered[:, split:].contiguous().view(dtypes.fp8_e8m0)
    return values, scales