| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Inference-only LongCat-Flash-Lite (n-gram embedding) model. |
| 4 | |
| 5 | ``LongcatFlashNgramForCausalLM`` is LongCat-Flash (MLA dual-attention + |
| 6 | zero-expert MoE + YaRN) plus an n-gram embedding input layer: each position's |
| 7 | embedding fuses the token embedding with hashed embeddings of the preceding |
| 8 | ``n`` tokens. That per-request token history is isolated in a Model-Runner-V2 |
| 9 | :class:`LongcatNgramModelState` (mirroring ``DiffusionGemmaModelState``), so |
| 10 | ``get_model_state_cls`` makes the model MRV2-only. |
| 11 | """ |
| 12 | |
| 13 | from collections.abc import Iterable |
| 14 | from typing import Any |
| 15 | |
| 16 | import torch |
| 17 | from torch import nn |
| 18 | |
| 19 | from vllm import _custom_ops as ops |
| 20 | from vllm.config import VllmConfig |
| 21 | from vllm.distributed import get_pp_group |
| 22 | from vllm.model_executor.layers.logits_processor import LogitsProcessor |
| 23 | from vllm.model_executor.layers.vocab_parallel_embedding import ( |
| 24 | ParallelLMHead, |
| 25 | VocabParallelEmbedding, |
| 26 | ) |
| 27 | from vllm.v1.core.sched.output import NewRequestData |
| 28 | from vllm.v1.worker.gpu.input_batch import InputBatch |
| 29 | from vllm.v1.worker.gpu.model_states.default import DefaultModelState |
| 30 | from vllm.v1.worker.gpu.states import RequestState |
| 31 | |
| 32 | from .interfaces import SupportsLoRA, SupportsPP |
| 33 | from .longcat_flash import FlashConfig, FlashModel |
| 34 | from .utils import AutoWeightsLoader, PPMissingLayer, WeightsMapper, maybe_prefix |
| 35 | |
| 36 | |
| 37 | def uses_ngram_embedding(config: FlashConfig) -> bool: |
| 38 | return getattr(config, "ngram_vocab_size_ratio", None) is not None |
| 39 | |
| 40 | |
| 41 | def _config_dtype(config: FlashConfig) -> torch.dtype: |
| 42 | dt = getattr(config, "torch_dtype", None) or getattr(config, "dtype", None) |
| 43 | if isinstance(dt, torch.dtype): |
| 44 | return dt |
| 45 | return getattr(torch, str(dt), None) or torch.bfloat16 |
| 46 | |
| 47 | |
| 48 | class NgramEmbedding(nn.Module): |
| 49 | """Token embedding fused with hashed n-gram embeddings. |
| 50 | |
| 51 | TP-sharded: the ``k*(n-1)`` per-embedder tables are concatenated into one |
| 52 | :class:`VocabParallelEmbedding` (``oe_embedder``) with per-embedder offsets, |
| 53 | and the projections are stacked into one ``oe_projection`` applied with a |
| 54 | single ``bmm``. Hashing math is ported from the HF reference. |
| 55 | """ |
| 56 | |
| 57 | def __init__(self, config: FlashConfig, base_embeddings: nn.Module) -> None: |
| 58 | super().__init__() |
| 59 | self.config = config |
| 60 | self.word_embeddings = base_embeddings |
| 61 | |
| 62 | self.m = config.ngram_vocab_size_ratio * config.vocab_size |
| 63 | self.k = config.emb_split_num |
| 64 | self.n = config.emb_neighbor_num |
| 65 | self.pad_id = config.pad_token_id |
| 66 | self.eos_token_id = config.eos_token_id |
| 67 | self._dtype = _config_dtype(config) |
| 68 | |
| 69 | self._init_ngram_embeddings() |
| 70 | |
| 71 | def _init_ngram_embeddings(self) -> None: |
| 72 | self.num_embedders = self.k * (self.n - 1) |
| 73 | oe_dim = self.config.hidden_size // self.num_embedders |
| 74 | self.oe_dim = oe_dim |
| 75 | |
| 76 | # Exclusive prefix sums of per-embedder table sizes; each embedder's |
| 77 | # local id is offset into the single concatenated table. |
| 78 | sizes = [int(self.m + i * 2 + 1) for i in range(self.num_embedders)] |
| 79 | offsets = [0] |
| 80 | for s in sizes: |
| 81 | offsets.append(offsets[-1] + s) |
| 82 | self._offsets = offsets # len num_embedders + 1 |
| 83 | self._sizes = sizes |
| 84 | |
| 85 | self.oe_embedder = VocabParallelEmbedding( |
| 86 | offsets[-1], oe_dim, params_dtype=self._dtype |
| 87 | ) |
| 88 | # Stacked projections: oe_projection[i] = post_projs[i].weight.T |
| 89 | self.oe_projection = nn.Parameter( |
| 90 | torch.empty( |
| 91 | self.num_embedders, oe_dim, self.config.hidden_size, dtype=self._dtype |
| 92 | ), |
| 93 | requires_grad=False, |
| 94 | ) |
| 95 | |
| 96 | # Precomputed tables for the CUDA n-gram id kernel (ngram_embedding |
| 97 | # _kernels.cu): ne_weights[i][j][delta] = vocab^delta mod ne_mods[i][j], |
| 98 | # ne_mods[i][j] = m + 2*(i*k+j) + 1. Registered as non-persistent buffers |
| 99 | # so they follow the module to the device (not part of the checkpoint). |
| 100 | vocab = self.config.vocab_size |
| 101 | ne_weights = torch.zeros(self.n - 1, self.k, self.n, dtype=torch.int32) |
| 102 | ne_mods = torch.zeros(self.n - 1, self.k, dtype=torch.int32) |
| 103 | for i in range(self.n - 1): |
| 104 | for j in range(self.k): |
| 105 | mod = int(self.m + 2 * (i * self.k + j) + 1) |
| 106 | ne_mods[i, j] = mod |
| 107 | for delta in range(self.n): |
| 108 | ne_weights[i, j, delta] = pow(vocab, delta, mod) |
| 109 | self.register_buffer("ne_weights", ne_weights, persistent=False) |
| 110 | self.register_buffer("ne_mods", ne_mods, persistent=False) |
| 111 | self.register_buffer( |
| 112 | "exclusive_sizes", |
| 113 | torch.tensor(offsets, dtype=torch.int32), |
| 114 | persistent=False, |
| 115 | ) |
| 116 | |
| 117 | def load_weight(self, weight_name: str, loaded_weight: torch.Tensor) -> str: |
| 118 | """Split a per-embedder checkpoint weight into the sharded layout. |
| 119 | |
| 120 | Returns the destination parameter's qualified name (relative to the |
| 121 | enclosing model) so the caller can mark it loaded for completeness |
| 122 | checks. |
| 123 | """ |
| 124 | if "ngram_embeddings.embedders." in weight_name: |
| 125 | index = int( |
| 126 | weight_name.split("ngram_embeddings.embedders.")[1].split(".")[0] |
| 127 | ) |
| 128 | lo, hi = self._offsets[index], self._offsets[index + 1] |
| 129 | assert hi - lo == loaded_weight.shape[0], ( |
| 130 | f"{hi - lo=} {loaded_weight.shape[0]=}" |
| 131 | ) |
| 132 | shard = self.oe_embedder.shard_indices |
| 133 | tp_start, tp_end = shard.org_vocab_start_index, shard.org_vocab_end_index |
| 134 | load_start, load_end = max(lo, tp_start), min(hi, tp_end) |
| 135 | if load_start < load_end: |
| 136 | self.oe_embedder.weight.data[ |
| 137 | load_start - tp_start : load_end - tp_start |
| 138 | ] = loaded_weight[load_start - lo : load_end - lo] |
| 139 | return "ngram_embeddings.oe_embedder.weight" |
| 140 | elif "ngram_embeddings.post_projs." in weight_name: |
| 141 | index = int( |
| 142 | weight_name.split("ngram_embeddings.post_projs.")[1].split(".")[0] |
| 143 | ) |
| 144 | self.oe_projection.data[index].copy_(loaded_weight.t()) |
| 145 | return "ngram_embeddings.oe_projection" |
| 146 | else: |
| 147 | raise AssertionError(f"Unexpected ngram weight: {weight_name}") |
| 148 | |
| 149 | def embed_batched( |
| 150 | self, input_ids: torch.Tensor, oe_ids: torch.Tensor |
| 151 | ) -> torch.Tensor: |
| 152 | """Fused n-gram embedding for a flat batch given precomputed ids. |
| 153 | |
| 154 | Args: |
| 155 | input_ids: ``[num_tokens]`` current token per position. |
| 156 | oe_ids: ``[num_tokens, num_embedders]`` global (offset) n-gram ids, |
| 157 | as produced by the ``ngram_compute_n_gram_ids`` kernel. |
| 158 | Returns: ``[num_tokens, hidden]``. |
| 159 | |
| 160 | """ |
| 161 | word = self.word_embeddings(input_ids) # [N, H] |
| 162 | flat = oe_ids.permute(1, 0).contiguous() # [num_embedders, N] |
| 163 | oe = self.oe_embedder(flat) # [num_embedders, N, oe_dim] |
| 164 | proj = torch.bmm(oe, self.oe_projection) # [num_embedders, N, H] |
| 165 | all_h = torch.cat([word.unsqueeze(0), proj], dim=0) # [ne+1, N, H] |
| 166 | return all_h.mean(dim=0) # [N, H] |
| 167 | |
| 168 | |
| 169 | class FlashNgramModel(FlashModel): |
| 170 | """FlashModel whose input embedding is an :class:`NgramEmbedding`.""" |
| 171 | |
| 172 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: |
| 173 | # Each FlashDecoderLayer is a *dual* layer (2 attentions), so the number |
| 174 | # of decoder layers is ``num_layers``. The ngram HF config sets |
| 175 | # ``num_hidden_layers`` to a multiple of that (attention-module count), |
| 176 | # which FlashModel would otherwise build as too many (dead) layers. |
| 177 | hf = vllm_config.model_config.hf_config |
| 178 | num_layers = getattr(hf, "num_layers", None) |
| 179 | if num_layers is not None and hf.num_hidden_layers != num_layers: |
| 180 | hf.num_hidden_layers = num_layers |
| 181 | super().__init__(vllm_config=vllm_config, prefix=prefix) |
| 182 | self.ngram_embeddings: NgramEmbedding | None |
| 183 | if get_pp_group().is_first_rank and uses_ngram_embedding(self.config): |
| 184 | self.ngram_embeddings = NgramEmbedding(self.config, self.embed_tokens) |
| 185 | else: |
| 186 | self.ngram_embeddings = None |
| 187 | |
| 188 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 189 | # Names arrive with the ``model.`` prefix already stripped (routed here |
| 190 | # by AutoWeightsLoader). Split the concatenated/sharded ngram tables and |
| 191 | # stacked projections; delegate everything else to FlashModel. |
| 192 | loaded: set[str] = set() |
| 193 | rest: list[tuple[str, torch.Tensor]] = [] |
| 194 | for name, w in weights: |
| 195 | if self.ngram_embeddings is not None and ( |
| 196 | "ngram_embeddings.embedders." in name |
| 197 | or "ngram_embeddings.post_projs." in name |
| 198 | ): |
| 199 | loaded.add(self.ngram_embeddings.load_weight(name, w)) |
| 200 | else: |
| 201 | rest.append((name, w)) |
| 202 | loaded |= super().load_weights(rest) |
| 203 | return loaded |
| 204 | |
| 205 | |
| 206 | class LongcatFlashNgramForCausalLM(nn.Module, SupportsLoRA, SupportsPP): |
| 207 | """LongCat-Flash-Lite for causal LM (MRV2-only, n-gram embedding).""" |
| 208 | |
| 209 | # MTP weights are not part of this model. |
| 210 | hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"model.mtp.": None}) |
| 211 | |
| 212 | packed_modules_mapping = { |
| 213 | "qkv_proj": ["q_proj", "k_proj", "v_proj"], |
| 214 | "gate_up_proj": ["gate_proj", "up_proj"], |
| 215 | } |
| 216 | |
| 217 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: |
| 218 | super().__init__() |
| 219 | if not vllm_config.use_v2_model_runner: |
| 220 | raise NotImplementedError( |
| 221 | "LongcatFlashNgramForCausalLM (LongCat-Flash-Lite) requires the " |
| 222 | "V2 model runner for its n-gram embedding state; it is selected " |
| 223 | "automatically unless VLLM_USE_V2_MODEL_RUNNER=0 is set." |
| 224 | ) |
| 225 | config = FlashConfig(**vllm_config.model_config.hf_config.__dict__) |
| 226 | config.intermediate_size = getattr( |
| 227 | config, "ffn_hidden_size", config.intermediate_size |
| 228 | ) |
| 229 | self.config = config |
| 230 | self.quant_config = vllm_config.quant_config |
| 231 | |
| 232 | self.model = FlashNgramModel( |
| 233 | vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") |
| 234 | ) |
| 235 | if get_pp_group().is_last_rank: |
| 236 | self.lm_head = ParallelLMHead( |
| 237 | config.vocab_size, |
| 238 | config.hidden_size, |
| 239 | quant_config=self.quant_config, |
| 240 | prefix=maybe_prefix(prefix, "lm_head"), |
| 241 | ) |
| 242 | else: |
| 243 | self.lm_head = PPMissingLayer() |
| 244 | self.logits_processor = LogitsProcessor(config.vocab_size) |
| 245 | self.make_empty_intermediate_tensors = ( |
| 246 | self.model.make_empty_intermediate_tensors |
| 247 | ) |
| 248 | |
| 249 | @staticmethod |
| 250 | def get_model_state_cls() -> type["LongcatNgramModelState"]: |
| 251 | return LongcatNgramModelState |
| 252 | |
| 253 | def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: |
| 254 | return self.model.embed_input_ids(input_ids) |
| 255 | |
| 256 | def forward( |
| 257 | self, |
| 258 | input_ids: torch.Tensor | None, |
| 259 | positions: torch.Tensor, |
| 260 | intermediate_tensors=None, |
| 261 | inputs_embeds: torch.Tensor | None = None, |
| 262 | ): |
| 263 | # inputs_embeds is produced by LongcatNgramModelState.prepare_inputs. |
| 264 | return self.model(input_ids, positions, intermediate_tensors, inputs_embeds) |
| 265 | |
| 266 | def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: |
| 267 | return self.logits_processor(self.lm_head, hidden_states) |
| 268 | |
| 269 | def get_expert_mapping(self): |
| 270 | return self.model.get_expert_mapping() |
| 271 | |
| 272 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 273 | # AutoWeightsLoader routes ``model.*`` to FlashNgramModel.load_weights |
| 274 | # (which handles the ngram split) and ``lm_head.*`` to the head. |
| 275 | loader = AutoWeightsLoader(self) |
| 276 | return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |
| 277 | |
| 278 | |
| 279 | class LongcatNgramModelState(DefaultModelState): |
| 280 | # prepare_inputs builds its own inputs_embeds from n-gram token embeddings. |
| 281 | supports_prompt_embeds = False |
| 282 | |
| 283 | """Per-request n-gram token history for LongCat-Flash-Lite. |
| 284 | |
| 285 | Maintains a small CPU-side per-slot context (last ``n-1`` processed tokens) |
| 286 | and a persistent ``inputs_embeds`` buffer. ``prepare_inputs`` computes the |
| 287 | fused n-gram embedding per request into the buffer, handed to the model |
| 288 | forward as ``inputs_embeds``. |
| 289 | """ |
| 290 | |
| 291 | def __init__(self, vllm_config, model, encoder_cache, device) -> None: |
| 292 | super().__init__(vllm_config, model, encoder_cache, device) |
| 293 | config = model.config |
| 294 | self.ngram = model.model.ngram_embeddings |
| 295 | self.n = int(config.emb_neighbor_num) |
| 296 | self.ctx_len = self.n - 1 |
| 297 | self.eos_id = int(config.eos_token_id) |
| 298 | |
| 299 | # Per-slot left-context: last ``n-1`` processed tokens, EOS negated. A |
| 300 | # negative entry (incl. the -1 fill) marks a context boundary that stops |
| 301 | # the n-gram walk (matches the kernel's EOS break / fresh-request start). |
| 302 | self.token_context = torch.full( |
| 303 | (self.max_num_reqs, self.ctx_len), -1, dtype=torch.int32, device=device |
| 304 | ) |
| 305 | |
| 306 | self._inputs_embeds_buf = torch.zeros( |
| 307 | self.max_num_tokens, |
| 308 | config.hidden_size, |
| 309 | dtype=self.dtype, |
| 310 | device=device, |
| 311 | ) |
| 312 | |
| 313 | def _neg_eos(self, toks: list[int]) -> list[int]: |
| 314 | return [-t if t == self.eos_id else t for t in toks] |
| 315 | |
| 316 | def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: |
| 317 | super().add_request(req_index, new_req_data) # rope positions |
| 318 | # Fresh request -> no left-context (-1 fill). On resume, seed from the |
| 319 | # already-processed token tail. Use prefill_token_ids (full processed |
| 320 | # sequence incl. generated tokens on v2 resume), like DefaultModelState; |
| 321 | # the prompt alone would be too short when resuming after decode. |
| 322 | ctx = [-1] * self.ctx_len |
| 323 | ncomp = new_req_data.num_computed_tokens |
| 324 | toks_src = new_req_data.prefill_token_ids or new_req_data.prompt_token_ids |
| 325 | if ncomp > 0 and toks_src is not None: |
| 326 | lo = max(0, ncomp - self.ctx_len) |
| 327 | toks = self._neg_eos(list(toks_src[lo:ncomp])) |
| 328 | ctx[self.ctx_len - len(toks) :] = toks |
| 329 | self.token_context[req_index] = torch.tensor( |
| 330 | ctx, dtype=torch.int32, device=self.token_context.device |
| 331 | ) |
| 332 | |
| 333 | def prepare_inputs( |
| 334 | self, input_batch: InputBatch, req_states: RequestState |
| 335 | ) -> dict[str, Any]: |
| 336 | model_inputs = super().prepare_inputs(input_batch, req_states) # positions |
| 337 | num_tokens = input_batch.num_tokens |
| 338 | num_padded = input_batch.num_tokens_after_padding |
| 339 | input_ids = input_batch.input_ids[:num_tokens] |
| 340 | embeds = self._inputs_embeds_buf[:num_padded] |
| 341 | |
| 342 | oe_ids = self._compute_oe_ids(input_batch) |
| 343 | embeds[:num_tokens].copy_(self.ngram.embed_batched(input_ids, oe_ids)) |
| 344 | model_inputs["inputs_embeds"] = embeds |
| 345 | return model_inputs |
| 346 | |
| 347 | def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]: |
| 348 | # FULL cudagraph replay reads only the captured buffers, so capture must |
| 349 | # reference the same persistent ``inputs_embeds`` buffer prepare_inputs |
| 350 | # re-fills (the base class wires this for multimodal models only). |
| 351 | model_inputs = super().prepare_dummy_inputs(num_reqs, num_tokens) # positions |
| 352 | model_inputs["inputs_embeds"] = self._inputs_embeds_buf[:num_tokens] |
| 353 | return model_inputs |
| 354 | |
| 355 | def _compute_oe_ids(self, input_batch: InputBatch) -> torch.Tensor: |
| 356 | """Batched global n-gram ids ``[num_tokens, num_embedders]``. |
| 357 | |
| 358 | Assembles an ephemeral per-request token table (``[n-1] context ++ |
| 359 | current tokens``, EOS-negated) and runs the ``ngram_compute_n_gram_ids`` |
| 360 | CUDA kernel for the whole batch, then rolls each slot's context forward. |
| 361 | """ |
| 362 | device = self.token_context.device |
| 363 | num_tokens = input_batch.num_tokens |
| 364 | num_reqs = input_batch.num_reqs |
| 365 | ctx_len = self.ctx_len |
| 366 | idx_mapping = input_batch.idx_mapping[:num_reqs].long() |
| 367 | qsl = input_batch.query_start_loc[: num_reqs + 1].to(torch.int32) |
| 368 | cur = input_batch.input_ids[:num_tokens].to(torch.int32) |
| 369 | |
| 370 | cur_neg = torch.where(cur == self.eos_id, -cur, cur) |
| 371 | req_lens = qsl[1:] - qsl[:-1] |
| 372 | max_len = int(req_lens.max().item()) |
| 373 | width = ctx_len + max_len |
| 374 | |
| 375 | # table[r] = [context(n-1) | current tokens | pad(-1)] |
| 376 | table = torch.full((num_reqs, width), -1, dtype=torch.int32, device=device) |
| 377 | table[:, :ctx_len] = self.token_context[idx_mapping] |
| 378 | tok_req = torch.repeat_interleave( |
| 379 | torch.arange(num_reqs, device=device), req_lens.long() |
| 380 | ) |
| 381 | col = ctx_len + ( |
| 382 | torch.arange(num_tokens, device=device) - qsl[:-1].long()[tok_req] |
| 383 | ) |
| 384 | table[tok_req, col] = cur_neg |
| 385 | |
| 386 | column_starts = torch.full( |
| 387 | (num_reqs,), ctx_len, dtype=torch.int32, device=device |
| 388 | ) |
| 389 | row_indices = torch.arange(num_reqs, dtype=torch.int64, device=device) |
| 390 | n_gram_ids = torch.empty( |
| 391 | num_tokens, self.ngram.num_embedders, dtype=torch.int32, device=device |
| 392 | ) |
| 393 | ops.ngram_compute_n_gram_ids( |
| 394 | self.n, |
| 395 | self.ngram.k, |
| 396 | self.ngram.ne_weights, |
| 397 | self.ngram.ne_mods, |
| 398 | self.ngram.exclusive_sizes, |
| 399 | qsl, |
| 400 | table, |
| 401 | row_indices, |
| 402 | column_starts, |
| 403 | n_gram_ids, |
| 404 | ) |
| 405 | |
| 406 | # Roll context: new context = last n-1 of [context | current] per slot. |
| 407 | gather = req_lens.long().unsqueeze(1) + torch.arange( |
| 408 | ctx_len, device=device |
| 409 | ).unsqueeze(0) |
| 410 | rows = torch.arange(num_reqs, device=device).unsqueeze(1) |
| 411 | self.token_context[idx_mapping] = table[rows, gather] |
| 412 | return n_gram_ids.long() |