wehub-resource-sync
406 lines
17 KiB
Python
406 lines
17 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Inference-only LongCat-Flash-Lite (n-gram embedding) model.
|
|
|
|
``LongcatFlashNgramForCausalLM`` is LongCat-Flash (MLA dual-attention +
|
|
zero-expert MoE + YaRN) plus an n-gram embedding input layer: each position's
|
|
embedding fuses the token embedding with hashed embeddings of the preceding
|
|
``n`` tokens. That per-request token history is isolated in a Model-Runner-V2
|
|
:class:`LongcatNgramModelState` (mirroring ``DiffusionGemmaModelState``), so
|
|
``get_model_state_cls`` makes the model MRV2-only.
|
|
"""
|
|
|
|
from collections.abc import Iterable
|
|
from typing import Any
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
from vllm import _custom_ops as ops
|
|
from vllm.config import VllmConfig
|
|
from vllm.distributed import get_pp_group
|
|
from vllm.model_executor.layers.logits_processor import LogitsProcessor
|
|
from vllm.model_executor.layers.vocab_parallel_embedding import (
|
|
ParallelLMHead,
|
|
VocabParallelEmbedding,
|
|
)
|
|
from vllm.v1.core.sched.output import NewRequestData
|
|
from vllm.v1.worker.gpu.input_batch import InputBatch
|
|
from vllm.v1.worker.gpu.model_states.default import DefaultModelState
|
|
from vllm.v1.worker.gpu.states import RequestState
|
|
|
|
from .interfaces import SupportsLoRA, SupportsPP
|
|
from .longcat_flash import FlashConfig, FlashModel
|
|
from .utils import AutoWeightsLoader, PPMissingLayer, maybe_prefix
|
|
|
|
|
|
def uses_ngram_embedding(config: FlashConfig) -> bool:
|
|
return getattr(config, "ngram_vocab_size_ratio", None) is not None
|
|
|
|
|
|
def _config_dtype(config: FlashConfig) -> torch.dtype:
|
|
dt = getattr(config, "torch_dtype", None) or getattr(config, "dtype", None)
|
|
if isinstance(dt, torch.dtype):
|
|
return dt
|
|
return getattr(torch, str(dt), None) or torch.bfloat16
|
|
|
|
|
|
class NgramEmbedding(nn.Module):
|
|
"""Token embedding fused with hashed n-gram embeddings.
|
|
|
|
TP-sharded: the ``k*(n-1)`` per-embedder tables are concatenated into one
|
|
:class:`VocabParallelEmbedding` (``oe_embedder``) with per-embedder offsets,
|
|
and the projections are stacked into one ``oe_projection`` applied with a
|
|
single ``bmm``. Hashing math is ported from the HF reference.
|
|
"""
|
|
|
|
def __init__(self, config: FlashConfig, base_embeddings: nn.Module) -> None:
|
|
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.pad_id = config.pad_token_id
|
|
self.eos_token_id = config.eos_token_id
|
|
self._dtype = _config_dtype(config)
|
|
|
|
self._init_ngram_embeddings()
|
|
|
|
def _init_ngram_embeddings(self) -> None:
|
|
self.num_embedders = self.k * (self.n - 1)
|
|
oe_dim = self.config.hidden_size // self.num_embedders
|
|
self.oe_dim = oe_dim
|
|
|
|
# Exclusive prefix sums of per-embedder table sizes; each embedder's
|
|
# local id is offset into the single concatenated table.
|
|
sizes = [int(self.m + i * 2 + 1) for i in range(self.num_embedders)]
|
|
offsets = [0]
|
|
for s in sizes:
|
|
offsets.append(offsets[-1] + s)
|
|
self._offsets = offsets # len num_embedders + 1
|
|
self._sizes = sizes
|
|
|
|
self.oe_embedder = VocabParallelEmbedding(
|
|
offsets[-1], oe_dim, params_dtype=self._dtype
|
|
)
|
|
# Stacked projections: oe_projection[i] = post_projs[i].weight.T
|
|
self.oe_projection = nn.Parameter(
|
|
torch.empty(
|
|
self.num_embedders, oe_dim, self.config.hidden_size, dtype=self._dtype
|
|
),
|
|
requires_grad=False,
|
|
)
|
|
|
|
# Precomputed tables for the CUDA n-gram id kernel (ngram_embedding
|
|
# _kernels.cu): ne_weights[i][j][delta] = vocab^delta mod ne_mods[i][j],
|
|
# ne_mods[i][j] = m + 2*(i*k+j) + 1. Registered as non-persistent buffers
|
|
# so they follow the module to the device (not part of the checkpoint).
|
|
vocab = self.config.vocab_size
|
|
ne_weights = torch.zeros(self.n - 1, self.k, self.n, dtype=torch.int32)
|
|
ne_mods = torch.zeros(self.n - 1, self.k, dtype=torch.int32)
|
|
for i in range(self.n - 1):
|
|
for j in range(self.k):
|
|
mod = int(self.m + 2 * (i * self.k + j) + 1)
|
|
ne_mods[i, j] = mod
|
|
for delta in range(self.n):
|
|
ne_weights[i, j, delta] = pow(vocab, delta, mod)
|
|
self.register_buffer("ne_weights", ne_weights, persistent=False)
|
|
self.register_buffer("ne_mods", ne_mods, persistent=False)
|
|
self.register_buffer(
|
|
"exclusive_sizes",
|
|
torch.tensor(offsets, dtype=torch.int32),
|
|
persistent=False,
|
|
)
|
|
|
|
def load_weight(self, weight_name: str, loaded_weight: torch.Tensor) -> str:
|
|
"""Split a per-embedder checkpoint weight into the sharded layout.
|
|
|
|
Returns the destination parameter's qualified name (relative to the
|
|
enclosing model) so the caller can mark it loaded for completeness
|
|
checks.
|
|
"""
|
|
if "ngram_embeddings.embedders." in weight_name:
|
|
index = int(
|
|
weight_name.split("ngram_embeddings.embedders.")[1].split(".")[0]
|
|
)
|
|
lo, hi = self._offsets[index], self._offsets[index + 1]
|
|
assert hi - lo == loaded_weight.shape[0], (
|
|
f"{hi - lo=} {loaded_weight.shape[0]=}"
|
|
)
|
|
shard = self.oe_embedder.shard_indices
|
|
tp_start, tp_end = shard.org_vocab_start_index, shard.org_vocab_end_index
|
|
load_start, load_end = max(lo, tp_start), min(hi, tp_end)
|
|
if load_start < load_end:
|
|
self.oe_embedder.weight.data[
|
|
load_start - tp_start : load_end - tp_start
|
|
] = loaded_weight[load_start - lo : load_end - lo]
|
|
return "ngram_embeddings.oe_embedder.weight"
|
|
elif "ngram_embeddings.post_projs." in weight_name:
|
|
index = int(
|
|
weight_name.split("ngram_embeddings.post_projs.")[1].split(".")[0]
|
|
)
|
|
self.oe_projection.data[index].copy_(loaded_weight.t())
|
|
return "ngram_embeddings.oe_projection"
|
|
else:
|
|
raise AssertionError(f"Unexpected ngram weight: {weight_name}")
|
|
|
|
def embed_batched(
|
|
self, input_ids: torch.Tensor, oe_ids: torch.Tensor
|
|
) -> torch.Tensor:
|
|
"""Fused n-gram embedding for a flat batch given precomputed ids.
|
|
|
|
Args:
|
|
input_ids: ``[num_tokens]`` current token per position.
|
|
oe_ids: ``[num_tokens, num_embedders]`` global (offset) n-gram ids,
|
|
as produced by the ``ngram_compute_n_gram_ids`` kernel.
|
|
Returns: ``[num_tokens, hidden]``.
|
|
"""
|
|
word = self.word_embeddings(input_ids) # [N, H]
|
|
flat = oe_ids.permute(1, 0).contiguous() # [num_embedders, N]
|
|
oe = self.oe_embedder(flat) # [num_embedders, N, oe_dim]
|
|
proj = torch.bmm(oe, self.oe_projection) # [num_embedders, N, H]
|
|
all_h = torch.cat([word.unsqueeze(0), proj], dim=0) # [ne+1, N, H]
|
|
return all_h.mean(dim=0) # [N, H]
|
|
|
|
|
|
class FlashNgramModel(FlashModel):
|
|
"""FlashModel whose input embedding is an :class:`NgramEmbedding`."""
|
|
|
|
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
|
|
# Each FlashDecoderLayer is a *dual* layer (2 attentions), so the number
|
|
# of decoder layers is ``num_layers``. The ngram HF config sets
|
|
# ``num_hidden_layers`` to a multiple of that (attention-module count),
|
|
# which FlashModel would otherwise build as too many (dead) layers.
|
|
hf = vllm_config.model_config.hf_config
|
|
num_layers = getattr(hf, "num_layers", None)
|
|
if num_layers is not None and hf.num_hidden_layers != num_layers:
|
|
hf.num_hidden_layers = num_layers
|
|
super().__init__(vllm_config=vllm_config, prefix=prefix)
|
|
if get_pp_group().is_first_rank and uses_ngram_embedding(self.config):
|
|
self.ngram_embeddings = NgramEmbedding(self.config, self.embed_tokens)
|
|
else:
|
|
self.ngram_embeddings = None
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
# Names arrive with the ``model.`` prefix already stripped (routed here
|
|
# by AutoWeightsLoader). Split the concatenated/sharded ngram tables and
|
|
# stacked projections; delegate everything else to FlashModel.
|
|
loaded: set[str] = set()
|
|
rest: list[tuple[str, torch.Tensor]] = []
|
|
for name, w in weights:
|
|
if self.ngram_embeddings is not None and (
|
|
"ngram_embeddings.embedders." in name
|
|
or "ngram_embeddings.post_projs." in name
|
|
):
|
|
loaded.add(self.ngram_embeddings.load_weight(name, w))
|
|
else:
|
|
rest.append((name, w))
|
|
loaded |= super().load_weights(rest)
|
|
return loaded
|
|
|
|
|
|
class LongcatFlashNgramForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
|
|
"""LongCat-Flash-Lite for causal LM (MRV2-only, n-gram embedding)."""
|
|
|
|
packed_modules_mapping = {
|
|
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
|
|
"gate_up_proj": ["gate_proj", "up_proj"],
|
|
}
|
|
|
|
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
|
|
super().__init__()
|
|
if not vllm_config.use_v2_model_runner:
|
|
raise NotImplementedError(
|
|
"LongcatFlashNgramForCausalLM (LongCat-Flash-Lite) requires the "
|
|
"V2 model runner for its n-gram embedding state; it is selected "
|
|
"automatically unless VLLM_USE_V2_MODEL_RUNNER=0 is set."
|
|
)
|
|
config = FlashConfig(**vllm_config.model_config.hf_config.__dict__)
|
|
config.intermediate_size = getattr(
|
|
config, "ffn_hidden_size", config.intermediate_size
|
|
)
|
|
self.config = config
|
|
self.quant_config = vllm_config.quant_config
|
|
|
|
self.model = FlashNgramModel(
|
|
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
|
|
)
|
|
if get_pp_group().is_last_rank:
|
|
self.lm_head = ParallelLMHead(
|
|
config.vocab_size,
|
|
config.hidden_size,
|
|
quant_config=self.quant_config,
|
|
prefix=maybe_prefix(prefix, "lm_head"),
|
|
)
|
|
else:
|
|
self.lm_head = PPMissingLayer()
|
|
self.logits_processor = LogitsProcessor(config.vocab_size)
|
|
self.make_empty_intermediate_tensors = (
|
|
self.model.make_empty_intermediate_tensors
|
|
)
|
|
|
|
@staticmethod
|
|
def get_model_state_cls() -> type["LongcatNgramModelState"]:
|
|
return LongcatNgramModelState
|
|
|
|
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
|
|
return self.model.embed_input_ids(input_ids)
|
|
|
|
def forward(
|
|
self,
|
|
input_ids: torch.Tensor | None,
|
|
positions: torch.Tensor,
|
|
intermediate_tensors=None,
|
|
inputs_embeds: torch.Tensor | None = None,
|
|
):
|
|
# inputs_embeds is produced by LongcatNgramModelState.prepare_inputs.
|
|
return self.model(input_ids, positions, intermediate_tensors, inputs_embeds)
|
|
|
|
def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
|
|
return self.logits_processor(self.lm_head, hidden_states)
|
|
|
|
def get_expert_mapping(self):
|
|
return self.model.get_expert_mapping()
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
# AutoWeightsLoader routes ``model.*`` to FlashNgramModel.load_weights
|
|
# (which handles the ngram split) and ``lm_head.*`` to the head. MTP
|
|
# weights are not part of this model.
|
|
loader = AutoWeightsLoader(self, skip_prefixes=["model.mtp."])
|
|
return loader.load_weights(weights)
|
|
|
|
|
|
class LongcatNgramModelState(DefaultModelState):
|
|
"""Per-request n-gram token history for LongCat-Flash-Lite.
|
|
|
|
Maintains a small CPU-side per-slot context (last ``n-1`` processed tokens)
|
|
and a persistent ``inputs_embeds`` buffer. ``prepare_inputs`` computes the
|
|
fused n-gram embedding per request into the buffer, handed to the model
|
|
forward as ``inputs_embeds``.
|
|
"""
|
|
|
|
def __init__(self, vllm_config, model, encoder_cache, device) -> None:
|
|
super().__init__(vllm_config, model, encoder_cache, device)
|
|
config = model.config
|
|
self.ngram = model.model.ngram_embeddings
|
|
self.n = int(config.emb_neighbor_num)
|
|
self.ctx_len = self.n - 1
|
|
self.eos_id = int(config.eos_token_id)
|
|
|
|
# Per-slot left-context: last ``n-1`` processed tokens, EOS negated. A
|
|
# negative entry (incl. the -1 fill) marks a context boundary that stops
|
|
# the n-gram walk (matches the kernel's EOS break / fresh-request start).
|
|
self.token_context = torch.full(
|
|
(self.max_num_reqs, self.ctx_len), -1, dtype=torch.int32, device=device
|
|
)
|
|
|
|
self._inputs_embeds_buf = torch.zeros(
|
|
self.max_num_tokens,
|
|
config.hidden_size,
|
|
dtype=self.dtype,
|
|
device=device,
|
|
)
|
|
|
|
def _neg_eos(self, toks: list[int]) -> list[int]:
|
|
return [-t if t == self.eos_id else t for t in toks]
|
|
|
|
def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
|
|
super().add_request(req_index, new_req_data) # rope positions
|
|
# Fresh request -> no left-context (-1 fill). On resume, seed from the
|
|
# already-processed token tail. Use prefill_token_ids (full processed
|
|
# sequence incl. generated tokens on v2 resume), like DefaultModelState;
|
|
# the prompt alone would be too short when resuming after decode.
|
|
ctx = [-1] * self.ctx_len
|
|
ncomp = new_req_data.num_computed_tokens
|
|
toks_src = new_req_data.prefill_token_ids or new_req_data.prompt_token_ids
|
|
if ncomp > 0 and toks_src is not None:
|
|
lo = max(0, ncomp - self.ctx_len)
|
|
toks = self._neg_eos(list(toks_src[lo:ncomp]))
|
|
ctx[self.ctx_len - len(toks) :] = toks
|
|
self.token_context[req_index] = torch.tensor(
|
|
ctx, dtype=torch.int32, device=self.token_context.device
|
|
)
|
|
|
|
def prepare_inputs(
|
|
self, input_batch: InputBatch, req_states: RequestState
|
|
) -> dict[str, Any]:
|
|
model_inputs = super().prepare_inputs(input_batch, req_states) # positions
|
|
num_tokens = input_batch.num_tokens
|
|
num_padded = input_batch.num_tokens_after_padding
|
|
input_ids = input_batch.input_ids[:num_tokens]
|
|
embeds = self._inputs_embeds_buf[:num_padded]
|
|
|
|
oe_ids = self._compute_oe_ids(input_batch)
|
|
embeds[:num_tokens].copy_(self.ngram.embed_batched(input_ids, oe_ids))
|
|
model_inputs["inputs_embeds"] = embeds
|
|
return model_inputs
|
|
|
|
def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]:
|
|
# FULL cudagraph replay reads only the captured buffers, so capture must
|
|
# reference the same persistent ``inputs_embeds`` buffer prepare_inputs
|
|
# re-fills (the base class wires this for multimodal models only).
|
|
model_inputs = super().prepare_dummy_inputs(num_reqs, num_tokens) # positions
|
|
model_inputs["inputs_embeds"] = self._inputs_embeds_buf[:num_tokens]
|
|
return model_inputs
|
|
|
|
def _compute_oe_ids(self, input_batch: InputBatch) -> torch.Tensor:
|
|
"""Batched global n-gram ids ``[num_tokens, num_embedders]``.
|
|
|
|
Assembles an ephemeral per-request token table (``[n-1] context ++
|
|
current tokens``, EOS-negated) and runs the ``ngram_compute_n_gram_ids``
|
|
CUDA kernel for the whole batch, then rolls each slot's context forward.
|
|
"""
|
|
device = self.token_context.device
|
|
num_tokens = input_batch.num_tokens
|
|
num_reqs = input_batch.num_reqs
|
|
ctx_len = self.ctx_len
|
|
idx_mapping = input_batch.idx_mapping[:num_reqs].long()
|
|
qsl = input_batch.query_start_loc[: num_reqs + 1].to(torch.int32)
|
|
cur = input_batch.input_ids[:num_tokens].to(torch.int32)
|
|
|
|
cur_neg = torch.where(cur == self.eos_id, -cur, cur)
|
|
req_lens = qsl[1:] - qsl[:-1]
|
|
max_len = int(req_lens.max().item())
|
|
width = ctx_len + max_len
|
|
|
|
# table[r] = [context(n-1) | current tokens | pad(-1)]
|
|
table = torch.full((num_reqs, width), -1, dtype=torch.int32, device=device)
|
|
table[:, :ctx_len] = self.token_context[idx_mapping]
|
|
tok_req = torch.repeat_interleave(
|
|
torch.arange(num_reqs, device=device), req_lens.long()
|
|
)
|
|
col = ctx_len + (
|
|
torch.arange(num_tokens, device=device) - qsl[:-1].long()[tok_req]
|
|
)
|
|
table[tok_req, col] = cur_neg
|
|
|
|
column_starts = torch.full(
|
|
(num_reqs,), ctx_len, dtype=torch.int32, device=device
|
|
)
|
|
row_indices = torch.arange(num_reqs, dtype=torch.int64, device=device)
|
|
n_gram_ids = torch.empty(
|
|
num_tokens, self.ngram.num_embedders, dtype=torch.int32, device=device
|
|
)
|
|
ops.ngram_compute_n_gram_ids(
|
|
self.n,
|
|
self.ngram.k,
|
|
self.ngram.ne_weights,
|
|
self.ngram.ne_mods,
|
|
self.ngram.exclusive_sizes,
|
|
qsl,
|
|
table,
|
|
row_indices,
|
|
column_starts,
|
|
n_gram_ids,
|
|
)
|
|
|
|
# Roll context: new context = last n-1 of [context | current] per slot.
|
|
gather = req_lens.long().unsqueeze(1) + torch.arange(
|
|
ctx_len, device=device
|
|
).unsqueeze(0)
|
|
rows = torch.arange(num_reqs, device=device).unsqueeze(1)
|
|
self.token_context[idx_mapping] = table[rows, gather]
|
|
return n_gram_ids.long()
|