iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Feature] Skip MLA weight absorption when weights are pre-processed (#57320)

Browse Source 
Co-authored-by: Xun Sun <UNIDY2002@outlook.com>
This commit is contained in:
siyu
2026-10-17 23:10:04 +08:00
committed by GitHub
co-authored by Xun Sun
parent 237d99b4af
commit c61a1e73c6
1 changed files with 16 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/models/kimi_k3/nvidia/mla.py
+16 -13
View File
@@ -75,7 +75,7 @@ from vllm.model_executor.layers.quantization import (
resolve_quant_method, resolve_quant_method,
) )
from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding, get_rope from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding, get_rope
from vllm.model_executor.utils import replace_parameter from vllm.model_executor.utils import is_weights_pre_processed, replace_parameter
from vllm.models.common.ops import fused_q_kv_rmsnorm 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.low_latency_gemm import try_low_latency_gemm
from vllm.models.kimi_k3.nvidia.ops.fused_mla_key_concat_kv_cache import ( from vllm.models.kimi_k3.nvidia.ops.fused_mla_key_concat_kv_cache import (
@@ -440,18 +440,21 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
projected into latent space by ``W_UK_T`` and the attention output is projected into latent space by ``W_UK_T`` and the attention output is
projected back to ``v`` by ``W_UV`` -- avoiding materializing full K/V. projected back to ``v`` by ``W_UV`` -- avoiding materializing full K/V.
""" """
W_UK, W_UV = split_kv_b_proj( # The weight cache daemon exports the absorbed weights. Repeating the
self.kv_b_proj, # absorption would mutate its shared CUDA allocation.
act_dtype, if not (is_weights_pre_processed() and getattr(self, "W_UV", None) is not None):
self.kv_lora_rank, W_UK, W_UV = split_kv_b_proj(
self.num_local_heads, self.kv_b_proj,
self.qk_nope_head_dim, act_dtype,
self.v_head_dim, self.kv_lora_rank,
) self.num_local_heads,
# (L, N, V) -> (N, L, V) self.qk_nope_head_dim,
replace_parameter(self, "W_UV", W_UV.transpose(0, 1), prefer_copy=True) self.v_head_dim,
# (L, N, P) -> (N, P, L) )
replace_parameter(self, "W_UK_T", W_UK.permute(1, 2, 0), prefer_copy=True) # (L, N, V) -> (N, L, V)
replace_parameter(self, "W_UV", W_UV.transpose(0, 1), prefer_copy=True)
# (L, N, P) -> (N, P, L)
replace_parameter(self, "W_UK_T", W_UK.permute(1, 2, 0), prefer_copy=True)
quant_method = ( quant_method = (
resolve_quant_method(self.quant_config, self, prefix=self.layer_name) resolve_quant_method(self.quant_config, self, prefix=self.layer_name)