# -*- coding: utf-8 -*- # Copyright (c) 2025 Meituan # This code is licensed under the MIT License, for details, see the ./LICENSE file. from typing import Any, Optional, Tuple, Dict, List import math import torch import torch.nn.functional as F from torch import nn from transformers.cache_utils import Cache, DynamicCache from transformers.conversion_mapping import register_checkpoint_conversion_mapping from transformers.core_model_loading import ( Concatenate, ConversionOps, MergeModulelist, WeightConverter, WeightRenaming, ) from transformers.masking_utils import create_causal_mask from transformers.modeling_outputs import BaseModelOutputWithPast from transformers.processing_utils import Unpack from transformers.utils import auto_docstring, logging from transformers.models.longcat_flash.modeling_longcat_flash import ( LongcatFlashForCausalLM, LongcatFlashModel, LongcatFlashRMSNorm, LongcatFlashRotaryEmbedding, LongcatFlashDecoderLayer, LongcatFlashMLA, LongcatFlashMLP, LongcatFlashPreTrainedModel, apply_rotary_pos_emb_interleave, ) from .configuration_longcat_flash_sparse import LongcatFlashSparseConfig logger = logging.get_logger(__name__) @auto_docstring class LongcatFlashSparsePreTrainedModel(LongcatFlashPreTrainedModel): pass class NgramCache(DynamicCache): """Dynamic MLA cache plus N-gram and LSA-indexer history.""" def __init__(self, config=None): # Keep the same lazy DynamicCache behavior as the validated Heretic port. super().__init__() self.ngram_context = None self.max_context_len = config.emb_neighbor_num - 1 # Owner attention layer id -> post-RMSNorm/post-RoPE indexer K states, # shape [batch, seq, index_head_dim]. Reuse/CLI layers intentionally have # no independent indexer cache. self.indexer_key_cache: Dict[int, torch.Tensor] = {} def update_ngram_context(self, new_tokens: torch.Tensor) -> None: if self.ngram_context is None: self.ngram_context = new_tokens.clone() else: self.ngram_context = torch.cat( [self.ngram_context, new_tokens], dim=-1 ) if self.ngram_context.size(-1) > self.max_context_len: self.ngram_context = self.ngram_context[..., -self.max_context_len :] def update_indexer_keys( self, layer_idx: int, new_keys: torch.Tensor ) -> tuple[int, torch.Tensor]: """Append indexer K states and return (old_length, full_history).""" previous = self.indexer_key_cache.get(layer_idx) if previous is None: old_length = 0 full = new_keys else: old_length = previous.shape[1] full = torch.cat([previous, new_keys], dim=1) self.indexer_key_cache[layer_idx] = full return old_length, full def reorder_cache(self, beam_idx: torch.LongTensor) -> "Cache": super().reorder_cache(beam_idx) if self.ngram_context is not None: self.ngram_context = self.ngram_context.index_select( 0, beam_idx.to(self.ngram_context.device) ) for layer_idx, value in list(self.indexer_key_cache.items()): self.indexer_key_cache[layer_idx] = value.index_select( 0, beam_idx.to(value.device) ) return self def crop(self, max_length: int): super().crop(max_length) if max_length < 0: # DynamicCache has already resolved the negative crop relative to its # own length. Mirror the resulting length from the first MLA layer. target = self.get_seq_length() else: target = max_length for layer_idx, value in list(self.indexer_key_cache.items()): if value.shape[1] > target: self.indexer_key_cache[layer_idx] = value[:, :target, :] def batch_repeat_interleave(self, repeats: int): super().batch_repeat_interleave(repeats) if self.ngram_context is not None: self.ngram_context = self.ngram_context.repeat_interleave( repeats, dim=0 ) for layer_idx, value in list(self.indexer_key_cache.items()): self.indexer_key_cache[layer_idx] = value.repeat_interleave( repeats, dim=0 ) def batch_select_indices(self, indices: torch.Tensor): super().batch_select_indices(indices) if self.ngram_context is not None: self.ngram_context = self.ngram_context[indices.to(self.ngram_context.device)] for layer_idx, value in list(self.indexer_key_cache.items()): self.indexer_key_cache[layer_idx] = value[ indices.to(value.device) ] class NgramEmbedding(nn.Module): """ Computes embeddings enriched with N-gram features without maintaining internal state. """ def __init__(self, config, base_embeddings): super().__init__() self.config = config self.word_embeddings = base_embeddings self.m = config.ngram_vocab_size_ratio * config.vocab_size self.k = config.emb_split_num self.n = config.emb_neighbor_num self._init_ngram_embeddings() self._vocab_mods_cache = None def _init_ngram_embeddings(self) -> None: """Initialize N-gram embedding and projection layers.""" num_embedders = self.k * (self.n - 1) emb_dim = self.config.hidden_size // num_embedders embedders = [] post_projs = [] for i in range(num_embedders): vocab_size = int(self.m + i * 2 + 1) emb = nn.Embedding(vocab_size, emb_dim, padding_idx=self.config.pad_token_id) proj = nn.Linear(emb_dim, self.config.hidden_size, bias=False) embedders.append(emb) post_projs.append(proj) self.embedders = nn.ModuleList(embedders) self.post_projs = nn.ModuleList(post_projs) def _shift_right_ignore_eos(self, tensor: torch.Tensor, n: int, eos_token_id: int = 2) -> torch.Tensor: """Shift tensor right by n positions, resetting at EOS tokens.""" batch_size, seq_len = tensor.shape result = torch.zeros_like(tensor) eos_mask = (tensor == eos_token_id) for i in range(batch_size): eos_positions = eos_mask[i].nonzero(as_tuple=True)[0] prev_idx = 0 for eos_idx in eos_positions: end_idx = eos_idx.item() + 1 if end_idx - prev_idx > n: result[i, prev_idx+n:end_idx] = tensor[i, prev_idx:end_idx-n] prev_idx = end_idx if prev_idx < seq_len and seq_len - prev_idx > n: result[i, prev_idx+n:seq_len] = tensor[i, prev_idx:seq_len-n] return result def _precompute_vocab_mods(self) -> Dict[Tuple[int, int], List[int]]: """Precompute modular arithmetic values for vocabulary.""" if self._vocab_mods_cache is not None: return self._vocab_mods_cache vocab_mods = {} vocab_size = self.config.vocab_size for i in range(2, self.n + 1): for j in range(self.k): index = (i - 2) * self.k + j emb_vocab_dim = int(self.m + index * 2 + 1) mods = [] power_mod = 1 for _ in range(i - 1): power_mod = (power_mod * vocab_size) % emb_vocab_dim mods.append(power_mod) vocab_mods[(i, j)] = mods self._vocab_mods_cache = vocab_mods return vocab_mods def _get_ngram_ids( self, input_ids: torch.Tensor, shifted_ids: Dict[int, torch.Tensor], vocab_mods: List[int], ngram: int ) -> torch.Tensor: """Compute N-gram hash IDs using polynomial rolling hash.""" ngram_ids = input_ids.clone() for k in range(2, ngram + 1): ngram_ids = ngram_ids + shifted_ids[k] * vocab_mods[k - 2] return ngram_ids def forward( self, input_ids: torch.Tensor, ngram_context: Optional[torch.Tensor] = None ) -> torch.Tensor: """ Stateless forward pass. Args: input_ids: Current input token IDs of shape (batch_size, seq_len) ngram_context: Optional historical context of shape (batch_size, context_len) Returns: Embedding tensor of shape (batch_size, seq_len, hidden_size) """ seq_len = input_ids.size(-1) # Determine complete context if ngram_context is not None: context = torch.cat([ngram_context[..., -(self.n-1):], input_ids], dim=-1) else: context = input_ids # Base word embeddings device = self.word_embeddings.weight.device x = self.word_embeddings(input_ids.to(device)).clone() # Precompute modular values vocab_mods = self._precompute_vocab_mods() # Compute shifted IDs shifted_ids = {} for i in range(2, self.n + 1): shifted_ids[i] = self._shift_right_ignore_eos( context, i - 1, eos_token_id=self.config.eos_token_id ) # Add N-gram embeddings for i in range(2, self.n + 1): for j in range(self.k): index = (i - 2) * self.k + j emb_vocab_dim = int(self.m + index * 2 + 1) ngram_ids = self._get_ngram_ids(context, shifted_ids, vocab_mods[(i, j)], ngram=i) new_ids = (ngram_ids % emb_vocab_dim)[..., -seq_len:] embedder_device = self.embedders[index].weight.device x_ngram = self.embedders[index](new_ids.to(embedder_device)) proj_device = self.post_projs[index].weight.device x_proj = self.post_projs[index](x_ngram.to(proj_device)) x = x + x_proj.to(x.device) # Normalize x = x / (1 + self.k * (self.n - 1)) return x class LongcatPadZeroExperts(ConversionOps): """Expand 256 checkpoint experts to the 384 runtime gate/up slots.""" @torch.no_grad() def convert( self, input_dict: dict[str, torch.Tensor | list[torch.Tensor]], source_patterns: list[str], target_patterns: list[str], config, **kwargs, ) -> dict[str, torch.Tensor]: if len(input_dict) != 1: raise ValueError("Expected exactly one LongCat expert tensor") key, value = next(iter(input_dict.items())) tensor = value[0] if isinstance(value, list) else value routed = int(config.n_routed_experts) zeros = int(config.zero_expert_num or 0) if tensor.shape[0] != routed: raise ValueError( f"LongCat routed expert mismatch: {tensor.shape[0]} != {routed}" ) if zeros: tensor = torch.cat( [tensor, tensor.new_zeros((zeros, *tensor.shape[1:]))], dim=0 ) return {key: tensor} @property def reverse_op(self) -> ConversionOps: return LongcatStripZeroExperts() class LongcatStripZeroExperts(ConversionOps): """Restore the compact 256-expert checkpoint layout on save.""" @torch.no_grad() def convert( self, input_dict: dict[str, torch.Tensor | list[torch.Tensor]], source_patterns: list[str], target_patterns: list[str], config, **kwargs, ) -> dict[str, torch.Tensor]: if len(input_dict) != 1: raise ValueError("Expected exactly one fused LongCat expert tensor") key, value = next(iter(input_dict.items())) tensor = value[0] if isinstance(value, list) else value routed = int(config.n_routed_experts) total = routed + int(config.zero_expert_num or 0) if tensor.shape[0] not in (routed, total): raise ValueError( f"LongCat runtime expert mismatch: {tensor.shape[0]} not in " f"({routed}, {total})" ) return {key: tensor[:routed].contiguous()} @property def reverse_op(self) -> ConversionOps: return LongcatPadZeroExperts() def _longcat_sparse_checkpoint_conversions() -> list[Any]: """Checkpoint <-> runtime mappings required by the standalone HF model.""" return [ WeightRenaming( source_patterns=r"oe_embed_tokens(\d+)\.weight", target_patterns=r"ngram_embeddings.embedders.\1.weight", ), WeightRenaming( source_patterns=r"oe_embed_proj(\d+)\.weight", target_patterns=r"ngram_embeddings.post_projs.\1.weight", ), WeightConverter( source_patterns=[ "mlp.experts.*.gate_proj.weight", "mlp.experts.*.up_proj.weight", ], target_patterns="mlp.experts.gate_up_proj", operations=[ MergeModulelist(dim=0), Concatenate(dim=1), LongcatPadZeroExperts(), ], ), WeightConverter( source_patterns="mlp.experts.*.down_proj.weight", target_patterns="mlp.experts.down_proj", operations=[MergeModulelist(dim=0)], ), ] def _register_longcat_sparse_checkpoint_conversions() -> None: # AutoModel imports this module before loading checkpoint shards, so module # import is early enough for Transformers' conversion cache. for class_name in ( "LongcatFlashSparseModel", "LongcatFlashSparseForCausalLM", ): register_checkpoint_conversion_mapping( class_name, _longcat_sparse_checkpoint_conversions(), overwrite=True, ) class LongcatSparseIndexer(nn.Module): """Pure-PyTorch reference implementation of the Lite-Sparse LSA indexer. It mirrors the trained SI/CLI semantics, not SGLang's FP8/DeepGEMM serving optimizations. Sink and local-window positions are forced with +inf *inside* the fixed ``index_topk`` result width, matching the Lite-Sparse SGLang path. """ def __init__(self, config: LongcatFlashSparseConfig, layer_idx: int): super().__init__() if config.q_lora_rank is None: raise ValueError("Lite-Sparse LSA requires q_lora_rank") if getattr(config, "index_k_norm_type", "rms") != "rms": raise NotImplementedError( "This Lite-Sparse reference path currently supports RMS index K norm only" ) if not getattr(config, "indexer_rope_interleave", True): raise NotImplementedError( "This Lite-Sparse checkpoint requires interleaved indexer RoPE" ) self.layer_idx = layer_idx self.n_heads = int(config.index_n_heads) self.head_dim = int(config.index_head_dim) self.rope_head_dim = int(config.qk_rope_head_dim) self.topk = int(config.index_topk) self.num_init_tokens = int(config.index_init_tokens) self.num_local_tokens = int(config.index_local_tokens) self.wq_b = nn.Linear( config.q_lora_rank, self.n_heads * self.head_dim, bias=False, ) self.wk = nn.Linear(config.hidden_size, self.head_dim, bias=False) # The checkpoint stores this trained head-gating projection in FP32, # and the reference SGLang implementation explicitly keeps it FP32. # Pin the parameter dtype so a global BF16 model-construction context # cannot silently downcast it during from_pretrained(). self.weights_proj = nn.Linear( config.hidden_size, self.n_heads, bias=False, dtype=torch.float32, ) # SGLang constructs Indexer RMSNorm(head_dim) without an explicit eps; # its RMSNorm default is 1e-6. Do not use the model-wide 1e-5 here. self.k_norm = LongcatFlashRMSNorm(self.head_dim, eps=1e-6) # Bounded scratch sizes. These affect speed only, never the selected set. self.query_chunk_size = int( getattr(config, "lsa_index_query_chunk_size", 8) ) self.key_chunk_size = int( getattr(config, "lsa_index_key_chunk_size", 4096) ) def _rope_q( self, query: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor], ) -> torch.Tensor: # [B,Q,H,D] -> [B,H,Q,D] query = query.transpose(1, 2) q_rope, q_rest = torch.split( query, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1, ) cos, sin = position_embeddings q_rope, _ = apply_rotary_pos_emb_interleave( q_rope, q_rope, cos, sin ) return torch.cat([q_rope, q_rest], dim=-1).transpose(1, 2) def project_key( self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor], ) -> torch.Tensor: key = self.k_norm(self.wk(hidden_states)) k_rope, k_rest = torch.split( key, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1, ) # apply_rotary_pos_emb_interleave expects [B,H,Q,D]. Index K is MQA. k_rope = k_rope.unsqueeze(1) cos, sin = position_embeddings _, k_rope = apply_rotary_pos_emb_interleave( k_rope, k_rope, cos, sin ) return torch.cat([k_rope.squeeze(1), k_rest], dim=-1) @staticmethod def _normalize_attention_mask( attention_mask: Optional[torch.Tensor], *, batch_size: int, total_kv_len: int, device: torch.device, ) -> torch.Tensor: if attention_mask is None: return torch.ones( (batch_size, total_kv_len), dtype=torch.bool, device=device ) if attention_mask.dim() != 2: raise ValueError( "LongCat LSA sparse selection requires a 2D token attention_mask " "when padding is present." ) if attention_mask.shape[0] != batch_size: raise ValueError("attention_mask batch dimension mismatch") if attention_mask.shape[1] != total_kv_len: raise ValueError( "attention_mask must cover the complete cached KV history: " f"mask={attention_mask.shape[1]}, kv={total_kv_len}" ) return attention_mask.to(device=device, dtype=torch.bool) def select( self, hidden_states: torch.Tensor, q_lora: torch.Tensor, index_keys: torch.Tensor, *, past_len: int, position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor], ) -> torch.Tensor: """Return absolute KV indices [B,Q,index_topk], padded with -1.""" batch_size, query_len, _ = hidden_states.shape total_kv_len = index_keys.shape[1] if total_kv_len <= self.topk: # Exact all-candidate identity: every causal candidate is selected. # This branch is mostly useful for direct indexer tests; the attention # module itself bypasses scoring entirely at this boundary. out = torch.full( (batch_size, query_len, self.topk), -1, dtype=torch.long, device=hidden_states.device, ) mask = self._normalize_attention_mask( attention_mask, batch_size=batch_size, total_kv_len=total_kv_len, device=hidden_states.device, ) for q in range(query_len): end = min(past_len + q + 1, total_kv_len) for b in range(batch_size): idx = torch.nonzero(mask[b, :end], as_tuple=False).flatten() out[b, q, : idx.numel()] = idx return out query = self.wq_b(q_lora).view( batch_size, query_len, self.n_heads, self.head_dim ) query = self._rope_q(query, position_embeddings) # SGLang explicitly accumulates weights_proj in FP32. Preserve that # detail rather than computing BF16 output and merely casting afterwards. gates = F.linear( hidden_states.float(), self.weights_proj.weight.float(), bias=None ) gates = gates * (self.n_heads ** -0.5) # The dedicated SGLang kernels quantize Q/K before scoring. Here we retain # BF16 projections and accumulate dot products/scores in FP32. Positive # global scale factors match the reference and do not alter top-K order. query = query.float() index_keys = index_keys.float() valid_tokens = self._normalize_attention_mask( attention_mask, batch_size=batch_size, total_kv_len=total_kv_len, device=hidden_states.device, ) valid_rank = valid_tokens.to(torch.int64).cumsum(dim=-1) result = torch.full( (batch_size, query_len, self.topk), -1, dtype=torch.long, device=hidden_states.device, ) q_chunk_size = max(1, self.query_chunk_size) k_chunk_size = max(1, self.key_chunk_size) for qs in range(0, query_len, q_chunk_size): qe = min(qs + q_chunk_size, query_len) q_chunk = query[:, qs:qe] gate_chunk = gates[:, qs:qe] q_positions = torch.arange( qs, qe, device=hidden_states.device, dtype=torch.long ) + past_len q_positions = q_positions.clamp_max(total_kv_len - 1) # Number of valid sequence tokens available to each query. Shape B,C. valid_to_q = valid_rank[:, q_positions] query_is_valid = valid_tokens[:, q_positions] best_scores = None best_indices = None for ks in range(0, total_kv_len, k_chunk_size): ke = min(ks + k_chunk_size, total_kv_len) k_chunk = index_keys[:, ks:ke] # [B,C,H,K]. ReLU is part of the trained indexer definition. dots = torch.einsum("bchd,bkd->bchk", q_chunk, k_chunk) dots.mul_(self.head_dim ** -0.5).relu_() scores = ( dots * gate_chunk.float().unsqueeze(-1) ).sum(dim=2) del dots key_positions = torch.arange( ks, ke, device=hidden_states.device, dtype=torch.long ) base_valid = valid_tokens[:, None, ks:ke] causal = key_positions[None, None, :] <= q_positions[None, :, None] valid = base_valid & causal & query_is_valid[:, :, None] key_rank = valid_rank[:, None, ks:ke] sink = valid & (key_rank <= self.num_init_tokens) local_cutoff = ( valid_to_q[:, :, None] - self.num_local_tokens ) local = valid & (key_rank > local_cutoff) scores.masked_fill_(~valid, float("-inf")) scores.masked_fill_(sink | local, float("inf")) take = min(self.topk, ke - ks) chunk_scores, chunk_local_idx = torch.topk( scores, k=take, dim=-1, sorted=False ) chunk_indices = chunk_local_idx.to(torch.long) + ks if best_scores is None: if take < self.topk: pad_shape = ( batch_size, qe - qs, self.topk - take, ) chunk_scores = torch.cat( [ chunk_scores, chunk_scores.new_full( pad_shape, float("-inf") ), ], dim=-1, ) chunk_indices = torch.cat( [ chunk_indices, chunk_indices.new_full(pad_shape, -1), ], dim=-1, ) best_scores = chunk_scores best_indices = chunk_indices else: merged_scores = torch.cat( [best_scores, chunk_scores], dim=-1 ) merged_indices = torch.cat( [best_indices, chunk_indices], dim=-1 ) best_scores, pick = torch.topk( merged_scores, k=self.topk, dim=-1, sorted=False, ) best_indices = torch.gather( merged_indices, -1, pick ) assert best_scores is not None and best_indices is not None best_indices = best_indices.masked_fill( best_scores == float("-inf"), -1 ) result[:, qs:qe] = best_indices return result class LongcatFlashSparseMLA(LongcatFlashMLA): """MLA with LongCat LSA selection and CLI top-K reuse.""" def __init__( self, config: LongcatFlashSparseConfig, layer_idx: int, *, indexer_owner: bool, ): super().__init__(config=config, layer_idx=layer_idx) self.indexer_owner = indexer_owner self.index_topk = int(config.index_topk) self.sparse_query_chunk_size = int( getattr(config, "lsa_attention_query_chunk_size", 4) ) # Diagnostic-only state, useful for release smoke tests. It is not part # of the checkpoint and has no effect on model outputs. self.last_lsa_mode: Optional[str] = None self.last_lsa_valid_topk_range: Optional[tuple[int, int]] = None if indexer_owner: self.indexer = LongcatSparseIndexer(config, layer_idx) else: self.indexer = None @staticmethod def _gather_selected( values: torch.Tensor, indices: torch.Tensor, ) -> torch.Tensor: """Gather [B,K,D] with [B,Q,T] -> [B,Q,T,D].""" batch_size, _, dim = values.shape _, query_len, topk = indices.shape safe = indices.clamp_min(0) expanded = values[:, None, :, :].expand( batch_size, query_len, -1, dim ) return torch.gather( expanded, 2, safe[..., None].expand(batch_size, query_len, topk, dim), ) def _forward_sparse_mla( self, hidden_states: torch.Tensor, q_lora: torch.Tensor, topk_indices: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor], past_key_values: Optional[Cache], ) -> torch.Tensor: batch_size, query_len, _ = hidden_states.shape query_shape = (batch_size, query_len, -1, self.qk_head_dim) q_states = self.q_b_proj(q_lora) q_states = q_states.view(query_shape).transpose(1, 2) q_pass, q_rot = torch.split( q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1, ) compressed_kv = self.kv_a_proj_with_mqa(hidden_states) kv_pass, k_rot = torch.split( compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1, ) k_pass = self.kv_a_layernorm(kv_pass) q_pass = q_pass * self.mla_scale_q_lora q_rot = q_rot * self.mla_scale_q_lora k_pass = k_pass * self.mla_scale_kv_lora k_pass = k_pass.view( batch_size, 1, query_len, self.kv_lora_rank ) k_rot = k_rot.view( batch_size, 1, query_len, self.qk_rope_head_dim ) cos, sin = position_embeddings q_rot, k_rot = apply_rotary_pos_emb_interleave( q_rot, k_rot, cos, sin ) if past_key_values is not None: k_pass, k_rot = past_key_values.update( k_pass, k_rot, self.layer_idx ) query_states = torch.cat((q_pass, q_rot), dim=-1) k_pass_base = k_pass.squeeze(1) k_rot_base = k_rot.squeeze(1) outputs = [] q_chunk_size = max(1, self.sparse_query_chunk_size) for qs in range(0, query_len, q_chunk_size): qe = min(qs + q_chunk_size, query_len) idx = topk_indices[:, qs:qe] invalid = idx < 0 selected_k_pass = self._gather_selected(k_pass_base, idx) selected_k_rot = self._gather_selected(k_rot_base, idx) # Expand only the selected latent K/V. This avoids materializing # dense full-history K/V at long context. expanded = self.kv_b_proj(selected_k_pass) expanded = expanded.view( batch_size, qe - qs, self.index_topk, self.num_heads, self.qk_nope_head_dim + self.v_head_dim, ) k_nope, value_states = torch.split( expanded, [self.qk_nope_head_dim, self.v_head_dim], dim=-1, ) # B,C,T,H,D -> B,C,H,T,D k_nope = k_nope.permute(0, 1, 3, 2, 4) value_states = value_states.permute(0, 1, 3, 2, 4) k_rope = selected_k_rot[:, :, None, :, :].expand( -1, -1, self.num_heads, -1, -1 ) key_states = torch.cat((k_nope, k_rope), dim=-1) q = query_states[:, :, qs:qe, :].permute(0, 2, 1, 3) attn_weights = torch.einsum( "bchd,bchtd->bcht", q, key_states ) * self.scaling attn_weights.masked_fill_( invalid[:, :, None, :], float("-inf") ) query_has_any_key = (~invalid).any(dim=-1) # Padding-only query rows would otherwise softmax an all--inf row. if not bool(query_has_any_key.all()): attn_weights = attn_weights.masked_fill( (~query_has_any_key)[:, :, None, None], 0.0 ) attn_weights = F.softmax( attn_weights, dim=-1, dtype=torch.float32 ).to(q.dtype) if self.training and self.attention_dropout: attn_weights = F.dropout( attn_weights, p=self.attention_dropout, training=True, ) attn_output = torch.einsum( "bcht,bchtd->bchd", attn_weights, value_states ) if not bool(query_has_any_key.all()): attn_output = attn_output.masked_fill( (~query_has_any_key)[:, :, None, None], 0.0 ) outputs.append( attn_output.reshape(batch_size, qe - qs, -1) ) attn_output = torch.cat(outputs, dim=1).contiguous() return self.o_proj(attn_output) def forward( self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor], past_key_values: Optional[Cache] = None, *, sequence_attention_mask: Optional[torch.Tensor] = None, topk_indices: Optional[torch.Tensor] = None, **kwargs, ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: batch_size, query_len, _ = hidden_states.shape if self.indexer_owner: if not isinstance(past_key_values, NgramCache) and past_key_values is not None: raise TypeError( "LongCat Lite-Sparse generation requires NgramCache so the " "LSA indexer K history is preserved." ) assert self.indexer is not None new_index_keys = self.indexer.project_key( hidden_states, position_embeddings ) if isinstance(past_key_values, NgramCache): past_len, full_index_keys = past_key_values.update_indexer_keys( self.layer_idx, new_index_keys ) else: past_len, full_index_keys = 0, new_index_keys total_kv_len = full_index_keys.shape[1] if total_kv_len <= self.index_topk: self.last_lsa_mode = "full-owner" self.last_lsa_valid_topk_range = None # Exact reference fast path: indexer scores cannot exclude any # causal token. We still stored K above for a later boundary crossing. out, weights = super().forward( hidden_states=hidden_states, position_embeddings=position_embeddings, attention_mask=attention_mask, past_key_values=past_key_values, **kwargs, ) return out, weights, None q_lora = self.q_a_layernorm(self.q_a_proj(hidden_states)) topk_indices = self.indexer.select( hidden_states, q_lora, full_index_keys, past_len=past_len, position_embeddings=position_embeddings, attention_mask=sequence_attention_mask, ) valid_counts = (topk_indices >= 0).sum(dim=-1) self.last_lsa_mode = "sparse-owner" self.last_lsa_valid_topk_range = ( int(valid_counts.min().item()), int(valid_counts.max().item()), ) out = self._forward_sparse_mla( hidden_states, q_lora, topk_indices, position_embeddings, past_key_values, ) return out, None, topk_indices # CLI reuse layer. No indexer weights exist here; the first attention # sublayer in the logical layer supplies the selection. if topk_indices is None: # If the owner used the <=index_topk identity fast path, reuse is # also mathematically full attention. total_kv_len = query_len if past_key_values is not None: total_kv_len += past_key_values.get_seq_length(self.layer_idx) if total_kv_len > self.index_topk: raise RuntimeError( "Lite-Sparse CLI reuse layer received no owner top-K above " f"index_topk={self.index_topk}. Refusing dense fallback." ) self.last_lsa_mode = "full-reuse" self.last_lsa_valid_topk_range = None out, weights = super().forward( hidden_states=hidden_states, position_embeddings=position_embeddings, attention_mask=attention_mask, past_key_values=past_key_values, **kwargs, ) return out, weights, None self.last_lsa_mode = "sparse-reuse" valid_counts = (topk_indices >= 0).sum(dim=-1) self.last_lsa_valid_topk_range = ( int(valid_counts.min().item()), int(valid_counts.max().item()), ) q_lora = self.q_a_layernorm(self.q_a_proj(hidden_states)) out = self._forward_sparse_mla( hidden_states, q_lora, topk_indices, position_embeddings, past_key_values, ) return out, None, topk_indices class LongcatFlashSparseDecoderLayer(LongcatFlashDecoderLayer): """LongCat logical layer with two physical attention layers and CLI=2.""" def __init__(self, config: LongcatFlashSparseConfig, layer_idx: int): if int(config.cli_factor) != 2: raise NotImplementedError( "This release implementation is validated for Lite-Sparse cli_factor=2" ) super().__init__(config, layer_idx) self.self_attn = nn.ModuleList( [ LongcatFlashSparseMLA( config, layer_idx=layer_idx * 2, indexer_owner=True, ), LongcatFlashSparseMLA( config, layer_idx=layer_idx * 2 + 1, indexer_owner=False, ), ] ) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, use_cache: Optional[bool] = False, position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, sequence_attention_mask: Optional[torch.Tensor] = None, **kwargs, ) -> torch.Tensor: residual = hidden_states hidden_states = self.input_layernorm[0](hidden_states) hidden_states, _, topk_indices = self.self_attn[0]( hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, position_embeddings=position_embeddings, sequence_attention_mask=sequence_attention_mask, **kwargs, ) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.post_attention_layernorm[0](hidden_states) shortcut_mlp_output = self.mlp(hidden_states) hidden_states = self.mlps[0](hidden_states) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.input_layernorm[1](hidden_states) hidden_states, _, _ = self.self_attn[1]( hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, position_embeddings=position_embeddings, sequence_attention_mask=sequence_attention_mask, topk_indices=topk_indices, **kwargs, ) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.post_attention_layernorm[1](hidden_states) hidden_states = self.mlps[1](hidden_states) hidden_states = residual + hidden_states + shortcut_mlp_output return hidden_states class _LongcatMTPNormWrapper(nn.Module): """Checkpoint-compatible wrapper for ``enorm.m`` / ``hnorm.m``.""" def __init__(self, hidden_size: int, eps: float): super().__init__() self.m = LongcatFlashRMSNorm(hidden_size, eps=eps) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return self.m(hidden_states) class _LongcatMTPTransformerLayer(nn.Module): """Namespace wrapper matching ``transformer_layer.mlp.*`` checkpoint keys.""" def __init__(self, config: LongcatFlashSparseConfig): super().__init__() self.mlp = LongcatFlashMLP(config) class LongcatFlashSparseMTPMLA(LongcatFlashSparseMLA): """One physical MTP attention module reused across three conceptual steps. Step 0 owns the MTP indexer. Steps 1 and 2 execute the same MLA weights but reuse step 0's absolute top-K positions, matching ``dsa_mtp_cli=True``. Each conceptual MTP step must use its own KV cache because parameter sharing does not imply state/cache sharing. """ def __init__(self, config: LongcatFlashSparseConfig): # The physical MTP checkpoint has one self_attn.indexer.* parameter set. super().__init__(config, layer_idx=0, indexer_owner=True) self.last_mtp_step_index: Optional[int] = None def forward_mtp( self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor], past_key_values: Optional[Cache], *, sequence_attention_mask: Optional[torch.Tensor], mtp_step_index: int, topk_indices: Optional[torch.Tensor], **kwargs, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: if mtp_step_index not in (0, 1, 2): raise ValueError(f"mtp_step_index must be 0, 1, or 2; got {mtp_step_index}") self.last_mtp_step_index = mtp_step_index batch_size, query_len, _ = hidden_states.shape if mtp_step_index == 0: out, _, owner_topk = super().forward( hidden_states=hidden_states, position_embeddings=position_embeddings, attention_mask=attention_mask, past_key_values=past_key_values, sequence_attention_mask=sequence_attention_mask, **kwargs, ) if self.last_lsa_mode == "full-owner": self.last_lsa_mode = "full-mtp-owner" elif self.last_lsa_mode == "sparse-owner": self.last_lsa_mode = "sparse-mtp-owner" return out, owner_topk # Steps 2/3 reuse the exact positional selection generated by MTP step 1. # They deliberately DO NOT run or populate their own indexer K cache. past_len = ( past_key_values.get_seq_length(self.layer_idx) if past_key_values is not None else 0 ) total_kv_len = past_len + query_len if total_kv_len <= self.index_topk: if topk_indices is not None: raise RuntimeError( "MTP CLI reuse received top-K indices on the exact full-attention " f"path (kv={total_kv_len} <= {self.index_topk})." ) self.last_lsa_mode = "full-mtp-reuse" self.last_lsa_valid_topk_range = None out, _ = LongcatFlashMLA.forward( self, hidden_states=hidden_states, position_embeddings=position_embeddings, attention_mask=attention_mask, past_key_values=past_key_values, **kwargs, ) return out, None if topk_indices is None: raise RuntimeError( "MTP step 2/3 requires the step-1 CLI top-K above " f"index_topk={self.index_topk}; refusing to recompute or fall back dense." ) if topk_indices.shape[:2] != (batch_size, query_len): raise RuntimeError( "MTP CLI top-K query shape mismatch: " f"topk={tuple(topk_indices.shape)}, hidden={tuple(hidden_states.shape)}" ) valid_counts = (topk_indices >= 0).sum(dim=-1) self.last_lsa_mode = "sparse-mtp-reuse" self.last_lsa_valid_topk_range = ( int(valid_counts.min().item()), int(valid_counts.max().item()), ) q_lora = self.q_a_layernorm(self.q_a_proj(hidden_states)) out = self._forward_sparse_mla( hidden_states, q_lora, topk_indices, position_embeddings, past_key_values, ) return out, topk_indices class LongcatFlashSparseMTPLayer(nn.Module): """The single physical ``model.mtp.layers.0`` parameterized module.""" def __init__(self, config: LongcatFlashSparseConfig): super().__init__() self.config = config self.eh_proj = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False) self.enorm = _LongcatMTPNormWrapper(config.hidden_size, config.rms_norm_eps) self.hnorm = _LongcatMTPNormWrapper(config.hidden_size, config.rms_norm_eps) self.input_layernorm = LongcatFlashRMSNorm( config.hidden_size, eps=config.rms_norm_eps ) self.self_attn = LongcatFlashSparseMTPMLA(config) self.post_attention_layernorm = LongcatFlashRMSNorm( config.hidden_size, eps=config.rms_norm_eps ) self.transformer_layer = _LongcatMTPTransformerLayer(config) def forward( self, *, input_ids: torch.LongTensor, previous_hidden_states: torch.Tensor, position_ids: torch.LongTensor, shared_inputs_embeds: torch.Tensor, past_key_values: Optional[Cache], attention_mask: Optional[torch.Tensor], mtp_step_index: int, topk_indices: Optional[torch.Tensor], position_embeddings: tuple[torch.Tensor, torch.Tensor], causal_mask: Optional[torch.Tensor], **kwargs, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: if shared_inputs_embeds.shape != previous_hidden_states.shape: raise ValueError( "MTP input embedding and previous hidden-state shapes must match: " f"embed={tuple(shared_inputs_embeds.shape)}, " f"hidden={tuple(previous_hidden_states.shape)}" ) # LongCat NextN uses the target model's token embedding directly here: # eh_proj(cat(enorm(embed), hnorm(target_hidden))). Do not apply the # DeepSeek-style absolute-position-0 embedding suppression. embed_normed = self.enorm(shared_inputs_embeds) hidden_normed = self.hnorm(previous_hidden_states) hidden_states = self.eh_proj( torch.cat([embed_normed, hidden_normed], dim=-1) ) residual = hidden_states attn_input = self.input_layernorm(hidden_states) attn_output, returned_topk = self.self_attn.forward_mtp( hidden_states=attn_input, position_embeddings=position_embeddings, attention_mask=causal_mask, past_key_values=past_key_values, sequence_attention_mask=attention_mask, mtp_step_index=mtp_step_index, topk_indices=topk_indices, **kwargs, ) hidden_states = residual + attn_output residual = hidden_states mlp_input = self.post_attention_layernorm(hidden_states) mlp_output = self.transformer_layer.mlp(mlp_input) hidden_states = residual + mlp_output return hidden_states, returned_topk class LongcatFlashSparseMTP(nn.Module): """Native 3-step LongCat MTP runtime with one shared physical module. Checkpoint layout: model.mtp.embed_tokens.weight model.mtp.layers.0.* model.mtp.norm.weight Runtime semantics: * exactly one parameterized layer object; * conceptual MTP steps = config.mtp_num_layers (=3); * all three calls reuse that one layer's parameters; * step 1 owns the sparse indexer; * steps 2 and 3 reuse step 1's positional top-K when sparse; * each conceptual step receives a distinct KV cache. """ def __init__(self, config: LongcatFlashSparseConfig): super().__init__() if int(getattr(config, "mtp_num_layers", 0)) != 3: raise NotImplementedError( "This release runtime is validated for mtp_num_layers=3." ) if getattr(config, "mtp_replicate_modules", False) is not True: raise NotImplementedError( "This release runtime requires mtp_replicate_modules=True." ) if getattr(config, "dsa_mtp_cli", False) is not True: raise NotImplementedError( "This release runtime requires dsa_mtp_cli=True." ) self.config = config self.num_steps = 3 # Retain the physical checkpoint tensor under its native key. The parent # CausalLM API deliberately uses the MAIN model's embedding for execution, # matching standard MTP parameter-sharing semantics. Keeping this module # here makes the full 21-key checkpoint load/export contract explicit. self.embed_tokens = nn.Embedding( config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id, ) self.layers = nn.ModuleList([LongcatFlashSparseMTPLayer(config)]) self.norm = LongcatFlashRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = LongcatFlashRotaryEmbedding(config=config) self.last_step_modes: List[Optional[str]] = [None, None, None] self.last_step_valid_topk_ranges: List[Optional[tuple[int, int]]] = [ None, None, None ] @property def physical_layer(self) -> LongcatFlashSparseMTPLayer: return self.layers[0] def create_step_caches(self) -> List[NgramCache]: # Three conceptual layers share parameters but MUST NOT share attention # state. Each gets an independent one-layer DynamicCache. return [NgramCache(config=self.config) for _ in range(self.num_steps)] def forward_step( self, *, input_ids: torch.LongTensor, previous_hidden_states: torch.Tensor, mtp_step_index: int, shared_inputs_embeds: Optional[torch.Tensor] = None, past_key_values: Optional[NgramCache] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, cache_position: Optional[torch.LongTensor] = None, topk_indices: Optional[torch.Tensor] = None, use_cache: bool = True, **kwargs, ) -> Dict[str, Any]: if mtp_step_index not in range(self.num_steps): raise ValueError( f"mtp_step_index must be 0..{self.num_steps - 1}; got {mtp_step_index}" ) if shared_inputs_embeds is None: raise ValueError( "LongCat MTP requires target-model shared_inputs_embeds. " "Call LongcatFlashSparseForCausalLM.mtp_forward_step(), or pass " "embeddings from the target model's model.embed_tokens explicitly." ) if use_cache and past_key_values is None: past_key_values = NgramCache(config=self.config) if past_key_values is not None and not isinstance(past_key_values, NgramCache): raise TypeError("LongCat Sparse MTP requires NgramCache-compatible step caches.") if cache_position is None: past_seen_tokens = ( past_key_values.get_seq_length(0) if past_key_values is not None else 0 ) cache_position = ( torch.arange( input_ids.shape[1], device=shared_inputs_embeds.device, dtype=torch.long, ) + past_seen_tokens ) if position_ids is None: position_ids = cache_position.unsqueeze(0) causal_mask = create_causal_mask( config=self.config, inputs_embeds=shared_inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, position_ids=position_ids, ) position_embeddings = self.rotary_emb( previous_hidden_states, position_ids ) hidden_states, returned_topk = self.physical_layer( input_ids=input_ids, previous_hidden_states=previous_hidden_states, position_ids=position_ids, shared_inputs_embeds=shared_inputs_embeds, past_key_values=past_key_values, attention_mask=attention_mask, mtp_step_index=mtp_step_index, topk_indices=topk_indices, position_embeddings=position_embeddings, causal_mask=causal_mask, cache_position=cache_position, **kwargs, ) hidden_states = self.norm(hidden_states) mode = self.physical_layer.self_attn.last_lsa_mode valid_range = self.physical_layer.self_attn.last_lsa_valid_topk_range self.last_step_modes[mtp_step_index] = mode self.last_step_valid_topk_ranges[mtp_step_index] = valid_range return { "hidden_states": hidden_states, "past_key_values": past_key_values, "topk_indices": returned_topk, "mtp_step_index": mtp_step_index, "lsa_mode": mode, "valid_topk_range": valid_range, } class LongcatFlashSparseModel(LongcatFlashModel): """LongCat-Flash-Lite-Sparse with genuine SI+CLI LSA selection.""" _keys_to_ignore_on_load_unexpected = [] config_class = LongcatFlashSparseConfig def __init__(self, config): super().__init__(config) self.embed_tokens = nn.Embedding( config.vocab_size, config.hidden_size, self.padding_idx ) self.ngram_embeddings = NgramEmbedding(config, self.embed_tokens) self.layers = nn.ModuleList( [ LongcatFlashSparseDecoderLayer(config, layer_idx) for layer_idx in range(config.num_layers) ] ) self.head_dim = config.head_dim self.config.num_hidden_layers = 2 * config.num_layers self.norm = LongcatFlashRMSNorm( config.hidden_size, eps=config.rms_norm_eps ) self.rotary_emb = LongcatFlashRotaryEmbedding(config=config) self.mtp = LongcatFlashSparseMTP(config) self.gradient_checkpointing = False self.post_init() def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, inputs_embeds: Optional[torch.FloatTensor] = None, cache_position: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_hidden_states: Optional[bool] = None, **kwargs, ) -> BaseModelOutputWithPast: if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError( "You must specify exactly one of input_ids or inputs_embeds" ) output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) all_hidden_states = [] if output_hidden_states else None ngram_context = None if isinstance(past_key_values, NgramCache): ngram_context = past_key_values.ngram_context if inputs_embeds is None: inputs_embeds = self.ngram_embeddings( input_ids, ngram_context=ngram_context ) if use_cache and past_key_values is None: past_key_values = NgramCache(config=self.config) if use_cache and isinstance(past_key_values, NgramCache): past_key_values.update_ngram_context(input_ids) if cache_position is None: past_seen_tokens = ( past_key_values.get_seq_length() if past_key_values is not None else 0 ) cache_position = torch.arange( inputs_embeds.shape[1], device=inputs_embeds.device ) + past_seen_tokens if position_ids is None: position_ids = cache_position.unsqueeze(0) causal_mask = create_causal_mask( config=self.config, inputs_embeds=inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, position_ids=position_ids, ) hidden_states = inputs_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids) for decoder_layer in self.layers[: self.config.num_layers]: if all_hidden_states is not None: all_hidden_states.append(hidden_states) hidden_states = decoder_layer( hidden_states, attention_mask=causal_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, sequence_attention_mask=attention_mask, **kwargs, ) hidden_states = self.norm(hidden_states) if all_hidden_states is not None: all_hidden_states.append(hidden_states) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=past_key_values, hidden_states=( tuple(all_hidden_states) if all_hidden_states is not None else None ), attentions=None, ) class LongcatFlashSparseForCausalLM(LongcatFlashForCausalLM): """Standalone HF causal LM for LongCat-Flash-Lite-Sparse.""" _keys_to_ignore_on_load_unexpected = [] config_class = LongcatFlashSparseConfig _tied_weights_keys = { "lm_head.weight": "model.embed_tokens.weight", "model.ngram_embeddings.word_embeddings.weight": "model.embed_tokens.weight", } def __init__(self, config): super().__init__(config) self.model = LongcatFlashSparseModel(config) @torch.no_grad() def mtp_forward_step( self, *, input_ids: torch.LongTensor, previous_hidden_states: torch.Tensor, mtp_step_index: int, past_key_values: Optional[NgramCache] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, cache_position: Optional[torch.LongTensor] = None, topk_indices: Optional[torch.Tensor] = None, use_cache: bool = True, return_logits: bool = True, **kwargs, ) -> Dict[str, Any]: """Execute one native LongCat MTP conceptual step. This is a drafter primitive, not a speculative-verification scheduler. A serving engine can call it for steps 0..2 while maintaining one cache per conceptual step and passing step-0 ``topk_indices`` into steps 1/2. """ shared_inputs_embeds = self.model.embed_tokens( input_ids.to(self.model.embed_tokens.weight.device) ) result = self.model.mtp.forward_step( input_ids=input_ids, previous_hidden_states=previous_hidden_states, mtp_step_index=mtp_step_index, shared_inputs_embeds=shared_inputs_embeds, past_key_values=past_key_values, attention_mask=attention_mask, position_ids=position_ids, cache_position=cache_position, topk_indices=topk_indices, use_cache=use_cache, **kwargs, ) if return_logits: result["logits"] = self.lm_head(result["hidden_states"]) return result @torch.no_grad() def generate(self, inputs=None, generation_config=None, **kwargs): if "past_key_values" not in kwargs or kwargs["past_key_values"] is None: kwargs["past_key_values"] = NgramCache(config=self.config) return super().generate( inputs=inputs, generation_config=generation_config, **kwargs, ) _register_longcat_sparse_checkpoint_conversions() __all__ = [ "LongcatFlashSparsePreTrainedModel", "LongcatFlashSparseModel", "LongcatFlashSparseForCausalLM", "LongcatFlashSparseMTP", "LongcatFlashSparseMTPLayer", ]