@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 845adb3ec6

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 14, 05:30 UTC

[Model] Add LongCat-Flash (#23991)

Signed-off-by: yangxurui <yangxurui@meituan.com> Co-authored-by: yangxurui <yangxurui@meituan.com>

XuruiYang committed Sep 25, 2025 at 12:53 UTC 845adb3ec6d7123a4f4251c121d406524aa76548

31 files changed +1357 -66

csrc/moe/moe_align_sum_kernels.cu

+8 -2

@@ -44,6 +44,9 @@ __global__ void moe_align_block_size_kernel(

45 for (size_t i = tid; i < numel; i += stride) {

46 int expert_id = topk_ids[i];

47 + if (expert_id >= num_experts) {

48 + continue;

49 + }

50 int warp_idx = expert_id / experts_per_warp;

51 int expert_offset = expert_id % experts_per_warp;

52 atomicAdd(&shared_counts[warp_idx * experts_per_warp + expert_offset], 1);

@@ -95,12 +98,15 @@ template <typename scalar_t>

98 __global__ void count_and_sort_expert_tokens_kernel(

99 const scalar_t* __restrict__ topk_ids,

100 int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer,

98 - size_t numel) {

101 + size_t numel, int32_t num_experts) {

102 const size_t tid = blockIdx.x * blockDim.x + threadIdx.x;

103 const size_t stride = blockDim.x * gridDim.x;

104

105 for (size_t i = tid; i < numel; i += stride) {

106 int32_t expert_id = topk_ids[i];

107 + if (expert_id >= num_experts) {

108 + continue;

109 + }

110 int32_t rank_post_pad = atomicAdd(&cumsum_buffer[expert_id], 1);

111 sorted_token_ids[rank_post_pad] = i;

112 }

@@ -269,7 +275,7 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,

275 sort_kernel<<<actual_blocks, block_threads, 0, stream>>>(

276 topk_ids.data_ptr<scalar_t>(),

277 sorted_token_ids.data_ptr<int32_t>(),

272 - cumsum_buffer.data_ptr<int32_t>(), topk_ids.numel());

278 + cumsum_buffer.data_ptr<int32_t>(), topk_ids.numel(), num_experts);

279 }

280 });

281 }

docs/models/supported_models.md

