[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>

.buildkite/image_build/image_build_arm64.sh

@@ -19,13 +19,14 @@ if docker manifest inspect "$IMAGE" >/dev/null 2>&1; then
   echo "Image found"
 else
   echo "Image not found, proceeding with build..."
-  # build for arm64 GPU targets: Grace/GH200 (sm_90) and DGX Spark/GB10
+  # build for arm64 GPU targets: Grace/GH200 (sm_90),
+  # Blackwell/Thor (sm_100/sm_103/sm_110), and DGX Spark/GB10
   # (sm_121, family-covered by 12.0 under CUDA 13)
   docker build --file docker/Dockerfile \
     --platform linux/arm64 \
     --build-arg max_jobs=16 \
     --build-arg nvcc_threads=4 \
-    --build-arg torch_cuda_arch_list="9.0 12.0" \
+    --build-arg torch_cuda_arch_list="9.0 10.0 11.0 12.0" \
     --build-arg USE_SCCACHE=1 \
     --build-arg buildkite_commit="$BUILDKITE_COMMIT" \
     --tag "$IMAGE" \

CMakeLists.txt

@@ -390,6 +390,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     "csrc/libtorch_stable/cuda_view.cu"
     "csrc/libtorch_stable/cuda_utils_kernels.cu"
     "csrc/libtorch_stable/activation_kernels.cu"
+    "csrc/libtorch_stable/ngram_embedding_kernels.cu"
     "csrc/libtorch_stable/quantization/activation_kernels.cu"
     "csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
     "csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"

csrc/libtorch_stable/ngram_embedding_kernels.cu

@@ -0,0 +1,96 @@
+// N-gram embedding index kernel for LongCat-Flash (n-gram embedding variant).
+//
+// Adapted from SGLang:
+// https://github.com/sgl-project/sglang/blob/main/python/sglang/jit_kernel/csrc/ngram_embedding.cuh
+//
+// For each position, computes the hashed n-gram embedding ids that index the
+// concatenated embedder table. Integer tensors are int32 except ``row_indices``
+// (int64); the token table is ``[max_running_reqs, max_context_len]`` int32,
+// where a negative entry marks an ignored token (e.g. an EOS boundary).
+
+#include "torch_utils.h"
+
+#include "ops.h"
+
+#include <cstdint>
+
+namespace vllm::ngram_embedding {
+
+constexpr int kBlockThreads = 256;
+
+__global__ void ComputeNGramIdsKernel(
+    int batch_size, int ne_n, int ne_k,
+    int* ne_weights,                       // [ne_n-1, ne_k, ne_n]
+    int* ne_mods,                          // [ne_n-1, ne_k]
+    int* exclusive_ne_embedder_size_sums,  // [(ne_n-1)*ne_k + 1]
+    int* exclusive_req_len_sums,           // [batch_size + 1]
+    int* ne_token_table,  // [max_running_reqs, max_context_len]
+    int max_context_len,
+    const int64_t* __restrict__ row_indices,  // [batch_size]
+    int* column_starts,                       // [batch_size]
+    int* n_gram_ids                           // [token_num, (ne_n-1)*ne_k]
+) {
+  const int req_id = blockIdx.x % batch_size;
+  const int config_id = (blockIdx.x - req_id) / batch_size;
+  // n and k are offset from their physical meaning: n = real_n - 2, k = real_k
+  // - 1 (they index into ne_weights / ne_mods).
+  const int k = config_id % ne_k;
+  const int n = (config_id - config_id % ne_k) / ne_k;
+  const int ne_weight_base_idx = n * ne_k * ne_n + k * ne_n;
+  const int ne_mod = ne_mods[n * ne_k + k];
+  for (int i = exclusive_req_len_sums[req_id] + threadIdx.x;
+       i < exclusive_req_len_sums[req_id + 1]; i += blockDim.x) {
+    uint64_t n_gram_id = 0;
+    const int64_t current_token_offset = i - exclusive_req_len_sums[req_id];
+    const int64_t req_token_table_index =
+        row_indices[req_id] * static_cast<int64_t>(max_context_len);
+    const int64_t current_token_table_index =
+        req_token_table_index + column_starts[req_id] + current_token_offset;
+    for (int j = 0; j < n + 2; j++) {
+      if (current_token_table_index - j < req_token_table_index) {
+        break;  // outside this request's range
+      }
+      if (ne_token_table[current_token_table_index - j] < 0) {
+        break;  // ignored token
+      }
+      const uint64_t term =
+          (uint64_t)ne_token_table[current_token_table_index - j] *
+          (uint64_t)ne_weights[ne_weight_base_idx + j];
+      n_gram_id += term % ne_mod;
+    }
+    n_gram_id %= ne_mod;
+    n_gram_id += exclusive_ne_embedder_size_sums[n * ne_k + k];
+    n_gram_ids[i * (ne_n - 1) * ne_k + n * ne_k + k] = (int)(n_gram_id);
+  }
+}
+
+}  // namespace vllm::ngram_embedding
+
+void ngram_compute_n_gram_ids(
+    int64_t ne_n, int64_t ne_k, torch::stable::Tensor& ne_weights,
+    torch::stable::Tensor& ne_mods,
+    torch::stable::Tensor& exclusive_ne_embedder_size_sums,
+    torch::stable::Tensor& exclusive_req_len_sums,
+    torch::stable::Tensor& ne_token_table, torch::stable::Tensor& row_indices,
+    torch::stable::Tensor& column_starts, torch::stable::Tensor& n_gram_ids) {
+  const int batch_size = static_cast<int>(exclusive_req_len_sums.size(0) - 1);
+  const int max_context_len = static_cast<int>(ne_token_table.size(1));
+  const int num_configs = (static_cast<int>(ne_n) - 1) * static_cast<int>(ne_k);
+  const int grid_size = num_configs * batch_size;
+  if (grid_size <= 0) return;
+
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      ne_weights.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
+  vllm::ngram_embedding::ComputeNGramIdsKernel<<<
+      grid_size, vllm::ngram_embedding::kBlockThreads, 0, stream>>>(
+      batch_size, static_cast<int>(ne_n), static_cast<int>(ne_k),
+      ne_weights.mutable_data_ptr<int32_t>(),
+      ne_mods.mutable_data_ptr<int32_t>(),
+      exclusive_ne_embedder_size_sums.mutable_data_ptr<int32_t>(),
+      exclusive_req_len_sums.mutable_data_ptr<int32_t>(),
+      ne_token_table.mutable_data_ptr<int32_t>(), max_context_len,
+      row_indices.const_data_ptr<int64_t>(),
+      column_starts.mutable_data_ptr<int32_t>(),
+      n_gram_ids.mutable_data_ptr<int32_t>());
+}

csrc/libtorch_stable/ops.h

@@ -554,3 +554,12 @@ void cp_gather_indexer_k_quant_cache(
                                                 // quant_block_size * 4]
     const torch::stable::Tensor& block_table,   // [batch_size, num_blocks]
     const torch::stable::Tensor& cu_seq_lens);  // [batch_size + 1]
+
+// LongCat n-gram embedding index kernel (see ngram_embedding_kernels.cu).
+void ngram_compute_n_gram_ids(
+    int64_t ne_n, int64_t ne_k, torch::stable::Tensor& ne_weights,
+    torch::stable::Tensor& ne_mods,
+    torch::stable::Tensor& exclusive_ne_embedder_size_sums,
+    torch::stable::Tensor& exclusive_req_len_sums,
+    torch::stable::Tensor& ne_token_table, torch::stable::Tensor& row_indices,
+    torch::stable::Tensor& column_starts, torch::stable::Tensor& n_gram_ids);

csrc/libtorch_stable/torch_bindings.cpp

@@ -598,9 +598,22 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "Tensor? initial_state_idx,"
       "Tensor? cu_chunk_seqlen,"
       "Tensor? last_chunk_indices) -> ()");
