@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 6ec92bcbc8

[Kimi-K3] Merge MLA gate into QKV-A projection (#54015)

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>

Thien Tran committed Aug 28, 2026 at 08:47 UTC 6ec92bcbc8ef9af25cb4834ba3d922115792fbeb

4 files changed +280 -103

vllm/model_executor/layers/linear.py

+90 -1

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 from abc import abstractmethod

5 -from collections.abc import Iterable

5 +from collections.abc import Callable, Iterable

6 from typing import Any

8 import torch

@@ -1512,6 +1512,95 @@ class MinimaxM3QKVParallelLinearWithIndexer(QKVParallelLinear):

1512 param_data.copy_(loaded_weight)

1513

1514

1515 +class KimiK3MergedQKVGateLinear(MergedColumnParallelLinear):

1516 + """Kimi-K3 QKV-A projection fused with its output-gate projection.

1517 +

1518 + NOTE: Kimi-K3-specific. This is tailored to its replicated Q/KV latent

1519 + projections and tensor-parallel output gate; it is not a general-purpose

1520 + linear layer. It lives here alongside ``MergedColumnParallelLinear``, whose

1521 + sharding and weight-loading machinery it reuses.

1522 +

1523 + The per-rank shard layout is ``[q_a | kv_a | gate]``. The latent shards

1524 + stay replicated across tensor-parallel ranks while the gate shard is

1525 + tensor-parallel, matching the standalone projections they replace.

1526 + """

1527 +

1528 + def __init__(

1529 + self,

1530 + hidden_size: int,

1531 + q_lora_rank: int,

1532 + kv_lora_rank: int,

1533 + qk_rope_head_dim: int,

1534 + total_num_heads: int,

1535 + v_head_dim: int,

1536 + bias: bool = False,

1537 + quant_config: QuantizationConfig | None = None,

1538 + prefix: str = "",

1539 + ) -> None:

1540 + tp_size = get_tensor_model_parallel_world_size()

1541 +

1542 + # MergedColumnParallelLinear divides every logical shard by TP size.

1543 + # Inflate the replicated latent shards to retain their checkpoint widths.

1544 + output_sizes = [

1545 + q_lora_rank * tp_size,

1546 + (kv_lora_rank + qk_rope_head_dim) * tp_size,

1547 + total_num_heads * v_head_dim,

1548 + ]

1549 + super().__init__(

1550 + input_size=hidden_size,

1551 + output_sizes=output_sizes,

1552 + bias=bias,

1553 + quant_config=quant_config,

1554 + prefix=prefix,

1555 + )

1556 +

1557 + def _load_with_replicated_latents(

1558 + self,

1559 + loader: Callable[..., None],

1560 + param: Parameter,

1561 + loaded_weight: torch.Tensor,

1562 + loaded_shard_id: tuple[int, ...] | int | None,

1563 + ) -> None:

1564 + is_replicated = isinstance(loaded_shard_id, int) and loaded_shard_id < 2

1565 + if not is_replicated:

1566 + loader(param, loaded_weight, loaded_shard_id)

1567 + return

1568 +

1569 + # The base loader narrows every checkpoint shard using its TP rank.

1570 + # Replicated latent shards instead load the rank-zero slice on every rank.

1571 + tp_rank = self.tp_rank

1572 + param_tp_rank = getattr(param, "tp_rank", None)

1573 + self.tp_rank = 0

1574 + if param_tp_rank is not None:

1575 + param.tp_rank = 0

1576 + try:

1577 + loader(param, loaded_weight, loaded_shard_id)

1578 + finally:

1579 + self.tp_rank = tp_rank

1580 + if param_tp_rank is not None:

1581 + param.tp_rank = param_tp_rank

1582 +

1583 + def weight_loader(

1584 + self,

1585 + param: Parameter,

1586 + loaded_weight: torch.Tensor,

1587 + loaded_shard_id: tuple[int, ...] | int | None = None,

1588 + ) -> None:

1589 + self._load_with_replicated_latents(

1590 + super().weight_loader, param, loaded_weight, loaded_shard_id

1591 + )

1592 +

1593 + def weight_loader_v2(

1594 + self,

1595 + param: BasevLLMParameter,

1596 + loaded_weight: torch.Tensor,

1597 + loaded_shard_id: tuple[int, ...] | int | None = None,

1598 + ) -> None:

1599 + self._load_with_replicated_latents(

1600 + super().weight_loader_v2, param, loaded_weight, loaded_shard_id

1601 + )

1602 +

1603 +

1604 # --8<-- [start:row_parallel_linear]

1605 @PluggableLayer.register("row_parallel_linear")

1606 class RowParallelLinear(LinearBase):

vllm/models/kimi_k3/nvidia/mla.py

+167 -94

@@ -5,7 +5,8 @@

5 This is a self-contained MLA layer that owns the full attention path:

7 hidden_states

8 - -> fused pre-attention ops (fused_qkv_a_proj / norms / q_b_proj)

8 + -> fused pre-attention ops (fused_qkv_a_proj / norms / q_b_proj;

9 + with an output gate the rows use fused_qkv_a_g_proj instead)

10 -> explicit prefill / decode split

11 prefill: fused key-concat + cache-insert kernel -> run_prefill_new_tokens

12 (+ 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

24 ``DeepseekV4Attention`` structure.

26 K3 specifics: optional rotary embedding (disabled for the target model's NoPE

26 -layers, enabled for DSpark) and an optional sigmoid output gate (``g_proj``).

27 +layers, enabled for DSpark) and an optional sigmoid output gate (``g_proj``,

28 +merged into ``fused_qkv_a_g_proj`` when the q-LoRA front-end is present).

30 Out of scope (extension points, not wired here): prefill context parallelism

31 (PCP), sparse/indexer MLA, and the ROCm/aiter fp8/fp4 BMM fast paths.

@@ -33,6 +35,7 @@ import math

35 from typing import TYPE_CHECKING, cast

37 import torch

38 +import torch.nn.functional as F

39 from torch import nn

41 from vllm import _custom_ops as ops

@@ -60,9 +63,11 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

63 from vllm.model_executor.layers.layernorm import RMSNorm

64 from vllm.model_executor.layers.linear import (

65 ColumnParallelLinear,

66 + KimiK3MergedQKVGateLinear,

67 MergedColumnParallelLinear,

68 ReplicatedLinear,

69 RowParallelLinear,

70 + UnquantizedLinearMethod,

71 )

72 from vllm.model_executor.layers.quantization import QuantizationConfig

73 from vllm.model_executor.layers.quantization.utils.quant_utils import (

@@ -71,6 +76,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

76 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding, get_rope

77 from vllm.model_executor.utils import replace_parameter

78 from vllm.models.common.ops import fused_q_kv_rmsnorm

79 +from vllm.models.kimi_k3.nvidia.low_latency_gemm import try_low_latency_gemm

80 from vllm.models.kimi_k3.nvidia.ops.fused_mla_key_concat_kv_cache import (

81 fused_mla_decode_q_concat_kv_cache_insert,

82 fused_mla_key_concat_ds_mla_insert,

@@ -102,15 +108,14 @@ if TYPE_CHECKING:

108

109 logger = init_logger(__name__)

110

105 -# Below this many tokens, overlap the g_proj GEMM on the aux stream with the

106 -# attention front-end (the GEMM is small and launch-bound, so the overlap

107 -# hides it); at or above it, run the gate on the main stream.

111 +# Below this conservative threshold, overlap the gate projection with attention

112 +# on the auxiliary stream.

113 _GATE_MULTI_STREAM_TOKEN_THRESHOLD = 512

114

115

116 @torch.compile(backend=current_platform.simple_compile_backend)

117 def _gate_sigmoid_mul(attn_out: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:

113 - """Apply the sigmoid output gate to a precomputed ``g_proj`` projection."""

118 + """Apply the sigmoid output gate to a precomputed gate projection."""

119 return attn_out * gate.sigmoid()

120

121

@@ -144,6 +149,7 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):

149 self.v_head_dim = v_head_dim

150 self.q_lora_rank = q_lora_rank

151 self.kv_lora_rank = kv_lora_rank

152 + self.use_output_gate = use_output_gate

153 self.non_causal_multi_token_decode = non_causal_multi_token_decode

154 # Latent "head" seen by the attention kernel / KV cache.

155 self.head_size = kv_lora_rank + qk_rope_head_dim

@@ -192,22 +198,52 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):

198 self.num_local_heads = num_heads // tp_size

199

200 # ---- Pre-attention projections (fusable front-end) ----

195 - # Two query variants: a low-rank q-LoRA path (Kimi-K3) fused with the

196 - # kv-down proj, or an uncompressed q path (Kimi-Linear, ``q_lora_rank``

197 - # None) with a standalone ``q_proj`` and separate ``kv_a_proj_with_mqa``.

201 + # Exactly one query layout is constructed below. Kimi-K3 uses q-LoRA;

202 + # Kimi-Linear uses independent full-rank Q and KV projections.

203 + self.fused_qkv_a_proj: MergedColumnParallelLinear | None = None

204 + self.fused_qkv_a_g_proj: KimiK3MergedQKVGateLinear | None = None

205 + self.q_proj: ColumnParallelLinear | None = None

206 + self.kv_a_proj_with_mqa: ReplicatedLinear | None = None

207 + self.g_proj: ColumnParallelLinear | None = None

208 + self.aux_stream: torch.cuda.Stream | None = None

209 + self._gate_events: tuple[torch.cuda.Event, torch.cuda.Event] | None = None

210 +

211 + # Kimi-K3: low-rank Q and KV projections share one input projection.

212 if self.q_lora_rank is not None:

199 - # Fused q-down + kv-down projection. Replicated (disable_tp) because

200 - # the low-rank latents are shared across TP ranks; TP splitting

201 - # happens at q_b_proj / kv_b_proj. Checkpoint weights ``q_a_proj``

202 - # and ``kv_a_proj_with_mqa`` map onto shards 0 and 1 respectively.

203 - self.fused_qkv_a_proj = MergedColumnParallelLinear(

204 - self.hidden_size,

205 - [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],

206 - bias=False,

207 - quant_config=quant_config,

208 - prefix=f"{prefix}.fused_qkv_a_proj",

209 - disable_tp=True,

210 - )

213 + # The latent Q and KV shards are replicated across TP ranks. Their

214 + # TP splitting happens in q_b_proj and kv_b_proj instead.

215 + if use_output_gate:

216 + # The output gate is shard 2 of the same input projection.

217 + self.fused_qkv_a_g_proj = KimiK3MergedQKVGateLinear(

218 + hidden_size=self.hidden_size,

219 + q_lora_rank=self.q_lora_rank,

220 + kv_lora_rank=self.kv_lora_rank,

221 + qk_rope_head_dim=self.qk_rope_head_dim,

222 + total_num_heads=num_heads,

223 + v_head_dim=self.v_head_dim,

224 + bias=False,

225 + quant_config=quant_config,

226 + prefix=f"{prefix}.fused_qkv_a_g_proj",

227 + )

228 + # Only supports unquantized weights for sliced side-stream path

229 + # for now.

230 + if aux_stream is not None and isinstance(

231 + self.fused_qkv_a_g_proj.quant_method, UnquantizedLinearMethod

232 + ):

233 + self.aux_stream = aux_stream

234 + self._gate_events = (torch.cuda.Event(), torch.cuda.Event())

235 + else:

236 + self.fused_qkv_a_proj = MergedColumnParallelLinear(

237 + self.hidden_size,

238 + [

239 + self.q_lora_rank,

240 + self.kv_lora_rank + self.qk_rope_head_dim,

241 + ],

242 + bias=False,

243 + quant_config=quant_config,

244 + prefix=f"{prefix}.fused_qkv_a_proj",

245 + disable_tp=True,

246 + )

247 self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)

248 self.q_b_proj = ColumnParallelLinear(

249 self.q_lora_rank,

@@ -216,9 +252,9 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):

252 quant_config=quant_config,

253 prefix=f"{prefix}.q_b_proj",

254 )

