Compare commits

...
1 Commits
Author SHA1 Message Date
mgoinandClaude Fable 5 62a0750892 [Model] Add LongCat-2.0 (sparse attention, MTP, n-gram fixes)
Support meituan-longcat/LongCat-2.0-FP8: the LongCat-Flash-Lite backbone
plus a DeepSeek-V3.2-style sparse attention indexer, reusing vLLM's
existing DSA infrastructure.

- Resolve the checkpoint's null model_type via a registered config class
  (arch LongcatCausalLM); alias the oe_* n-gram config fields.
- Wire the indexer into the dual-attention layer: only the first
  attention of each layer computes top-k indices (cli_factor), the
  second reuses them via the shared buffer. Adds a generic skip_topk
  override to DeepseekV2MLAAttention; the schedule lives in the model.
- Streaming-aware indexing: force the first index_init_tokens and last
  index_local_tokens into the top-k set (sparse_attn_indexer).
- MTP speculative decoding (one module iterated up to 3 draft steps,
  plain token embedding, FP8 indexer weights).
- Fix a latent LongCat weight-loading bug: the post-load
  mla_scale_q_lora/kv_lora fold breaks under incremental load_weights
  calls; fold at weight-load time instead.
- Fix two n-gram embedding bugs (also affect LongCat-Flash-Lite):
  rebuild the left-context from the accepted-token history so rejected
  draft tokens don't pollute it, and hash EOS-current positions with
  full look-back per the HF reference (unit test included).