+
+  // LongCat n-gram embedding index kernel. All tensor args are marked mutable
+  // to match the (non-const) stable-Tensor& C++ signature; only ne_token_table
+  // and n_gram_ids are actually written in place.
+  ops.def(
+      "ngram_compute_n_gram_ids(int ne_n, int ne_k, Tensor(a!) ne_weights, "
+      "Tensor(b!) ne_mods, Tensor(c!) exclusive_ne_embedder_size_sums, "
+      "Tensor(d!) exclusive_req_len_sums, Tensor(e!) ne_token_table, "
+      "Tensor(f!) row_indices, Tensor(g!) column_starts, "
+      "Tensor(h!) n_gram_ids) -> ()");
 }
 
 STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
+  // LongCat n-gram embedding index kernel.
+  ops.impl("ngram_compute_n_gram_ids", TORCH_BOX(&ngram_compute_n_gram_ids));
+
   // Per-token group quantization
   ops.impl("per_token_group_fp8_quant", TORCH_BOX(&per_token_group_quant_fp8));
   ops.impl("per_token_group_fp8_quant_packed",

docs/models/supported_models.md

@@ -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. | | ✅︎ |

tests/kernels/attention/test_rocm_triton_attn_dsv4.py

@@ -197,46 +197,6 @@ def _ragged_from_rows(
     )
 
 
-def _ref_combine_topk_swa_ragged(
-    device: torch.device,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    expected_ragged = torch.tensor(
-        [
-            100,
-            101,
-            7,
-            8,
-            9,
-            110,
-            111,
-            8,
-            9,
-            10,
-            120,
-            121,
-            122,
-            9,
-            10,
-            11,
-            150,
-            27,
-            28,
-            29,
-            160,
-            161,
-            28,
-            29,
-            30,
-        ],
-        dtype=torch.int32,
-        device=device,
-    )
-    expected_lens = torch.tensor([5, 5, 6, 4, 5], dtype=torch.int32, device=device)
-    expected_indptr = torch.zeros(6, dtype=torch.int32, device=device)
-    torch.cumsum(expected_lens, dim=0, out=expected_indptr[1:])
-    return expected_ragged, expected_indptr, expected_lens
-
-
 @torch.inference_mode()
 def test_compute_global_topk_ragged_indices_and_indptr() -> None:
     from vllm.models.deepseek_v4.amd.rocm import (
@@ -369,55 +329,6 @@ def test_sparse_attn_decode_ragged_kernel() -> None:
     torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
 
 
-@torch.inference_mode()
-def test_combine_topk_swa_indices_ragged() -> None:
-    from vllm.models.deepseek_v4.amd.rocm import (
-        combine_topk_swa_indices_ragged,
-    )
-
-    device = torch.device("cuda")
-    topk_indices = torch.tensor(
-        [
-            [100, 101, 102, 103],
-            [110, 111, 112, 113],
-            [120, 121, 122, 123],
-            [130, 131, 132, 133],
-            [140, 141, 142, 143],
-        ],
-        dtype=torch.int32,
-        device=device,
-    )
-    query_start_loc = torch.tensor([0, 3, 5], dtype=torch.int32, device=device)
-    seq_lens = torch.tensor([6, 4], dtype=torch.int32, device=device)
-    gather_lens = torch.tensor([4, 3], dtype=torch.int32, device=device)
-    window_size = 3
-    compress_ratio = 2
-    topk = 4
-    M = 20
-    N = 8
-
-    actual_ragged, actual_indptr, actual_lens = combine_topk_swa_indices_ragged(
-        topk_indices,
-        query_start_loc,
-        seq_lens,
-        gather_lens,
-        window_size,
-        compress_ratio,
-        topk,
-        M,
-        N,
-    )
-    expected_ragged, expected_indptr, expected_lens = _ref_combine_topk_swa_ragged(
-        device
-    )
-
-    torch.testing.assert_close(
-        actual_ragged[: expected_ragged.numel()], expected_ragged
-    )
-    torch.testing.assert_close(actual_indptr, expected_indptr)
-    torch.testing.assert_close(actual_lens, expected_lens)
-
-
 @requires_gfx950
 @torch.inference_mode()
 def test_decode_num_splits_heuristic(monkeypatch) -> None:

tests/kernels/test_ngram_embedding.py

@@ -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)

tests/model_executor/model_loader/test_filter_duplicate_safetensors.py

@@ -0,0 +1,79 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+import os
+import tempfile
+
+import pytest
+
+from vllm.model_executor.model_loader.weight_utils import (
+    filter_duplicate_safetensors_files,
+)
+
+
+def test_filter_duplicate_safetensors_files_missing_weight():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        existing_file = os.path.join(tmpdir, "model-00001-of-00002.safetensors")
+        with open(existing_file, "wb") as f:
+            f.write(b"")
+
+        existing_file2 = os.path.join(tmpdir, "model-00002-of-00002.safetensors")
+        with open(existing_file2, "wb") as f:
+            f.write(b"")
+
+        index_file = os.path.join(tmpdir, "model.safetensors.index.json")
+        index_content = {
+            "weight_map": {
+                "layer.0.weight": "model-00001-of-00002.safetensors",
+                "layer.1.weight": "model-00002-of-00002.safetensors",
+                "layer.2.weight": "model-00003-of-00002.safetensors",
+            }
+        }
+        with open(index_file, "w") as f:
+            json.dump(index_content, f)
+
+        hf_weights_files = [
+            os.path.join(tmpdir, "model-00001-of-00002.safetensors"),
+            os.path.join(tmpdir, "model-00002-of-00002.safetensors"),
+        ]
+
+        with pytest.raises(FileNotFoundError) as exc_info:
+            filter_duplicate_safetensors_files(
+                hf_weights_files=hf_weights_files,
+                hf_folder=tmpdir,
+                index_file="model.safetensors.index.json",
+            )
+
+        assert "model-00003-of-00002.safetensors" in str(exc_info.value)
+
+
+def test_filter_duplicate_safetensors_files_all_exist():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        existing_files = []
+        for i in range(1, 3):
+            file_path = os.path.join(tmpdir, f"model-0000{i}-of-00002.safetensors")
+            with open(file_path, "wb") as f:
+                f.write(b"")
+            existing_files.append(file_path)
+
+        index_file = os.path.join(tmpdir, "model.safetensors.index.json")
+        index_content = {
+            "weight_map": {
+                "layer.0.weight": "model-00001-of-00002.safetensors",
+                "layer.1.weight": "model-00002-of-00002.safetensors",
+            }
+        }
+        with open(index_file, "w") as f:
+            json.dump(index_content, f)
+
+        filter_duplicate_safetensors_files(
+            hf_weights_files=existing_files,
+            hf_folder=tmpdir,
+            index_file="model.safetensors.index.json",
+        )
+
+
+if __name__ == "__main__":
+    test_filter_duplicate_safetensors_files_missing_weight()
+    test_filter_duplicate_safetensors_files_all_exist()

tests/models/registry.py

@@ -382,6 +382,19 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
     "LongcatFlashForCausalLM": _HfExamplesInfo(
         "meituan-longcat/LongCat-Flash-Chat", trust_remote_code=True
     ),
+    "LongcatFlashNgramForCausalLM": _HfExamplesInfo(
+        "meituan-longcat/LongCat-Flash-Lite",
+        trust_remote_code=True,
+        # Shrink the ~62GB n-gram tables (ngram_vocab_size_ratio * vocab_size)
+        # 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",

tests/models/test_registry.py

@@ -48,9 +48,11 @@ def test_registry_imports(model_arch):
                 "(see #41376)"
             )
 
-    # DSpark draft model is NVIDIA-only; class is stubbed to None on ROCm/XPU.
-    if model_arch == "DSparkDraftModel" and not current_platform.is_cuda():
-        pytest.skip("DSparkDraftModel is only supported on CUDA")
+    # DSpark draft model is supported on CUDA and ROCm; stubbed to None on XPU.
+    if model_arch == "DSparkDraftModel" and not (
+        current_platform.is_cuda() or current_platform.is_rocm()
+    ):
+        pytest.skip("DSparkDraftModel is only supported on CUDA and ROCm")
 
     # Ensure all model classes can be imported successfully
     model_cls = ModelRegistry._try_load_model_cls(model_arch)

tests/models/utils.py

@@ -529,8 +529,13 @@ def dummy_hf_overrides(
             }
         )
 
-    # Update num_hidden_layers for non-Longcat architectures
-    if model_arch != "LongcatFlashForCausalLM" and model_arch != "LongCatFlashMTPModel":
+    # Update num_hidden_layers for non-Longcat architectures (Longcat derives it
+    # from num_layers for its dual-attention layers).
+    if model_arch not in (
+        "LongcatFlashForCausalLM",
+        "LongCatFlashMTPModel",
+        "LongcatFlashNgramForCausalLM",
+    ):
         update_dict["num_hidden_layers"] = num_hidden_layers
 
     text_config.update(update_dict)

vllm/_custom_ops.py

@@ -235,6 +235,37 @@ def rms_norm(
     torch.ops._C.rms_norm(out, input, weight, epsilon)
 
 
+# LongCat n-gram embedding index kernel (see csrc/.../ngram_embedding_kernels.cu).
+def ngram_compute_n_gram_ids(
+    ne_n: int,
+    ne_k: int,
+    ne_weights: torch.Tensor,
+    ne_mods: torch.Tensor,
+    exclusive_ne_embedder_size_sums: torch.Tensor,
+    exclusive_req_len_sums: torch.Tensor,
+    ne_token_table: torch.Tensor,
+    row_indices: torch.Tensor,
+    column_starts: torch.Tensor,
+    n_gram_ids: torch.Tensor,
+) -> None:
+    """Compute concatenated (offset) n-gram ids for a ragged prefill batch.
+
+    Writes ``n_gram_ids`` of shape ``[token_num, (ne_n-1)*ne_k]``.
+    """
+    torch.ops._C.ngram_compute_n_gram_ids(
+        ne_n,
+        ne_k,
+        ne_weights,
+        ne_mods,
+        exclusive_ne_embedder_size_sums,
+        exclusive_req_len_sums,
+        ne_token_table,
+        row_indices,
+        column_starts,
+        n_gram_ids,
+    )
+
+
 def fused_add_rms_norm(
     input: torch.Tensor,
     residual: torch.Tensor,

vllm/config/speculative.py

@@ -518,9 +518,15 @@ class SpeculativeConfig:
                     "architectures": ["Qwen3_5MoeMTP" if is_moe else "Qwen3_5MTP"],
                 }
             )
-        if hf_config.model_type == "longcat_flash":
+        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"]}
             )
