From 6ec92bcbc8ef9af25cb4834ba3d922115792fbeb Mon Sep 17 00:00:00 2001 From: Thien Tran Date: Fri, 28 Aug 2026 08:47:18 +0800 Subject: [PATCH] [Kimi-K3] Merge MLA gate into QKV-A projection (#54015) Signed-off-by: Thien Tran --- vllm/model_executor/layers/linear.py | 91 ++++++++- vllm/models/kimi_k3/nvidia/mla.py | 267 +++++++++++++++++---------- vllm/models/kimi_k3/nvidia/model.py | 16 +- vllm/models/kimi_k3/nvidia/mtp.py | 15 +- 4 files changed, 283 insertions(+), 106 deletions(-) diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 39033f99f2..f51d16c071 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from abc import abstractmethod -from collections.abc import Iterable +from collections.abc import Callable, Iterable from typing import Any import torch @@ -1512,6 +1512,95 @@ class MinimaxM3QKVParallelLinearWithIndexer(QKVParallelLinear): param_data.copy_(loaded_weight) +class KimiK3MergedQKVGateLinear(MergedColumnParallelLinear): + """Kimi-K3 QKV-A projection fused with its output-gate projection. + + NOTE: Kimi-K3-specific. This is tailored to its replicated Q/KV latent + projections and tensor-parallel output gate; it is not a general-purpose + linear layer. It lives here alongside ``MergedColumnParallelLinear``, whose + sharding and weight-loading machinery it reuses. + + The per-rank shard layout is ``[q_a | kv_a | gate]``. The latent shards + stay replicated across tensor-parallel ranks while the gate shard is + tensor-parallel, matching the standalone projections they replace. + """ + + def __init__( + self, + hidden_size: int, + q_lora_rank: int, + kv_lora_rank: int, + qk_rope_head_dim: int, + total_num_heads: int, + v_head_dim: int, + bias: bool = False, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + tp_size = get_tensor_model_parallel_world_size() + + # MergedColumnParallelLinear divides every logical shard by TP size. + # Inflate the replicated latent shards to retain their checkpoint widths. + output_sizes = [ + q_lora_rank * tp_size, + (kv_lora_rank + qk_rope_head_dim) * tp_size, + total_num_heads * v_head_dim, + ] + super().__init__( + input_size=hidden_size, + output_sizes=output_sizes, + bias=bias, + quant_config=quant_config, + prefix=prefix, + ) + + def _load_with_replicated_latents( + self, + loader: Callable[..., None], + param: Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: tuple[int, ...] | int | None, + ) -> None: + is_replicated = isinstance(loaded_shard_id, int) and loaded_shard_id < 2 + if not is_replicated: + loader(param, loaded_weight, loaded_shard_id) + return + + # The base loader narrows every checkpoint shard using its TP rank. + # Replicated latent shards instead load the rank-zero slice on every rank. + tp_rank = self.tp_rank + param_tp_rank = getattr(param, "tp_rank", None) + self.tp_rank = 0 + if param_tp_rank is not None: + param.tp_rank = 0 + try: + loader(param, loaded_weight, loaded_shard_id) + finally: + self.tp_rank = tp_rank + if param_tp_rank is not None: + param.tp_rank = param_tp_rank + + def weight_loader( + self, + param: Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: tuple[int, ...] | int | None = None, + ) -> None: + self._load_with_replicated_latents( + super().weight_loader, param, loaded_weight, loaded_shard_id + ) + + def weight_loader_v2( + self, + param: BasevLLMParameter, + loaded_weight: torch.Tensor, + loaded_shard_id: tuple[int, ...] | int | None = None, + ) -> None: + self._load_with_replicated_latents( + super().weight_loader_v2, param, loaded_weight, loaded_shard_id + ) + + # --8<-- [start:row_parallel_linear] @PluggableLayer.register("row_parallel_linear") class RowParallelLinear(LinearBase): diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index 21be0c398f..45ae01606c 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -5,7 +5,8 @@ This is a self-contained MLA layer that owns the full attention path: hidden_states - -> fused pre-attention ops (fused_qkv_a_proj / norms / q_b_proj) + -> fused pre-attention ops (fused_qkv_a_proj / norms / q_b_proj; + with an output gate the rows use fused_qkv_a_g_proj instead) -> explicit prefill / decode split prefill: fused key-concat + cache-insert kernel -> run_prefill_new_tokens (+ chunked-context merge, whose per-chunk gather -> kv_b_proj @@ -23,7 +24,8 @@ KV cache, and absorbs ``kv_b_proj`` into ``W_UK_T`` / ``W_UV`` -- mirroring the ``DeepseekV4Attention`` structure. K3 specifics: optional rotary embedding (disabled for the target model's NoPE -layers, enabled for DSpark) and an optional sigmoid output gate (``g_proj``). +layers, enabled for DSpark) and an optional sigmoid output gate (``g_proj``, +merged into ``fused_qkv_a_g_proj`` when the q-LoRA front-end is present). Out of scope (extension points, not wired here): prefill context parallelism (PCP), sparse/indexer MLA, and the ROCm/aiter fp8/fp4 BMM fast paths. @@ -33,6 +35,7 @@ import math from typing import TYPE_CHECKING, cast import torch +import torch.nn.functional as F from torch import nn from vllm import _custom_ops as ops @@ -60,9 +63,11 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, + KimiK3MergedQKVGateLinear, MergedColumnParallelLinear, ReplicatedLinear, RowParallelLinear, + UnquantizedLinearMethod, ) from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.quantization.utils.quant_utils import ( @@ -71,6 +76,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding, get_rope from vllm.model_executor.utils import replace_parameter from vllm.models.common.ops import fused_q_kv_rmsnorm +from vllm.models.kimi_k3.nvidia.low_latency_gemm import try_low_latency_gemm from vllm.models.kimi_k3.nvidia.ops.fused_mla_key_concat_kv_cache import ( fused_mla_decode_q_concat_kv_cache_insert, fused_mla_key_concat_ds_mla_insert, @@ -102,15 +108,14 @@ if TYPE_CHECKING: logger = init_logger(__name__) -# Below this many tokens, overlap the g_proj GEMM on the aux stream with the -# attention front-end (the GEMM is small and launch-bound, so the overlap -# hides it); at or above it, run the gate on the main stream. +# Below this conservative threshold, overlap the gate projection with attention +# on the auxiliary stream. _GATE_MULTI_STREAM_TOKEN_THRESHOLD = 512 @torch.compile(backend=current_platform.simple_compile_backend) def _gate_sigmoid_mul(attn_out: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: - """Apply the sigmoid output gate to a precomputed ``g_proj`` projection.""" + """Apply the sigmoid output gate to a precomputed gate projection.""" return attn_out * gate.sigmoid() @@ -144,6 +149,7 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase): self.v_head_dim = v_head_dim self.q_lora_rank = q_lora_rank self.kv_lora_rank = kv_lora_rank + self.use_output_gate = use_output_gate self.non_causal_multi_token_decode = non_causal_multi_token_decode # Latent "head" seen by the attention kernel / KV cache. self.head_size = kv_lora_rank + qk_rope_head_dim @@ -192,22 +198,52 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase): self.num_local_heads = num_heads // tp_size # ---- Pre-attention projections (fusable front-end) ---- - # Two query variants: a low-rank q-LoRA path (Kimi-K3) fused with the - # kv-down proj, or an uncompressed q path (Kimi-Linear, ``q_lora_rank`` - # None) with a standalone ``q_proj`` and separate ``kv_a_proj_with_mqa``. + # Exactly one query layout is constructed below. Kimi-K3 uses q-LoRA; + # Kimi-Linear uses independent full-rank Q and KV projections. + self.fused_qkv_a_proj: MergedColumnParallelLinear | None = None + self.fused_qkv_a_g_proj: KimiK3MergedQKVGateLinear | None = None + self.q_proj: ColumnParallelLinear | None = None + self.kv_a_proj_with_mqa: ReplicatedLinear | None = None + self.g_proj: ColumnParallelLinear | None = None + self.aux_stream: torch.cuda.Stream | None = None + self._gate_events: tuple[torch.cuda.Event, torch.cuda.Event] | None = None + + # Kimi-K3: low-rank Q and KV projections share one input projection. if self.q_lora_rank is not None: - # Fused q-down + kv-down projection. Replicated (disable_tp) because - # the low-rank latents are shared across TP ranks; TP splitting - # happens at q_b_proj / kv_b_proj. Checkpoint weights ``q_a_proj`` - # and ``kv_a_proj_with_mqa`` map onto shards 0 and 1 respectively. - 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, - ) + # The latent Q and KV shards are replicated across TP ranks. Their + # TP splitting happens in q_b_proj and kv_b_proj instead. + if use_output_gate: + # The output gate is shard 2 of the same input projection. + self.fused_qkv_a_g_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_g_proj", + ) + # Only supports unquantized weights for sliced side-stream path + # for now. + if aux_stream is not None and isinstance( + self.fused_qkv_a_g_proj.quant_method, UnquantizedLinearMethod + ): + self.aux_stream = aux_stream + self._gate_events = (torch.cuda.Event(), torch.cuda.Event()) + else: + 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, + ) self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps) self.q_b_proj = ColumnParallelLinear( self.q_lora_rank, @@ -216,9 +252,9 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase): quant_config=quant_config, prefix=f"{prefix}.q_b_proj", ) + + # Kimi-Linear: Q and KV use independent full-rank input projections. else: - # Uncompressed query: full-rank q_proj (TP-split over heads) plus a - # replicated kv-down projection (shared latent across TP ranks). self.q_proj = ColumnParallelLinear( self.hidden_size, self.num_heads * self.qk_head_dim, @@ -233,6 +269,16 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase): quant_config=quant_config, prefix=f"{prefix}.kv_a_proj_with_mqa", ) + if use_output_gate: + # Keep the gate standalone for the non-LoRA layout. + self.g_proj = ColumnParallelLinear( + self.hidden_size, + self.num_heads * self.v_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.g_proj", + ) + self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps) self.kv_b_proj = ColumnParallelLinear( self.kv_lora_rank, @@ -242,28 +288,7 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase): prefix=f"{prefix}.kv_b_proj", ) - # ---- Post-attention projections ---- - self.use_output_gate = use_output_gate - self.g_proj = ( - ColumnParallelLinear( - self.hidden_size, - self.num_heads * self.v_head_dim, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.g_proj", - ) - if use_output_gate - else None - ) - # Aux stream (created at the model level, DeepseekV4 convention) for - # overlapping the g_proj GEMM with the attention front-end. None on - # ROCm/non-cuda -> maybe_execute_in_parallel falls back to sequential. - self.aux_stream = aux_stream - self._gate_events = ( - [torch.cuda.Event(), torch.cuda.Event()] - if self.g_proj is not None and current_platform.is_cuda_alike() - else None - ) + # ---- Post-attention projection ---- self.o_proj = RowParallelLinear( self.num_heads * self.v_head_dim, self.hidden_size, @@ -483,76 +508,124 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase): # ------------------------------------------------------------------ # Forward # ------------------------------------------------------------------ - def _forward_attn( + def _unquantized_gemm( + self, hidden_states: torch.Tensor, weight: torch.Tensor + ) -> torch.Tensor: + out = try_low_latency_gemm(hidden_states, weight) + return out if out is not None else F.linear(hidden_states, weight) + + def _apply_q_lora_attention( self, positions: torch.Tensor, hidden_states: torch.Tensor, + qkv_lora: torch.Tensor, ) -> torch.Tensor: - """Attention front-end: fused qkv-a proj -> norms -> q_b -> attention. - - Returns the pre-gate attention output ``[num_tokens, - num_local_heads * v_head_dim]``. On a profile/dummy run - it returns a zeroed buffer. - """ - if self.q_lora_rank is not None: - qkv_lora = self.fused_qkv_a_proj(hidden_states)[0] - q_c, kv_c, k_pe = qkv_lora.split( - [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 - ) - q_c, kv_c_normed = fused_q_kv_rmsnorm( - q_c, - kv_c, - self.q_a_layernorm.weight.data, - self.kv_a_layernorm.weight.data, - self.rms_norm_eps, - ) - q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim) - else: - # Uncompressed query: project directly (no q-LoRA, no q norm) and - # normalize only the kv latent. - q = self.q_proj(hidden_states)[0].view( - -1, self.num_local_heads, self.qk_head_dim - ) - kv_lora = self.kv_a_proj_with_mqa(hidden_states)[0] - kv_c, k_pe = kv_lora.split( - [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 - ) - kv_c_normed = self.kv_a_layernorm(kv_c) - k_pe = k_pe.unsqueeze(1) + """Expand q-LoRA output and run attention.""" + q_lora_rank = self.q_lora_rank + assert q_lora_rank is not None + q_c, kv_c, k_pe = qkv_lora.split( + [q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + q_c, kv_c_normed = fused_q_kv_rmsnorm( + q_c, + kv_c, + self.q_a_layernorm.weight.data, + self.kv_a_layernorm.weight.data, + self.rms_norm_eps, + ) + q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim) attn_out = torch.empty( (hidden_states.shape[0], self.num_local_heads * self.v_head_dim), dtype=hidden_states.dtype, device=hidden_states.device, ) - self._attention(positions, q, kv_c_normed, k_pe, attn_out) + self._attention(positions, q, kv_c_normed, k_pe.unsqueeze(1), attn_out) return attn_out + def _forward_q_lora( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Run the q-LoRA path, including gate overlap when eligible.""" + q_lora_rank = self.q_lora_rank + assert q_lora_rank is not None + qkv_a_rows = q_lora_rank + self.kv_lora_rank + self.qk_rope_head_dim + + # run g_proj in aux_stream + if ( + self._gate_events is not None + and hidden_states.shape[0] < _GATE_MULTI_STREAM_TOKEN_THRESHOLD + ): + fused_qkv_a_g_proj = self.fused_qkv_a_g_proj + assert fused_qkv_a_g_proj is not None + qkv_a_weight = fused_qkv_a_g_proj.weight[:qkv_a_rows] + gate_weight = fused_qkv_a_g_proj.weight[qkv_a_rows:] + + start_event, end_event = self._gate_events + return maybe_execute_in_parallel( + lambda: self._apply_q_lora_attention( + positions, + hidden_states, + self._unquantized_gemm(hidden_states, qkv_a_weight), + ), + lambda: self._unquantized_gemm(hidden_states, gate_weight), + start_event, + end_event, + self.aux_stream, + ) + + # run g_proj together with qkv_a + if self.use_output_gate: + qkv_a_proj = self.fused_qkv_a_g_proj + assert qkv_a_proj is not None + qkv_lora, gate = qkv_a_proj(hidden_states)[0].split( + [qkv_a_rows, self.num_local_heads * self.v_head_dim], dim=-1 + ) + else: + qkv_a_proj = self.fused_qkv_a_proj + assert qkv_a_proj is not None + qkv_lora = qkv_a_proj(hidden_states)[0] + gate = None + + attn_out = self._apply_q_lora_attention(positions, hidden_states, qkv_lora) + return attn_out, gate + + def _forward_full_rank_q( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Run the full-rank Q path with an optional sequential gate.""" + q_proj = self.q_proj + kv_a_proj_with_mqa = self.kv_a_proj_with_mqa + assert q_proj is not None + assert kv_a_proj_with_mqa is not None + + q = q_proj(hidden_states)[0].view(-1, self.num_local_heads, self.qk_head_dim) + kv_lora = kv_a_proj_with_mqa(hidden_states)[0] + kv_c, k_pe = kv_lora.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + kv_c_normed = self.kv_a_layernorm(kv_c) + attn_out = torch.empty( + (hidden_states.shape[0], self.num_local_heads * self.v_head_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + self._attention(positions, q, kv_c_normed, k_pe.unsqueeze(1), attn_out) + + gate = self.g_proj(hidden_states)[0] if self.g_proj is not None else None + return attn_out, gate + def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, ) -> torch.Tensor: - # Both branches produce (attn_out, gate); they differ only in whether - # the g_proj GEMM is overlapped on the aux stream. - g_proj = self.g_proj - events = self._gate_events - if ( - g_proj is not None - and events is not None - and self.aux_stream is not None - and hidden_states.shape[0] < _GATE_MULTI_STREAM_TOKEN_THRESHOLD - ): - attn_out, gate = maybe_execute_in_parallel( - lambda: self._forward_attn(positions, hidden_states), - lambda: g_proj(hidden_states)[0], - events[0], - events[1], - self.aux_stream, - ) + if self.q_lora_rank is None: + attn_out, gate = self._forward_full_rank_q(positions, hidden_states) else: - attn_out = self._forward_attn(positions, hidden_states) - gate = g_proj(hidden_states)[0] if g_proj is not None else None + attn_out, gate = self._forward_q_lora(positions, hidden_states) if gate is not None: attn_out = _gate_sigmoid_mul(attn_out, gate) diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 15ac51cd41..dbf998fc27 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -1120,6 +1120,7 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant): "in_proj_qkvgfab": ["q_proj", "k_proj", "v_proj", "b_proj", "f_a_proj"], "conv1d": ["q_conv1d", "k_conv1d", "v_conv1d"], "fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"], + "fused_qkv_a_g_proj": ["q_a_proj", "kv_a_proj_with_mqa", "g_proj"], } def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -1460,10 +1461,17 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant): if use_full_rank_gate: stacked_params_mapping.append((".in_proj_qkvgfab", ".g_proj", 3)) if getattr(self.config, "q_lora_rank", None) is not None: - stacked_params_mapping += [ - (".fused_qkv_a_proj", ".q_a_proj", 0), - (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1), - ] + if self.config.mla_use_output_gate: + stacked_params_mapping += [ + (".fused_qkv_a_g_proj", ".q_a_proj", 0), + (".fused_qkv_a_g_proj", ".kv_a_proj_with_mqa", 1), + (".fused_qkv_a_g_proj", ".g_proj", 2), + ] + else: + stacked_params_mapping += [ + (".fused_qkv_a_proj", ".q_a_proj", 0), + (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1), + ] use_mega_moe = any( module.use_mega_moe for module in self.modules() diff --git a/vllm/models/kimi_k3/nvidia/mtp.py b/vllm/models/kimi_k3/nvidia/mtp.py index 1677700202..038c8db6b3 100644 --- a/vllm/models/kimi_k3/nvidia/mtp.py +++ b/vllm/models/kimi_k3/nvidia/mtp.py @@ -272,10 +272,17 @@ class KimiK3MTP(nn.Module): if use_full_rank_gate: stacked_params_mapping.append((".in_proj_qkvgfab", ".g_proj", 3)) if getattr(self.config, "q_lora_rank", None) is not None: - stacked_params_mapping += [ - (".fused_qkv_a_proj", ".q_a_proj", 0), - (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1), - ] + if self.config.mla_use_output_gate: + stacked_params_mapping += [ + (".fused_qkv_a_g_proj", ".q_a_proj", 0), + (".fused_qkv_a_g_proj", ".kv_a_proj_with_mqa", 1), + (".fused_qkv_a_g_proj", ".g_proj", 2), + ] + else: + stacked_params_mapping += [ + (".fused_qkv_a_proj", ".q_a_proj", 0), + (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1), + ] use_mega_moe = any( module.use_mega_moe