Skip to content

[Model] LongCat-Flash: scale MLA norms while loading and drop the post-load sweep #60080

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
aoshen02 merged 3 commits into from
Oct 6, 2026
Merged

[Model] LongCat-Flash: scale MLA norms while loading and drop the post-load sweep #60080

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
acb7cf8
[Model] LongCat-Flash: scale MLA norms while loading and drop the pos…
aoshen02 Oct 5, 2026
0b3832e
Merge branch 'main' into wu/longcat-scale-norms-on-load
aoshen02 Oct 6, 2026
3f57702
Merge branch 'main' into wu/longcat-scale-norms-on-load
aoshen02 Oct 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
87 changes: 28 additions & 59 deletions vllm/model_executor/models/longcat_flash.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,13 +58,17 @@
)
from vllm.model_executor.layers.logits_processor import LogitsProcessor
from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.model_executor.layers.quantization.utils.int8_utils import block_dequant
from vllm.model_executor.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
)
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm.model_executor.model_loader.weight_utils import (
LoaderFunction,
composed_weight_loader,
default_weight_loader,
)
from vllm.model_executor.models.deepseek_v2 import DeepseekV2MLAAttention
from vllm.model_executor.utils import set_weight_attrs
from vllm.sequence import IntermediateTensors

from .interfaces import SupportsLoRA, SupportsPP
Expand All @@ -80,6 +84,23 @@
logger = init_logger(__name__)


def _scaled_weight_loader(scale: float) -> LoaderFunction:
return composed_weight_loader(default_weight_loader, lambda x: x * scale)


def scale_mla_lora_norms_on_load(self_attn: nn.Module, config) -> None:
"""Fold LongCat's MLA LoRA scaling into the q/kv_a_layernorm weight loaders."""
for norm, rank, enabled in (
(self_attn.q_a_layernorm, config.q_lora_rank, config.mla_scale_q_lora),
(self_attn.kv_a_layernorm, config.kv_lora_rank, config.mla_scale_kv_lora),
):
if enabled:
scale = (config.hidden_size / rank) ** 0.5
set_weight_attrs(
norm.weight, {"weight_loader": _scaled_weight_loader(scale)}
)


class FlashMLP(nn.Module):
"""Flash MLP layer."""

Expand Down Expand Up @@ -265,6 +286,8 @@ def __init__(
for i in range(2)
]
)
for self_attn in self.self_attn:
scale_mla_lora_norms_on_load(self_attn, config)
self.input_layernorm = nn.ModuleList(
[RMSNorm(config.hidden_size, eps=config.rms_norm_eps) for i in range(2)]
)
Expand Down Expand Up @@ -450,6 +473,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
for name, loaded_weight in weights:
if "rotary_emb.inv_freq" in name:
continue
# The MTP layers are loaded by LongCatFlashMTP.
if name.startswith("mtp."):
continue
for param_name, weight_name, shard_id in stacked_params_mapping:
if weight_name not in name:
continue
Expand All @@ -461,9 +487,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
name.endswith(".bias") or name.endswith("_bias")
) and name not in params_dict:
continue
# Skip mtp
if ".mtp." in name:
continue
if is_pp_missing_parameter(name, self):
continue
param = params_dict[name]
Expand All @@ -478,9 +501,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
continue
is_expert_weight = True
name_mapped = name.replace(weight_name, param_name)
# Skip mtp
if ".mtp." in name_mapped:
continue
if (
name_mapped.endswith(".bias") or name_mapped.endswith("_bias")
) and name not in params_dict:
Expand Down Expand Up @@ -515,9 +535,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
# Skip loading kv_scale from ckpts towards new design.
if name.endswith(".kv_scale") and name not in params_dict:
continue
# Skip mtp
if ".mtp." in name:
continue
if name is None:
continue
if is_pp_missing_parameter(name, self):
Expand All @@ -528,54 +545,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
)
weight_loader(param, loaded_weight)
loaded_params.add(name)
for layer_id in range(self.config.num_layers):
for i in range(2):
if isinstance(self.layers[layer_id], PPMissingLayer):
continue
self_attn = self.layers[layer_id].self_attn[i]
if (
self.quant_config is not None
and hasattr(self.quant_config, "weight_block_size")
and self_attn.kv_b_proj.weight.dtype
in (
torch.float8_e4m3fn,
torch.float8_e4m3fnuz,
)
):
weight_block_size = self.quant_config.weight_block_size
if weight_block_size is not None:
assert hasattr(self_attn.kv_b_proj, "weight_scale_inv")
dtype = torch.get_default_dtype()
w = block_dequant(
self_attn.kv_b_proj.weight,
self_attn.kv_b_proj.weight_scale_inv,
weight_block_size,
).to(dtype)
else:
w = self_attn.kv_b_proj.weight