@@ -942,6 +948,31 @@ class SpeculativeConfig:
                         "`num_speculative_tokens` was not provided"
                     )
 
+                if self.method == "dspark":
+                    # DSpark is a semi-autoregressive *block* drafter. A
+                    # speculative length smaller than the checkpoint's block
+                    # feeds the block / Markov-head machinery an unsupported
+                    # layout and yields incorrect (garbled) output rather than
+                    # merely lower acceptance. Require num_speculative_tokens to
+                    # be at least the block size (e.g. 5 or 7 for DeepSeek-V4).
+                    dspark_block_size = getattr(
+                        self.draft_model_config.hf_config,
+                        "dspark_block_size",
+                        None,
+                    )
+                    if (
+                        dspark_block_size is not None
+                        and self.num_speculative_tokens < dspark_block_size
+                    ):
+                        raise ValueError(
+                            "DSpark requires num_speculative_tokens >= "
+                            f"dspark_block_size ({dspark_block_size}); got "
+                            f"{self.num_speculative_tokens}. Smaller values "
+                            "produce incorrect output. Use "
+                            f"num_speculative_tokens={dspark_block_size} or "
+                            "larger (e.g. 7)."
+                        )
+
                 self.draft_tensor_parallel_size = (
                     SpeculativeConfig._verify_and_get_draft_tp(
                         self.target_parallel_config,

vllm/config/vllm.py

@@ -70,6 +70,8 @@ DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset(
         "DeepseekV2ForCausalLM",
         "Qwen2MoeForCausalLM",
         "GraniteMoeForCausalLM",
+        "LongcatFlashNgramForCausalLM",
+        "LongcatCausalLM",
     }
 )
 

vllm/model_executor/layers/sparse_attn_indexer.py

@@ -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,

vllm/model_executor/model_loader/weight_utils.py

@@ -595,6 +595,14 @@ def filter_duplicate_safetensors_files(
     weight_files_in_index = set()
     for weight_name in weight_map:
         weight_files_in_index.add(os.path.join(hf_folder, weight_map[weight_name]))
+    # Check if files referenced in model.safetensors.index.json actually exist.
+    # Raise error if any file is missing.
+    hf_weights_files_set = set(hf_weights_files)
+    missing_files = weight_files_in_index - hf_weights_files_set
+    if missing_files:
+        raise FileNotFoundError(
+            f"Weight files referenced in index but missing: {missing_files}"
+        )
     # Filter out any fields that are not found in the index file.
     hf_weights_files = [f for f in hf_weights_files if f in weight_files_in_index]
     return hf_weights_files

vllm/model_executor/models/config.py

@@ -499,7 +499,7 @@ class LlamaBidirectionalConfig(VerifyAndUpdateConfig):
             "last": "LAST",
         }
 
-        pooling_type = pooling_type_map.get(hf_config.pooling, None)
+        pooling_type = pooling_type_map.get(hf_config.pooling)
         if pooling_type is None:
             raise ValueError(f"pool_type {hf_config.pooling!r} not supported")
 
@@ -809,6 +809,32 @@ class VoyageQwen3BidirectionalEmbedModelConfig(VerifyAndUpdateConfig):
         model_config.hf_config.embedding_size = model_config.hf_config.num_labels
 
 
+class LongcatFlashNgramForCausalLMConfig(VerifyAndUpdateConfig):
+    @staticmethod
+    def verify_and_update_config(vllm_config: "VllmConfig") -> None:
+        # LongCat-Flash-Lite's zero-expert MoE trips a data-dependent assert
+        # under torch.compile, and its n-gram inputs_embeds are only wired for
+        # FULL cudagraph capture (PIECEWISE prefill drops them). Default to
+        # no-compile + FULL cudagraph (prefill runs eager) unless the user
+        # configured compilation explicitly.
+        from vllm.config.compilation import CompilationMode, CUDAGraphMode
+
+        compilation_config = vllm_config.compilation_config
+        if compilation_config.mode is None:
+            compilation_config.mode = CompilationMode.NONE
+        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,
     "ColQwen3_5": ColQwen3_5Config,
@@ -822,6 +848,8 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
     "Gemma4ForConditionalGeneration": Gemma4Config,
     "Gemma4UnifiedForConditionalGeneration": Gemma4Config,
     "GptOssForCausalLM": GptOssForCausalLMConfig,
+    "LongcatFlashNgramForCausalLM": LongcatFlashNgramForCausalLMConfig,
+    "LongcatCausalLM": LongcatFlashNgramForCausalLMConfig,
     "GteModel": SnowflakeGteNewModelConfig,
     "GteNewForSequenceClassification": GteNewModelConfig,
     "GteNewModel": GteNewModelConfig,

vllm/model_executor/models/deepseek_v2.py

@@ -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)