255 +

256 + # Kimi-Linear: Q and KV use independent full-rank input projections.

257 else:

220 - # Uncompressed query: full-rank q_proj (TP-split over heads) plus a

221 - # replicated kv-down projection (shared latent across TP ranks).

258 self.q_proj = ColumnParallelLinear(

259 self.hidden_size,

260 self.num_heads * self.qk_head_dim,

@@ -233,6 +269,16 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):

269 quant_config=quant_config,

270 prefix=f"{prefix}.kv_a_proj_with_mqa",

271 )

272 + if use_output_gate:

273 + # Keep the gate standalone for the non-LoRA layout.

274 + self.g_proj = ColumnParallelLinear(

275 + self.hidden_size,

276 + self.num_heads * self.v_head_dim,

277 + bias=False,

278 + quant_config=quant_config,

279 + prefix=f"{prefix}.g_proj",

280 + )

281 +

282 self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)

283 self.kv_b_proj = ColumnParallelLinear(

284 self.kv_lora_rank,

@@ -242,28 +288,7 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):

288 prefix=f"{prefix}.kv_b_proj",

289 )

290

245 - # ---- Post-attention projections ----

246 - self.use_output_gate = use_output_gate

247 - self.g_proj = (

248 - ColumnParallelLinear(

249 - self.hidden_size,

250 - self.num_heads * self.v_head_dim,

251 - bias=False,

252 - quant_config=quant_config,

253 - prefix=f"{prefix}.g_proj",

254 - )

255 - if use_output_gate

256 - else None

257 - )

