@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 08dfd68610

[Model] Add LongCat-Flash-Lite (n-gram embedding) (#47857)

Signed-off-by: mgoin <mgoin64@gmail.com>

Michael Goin committed Jul 10, 2026 at 10:17 UTC 08dfd68610d2e05a0d8ddc99c23488da6163df3f

16 files changed +630 -12

CMakeLists.txt

@@ -390,6 +390,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")

390 "csrc/libtorch_stable/cuda_view.cu"

391 "csrc/libtorch_stable/cuda_utils_kernels.cu"

392 "csrc/libtorch_stable/activation_kernels.cu"

393 + "csrc/libtorch_stable/ngram_embedding_kernels.cu"

394 "csrc/libtorch_stable/quantization/activation_kernels.cu"

395 "csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"

396 "csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"

csrc/libtorch_stable/ngram_embedding_kernels.cu new

+96

@@ -0,0 +1,96 @@

1 +// N-gram embedding index kernel for LongCat-Flash (n-gram embedding variant).

2 +//

3 +// Adapted from SGLang:

4 +// https://github.com/sgl-project/sglang/blob/main/python/sglang/jit_kernel/csrc/ngram_embedding.cuh

5 +//

6 +// For each position, computes the hashed n-gram embedding ids that index the

7 +// concatenated embedder table. Integer tensors are int32 except ``row_indices``

8 +// (int64); the token table is ``[max_running_reqs, max_context_len]`` int32,

9 +// where a negative entry marks an ignored token (e.g. an EOS boundary).

10 +

11 +#include "torch_utils.h"

12 +

13 +#include "ops.h"

14 +

15 +#include <cstdint>

16 +

17 +namespace vllm::ngram_embedding {

18 +

19 +constexpr int kBlockThreads = 256;

20 +

21 +__global__ void ComputeNGramIdsKernel(

22 + int batch_size, int ne_n, int ne_k,

23 + int* ne_weights, // [ne_n-1, ne_k, ne_n]

24 + int* ne_mods, // [ne_n-1, ne_k]

25 + int* exclusive_ne_embedder_size_sums, // [(ne_n-1)*ne_k + 1]

26 + int* exclusive_req_len_sums, // [batch_size + 1]

27 + int* ne_token_table, // [max_running_reqs, max_context_len]

28 + int max_context_len,

29 + const int64_t* __restrict__ row_indices, // [batch_size]

30 + int* column_starts, // [batch_size]

31 + int* n_gram_ids // [token_num, (ne_n-1)*ne_k]

32 +) {

33 + const int req_id = blockIdx.x % batch_size;

34 + const int config_id = (blockIdx.x - req_id) / batch_size;

35 + // n and k are offset from their physical meaning: n = real_n - 2, k = real_k

36 + // - 1 (they index into ne_weights / ne_mods).

37 + const int k = config_id % ne_k;

38 + const int n = (config_id - config_id % ne_k) / ne_k;

39 + const int ne_weight_base_idx = n * ne_k * ne_n + k * ne_n;

40 + const int ne_mod = ne_mods[n * ne_k + k];

41 + for (int i = exclusive_req_len_sums[req_id] + threadIdx.x;

42 + i < exclusive_req_len_sums[req_id + 1]; i += blockDim.x) {

43 + uint64_t n_gram_id = 0;

44 + const int64_t current_token_offset = i - exclusive_req_len_sums[req_id];

45 + const int64_t req_token_table_index =

46 + row_indices[req_id] * static_cast<int64_t>(max_context_len);

47 + const int64_t current_token_table_index =

48 + req_token_table_index + column_starts[req_id] + current_token_offset;

49 + for (int j = 0; j < n + 2; j++) {

50 + if (current_token_table_index - j < req_token_table_index) {

51 + break; // outside this request's range

52 + }

53 + if (ne_token_table[current_token_table_index - j] < 0) {

54 + break; // ignored token

55 + }

56 + const uint64_t term =

57 + (uint64_t)ne_token_table[current_token_table_index - j] *

58 + (uint64_t)ne_weights[ne_weight_base_idx + j];

59 + n_gram_id += term % ne_mod;

60 + }

61 + n_gram_id %= ne_mod;

62 + n_gram_id += exclusive_ne_embedder_size_sums[n * ne_k + k];

63 + n_gram_ids[i * (ne_n - 1) * ne_k + n * ne_k + k] = (int)(n_gram_id);

64 + }

65 +}

66 +

67 +} // namespace vllm::ngram_embedding

68 +

69 +void ngram_compute_n_gram_ids(

70 + int64_t ne_n, int64_t ne_k, torch::stable::Tensor& ne_weights,

71 + torch::stable::Tensor& ne_mods,

72 + torch::stable::Tensor& exclusive_ne_embedder_size_sums,

73 + torch::stable::Tensor& exclusive_req_len_sums,

74 + torch::stable::Tensor& ne_token_table, torch::stable::Tensor& row_indices,

75 + torch::stable::Tensor& column_starts, torch::stable::Tensor& n_gram_ids) {

76 + const int batch_size = static_cast<int>(exclusive_req_len_sums.size(0) - 1);

77 + const int max_context_len = static_cast<int>(ne_token_table.size(1));

78 + const int num_configs = (static_cast<int>(ne_n) - 1) * static_cast<int>(ne_k);

79 + const int grid_size = num_configs * batch_size;

80 + if (grid_size <= 0) return;

81 +

82 + const torch::stable::accelerator::DeviceGuard device_guard(

83 + ne_weights.get_device_index());

84 + const cudaStream_t stream = get_current_cuda_stream();

85 + vllm::ngram_embedding::ComputeNGramIdsKernel<<<

86 + grid_size, vllm::ngram_embedding::kBlockThreads, 0, stream>>>(

87 + batch_size, static_cast<int>(ne_n), static_cast<int>(ne_k),

88 + ne_weights.mutable_data_ptr<int32_t>(),

89 + ne_mods.mutable_data_ptr<int32_t>(),

90 + exclusive_ne_embedder_size_sums.mutable_data_ptr<int32_t>(),

91 + exclusive_req_len_sums.mutable_data_ptr<int32_t>(),

92 + ne_token_table.mutable_data_ptr<int32_t>(), max_context_len,

93 + row_indices.const_data_ptr<int64_t>(),

94 + column_starts.mutable_data_ptr<int32_t>(),

95 + n_gram_ids.mutable_data_ptr<int32_t>());

96 +}

csrc/libtorch_stable/ops.h

@@ -554,3 +554,12 @@ void cp_gather_indexer_k_quant_cache(

554 // quant_block_size * 4]

555 const torch::stable::Tensor& block_table, // [batch_size, num_blocks]

556 const torch::stable::Tensor& cu_seq_lens); // [batch_size + 1]

557 +

558 +// LongCat n-gram embedding index kernel (see ngram_embedding_kernels.cu).

559 +void ngram_compute_n_gram_ids(

560 + int64_t ne_n, int64_t ne_k, torch::stable::Tensor& ne_weights,

561 + torch::stable::Tensor& ne_mods,

562 + torch::stable::Tensor& exclusive_ne_embedder_size_sums,

563 + torch::stable::Tensor& exclusive_req_len_sums,

564 + torch::stable::Tensor& ne_token_table, torch::stable::Tensor& row_indices,

565 + torch::stable::Tensor& column_starts, torch::stable::Tensor& n_gram_ids);

csrc/libtorch_stable/torch_bindings.cpp

+13

@@ -598,9 +598,22 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {

598 "Tensor? initial_state_idx,"

599 "Tensor? cu_chunk_seqlen,"

600 "Tensor? last_chunk_indices) -> ()");

601 +

602 + // LongCat n-gram embedding index kernel. All tensor args are marked mutable

603 + // to match the (non-const) stable-Tensor& C++ signature; only ne_token_table

604 + // and n_gram_ids are actually written in place.

605 + ops.def(

606 + "ngram_compute_n_gram_ids(int ne_n, int ne_k, Tensor(a!) ne_weights, "

607 + "Tensor(b!) ne_mods, Tensor(c!) exclusive_ne_embedder_size_sums, "

608 + "Tensor(d!) exclusive_req_len_sums, Tensor(e!) ne_token_table, "

609 + "Tensor(f!) row_indices, Tensor(g!) column_starts, "

610 + "Tensor(h!) n_gram_ids) -> ()");

611 }