vllm/model_executor/models/longcat_flash.py

@@ -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,
@@ -318,7 +323,7 @@ class LongcatMoe(nn.Module):
         hidden_states = hidden_states.view(-1, hidden_dim)
 
         # Align to FusedMoE padded hidden size to avoid dim mismatch
-        padded_hidden = self.experts.hidden_size
+        padded_hidden = self.experts.moe_config.hidden_dim
         if hidden_dim < padded_hidden:
             hidden_states_padded = torch.nn.functional.pad(
                 hidden_states,
@@ -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,14 +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)
-                if self.config.mla_scale_q_lora:
-                    self_attn.q_a_layernorm.weight.data *= (
-                        self.config.hidden_size / self.config.q_lora_rank
-                    ) ** 0.5
-                if self.config.mla_scale_kv_lora:
-                    self_attn.kv_a_layernorm.weight.data *= (
-                        self.config.hidden_size / self.config.kv_lora_rank
-                    ) ** 0.5
         return loaded_params
 
 

vllm/model_executor/models/longcat_flash_mtp.py

@@ -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,
@@ -126,8 +144,10 @@ class LongCatMultiTokenPredictor(nn.Module):
 class LongCatFlashMTP(nn.Module):
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
-        # LongCat MTP without MoE layers
-        vllm_config.model_config.hf_config.n_routed_experts = None
+        # LongCat MTP has no MoE layers: clear n_routed_experts so the predictor
+        # builds a dense MLP. object.__setattr__ bypasses the ngram remote
+        # config's strict validation (it rejects setting the int field to None).
+        object.__setattr__(vllm_config.model_config.hf_config, "n_routed_experts", None)
         self.config = FlashConfig(**vllm_config.model_config.hf_config.__dict__)
         self.quant_config = (
             None
@@ -176,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 = {
@@ -186,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
@@ -209,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:
@@ -258,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)
@@ -287,14 +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)
-        if self.config.mla_scale_q_lora:
-            self_attn.q_a_layernorm.weight.data *= (
-                self.config.hidden_size / self.config.q_lora_rank
-            ) ** 0.5
-        if self.config.mla_scale_kv_lora:
-            self_attn.kv_a_layernorm.weight.data *= (
-                self.config.hidden_size / self.config.kv_lora_rank
-            ) ** 0.5
         return loaded_params
 
     def _rewrite_spec_layer_name(

vllm/model_executor/models/longcat_flash_ngram.py

@@ -0,0 +1,435 @@
+# 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.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 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):
+        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 "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)
+        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):
+    """n-gram input embedding state for LongCat-Flash models.
+
+    ``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:
+        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)
+        self.device = device
+
+        self._inputs_embeds_buf = torch.zeros(
+            self.max_num_tokens,
+            config.hidden_size,
+            dtype=self.dtype,
+            device=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, req_states)
+        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, 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. 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.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
+
+        # 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] = ctx
+        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,
+        )
+
+        # 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()

vllm/model_executor/models/registry.py

@@ -145,6 +145,15 @@ _TEXT_GENERATION_MODELS = {
     # For decapoda-research/llama-*
     "LLaMAForCausalLM": ("llama", "LlamaForCausalLM"),
     "LongcatFlashForCausalLM": ("longcat_flash", "LongcatFlashForCausalLM"),
+    "LongcatFlashNgramForCausalLM": (
+        "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"),

vllm/models/deepseek_v4/__init__.py

@@ -15,16 +15,16 @@ from .quant_config import DeepseekV4FP8Config
 # default that mypy sees; the ROCm/XPU branches override at runtime and are
 # kept type-compatible via ``# type: ignore[assignment]``.
 if current_platform.is_rocm():
+    from .amd.dspark import (  # type: ignore[assignment]
+        DSparkDeepseekV4ForCausalLM,
+    )
     from .amd.model import DeepseekV4ForCausalLM
     from .amd.mtp import DeepSeekV4MTP
-
-    # DSpark is NVIDIA-only for now.
-    DSparkDeepseekV4ForCausalLM = None  # type: ignore[assignment]
 elif current_platform.is_xpu():
     from .xpu.model import DeepseekV4ForCausalLM  # type: ignore[assignment]
     from .xpu.mtp import DeepSeekV4MTP  # type: ignore[assignment]
 
-    DSparkDeepseekV4ForCausalLM = None  # type: ignore[assignment]
+    DSparkDeepseekV4ForCausalLM = None  # type: ignore[assignment, misc]
 else:
     from .nvidia.dspark import (  # type: ignore[assignment]
         DSparkDeepseekV4ForCausalLM,

vllm/models/deepseek_v4/amd/dspark.py

@@ -0,0 +1,499 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""DSpark draft model for DeepSeek-V4 on ROCm/AMD (gfx950).
+
+ROCm port of ``nvidia/dspark.py``. Follows the same nvidia->amd recipe used for
+``amd/mtp.py``:
+
+  * import ``DeepseekV4DecoderLayer`` from the AMD ``.model`` (aiter/triton
+    attention + MHC CustomOp path) instead of the nvidia one;
+  * route the MHC head through the ``HCHeadOp`` CustomOp dispatcher (aiter /
+    tilelang / triton / torch) instead of calling the tilelang kernels directly,
+    and gate the trailing ``mhc_post`` on ``use_fused_mhc`` (False on the aiter
+    path, where the decoder layer already applies hc_post in-layer);
+  * drop the mega-MoE weight path (``make_deepseek_v4_expert_params_mapping`` /
+    ``use_mega_moe`` / ``finalize_mega_moe_weights`` do not exist in amd/model.py).
+
+Everything else — the semi-autoregressive drafting hooks, the Markov head, the
+sliding-window context-KV insert, and the checkpoint ``mtp.*`` weight remap — is
+pure torch / Triton and shared with the nvidia implementation unchanged.
+"""
+
+from collections.abc import Iterable
+
+import regex as re
+import torch
+import torch.nn as nn
+
+from vllm.config import VllmConfig, get_current_vllm_config
+from vllm.distributed import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+)
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe import (
+    fused_moe_make_expert_params_mapping,
+)
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import ReplicatedLinear
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.mhc import HCHeadOp
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.models.qwen3_dspark import (
+    DSparkMarkovHead,
+)
+from vllm.model_executor.models.utils import maybe_prefix
+
+from .model import (
+    DeepseekV4DecoderLayer,
+)
+
+logger = init_logger(__name__)
+
+# MoE expert scale suffix differs by expert dtype (mirrors deepseek_v4 loaders):
+# fp4 experts register ``.weight_scale``; block-fp8 experts ``.weight_scale_inv``.
+_EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$")
+
+
+class DSparkDeepseekV4Model(nn.Module):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        assert vllm_config.speculative_config is not None
+        config = vllm_config.speculative_config.draft_model_config.hf_config
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.hc_mult = config.hc_mult
+        self.hc_eps = config.hc_eps
+        self.rms_norm_eps = config.rms_norm_eps
+        self.num_hidden_layers = config.num_hidden_layers
+        self.target_layer_ids = tuple(config.dspark_target_layer_ids)
+
+        self.num_dspark_layers = getattr(config, "n_mtp_layers", None) or 3
+
+        # Shared with the target (aliased by the speculator's loading utility).
+        self.embed_tokens = VocabParallelEmbedding(
+            config.vocab_size,
+            config.hidden_size,
+            prefix=maybe_prefix(prefix, "embed_tokens"),
+        )
+
+        self.main_proj = ReplicatedLinear(
+            config.hidden_size * len(self.target_layer_ids),
+            config.hidden_size,
+            bias=False,
+            return_bias=False,
+            quant_config=vllm_config.quant_config,
+            prefix=maybe_prefix(prefix, "main_proj"),
+        )
+        self.main_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        current_vllm_config = get_current_vllm_config()
+        self.layers = nn.ModuleList(
+            [
+                DeepseekV4DecoderLayer(
+                    current_vllm_config,
+                    prefix=maybe_prefix(prefix, f"layers.{self.num_hidden_layers + i}"),
+                )
+                for i in range(self.num_dspark_layers)
+            ]
+        )
+
+        # Heads: final norm + hc_head, and the Markov head
+        # Loaded from the "final" MTP layer weights (mtp.*) in the target checkpoint
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        hc_dim = self.hc_mult * config.hidden_size
+        self.hc_head_fn = nn.Parameter(
+            torch.empty(self.hc_mult, hc_dim, dtype=torch.float32),
+            requires_grad=False,
+        )
+        self.hc_head_base = nn.Parameter(
+            torch.empty(self.hc_mult, dtype=torch.float32), requires_grad=False
+        )
+        self.hc_head_scale = nn.Parameter(
+            torch.empty(1, dtype=torch.float32), requires_grad=False
+        )
+        draft_vocab_size = (
+            getattr(config, "draft_vocab_size", None) or config.vocab_size
+        )
+        self.markov_head = DSparkMarkovHead(
+            config.vocab_size,
+            draft_vocab_size,
+            config.dspark_markov_rank,
+            prefix=maybe_prefix(prefix, "markov_head"),
+        )
+
+        # MHC head CustomOp dispatcher (aiter / tilelang / triton / torch),
+        # replacing the direct nvidia tilelang kernel call.
+        self.hc_head_op = HCHeadOp()
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def combine_hidden_states(self, aux_hidden_states: torch.Tensor) -> torch.Tensor:
+        """main_x = main_norm(main_proj(concat of target aux hidden states)).
+
+        ``aux_hidden_states`` is [T, hidden_size * len(target_layer_ids)].
+        """
+        return self.main_norm(self.main_proj(aux_hidden_states))
+
+    @torch.inference_mode()
+    def precompute_and_store_context_kv(
+        self,
+        main_x: torch.Tensor,
+        context_positions: torch.Tensor,
+        context_slot_mappings: list[torch.Tensor | None] | None = None,
+    ) -> None:
+        """Insert the sliding-window context KV for every draft layer.
+
+        Mirrors the reference DSparkAttention: each layer derives its context KV
+        from the SAME projected target hidden ``main_x``, via that layer's own
+        ``wkv`` + ``kv_norm`` + RoPE + quant, then writes it at the
+        layer's context slots.
+
+        ``context_slot_mappings`` is a per-layer list (each entry is the context
+        slot mapping for that layer's kv-cache group, since the hybrid manager may
+        place draft layers in different groups). ``None`` (or a ``None`` entry)
+        runs the projection to reserve workspace but writes nothing (profiling).
+        """
+        for i, layer in enumerate(self.layers):
+            slot_mapping = (
+                None if context_slot_mappings is None else context_slot_mappings[i]
+            )
+            attn = layer.attn
+            # Optimized DSV4 MLA path: wkv part of the fused wq_a|wkv projection
+            # (q_lora part discarded), then RoPE/quant/insert via the fused op.
+            qr_kv, _ = attn.fused_wqa_wkv(main_x)
+            kv = qr_kv[..., attn.q_lora_rank :]
+            kv = attn.kv_norm(kv)
+            if slot_mapping is None:
+                continue
+            _insert_context_kv(attn, kv, context_positions, slot_mapping)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_input_ids(input_ids)
+        # Expand to hc_mult copies for hyper-connections ([T, H] -> [T, hc, H]).
+        hidden_states = inputs_embeds.unsqueeze(-2).repeat(1, self.hc_mult, 1)
+
+        residual = post_mix = res_mix = None
+        for layer in self.layers:
+            hidden_states, residual, post_mix, res_mix = layer(
+                hidden_states,
+                positions,
+                input_ids,
+                post_mix,
+                res_mix,
+                residual,
+            )
+        # On the fused-MHC path the trailing hc_post must be applied here; on the
+        # aiter unfused path (ROCm default) the decoder layer already applied
+        # hc_post in-layer and returned None mixes, so this is skipped. Mirrors
+        # amd/mtp.py.
+        last_layer = self.layers[-1]
+        if last_layer.use_fused_mhc:
+            hidden_states = last_layer.hc_post(
+                hidden_states, residual, post_mix, res_mix
+            )
+        # hc_head reduces the hc copies; return the PRE-norm head hidden.
+        hidden_states = self.hc_head_op(
+            hidden_states,
+            self.hc_head_fn,
+            self.hc_head_scale,
+            self.hc_head_base,
+            self.rms_norm_eps,
+            self.hc_eps,
+        )
+        return hidden_states
+
+
+def _insert_context_kv(
+    attn: nn.Module,
+    kv: torch.Tensor,
+    positions: torch.Tensor,
+    slot_mapping: torch.Tensor,
+) -> None:
+    """RoPE + quant + paged-cache insert of (already kv_norm'd) context KV.
+
+    Reuses the DSV4 fused insert ops (which also process a query; we pass a dummy
+    query and discard it, since context tokens have no query). Mirrors
+    ``DeepseekV4Attention._fused_qnorm_rope_kv_insert``.
+    """
+    swa_cache = attn.swa_cache_layer.kv_cache
+    block_size = attn.swa_cache_layer.block_size
+    cos_sin_cache = attn.rotary_emb.cos_sin_cache
+    cache_dtype = swa_cache.dtype
+    n_ctx = kv.shape[0]
+    dummy_q = torch.zeros(
+        (n_ctx, attn.n_local_heads, attn.head_dim),
+        dtype=kv.dtype,
+        device=kv.device,
+    )
+    if cache_dtype == torch.uint8:
+        # fp8_ds_mla UE8M0 paged layout (the gfx950 aiter SWA cache path).
+        swa_2d = swa_cache.view(swa_cache.shape[0], -1)
+        torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+            dummy_q,
+            kv,
+            swa_2d,
+            slot_mapping,
+            positions,
+            cos_sin_cache,
+            attn.padded_heads,
+            attn.eps,
+            block_size,
+        )
+    elif cache_dtype == torch.bfloat16:
+        swa_3d = swa_cache.view(-1, block_size, attn.head_dim)
+        torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
+            dummy_q,
+            kv,
+            swa_3d,
+            slot_mapping,
+            positions,
+            cos_sin_cache,
+            attn.eps,
+            block_size,
+        )
+    else:  # per-tensor fp8 (torch.float8_e4m3fn)
+        # NOTE(rocm): unreachable on ROCm/aiter, where the SWA cache dtype is
+        # uint8 (fp8_ds_mla) or bfloat16. This branch relies on FlashInfer-only
+        # attributes (``_flashinfer_fp8_*``) that the aiter attention layer does
+        # not define; kept for parity with the nvidia path.
+        swa_3d = swa_cache.view(-1, block_size, attn.head_dim)
+        dummy_q_fp8 = torch.zeros_like(dummy_q, dtype=torch.float8_e4m3fn)
+        torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
+            dummy_q,
+            kv,
+            dummy_q_fp8,
+            swa_3d,
+            slot_mapping,
+            positions,
+            cos_sin_cache,
+            attn._flashinfer_fp8_kv_scale,
+            attn._flashinfer_fp8_q_scale_inv,
+            attn.eps,
+            block_size,
+        )
+
+
+class DSparkDeepseekV4ForCausalLM(nn.Module):
+    # Draft weights ship in the target checkpoint (mtp.*) without embed/head, so
+    # load_dspark_model always aliases the target's.
+    has_own_embed_tokens = False
+    has_own_lm_head = False
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        assert vllm_config.speculative_config is not None
+        self.draft_model_config = vllm_config.speculative_config.draft_model_config
+        self.config = self.draft_model_config.hf_config
+        self.model = DSparkDeepseekV4Model(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+        )
+        # Shared with the target (aliased by the speculator's load utility).
+        self.lm_head = ParallelLMHead(
+            self.config.vocab_size,
+            self.config.hidden_size,
+            prefix=maybe_prefix(prefix, "lm_head"),
+        )
+        self.logits_processor = LogitsProcessor(self.config.vocab_size)
+
+    # --- Hooks used by the speculator -------------------------------------
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids)
+
+    def combine_hidden_states(self, aux_hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.model.combine_hidden_states(aux_hidden_states)
+
+    def get_draft_kv_cache_layer_names(self) -> list[str]:
+        # DSV4 MLA path: each draft layer's sliding-window cache is a separate
+        # layer, named by its prefix.
+        return [layer.attn.swa_cache_layer.prefix for layer in self.model.layers]
+
+    def precompute_and_store_context_kv(
+        self,
+        context_states: torch.Tensor,
+        context_positions: torch.Tensor,
+        context_slot_mappings: list[torch.Tensor | None] | None = None,
+    ) -> None:
+        self.model.precompute_and_store_context_kv(
+            context_states, context_positions, context_slot_mappings
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # Returns the pre-norm hc_head hidden ([T, hidden_size]).
+        return self.model(input_ids, positions, inputs_embeds)
+
+    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Base logits U_k = lm_head(norm(head_hidden))."""
+        return self.logits_processor(self.lm_head, self.model.norm(hidden_states))
+
+    def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # Full-vocab draft: base logits, no d2t scatter.
+        return self.compute_logits(hidden_states)
+
+    def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor:
+        return draft_ids  # full-vocab: draft ids are target ids
+
+    def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.markov_head.embed(token_ids)
+
+    def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor:
+        return self.model.markov_head.bias(markov_embed, self.logits_processor)
+
+    # --- Weight loading ----------------------------------------------------
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Load the ``mtp.{0,1,2}.*`` draft weights from the target checkpoint.
+
+        Non-mtp weights (embed/head/main layers) belong to the target model and
+        are skipped here. ``embed_tokens``/``lm_head`` are aliased from the target.
+        """
+        # AMD DeepseekV4MoE has no mega-MoE path; always use the standard
+        # per-expert fused-MoE mapping (mirrors amd/mtp.py).
+        expert_mapping = fused_moe_make_expert_params_mapping(
+            self,
+            ckpt_gate_proj_name="w1",
+            ckpt_down_proj_name="w2",
+            ckpt_up_proj_name="w3",
+            num_experts=self.config.n_routed_experts,
+        )
+        expert_scale_suffix = (
+            ".weight_scale"
+            if getattr(self.config, "expert_dtype", "fp4") == "fp4"
+            else ".weight_scale_inv"
+        )
+
+        # (param_name, ckpt_shard_name, shard_id) for non-expert stacked params.
+        stacked_params_mapping = [
+            ("gate_up_proj", "w1", 0),
+            ("gate_up_proj", "w3", 1),
+            ("attn.fused_wqa_wkv", "attn.wq_a", 0),
+            ("attn.fused_wqa_wkv", "attn.wkv", 1),
+        ]
+
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+
+        tp_size = get_tensor_model_parallel_world_size()
+        tp_rank = get_tensor_model_parallel_rank()
+        n_local_head = self.config.num_attention_heads // tp_size
+        head_start = n_local_head * tp_rank
+        head_end = n_local_head * (tp_rank + 1)
+
+        for name, loaded_weight in weights:
+            mapped = self._remap_dspark_name(name)
+            if mapped is None:
+                continue
+            name = mapped
+
+            # ``.scale`` -> per-method scale suffix.
+            if name.endswith(".scale"):
+                suffix = (
+                    expert_scale_suffix
+                    if _EXPERT_SCALE_RE.search(name)
+                    else ".weight_scale_inv"
+                )
+                name = name.removesuffix(".scale") + suffix
+
+            # E8M0 expert scales: keep raw exponent bytes.
+            if ".experts." in name:
+                if (
+                    "weight_scale" in name
+                    and loaded_weight.dtype == torch.float8_e8m0fnu
+                ):
+                    loaded_weight = loaded_weight.view(torch.uint8)
+                for param_name, weight_name, expert_id, shard_id in expert_mapping:
+                    if weight_name not in name:
+                        continue
+                    name_mapped = name.replace(weight_name, param_name)
+                    param = params_dict[name_mapped]
+                    success = param.weight_loader(
+                        param,
+                        loaded_weight,
+                        name_mapped,
+                        shard_id=shard_id,
+                        expert_id=expert_id,
+                        return_success=True,
+                    )
+                    if success:
+                        loaded_params.add(name_mapped)
+                        break
+                continue
+
+            # Stacked rules only apply to decoder-layer weights. Head-stack params
+            # (main_proj/norm/hc_head/markov_head) load directly — otherwise e.g.
+            # "markov_w1" would collide with the "w1" shard rule.
+            is_layer_param = name.startswith("model.layers.")
+            for param_name, weight_name, stacked_shard_id in stacked_params_mapping:
+                if not is_layer_param or weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                param = params_dict[name]
+                param.weight_loader(param, loaded_weight, stacked_shard_id)
+                loaded_params.add(name)
+                break
+            else:
+                if "attn_sink" in name:
+                    narrow = loaded_weight[head_start:head_end]
+                    params_dict[name][: narrow.shape[0]].copy_(narrow)
+                    loaded_params.add(name)
+                    continue
+                if ".shared_experts.w2" in name:
+                    name = name.replace(
+                        ".shared_experts.w2", ".shared_experts.down_proj"
+                    )
+                if name.endswith(".ffn.gate.bias"):
+                    name = name.replace(
+                        ".ffn.gate.bias", ".ffn.gate.e_score_correction_bias"
+                    )
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+                loaded_params.add(name)
+
+        logger.info_once("DSpark draft model loaded: %d params", len(loaded_params))
+        return loaded_params
+
+    def _remap_dspark_name(self, name: str) -> str | None:
+        """Map a checkpoint ``mtp.{i}.*`` name to this model's parameter path.
+
+        Returns None for non-mtp weights (owned by the target model).
+        """
+        m = re.match(r"mtp\.(\d+)\.(.*)", name)
+        if m is None:
+            return None
+        stage = int(m.group(1))
+        rest = m.group(2)
+        # The confidence head is not wired into inference yet; drop its weights.
+        if rest.startswith("confidence_head."):
+            return None
+        # Head-stack params live at model level (mtp.last), context combiner at
+        # model level (mtp.0); everything else is a per-layer decoder block.
+        head_prefixes = (
+            "norm.",
+            "hc_head_fn",
+            "hc_head_base",
+            "hc_head_scale",
+            "markov_head.",
+        )
+        if rest.startswith(("main_proj.", "main_norm.")) or rest.startswith(
+            head_prefixes
+        ):
+            return f"model.{rest}"
+        return f"model.layers.{stage}.{rest}"

vllm/models/deepseek_v4/amd/model.py

@@ -40,7 +40,11 @@ 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.interfaces import SupportsPP
+from vllm.model_executor.models.interfaces import (
+    EagleModelMixin,
+    SupportsEagle3,
+    SupportsPP,
+)
 from vllm.model_executor.models.utils import (
     AutoWeightsLoader,
     PPMissingLayer,
@@ -437,7 +441,7 @@ class DeepseekV4DecoderLayer(nn.Module):
         )
 
 
-class DeepseekV4Model(nn.Module):
+class DeepseekV4Model(nn.Module, EagleModelMixin):
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
 
@@ -573,7 +577,19 @@ class DeepseekV4Model(nn.Module):
             hidden_states = intermediate_tensors["hidden_states"]
 
         residual, post_mix, res_mix = None, None, None
-        for layer in islice(self.layers, self.start_layer, self.end_layer):
+        # EAGLE3 / DSpark / DFlash aux hidden states: reconstructed (post-mhc)
+        # hidden state at the configured target layers, averaged over the
+        # hc_mult streams to [T, hidden_size]. Empty unless a draft model set
+        # aux_hidden_state_layers.
+        aux_hidden_states: list[torch.Tensor] = []
+        # On the fused path the final layer's hc_post output is reused below
+        # (avoids computing hc_post twice when the last layer is also an aux
+        # layer).
+        final_aux_recon: torch.Tensor | None = None
+        for idx, layer in enumerate(
+            islice(self.layers, self.start_layer, self.end_layer),
+            start=self.start_layer,
+        ):
             hidden_states, residual, post_mix, res_mix = layer(
                 hidden_states,
                 positions,
@@ -582,8 +598,30 @@ class DeepseekV4Model(nn.Module):
                 res_mix,
                 residual,
             )
+            if (idx + 1) in self.aux_hidden_state_layers:
+                # On the unfused (aiter) path the layer already applied hc_post,
+                # so hidden_states is the reconstructed stream; on the fused
+                # path reconstruct it via hc_post before averaging.
+                if layer.use_fused_mhc:
+                    aux_recon = layer.hc_post(
+                        hidden_states, residual, post_mix, res_mix
+                    )
+                    final_aux_recon = aux_recon
+                else:
+                    aux_recon = hidden_states
+                aux_hidden_states.append(aux_recon.mean(dim=1))
         if layer is not None and layer.use_fused_mhc:
-            hidden_states = layer.hc_post(hidden_states, residual, post_mix, res_mix)
+            # Reuse the last layer's hc_post output if it was already computed
+            # for the aux hidden state above; otherwise compute it now.
+            if (
+                final_aux_recon is not None
+                and self.end_layer in self.aux_hidden_state_layers
+            ):
+                hidden_states = final_aux_recon
+            else:
+                hidden_states = layer.hc_post(
+                    hidden_states, residual, post_mix, res_mix
+                )
 
         if not get_pp_group().is_last_rank:
             return IntermediateTensors({"hidden_states": hidden_states})
@@ -601,6 +639,8 @@ class DeepseekV4Model(nn.Module):
             self.hc_eps,
         )
         hidden_states = self.norm(hidden_states)
+        if len(aux_hidden_states) > 0:
+            return hidden_states, aux_hidden_states
         return hidden_states
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
@@ -751,7 +791,7 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
     )
 
 
