@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 17808394bc
Fork 0

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 6, 12:00 UTC

standardize load_weights using AutoWeightsLoader for kimi_linear and minimax_text_01 (#37371)

Signed-off-by: XuLiu <xuliu40@gmail.com> Co-authored-by: XuLiu <xuliu40@gmail.com>

XLiu-2000 committed Mar 18, 2026 at 23:05 UTC 17808394bc48b7568a471ad717a15aab885b0349
2 files changed +235 -219
vllm/model_executor/models/kimi_linear.py
+97 -88
@@ -46,6 +46,7 @@ from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
46
47 from .interfaces import HasInnerState, IsHybrid, MixtureOfExperts, SupportsPP
48 from .utils import (
49 + AutoWeightsLoader,
50 PPMissingLayer,
51 is_pp_missing_parameter,
52 make_layers,
@@ -472,94 +473,7 @@ class KimiLinearModel(nn.Module):
473 hidden_states, _ = self.norm(hidden_states, residual)
474 return hidden_states
475
475 -
476 -class KimiLinearForCausalLM(
477 - nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid
478 -):
479 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
480 - super().__init__()
481 - self.model_config = vllm_config.model_config
482 - self.vllm_config = vllm_config
483 - self.config = self.model_config.hf_config
484 - quant_config = vllm_config.quant_config
485 - self.quant_config = quant_config
486 - self.model = KimiLinearModel(
487 - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
488 - )
489 - if get_pp_group().is_last_rank:
490 - self.lm_head = ParallelLMHead(
491 - self.config.vocab_size,
492 - self.config.hidden_size,
493 - quant_config=quant_config,
494 - prefix=maybe_prefix(prefix, "lm_head"),
495 - )
496 - else:
497 - self.lm_head = PPMissingLayer()
498 - logit_scale = getattr(self.config, "logit_scale", 1.0)
499 - self.logits_processor = LogitsProcessor(
500 - self.config.vocab_size, scale=logit_scale
501 - )
502 -
503 - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
504 - return self.model.embed_input_ids(input_ids)
505 -
506 - def forward(
507 - self,
508 - input_ids: torch.Tensor | None,
509 - positions: torch.Tensor,
510 - intermediate_tensors: IntermediateTensors | None = None,
511 - inputs_embeds: torch.Tensor | None = None,
512 - **kwargs,
513 - ) -> torch.Tensor | IntermediateTensors:
514 - hidden_states = self.model(
515 - input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs
516 - )
517 - return hidden_states
518 -
519 - @classmethod
520 - def get_mamba_state_dtype_from_config(
521 - cls,
522 - vllm_config: "VllmConfig",
523 - ) -> tuple[torch.dtype, torch.dtype, torch.dtype, torch.dtype]:
524 - return MambaStateDtypeCalculator.kda_state_dtype(
525 - vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype
526 - )
527 -
528 - @classmethod
529 - def get_mamba_state_shape_from_config(
530 - cls, vllm_config: "VllmConfig"
531 - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
532 - parallel_config = vllm_config.parallel_config
533 - hf_config = vllm_config.model_config.hf_config
534 - tp_size = parallel_config.tensor_parallel_size
535 - num_spec = (
536 - vllm_config.speculative_config.num_speculative_tokens
537 - if vllm_config.speculative_config
538 - else 0
539 - )
540 - return MambaStateShapeCalculator.kda_state_shape(
541 - tp_size,
542 - hf_config.linear_attn_config["num_heads"],
543 - hf_config.linear_attn_config["head_dim"],
544 - conv_kernel_size=hf_config.linear_attn_config["short_conv_kernel_size"],
545 - num_spec=num_spec,
546 - )
547 -
548 - @classmethod
549 - def get_mamba_state_copy_func(
550 - cls,
551 - ) -> tuple[
552 - MambaStateCopyFunc, MambaStateCopyFunc, MambaStateCopyFunc, MambaStateCopyFunc
553 - ]:
554 - return MambaStateCopyFuncCalculator.kda_state_copy_func()
555 -
556 - def compute_logits(
557 - self,
558 - hidden_states: torch.Tensor,
559 - ) -> torch.Tensor | None:
560 - return self.logits_processor(self.lm_head, hidden_states)
561 -
562 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
476 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
477 stacked_params_mapping = [
478 # (param_name, shard_name, shard_id)
479 (".gate_up_proj", ".gate_proj", 0),
@@ -653,6 +567,101 @@ class KimiLinearForCausalLM(
567 )
568 weight_loader(param, loaded_weight, **kwargs)
569 loaded_params.add(name)
570 + return loaded_params
571 +
572 +
573 +class KimiLinearForCausalLM(
574 + nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid
575 +):
576 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
577 + super().__init__()
578 + self.model_config = vllm_config.model_config
579 + self.vllm_config = vllm_config
580 + self.config = self.model_config.hf_config
581 + quant_config = vllm_config.quant_config
582 + self.quant_config = quant_config
583 + self.model = KimiLinearModel(
584 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
585 + )
586 + if get_pp_group().is_last_rank:
587 + self.lm_head = ParallelLMHead(
588 + self.config.vocab_size,
589 + self.config.hidden_size,
590 + quant_config=quant_config,
591 + prefix=maybe_prefix(prefix, "lm_head"),
592 + )
593 + else:
594 + self.lm_head = PPMissingLayer()
595 + logit_scale = getattr(self.config, "logit_scale", 1.0)
596 + self.logits_processor = LogitsProcessor(
597 + self.config.vocab_size, scale=logit_scale
598 + )
599 +
600 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
601 + return self.model.embed_input_ids(input_ids)
602 +
603 + def forward(
604 + self,
605 + input_ids: torch.Tensor | None,
606 + positions: torch.Tensor,
607 + intermediate_tensors: IntermediateTensors | None = None,
608 + inputs_embeds: torch.Tensor | None = None,
609 + **kwargs,
610 + ) -> torch.Tensor | IntermediateTensors:
611 + hidden_states = self.model(
612 + input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs
613 + )
614 + return hidden_states
615 +
616 + @classmethod
617 + def get_mamba_state_dtype_from_config(
618 + cls,
619 + vllm_config: "VllmConfig",
620 + ) -> tuple[torch.dtype, torch.dtype, torch.dtype, torch.dtype]:
621 + return MambaStateDtypeCalculator.kda_state_dtype(
622 + vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype
623 + )
624 +
625 + @classmethod
626 + def get_mamba_state_shape_from_config(
627 + cls, vllm_config: "VllmConfig"
628 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
629 + parallel_config = vllm_config.parallel_config
630 + hf_config = vllm_config.model_config.hf_config
631 + tp_size = parallel_config.tensor_parallel_size
632 + num_spec = (
633 + vllm_config.speculative_config.num_speculative_tokens
634 + if vllm_config.speculative_config
635 + else 0
636 + )
637 + return MambaStateShapeCalculator.kda_state_shape(
638 + tp_size,
639 + hf_config.linear_attn_config["num_heads"],
640 + hf_config.linear_attn_config["head_dim"],
641 + conv_kernel_size=hf_config.linear_attn_config["short_conv_kernel_size"],
642 + num_spec=num_spec,
643 + )
644 +
645 + @classmethod
646 + def get_mamba_state_copy_func(
647 + cls,
648 + ) -> tuple[
649 + MambaStateCopyFunc, MambaStateCopyFunc, MambaStateCopyFunc, MambaStateCopyFunc
650 + ]:
651 + return MambaStateCopyFuncCalculator.kda_state_copy_func()
652 +
653 + def compute_logits(
654 + self,
655 + hidden_states: torch.Tensor,
656 + ) -> torch.Tensor | None:
657 + return self.logits_processor(self.lm_head, hidden_states)
658 +
659 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
660 + loader = AutoWeightsLoader(
661 + self,
662 + skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
663 + )
664 + return loader.load_weights(weights)
665
666
667 def get_spec_layer_idx_from_weight_name(
vllm/model_executor/models/minimax_text_01.py
+138 -131
@@ -52,7 +52,12 @@ from vllm.sequence import IntermediateTensors
52 from vllm.v1.attention.backend import AttentionMetadata
53
54 from .interfaces import HasInnerState, IsHybrid
55 -from .utils import PPMissingLayer, is_pp_missing_parameter, make_layers
55 +from .utils import (
56 + AutoWeightsLoader,
57 + PPMissingLayer,
58 + is_pp_missing_parameter,
59 + make_layers,
60 +)
61
62
63 def replace_weight_name(
@@ -494,6 +499,8 @@ class MiniMaxText01Model(nn.Module):
499 quant_config = vllm_config.quant_config
500 cache_config = vllm_config.cache_config
501 scheduler_config = vllm_config.scheduler_config
502 + self.config = config
503 + self.CONCAT_FFN = True
504
505 self.vocab_size = config.vocab_size
506
@@ -620,128 +627,6 @@ class MiniMaxText01Model(nn.Module):
627 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
628 return self.embed_tokens(input_ids)
629
623 - def forward(
624 - self,
625 - input_ids: torch.Tensor | None,
626 - positions: torch.Tensor,
627 - intermediate_tensors: IntermediateTensors | None = None,
628 - inputs_embeds: torch.Tensor | None = None,
629 - **kwargs,
630 - ) -> torch.Tensor | IntermediateTensors:
631 - forward_context = get_forward_context()
632 - attn_metadata = forward_context.attn_metadata
633 -
634 - if get_pp_group().is_first_rank:
635 - if inputs_embeds is None:
636 - hidden_states = self.embed_scale * self.embed_tokens(input_ids)
637 - else:
638 - hidden_states = inputs_embeds
639 - residual = None
640 - else:
641 - assert intermediate_tensors is not None
642 - hidden_states = intermediate_tensors["hidden_states"]
643 - residual = intermediate_tensors["residual"]
644 -
645 - for layer in islice(self.layers, self.start_layer, self.end_layer):
646 - hidden_states, residual = layer(
647 - hidden_states=hidden_states,
648 - positions=positions,
649 - attn_metadata=attn_metadata,
650 - residual=residual,
651 - )
652 - if not get_pp_group().is_last_rank:
653 - return IntermediateTensors(
654 - {"hidden_states": hidden_states, "residual": residual}
655 - )
656 - if residual is not None:
657 - hidden_states, _ = self.norm(hidden_states, residual)
658 - else:
659 - hidden_states = self.norm(hidden_states)
660 -
661 - return hidden_states
662 -
663 -
664 -class MiniMaxText01ForCausalLM(nn.Module, HasInnerState, IsHybrid):
665 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
666 - super().__init__()
667 - config = vllm_config.model_config.hf_config
668 -
669 - self.config = config
670 -
671 - if not hasattr(config, "sliding_window"):
672 - config.sliding_window = None
673 -
674 - self.CONCAT_FFN = True
675 -
676 - if hasattr(vllm_config.model_config, "max_model_len"):
677 - self.config.max_model_len = vllm_config.model_config.max_model_len
678 - self.model = MiniMaxText01Model(
679 - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
680 - )
681 - if get_pp_group().is_last_rank:
682 - self.lm_head = ParallelLMHead(
683 - config.vocab_size,
684 - self.config.hidden_size,
685 - prefix=maybe_prefix(prefix, "lm_head"),
686 - )
687 -
688 - self.logits_processor = LogitsProcessor(
689 - config.vocab_size, self.config.vocab_size
690 - )
691 -
692 - else:
693 - self.lm_head = PPMissingLayer()
694 - self.lm_head.float()
695 - flash_layer_count = sum(
696 - 1 for attn_type in self.model.decoder_attention_types if attn_type == 1
697 - )
698 - self.kv_cache = [torch.tensor([]) for _ in range(flash_layer_count)]
699 - return
700 -
701 - def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
702 - return self.model.minimax_cache.copy_inputs_before_cuda_graphs(
703 - input_buffers, **kwargs
704 - )
705 -
706 - def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
707 - return self.model.minimax_cache.get_seqlen_agnostic_capture_inputs(batch_size)
708 -
709 - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
710 - return self.model.embed_input_ids(input_ids)
711 -
712 - def forward(
713 - self,
714 - input_ids: torch.Tensor | None,
715 - positions: torch.Tensor,
716 - intermediate_tensors: IntermediateTensors | None = None,
717 - inputs_embeds: torch.Tensor | None = None,
718 - **kwargs,
719 - ) -> torch.Tensor:
720 - hidden_states = self.model(
721 - input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs
722 - )
723 -
724 - return hidden_states
725 -
726 - def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
727 - logits = self.logits_processor(self.lm_head, hidden_states.float())
728 -
729 - return logits
730 -
731 - def make_empty_intermediate_tensors(
732 - self, batch_size: int, dtype: torch.dtype, device: torch.device
733 - ) -> IntermediateTensors:
734 - return IntermediateTensors(
735 - {
736 - "hidden_states": torch.zeros(
737 - (batch_size, self.config.hidden_size), dtype=dtype, device=device
738 - ),
739 - "residual": torch.zeros(
740 - (batch_size, self.config.hidden_size), dtype=dtype, device=device
741 - ),
742 - }
743 - )
744 -
630 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
631 params_dict = dict(self.named_parameters())
632 loaded_params: set[str] = set()
@@ -753,17 +638,15 @@ class MiniMaxText01ForCausalLM(nn.Module, HasInnerState, IsHybrid):
638 return None
639
640 def is_linear_attn_layer(layer_idx: int) -> bool:
756 - if layer_idx is None or layer_idx >= len(
757 - self.model.decoder_attention_types
758 - ):
641 + if layer_idx is None or layer_idx >= len(self.decoder_attention_types):
642 return False
760 - return self.model.decoder_attention_types[layer_idx] == 0
643 + return self.decoder_attention_types[layer_idx] == 0
644
645 def is_moe_weight(name: str) -> bool:
646 return "block_sparse_moe" in name and not name.endswith(".bias")
647
648 def get_expert_id(param_name):
766 - pattern = r"model\.layers\.\d+\.block_sparse_moe\.experts\.(\d+)\."
649 + pattern = r"layers\.\d+\.block_sparse_moe\.experts\.(\d+)\."
650 match = re.search(pattern, param_name)
651 if match:
652 return match.group(1)
@@ -948,9 +831,7 @@ class MiniMaxText01ForCausalLM(nn.Module, HasInnerState, IsHybrid):
831
832 for name, loaded_weight in weights:
833 weight_at_layer = which_layer(name)
951 - if weight_at_layer and weight_at_layer >= len(
952 - self.model.decoder_attention_types
953 - ):
834 + if weight_at_layer and weight_at_layer >= len(self.decoder_attention_types):
835 continue
836
837 if is_layer_norm_weight(name):
@@ -975,6 +856,128 @@ class MiniMaxText01ForCausalLM(nn.Module, HasInnerState, IsHybrid):
856 load_basic_weight(name, loaded_weight, self)
857 return loaded_params
858
859 + def forward(
860 + self,
861 + input_ids: torch.Tensor | None,
862 + positions: torch.Tensor,
863 + intermediate_tensors: IntermediateTensors | None = None,
864 + inputs_embeds: torch.Tensor | None = None,
865 + **kwargs,
866 + ) -> torch.Tensor | IntermediateTensors:
867 + forward_context = get_forward_context()
868 + attn_metadata = forward_context.attn_metadata
869 +
870 + if get_pp_group().is_first_rank:
871 + if inputs_embeds is None:
872 + hidden_states = self.embed_scale * self.embed_tokens(input_ids)
873 + else:
874 + hidden_states = inputs_embeds
875 + residual = None
876 + else:
877 + assert intermediate_tensors is not None
878 + hidden_states = intermediate_tensors["hidden_states"]
879 + residual = intermediate_tensors["residual"]
880 +
881 + for layer in islice(self.layers, self.start_layer, self.end_layer):
882 + hidden_states, residual = layer(
883 + hidden_states=hidden_states,
884 + positions=positions,
885 + attn_metadata=attn_metadata,
886 + residual=residual,
887 + )
888 + if not get_pp_group().is_last_rank:
889 + return IntermediateTensors(
890 + {"hidden_states": hidden_states, "residual": residual}
891 + )
892 + if residual is not None:
893 + hidden_states, _ = self.norm(hidden_states, residual)
894 + else:
895 + hidden_states = self.norm(hidden_states)
896 +
897 + return hidden_states
898 +
899 +
900 +class MiniMaxText01ForCausalLM(nn.Module, HasInnerState, IsHybrid):
901 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
902 + super().__init__()
903 + config = vllm_config.model_config.hf_config
904 +
905 + self.config = config
906 +
907 + if not hasattr(config, "sliding_window"):
908 + config.sliding_window = None
909 +
910 + self.CONCAT_FFN = True
911 +
912 + if hasattr(vllm_config.model_config, "max_model_len"):
913 + self.config.max_model_len = vllm_config.model_config.max_model_len
914 + self.model = MiniMaxText01Model(
915 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
916 + )
917 + if get_pp_group().is_last_rank:
918 + self.lm_head = ParallelLMHead(
919 + config.vocab_size,
920 + self.config.hidden_size,
921 + prefix=maybe_prefix(prefix, "lm_head"),
922 + )
923 +
924 + self.logits_processor = LogitsProcessor(
925 + config.vocab_size, self.config.vocab_size
926 + )
927 +
928 + else:
929 + self.lm_head = PPMissingLayer()
930 + self.lm_head.float()
931 + flash_layer_count = sum(
932 + 1 for attn_type in self.model.decoder_attention_types if attn_type == 1
933 + )
934 + self.kv_cache = [torch.tensor([]) for _ in range(flash_layer_count)]
935 + return
936 +
937 + def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
938 + return self.model.minimax_cache.copy_inputs_before_cuda_graphs(
939 + input_buffers, **kwargs
940 + )
941 +
942 + def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
943 + return self.model.minimax_cache.get_seqlen_agnostic_capture_inputs(batch_size)
944 +
945 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
946 + return self.model.embed_input_ids(input_ids)
947 +
948 + def forward(
949 + self,
950 + input_ids: torch.Tensor | None,
951 + positions: torch.Tensor,
952 + intermediate_tensors: IntermediateTensors | None = None,
953 + inputs_embeds: torch.Tensor | None = None,
954 + **kwargs,
955 + ) -> torch.Tensor:
956 + hidden_states = self.model(
957 + input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs
958 + )
959 +
960 + return hidden_states
961 +
962 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
963 + logits = self.logits_processor(self.lm_head, hidden_states.float())
964 +
965 + return logits
966 +
967 + def make_empty_intermediate_tensors(
968 + self, batch_size: int, dtype: torch.dtype, device: torch.device
969 + ) -> IntermediateTensors:
970 + return IntermediateTensors(
971 + {
972 + "hidden_states": torch.zeros(
973 + (batch_size, self.config.hidden_size), dtype=dtype, device=device
974 + ),
975 + "residual": torch.zeros(
976 + (batch_size, self.config.hidden_size), dtype=dtype, device=device
977 + ),
978 + }
979 + )
980 +
981 @classmethod
982 def get_mamba_state_dtype_from_config(
983 cls,
@@ -1011,3 +1014,7 @@ class MiniMaxText01ForCausalLM(nn.Module, HasInnerState, IsHybrid):
1014 @classmethod
1015 def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc]:
1016 return MambaStateCopyFuncCalculator.linear_attention_state_copy_func()
1017 +
1018 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1019 + loader = AutoWeightsLoader(self)
1020 + return loader.load_weights(weights)