@@ -428,6 +428,7 @@ th {

431 +| `LongcatFlashForCausalLM` | LongCat-Flash | `meituan-longcat/LongCat-Flash-Chat`, `meituan-longcat/LongCat-Flash-Chat-FP8` | ✅︎ |✅︎ | ✅︎ |

432

433 Some models are supported only via the [Transformers backend](#transformers). The purpose of the table below is to acknowledge models which we officially support in this way. The logs will say that the Transformers backend is being used, and you will see no warning that this is fallback behaviour. This means that, if you have issues with any of the models listed below, please [make an issue](https://github.com/vllm-project/vllm/issues/new/choose) and we'll do our best to fix it!

434

tests/kernels/moe/test_flashinfer.py

+2 -2

@@ -138,7 +138,7 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph(

138 td = TestData.make_moe_tensors_8bit(m, k, n, e, reorder=True)

139

140 score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16)

141 - topk_weights, topk_ids = FusedMoE.select_experts(

141 + topk_weights, topk_ids, _ = FusedMoE.select_experts(

142 hidden_states=td.hidden_states,

143 router_logits=score,

144 use_grouped_topk=False,

@@ -206,7 +206,7 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(

206 td = TestData.make_moe_tensors_8bit(m, k, n, e, reorder=False)

207

208 score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16)

209 - topk_weights, topk_ids = FusedMoE.select_experts(

209 + topk_weights, topk_ids, _ = FusedMoE.select_experts(

210 hidden_states=td.hidden_states,

211 router_logits=score,

212 use_grouped_topk=False,

tests/models/registry.py

@@ -273,6 +273,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

273 is_available_online=False),

274 "Llama4ForCausalLM": _HfExamplesInfo("meta-llama/Llama-4-Scout-17B-16E-Instruct", # noqa: E501

275 is_available_online=False),

276 + "LongcatFlashForCausalLM": _HfExamplesInfo

277 + ("meituan-longcat/LongCat-Flash-Chat", trust_remote_code=True),

278 "MambaForCausalLM": _HfExamplesInfo("state-spaces/mamba-130m-hf"),

279 "Mamba2ForCausalLM": _HfExamplesInfo("mistralai/Mamba-Codestral-7B-v0.1",

280 min_transformers_version="4.55.3",

@@ -639,6 +641,10 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {

641 speculative_model="zai-org/GLM-4.5",

642 min_transformers_version="4.54",

643 is_available_online=False),

644 + "LongCatFlashMTPModel": _HfExamplesInfo(

645 + "meituan-longcat/LongCat-Flash-Chat",

646 + trust_remote_code=True,

647 + speculative_model="meituan-longcat/LongCat-Flash-Chat"),

648 "MiMoMTPModel": _HfExamplesInfo("XiaomiMiMo/MiMo-7B-RL",

649 trust_remote_code=True,

650 speculative_model="XiaomiMiMo/MiMo-7B-RL"),

tests/models/utils.py

+9 -3

@@ -428,9 +428,8 @@ def dummy_hf_overrides(

428 num_hidden_layers = (3 if model_arch

429 == "Gemma3nForConditionalGeneration" else 1)

430

431 - text_config.update({

431 + update_dict = {

432 "num_layers": num_layers,

433 - "num_hidden_layers": num_hidden_layers,

433 "num_experts": num_experts,

434 "num_experts_per_tok": 2,

435 "num_local_experts": num_experts,

@@ -440,7 +439,14 @@ def dummy_hf_overrides(

439 "n_routed_experts": num_experts,

440 # For Gemma-3n

441 "num_kv_shared_layers": 1,

443 - })

442 + }

443 +

444 + # Update num_hidden_layers for non-Longcat architectures

445 + if model_arch != "LongcatFlashForCausalLM" \

446 + and model_arch != "LongCatFlashMTPModel":

447 + update_dict["num_hidden_layers"] = num_hidden_layers

448 +

449 + text_config.update(update_dict)

450

451 if hasattr(hf_config, "vision_config"):

452 hf_config.vision_config.update({

tests/test_routing_simulator.py

+1 -1

@@ -96,7 +96,7 @@ def test_routing_strategy_integration(monkeypatch, device):

96 envs.environment_variables[env_name] = lambda s=strategy: s

98 # Test the select_experts method

99 - topk_weights, topk_ids = FusedMoE.select_experts(

99 + topk_weights, topk_ids, _ = FusedMoE.select_experts(

100 hidden_states=hidden_states,

101 router_logits=router_logits,

102 top_k=top_k,

vllm/config/model.py

+5 -1

@@ -1131,7 +1131,8 @@ class ModelConfig:

1131 if not hasattr(self.hf_text_config, "model_type"):

1132 return False

1133 elif self.hf_text_config.model_type in \

1134 - ('deepseek_v2', 'deepseek_v3', 'deepseek_mtp', 'kimi_k2'):

1134 + ('deepseek_v2', 'deepseek_v3', 'deepseek_mtp',

1135 + 'kimi_k2', 'longcat_flash'):

1136 return self.hf_text_config.kv_lora_rank is not None

1137 elif self.hf_text_config.model_type == 'eagle':

1138 # if the model is an EAGLE module, check for the

@@ -1257,6 +1258,9 @@ class ModelConfig:

1258 or self.hf_config.model_type == "qwen3_next_mtp"):

1259 total_num_hidden_layers = getattr(self.hf_text_config,

1260 "num_nextn_predict_layers", 0)

1261 + elif (self.hf_config.model_type == "longcat_flash_mtp"):

1262 + total_num_hidden_layers = getattr(self.hf_text_config,

1263 + "num_nextn_predict_layers", 1)

1264 else:

1265 total_num_hidden_layers = getattr(self.hf_text_config,

1266 "num_hidden_layers", 0)

vllm/config/speculative.py

+19 -2

@@ -31,7 +31,8 @@ logger = init_logger(__name__)

32 SpeculativeMethod = Literal["ngram", "eagle", "eagle3", "medusa",

33 "mlp_speculator", "draft_model", "deepseek_mtp",

34 - "ernie_mtp", "qwen3_next_mtp", "mimo_mtp"]

34 + "ernie_mtp", "qwen3_next_mtp", "mimo_mtp",

35 + "longcat_flash_mtp"]

38 @config

@@ -186,6 +187,13 @@ class SpeculativeConfig:

187 "n_predict": n_predict,

188 "architectures": ["Qwen3NextMTP"]

189 })

190 + if hf_config.model_type == "longcat_flash":

191 + hf_config.model_type = "longcat_flash_mtp"

192 + n_predict = getattr(hf_config, "num_nextn_predict_layers", 1)

193 + hf_config.update({

194 + "n_predict": n_predict,

195 + "architectures": ["LongCatFlashMTPModel"]

196 + })

197

198 return hf_config

199

@@ -332,6 +340,15 @@ class SpeculativeConfig:

340 "one layer. Might need some code changes " \

341 "to support multiple layers."

342 )

343 + elif (self.draft_model_config.hf_config.model_type

344 + in ("longcat_flash_mtp")):

345 + self.method = "longcat_flash_mtp"

346 + if self.num_speculative_tokens > 1:

347 + logger.warning(

348 + "LongCat MTP models only have " \

349 + "one layer. Might need some code changes " \

350 + "to support multiple layers."

351 + )

352 else:

353 self.method = "draft_model"

354 raise NotImplementedError(

@@ -548,7 +565,7 @@ class SpeculativeConfig:

565

566 def use_eagle(self) -> bool:

567 return self.method in ("eagle", "eagle3", "deepseek_mtp", "ernie_mtp",

551 - "qwen3_next_mtp")

568 + "qwen3_next_mtp", "longcat_flash_mtp")

569

570 def __repr__(self) -> str:

571 method = self.method

vllm/model_executor/layers/fused_moe/fused_moe.py

+89

@@ -664,6 +664,76 @@ def invoke_fused_moe_kernel(A: torch.Tensor,

664 )

665

666

667 +@triton.jit

668 +def compute_identity_kernel(

669 + top_k: int,

670 + hidden_states_ptr: tl.tensor,

671 + expert_scales_ptr: tl.tensor,

672 + num_tokens: int,

673 + output_ptr: tl.tensor,

674 + hidden_dim: int,

675 + scales_stride: int,

676 + BLOCK_SIZE: tl.constexpr,

677 +) -> None:

678 + pid = tl.program_id(0)

679 +

680 + batch_id = pid // (hidden_dim // BLOCK_SIZE)

681 + dim_offset = pid % (hidden_dim // BLOCK_SIZE) * BLOCK_SIZE

682 +

683 + if batch_id >= num_tokens or dim_offset >= hidden_dim:

684 + return

685 +

686 + h = tl.load(hidden_states_ptr + batch_id * hidden_dim + dim_offset +

687 + tl.arange(0, BLOCK_SIZE),

688 + mask=(dim_offset + tl.arange(0, BLOCK_SIZE)) < hidden_dim)

689 +

690 + result = tl.zeros([BLOCK_SIZE], dtype=tl.float32)

691 + for i in range(top_k):

692 + scale = tl.load(expert_scales_ptr + batch_id * scales_stride + i)

693 + result += h * scale

694 +

695 + tl.store(output_ptr + batch_id * hidden_dim + dim_offset +

696 + tl.arange(0, BLOCK_SIZE),

697 + result,

698 + mask=(dim_offset + tl.arange(0, BLOCK_SIZE)) < hidden_dim)

699 +

700 +

701 +def zero_experts_compute_triton(expert_indices: torch.Tensor,

702 + expert_scales: torch.Tensor, num_experts: int,

703 + zero_expert_type: str,

704 + hidden_states: torch.Tensor) -> torch.Tensor:

705 + N = expert_indices.numel()

706 + top_k = expert_indices.size(-1)

707 + grid = lambda meta: (triton.cdiv(N, meta['BLOCK_SIZE']), )

708 +

709 + if zero_expert_type == "identity":

710 + zero_expert_mask = expert_indices < num_experts

711 + zero_expert_scales = expert_scales.clone()

712 + zero_expert_scales[zero_expert_mask] = 0.0

713 +

714 + normal_expert_mask = expert_indices >= num_experts

715 + expert_indices[normal_expert_mask] = 0

716 + expert_scales[normal_expert_mask] = 0.0

717 +

718 + output = torch.zeros_like(hidden_states).to(hidden_states.device)

719 + hidden_dim = hidden_states.size(-1)

720 + num_tokens = hidden_states.size(0)

721 +

722 + grid = lambda meta: (num_tokens * (hidden_dim // meta['BLOCK_SIZE']), )

723 + compute_identity_kernel[grid](

724 + top_k,

725 + hidden_states,

726 + zero_expert_scales,

727 + num_tokens,

728 + output,

729 + hidden_dim,

730 + zero_expert_scales.stride(0),

731 + BLOCK_SIZE=256,

732 + )

733 +

734 + return output

735 +

736 +

737 # Adapted from: https://github.com/sgl-project/sglang/pull/2628

738 def get_config_file_name(E: int,

739 N: int,

@@ -940,6 +1010,25 @@ def fused_topk(

1010 return topk_weights, topk_ids, token_expert_indices

1011

1012

1013 +def fused_topk_bias(

1014 + hidden_states: torch.Tensor,

1015 + gating_output: torch.Tensor,

1016 + e_score_correction_bias: torch.Tensor,

1017 + topk: int,

1018 + renormalize: bool,

1019 +):

1020 + n_routed_experts = gating_output.shape[-1]

1021 + scores = gating_output.softmax(dim=-1)

1022 + scores_for_choice = scores.view(

1023 + -1, n_routed_experts) + e_score_correction_bias.unsqueeze(0)

1024 + topk_indices = torch.topk(scores_for_choice, k=topk, dim=-1,

1025 + sorted=False)[1]

1026 + topk_weights = scores.gather(1, topk_indices)

1027 + if renormalize:

1028 + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)

1029 + return topk_weights.to(torch.float32), topk_indices.to(torch.int32)

1030 +

1031 +

1032 # This is used by the Deepseek-V2 and Deepseek-V3 model

1033 @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)

1034 def grouped_topk(

vllm/model_executor/layers/fused_moe/layer.py

+72 -15

@@ -24,6 +24,8 @@ from vllm.model_executor.custom_op import CustomOp

24 from vllm.model_executor.layers.fused_moe.config import (

25 FUSED_MOE_UNQUANTIZED_CONFIG, FusedMoEConfig, FusedMoEParallelConfig,

26 FusedMoEQuantConfig, biased_moe_quant_config)

27 +from vllm.model_executor.layers.fused_moe.fused_moe import (

28 + zero_experts_compute_triton)

29 # yapf: enable

30 from vllm.model_executor.layers.fused_moe.modular_kernel import (

31 FusedMoEActivationFormat, FusedMoEModularKernel,

@@ -548,7 +550,10 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

550 logical_replica_count: Optional[torch.Tensor] = None,

551 ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:

552

551 - topk_weights, topk_ids = FusedMoE.select_experts(

553 + zero_expert_num = getattr(layer, 'zero_expert_num', 0)

554 + zero_expert_type = getattr(layer, 'zero_expert_type', None)

555 +

556 + topk_weights, topk_ids, zero_expert_result = FusedMoE.select_experts(

557 hidden_states=x,

558 router_logits=router_logits,

559 use_grouped_topk=use_grouped_topk,

@@ -565,11 +570,14 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

570 expert_map=expert_map,

571 expert_load_view=expert_load_view,

572 logical_to_physical_map=logical_to_physical_map,

568 - logical_replica_count=logical_replica_count)

573 + logical_replica_count=logical_replica_count,

574 + global_num_experts=global_num_experts,

575 + zero_expert_num=zero_expert_num,

576 + zero_expert_type=zero_expert_type)

577

578 if self.rocm_aiter_moe_enabled:

579 assert self.fused_experts is None

572 - return self.rocm_aiter_fused_experts(

580 + result = self.rocm_aiter_fused_experts(

581 hidden_states=x,

582 w1=layer.w13_weight,

583 w2=layer.w2_weight,

@@ -591,7 +599,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

599 if self.moe.has_bias:

600 raise ValueError(

601 "FusedMoEModularKernel does not support bias.")

594 - return self.fused_experts(

602 + result = self.fused_experts(

603 hidden_states=x,

604 w1=layer.w13_weight,

605 w2=layer.w2_weight,

@@ -605,7 +613,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

613 )

614 else:

615 assert fused_experts is not None

608 - return fused_experts(

616 + result = fused_experts(

617 hidden_states=x,

618 w1=layer.w13_weight,

619 w2=layer.w2_weight,

@@ -619,6 +627,13 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

627 expert_map=expert_map,

628 )

629

630 + if zero_expert_num != 0 and zero_expert_type is not None:

631 + assert not isinstance(result, tuple), \

632 + "Shared + zero experts are mutually exclusive not yet supported"

633 + return result, zero_expert_result

634 + else:

635 + return result

636 +

637 def forward_cpu(

638 self,

639 layer: torch.nn.Module,

@@ -942,6 +957,8 @@ class FusedMoE(CustomOp):

957 num_redundant_experts: int = 0,

958 has_bias: bool = False,

959 is_sequence_parallel=False,

960 + zero_expert_num: Optional[int] = 0,

961 + zero_expert_type: Optional[str] = None,

962 ):

963 super().__init__()

964 if params_dtype is None:

@@ -976,6 +993,8 @@ class FusedMoE(CustomOp):

993 vllm_parallel_config=vllm_config.parallel_config))

994

995 self.global_num_experts = num_experts + num_redundant_experts

996 + self.zero_expert_num = zero_expert_num

997 + self.zero_expert_type = zero_expert_type

998

999 # Round up hidden size if needed.

1000 hidden_size = maybe_roundup_hidden_size(hidden_size, moe_in_dtype,

@@ -1656,25 +1675,30 @@ class FusedMoE(CustomOp):

1675 expert_load_view: Optional[torch.Tensor] = None,

1676 logical_to_physical_map: Optional[torch.Tensor] = None,

1677 logical_replica_count: Optional[torch.Tensor] = None,

1659 - ) -> tuple[torch.Tensor, torch.Tensor]:

1678 + global_num_experts: Optional[int] = None,

1679 + zero_expert_num: Optional[int] = None,

1680 + zero_expert_type: Optional[str] = None,

1681 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

1682 """

1683 Route the input hidden states to the top-k experts based on the

1684 router logits.

1685

1686 Returns:

1665 - (topk_weights, topk_ids) (tuple[torch.Tensor, torch.Tensor]):

1666 - The weights and *global physical* expert ids of the top-k experts.

1687 + (topk_weights, topk_ids, zero_expert_result)

1688 + (tuple[torch.Tensor, torch.Tensor, torch.Tensor]):

1689 + The weights, expert ids, and zero expert computation result.

1690

1691 **Compatibility**: When EPLB is not enabled, the returned ids are

1692 equivalent to global logical ids, so should be compatible with

1693 plain MoE implementations without redundant experts.

1694 """

1672 - from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk

1695 + from vllm.model_executor.layers.fused_moe.fused_moe import (

1696 + fused_topk, fused_topk_bias)

1697

1698 # Check if we should use a routing simulation strategy

1699 routing_strategy = envs.VLLM_MOE_ROUTING_SIMULATION_STRATEGY

1700 if routing_strategy != "":

1677 - return RoutingSimulator.simulate_routing(

1701 + topk_weights, topk_ids = RoutingSimulator.simulate_routing(

1702 hidden_states=hidden_states,

1703 router_logits=router_logits,

1704 strategy_name=routing_strategy,

@@ -1697,6 +1721,16 @@ class FusedMoE(CustomOp):

1721 e_score_correction_bias=e_score_correction_bias)

1722 if indices_type is not None:

1723 topk_ids = topk_ids.to(dtype=indices_type)

1724 + elif e_score_correction_bias is not None:

1725 + topk_weights, topk_ids = fused_topk_bias(

1726 + hidden_states=hidden_states,

1727 + gating_output=router_logits,

1728 + e_score_correction_bias=e_score_correction_bias.data,

1729 + topk=top_k,

1730 + renormalize=renormalize,

1731 + )

1732 + if routed_scaling_factor is not None:

1733 + topk_weights *= routed_scaling_factor

1734 elif custom_routing_function is None:

1735 topk_weights, topk_ids, token_expert_indices = fused_topk(

1736 hidden_states=hidden_states,

@@ -1729,7 +1763,20 @@ class FusedMoE(CustomOp):

1763

1764 assert topk_ids.dtype == indices_type or indices_type is None

1765

1732 - return topk_weights, topk_ids

1766 + # Compute zero expert result if needed

1767 + if (zero_expert_num is not None and zero_expert_num > 0

1768 + and zero_expert_type is not None

1769 + and global_num_experts is not None):

1770 + zero_expert_result = zero_experts_compute_triton(

1771 + expert_indices=topk_ids,

1772 + expert_scales=topk_weights,

1773 + num_experts=global_num_experts,

1774 + zero_expert_type=zero_expert_type,

1775 + hidden_states=hidden_states,

1776 + )

1777 + else:

1778 + zero_expert_result = None

1779 + return topk_weights, topk_ids, zero_expert_result

1780

1781 def must_reduce_shared_expert_outputs(self) -> bool:

1782 """

@@ -1878,6 +1925,11 @@ class FusedMoE(CustomOp):

1925 assert self.shared_experts is None or isinstance(

1926 final_hidden_states, tuple)

1927

1928 + if isinstance(final_hidden_states, tuple):

1929 + final_hidden_states, zero_expert_result = final_hidden_states

1930 + if zero_expert_result is not None:

1931 + final_hidden_states += zero_expert_result

1932 +

1933 if not skip_result_store:

1934 if self.shared_experts is None:

1935 full_fused_final_hidden_states[

@@ -1992,6 +2044,9 @@ class FusedMoE(CustomOp):

2044 shared_output,

2045 final_hidden_states,

2046 )

2047 + elif self.zero_expert_num is not None and self.zero_expert_num > 0:

2048 + assert isinstance(final_hidden_states, tuple)

2049 + final_hidden_states, zero_expert_result = final_hidden_states

2050

2051 def reduce_output(states: torch.Tensor,

2052 do_combine: bool = True) -> torch.Tensor:

@@ -2003,14 +2058,16 @@ class FusedMoE(CustomOp):

2058

2059 return states

2060

2006 - if self.shared_experts is None:

2007 - assert not isinstance(final_hidden_states, tuple)

2008 - return reduce_output(final_hidden_states)

2009 - else:

2061 + if self.shared_experts is not None:

2062 return (

2063 reduce_output(final_hidden_states[0], do_combine=False),

2064 reduce_output(final_hidden_states[1]),

2065 )

2066 + elif self.zero_expert_num is not None and self.zero_expert_num > 0:

2067 + assert isinstance(final_hidden_states, torch.Tensor)

2068 + return reduce_output(final_hidden_states) + zero_expert_result

2069 + else:

2070 + return reduce_output(final_hidden_states)

2071

2072 @classmethod

2073 def make_expert_params_mapping(

vllm/model_executor/layers/mla.py

-1

@@ -103,7 +103,6 @@ class MultiHeadLatentAttention(CustomOp):

103 )

104

105 self.prefix = prefix

106 - self.debug_layer_idx = int(self.prefix.split(".")[-2])

106

107 def forward_native(

108 self,

vllm/model_executor/layers/quantization/awq_marlin.py

+1 -1

@@ -520,7 +520,7 @@ class AWQMoEMethod(FusedMoEMethodBase):

520

521 assert activation == "silu", "Only SiLU activation is supported."

522

523 - topk_weights, topk_ids = FusedMoE.select_experts(

523 + topk_weights, topk_ids, _ = FusedMoE.select_experts(

524 hidden_states=x,

525 router_logits=router_logits,

526 use_grouped_topk=use_grouped_topk,

vllm/model_executor/layers/quantization/bitsandbytes.py

+1 -1

@@ -486,7 +486,7 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase):

486 if enable_eplb:

487 raise NotImplementedError(

488 "EPLB not supported for `BitsAndBytesMoEMethod` yet.")

489 - topk_weights, topk_ids = FusedMoE.select_experts(

489 + topk_weights, topk_ids, _ = FusedMoE.select_experts(

490 hidden_states=x,

491 router_logits=router_logits,

492 use_grouped_topk=use_grouped_topk,

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py

+5 -5

@@ -385,7 +385,7 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):

385 "`CompressedTensorsW4A4MoeMethod` yet.")

386 assert activation == "silu", "Only SiLU activation is supported."

387

388 - topk_weights, topk_ids = FusedMoE.select_experts(

388 + topk_weights, topk_ids, _ = FusedMoE.select_experts(

389 hidden_states=x,

390 router_logits=router_logits,

391 use_grouped_topk=use_grouped_topk,

@@ -934,7 +934,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

934 "EPLB not supported for "

935 "`CompressedTensorsW8A8Fp8MoEMethod` yet.")

936

937 - topk_weights, topk_ids = FusedMoE.select_experts(

937 + topk_weights, topk_ids, _ = FusedMoE.select_experts(

938 hidden_states=x,

939 router_logits=router_logits,

940 use_grouped_topk=use_grouped_topk,

@@ -1195,7 +1195,7 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):

1195

1196 from vllm.model_executor.layers.fused_moe import fused_experts

1197

1198 - topk_weights, topk_ids = FusedMoE.select_experts(

1198 + topk_weights, topk_ids, _ = FusedMoE.select_experts(

1199 hidden_states=x,

1200 router_logits=router_logits,

1201 use_grouped_topk=use_grouped_topk,

@@ -1502,7 +1502,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

1502 assert activation == "silu", (

1503 f"{activation} not supported for Marlin MoE.")

1504

1505 - topk_weights, topk_ids = FusedMoE.select_experts(

1505 + topk_weights, topk_ids, _ = FusedMoE.select_experts(

1506 hidden_states=x,

1507 router_logits=router_logits,

1508 use_grouped_topk=use_grouped_topk,

@@ -1747,7 +1747,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):

1747

1748 from vllm.model_executor.layers.fused_moe import fused_experts

1749

1750 - topk_weights, topk_ids = FusedMoE.select_experts(

1750 + topk_weights, topk_ids, _ = FusedMoE.select_experts(

1751 hidden_states=x,

1752 router_logits=router_logits,

1753 use_grouped_topk=use_grouped_topk,

vllm/model_executor/layers/quantization/experts_int8.py

+1 -1

@@ -146,7 +146,7 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):

146

147 from vllm.model_executor.layers.fused_moe import fused_experts

148

149 - topk_weights, topk_ids = FusedMoE.select_experts(

149 + topk_weights, topk_ids, _ = FusedMoE.select_experts(

150 hidden_states=x,

151 router_logits=router_logits,

152 use_grouped_topk=use_grouped_topk,

vllm/model_executor/layers/quantization/fp8.py

+29 -9

@@ -18,6 +18,8 @@ from vllm.model_executor.layers.fused_moe import (

18 FusedMoeWeightScaleSupported)

19 from vllm.model_executor.layers.fused_moe.config import (

20 FusedMoEQuantConfig, fp8_w8a8_moe_quant_config)

21 +from vllm.model_executor.layers.fused_moe.layer import (

22 + UnquantizedFusedMoEMethod)

23 from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,

24 UnquantizedLinearMethod)

25 from vllm.model_executor.layers.quantization import QuantizationMethods

@@ -174,6 +176,10 @@ class Fp8Config(QuantizationConfig):

176 return UnquantizedLinearMethod()

177 return Fp8LinearMethod(self)

178 elif isinstance(layer, FusedMoE):

179 + if is_layer_skipped(prefix=prefix,

180 + ignored_layers=self.ignored_layers,

181 + fused_mapping=self.packed_modules_mapping):

182 + return UnquantizedFusedMoEMethod(layer.moe_config)

183 return Fp8MoEMethod(self, layer)

184 elif isinstance(layer, Attention):

185 return Fp8KVCacheMethod(self)

@@ -927,6 +933,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

933 logical_to_physical_map: Optional[torch.Tensor] = None,

934 logical_replica_count: Optional[torch.Tensor] = None,

935 ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:

936 +

937 if enable_eplb:

938 assert expert_load_view is not None

939 assert logical_to_physical_map is not None

@@ -943,8 +950,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

950 import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: E501, F401

951 assert (renormalize and use_grouped_topk

952 and custom_routing_function is None)

946 -

947 - return torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8(

953 + result = torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8(

954 routing_logits=router_logits.to(torch.float32),

955 routing_bias=e_score_correction_bias,

956 x=x,

@@ -965,7 +971,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

971 else:

972 assert (not renormalize

973 and custom_routing_function is not None)

968 - return apply_flashinfer_per_tensor_scale_fp8(

974 + result = apply_flashinfer_per_tensor_scale_fp8(

975 layer=layer,

976 hidden_states=x,

977 router_logits=router_logits,

@@ -976,7 +982,10 @@ class Fp8MoEMethod(FusedMoEMethodBase):

982 topk_group=topk_group,

983 apply_router_weight_on_input=apply_router_weight_on_input)

984

979 - topk_weights, topk_ids = FusedMoE.select_experts(

985 + zero_expert_num = getattr(layer, 'zero_expert_num', 0)

986 + zero_expert_type = getattr(layer, 'zero_expert_type', None)

987 +

988 + select_result = FusedMoE.select_experts(

989 hidden_states=x,

990 router_logits=router_logits,

991 use_grouped_topk=use_grouped_topk,

@@ -994,17 +1003,22 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1003 expert_load_view=expert_load_view,

1004 logical_to_physical_map=logical_to_physical_map,

1005 logical_replica_count=logical_replica_count,

1006 + global_num_experts=global_num_experts,

1007 + zero_expert_num=zero_expert_num,

1008 + zero_expert_type=zero_expert_type,

1009 )

1010

1011 #

1012 # Note: the order of checks is important since self.fused_experts

1013 # can override fused_experts or cutlass but not rocm or marlin.

1014 #

1015 + topk_weights, topk_ids, zero_expert_result = select_result

1016 +

1017 if self.rocm_aiter_moe_enabled:

1018 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( # noqa: E501

1019 rocm_aiter_fused_experts)

1020 assert self.fused_experts is None

1007 - return rocm_aiter_fused_experts(

1021 + result = rocm_aiter_fused_experts(

1022 x,

1023 layer.w13_weight,

1024 layer.w2_weight,

@@ -1018,7 +1032,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1032 assert activation == "silu", (

1033 f"{activation} not supported for Marlin MoE.")

1034 assert self.fused_experts is None

1021 - return torch.ops.vllm.fused_marlin_moe(

1035 + result = torch.ops.vllm.fused_marlin_moe(

1036 x,

1037 layer.w13_weight,

1038 layer.w2_weight,

@@ -1035,7 +1049,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1049 expert_map=expert_map,

1050 workspace=layer.workspace)

1051 elif self.fused_experts:

1038 - return self.fused_experts(

1052 + result = self.fused_experts(

1053 hidden_states=x,

1054 w1=layer.w13_weight,

1055 w2=layer.w2_weight,

@@ -1055,7 +1069,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1069 assert scoring_func == 'sigmoid', (

1070 f"Expected 'sigmoid' scoring func but got {scoring_func}")

1071

1058 - return flashinfer_cutlass_moe_fp8(

1072 + result = flashinfer_cutlass_moe_fp8(

1073 x,

1074 layer,

1075 topk_weights,

@@ -1068,7 +1082,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1082 )

1083 else:

1084 from vllm.model_executor.layers.fused_moe import fused_experts

1071 - return fused_experts(

1085 + result = fused_experts(

1086 hidden_states=x,

1087 w1=layer.w13_weight,

1088 w2=layer.w2_weight,

@@ -1083,6 +1097,12 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1097 allow_deep_gemm=self.allow_deep_gemm,

1098 allow_cutlass_block_scaled_grouped_gemm=(

1099 self.allow_cutlass_block_scaled_grouped_gemm))

1100 + if zero_expert_num != 0 and zero_expert_type is not None:

1101 + assert not isinstance(result, tuple), \

1102 + "Shared + zero experts are mutually exclusive not yet supported"

1103 + return result, zero_expert_result

1104 + else:

1105 + return result

1106

1107

1108 class Fp8KVCacheMethod(BaseKVCacheMethod):

vllm/model_executor/layers/quantization/gguf.py

+1 -1

@@ -555,7 +555,7 @@ class GGUFMoEMethod(FusedMoEMethodBase):

555 "Apply router weight on input is not supported for"

556 "fused GGUF MoE method.")

557

558 - topk_weights, topk_ids = FusedMoE.select_experts(

558 + topk_weights, topk_ids, _ = FusedMoE.select_experts(

559 hidden_states=x,

560 router_logits=router_logits,

561 use_grouped_topk=use_grouped_topk,

vllm/model_executor/layers/quantization/gptq_marlin.py

+1 -1

@@ -669,7 +669,7 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):

669

670 assert activation == "silu", "Only SiLU activation is supported."

671

672 - topk_weights, topk_ids = FusedMoE.select_experts(

672 + topk_weights, topk_ids, _ = FusedMoE.select_experts(

673 hidden_states=x,

674 router_logits=router_logits,

675 use_grouped_topk=use_grouped_topk,

vllm/model_executor/layers/quantization/modelopt.py

+2 -2

@@ -543,7 +543,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

543 apply_router_weight_on_input=apply_router_weight_on_input)

544

545 # Expert selection

546 - topk_weights, topk_ids = FusedMoE.select_experts(

546 + topk_weights, topk_ids, _ = FusedMoE.select_experts(

547 hidden_states=x,

548 router_logits=router_logits,

549 use_grouped_topk=use_grouped_topk,

@@ -1491,7 +1491,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1491 )[0]

1492 return out

1493

1494 - topk_weights, topk_ids = FusedMoE.select_experts(

1494 + topk_weights, topk_ids, _ = FusedMoE.select_experts(

1495 hidden_states=x,

1496 router_logits=router_logits,

1497 use_grouped_topk=use_grouped_topk,

vllm/model_executor/layers/quantization/moe_wna16.py

+1 -1

@@ -332,7 +332,7 @@ class MoeWNA16Method(FusedMoEMethodBase):

332

333 from vllm.model_executor.layers.fused_moe import fused_experts

334 assert activation == "silu", "Only SiLU activation is supported."

335 - topk_weights, topk_ids = FusedMoE.select_experts(

335 + topk_weights, topk_ids, _ = FusedMoE.select_experts(

336 hidden_states=x,

337 router_logits=router_logits,

338 use_grouped_topk=use_grouped_topk,

vllm/model_executor/layers/quantization/mxfp4.py

+3 -3

@@ -718,7 +718,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

718

719 assert isinstance(self.fused_experts, mk.FusedMoEModularKernel)

720

721 - topk_weights, topk_ids = FusedMoE.select_experts(

721 + topk_weights, topk_ids, _ = FusedMoE.select_experts(

722 hidden_states=x,

723 router_logits=router_logits,

724 use_grouped_topk=use_grouped_topk,

@@ -783,7 +783,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

783 raise NotImplementedError("EPLB is not supported for mxfp4")

784

785 if self.mxfp4_backend == Mxfp4Backend.MARLIN:

786 - topk_weights, topk_ids = FusedMoE.select_experts(

786 + topk_weights, topk_ids, _ = FusedMoE.select_experts(

787 hidden_states=x,

788 router_logits=router_logits,

789 use_grouped_topk=use_grouped_topk,

@@ -894,7 +894,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

894 or self.mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16):

895 from vllm.utils.flashinfer import flashinfer_cutlass_fused_moe

896

897 - topk_weights, topk_ids = FusedMoE.select_experts(

897 + topk_weights, topk_ids, _ = FusedMoE.select_experts(

898 hidden_states=x,

899 router_logits=router_logits,

900 use_grouped_topk=use_grouped_topk,

vllm/model_executor/layers/quantization/quark/quark_moe.py

+2 -2

@@ -329,7 +329,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

329 raise NotImplementedError(

330 "EPLB not supported for `QuarkW8A8Fp8MoEMethod` yet.")

331

332 - topk_weights, topk_ids = FusedMoE.select_experts(

332 + topk_weights, topk_ids, _ = FusedMoE.select_experts(

333 hidden_states=x,

334 router_logits=router_logits,

335 use_grouped_topk=use_grouped_topk,

@@ -531,7 +531,7 @@ class QuarkW4A4MXFp4MoEMethod(QuarkMoEMethod):

531

532 from vllm.model_executor.layers.fused_moe import fused_experts

533

534 - topk_weights, topk_ids = FusedMoE.select_experts(

534 + topk_weights, topk_ids, _ = FusedMoE.select_experts(

535 hidden_states=x,

536 router_logits=router_logits,

537 use_grouped_topk=use_grouped_topk,

vllm/model_executor/layers/quantization/rtn.py

+1 -1

@@ -318,7 +318,7 @@ class RTNMoEMethod(FusedMoEMethodBase):

318

319 from vllm.model_executor.layers.fused_moe import fused_experts

320

321 - topk_weights, topk_ids = FusedMoE.select_experts(

321 + topk_weights, topk_ids, _ = FusedMoE.select_experts(

322 hidden_states=x,

323 router_logits=router_logits,

324 use_grouped_topk=use_grouped_topk,

vllm/model_executor/layers/quantization/utils/quant_utils.py

@@ -292,6 +292,11 @@ def is_layer_skipped(

292 f"Detected some but not all shards of {prefix} "

293 "are quantized. All shards of fused layers "

294 "to have the same precision.")

295 + elif "experts" in prefix:

296 + return any([

297 + prefix in layer_name for layer_name in ignored_layers

298 + if "experts" in layer_name

299 + ])

300 else:

301 is_skipped = prefix in ignored_layers

302

vllm/model_executor/models/longcat_flash.py new

+712

@@ -0,0 +1,712 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

3 +# Apache License, Version 2.0:

4 +# Licensed under the Apache License, Version 2.0 (the "License");

5 +# you may not use this file except in compliance with the License.

6 +# You may obtain a copy of the License at

7 +#

8 +# http://www.apache.org/licenses/LICENSE-2.0

9 +#

10 +# Unless required by applicable law or agreed to in writing, software

11 +# distributed under the License is distributed on an "AS IS" BASIS,

12 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

13 +# See the License for the specific language governing permissions and

14 +# limitations under the License.

15 +#

16 +# MIT License:

17 +# Permission is hereby granted, free of charge, to any person obtaining a copy

18 +# of this software and associated documentation files (the "Software"), to deal

19 +# in the Software without restriction, including without limitation the rights

20 +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell

21 +# copies of the Software, and to permit persons to whom the Software is

22 +# furnished to do so, subject to the following conditions:

23 +#

24 +# The above copyright notice and this permission notice shall be included in all

25 +# copies or substantial portions of the Software.

26 +#

27 +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR

28 +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,

29 +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE

30 +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER

31 +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,

32 +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE

33 +# SOFTWARE.

34 +"""Inference-only Flash model compatible with HuggingFace weights."""

35 +import typing

36 +from collections.abc import Callable, Iterable

37 +from typing import Optional, Union

38 +

39 +import torch

40 +from torch import nn

41 +from transformers import PretrainedConfig

42 +

43 +from vllm.compilation.decorators import support_torch_compile

44 +from vllm.config import CacheConfig, VllmConfig

45 +from vllm.distributed import get_pp_group

46 +from vllm.logger import init_logger

47 +from vllm.model_executor.layers.activation import SiluAndMul

48 +from vllm.model_executor.layers.fused_moe import FusedMoE

49 +from vllm.model_executor.layers.layernorm import RMSNorm

50 +from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,

51 + ReplicatedLinear,

52 + RowParallelLinear)

53 +from vllm.model_executor.layers.logits_processor import LogitsProcessor

54 +from vllm.model_executor.layers.quantization import QuantizationConfig

55 +from vllm.model_executor.layers.quantization.utils.int8_utils import (

56 + block_dequant)

57 +from vllm.model_executor.layers.vocab_parallel_embedding import (

58 + ParallelLMHead, VocabParallelEmbedding)

59 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader

60 +from vllm.model_executor.models.deepseek_v2 import DeepseekV2MLAAttention

61 +from vllm.sequence import IntermediateTensors

62 +

63 +from .interfaces import SupportsLoRA, SupportsPP

64 +from .utils import (PPMissingLayer, is_pp_missing_parameter,

65 + make_empty_intermediate_tensors_factory, make_layers,

66 + maybe_prefix)

67 +

68 +logger = init_logger(__name__)

69 +

70 +

71 +class FlashConfig(PretrainedConfig):

72 + """Flash model configuration."""

73 + model_type = "longcat_flash"

74 + keys_to_ignore_at_inference = ["past_key_values"]

75 +

76 + def __init__(

77 + self,

78 + vocab_size=131072,

79 + hidden_size=4096,

80 + intermediate_size=8192,

81 + num_layers=28,

82 + num_hidden_layers=None,

83 + num_attention_heads=96,

84 + num_key_value_heads=128,

85 + ep_size=1,

86 + kv_lora_rank=512,

87 + q_lora_rank=1536,

88 + qk_rope_head_dim=64,

89 + v_head_dim=128,

90 + qk_nope_head_dim=128,

91 + num_experts_per_tok=None,

92 + norm_topk_prob=False,

93 + max_position_embeddings=8192,

94 + initializer_range=0.02,

95 + rms_norm_eps=1e-05,

96 + use_cache=True,

97 + pad_token_id=None,

98 + bos_token_id=100000,

99 + eos_token_id=100001,

100 + pretraining_tp=1,

101 + tie_word_embeddings=False,

102 + rope_theta=1000000.0,

103 + rope_scaling=None,

104 + attention_bias=False,

105 + attention_dropout=0.0,

106 + mla_scale_q_lora=False,

107 + mla_scale_kv_lora=False,

108 + torch_dtype="bfloat16",

109 + params_dtype="bfloat16",

110 + router_dtype="float32",

111 + router_bias=False,

112 + topk_method=None,

113 + routed_scaling_factor=None,

114 + zero_expert_num=0,

115 + zero_expert_type=None,

116 + nextn_use_scmoe=False,

117 + **kwargs,

118 + ):

119 + super().__init__(

120 + pad_token_id=pad_token_id,

121 + bos_token_id=bos_token_id,

122 + eos_token_id=eos_token_id,

123 + tie_word_embeddings=tie_word_embeddings,

124 + torch_dtype=torch_dtype,

125 + params_dtype=params_dtype,

126 + router_dtype=router_dtype,

127 + topk_method=topk_method,

128 + router_bias=router_bias,

129 + nextn_use_scmoe=nextn_use_scmoe,

130 + **kwargs,

131 + )

132 + self.vocab_size = vocab_size

133 + self.max_position_embeddings = max_position_embeddings

134 + self.hidden_size = hidden_size

135 + self.num_hidden_layers = (num_hidden_layers if num_hidden_layers

136 + is not None else num_layers)

137 + self.num_attention_heads = num_attention_heads

138 + self.ep_size = ep_size

139 + self.kv_lora_rank = kv_lora_rank

140 + self.q_lora_rank = q_lora_rank

141 + self.qk_rope_head_dim = qk_rope_head_dim

142 + self.v_head_dim = v_head_dim

143 + self.qk_nope_head_dim = qk_nope_head_dim

144 + self.num_experts_per_tok = num_experts_per_tok

145 + self.norm_topk_prob = norm_topk_prob

146 + # for backward compatibility

147 + if num_key_value_heads is None:

148 + num_key_value_heads = num_attention_heads

149 +

150 + self.num_key_value_heads = num_key_value_heads

151 + self.initializer_range = initializer_range

152 + self.rms_norm_eps = rms_norm_eps

153 + self.pretraining_tp = pretraining_tp

154 + self.use_cache = use_cache

155 + self.rope_theta = rope_theta

156 + self.rope_scaling = rope_scaling

157 + self.attention_bias = attention_bias

158 + self.attention_dropout = attention_dropout

159 + self.mla_scale_q_lora = mla_scale_q_lora

160 + self.mla_scale_kv_lora = mla_scale_kv_lora

161 + self.zero_expert_num = zero_expert_num

162 + self.zero_expert_type = zero_expert_type

163 + self.routed_scaling_factor = routed_scaling_factor

164 + self.hidden_act = "silu"

165 + self.intermediate_size = self.ffn_hidden_size if hasattr(

166 + self, "ffn_hidden_size") else self.intermediate_size

167 + if hasattr(self, "moe_intermediate_size"):

168 + self.moe_intermediate_size = self.moe_intermediate_size

169 + elif hasattr(self, "expert_ffn_hidden_size"):

170 + self.moe_intermediate_size = self.expert_ffn_hidden_size

171 + else:

172 + self.moe_intermediate_size = self.intermediate_size

173 +

174 +

175 +class FlashMLP(nn.Module):

176 + """Flash MLP layer."""

177 +

178 + def __init__(

179 + self,

180 + hidden_size: int,

181 + intermediate_size: int,

182 + hidden_act: str,

183 + quant_config: Optional[QuantizationConfig] = None,

184 + reduce_results: bool = True,

185 + prefix: str = "",

186 + ) -> None:

187 + super().__init__()

188 + self.gate_up_proj = MergedColumnParallelLinear(

189 + hidden_size,

190 + [intermediate_size] * 2,

191 + bias=False,

192 + quant_config=quant_config,

193 + prefix=f"{prefix}.gate_up_proj",

194 + )

195 + self.down_proj = RowParallelLinear(

196 + intermediate_size,

197 + hidden_size,

198 + bias=False,

199 + quant_config=quant_config,

200 + reduce_results=reduce_results,

201 + prefix=f"{prefix}.down_proj",

202 + )

203 + if hidden_act != "silu":

204 + raise ValueError(f"Unsupported activation: {hidden_act}. "

205 + "Only silu is supported for now.")

206 + self.act_fn = SiluAndMul()

207 +

208 + def forward(self, x: torch.Tensor) -> torch.Tensor:

209 + if x.numel() == 0:

210 + return x

211 +

212 + gate_up, _ = self.gate_up_proj(x)

213 + x = self.act_fn(gate_up)

214 + x, _ = self.down_proj(x)

215 + return x

216 +

217 +

218 +class LongcatRouter(nn.Module):

219 +

220 + def __init__(self,

221 + config,

222 + zero_expert_num=0,

223 + rounter_params_dtype=torch.bfloat16,

224 + prefix: str = ""):

225 + super().__init__()

226 + self.n_routed_experts = config.n_routed_experts if hasattr(

227 + config, "n_routed_experts") else config.num_experts[0]

228 + self.n_routed_experts = self.n_routed_experts + zero_expert_num

229 + self.classifier = ReplicatedLinear(

230 + config.hidden_size,

231 + self.n_routed_experts,

232 + bias=config.router_bias,

233 + params_dtype=rounter_params_dtype,

234 + quant_config=None,

235 + prefix=f"{prefix}.classifier",

236 + )

237 + self.e_score_correction_bias = nn.Parameter(

238 + torch.zeros((self.n_routed_experts), dtype=rounter_params_dtype))

239 +

240 + def forward(self, hidden_states):

241 + logits, _ = self.classifier(hidden_states)

242 + return logits

243 +

244 +

245 +class LongcatMoe(nn.Module):

246 +

247 + def __init__(

248 + self,

249 + config: FlashConfig,

250 + num_experts: int,

251 + top_k: int,

252 + hidden_size: int,

253 + intermediate_size: int,

254 + params_dtype: Optional[torch.dtype] = None,

255 + quant_config: Optional[QuantizationConfig] = None,

256 + prefix: str = "",

257 + enable_eplb: bool = False,

258 + ):

259 + super().__init__()

260 + self.hidden_size = hidden_size

261 + self.zero_expert_num = config.zero_expert_num

262 + self.zero_expert_type = config.zero_expert_type

263 + self.routed_scaling_factor = config.routed_scaling_factor

264 + self.enable_eplb = enable_eplb

265 + # Gate always runs at half / full precision for now.

266 + self.rounter_params_dtype = params_dtype

267 + if config.router_dtype == "float32":

268 + self.rounter_params_dtype = torch.float32

269 +

270 + self.router = LongcatRouter(

271 + config=config,

272 + zero_expert_num=self.zero_expert_num,

273 + rounter_params_dtype=self.rounter_params_dtype,

274 + prefix=f"{prefix}.gate")

275 +

276 + self.experts = FusedMoE(

277 + num_experts=num_experts,

278 + top_k=top_k,

279 + hidden_size=hidden_size,

280 + intermediate_size=intermediate_size,

281 + reduce_results=True,

282 + params_dtype=params_dtype,

283 + e_score_correction_bias=self.router.e_score_correction_bias,

284 + renormalize=False,

285 + quant_config=quant_config,

286 + prefix=f"{prefix}.experts",

287 + zero_expert_num=self.zero_expert_num,

288 + zero_expert_type=self.zero_expert_type,

289 + enable_eplb=self.enable_eplb,

290 + routed_scaling_factor=config.routed_scaling_factor,

291 + )

292 +

293 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

294 +

295 + num_tokens, hidden_dim = hidden_states.shape

296 + hidden_states = hidden_states.view(-1, hidden_dim)

297 +

298 + router_logits = self.router(hidden_states.to(

299 + self.rounter_params_dtype))

300 + final_hidden_states = self.experts(hidden_states=hidden_states,

301 + router_logits=router_logits)

302 +

303 + return final_hidden_states.view(num_tokens, hidden_dim)

304 +

305 +

306 +class FlashDecoderLayer(nn.Module):

307 + """Flash decoder layer with dual attention and MLP structure."""

308 +

309 + def __init__(

310 + self,

311 + config: FlashConfig,

312 + cache_config: Optional[CacheConfig] = None,

313 + quant_config: Optional[QuantizationConfig] = None,

314 + prefix: str = "",

315 + enable_eplb: bool = False,

316 + ) -> None:

317 + super().__init__()

318 + self.layer_idx = int(prefix.split(sep='.')[-1])

319 + self.hidden_size = config.hidden_size

320 + rope_theta = getattr(config, "rope_theta", 10000)

321 + rope_scaling = getattr(config, "rope_scaling", None)

322 + max_position_embeddings = getattr(config, "max_position_embeddings",

323 + 8192)

324 + if rope_scaling is not None and getattr(

325 + config, "original_max_position_embeddings", None):

326 + rope_scaling["original_max_position_embeddings"] = (

327 + config.original_max_position_embeddings)

328 +

329 + # Dual attention structure

330 + self.self_attn = nn.ModuleList([

331 + DeepseekV2MLAAttention(

332 + config=config,

333 + hidden_size=self.hidden_size,

334 + num_heads=config.num_attention_heads,

335 + qk_nope_head_dim=config.qk_nope_head_dim,

336 + qk_rope_head_dim=config.qk_rope_head_dim,

337 + v_head_dim=config.v_head_dim,

338 + q_lora_rank=(config.q_lora_rank if hasattr(

339 + config, "q_lora_rank") else None),

340 + kv_lora_rank=config.kv_lora_rank,

341 + rope_theta=rope_theta,

342 + rope_scaling=rope_scaling,

343 + max_position_embeddings=max_position_embeddings,

344 + cache_config=cache_config,

345 + quant_config=None if "self_attn" in getattr(

346 + config, "disable_quant_module", []) else quant_config,

347 + prefix=f"{prefix}.self_attn.{i}",

348 + ) for i in range(2)

349 + ])

350 + self.input_layernorm = nn.ModuleList([

351 + RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

352 + for i in range(2)

353 + ])

354 + self.post_attention_layernorm = nn.ModuleList([

355 + RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

356 + for i in range(2)

357 + ])

358 +

359 + # Dual MLP structure

360 + self.mlps = nn.ModuleList([

361 + FlashMLP(

362 + hidden_size=self.hidden_size,

363 + intermediate_size=config.intermediate_size,

364 + hidden_act=config.hidden_act,

365 + quant_config=None if "mlps" in getattr(

366 + config, "disable_quant_module", []) else quant_config,

367 + prefix=f"{prefix}.mlps.{i}",

368 + ) for i in range(2)

369 + ])

370 +

371 + self.mlp = LongcatMoe(

372 + config=config,

373 + num_experts=config.n_routed_experts if hasattr(

374 + config, "n_routed_experts") else

375 + config.num_experts[self.layer_idx],

376 + top_k=config.moe_topk

377 + if hasattr(config, "moe_topk") else config.num_experts_per_tok,

378 + hidden_size=config.hidden_size,

379 + intermediate_size=config.moe_intermediate_size,

380 + quant_config=quant_config,

381 + prefix=(f"{prefix}.mlp"),

382 + )

383 +

384 + def forward(

385 + self,

386 + positions: torch.Tensor,

387 + hidden_states: torch.Tensor,

388 + residual: Optional[torch.Tensor],

389 + ) -> tuple[torch.Tensor, torch.Tensor]:

390 +

391 + if residual is None:

392 + residual = hidden_states

393 + hidden_states = self.input_layernorm[0](hidden_states)

394 + else:

395 + hidden_states, residual = self.input_layernorm[0](hidden_states,

396 + residual)

397 +

398 + hidden_states = self.self_attn[0](

399 + positions=positions,

400 + hidden_states=hidden_states,

401 + )

402 +

403 + hidden_states, residual = self.post_attention_layernorm[0](

404 + hidden_states, residual)

405 +

406 + # moe

407 + hidden_states_copy = hidden_states.clone()

408 + moe_hidden_states = self.mlp(hidden_states_copy)

409 +

410 + # first mlp

411 + hidden_states = self.mlps[0](hidden_states)

412 +

413 + hidden_states, residual = self.input_layernorm[1](hidden_states,

414 + residual)

415 +

416 + # second_attn

417 + hidden_states = self.self_attn[1](

418 + positions=positions,

419 + hidden_states=hidden_states,

420 + )

421 + hidden_states, residual = self.post_attention_layernorm[1](

422 + hidden_states, residual)

423 +

424 + # second_mlp

425 + hidden_states = self.mlps[1](hidden_states)

426 +

427 + hidden_states = hidden_states + moe_hidden_states

428 +

429 + return hidden_states, residual

430 +

431 +

432 +@support_torch_compile

433 +class FlashModel(nn.Module):

434 + """Flash model."""

435 +

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

437 + super().__init__()

438 + config = FlashConfig(**vllm_config.model_config.hf_config.__dict__)

439 + cache_config = vllm_config.cache_config

440 + quant_config = vllm_config.quant_config

441 + self.config = config

442 +

443 + self.padding_idx = getattr(config, "pad_token_id", None)

444 + self.vocab_size = config.vocab_size

445 +

446 + if get_pp_group().is_first_rank:

447 + self.embed_tokens = VocabParallelEmbedding(

448 + config.vocab_size,

449 + config.hidden_size,

450 + prefix=maybe_prefix(prefix, "embed_tokens"),

451 + )

452 + else:

453 + self.embed_tokens = PPMissingLayer()

454 + self.start_layer, self.end_layer, self.layers = make_layers(

455 + config.num_hidden_layers,

456 + lambda prefix: FlashDecoderLayer(

457 + config,

458 + cache_config=cache_config,

459 + quant_config=quant_config,

460 + prefix=prefix,

461 + ),

462 + prefix=f"{prefix}.layers")

463 + if get_pp_group().is_last_rank:

464 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

465 + else:

466 + self.norm = PPMissingLayer()

467 + self.make_empty_intermediate_tensors = (

468 + make_empty_intermediate_tensors_factory(

469 + ["hidden_states", "residual"], config.hidden_size))

470 +

471 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

472 + return self.embed_tokens(input_ids)

473 +

474 + def forward(

475 + self,

476 + input_ids: torch.Tensor,

477 + positions: torch.Tensor,

478 + intermediate_tensors: Optional[IntermediateTensors] = None,

479 + inputs_embeds: Optional[torch.Tensor] = None,

480 + ) -> Union[torch.Tensor, IntermediateTensors]:

481 + if get_pp_group().is_first_rank:

482 + if inputs_embeds is not None:

483 + hidden_states = inputs_embeds

484 + else:

485 + hidden_states = self.get_input_embeddings(input_ids)

486 + residual = None

487 + else:

488 + assert intermediate_tensors is not None

489 + hidden_states = intermediate_tensors["hidden_states"]

490 + residual = intermediate_tensors["residual"]

491 +

492 + for i in range(self.start_layer, self.end_layer):

493 + layer = self.layers[i]

494 + hidden_states, residual = layer(

495 + positions,

496 + hidden_states,

497 + residual,

498 + )

499 +

500 + if not get_pp_group().is_last_rank:

501 + return IntermediateTensors({

502 + "hidden_states": hidden_states,

503 + "residual": residual

504 + })

505 +

506 + hidden_states, _ = self.norm(hidden_states, residual)

507 + return hidden_states

508 +

509 +

510 +class LongcatFlashForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

511 + """Flash model for causal language modeling."""

512 +

513 + packed_modules_mapping = {

514 + "qkv_proj": [

515 + "q_proj",

516 + "k_proj",

517 + "v_proj",

518 + ],

519 + "gate_up_proj": [

520 + "gate_proj",

521 + "up_proj",

522 + ],

523 + }

524 +

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

526 + super().__init__()

527 + config = FlashConfig(**vllm_config.model_config.hf_config.__dict__)

528 + quant_config = vllm_config.quant_config

529 + lora_config = vllm_config.lora_config

530 +

531 + self.config = config

532 + config.intermediate_size = config.ffn_hidden_size if hasattr(

533 + config, "ffn_hidden_size") else config.intermediate_size

534 + self.lora_config = lora_config

535 + self.quant_config = quant_config

536 +

537 + self.model = FlashModel(vllm_config=vllm_config,

538 + prefix=maybe_prefix(prefix, "model"))

539 +

540 + if get_pp_group().is_last_rank:

541 + self.lm_head = ParallelLMHead(config.vocab_size,

542 + config.hidden_size,

543 + quant_config=quant_config,

544 + prefix=maybe_prefix(

545 + prefix, "lm_head"))

546 + else:

547 + self.lm_head = PPMissingLayer()

548 +

549 + self.logits_processor = LogitsProcessor(config.vocab_size)

550 + self.make_empty_intermediate_tensors = (

551 + self.model.make_empty_intermediate_tensors)

552 +

553 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

554 + return self.model.get_input_embeddings(input_ids)

555 +

556 + def forward(

557 + self,

558 + input_ids: torch.Tensor,

559 + positions: torch.Tensor,

560 + intermediate_tensors: Optional[IntermediateTensors] = None,

561 + inputs_embeds: Optional[torch.Tensor] = None,

562 + ) -> Union[torch.Tensor, IntermediateTensors]:

563 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

564 + inputs_embeds)

565 + return hidden_states

566 +

567 + def compute_logits(

568 + self,

569 + hidden_states: torch.Tensor,

570 + ) -> Optional[torch.Tensor]:

571 + logits = self.logits_processor(self.lm_head, hidden_states)

572 + return logits

573 +

574 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

575 + # Params for weights, fp8 weight scales, fp8 activation scales

576 + # (param_name, weight_name, expert_id, shard_id)

577 + return FusedMoE.make_expert_params_mapping(

578 + ckpt_gate_proj_name="gate_proj",

579 + ckpt_down_proj_name="down_proj",

580 + ckpt_up_proj_name="up_proj",

581 + num_experts=self.config.n_routed_experts if hasattr(

582 + self.config, "n_routed_experts") else

583 + self.config.num_experts[0],

584 + )

585 +

586 + def load_weights(self, weights: Iterable[tuple[str,

587 + torch.Tensor]]) -> set[str]:

588 +

589 + stacked_params_mapping = [

590 + ("fused_qkv_a_proj", "q_a_proj", 0),

591 + ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),

592 + (".gate_up_proj", ".gate_proj", 0),

593 + (".gate_up_proj", ".up_proj", 1),

594 + ]

595 +

596 + expert_params_mapping = self.get_expert_mapping()

597 + loaded_params: set[str] = set()

598 +

599 + params_dict = dict(self.named_parameters())

600 + for name, loaded_weight in weights:

601 + if "rotary_emb.inv_freq" in name:

602 + continue

603 + for param_name, weight_name, shard_id in stacked_params_mapping:

604 + if weight_name not in name:

605 + continue

606 + if "mlp" in name and "mlps" not in name:

607 + continue

608 + name = name.replace(weight_name, param_name)

609 + # Skip loading extra bias for GPTQ models.

610 + if (name.endswith(".bias")

611 + or name.endswith("_bias")) and name not in params_dict:

612 + continue

613 + # Skip mtp

614 + if ".mtp." in name:

615 + continue

616 + if is_pp_missing_parameter(name, self):

617 + continue

618 + param = params_dict[name]

619 + weight_loader = param.weight_loader

620 + weight_loader(param, loaded_weight, shard_id)

621 + break

622 + else:

623 + is_expert_weight = False

624 + for mapping in expert_params_mapping:

625 + param_name, weight_name, expert_id, shard_id = mapping

626 + if weight_name not in name:

627 + continue

628 + is_expert_weight = True

629 + name_mapped = name.replace(weight_name, param_name)

630 + # Skip mtp

631 + if ".mtp." in name_mapped:

632 + continue

633 + if (name_mapped.endswith(".bias")

634 + or name_mapped.endswith("_bias")

635 + ) and name not in params_dict:

636 + continue

637 + if is_pp_missing_parameter(name, self):

638 + continue

639 + param = params_dict[name_mapped]

640 + weight_loader = param.weight_loader

641 + weight_loader = typing.cast(Callable[..., bool],

642 + param.weight_loader)

643 + success = weight_loader(param,

644 + loaded_weight,

645 + name_mapped,

646 + shard_id=shard_id,

647 + expert_id=expert_id,

648 + return_success=True)

649 + if success:

650 + name = name_mapped

651 + break

652 + else:

653 + if is_expert_weight:

654 + # We've checked that this is an expert weight

655 + # However it's not mapped locally to this rank

656 + # So we simply skip it

657 + continue

658 + # Skip loading extra bias for GPTQ models.

659 + if name.endswith(".bias") and name not in params_dict:

660 + continue

661 + # Skip loading kv_scale from ckpts towards new design.

662 + if name.endswith(".kv_scale") and name not in params_dict:

663 + continue

664 + # Skip mtp

665 + if ".mtp." in name:

666 + continue

667 + if name is None:

668 + continue

669 + if is_pp_missing_parameter(name, self):

670 + continue

671 + param = params_dict[name]

672 + weight_loader = getattr(param, "weight_loader",

673 + default_weight_loader)

674 + weight_loader(param, loaded_weight)

675 + loaded_params.add(name)

676 + for layer_id in range(self.config.num_hidden_layers):

677 + for i in range(2):

678 + if isinstance(self.model.layers[layer_id], PPMissingLayer):

679 + continue

680 + self_attn = self.model.layers[layer_id].self_attn[i]

681 + if hasattr(self.quant_config, "weight_block_size"

682 + ) and self_attn.kv_b_proj.weight.dtype in (

683 + torch.float8_e4m3fn,

684 + torch.float8_e4m3fnuz,

685 + ):

686 + weight_block_size = self.quant_config.weight_block_size

687 + if weight_block_size is not None:

688 + assert hasattr(self_attn.kv_b_proj, "weight_scale_inv")

689 + dtype = torch.get_default_dtype()

690 + w = block_dequant(self_attn.kv_b_proj.weight,

691 + self_attn.kv_b_proj.weight_scale_inv,

692 + weight_block_size).to(dtype)

693 + else:

694 + w = self_attn.kv_b_proj.weight

695 +

696 + w_kc, w_vc = w.unflatten(

697 + 0,

698 + (-1,

699 + self_attn.qk_nope_head_dim + self_attn.v_head_dim)).split(

700 + [self_attn.qk_nope_head_dim, self_attn.v_head_dim],

701 + dim=1)

702 + self_attn.w_kc = w_kc.transpose(1, 2).contiguous().transpose(

703 + 1, 2)

704 + self_attn.w_vc = w_vc.contiguous().transpose(1, 2)

705 + if self.config.mla_scale_q_lora:

706 + self_attn.q_a_layernorm.weight.data *= (

707 + self.config.hidden_size / self.config.q_lora_rank)**0.5

708 + if self.config.mla_scale_kv_lora:

709 + self_attn.kv_a_layernorm.weight.data *= (

710 + self.config.hidden_size /

711 + self.config.kv_lora_rank)**0.5

712 + return loaded_params

vllm/model_executor/models/longcat_flash_mtp.py new

+352

@@ -0,0 +1,352 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

3 +

4 +# Adapted from

5 +# https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/deepseek_mtp.py

6 +from collections.abc import Iterable

7 +from typing import Optional

8 +

9 +import torch

10 +import torch.nn as nn

11 +from transformers import PretrainedConfig

12 +

13 +from vllm.config import VllmConfig

14 +from vllm.model_executor.layers.layernorm import RMSNorm

15 +from vllm.model_executor.layers.linear import ReplicatedLinear

16 +from vllm.model_executor.layers.logits_processor import LogitsProcessor

17 +from vllm.model_executor.layers.quantization import QuantizationConfig

18 +from vllm.model_executor.layers.quantization.utils.int8_utils import (

19 + block_dequant)

20 +from vllm.model_executor.layers.vocab_parallel_embedding import (

21 + ParallelLMHead, VocabParallelEmbedding)

22 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader

23 +from vllm.model_executor.models.longcat_flash import FlashConfig

24 +from vllm.sequence import IntermediateTensors

25 +

26 +from .deepseek_v2 import DeepseekV2DecoderLayer

27 +from .interfaces import SupportsPP

28 +from .utils import maybe_prefix

29 +

30 +

31 +class LongCatMultiTokenPredictorLayer(nn.Module):

32 +

33 + def __init__(

34 + self,

35 + config: PretrainedConfig,

36 + prefix: str,

37 + vllm_config: VllmConfig,

38 + quant_config: Optional[QuantizationConfig] = None,

39 + ) -> None:

40 + super().__init__()

41 + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

42 + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

43 + self.eh_proj = ReplicatedLinear(2 * config.hidden_size,

44 + config.hidden_size,

45 + bias=False,

46 + quant_config=quant_config,

47 + prefix="eh_proj")

48 + self.mtp_block = DeepseekV2DecoderLayer(vllm_config, prefix)

49 + self.final_layernorm = RMSNorm(config.hidden_size,

50 + eps=config.rms_norm_eps)

51 +

52 + def forward(

53 + self,

54 + input_ids: torch.Tensor,

55 + positions: torch.Tensor,

56 + previous_hidden_states: torch.Tensor,

57 + inputs_embeds: Optional[torch.Tensor] = None,

58 + spec_step_index: int = 0,

59 + ) -> torch.Tensor:

60 + assert inputs_embeds is not None

61 + inputs_embeds = self.enorm(inputs_embeds)

62 + previous_hidden_states = self.hnorm(previous_hidden_states)

63 +

64 + hidden_states, _ = self.eh_proj(

65 + torch.cat([inputs_embeds, previous_hidden_states], dim=-1))

66 +

67 + hidden_states, residual = self.mtp_block(positions=positions,

68 + hidden_states=hidden_states,

69 + residual=None)

70 + hidden_states, _ = self.final_layernorm(hidden_states, residual)

71 + return hidden_states

72 +

73 +

74 +class LongCatMultiTokenPredictor(nn.Module):

75 +

76 + def __init__(self,

77 + *,

78 + vllm_config: VllmConfig,

79 + quant_config: Optional[QuantizationConfig] = None,

80 + prefix: str = ""):

81 + super().__init__()

82 + config = FlashConfig(**vllm_config.model_config.hf_config.__dict__)

83 + vllm_config.model_config.hf_config.intermediate_size \

84 + = config.intermediate_size

85 + self.mtp_start_layer_idx = config.num_hidden_layers * 2

86 + self.num_mtp_layers = 1

87 + self.layers = torch.nn.ModuleDict({

88 + str(idx):

89 + LongCatMultiTokenPredictorLayer(

90 + config,

91 + prefix=f"{prefix}.layers.{idx}",

92 + vllm_config=vllm_config,

93 + quant_config=quant_config,

94 + )

95 + for idx in range(self.mtp_start_layer_idx,

96 + self.mtp_start_layer_idx + self.num_mtp_layers)

97 + })

98 + self.embed_tokens = VocabParallelEmbedding(

99 + config.vocab_size,

100 + config.hidden_size,

101 + )

102 +

103 + def forward(

104 + self,

105 + input_ids: torch.Tensor,

106 + positions: torch.Tensor,

107 + previous_hidden_states: torch.Tensor,

108 + inputs_embeds: Optional[torch.Tensor] = None,

109 + spec_step_idx: int = 0,

110 + ) -> torch.Tensor:

111 + if inputs_embeds is None:

112 + inputs_embeds = self.embed_tokens(input_ids)

113 + current_step_idx = (spec_step_idx % self.num_mtp_layers)

114 + return self.layers[str(self.mtp_start_layer_idx + current_step_idx)](

115 + input_ids,

116 + positions,

117 + previous_hidden_states,

118 + inputs_embeds,

119 + current_step_idx,

120 + )

121 +

122 +

123 +class LongCatFlashMTP(nn.Module, SupportsPP):

124 +

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

126 + super().__init__()

127 + # LongCat MTP without MoE layers

128 + vllm_config.model_config.hf_config.n_routed_experts = None

129 + self.config = FlashConfig(

130 + **vllm_config.model_config.hf_config.__dict__)

131 + self.quant_config = None if "mtp" in getattr(

132 + self.config, "disable_quant_module",

133 + []) else vllm_config.quant_config

134 +

135 + self.model = LongCatMultiTokenPredictor(vllm_config=vllm_config,

136 + quant_config=self.quant_config,

137 + prefix=maybe_prefix(

138 + prefix, "model"))

139 + self.lm_head = ParallelLMHead(

140 + self.config.vocab_size,

141 + self.config.hidden_size,

142 + quant_config=self.quant_config,

143 + )

144 + self.logits_processor = LogitsProcessor(self.config.vocab_size)

145 +

146 + def forward(

147 + self,

148 + input_ids: torch.Tensor,

149 + positions: torch.Tensor,

150 + hidden_states: torch.Tensor,

151 + intermediate_tensors: Optional[IntermediateTensors] = None,

152 + inputs_embeds: Optional[torch.Tensor] = None,

153 + spec_step_idx: int = 0,

154 + ) -> torch.Tensor:

155 + hidden_states = self.model(input_ids, positions, hidden_states,

156 + inputs_embeds, spec_step_idx)

157 + return hidden_states

158 +

159 + def compute_logits(

160 + self,

161 + hidden_states: torch.Tensor,

162 + spec_step_idx: int = 0,

163 + ) -> Optional[torch.Tensor]:

164 +

165 + logits = self.logits_processor(self.lm_head, hidden_states)

166 + return logits

167 +

168 + def load_weights(self, weights: Iterable[tuple[str,

169 + torch.Tensor]]) -> set[str]:

170 + stacked_params_mapping = [

171 + ("gate_up_proj", "gate_proj", 0),

172 + ("gate_up_proj", "up_proj", 1),

173 + ("fused_qkv_a_proj", "q_a_proj", 0),

174 + ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),

175 + ]

176 +

177 + new_to_old_names_mapping = {

178 + "model.mtp.embed_tokens.weight":

179 + "model.layers.0.embed_tokens.weight",

180 + "model.mtp.layers.0.eh_proj.weight": "eh_proj.weight",

181 + "model.mtp.layers.0.eh_proj.weight_scale_inv":

182 + "eh_proj.weight_scale_inv",

183 + "model.mtp.layers.0.enorm.m.weight": "enorm.weight",

184 + "model.mtp.layers.0.hnorm.m.weight": "hnorm.weight",

185 + "model.mtp.layers.0.input_layernorm.weight":

186 + "model.layers.0.input_layernorm.weight",

187 + "model.mtp.layers.0.post_attention_layernorm.weight":

188 + "model.layers.0.post_attention_layernorm.weight",

189 + "model.mtp.layers.0.self_attn.kv_a_layernorm.weight":

190 + "model.layers.0.self_attn.kv_a_layernorm.weight",

191 + "model.mtp.layers.0.self_attn.kv_a_proj_with_mqa.weight":

192 + "model.layers.0.self_attn.kv_a_proj_with_mqa.weight",

193 + "model.mtp.layers.0.self_attn.kv_a_proj_with_mqa.weight_scale_inv":

194 + "model.layers.0.self_attn.kv_a_proj_with_mqa.weight_scale_inv",

195 + "model.mtp.layers.0.self_attn.kv_b_proj.weight":

196 + "model.layers.0.self_attn.kv_b_proj.weight",

197 + "model.mtp.layers.0.self_attn.kv_b_proj.weight_scale_inv":

198 + "model.layers.0.self_attn.kv_b_proj.weight_scale_inv",

199 + "model.mtp.layers.0.self_attn.o_proj.weight":

200 + "model.layers.0.self_attn.o_proj.weight",

201 + "model.mtp.layers.0.self_attn.o_proj.weight_scale_inv":

202 + "model.layers.0.self_attn.o_proj.weight_scale_inv",

203 + "model.mtp.layers.0.self_attn.q_a_layernorm.weight":

204 + "model.layers.0.self_attn.q_a_layernorm.weight",

205 + "model.mtp.layers.0.self_attn.q_a_proj.weight":

206 + "model.layers.0.self_attn.q_a_proj.weight",

207 + "model.mtp.layers.0.self_attn.q_a_proj.weight_scale_inv":

208 + "model.layers.0.self_attn.q_a_proj.weight_scale_inv",

209 + "model.mtp.layers.0.self_attn.q_b_proj.weight":

210 + "model.layers.0.self_attn.q_b_proj.weight",

211 + "model.mtp.layers.0.self_attn.q_b_proj.weight_scale_inv":

212 + "model.layers.0.self_attn.q_b_proj.weight_scale_inv",

213 + "model.mtp.layers.0.transformer_layer.mlp.down_proj.weight":

214 + "model.layers.0.mlp.down_proj.weight",

215 + "model.mtp.layers.0.transformer_layer.mlp.down_proj.weight_scale_inv":

216 + "model.layers.0.mlp.down_proj.weight_scale_inv",

217 + "model.mtp.layers.0.transformer_layer.mlp.gate_proj.weight":

218 + "model.layers.0.mlp.gate_proj.weight",

219 + "model.mtp.layers.0.transformer_layer.mlp.gate_proj.weight_scale_inv":

220 + "model.layers.0.mlp.gate_proj.weight_scale_inv",

221 + "model.mtp.layers.0.transformer_layer.mlp.up_proj.weight":

222 + "model.layers.0.mlp.up_proj.weight",

223 + "model.mtp.layers.0.transformer_layer.mlp.up_proj.weight_scale_inv":

224 + "model.layers.0.mlp.up_proj.weight_scale_inv",

225 + "model.mtp.norm.weight": "final_layernorm.weight",

226 + }

227 +

228 + params_dict = dict(self.named_parameters())

229 + loaded_params: set[str] = set()

230 + for name, loaded_weight in weights:

231 + if "rotary_emb.inv_freq" in name:

232 + continue

233 + spec_layer = self.get_spec_layer_idx_from_weight_name(

234 + self.config, name)

235 + if spec_layer is None:

236 + continue

237 + name = self._rewrite_spec_layer_name(spec_layer, name,

238 + new_to_old_names_mapping)

239 + for (param_name, weight_name, shard_id) in stacked_params_mapping:

240 + # Skip non-stacked layers and experts (experts handled below).

241 + if weight_name not in name:

242 + continue

243 + # We have mlp.experts[0].gate_proj in the checkpoint.

244 + # Since we handle the experts below in expert_params_mapping,

245 + # we need to skip here BEFORE we update the name, otherwise

246 + # name will be updated to mlp.experts[0].gate_up_proj, which

247 + # will then be updated below in expert_params_mapping

248 + # for mlp.experts[0].gate_gate_up_proj, which breaks load.

249 + if (("mlp.experts." in name) and name not in params_dict):

250 + continue

251 + name = name.replace(weight_name, param_name)

252 +

253 + # QKV fusion is optional, fall back to normal

254 + # weight loading if it's not enabled

255 + if ((param_name == "fused_qkv_a_proj")

256 + and name not in params_dict):

257 + continue

258 +

259 + # Skip loading extra bias for GPTQ models.

260 + if name.endswith(".bias") and name not in params_dict:

261 + continue

262 +

263 + param = params_dict[name]

264 + weight_loader = param.weight_loader

265 + weight_loader(param, loaded_weight, shard_id)

266 + break

267 + else:

268 + # Skip loading extra bias for GPTQ models.

269 + if name.endswith(".bias") and name not in params_dict:

270 + continue

271 +

272 + # According to DeepSeek-V3 Technical Report, MTP modules

273 + # shares embedding layer. We only load the first weights.

274 + if (spec_layer != self.model.mtp_start_layer_idx

275 + and ".layers" not in name):

276 + continue

277 +

278 + param = params_dict[name]

279 + weight_loader = getattr(param, "weight_loader",

280 + default_weight_loader)

281 + weight_loader(param, loaded_weight)

282 + loaded_params.add(name)

283 + spec_layer_id = self.config.num_hidden_layers * 2

284 + self_attn = self.model.layers[str(spec_layer_id)].mtp_block.self_attn

285 + if hasattr(

286 + self.quant_config,

287 + "weight_block_size") and self_attn.kv_b_proj.weight.dtype in (

288 + torch.float8_e4m3fn,

289 + torch.float8_e4m3fnuz,

290 + ):

291 + weight_block_size = self.quant_config.weight_block_size

292 + if weight_block_size is not None:

293 + dtype = torch.get_default_dtype()

294 + w = block_dequant(self_attn.kv_b_proj.weight,

295 + self_attn.kv_b_proj.weight_scale_inv,

296 + weight_block_size).to(dtype)

297 + else:

298 + w = self_attn.kv_b_proj.weight

299 + else:

300 + w = self_attn.kv_b_proj.weight

301 + w_kc, w_vc = w.unflatten(

302 + 0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim)).split(

303 + [self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1)

304 + self_attn.w_kc = w_kc.transpose(1, 2).contiguous().transpose(1, 2)

305 + self_attn.w_vc = w_vc.contiguous().transpose(1, 2)

306 + if self.config.mla_scale_q_lora:

307 + self_attn.q_a_layernorm.weight.data *= (

308 + self.config.hidden_size / self.config.q_lora_rank)**0.5

309 + if self.config.mla_scale_kv_lora:

310 + self_attn.kv_a_layernorm.weight.data *= (

311 + self.config.hidden_size / self.config.kv_lora_rank)**0.5

312 + return loaded_params

313 +

314 + def _rewrite_spec_layer_name(self, spec_layer: int, name: str,

315 + new_to_old_names_mapping: dict) -> str:

316 + """

317 + Rewrite the weight name to match the format of the original model.

318 + Add .mtp_block for modules in transformer layer block for spec layer

319 + and rename shared layer weights to be top level.

320 + """

321 + if name in new_to_old_names_mapping:

322 + name = new_to_old_names_mapping[name]

323 + spec_layer_weight_names = [

324 + "embed_tokens", "enorm", "hnorm", "eh_proj", "shared_head"

325 + ]

326 + if name.startswith("enorm") or name.startswith(

327 + "hnorm") or name.startswith("eh_proj") or name.startswith(

328 + "final_layernorm"):

329 + name = "model.layers." + str(spec_layer) + "." + name

330 + shared_weight_names = ["embed_tokens"]

331 + spec_layer_weight = False

332 + shared_weight = False

333 + for weight_name in spec_layer_weight_names:

334 + if weight_name in name:

335 + spec_layer_weight = True

336 + if weight_name in shared_weight_names:

337 + shared_weight = True

338 + break

339 + if not spec_layer_weight:

340 + # treat rest weights as weights for transformer layer block

341 + name = name.replace("model.layers.0.",

342 + f"model.layers.{spec_layer}.mtp_block.")

343 + elif shared_weight:

344 + # treat shared weights as top level weights

345 + name = name.replace("model.layers.0.", "model.")

346 + return name

347 +

348 + def get_spec_layer_idx_from_weight_name(self, config: PretrainedConfig,

349 + weight_name: str) -> Optional[int]:

350 + if "model.mtp" in weight_name:

351 + return config.num_hidden_layers * 2

352 + return None

vllm/model_executor/models/registry.py

@@ -109,6 +109,7 @@ _TEXT_GENERATION_MODELS = {

109 "Llama4ForCausalLM": ("llama4", "Llama4ForCausalLM"), # noqa: E501

110 # For decapoda-research/llama-*

111 "LLaMAForCausalLM": ("llama", "LlamaForCausalLM"),

112 + "LongcatFlashForCausalLM": ("longcat_flash", "LongcatFlashForCausalLM"),

113 "MambaForCausalLM": ("mamba", "MambaForCausalLM"),

114 "FalconMambaForCausalLM": ("mamba", "MambaForCausalLM"),

115 "FalconH1ForCausalLM":("falcon_h1", "FalconH1ForCausalLM"),

@@ -287,6 +288,7 @@ _SPECULATIVE_DECODING_MODELS = {

288 "EagleDeepSeekMTPModel": ("deepseek_eagle", "EagleDeepseekV3ForCausalLM"),

289 "DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),

290 "ErnieMTPModel": ("ernie_mtp", "ErnieMTP"),

291 + "LongCatFlashMTPModel": ("longcat_flash_mtp", "LongCatFlashMTP"),

292 "Glm4MoeMTPModel": ("glm4_moe_mtp", "Glm4MoeMTP"),

293 "MedusaModel": ("medusa", "Medusa"),

294 "Qwen3NextMTP": ("qwen3_next_mtp", "Qwen3NextMTP"),

vllm/model_executor/models/utils.py

+13 -5

@@ -691,14 +691,14 @@ def maybe_prefix(prefix: str, name: str) -> str:

691 return name if not prefix else f"{prefix}.{name}"

692

693

694 -def extract_layer_index(layer_name: str) -> int:

694 +def extract_layer_index(layer_name: str, num_attn_module: int = 1) -> int:

695 """

696 Extract the layer index from the module name.

697 Examples:

698 - "encoder.layers.0" -> 0

699 - "encoder.layers.1.self_attn" -> 1

700 - "2.self_attn" -> 2

701 - - "model.encoder.layers.0.sub.1" -> ValueError

701 + - "model.encoder.layers.0.sub.1" -> ValueError if num_attn_module == 1

702 """

703 subnames = layer_name.split(".")

704 int_vals: list[int] = []

@@ -707,9 +707,17 @@ def extract_layer_index(layer_name: str) -> int:

707 int_vals.append(int(subname))

708 except ValueError:

709 continue

710 - assert len(int_vals) == 1, (f"layer name {layer_name} should"

711 - " only contain one integer")

712 - return int_vals[0]

710 + if num_attn_module == 1 or "attn" not in layer_name:

711 + assert len(int_vals) == 1, (f"layer name {layer_name} should"

712 + " only contain one integer")

713 +

714 + return int_vals[0]

715 + else:

716 + assert len(int_vals) <= 2, (f"layer name {layer_name} should"

717 + " contain most two integers")

718 + layer_index = int_vals[0] * num_attn_module + int_vals[1] if len(

719 + int_vals) == 2 else int_vals[0]

720 + return layer_index

721

722

723 def cast_overflow_tensors(

vllm/v1/spec_decode/eagle.py

+7 -4

@@ -169,7 +169,6 @@ class EagleProposer:

169 target_hidden_states = self.model.combine_hidden_states(

170 target_hidden_states)

171 assert target_hidden_states.shape[-1] == self.hidden_size

172 -

172 # Shift the input ids by one token.

173 # E.g., [a1, b1, b2, c1, c2, c3] -> [b1, b2, c1, c2, c3, c3]

174 self.input_ids[:num_tokens - 1] = target_token_ids[1:]

@@ -223,7 +222,8 @@ class EagleProposer:

222 hidden_states=self.hidden_states[:num_input_tokens],

223 inputs_embeds=inputs_embeds,

224 )

226 - if self.method in ("deepseek_mtp", "ernie_mtp", "qwen3_next_mtp"):

225 + if self.method in ("deepseek_mtp", "ernie_mtp", "qwen3_next_mtp",

226 + "longcat_flash_mtp"):

227 last_hidden_states = ret_hidden_states

228 hidden_states = last_hidden_states

229 else:

@@ -237,7 +237,10 @@ class EagleProposer:

237 return draft_token_ids.view(-1, 1)

238

239 positions = target_positions[last_token_indices]

240 - hidden_states = hidden_states[last_token_indices]

240 + if self.method in ("deepseek_mtp", "ernie_mtp", "longcat_flash_mtp"):

241 + hidden_states = self.hidden_states[last_token_indices]

242 + else:

243 + hidden_states = hidden_states[last_token_indices]

244

245 if isinstance(attn_metadata, TreeAttentionMetadata):

246 # Draft using tree attention.

@@ -350,7 +353,7 @@ class EagleProposer:

353 inputs_embeds=inputs_embeds,

354 )

355 if self.method in ("deepseek_mtp", "ernie_mtp",

353 - "qwen3_next_mtp"):

356 + "qwen3_next_mtp", "longcat_flash_mtp"):

357 last_hidden_states = ret_hidden_states

358 hidden_states = ret_hidden_states

359 else:

vllm/v1/worker/gpu_model_runner.py

+3 -1

@@ -3840,9 +3840,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3840 target_layer_name)

3841 kv_caches[layer_name] = kv_caches[target_layer_name]

3842

3843 + num_attn_module = 2 \

3844 + if self.model_config.hf_config.model_type == "longcat_flash" else 1

3845 bind_kv_cache(kv_caches,

3846 self.compilation_config.static_forward_context,

3845 - self.kv_caches)

3847 + self.kv_caches, num_attn_module)

3848 return kv_caches

3849

3850 def maybe_add_kv_sharing_layers_to_kv_cache_groups(

vllm/v1/worker/utils.py

+3 -1

@@ -266,6 +266,7 @@ def bind_kv_cache(

266 kv_caches: dict[str, torch.Tensor],

267 forward_context: dict[str, "Attention"],

268 runner_kv_caches: list[torch.Tensor],

269 + num_attn_module: Optional[int] = 1,

270 ) -> None:

271 """

272 Bind the allocated KV cache to both ModelRunner and forward context so

@@ -289,7 +290,8 @@ def bind_kv_cache(

290 # Convert kv_caches dict to a list of tensors in the order of layer_index.

291 index2name = defaultdict(list)

292 for layer_name in kv_caches:

292 - index2name[extract_layer_index(layer_name)].append(layer_name)

293 + index2name[extract_layer_index(layer_name,

294 + num_attn_module)].append(layer_name)

295

296 for layer_index in sorted(index2name.keys()):

297 layer_names = index2name[layer_index]