-class DeepseekV4ForCausalLM(nn.Module, SupportsPP):
+class DeepseekV4ForCausalLM(nn.Module, SupportsPP, SupportsEagle3):
     model_cls = DeepseekV4Model
 
     # Default mapper assumes the original FP4-expert checkpoint layout.

vllm/models/deepseek_v4/amd/rocm.py

@@ -272,153 +272,6 @@ def compute_global_topk_ragged_indices_and_indptr(
     return global_topk_ragged, topk_indptr, topk_lens
 
 
-@triton.jit
-def _compute_combined_lens_kernel(
-    combined_lens_ptr,
-    query_start_loc_ptr,
-    seq_lens_ptr,
-    TOP_K: tl.constexpr,
-    COMPRESS_RATIO: tl.constexpr,
-    WINDOW_SIZE: tl.constexpr,
-):
-    batch_idx = tl.program_id(0)
-    worker_id = tl.program_id(1)
-    num_workers = tl.num_programs(1)
-
-    base = tl.load(query_start_loc_ptr)
-    query_start = tl.load(query_start_loc_ptr + batch_idx) - base
-    query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base
-    query_len = query_end - query_start
-    seq_len = tl.load(seq_lens_ptr + batch_idx)
-    start_pos = seq_len - query_len
-
-    for token_idx in range(query_start + worker_id, query_end, num_workers):
-        token_idx_in_query = token_idx - query_start
-        pos = start_pos + token_idx_in_query
-        topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
-        swa_len = tl.minimum(pos + 1, WINDOW_SIZE)
-        tl.store(combined_lens_ptr + token_idx, topk_len + swa_len)
-
-
-@triton.jit
-def _combine_topk_swa_indices_ragged_kernel(
-    combined_ragged_ptr,
-    combined_indptr_ptr,
-    topk_indices_ptr,
-    topk_indices_stride,
-    query_start_loc_ptr,
-    seq_lens_ptr,
-    gather_lens_ptr,
-    M,
-    N,
-    topk_width,
-    TOP_K: tl.constexpr,
-    COMPRESS_RATIO: tl.constexpr,
-    WINDOW_SIZE: tl.constexpr,
-    BLOCK_SIZE: tl.constexpr,
-):
-    batch_idx = tl.program_id(0)
-    worker_id = tl.program_id(1)
-    block_idx = tl.program_id(2)
-    num_workers = tl.num_programs(1)
-
-    base = tl.load(query_start_loc_ptr)
-    query_start = tl.load(query_start_loc_ptr + batch_idx) - base
-    query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base
-    query_len = query_end - query_start
-    seq_len = tl.load(seq_lens_ptr + batch_idx)
-    gather_len = tl.load(gather_lens_ptr + batch_idx)
-    start_pos = seq_len - query_len
-    gather_start = seq_len - gather_len
-
-    for token_idx in range(query_start + worker_id, query_end, num_workers):
-        token_idx_in_query = token_idx - query_start
-        pos = start_pos + token_idx_in_query
-        topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
-        swa_len = tl.minimum(pos + 1, WINDOW_SIZE)
-        combined_len = topk_len + swa_len
-
-        offset = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
-        if block_idx * BLOCK_SIZE < combined_len:
-            out_start = tl.load(combined_indptr_ptr + token_idx)
-            topk_mask = (offset < topk_len) & (offset < topk_width)
-            topk_vals = tl.load(
-                topk_indices_ptr + token_idx * topk_indices_stride + offset,
-                mask=topk_mask,
-                other=-1,
-            )
-            tl.store(
-                combined_ragged_ptr + out_start + offset,
-                topk_vals + M * batch_idx,
-                mask=topk_mask,
-            )
-
-            swa_offset = offset - topk_len
-            swa_mask = (offset >= topk_len) & (swa_offset < swa_len)
-            tl.store(
-                combined_ragged_ptr + out_start + offset,
-                M * batch_idx + N + swa_offset + pos - swa_len + 1 - gather_start,
-                mask=swa_mask,
-            )
-
-
-def combine_topk_swa_indices_ragged(
-    topk_indices: torch.Tensor,
-    query_start_loc: torch.Tensor,
-    seq_lens: torch.Tensor,
-    gather_lens: torch.Tensor,
-    window_size: int,
-    compress_ratio: int,
-    topk: int,
-    M: int,
-    N: int,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    topk_indices = topk_indices.reshape(topk_indices.shape[0], -1).contiguous()
-    num_tokens = topk_indices.shape[0]
-    num_reqs = seq_lens.shape[0]
-    combined_lens = torch.empty(
-        num_tokens, dtype=torch.int32, device=topk_indices.device
-    )
-
-    num_workers = 128
-    _compute_combined_lens_kernel[(num_reqs, num_workers)](
-        combined_lens,
-        query_start_loc,
-        seq_lens,
-        TOP_K=topk,
-        COMPRESS_RATIO=compress_ratio,
-        WINDOW_SIZE=window_size,
-    )
-
-    combined_indptr = _build_indptr_from_lengths(combined_lens)
-    combined_ragged = torch.empty(
-        num_tokens * (topk + window_size),
-        dtype=torch.int32,
-        device=topk_indices.device,
-    )
-    if combined_ragged.numel() > 0:
-        block = 128
-        _combine_topk_swa_indices_ragged_kernel[
-            (num_reqs, num_workers, triton.cdiv(topk + window_size, block))
-        ](
-            combined_ragged,
-            combined_indptr,
-            topk_indices,
-            topk_indices.stride(0),
-            query_start_loc,
-            seq_lens,
-            gather_lens,
-            M,
-            N,
-            topk_indices.shape[-1],
-            TOP_K=topk,
-            COMPRESS_RATIO=compress_ratio,
-            WINDOW_SIZE=window_size,
-            BLOCK_SIZE=block,
-        )
-    return combined_ragged, combined_indptr, combined_lens
-
-
 def _copy_ragged_to_graph_buffers(
     ragged_indices: torch.Tensor,
     ragged_indptr: torch.Tensor,
@@ -518,8 +371,12 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuild
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         max_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens
+        # The non-causal (DSpark draft) path widens each token's SWA index list
+        # to ``noncausal_index_width`` (>= window_size), so size the persistent
+        # ragged buffer to the wider bound to cover both causal and non-causal.
+        swa_index_width = max(self.window_size, self.noncausal_index_width)
         self.decode_swa_ragged_indices_buffer = torch.empty(
-            max_tokens * self.window_size,
+            max_tokens * swa_index_width,
             dtype=torch.int32,
             device=self.device,
         )
@@ -558,7 +415,9 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuild
                 self.decode_swa_ragged_indices_buffer,
                 self.decode_swa_ragged_indptr_buffer,
                 base.num_decode_tokens,
-                self.window_size,
+                # Actual dense width for this build: window_size (causal) or
+                # noncausal_index_width (DSpark non-causal draft).
+                base.decode_swa_indices.shape[-1],
             )
 
         return DeepseekV4ROCMAiterSparseSWAMetadata(

vllm/reasoning/__init__.py

@@ -84,6 +84,10 @@ _REASONING_PARSERS_TO_REGISTER = {
         "kimi_k2_reasoning_parser",
         "KimiK2ReasoningParser",
     ),
+    "longcat": (
+        "longcat_reasoning_parser",
+        "LongcatReasoningParser",
+    ),
     "mimo": (
         "qwen3_engine_reasoning_parser",
         "Qwen3ParserReasoningAdapter",

vllm/reasoning/longcat_reasoning_parser.py

@@ -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

vllm/transformers_utils/config.py

@@ -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(

vllm/transformers_utils/configs/__init__.py

@@ -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",

vllm/transformers_utils/configs/longcat_flash.py

@@ -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

vllm/transformers_utils/model_arch_config_convertor.py

@@ -268,6 +268,7 @@ class ModelArchConfigConvertorBase:
             "kimi_k2",
             "kimi_linear",
             "longcat_flash",
+            "longcat_flash_ngram",
             "pangu_ultra_moe",
             "pangu_ultra_moe_mtp",
             "bailing_hybrid",

vllm/utils/deep_gemm.py

@@ -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",
 }
 
 

vllm/v1/attention/ops/triton_turboquant_store.py

@@ -188,7 +188,7 @@ def _tq_fused_store_fp8(
     # ── FP8 KEY: cast to FP8 in-kernel and store ─────────────────
     d_offs = tl.arange(0, BLOCK_D)
     d_mask = d_offs < D
-    k_vals = tl.load(Key_ptr + base + d_offs, mask=d_mask, other=0.0)
+    k_vals = tl.load(Key_ptr + base + d_offs, mask=d_mask, other=0.0).to(tl.float32)
     k_fp8 = k_vals.to(tl.float8e4b15) if FP8_E4B15 else k_vals.to(tl.float8e4nv)
     k_bytes = k_fp8.to(tl.uint8, bitcast=True)
     tl.store(KV_cache_ptr + slot_base + d_offs, k_bytes, mask=d_mask)

vllm/v1/worker/gpu/attn_utils.py

@@ -540,7 +540,15 @@ def init_kv_cache(
         shared_kv_cache_layers=shared_kv_cache_layers,
         kv_cache_config=kv_cache_config,
     )
-    bind_kv_cache(kv_caches, forward_context, runner_kv_caches)
+    # Dual-attention models (e.g. LongCat-Flash) put two Attention modules per
+    # decoder layer, so a layer name carries two integers (layer + module index).
+    num_attn_module = (
+        2
+        if vllm_config.model_config.hf_config.model_type
+        in ("longcat_flash", "longcat_flash_ngram")
+        else 1
+    )
+    bind_kv_cache(kv_caches, forward_context, runner_kv_caches, num_attn_module)
     return kv_caches