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 4, 10:00 UTC
Refractor longcat loading to use AutoWeightsLoader (#41448)
Signed-off-by: George-ao <yuyiao772@gmail.com>
Yuyi Ao committed
May 1, 2026 at 05:07 UTC
0dbaf9daad2031235344428d2a574496bb4d9a3b
1 file changed
+82
-73
vllm/model_executor/models/longcat_flash.py
+82
-73
@@ -69,6 +69,7 @@ from vllm.sequence import IntermediateTensors
69
70
from .interfaces import SupportsLoRA, SupportsPP
71
from .utils import (
72
+ AutoWeightsLoader,
73
PPMissingLayer,
74
is_pp_missing_parameter,
75
make_empty_intermediate_tensors_factory,
@@ -485,6 +486,7 @@ class FlashModel(nn.Module):
486
cache_config = vllm_config.cache_config
487
quant_config = vllm_config.quant_config
488
self.config = config
489
+ self.quant_config = quant_config
490
491
self.vocab_size = config.vocab_size
492
@@ -551,77 +553,6 @@ class FlashModel(nn.Module):
553
hidden_states, _ = self.norm(hidden_states, residual)
554
return hidden_states
555
554
-
555
-class LongcatFlashForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
556
- """Flash model for causal language modeling."""
557
-
558
- packed_modules_mapping = {
559
- "qkv_proj": [
560
- "q_proj",
561
- "k_proj",
562
- "v_proj",
563
- ],
564
- "gate_up_proj": [
565
- "gate_proj",
566
- "up_proj",
567
- ],
568
- }
569
-
570
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
571
- super().__init__()
572
- config = FlashConfig(**vllm_config.model_config.hf_config.__dict__)
573
- quant_config = vllm_config.quant_config
574
-
575
- self.config = config
576
- config.intermediate_size = (
577
- config.ffn_hidden_size
578
- if hasattr(config, "ffn_hidden_size")
579
- else config.intermediate_size
580
- )
581
-
582
- self.quant_config = quant_config
583
-
584
- self.model = FlashModel(
585
- vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
586
- )
587
-
588
- if get_pp_group().is_last_rank:
589
- self.lm_head = ParallelLMHead(
590
- config.vocab_size,
591
- config.hidden_size,
592
- quant_config=quant_config,
593
- prefix=maybe_prefix(prefix, "lm_head"),
594
- )
595
- else:
596
- self.lm_head = PPMissingLayer()
597
-
598
- self.logits_processor = LogitsProcessor(config.vocab_size)
599
- self.make_empty_intermediate_tensors = (
600
- self.model.make_empty_intermediate_tensors
601
- )
602
-
603
- def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
604
- return self.model.embed_input_ids(input_ids)
605
-
606
- def forward(
607
- self,
608
- input_ids: torch.Tensor | None,
609
- positions: torch.Tensor,
610
- intermediate_tensors: IntermediateTensors | None = None,
611
- inputs_embeds: torch.Tensor | None = None,
612
- ) -> torch.Tensor | IntermediateTensors:
613
- hidden_states = self.model(
614
- input_ids, positions, intermediate_tensors, inputs_embeds
615
- )
616
- return hidden_states
617
-
618
- def compute_logits(
619
- self,
620
- hidden_states: torch.Tensor,
621
- ) -> torch.Tensor | None:
622
- logits = self.logits_processor(self.lm_head, hidden_states)
623
- return logits
624
-
556
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
557
# Params for weights, fp8 weight scales, fp8 activation scales
558
# (param_name, weight_name, expert_id, shard_id)
@@ -730,9 +661,9 @@ class LongcatFlashForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
661
loaded_params.add(name)
662
for layer_id in range(self.config.num_hidden_layers):
663
for i in range(2):
733
- if isinstance(self.model.layers[layer_id], PPMissingLayer):
664
+ if isinstance(self.layers[layer_id], PPMissingLayer):
665
continue
735
- self_attn = self.model.layers[layer_id].self_attn[i]
666
+ self_attn = self.layers[layer_id].self_attn[i]
667
if hasattr(
668
self.quant_config, "weight_block_size"
669
) and self_attn.kv_b_proj.weight.dtype in (
@@ -765,3 +696,81 @@ class LongcatFlashForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
696
self.config.hidden_size / self.config.kv_lora_rank
697
) ** 0.5
698
return loaded_params
699
+
700
+
701
+class LongcatFlashForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
702
+ """Flash model for causal language modeling."""
703
+
704
+ packed_modules_mapping = {
705
+ "qkv_proj": [
706
+ "q_proj",
707
+ "k_proj",
708
+ "v_proj",
709
+ ],
710
+ "gate_up_proj": [
711
+ "gate_proj",
712
+ "up_proj",
713
+ ],
714
+ }
715
+
716
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
717
+ super().__init__()
718
+ config = FlashConfig(**vllm_config.model_config.hf_config.__dict__)
719
+ quant_config = vllm_config.quant_config
720
+
721
+ self.config = config
722
+ config.intermediate_size = (
723
+ config.ffn_hidden_size
724
+ if hasattr(config, "ffn_hidden_size")
725
+ else config.intermediate_size
726
+ )
727
+
728
+ self.quant_config = quant_config
729
+
730
+ self.model = FlashModel(
731
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
732
+ )
733
+
734
+ if get_pp_group().is_last_rank:
735
+ self.lm_head = ParallelLMHead(
736
+ config.vocab_size,
737
+ config.hidden_size,
738
+ quant_config=quant_config,
739
+ prefix=maybe_prefix(prefix, "lm_head"),
740
+ )
741
+ else:
742
+ self.lm_head = PPMissingLayer()
743
+
744
+ self.logits_processor = LogitsProcessor(config.vocab_size)
745
+ self.make_empty_intermediate_tensors = (
746
+ self.model.make_empty_intermediate_tensors
747
+ )
748
+
749
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
750
+ return self.model.embed_input_ids(input_ids)
751
+
752
+ def forward(
753
+ self,
754
+ input_ids: torch.Tensor | None,
755
+ positions: torch.Tensor,
756
+ intermediate_tensors: IntermediateTensors | None = None,
757
+ inputs_embeds: torch.Tensor | None = None,
758
+ ) -> torch.Tensor | IntermediateTensors:
759
+ hidden_states = self.model(
760
+ input_ids, positions, intermediate_tensors, inputs_embeds
761
+ )
762
+ return hidden_states
763
+
764
+ def compute_logits(
765
+ self,
766
+ hidden_states: torch.Tensor,
767
+ ) -> torch.Tensor | None:
768
+ logits = self.logits_processor(self.lm_head, hidden_states)
769
+ return logits
770
+
771
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
772
+ return self.model.get_expert_mapping()
773
+
774
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
775
+ loader = AutoWeightsLoader(self)
776
+ return loader.load_weights(weights)