- Add --reasoning-parser longcat (<longcat_think> tags) and docs rows.
- Add longcat model types to the DeepGEMM Blackwell blocklist (non-ue8m0
  scales; measured GSM8K-neutral on vLLM but matches SGLang's guard).

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: mgoin <mgoin64@gmail.com>
2026-07-11 02:12:49 +00:00
18 changed files with 625 additions and 104 deletions
+2
View File
@@ -405,6 +405,8 @@ th {
| `Lfm2MoeForCausalLM` | LFM2MoE | `LiquidAI/LFM2-8B-A1B-preview`, etc. | ✅︎ | ✅︎ |
| `LlamaForCausalLM` | Llama 3.1, Llama 3, Llama 2, LLaMA, Yi | `meta-llama/Meta-Llama-3.1-405B-Instruct`, `meta-llama/Meta-Llama-3.1-70B`, `meta-llama/Meta-Llama-3-70B-Instruct`, `meta-llama/Llama-2-70b-hf`, `01-ai/Yi-34B`, etc. | ✅︎ | ✅︎ |
| `LongcatFlashForCausalLM` | LongCat-Flash | `meituan-longcat/LongCat-Flash-Chat`, `meituan-longcat/LongCat-Flash-Chat-FP8` | ✅︎ | ✅︎ |
| `LongcatFlashNgramForCausalLM` | LongCat-Flash-Lite | `meituan-longcat/LongCat-Flash-Lite` | ✅︎ | ✅︎ |
| `LongcatCausalLM` | LongCat-2.0 | `meituan-longcat/LongCat-2.0-FP8` | ✅︎ | ✅︎ |
| `MambaForCausalLM` | Mamba | `state-spaces/mamba-130m-hf`, `state-spaces/mamba-790m-hf`, `state-spaces/mamba-2.8b-hf`, etc. | | ✅︎ |
| `Mamba2ForCausalLM` | Mamba2 | `mistralai/Mamba-Codestral-7B-v0.1`, etc. | | ✅︎ |
| `MellumForCausalLM` | Mellum 2 | `JetBrains/Mellum2-12B-A2.5B-Base`, etc. | | ✅︎ |
+117
View File
@@ -0,0 +1,117 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""LongCat n-gram embedding id computation vs a pure-Python reference.
Guards the hash-id semantics of ``ngram_compute_n_gram_ids`` plus the
EOS-position fixup (``compute_eos_position_ngram_ids``): an EOS *current*
token hashes with its full look-back, while later positions' look-back stops
at the EOS boundary (LongCat reference behavior).
"""
from types import SimpleNamespace
import pytest
import torch
from vllm import _custom_ops as ops
from vllm.model_executor.models.longcat_flash_ngram import (
compute_eos_position_ngram_ids,
)
from vllm.platforms import current_platform
VOCAB = 163840
N, K = 5, 4 # oe_neighbor_num, oe_split_num (LongCat-2.0)
M = int(100.567 * VOCAB)
EOS = 2
NUM_EMB = K * (N - 1)
SIZES = [M + c * 2 + 1 for c in range(NUM_EMB)]
OFFSETS = [0]
for s in SIZES:
OFFSETS.append(OFFSETS[-1] + s)
def _ref_ids(tokens: list[int]) -> list[list[int]]:
"""Reference n-gram ids on raw tokens (HF LongCat semantics)."""
out = []
for pos in range(len(tokens)):
row = []
for i in range(N - 1):
n_real = i + 2
for j in range(K):
cfg = i * K + j
mod = SIZES[cfg]
h = 0
for delta in range(n_real):
p = pos - delta
if p < 0:
break
t = tokens[p]
if delta > 0 and t == EOS:
break # look-back stops at an EOS boundary
h += (t * pow(VOCAB, delta, mod)) % mod
row.append(h % mod + OFFSETS[cfg])
out.append(row)
return out
@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA kernel")
def test_ngram_ids_match_reference_with_eos():
device = "cuda"
torch.manual_seed(0)
tokens = torch.randint(3, VOCAB, (14,), dtype=torch.int32).tolist()
tokens[5] = EOS
tokens[6] = EOS # double EOS: second one's look-back stops at the first
tokens[13] = EOS # trailing EOS (chat-template turn boundary)
ctx_len = N - 1
toks_neg = [-t if t == EOS else t for t in tokens]
width = ctx_len + len(tokens)
table = torch.full((1, width), -1, dtype=torch.int32, device=device)
table[0, ctx_len:] = torch.tensor(toks_neg, dtype=torch.int32, device=device)
ne_weights = torch.zeros(N - 1, K, N, dtype=torch.int32)
ne_mods = torch.zeros(N - 1, K, dtype=torch.int32)
for i in range(N - 1):
for j in range(K):
mod = SIZES[i * K + j]
ne_mods[i, j] = mod
for delta in range(N):
ne_weights[i, j, delta] = pow(VOCAB, delta, mod)
ngram = SimpleNamespace(
n=N,
k=K,
num_embedders=NUM_EMB,
ne_weights=ne_weights.to(device),
ne_mods=ne_mods.to(device),
exclusive_sizes=torch.tensor(OFFSETS, dtype=torch.int32, device=device),
)
T = len(tokens)
qsl = torch.tensor([0, T], dtype=torch.int32, device=device)
row_indices = torch.zeros(1, dtype=torch.int64, device=device)
column_starts = torch.full((1,), ctx_len, dtype=torch.int32, device=device)
got = torch.empty(T, NUM_EMB, dtype=torch.int32, device=device)
ops.ngram_compute_n_gram_ids(
N,
K,
ngram.ne_weights,
ngram.ne_mods,
ngram.exclusive_sizes,
qsl,
table,
row_indices,
column_starts,
got,
)
cur = torch.tensor(tokens, dtype=torch.int32, device=device)
tok_req = torch.zeros(T, dtype=torch.int64, device=device)
col = ctx_len + torch.arange(T, device=device)
eos_tok = (cur == EOS).nonzero(as_tuple=True)[0]
got[eos_tok] = compute_eos_position_ngram_ids(
ngram, EOS, table, tok_req, col, eos_tok
)
want = torch.tensor(_ref_ids(tokens), dtype=torch.int32)
torch.testing.assert_close(got.cpu(), want, rtol=0, atol=0)
+6
View File
@@ -389,6 +389,12 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
# so the dummy-weight init test fits in CI memory.
hf_overrides={"ngram_vocab_size_ratio": 1},
),
"LongcatCausalLM": _HfExamplesInfo(
"meituan-longcat/LongCat-2.0-FP8",
# Shrink the huge n-gram tables (~264M rows at the checkpoint's
# oe_vocab_size_ratio=100.567) so dummy-weight init fits in CI memory.
hf_overrides={"ngram_vocab_size_ratio": 1},
),
"MambaForCausalLM": _HfExamplesInfo("state-spaces/mamba-130m-hf"),
"Mamba2ForCausalLM": _HfExamplesInfo(
"mistralai/Mamba-Codestral-7B-v0.1",
+7 -1
View File
@@ -520,7 +520,13 @@ class SpeculativeConfig:
)
if hf_config.model_type in ("longcat_flash", "longcat_flash_ngram"):
hf_config.model_type = "longcat_flash_mtp"
n_predict = getattr(hf_config, "num_nextn_predict_layers", 1)
# LongCat-2.0 ships one MTP module applied for up to
# mtp_num_layers draft steps (mtp_replicate_modules).
n_predict = (
getattr(hf_config, "num_nextn_predict_layers", None)
or getattr(hf_config, "mtp_num_layers", None)
or 1
)
hf_config.update(
{"n_predict": n_predict, "architectures": ["LongCatFlashMTPModel"]}
)
+1
View File
@@ -71,6 +71,7 @@ DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset(
"Qwen2MoeForCausalLM",
"GraniteMoeForCausalLM",
"LongcatFlashNgramForCausalLM",
"LongcatCausalLM",
}
)
@@ -250,6 +250,40 @@ def fused_indexer_q_rope_quant(
return q_fp8, weights_out
def _mask_init_and_local_tokens(
logits: torch.Tensor,
row_starts: torch.Tensor | None,
row_ends: torch.Tensor,
num_init_tokens: int,
num_local_tokens: int,
) -> None:
"""Force streaming tokens into the top-k set (streaming-aware indexing):
scatter +inf into the first ``num_init_tokens`` and last
``num_local_tokens`` columns of each row's valid range
``[row_starts, row_ends)``. Out-of-range writes are clamped into
``[row_starts, row_ends)`` so they never leak into other rows' ranges.
"""
device = logits.device
ends = row_ends.to(device=device, dtype=torch.int64).reshape(-1)
if row_starts is None:
starts = torch.zeros_like(ends)
else:
starts = row_starts.to(device=device, dtype=torch.int64).reshape(-1)
last = (ends - 1).clamp_max_(logits.shape[1] - 1).clamp_min_(starts)
if num_init_tokens > 0:
init_idx = starts[:, None] + torch.arange(
num_init_tokens, dtype=torch.int64, device=device
)
init_idx = torch.minimum(init_idx, last[:, None])
logits.scatter_(1, init_idx, float("inf"))
if num_local_tokens > 0:
local_idx = last[:, None] - torch.arange(
num_local_tokens, dtype=torch.int64, device=device
)
local_idx = torch.maximum(local_idx, starts[:, None])
logits.scatter_(1, local_idx, float("inf"))
def _gather_workspace_shapes(
total_seq_lens: int,
head_dim: int,
@@ -313,6 +347,8 @@ def sparse_attn_indexer(
dcp_world_size: int = 1,
cp_kv_cache_interleave_size: int = 1,
skip_topk_buffer_clear: bool = False,
num_init_tokens: int = 0,
num_local_tokens: int = 0,
) -> torch.Tensor:
# careful! this will be None in dummy run
attn_metadata = get_forward_context().attn_metadata
@@ -471,6 +507,14 @@ def sparse_attn_indexer(
cu_seqlen_ke,
clean_logits=False,
)
if num_init_tokens > 0 or num_local_tokens > 0:
_mask_init_and_local_tokens(
logits,
cu_seqlen_ks,
cu_seqlen_ke,
num_init_tokens,
num_local_tokens,
)
num_rows = logits.shape[0]
ops.top_k_per_row_prefill(
logits,
@@ -572,6 +616,17 @@ def sparse_attn_indexer(
num_rows = logits.shape[0]
topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
if num_init_tokens > 0 or num_local_tokens > 0:
# seq_lens is (B, next_n) whenever next_n > 1, so flattening
# yields the per-row context length.
_mask_init_and_local_tokens(
logits,
None,
seq_lens.reshape(-1)[:num_rows],
num_init_tokens,
num_local_tokens,
)
use_cooperative_topk = (
current_platform.is_cuda()
and topk_tokens in (512, 1024, 2048)
@@ -668,6 +723,8 @@ def sparse_attn_indexer_fake(
dcp_world_size: int = 1,
cp_kv_cache_interleave_size: int = 1,
skip_topk_buffer_clear: bool = False,
num_init_tokens: int = 0,
num_local_tokens: int = 0,
) -> torch.Tensor:
return topk_indices_buffer
@@ -706,6 +763,8 @@ class SparseAttnIndexer(CustomOp):
topk_indices_buffer: torch.Tensor,
skip_k_cache_insert: bool = False,
use_fp4_cache: bool = False,
num_init_tokens: int = 0,
num_local_tokens: int = 0,
):
super().__init__()
self.k_cache = k_cache
@@ -718,6 +777,8 @@ class SparseAttnIndexer(CustomOp):
self.topk_indices_buffer = topk_indices_buffer
self.skip_k_cache_insert = skip_k_cache_insert
self.use_fp4_cache = use_fp4_cache
self.num_init_tokens = num_init_tokens
self.num_local_tokens = num_local_tokens
# DCP scalars are constant for the run; resolve them here (config is set
# during model construction) and pass them into the custom op, rather
# than threading them through per-step metadata.
@@ -781,6 +842,8 @@ class SparseAttnIndexer(CustomOp):
self.dcp_rank,
self.dcp_world_size,
self.cp_kv_cache_interleave_size,
num_init_tokens=self.num_init_tokens,
num_local_tokens=self.num_local_tokens,
)
def forward_xpu(
@@ -803,6 +866,10 @@ class SparseAttnIndexer(CustomOp):
assert isinstance(q_quant, torch.Tensor), (
"AMD sparse_attn_indexer expects a single FP8 q_quant tensor"
)
assert self.num_init_tokens == 0 and self.num_local_tokens == 0, (
"Streaming-aware indexing (index_init_tokens/index_local_tokens) is "
"not supported on the ROCm sparse_attn_indexer path yet"
)
if rocm_aiter_ops.is_enabled():
return torch.ops.vllm.rocm_aiter_sparse_attn_indexer(
hidden_states,
+10
View File
@@ -825,6 +825,15 @@ class LongcatFlashNgramForCausalLMConfig(VerifyAndUpdateConfig):
if compilation_config.cudagraph_mode is None:
compilation_config.cudagraph_mode = CUDAGraphMode.FULL
# LongCat-2.0 sparse attention (DSA indexer) requires the same
# kv-cache dtype normalization as DeepSeek-V3.2.
hf_config = vllm_config.model_config.hf_config
if hasattr(hf_config, "index_topk"):
cache_config = vllm_config.cache_config
if cache_config.cache_dtype == "bfloat16":
cache_config.cache_dtype = "auto"
logger.info("Using bfloat16 kv-cache for LongCat sparse attention")
MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
"ColBERTJinaRobertaModel": JinaRobertaModelConfig,
@@ -840,6 +849,7 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
"Gemma4UnifiedForConditionalGeneration": Gemma4Config,
"GptOssForCausalLM": GptOssForCausalLMConfig,
"LongcatFlashNgramForCausalLM": LongcatFlashNgramForCausalLMConfig,
"LongcatCausalLM": LongcatFlashNgramForCausalLMConfig,
"GteModel": SnowflakeGteNewModelConfig,
"GteNewForSequenceClassification": GteNewModelConfig,
"GteNewModel": GteNewModelConfig,
+34 -17
View File
@@ -710,6 +710,8 @@ class Indexer(nn.Module):
self.max_model_len,
self.max_total_seq_len,
self.topk_indices_buffer,
num_init_tokens=getattr(config, "index_init_tokens", 0),
num_local_tokens=getattr(config, "index_local_tokens", 0),
)
self.is_inplace_rope = is_inplace_rope
@@ -820,45 +822,56 @@ def _try_load_fp8_indexer_wk(
name, tensor, buf, params_dict, loaded_params, pp_missing_layer_names
):
"""
We fuse the WK and weights_proj projections, but in some checkpoints WK is stored
in FP8 with a separate weight_scale_inv, while weights_proj is stored in BF16.
Upcasting to BF16 during loading enables the fusion. This function loads the FP8 WK
weights and scale, and when both are available, dequantizes to BF16 and stores into
the fused wk_weights_proj.weight parameter.
We fuse the WK and weights_proj projections, but in some checkpoints one
or both are stored in FP8 with a separate weight_scale_inv while the fused
parameter is BF16. Upcasting to BF16 during loading enables the fusion.
This function buffers the FP8 weight and scale, and when both are
available, dequantizes to BF16 and stores into the corresponding shard of
the fused wk_weights_proj.weight.
"""
if "indexer.wk." not in name or "wk_weights" in name:
return False # Weight is not an isolated WK weight for the indexer, ignore.
if "wk_weights" in name:
return False # Already-fused parameter name, ignore.
if "indexer.wk." in name:
sub_name, shard_id = ".wk.", 0
elif "indexer.weights_proj." in name:
sub_name, shard_id = ".weights_proj.", 1
else:
return False # Not an isolated indexer projection weight, ignore.
is_weight = name.endswith(".weight") and tensor.dtype == torch.float8_e4m3fn
is_scale = "weight_scale" in name
if not is_weight and not is_scale:
return False # WK is not in FP8 format, ignore.
return False # Projection is not in FP8 format, ignore.
# Buffer this tensor (weight or scale) until both have arrived.
layer_prefix = name.rsplit(".wk.", 1)[0] # e.g. "model.layers.0.self_attn.indexer"
# layer_prefix is e.g. "model.layers.0.self_attn.indexer"
layer_prefix = name.rsplit(sub_name, 1)[0]
fused_name = f"{layer_prefix}.wk_weights_proj.weight"
if any(
name.startswith(missing_layer_name)
for missing_layer_name in pp_missing_layer_names
):
return True
entry = buf.setdefault(layer_prefix, {})
entry = buf.setdefault((layer_prefix, shard_id), {})
entry["weight" if is_weight else "scale"] = tensor
if "weight" not in entry or "scale" not in entry:
return True # still waiting for the other param
# We have both weight and scale: dequantize FP8 to BF16.
# We have both weight and scale: dequantize FP8 to BF16. Derive the block
# shape per axis: narrow projections (e.g. a 32-row weights_proj under
# 128x128 quantization) span a partial row block.
weight_fp8, scale_inv = entry["weight"], entry["scale"]
del buf[layer_prefix]
block_size = weight_fp8.shape[1] // scale_inv.shape[1]
del buf[(layer_prefix, shard_id)]
row_block = weight_fp8.shape[0] // scale_inv.shape[0]
col_block = weight_fp8.shape[1] // scale_inv.shape[1]
weight_bf16 = scaled_dequantize(
weight_fp8,
scale_inv,
group_shape=GroupShape(block_size, block_size),
group_shape=GroupShape(row_block, col_block),
out_dtype=torch.bfloat16,
)
# Load the dequantized weight into shard 0 of the fused buffer.
# Load the dequantized weight into its shard of the fused buffer.
param = params_dict[fused_name]
param.weight_loader(param, weight_bf16, 0)
param.weight_loader(param, weight_bf16, shard_id)
loaded_params.add(fused_name)
return True
@@ -975,6 +988,7 @@ class DeepseekV2MLAAttention(nn.Module):
topk_indices_buffer: torch.Tensor | None = None,
input_size: int | None = None,
reduce_results: bool = True,
skip_topk: bool | None = None,
) -> None:
super().__init__()
self.hidden_size = hidden_size
@@ -1077,7 +1091,10 @@ class DeepseekV2MLAAttention(nn.Module):
# Refer: https://arxiv.org/abs/2603.12201 for more details.
_skip_topk = False
is_mtp_layer = False
if self.is_v32:
if self.is_v32 and skip_topk is not None:
# The caller decides the skip schedule directly.
_skip_topk = skip_topk
elif self.is_v32:
_index_topk_freq = getattr(config, "index_topk_freq", 1)
_index_topk_pattern = getattr(config, "index_topk_pattern", None)
_index_skip_topk_offset = getattr(config, "index_skip_topk_offset", 2)
+77 -18
View File
@@ -64,13 +64,18 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
VocabParallelEmbedding,
)
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm.model_executor.models.deepseek_v2 import DeepseekV2MLAAttention
from vllm.model_executor.models.deepseek_v2 import (
DeepseekV2MLAAttention,
_try_load_fp8_indexer_wk,
)
from vllm.platforms import current_platform
from vllm.sequence import IntermediateTensors
from .interfaces import SupportsLoRA, SupportsPP
from .utils import (
AutoWeightsLoader,
PPMissingLayer,
get_pp_missing_layer_names,
is_pp_missing_parameter,
make_empty_intermediate_tensors_factory,
make_layers,
@@ -348,6 +353,18 @@ class LongcatMoe(nn.Module):
return final_hidden_states.view(num_tokens, hidden_dim)
def maybe_replace_indexer_k_norm(
attn: DeepseekV2MLAAttention, config: PretrainedConfig
) -> None:
"""LongCat's DSA indexer normalizes K with RMSNorm (index_k_norm_type),
where the shared Indexer defaults to DeepSeek-V3.2's LayerNorm."""
if (
getattr(attn, "indexer", None) is not None
and getattr(config, "index_k_norm_type", None) == "rms"
):
attn.indexer.k_norm = RMSNorm(config.index_head_dim, eps=1e-6)
class FlashDecoderLayer(nn.Module):
"""Flash decoder layer with dual attention and MLP structure."""
@@ -359,12 +376,18 @@ class FlashDecoderLayer(nn.Module):
quant_config: QuantizationConfig | None = None,
prefix: str = "",
enable_eplb: bool = False,
topk_indices_buffer: torch.Tensor | None = None,
) -> None:
super().__init__()
self.layer_idx = int(prefix.split(sep=".")[-1])
self.hidden_size = config.hidden_size
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
# Cross-Layer Indexing: with cli_factor=2 the first attention of each
# layer computes DSA top-k indices and the second reuses them via the
# shared buffer (attention sublayer id = 2 * layer_idx + i).
cli_factor = getattr(config, "cli_factor", 1) or 1
# Dual attention structure
self.self_attn = nn.ModuleList(
[
@@ -386,10 +409,14 @@ class FlashDecoderLayer(nn.Module):
if "self_attn" in getattr(config, "disable_quant_module", [])
else quant_config,
prefix=f"{prefix}.self_attn.{i}",
topk_indices_buffer=topk_indices_buffer,
skip_topk=(2 * self.layer_idx + i) % cli_factor != 0,
)
for i in range(2)
]
)
for attn in self.self_attn:
maybe_replace_indexer_k_norm(attn, config)
self.input_layernorm = nn.ModuleList(
[RMSNorm(config.hidden_size, eps=config.rms_norm_eps) for i in range(2)]
)
@@ -490,6 +517,17 @@ class FlashModel(nn.Module):
self.vocab_size = config.vocab_size
self.is_v32 = hasattr(config, "index_topk")
if self.is_v32:
topk_indices_buffer = torch.empty(
vllm_config.scheduler_config.max_num_batched_tokens,
config.index_topk,
dtype=torch.int32,
device=current_platform.device_type,
)
else:
topk_indices_buffer = None
if get_pp_group().is_first_rank:
self.embed_tokens = VocabParallelEmbedding(
config.vocab_size,
@@ -506,6 +544,7 @@ class FlashModel(nn.Module):
cache_config=cache_config,
quant_config=quant_config,
prefix=prefix,
topk_indices_buffer=topk_indices_buffer,
),
prefix=f"{prefix}.layers",
)
@@ -572,15 +611,37 @@ class FlashModel(nn.Module):
("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
(".gate_up_proj", ".gate_proj", 0),
(".gate_up_proj", ".up_proj", 1),
# Fused indexer wk + weights_proj (shard 0 = wk, 1 = weights_proj)
("wk_weights_proj", "wk", 0),
("wk_weights_proj", "weights_proj", 1),
]
expert_params_mapping = self.get_expert_mapping()
loaded_params: set[str] = set()
pp_missing_layer_names = get_pp_missing_layer_names(self)
params_dict = dict(self.named_parameters())
_pending_wk_fp8: dict = {}
# Drop checkpoint indexer weights for sublayers without an indexer.
indexer_present_prefixes = {
n.rsplit(".indexer.", 1)[0] for n in params_dict if ".indexer." in n
}
for name, loaded_weight in weights:
if "rotary_emb.inv_freq" in name:
continue
if ".indexer." in name and (
name.rsplit(".indexer.", 1)[0] not in indexer_present_prefixes
):
continue
if _try_load_fp8_indexer_wk(
name,
loaded_weight,
_pending_wk_fp8,
params_dict,
loaded_params,
pp_missing_layer_names,
):
continue
for param_name, weight_name, shard_id in stacked_params_mapping:
if weight_name not in name:
continue
@@ -653,6 +714,21 @@ class FlashModel(nn.Module):
continue
if is_pp_missing_parameter(name, self):
continue
# Fold the MLA LoRA scaling at load time: load_weights may
# run incrementally, so a post-load fold can miss weights
# arriving in a later call.
if name.endswith(".q_a_layernorm.weight") and getattr(
self.config, "mla_scale_q_lora", False
):
loaded_weight = loaded_weight.float() * (
(self.config.hidden_size / self.config.q_lora_rank) ** 0.5
)
elif name.endswith(".kv_a_layernorm.weight") and getattr(
self.config, "mla_scale_kv_lora", False
):
loaded_weight = loaded_weight.float() * (
(self.config.hidden_size / self.config.kv_lora_rank) ** 0.5
)
param = params_dict[name]
weight_loader = getattr(
param, "weight_loader", default_weight_loader
@@ -687,23 +763,6 @@ class FlashModel(nn.Module):
).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
+60 -19
View File
@@ -20,10 +20,14 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
VocabParallelEmbedding,
)
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm.model_executor.models.longcat_flash import FlashConfig
from vllm.model_executor.models.longcat_flash import (
FlashConfig,
maybe_replace_indexer_k_norm,
)
from vllm.platforms import current_platform
from vllm.sequence import IntermediateTensors
from .deepseek_v2 import DeepseekV2DecoderLayer
from .deepseek_v2 import DeepseekV2DecoderLayer, _try_load_fp8_indexer_wk
from .utils import maybe_prefix
@@ -34,6 +38,7 @@ class LongCatMultiTokenPredictorLayer(nn.Module):
prefix: str,
vllm_config: VllmConfig,
quant_config: QuantizationConfig | None = None,
topk_indices_buffer: torch.Tensor | None = None,
) -> None:
super().__init__()
self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -45,7 +50,10 @@ class LongCatMultiTokenPredictorLayer(nn.Module):
quant_config=quant_config,
prefix="eh_proj",
)
self.mtp_block = DeepseekV2DecoderLayer(vllm_config, prefix)
self.mtp_block = DeepseekV2DecoderLayer(
vllm_config, prefix, topk_indices_buffer=topk_indices_buffer
)
maybe_replace_indexer_k_norm(self.mtp_block.self_attn, config)
self.final_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
def forward(
@@ -84,6 +92,15 @@ class LongCatMultiTokenPredictor(nn.Module):
vllm_config.model_config.hf_config.intermediate_size = config.intermediate_size
self.mtp_start_layer_idx = config.num_hidden_layers * 2
self.num_mtp_layers = 1
if hasattr(config, "index_topk"):
topk_indices_buffer = torch.empty(
vllm_config.scheduler_config.max_num_batched_tokens,
config.index_topk,
dtype=torch.int32,
device=current_platform.device_type,
)
else:
topk_indices_buffer = None
self.layers = torch.nn.ModuleDict(
{
str(idx): LongCatMultiTokenPredictorLayer(
@@ -91,6 +108,7 @@ class LongCatMultiTokenPredictor(nn.Module):
prefix=f"{prefix}.layers.{idx}",
vllm_config=vllm_config,
quant_config=quant_config,
topk_indices_buffer=topk_indices_buffer,
)
for idx in range(
self.mtp_start_layer_idx,
@@ -178,6 +196,9 @@ class LongCatFlashMTP(nn.Module):
("gate_up_proj", "up_proj", 1),
("fused_qkv_a_proj", "q_a_proj", 0),
("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
# Fused indexer wk + weights_proj (shard 0 = wk, 1 = weights_proj)
("wk_weights_proj", "wk", 0),
("wk_weights_proj", "weights_proj", 1),
]
new_to_old_names_mapping = {
@@ -188,6 +209,13 @@ class LongCatFlashMTP(nn.Module):
"model.mtp.layers.0.hnorm.m.weight": "hnorm.weight",
"model.mtp.layers.0.input_layernorm.weight": "model.layers.0.input_layernorm.weight", # noqa: E501
"model.mtp.layers.0.post_attention_layernorm.weight": "model.layers.0.post_attention_layernorm.weight", # noqa: E501
"model.mtp.layers.0.self_attn.indexer.k_norm.weight": "model.layers.0.self_attn.indexer.k_norm.weight", # noqa: E501
"model.mtp.layers.0.self_attn.indexer.wq_b.weight": "model.layers.0.self_attn.indexer.wq_b.weight", # noqa: E501
"model.mtp.layers.0.self_attn.indexer.wq_b.weight_scale_inv": "model.layers.0.self_attn.indexer.wq_b.weight_scale_inv", # noqa: E501
"model.mtp.layers.0.self_attn.indexer.wk.weight": "model.layers.0.self_attn.indexer.wk.weight", # noqa: E501
"model.mtp.layers.0.self_attn.indexer.wk.weight_scale_inv": "model.layers.0.self_attn.indexer.wk.weight_scale_inv", # noqa: E501
"model.mtp.layers.0.self_attn.indexer.weights_proj.weight": "model.layers.0.self_attn.indexer.weights_proj.weight", # noqa: E501
"model.mtp.layers.0.self_attn.indexer.weights_proj.weight_scale_inv": "model.layers.0.self_attn.indexer.weights_proj.weight_scale_inv", # noqa: E501
"model.mtp.layers.0.self_attn.kv_a_layernorm.weight": "model.layers.0.self_attn.kv_a_layernorm.weight", # noqa: E501
"model.mtp.layers.0.self_attn.kv_a_proj_with_mqa.weight": "model.layers.0.self_attn.kv_a_proj_with_mqa.weight", # noqa: E501
"model.mtp.layers.0.self_attn.kv_a_proj_with_mqa.weight_scale_inv": "model.layers.0.self_attn.kv_a_proj_with_mqa.weight_scale_inv", # noqa: E501
@@ -211,15 +239,29 @@ class LongCatFlashMTP(nn.Module):
params_dict = dict(self.named_parameters())
loaded_params: set[str] = set()
_pending_wk_fp8: dict = {}
for name, loaded_weight in weights:
if "rotary_emb.inv_freq" in name:
continue
# MTP embeds plain tokens (mtp_disable_over_tokenizer); its
# checkpoint n-gram tables are unused.
if "ngram_embeddings" in name:
continue
spec_layer = self.get_spec_layer_idx_from_weight_name(self.config, name)
if spec_layer is None:
continue
name = self._rewrite_spec_layer_name(
spec_layer, name, new_to_old_names_mapping
)
if _try_load_fp8_indexer_wk(
name,
loaded_weight,
_pending_wk_fp8,
params_dict,
loaded_params,
[],
):
continue
for param_name, weight_name, shard_id in stacked_params_mapping:
# Skip non-stacked layers and experts (experts handled below).
if weight_name not in name:
@@ -260,6 +302,21 @@ class LongCatFlashMTP(nn.Module):
):
continue
# Fold the MLA LoRA scaling at load time (see
# FlashModel.load_weights).
if name.endswith(".q_a_layernorm.weight") and getattr(
self.config, "mla_scale_q_lora", False
):
loaded_weight = loaded_weight.float() * (
(self.config.hidden_size / self.config.q_lora_rank) ** 0.5
)
elif name.endswith(".kv_a_layernorm.weight") and getattr(
self.config, "mla_scale_kv_lora", False
):
loaded_weight = loaded_weight.float() * (
(self.config.hidden_size / self.config.kv_lora_rank) ** 0.5
)
param = params_dict[name]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)
@@ -289,22 +346,6 @@ class LongCatFlashMTP(nn.Module):
).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(
@@ -24,7 +24,6 @@ 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
@@ -38,6 +37,42 @@ def uses_ngram_embedding(config: FlashConfig) -> bool:
return getattr(config, "ngram_vocab_size_ratio", None) is not None
def compute_eos_position_ngram_ids(
ngram: "NgramEmbedding",
eos_id: int,
table: torch.Tensor,
tok_req: torch.Tensor,
col: torch.Tensor,
eos_tok: torch.Tensor,
) -> torch.Tensor:
"""Hash ids for positions whose current token is EOS.
The ``ngram_compute_n_gram_ids`` kernel stops the n-gram walk at a negated
(EOS) table entry even at delta 0, but per the reference semantics an EOS
*current* token hashes normally with its full look-back; only later
positions' look-back stops at it. This recomputes those (rare) ids,
mirroring the kernel's per-config walk with delta 0 forced to the
un-negated EOS id.
"""
device = table.device
n, k = ngram.n, ngram.k
num_emb = ngram.num_embedders
deltas = torch.arange(n, device=device)
back = table[tok_req[eos_tok, None], col[eos_tok, None] - deltas] # [E, n]
back[:, 0] = eos_id
valid = (back >= 0).cumprod(dim=1).bool()
toks = torch.where(valid, back, 0).to(torch.int64) # [E, n]
w = ngram.ne_weights.view(num_emb, n).to(torch.int64) # [cfg, n]
m = ngram.ne_mods.view(num_emb).to(torch.int64) # [cfg]
# Config i*k+j is an (i+2)-gram: only deltas < i+2 participate.
cfg_n = torch.arange(num_emb, device=device) // k + 2
dmask = deltas[None, :] < cfg_n[:, None] # [cfg, n]
terms = (toks[:, None, :] * w[None]) % m[None, :, None] * dmask[None]
h = terms.sum(-1) % m[None]
return (h + ngram.exclusive_sizes[:-1].to(torch.int64)).to(torch.int32)
def _config_dtype(config: FlashConfig) -> torch.dtype:
dt = getattr(config, "torch_dtype", None) or getattr(config, "dtype", None)
if isinstance(dt, torch.dtype):
@@ -190,11 +225,11 @@ class FlashNgramModel(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))
if "ngram_embeddings." in name:
# Drop the checkpoint's n-gram tables when the module is
# disabled (e.g. ngram_vocab_size_ratio overridden to None).
if self.ngram_embeddings is not None:
loaded.add(self.ngram_embeddings.load_weight(name, w))
else:
rest.append((name, w))
loaded |= super().load_weights(rest)
@@ -273,12 +308,14 @@ class LongcatFlashNgramForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
class LongcatNgramModelState(DefaultModelState):
"""Per-request n-gram token history for LongCat-Flash-Lite.
"""n-gram input embedding state for LongCat-Flash models.
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``.
``prepare_inputs`` computes the fused n-gram embedding for the batch into
a persistent ``inputs_embeds`` buffer handed to the model forward. Each
position's left-context is gathered from the runner's authoritative
``all_token_ids`` history, which keeps it correct under chunked prefill,
request resumption, and speculative decoding (rejected draft tokens never
enter the history).
"""
def __init__(self, vllm_config, model, encoder_cache, device) -> None:
@@ -288,13 +325,7 @@ class LongcatNgramModelState(DefaultModelState):
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.device = device
self._inputs_embeds_buf = torch.zeros(
self.max_num_tokens,
@@ -303,26 +334,6 @@ class LongcatNgramModelState(DefaultModelState):
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]:
@@ -332,7 +343,7 @@ class LongcatNgramModelState(DefaultModelState):
input_ids = input_batch.input_ids[:num_tokens]
embeds = self._inputs_embeds_buf[:num_padded]
oe_ids = self._compute_oe_ids(input_batch)
oe_ids = self._compute_oe_ids(input_batch, req_states)
embeds[:num_tokens].copy_(self.ngram.embed_batched(input_ids, oe_ids))
model_inputs["inputs_embeds"] = embeds
return model_inputs
@@ -345,14 +356,18 @@ class LongcatNgramModelState(DefaultModelState):
model_inputs["inputs_embeds"] = self._inputs_embeds_buf[:num_tokens]
return model_inputs
def _compute_oe_ids(self, input_batch: InputBatch) -> torch.Tensor:
def _compute_oe_ids(
self, input_batch: InputBatch, req_states: RequestState
) -> 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.
CUDA kernel for the whole batch. The left-context is gathered from the
authoritative ``all_token_ids`` history each step (rather than rolled
incrementally) so rejected speculative-draft tokens never pollute it.
"""
device = self.token_context.device
device = self.device
num_tokens = input_batch.num_tokens
num_reqs = input_batch.num_reqs
ctx_len = self.ctx_len
@@ -365,9 +380,20 @@ class LongcatNgramModelState(DefaultModelState):
max_len = int(req_lens.max().item())
width = ctx_len + max_len
# Left-context: the ctx_len accepted tokens preceding this batch's
# first position, EOS-negated; -1 marks the sequence start.
p0 = req_states.num_computed_tokens.gpu[idx_mapping].long() # [R]
ctx_pos = p0[:, None] + torch.arange(-ctx_len, 0, device=device) # [R, C]
in_range = ctx_pos >= 0
ctx = req_states.all_token_ids.gpu[
idx_mapping[:, None], ctx_pos.clamp_min(0)
].to(torch.int32)
ctx = torch.where(ctx == self.eos_id, -ctx, ctx)
ctx = torch.where(in_range, ctx, ctx.new_full((), -1))
# 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]
table[:, :ctx_len] = ctx
tok_req = torch.repeat_interleave(
torch.arange(num_reqs, device=device), req_lens.long()
)
@@ -396,10 +422,14 @@ class LongcatNgramModelState(DefaultModelState):
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]
# The kernel stops the n-gram walk at a negated (EOS) table entry even
# at delta 0, but per the reference semantics an EOS *current* token
# hashes normally with its full look-back; only later positions'
# look-back stops at it. Recompute the (rare) EOS-position ids here.
eos_tok = (cur == self.eos_id).nonzero(as_tuple=True)[0]
if eos_tok.numel():
n_gram_ids[eos_tok] = compute_eos_position_ngram_ids(
self.ngram, self.eos_id, table, tok_req, col, eos_tok
)
return n_gram_ids.long()
+5
View File
@@ -149,6 +149,11 @@ _TEXT_GENERATION_MODELS = {
"longcat_flash_ngram",
"LongcatFlashNgramForCausalLM",
),
# LongCat-2.0 (LongCat-Flash + n-gram embedding + LongCat Sparse Attention)
"LongcatCausalLM": (
"longcat_flash_ngram",
"LongcatFlashNgramForCausalLM",
),
"MambaForCausalLM": ("mamba", "MambaForCausalLM"),
"Mamba2ForCausalLM": ("mamba2", "Mamba2ForCausalLM"),
"MellumForCausalLM": ("mellum", "MellumForCausalLM"),
+4
View File
@@ -84,6 +84,10 @@ _REASONING_PARSERS_TO_REGISTER = {
"kimi_k2_reasoning_parser",
"KimiK2ReasoningParser",
),
"longcat": (
"longcat_reasoning_parser",
"LongcatReasoningParser",
),
"mimo": (
"qwen3_engine_reasoning_parser",
"Qwen3ParserReasoningAdapter",
@@ -0,0 +1,69 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Sequence
from vllm.entrypoints.openai.engine.protocol import DeltaMessage
from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser
class LongcatReasoningParser(BaseThinkingReasoningParser):
"""
Reasoning parser for LongCat models.
LongCat delimits reasoning with <longcat_think>...</longcat_think>. In
thinking mode the chat template ends the generation prompt with the start
token, so the model output begins mid-reasoning without emitting it
(same convention as DeepSeek R1).
"""
@property
def start_token(self) -> str:
"""The token that starts reasoning content."""
return "<longcat_think>"
@property
def end_token(self) -> str:
"""The token that ends reasoning content."""
return "</longcat_think>"
def extract_reasoning_streaming(
self,
previous_text: str,
current_text: str,
delta_text: str,
previous_token_ids: Sequence[int],
current_token_ids: Sequence[int],
delta_token_ids: Sequence[int],
) -> DeltaMessage | None:
ret = super().extract_reasoning_streaming(
previous_text,
current_text,
delta_text,
previous_token_ids,
current_token_ids,
delta_token_ids,
)
if (
ret is not None
and self.start_token_id not in previous_token_ids
and self.start_token_id not in delta_token_ids
):
if self.end_token_id in delta_token_ids:
# end token in delta with more tokens,
# extract reasoning content and content
end_index = delta_text.find(self.end_token)
reasoning = delta_text[:end_index]
content = delta_text[end_index + len(self.end_token) :]
return DeltaMessage(
reasoning=reasoning,
content=content if content else None,
)
elif self.end_token_id in previous_token_ids:
# end token in previous, thinking content ends
return DeltaMessage(content=delta_text)
else:
# no end token in previous or delta, reasoning content continues
return DeltaMessage(reasoning=delta_text)
return ret
+17
View File
@@ -96,6 +96,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
isaac="IsaacConfig",
kimi_k2="DeepseekV3Config", # Kimi K2 uses same architecture as DeepSeek V3
kimi_linear="KimiLinearConfig",
longcat_flash_ngram="LongcatFlashNgramConfig",
kimi_vl="KimiVLConfig",
kimi_k25="KimiK25Config",
RefinedWeb="RWConfig", # For tiiuae/falcon-40b(-instruct)
@@ -231,6 +232,11 @@ class HFConfigParser(ConfigParserBase):
if config_dict.get("speculators_config") is not None
else model_type
)
if model_type is None and "LongcatCausalLM" in (
config_dict.get("architectures") or []
):
# LongCat-2.0 ships model_type: null without remote code.
model_type = "longcat_flash_ngram"
# Allow hf_overrides to override model_type before checking _CONFIG_REGISTRY
if (hf_overrides := kwargs.pop("hf_overrides", None)) is not None:
if isinstance(hf_overrides, dict) and "model_type" in hf_overrides:
@@ -274,6 +280,17 @@ class HFConfigParser(ConfigParserBase):
config_class.model_type = model_type
# Now that it is registered, it is not considered remote code anymore
trust_remote_code = False
if config_model_type is None:
# The checkpoint has no model_type (e.g. LongCat-2.0), so
# AutoConfig cannot dispatch on it; use the registry class
# directly.
config = config_class.from_pretrained(
model,
revision=revision,
code_revision=code_revision,
**kwargs,
)
return config_dict, _maybe_remap_hf_config_attrs(config)
try:
kwargs = _maybe_update_auto_config_kwargs(kwargs, model_type=model_type)
config = AutoConfig.from_pretrained(
@@ -68,6 +68,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
"KimiLinearConfig": "vllm.transformers_utils.configs.kimi_linear",
"KimiVLConfig": "vllm.transformers_utils.configs.kimi_vl",
"KimiK25Config": "vllm.transformers_utils.configs.kimi_k25",
"LongcatFlashNgramConfig": "vllm.transformers_utils.configs.longcat_flash",
"NemotronConfig": "vllm.transformers_utils.configs.nemotron",
"NemotronHConfig": "vllm.transformers_utils.configs.nemotron_h",
"OlmoHybridConfig": "vllm.transformers_utils.configs.olmo_hybrid",
@@ -144,6 +145,7 @@ __all__ = [
"KimiLinearConfig",
"KimiVLConfig",
"KimiK25Config",
"LongcatFlashNgramConfig",
"NemotronConfig",
"NemotronHConfig",
"OlmoHybridConfig",
@@ -0,0 +1,63 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""LongCat-Flash config for checkpoints without usable remote code.
``meituan-longcat/LongCat-2.0`` ships ``model_type: null`` with
``architectures: ["LongcatCausalLM"]`` and no ``auto_map``, so it cannot be
resolved by ``AutoConfig``. This subclass of the upstream transformers
``LongcatFlashConfig`` fills that gap and aliases the LongCat-2.0
``oe_{vocab_size_ratio,neighbor_num,split_num}`` field names to the n-gram
embedding fields (``ngram_vocab_size_ratio``/``emb_neighbor_num``/
``emb_split_num``) the vLLM model code uses.
"""
import dataclasses
from transformers.models.longcat_flash import (
LongcatFlashConfig as _HfLongcatFlashConfig,
)
def _coerce_float_fields(kwargs: dict) -> dict:
"""The upstream config is a strict dataclass; promote ints in the
checkpoint json (e.g. ``routed_scaling_factor: 9``) to float fields."""
if not dataclasses.is_dataclass(_HfLongcatFlashConfig):
return kwargs
for field in dataclasses.fields(_HfLongcatFlashConfig):
if field.type in (float, "float") and isinstance(kwargs.get(field.name), int):
kwargs[field.name] = float(kwargs[field.name])
return kwargs
class LongcatFlashNgramConfig(_HfLongcatFlashConfig):
model_type = "longcat_flash_ngram"
def __init__(
self,
oe_vocab_size_ratio=None,
oe_neighbor_num=None,
oe_split_num=None,
ngram_vocab_size_ratio=None,
emb_neighbor_num=None,
emb_split_num=None,
**kwargs,
):
self.ngram_vocab_size_ratio = (
ngram_vocab_size_ratio
if ngram_vocab_size_ratio is not None
else oe_vocab_size_ratio
)
self.emb_neighbor_num = (
emb_neighbor_num if emb_neighbor_num is not None else oe_neighbor_num
)
self.emb_split_num = (
emb_split_num if emb_split_num is not None else oe_split_num
)
super().__init__(**_coerce_float_fields(kwargs))
# ``num_hidden_layers`` counts attention sublayers (two per decoder
# layer); the model builds ``num_layers`` decoder layers
# (FlashNgramModel re-syncs it at build).
if kwargs.get("num_hidden_layers") is None:
self.num_hidden_layers = self.num_layers * 2
+5
View File
@@ -27,6 +27,11 @@ from vllm.utils.math_utils import cdiv
_DEEPGEMM_BLACKWELL_EXCLUDED_MODEL_TYPES: set[str] = {
"qwen3_5_text",
"qwen3_5_moe_text",
# LongCat FP8 checkpoints ship non-ue8m0 block scales; steer clear of the
# E8M0 requantization (see sgl-project/sglang#30275).
"longcat_flash",
"longcat_flash_ngram",
"longcat_flash_mtp",
}