258 - # Aux stream (created at the model level, DeepseekV4 convention) for

259 - # overlapping the g_proj GEMM with the attention front-end. None on

260 - # ROCm/non-cuda -> maybe_execute_in_parallel falls back to sequential.

261 - self.aux_stream = aux_stream

262 - self._gate_events = (

263 - [torch.cuda.Event(), torch.cuda.Event()]

264 - if self.g_proj is not None and current_platform.is_cuda_alike()

265 - else None

266 - )

291 + # ---- Post-attention projection ----

292 self.o_proj = RowParallelLinear(

293 self.num_heads * self.v_head_dim,

294 self.hidden_size,

@@ -483,76 +508,124 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):

508 # ------------------------------------------------------------------

509 # Forward

510 # ------------------------------------------------------------------

486 - def _forward_attn(

511 + def _unquantized_gemm(

512 + self, hidden_states: torch.Tensor, weight: torch.Tensor

513 + ) -> torch.Tensor:

514 + out = try_low_latency_gemm(hidden_states, weight)

515 + return out if out is not None else F.linear(hidden_states, weight)

516 +

517 + def _apply_q_lora_attention(

518 self,

519 positions: torch.Tensor,

520 hidden_states: torch.Tensor,

521 + qkv_lora: torch.Tensor,

522 ) -> torch.Tensor:

491 - """Attention front-end: fused qkv-a proj -> norms -> q_b -> attention.

492 -

493 - Returns the pre-gate attention output ``[num_tokens,

494 - num_local_heads * v_head_dim]``. On a profile/dummy run

495 - it returns a zeroed buffer.

496 - """

497 - if self.q_lora_rank is not None:

498 - qkv_lora = self.fused_qkv_a_proj(hidden_states)[0]

499 - q_c, kv_c, k_pe = qkv_lora.split(

500 - [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], dim=-1

501 - )

502 - q_c, kv_c_normed = fused_q_kv_rmsnorm(

503 - q_c,

504 - kv_c,

505 - self.q_a_layernorm.weight.data,

506 - self.kv_a_layernorm.weight.data,

507 - self.rms_norm_eps,

508 - )

509 - q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim)

