A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 9, 07:30 UTC
[Models] Refactor Kimi-K2.5 weight loading (#33346)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Isotr0py committed
Jan 30, 2026 at 13:31 UTC
8bfc8d5600ed2c3c22835225b70c72b5b39fdd9f
2 files changed
+40
-176
vllm/model_executor/models/kimi_k25.py
+38
-174
@@ -23,16 +23,7 @@ from transformers.processing_utils import ProcessorMixin
23
24
from vllm.config import VllmConfig
25
from vllm.config.multimodal import BaseDummyOptions
26
-from vllm.distributed import get_pp_group
26
from vllm.logger import init_logger
28
-from vllm.model_executor.layers.fused_moe import SharedFusedMoE
29
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
30
-from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
31
-from vllm.model_executor.model_loader.weight_utils import (
32
- default_weight_loader,
33
- maybe_remap_kv_scale_name,
34
-)
35
-from vllm.model_executor.models.deepseek_v2 import DeepseekV2Model
27
from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP
28
from vllm.model_executor.models.kimi_k25_vit import (
29
KimiK25MultiModalProjector,
@@ -64,7 +55,12 @@ from vllm.transformers_utils.configs import KimiK25Config
55
from vllm.transformers_utils.processor import cached_get_image_processor
56
from vllm.utils.tensor_schema import TensorSchema, TensorShape
57
67
-from .utils import PPMissingLayer, is_pp_missing_parameter, maybe_prefix
58
+from .utils import (
59
+ AutoWeightsLoader,
60
+ WeightsMapper,
61
+ init_vllm_registered_model,
62
+ maybe_prefix,
63
+)
64
65
logger = init_logger(__name__)
66
@@ -294,6 +290,13 @@ class KimiK25ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)
290
291
supports_encoder_tp_data = True
292
293
+ weights_mapper = WeightsMapper(
294
+ orig_to_new_prefix={
295
+ "mm_projector.proj.0": "mm_projector.linear_1",
296
+ "mm_projector.proj.2": "mm_projector.linear_2",
297
+ }
298
+ )
299
+
300
@classmethod
301
def get_placeholder_str(cls, modality: str, i: int) -> str | None:
302
# Kimi-K2.5 uses video_chunk for all media types
@@ -323,45 +326,39 @@ class KimiK25ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)
326
self.hidden_size = config.text_config.hidden_size
327
self.device = current_platform.current_device()
328
# Build vision tower directly with KimiK25VisionConfig
326
- self.vision_tower = MoonViT3dPretrainedModel(
327
- config.vision_config,
328
- prefix=maybe_prefix(prefix, "vision_tower"),
329
- )
330
- self.vision_tower = self.vision_tower.to(
331
- device=self.device, dtype=model_config.dtype
332
- )
329
+ with self._mark_tower_model(vllm_config, "vision_chunk"):
330
+ self.vision_tower = MoonViT3dPretrainedModel(
331
+ config.vision_config,
332
+ prefix=maybe_prefix(prefix, "vision_tower"),
333
+ )
334
+ self.vision_tower = self.vision_tower.to(
335
+ device=self.device, dtype=model_config.dtype
336
+ )
337
334
- self.mm_projector = KimiK25MultiModalProjector(
335
- config=config.vision_config,
336
- use_data_parallel=self.use_data_parallel,
337
- prefix=maybe_prefix(prefix, "mm_projector"),
338
- )
339
- self.mm_projector = self.mm_projector.to(
340
- device=self.device, dtype=model_config.dtype
341
- )
338
+ self.mm_projector = KimiK25MultiModalProjector(
339
+ config=config.vision_config,
340
+ use_data_parallel=self.use_data_parallel,
341
+ prefix=maybe_prefix(prefix, "mm_projector"),
342
+ )
343
+ self.mm_projector = self.mm_projector.to(
344
+ device=self.device, dtype=model_config.dtype
345
+ )
346
347
self.quant_config = quant_config
348
sub_vllm_config = copy.deepcopy(vllm_config)
349
sub_vllm_config.model_config.hf_config = (
350
sub_vllm_config.model_config.hf_config.text_config
351
)
348
- self.language_model = DeepseekV2Model(
349
- vllm_config=sub_vllm_config,
350
- prefix=maybe_prefix(prefix, "language_model"),
351
- )
352
- if get_pp_group().is_last_rank:
353
- self.lm_head = ParallelLMHead(
354
- config.vocab_size,
355
- config.text_config.hidden_size,
356
- prefix=maybe_prefix(prefix, "lm_head"),
352
+ with self._mark_language_model(vllm_config):
353
+ self.language_model = init_vllm_registered_model(
354
+ vllm_config=vllm_config,
355
+ hf_config=config.text_config,
356
+ prefix=maybe_prefix(prefix, "language_model"),
357
+ architectures=["DeepseekV2ForCausalLM"],
358
)
358
- else:
359
- self.lm_head = PPMissingLayer()
359
self.make_empty_intermediate_tensors = (
360
self.language_model.make_empty_intermediate_tensors
361
)
363
- logit_scale = getattr(config, "logit_scale", 1.0)
364
- self.logits_processor = LogitsProcessor(config.vocab_size, scale=logit_scale)
362
self.media_placeholder: int = self.config.media_placeholder_token_id
363
364
def _parse_and_validate_media_input(
@@ -421,9 +418,6 @@ class KimiK25ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)
418
vision_embeddings = self._process_media_input(media_input)
419
return vision_embeddings
420
424
- def get_language_model(self) -> torch.nn.Module:
425
- return self.language_model
426
-
421
def forward(
422
self,
423
input_ids: torch.Tensor,
@@ -444,139 +438,9 @@ class KimiK25ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)
438
return hidden_states
439
440
def compute_logits(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor:
447
- logits = self.logits_processor(self.lm_head, hidden_states, **kwargs)
441
+ logits = self.language_model.compute_logits(hidden_states)
442
return logits
443
450
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
451
- # Params for weights, fp8 weight scales, fp8 activation scales
452
- # (param_name, weight_name, expert_id, shard_id)
453
- config = self.config.text_config
454
- if not getattr(config, "n_routed_experts", None):
455
- return []
456
- return SharedFusedMoE.make_expert_params_mapping(
457
- self,
458
- ckpt_gate_proj_name="gate_proj",
459
- ckpt_down_proj_name="down_proj",
460
- ckpt_up_proj_name="up_proj",
461
- num_experts=config.n_routed_experts,
462
- num_redundant_experts=0,
463
- )
464
-
444
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
466
- config = self.config.text_config
467
- _KEYS_TO_MODIFY_MAPPING = {
468
- "language_model.lm_head": "lm_head",
469
- "language_model.model": "language_model",
470
- # mm_projector -> mm_projector mapping
471
- # "mm_projector": "mm_projector",
472
- "mm_projector.proj.0": "mm_projector.linear_1",
473
- "mm_projector.proj.2": "mm_projector.linear_2",
474
- }
475
- stacked_params_mapping = [
476
- (".gate_up_proj", ".gate_proj", 0),
477
- (".gate_up_proj", ".up_proj", 1),
478
- ]
479
- if getattr(config, "kv_lora_rank", None) and getattr(
480
- config, "q_lora_rank", None
481
- ):
482
- stacked_params_mapping += [
483
- (".fused_qkv_a_proj", ".q_a_proj", 0),
484
- (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),
485
- ]
486
- expert_params_mapping = self.get_expert_mapping()
487
-
488
- params_dict = dict(self.named_parameters())
489
-
490
- for args in weights:
491
- name, loaded_weight = args[:2]
492
- kwargs = args[2] if len(args) > 2 else {}
493
- if "rotary_emb.inv_freq" in name:
494
- continue
495
-
496
- spec_layer = get_spec_layer_idx_from_weight_name(config, name)
497
- if spec_layer is not None:
498
- continue # skip spec decode layers for main model
499
-
500
- if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
501
- continue
502
-
503
- for key_to_modify, new_key in _KEYS_TO_MODIFY_MAPPING.items():
504
- if key_to_modify in name:
505
- name = name.replace(key_to_modify, new_key)
506
-
507
- use_default_weight_loading = False
508
- if "vision" in name:
509
- if self.vision_tower is not None:
510
- use_default_weight_loading = True
511
- else:
512
- for param_name, weight_name, shard_id in stacked_params_mapping:
513
- if weight_name not in name:
514
- continue
515
- if ("mlp.experts." in name) and name not in params_dict:
516
- continue
517
- name = name.replace(weight_name, param_name)
518
- if name.endswith(".bias") and name not in params_dict:
519
- continue
520
-
521
- if is_pp_missing_parameter(name, self):
522
- continue
523
-
524
- param = params_dict[name]
525
- weight_loader = param.weight_loader
526
- weight_loader(param, loaded_weight, shard_id, **kwargs)
527
- break
528
- else:
529
- for _, (
530
- param_name,
531
- weight_name,
532
- expert_id,
533
- shard_id,
534
- ) in enumerate(expert_params_mapping):
535
- if weight_name not in name:
536
- continue
537
- name = name.replace(weight_name, param_name)
538
-
539
- if is_pp_missing_parameter(name, self):
540
- continue
541
-
542
- param = params_dict[name]
543
- weight_loader = param.weight_loader
544
- weight_loader(
545
- param,
546
- loaded_weight,
547
- name,
548
- expert_id=expert_id,
549
- shard_id=shard_id,
550
- **kwargs,
551
- )
552
- break
553
- else:
554
- use_default_weight_loading = True
555
-
556
- if use_default_weight_loading:
557
- if name.endswith(".bias") and name not in params_dict:
558
- continue
559
- name = maybe_remap_kv_scale_name(name, params_dict)
560
- if name is None:
561
- continue
562
-
563
- if is_pp_missing_parameter(name, self):
564
- continue
565
-
566
- param = params_dict[name]
567
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
568
- weight_loader(param, loaded_weight, **kwargs)
569
-
570
-
571
-def get_spec_layer_idx_from_weight_name(
572
- config: KimiK25Config, weight_name: str
573
-) -> int | None:
574
- if hasattr(config, "num_nextn_predict_layers") and (
575
- config.num_nextn_predict_layers > 0
576
- ):
577
- layer_idx = config.num_hidden_layers
578
- for i in range(config.num_nextn_predict_layers):
579
- # might start with language_model.model.layers
580
- if f"model.layers.{layer_idx + i}." in weight_name:
581
- return layer_idx + i
582
- return None
445
+ loader = AutoWeightsLoader(self)
446
+ return loader.load_weights(weights, mapper=self.weights_mapper)
vllm/model_executor/models/kimi_k25_vit.py
+2
-2
@@ -660,13 +660,13 @@ class KimiK25MultiModalProjector(nn.Module):
660
self.hidden_size,
661
self.hidden_size,
662
bias=True,
663
- prefix=maybe_prefix(prefix, "linear_1"),
663
+ prefix=f"{prefix}.linear_1",
664
)
665
self.linear_2 = ReplicatedLinear(
666
self.hidden_size,
667
config.mm_hidden_size,
668
bias=True,
669
- prefix=maybe_prefix(prefix, "linear_2"),
669
+ prefix=f"{prefix}.linear_2",
670
)
671
self.act = GELUActivation()
672