612

613 STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {

614 + // LongCat n-gram embedding index kernel.

615 + ops.impl("ngram_compute_n_gram_ids", TORCH_BOX(&ngram_compute_n_gram_ids));

616 +

617 // Per-token group quantization

618 ops.impl("per_token_group_fp8_quant", TORCH_BOX(&per_token_group_quant_fp8));

619 ops.impl("per_token_group_fp8_quant_packed",

tests/models/registry.py

@@ -382,6 +382,13 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

382 "LongcatFlashForCausalLM": _HfExamplesInfo(

383 "meituan-longcat/LongCat-Flash-Chat", trust_remote_code=True

384 ),

385 + "LongcatFlashNgramForCausalLM": _HfExamplesInfo(

386 + "meituan-longcat/LongCat-Flash-Lite",

387 + trust_remote_code=True,

388 + # Shrink the ~62GB n-gram tables (ngram_vocab_size_ratio * vocab_size)

389 + # so the dummy-weight init test fits in CI memory.

390 + hf_overrides={"ngram_vocab_size_ratio": 1},

391 + ),

392 "MambaForCausalLM": _HfExamplesInfo("state-spaces/mamba-130m-hf"),

393 "Mamba2ForCausalLM": _HfExamplesInfo(

394 "mistralai/Mamba-Codestral-7B-v0.1",

tests/models/utils.py

+7 -2

@@ -529,8 +529,13 @@ def dummy_hf_overrides(

529 }

530 )

531

532 - # Update num_hidden_layers for non-Longcat architectures

533 - if model_arch != "LongcatFlashForCausalLM" and model_arch != "LongCatFlashMTPModel":

532 + # Update num_hidden_layers for non-Longcat architectures (Longcat derives it

533 + # from num_layers for its dual-attention layers).

534 + if model_arch not in (

535 + "LongcatFlashForCausalLM",

536 + "LongCatFlashMTPModel",

537 + "LongcatFlashNgramForCausalLM",

538 + ):

539 update_dict["num_hidden_layers"] = num_hidden_layers

540

541 text_config.update(update_dict)

vllm/_custom_ops.py

+31

@@ -235,6 +235,37 @@ def rms_norm(

235 torch.ops._C.rms_norm(out, input, weight, epsilon)

236

237

238 +# LongCat n-gram embedding index kernel (see csrc/.../ngram_embedding_kernels.cu).

239 +def ngram_compute_n_gram_ids(

240 + ne_n: int,

241 + ne_k: int,

242 + ne_weights: torch.Tensor,

243 + ne_mods: torch.Tensor,

244 + exclusive_ne_embedder_size_sums: torch.Tensor,

245 + exclusive_req_len_sums: torch.Tensor,

246 + ne_token_table: torch.Tensor,

247 + row_indices: torch.Tensor,

248 + column_starts: torch.Tensor,

249 + n_gram_ids: torch.Tensor,

250 +) -> None:

251 + """Compute concatenated (offset) n-gram ids for a ragged prefill batch.

252 +

253 + Writes ``n_gram_ids`` of shape ``[token_num, (ne_n-1)*ne_k]``.

254 + """

255 + torch.ops._C.ngram_compute_n_gram_ids(

256 + ne_n,

257 + ne_k,

258 + ne_weights,

259 + ne_mods,

260 + exclusive_ne_embedder_size_sums,

261 + exclusive_req_len_sums,

262 + ne_token_table,

263 + row_indices,

264 + column_starts,

265 + n_gram_ids,

266 + )

267 +

268 +

269 def fused_add_rms_norm(

270 input: torch.Tensor,

271 residual: torch.Tensor,

vllm/config/speculative.py

+1 -1

@@ -518,7 +518,7 @@ class SpeculativeConfig:

518 "architectures": ["Qwen3_5MoeMTP" if is_moe else "Qwen3_5MTP"],

519 }

520 )

521 - if hf_config.model_type == "longcat_flash":

521 + if hf_config.model_type in ("longcat_flash", "longcat_flash_ngram"):

522 hf_config.model_type = "longcat_flash_mtp"

523 n_predict = getattr(hf_config, "num_nextn_predict_layers", 1)

524 hf_config.update(

vllm/config/vllm.py

@@ -70,6 +70,7 @@ DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset(

70 "DeepseekV2ForCausalLM",

71 "Qwen2MoeForCausalLM",

72 "GraniteMoeForCausalLM",

73 + "LongcatFlashNgramForCausalLM",

74 }

75 )

vllm/model_executor/models/config.py

+19 -1

@@ -499,7 +499,7 @@ class LlamaBidirectionalConfig(VerifyAndUpdateConfig):

499 "last": "LAST",

500 }

501

502 - pooling_type = pooling_type_map.get(hf_config.pooling, None)

502 + pooling_type = pooling_type_map.get(hf_config.pooling)

503 if pooling_type is None:

504 raise ValueError(f"pool_type {hf_config.pooling!r} not supported")

505

@@ -809,6 +809,23 @@ class VoyageQwen3BidirectionalEmbedModelConfig(VerifyAndUpdateConfig):

809 model_config.hf_config.embedding_size = model_config.hf_config.num_labels

810

811

812 +class LongcatFlashNgramForCausalLMConfig(VerifyAndUpdateConfig):

813 + @staticmethod

814 + def verify_and_update_config(vllm_config: "VllmConfig") -> None:

815 + # LongCat-Flash-Lite's zero-expert MoE trips a data-dependent assert

816 + # under torch.compile, and its n-gram inputs_embeds are only wired for

817 + # FULL cudagraph capture (PIECEWISE prefill drops them). Default to

818 + # no-compile + FULL cudagraph (prefill runs eager) unless the user

819 + # configured compilation explicitly.

820 + from vllm.config.compilation import CompilationMode, CUDAGraphMode

821 +

822 + compilation_config = vllm_config.compilation_config

823 + if compilation_config.mode is None:

824 + compilation_config.mode = CompilationMode.NONE

825 + if compilation_config.cudagraph_mode is None:

826 + compilation_config.cudagraph_mode = CUDAGraphMode.FULL

827 +

828 +

829 MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {

830 "ColBERTJinaRobertaModel": JinaRobertaModelConfig,

831 "ColQwen3_5": ColQwen3_5Config,

@@ -822,6 +839,7 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {

839 "Gemma4ForConditionalGeneration": Gemma4Config,

840 "Gemma4UnifiedForConditionalGeneration": Gemma4Config,

841 "GptOssForCausalLM": GptOssForCausalLMConfig,

842 + "LongcatFlashNgramForCausalLM": LongcatFlashNgramForCausalLMConfig,

843 "GteModel": SnowflakeGteNewModelConfig,

844 "GteNewForSequenceClassification": GteNewModelConfig,

845 "GteNewModel": GteNewModelConfig,

vllm/model_executor/models/longcat_flash.py

+12 -3

@@ -318,7 +318,7 @@ class LongcatMoe(nn.Module):

318 hidden_states = hidden_states.view(-1, hidden_dim)

319

320 # Align to FusedMoE padded hidden size to avoid dim mismatch

321 - padded_hidden = self.experts.hidden_size

321 + padded_hidden = self.experts.moe_config.hidden_dim

322 if hidden_dim < padded_hidden:

323 hidden_states_padded = torch.nn.functional.pad(

324 hidden_states,

@@ -687,14 +687,23 @@ class FlashModel(nn.Module):

687 ).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1)

688 self_attn.w_kc = w_kc.transpose(1, 2).contiguous().transpose(1, 2)

689 self_attn.w_vc = w_vc.contiguous().transpose(1, 2)

690 - if self.config.mla_scale_q_lora:

690 + # Guard against compounding on incremental load_weights calls:

691 + # the in-place ``*=`` would otherwise re-apply the MLA LoRA

692 + # scaling to the layernorm weights on each pass.

693 + if self.config.mla_scale_q_lora and not getattr(

694 + self_attn, "_mla_q_lora_scaled", False

695 + ):

696 self_attn.q_a_layernorm.weight.data *= (

697 self.config.hidden_size / self.config.q_lora_rank

698 ) ** 0.5

694 - if self.config.mla_scale_kv_lora:

699 + self_attn._mla_q_lora_scaled = True

700 + if self.config.mla_scale_kv_lora and not getattr(

701 + self_attn, "_mla_kv_lora_scaled", False

702 + ):

703 self_attn.kv_a_layernorm.weight.data *= (

704 self.config.hidden_size / self.config.kv_lora_rank

705 ) ** 0.5

706 + self_attn._mla_kv_lora_scaled = True

707 return loaded_params

708

709

vllm/model_executor/models/longcat_flash_mtp.py

+14 -4

@@ -126,8 +126,10 @@ class LongCatMultiTokenPredictor(nn.Module):

126 class LongCatFlashMTP(nn.Module):

127 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

128 super().__init__()

129 - # LongCat MTP without MoE layers

130 - vllm_config.model_config.hf_config.n_routed_experts = None

129 + # LongCat MTP has no MoE layers: clear n_routed_experts so the predictor

130 + # builds a dense MLP. object.__setattr__ bypasses the ngram remote

131 + # config's strict validation (it rejects setting the int field to None).

132 + object.__setattr__(vllm_config.model_config.hf_config, "n_routed_experts", None)

133 self.config = FlashConfig(**vllm_config.model_config.hf_config.__dict__)

134 self.quant_config = (

135 None

@@ -287,14 +289,22 @@ class LongCatFlashMTP(nn.Module):

289 ).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1)

290 self_attn.w_kc = w_kc.transpose(1, 2).contiguous().transpose(1, 2)

291 self_attn.w_vc = w_vc.contiguous().transpose(1, 2)

290 - if self.config.mla_scale_q_lora:

292 + # Guard against compounding on incremental load_weights calls (the

293 + # in-place *= would otherwise double-apply the LoRA scaling).

294 + if self.config.mla_scale_q_lora and not getattr(

295 + self_attn, "_mla_q_lora_scaled", False

296 + ):

297 self_attn.q_a_layernorm.weight.data *= (

298 self.config.hidden_size / self.config.q_lora_rank

299 ) ** 0.5

294 - if self.config.mla_scale_kv_lora:

300 + self_attn._mla_q_lora_scaled = True

301 + if self.config.mla_scale_kv_lora and not getattr(

302 + self_attn, "_mla_kv_lora_scaled", False

303 + ):

304 self_attn.kv_a_layernorm.weight.data *= (

305 self.config.hidden_size / self.config.kv_lora_rank

306 ) ** 0.5

307 + self_attn._mla_kv_lora_scaled = True

308 return loaded_params

309

310 def _rewrite_spec_layer_name(

vllm/model_executor/models/longcat_flash_ngram.py new

+405

@@ -0,0 +1,405 @@

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, 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 + word = self.word_embeddings(input_ids) # [N, H]

161 + flat = oe_ids.permute(1, 0).contiguous() # [num_embedders, N]

162 + oe = self.oe_embedder(flat) # [num_embedders, N, oe_dim]

163 + proj = torch.bmm(oe, self.oe_projection) # [num_embedders, N, H]

164 + all_h = torch.cat([word.unsqueeze(0), proj], dim=0) # [ne+1, N, H]

165 + return all_h.mean(dim=0) # [N, H]

166 +

167 +

168 +class FlashNgramModel(FlashModel):

169 + """FlashModel whose input embedding is an :class:`NgramEmbedding`."""

170 +

171 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:

172 + # Each FlashDecoderLayer is a *dual* layer (2 attentions), so the number

173 + # of decoder layers is ``num_layers``. The ngram HF config sets

174 + # ``num_hidden_layers`` to a multiple of that (attention-module count),

175 + # which FlashModel would otherwise build as too many (dead) layers.

176 + hf = vllm_config.model_config.hf_config

177 + num_layers = getattr(hf, "num_layers", None)

178 + if num_layers is not None and hf.num_hidden_layers != num_layers:

179 + hf.num_hidden_layers = num_layers

180 + super().__init__(vllm_config=vllm_config, prefix=prefix)

181 + if get_pp_group().is_first_rank and uses_ngram_embedding(self.config):

182 + self.ngram_embeddings = NgramEmbedding(self.config, self.embed_tokens)

183 + else:

184 + self.ngram_embeddings = None

185 +

186 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

187 + # Names arrive with the ``model.`` prefix already stripped (routed here

188 + # by AutoWeightsLoader). Split the concatenated/sharded ngram tables and

189 + # stacked projections; delegate everything else to FlashModel.

190 + loaded: set[str] = set()

191 + rest: list[tuple[str, torch.Tensor]] = []

192 + for name, w in weights:

193 + if self.ngram_embeddings is not None and (

194 + "ngram_embeddings.embedders." in name

195 + or "ngram_embeddings.post_projs." in name

196 + ):

197 + loaded.add(self.ngram_embeddings.load_weight(name, w))

198 + else:

199 + rest.append((name, w))

200 + loaded |= super().load_weights(rest)

201 + return loaded

202 +

203 +

204 +class LongcatFlashNgramForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

205 + """LongCat-Flash-Lite for causal LM (MRV2-only, n-gram embedding)."""

206 +

207 + packed_modules_mapping = {

208 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

209 + "gate_up_proj": ["gate_proj", "up_proj"],

210 + }

211 +

212 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:

213 + super().__init__()

214 + if not vllm_config.use_v2_model_runner:

215 + raise NotImplementedError(

216 + "LongcatFlashNgramForCausalLM (LongCat-Flash-Lite) requires the "

217 + "V2 model runner for its n-gram embedding state; it is selected "

218 + "automatically unless VLLM_USE_V2_MODEL_RUNNER=0 is set."

219 + )

220 + config = FlashConfig(**vllm_config.model_config.hf_config.__dict__)

221 + config.intermediate_size = getattr(

222 + config, "ffn_hidden_size", config.intermediate_size

223 + )

224 + self.config = config

225 + self.quant_config = vllm_config.quant_config

226 +

227 + self.model = FlashNgramModel(

228 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

229 + )

230 + if get_pp_group().is_last_rank:

231 + self.lm_head = ParallelLMHead(

232 + config.vocab_size,

233 + config.hidden_size,

234 + quant_config=self.quant_config,

235 + prefix=maybe_prefix(prefix, "lm_head"),

236 + )

237 + else:

238 + self.lm_head = PPMissingLayer()

239 + self.logits_processor = LogitsProcessor(config.vocab_size)

240 + self.make_empty_intermediate_tensors = (

241 + self.model.make_empty_intermediate_tensors

242 + )

243 +

244 + @staticmethod

245 + def get_model_state_cls() -> type["LongcatNgramModelState"]:

246 + return LongcatNgramModelState

247 +

248 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

249 + return self.model.embed_input_ids(input_ids)

250 +

251 + def forward(

252 + self,

253 + input_ids: torch.Tensor | None,

254 + positions: torch.Tensor,

255 + intermediate_tensors=None,

256 + inputs_embeds: torch.Tensor | None = None,

257 + ):

258 + # inputs_embeds is produced by LongcatNgramModelState.prepare_inputs.

259 + return self.model(input_ids, positions, intermediate_tensors, inputs_embeds)

260 +

261 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:

262 + return self.logits_processor(self.lm_head, hidden_states)

263 +

264 + def get_expert_mapping(self):

265 + return self.model.get_expert_mapping()

266 +

267 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

268 + # AutoWeightsLoader routes ``model.*`` to FlashNgramModel.load_weights

269 + # (which handles the ngram split) and ``lm_head.*`` to the head. MTP

270 + # weights are not part of this model.

271 + loader = AutoWeightsLoader(self, skip_prefixes=["model.mtp."])

272 + return loader.load_weights(weights)

273 +

274 +

275 +class LongcatNgramModelState(DefaultModelState):

276 + """Per-request n-gram token history for LongCat-Flash-Lite.

277 +

278 + Maintains a small CPU-side per-slot context (last ``n-1`` processed tokens)

279 + and a persistent ``inputs_embeds`` buffer. ``prepare_inputs`` computes the

280 + fused n-gram embedding per request into the buffer, handed to the model

281 + forward as ``inputs_embeds``.

282 + """

283 +

284 + def __init__(self, vllm_config, model, encoder_cache, device) -> None:

285 + super().__init__(vllm_config, model, encoder_cache, device)

286 + config = model.config

287 + self.ngram = model.model.ngram_embeddings

288 + self.n = int(config.emb_neighbor_num)

289 + self.ctx_len = self.n - 1

290 + self.eos_id = int(config.eos_token_id)

291 +

292 + # Per-slot left-context: last ``n-1`` processed tokens, EOS negated. A

293 + # negative entry (incl. the -1 fill) marks a context boundary that stops

294 + # the n-gram walk (matches the kernel's EOS break / fresh-request start).

295 + self.token_context = torch.full(

296 + (self.max_num_reqs, self.ctx_len), -1, dtype=torch.int32, device=device

297 + )

298 +

299 + self._inputs_embeds_buf = torch.zeros(

300 + self.max_num_tokens,

301 + config.hidden_size,

302 + dtype=self.dtype,

303 + device=device,

304 + )

305 +

306 + def _neg_eos(self, toks: list[int]) -> list[int]:

307 + return [-t if t == self.eos_id else t for t in toks]

308 +

309 + def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:

310 + super().add_request(req_index, new_req_data) # rope positions

311 + # Fresh request -> no left-context (-1 fill). On resume, seed from the

312 + # already-processed token tail. Use prefill_token_ids (full processed

313 + # sequence incl. generated tokens on v2 resume), like DefaultModelState;

314 + # the prompt alone would be too short when resuming after decode.

315 + ctx = [-1] * self.ctx_len

316 + ncomp = new_req_data.num_computed_tokens

317 + toks_src = new_req_data.prefill_token_ids or new_req_data.prompt_token_ids

318 + if ncomp > 0 and toks_src is not None:

319 + lo = max(0, ncomp - self.ctx_len)

320 + toks = self._neg_eos(list(toks_src[lo:ncomp]))

321 + ctx[self.ctx_len - len(toks) :] = toks

322 + self.token_context[req_index] = torch.tensor(

323 + ctx, dtype=torch.int32, device=self.token_context.device

324 + )

325 +

326 + def prepare_inputs(

327 + self, input_batch: InputBatch, req_states: RequestState

328 + ) -> dict[str, Any]:

329 + model_inputs = super().prepare_inputs(input_batch, req_states) # positions

330 + num_tokens = input_batch.num_tokens

331 + num_padded = input_batch.num_tokens_after_padding

332 + input_ids = input_batch.input_ids[:num_tokens]

333 + embeds = self._inputs_embeds_buf[:num_padded]

334 +

335 + oe_ids = self._compute_oe_ids(input_batch)

336 + embeds[:num_tokens].copy_(self.ngram.embed_batched(input_ids, oe_ids))

337 + model_inputs["inputs_embeds"] = embeds

338 + return model_inputs

339 +

340 + def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]:

341 + # FULL cudagraph replay reads only the captured buffers, so capture must

342 + # reference the same persistent ``inputs_embeds`` buffer prepare_inputs

343 + # re-fills (the base class wires this for multimodal models only).

344 + model_inputs = super().prepare_dummy_inputs(num_reqs, num_tokens) # positions

345 + model_inputs["inputs_embeds"] = self._inputs_embeds_buf[:num_tokens]

346 + return model_inputs

347 +

348 + def _compute_oe_ids(self, input_batch: InputBatch) -> torch.Tensor:

349 + """Batched global n-gram ids ``[num_tokens, num_embedders]``.

350 +

351 + Assembles an ephemeral per-request token table (``[n-1] context ++

352 + current tokens``, EOS-negated) and runs the ``ngram_compute_n_gram_ids``

353 + CUDA kernel for the whole batch, then rolls each slot's context forward.

354 + """

355 + device = self.token_context.device

356 + num_tokens = input_batch.num_tokens

357 + num_reqs = input_batch.num_reqs

358 + ctx_len = self.ctx_len

359 + idx_mapping = input_batch.idx_mapping[:num_reqs].long()

360 + qsl = input_batch.query_start_loc[: num_reqs + 1].to(torch.int32)

361 + cur = input_batch.input_ids[:num_tokens].to(torch.int32)

362 +

363 + cur_neg = torch.where(cur == self.eos_id, -cur, cur)

364 + req_lens = qsl[1:] - qsl[:-1]

365 + max_len = int(req_lens.max().item())

366 + width = ctx_len + max_len

367 +

368 + # table[r] = [context(n-1) | current tokens | pad(-1)]

369 + table = torch.full((num_reqs, width), -1, dtype=torch.int32, device=device)

370 + table[:, :ctx_len] = self.token_context[idx_mapping]

371 + tok_req = torch.repeat_interleave(

372 + torch.arange(num_reqs, device=device), req_lens.long()

373 + )

374 + col = ctx_len + (

375 + torch.arange(num_tokens, device=device) - qsl[:-1].long()[tok_req]

376 + )

377 + table[tok_req, col] = cur_neg

378 +

379 + column_starts = torch.full(

380 + (num_reqs,), ctx_len, dtype=torch.int32, device=device

381 + )

382 + row_indices = torch.arange(num_reqs, dtype=torch.int64, device=device)

383 + n_gram_ids = torch.empty(

384 + num_tokens, self.ngram.num_embedders, dtype=torch.int32, device=device

385 + )

386 + ops.ngram_compute_n_gram_ids(

387 + self.n,

388 + self.ngram.k,

389 + self.ngram.ne_weights,

390 + self.ngram.ne_mods,

391 + self.ngram.exclusive_sizes,

392 + qsl,

393 + table,

394 + row_indices,

395 + column_starts,

396 + n_gram_ids,

397 + )

398 +

399 + # Roll context: new context = last n-1 of [context | current] per slot.

400 + gather = req_lens.long().unsqueeze(1) + torch.arange(

401 + ctx_len, device=device

402 + ).unsqueeze(0)

403 + rows = torch.arange(num_reqs, device=device).unsqueeze(1)

404 + self.token_context[idx_mapping] = table[rows, gather]

405 + return n_gram_ids.long()

vllm/model_executor/models/registry.py

@@ -145,6 +145,10 @@ _TEXT_GENERATION_MODELS = {

145 # For decapoda-research/llama-*

146 "LLaMAForCausalLM": ("llama", "LlamaForCausalLM"),

147 "LongcatFlashForCausalLM": ("longcat_flash", "LongcatFlashForCausalLM"),

148 + "LongcatFlashNgramForCausalLM": (

149 + "longcat_flash_ngram",

150 + "LongcatFlashNgramForCausalLM",

151 + ),

152 "MambaForCausalLM": ("mamba", "MambaForCausalLM"),

153 "Mamba2ForCausalLM": ("mamba2", "Mamba2ForCausalLM"),

154 "MellumForCausalLM": ("mellum", "MellumForCausalLM"),

vllm/transformers_utils/model_arch_config_convertor.py

@@ -268,6 +268,7 @@ class ModelArchConfigConvertorBase:

268 "kimi_k2",

269 "kimi_linear",

270 "longcat_flash",

271 + "longcat_flash_ngram",

272 "pangu_ultra_moe",

273 "pangu_ultra_moe_mtp",

274 "bailing_hybrid",

vllm/v1/worker/gpu/attn_utils.py

+9 -1

@@ -540,7 +540,15 @@ def init_kv_cache(

540 shared_kv_cache_layers=shared_kv_cache_layers,

541 kv_cache_config=kv_cache_config,

542 )

543 - bind_kv_cache(kv_caches, forward_context, runner_kv_caches)

543 + # Dual-attention models (e.g. LongCat-Flash) put two Attention modules per

544 + # decoder layer, so a layer name carries two integers (layer + module index).

545 + num_attn_module = (

546 + 2

547 + if vllm_config.model_config.hf_config.model_type

548 + in ("longcat_flash", "longcat_flash_ngram")

549 + else 1

550 + )

551 + bind_kv_cache(kv_caches, forward_context, runner_kv_caches, num_attn_module)

552 return kv_caches

553

554