510 - else:

511 - # Uncompressed query: project directly (no q-LoRA, no q norm) and

512 - # normalize only the kv latent.

513 - q = self.q_proj(hidden_states)[0].view(

514 - -1, self.num_local_heads, self.qk_head_dim

515 - )

516 - kv_lora = self.kv_a_proj_with_mqa(hidden_states)[0]

517 - kv_c, k_pe = kv_lora.split(

518 - [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1

519 - )

520 - kv_c_normed = self.kv_a_layernorm(kv_c)

521 - k_pe = k_pe.unsqueeze(1)

523 + """Expand q-LoRA output and run attention."""

524 + q_lora_rank = self.q_lora_rank

525 + assert q_lora_rank is not None

526 + q_c, kv_c, k_pe = qkv_lora.split(

527 + [q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], dim=-1

528 + )

529 + q_c, kv_c_normed = fused_q_kv_rmsnorm(

530 + q_c,

531 + kv_c,

532 + self.q_a_layernorm.weight.data,

533 + self.kv_a_layernorm.weight.data,

534 + self.rms_norm_eps,

535 + )

536 + q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim)

537

538 attn_out = torch.empty(

539 (hidden_states.shape[0], self.num_local_heads * self.v_head_dim),

540 dtype=hidden_states.dtype,

541 device=hidden_states.device,

542 )