w_kc, w_vc = w.unflatten(
0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim)
).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1)
self_attn.w_kc = w_kc.transpose(1, 2).contiguous().transpose(1, 2)
self_attn.w_vc = w_vc.contiguous().transpose(1, 2)
# Guard against compounding on incremental load_weights calls:
# the in-place ``*=`` would otherwise re-apply the MLA LoRA
# scaling to the layernorm weights on each pass.
if self.config.mla_scale_q_lora and not getattr(
self_attn, "_mla_q_lora_scaled", False
):
self_attn.q_a_layernorm.weight.data *= (
self.config.hidden_size / self.config.q_lora_rank
) ** 0.5
self_attn._mla_q_lora_scaled = True
if self.config.mla_scale_kv_lora and not getattr(
self_attn, "_mla_kv_lora_scaled", False
):
self_attn.kv_a_layernorm.weight.data *= (
self.config.hidden_size / self.config.kv_lora_rank
) ** 0.5
self_attn._mla_kv_lora_scaled = True
return loaded_params


Expand Down
47 changes: 2 additions & 45 deletions vllm/model_executor/models/longcat_flash_mtp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@
from vllm.model_executor.layers.linear import ReplicatedLinear
from vllm.model_executor.layers.logits_processor import LogitsProcessor
from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.model_executor.layers.quantization.utils.int8_utils import block_dequant
from vllm.model_executor.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
Expand All @@ -23,6 +22,7 @@
from vllm.sequence import IntermediateTensors

from .deepseek_v2 import DeepseekV2DecoderLayer
from .longcat_flash import scale_mla_lora_norms_on_load
from .utils import maybe_prefix


Expand All @@ -45,6 +45,7 @@ def __init__(
prefix="eh_proj",
)
self.mtp_block = DeepseekV2DecoderLayer(vllm_config, prefix)
scale_mla_lora_norms_on_load(self.mtp_block.self_attn, config)
self.final_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

def forward(
Expand Down Expand Up @@ -263,50 +264,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)
loaded_params.add(name)
spec_layer_id = self.model.mtp_start_layer_idx
self_attn = self.model.layers[str(spec_layer_id)].mtp_block.self_attn
if (
self.quant_config is not None
and hasattr(self.quant_config, "weight_block_size")
and self_attn.kv_b_proj.weight.dtype
in (
torch.float8_e4m3fn,
torch.float8_e4m3fnuz,
)
):
weight_block_size = self.quant_config.weight_block_size
if weight_block_size is not None:
dtype = torch.get_default_dtype()
w = block_dequant(
self_attn.kv_b_proj.weight,
self_attn.kv_b_proj.weight_scale_inv,
weight_block_size,
).to(dtype)
else:
w = self_attn.kv_b_proj.weight
else:
w = self_attn.kv_b_proj.weight
w_kc, w_vc = w.unflatten(
0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim)
).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1)
self_attn.w_kc = w_kc.transpose(1, 2).contiguous().transpose(1, 2)
self_attn.w_vc = w_vc.contiguous().transpose(1, 2)
# Guard against compounding on incremental load_weights calls (the
# in-place *= would otherwise double-apply the LoRA scaling).
if self.config.mla_scale_q_lora and not getattr(
self_attn, "_mla_q_lora_scaled", False
):
self_attn.q_a_layernorm.weight.data *= (
self.config.hidden_size / self.config.q_lora_rank
) ** 0.5
self_attn._mla_q_lora_scaled = True
if self.config.mla_scale_kv_lora and not getattr(
self_attn, "_mla_kv_lora_scaled", False
):
self_attn.kv_a_layernorm.weight.data *= (
self.config.hidden_size / self.config.kv_lora_rank
) ** 0.5
self_attn._mla_kv_lora_scaled = True
return loaded_params

def _rewrite_spec_layer_name(
Expand Down
Loading