528 - self._attention(positions, q, kv_c_normed, k_pe, attn_out)

543 + self._attention(positions, q, kv_c_normed, k_pe.unsqueeze(1), attn_out)

544 return attn_out

545

531 - def forward(

546 + def _forward_q_lora(

547 self,

548 positions: torch.Tensor,

549 hidden_states: torch.Tensor,

535 - ) -> torch.Tensor:

536 - # Both branches produce (attn_out, gate); they differ only in whether

537 - # the g_proj GEMM is overlapped on the aux stream.

538 - g_proj = self.g_proj

539 - events = self._gate_events

550 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

551 + """Run the q-LoRA path, including gate overlap when eligible."""

552 + q_lora_rank = self.q_lora_rank

553 + assert q_lora_rank is not None

554 + qkv_a_rows = q_lora_rank + self.kv_lora_rank + self.qk_rope_head_dim

555 +

556 + # run g_proj in aux_stream

557 if (

541 - g_proj is not None

542 - and events is not None

543 - and self.aux_stream is not None

558 + self._gate_events is not None

559 and hidden_states.shape[0] < _GATE_MULTI_STREAM_TOKEN_THRESHOLD

560 ):

546 - attn_out, gate = maybe_execute_in_parallel(

547 - lambda: self._forward_attn(positions, hidden_states),

548 - lambda: g_proj(hidden_states)[0],

549 - events[0],

550 - events[1],

561 + fused_qkv_a_g_proj = self.fused_qkv_a_g_proj

562 + assert fused_qkv_a_g_proj is not None

563 + qkv_a_weight = fused_qkv_a_g_proj.weight[:qkv_a_rows]

564 + gate_weight = fused_qkv_a_g_proj.weight[qkv_a_rows:]

565 +

566 + start_event, end_event = self._gate_events

567 + return maybe_execute_in_parallel(

568 + lambda: self._apply_q_lora_attention(

569 + positions,

570 + hidden_states,

571 + self._unquantized_gemm(hidden_states, qkv_a_weight),

572 + ),

573 + lambda: self._unquantized_gemm(hidden_states, gate_weight),

574 + start_event,

575 + end_event,

576 self.aux_stream,

577 )

578 +

579 + # run g_proj together with qkv_a

580 + if self.use_output_gate:

581 + qkv_a_proj = self.fused_qkv_a_g_proj

582 + assert qkv_a_proj is not None

583 + qkv_lora, gate = qkv_a_proj(hidden_states)[0].split(

584 + [qkv_a_rows, self.num_local_heads * self.v_head_dim], dim=-1

585 + )

586 + else:

587 + qkv_a_proj = self.fused_qkv_a_proj

588 + assert qkv_a_proj is not None

589 + qkv_lora = qkv_a_proj(hidden_states)[0]

590 + gate = None

591 +

592 + attn_out = self._apply_q_lora_attention(positions, hidden_states, qkv_lora)

593 + return attn_out, gate

594 +

595 + def _forward_full_rank_q(

596 + self,

597 + positions: torch.Tensor,

598 + hidden_states: torch.Tensor,

599 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

600 + """Run the full-rank Q path with an optional sequential gate."""

601 + q_proj = self.q_proj

602 + kv_a_proj_with_mqa = self.kv_a_proj_with_mqa

603 + assert q_proj is not None

604 + assert kv_a_proj_with_mqa is not None

605 +

606 + q = q_proj(hidden_states)[0].view(-1, self.num_local_heads, self.qk_head_dim)

607 + kv_lora = kv_a_proj_with_mqa(hidden_states)[0]

608 + kv_c, k_pe = kv_lora.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)

609 + kv_c_normed = self.kv_a_layernorm(kv_c)

610 + attn_out = torch.empty(

611 + (hidden_states.shape[0], self.num_local_heads * self.v_head_dim),

612 + dtype=hidden_states.dtype,

613 + device=hidden_states.device,

614 + )

615 + self._attention(positions, q, kv_c_normed, k_pe.unsqueeze(1), attn_out)

616 +

617 + gate = self.g_proj(hidden_states)[0] if self.g_proj is not None else None

618 + return attn_out, gate

619 +

620 + def forward(

621 + self,

622 + positions: torch.Tensor,

623 + hidden_states: torch.Tensor,

624 + ) -> torch.Tensor:

625 + if self.q_lora_rank is None:

626 + attn_out, gate = self._forward_full_rank_q(positions, hidden_states)

627 else:

554 - attn_out = self._forward_attn(positions, hidden_states)

555 - gate = g_proj(hidden_states)[0] if g_proj is not None else None

628 + attn_out, gate = self._forward_q_lora(positions, hidden_states)

629

630 if gate is not None:

631 attn_out = _gate_sigmoid_mul(attn_out, gate)

vllm/models/kimi_k3/nvidia/model.py

+12 -4

@@ -1120,6 +1120,7 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant):

1120 "in_proj_qkvgfab": ["q_proj", "k_proj", "v_proj", "b_proj", "f_a_proj"],

1121 "conv1d": ["q_conv1d", "k_conv1d", "v_conv1d"],

1122 "fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"],

1123 + "fused_qkv_a_g_proj": ["q_a_proj", "kv_a_proj_with_mqa", "g_proj"],

1124 }

1125

1126 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

@@ -1460,10 +1461,17 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant):

1461 if use_full_rank_gate:

1462 stacked_params_mapping.append((".in_proj_qkvgfab", ".g_proj", 3))

1463 if getattr(self.config, "q_lora_rank", None) is not None:

1463 - stacked_params_mapping += [

1464 - (".fused_qkv_a_proj", ".q_a_proj", 0),

1465 - (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),

1466 - ]

1464 + if self.config.mla_use_output_gate:

1465 + stacked_params_mapping += [

1466 + (".fused_qkv_a_g_proj", ".q_a_proj", 0),

1467 + (".fused_qkv_a_g_proj", ".kv_a_proj_with_mqa", 1),

1468 + (".fused_qkv_a_g_proj", ".g_proj", 2),

1469 + ]

1470 + else:

1471 + stacked_params_mapping += [

1472 + (".fused_qkv_a_proj", ".q_a_proj", 0),

1473 + (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),

1474 + ]

1475 use_mega_moe = any(

1476 module.use_mega_moe

1477 for module in self.modules()

vllm/models/kimi_k3/nvidia/mtp.py

+11 -4

@@ -272,10 +272,17 @@ class KimiK3MTP(nn.Module):

272 if use_full_rank_gate:

273 stacked_params_mapping.append((".in_proj_qkvgfab", ".g_proj", 3))

274 if getattr(self.config, "q_lora_rank", None) is not None:

275 - stacked_params_mapping += [

276 - (".fused_qkv_a_proj", ".q_a_proj", 0),

277 - (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),

278 - ]

275 + if self.config.mla_use_output_gate:

276 + stacked_params_mapping += [

277 + (".fused_qkv_a_g_proj", ".q_a_proj", 0),

278 + (".fused_qkv_a_g_proj", ".kv_a_proj_with_mqa", 1),

279 + (".fused_qkv_a_g_proj", ".g_proj", 2),

280 + ]

281 + else:

282 + stacked_params_mapping += [

283 + (".fused_qkv_a_proj", ".q_a_proj", 0),

284 + (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),

285 + ]

286

287 use_mega_moe = any(

288 module.use_mega_moe