minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Model] Support LongCat 2.0 FP8 (#30275)

Browse Source 
Co-authored-by: sunjiaqi11 <sunjiaqi11@meituan.com>
Co-authored-by: Ke Bao <ispobaoke@gmail.com>
This commit is contained in:
Xiaoyu Zhang
2026-07-07 19:51:12 +08:00
committed by GitHub
co-authored by sunjiaqi11 Ke Bao
parent cfd3fdc54f
commit e339c83f82
23 changed files with 481 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/ngram_embedding.cuh
+25 -9
View File
@@ -30,7 +30,8 @@ __global__ void ComputeNGramIdsKernel(
int max_context_len, // max_context_len
const int64_t* __restrict__ row_indices, // [batch_size]
int* column_starts, // [batch_size]
int* n_gram_ids // [ne_n-1,ne_k,token_num]
int* n_gram_ids, // [ne_n-1,ne_k,token_num]
int eos_token_id // tokens before an eos are excluded from the n-gram context
) {
// Determine which n, k, and request this block handles.
/**
@@ -73,12 +74,17 @@ __global__ void ComputeNGramIdsKernel(
// Out of this request's range, stop computing n_gram_id
break;
}
if (ne_token_table[current_token_table_index - j] < 0) {
const int table_token = ne_token_table[current_token_table_index - j];
if (table_token < 0) {
// Token was marked as ignored during write
break;
}
const uint64_t term =
(uint64_t)ne_token_table[current_token_table_index - j] * (uint64_t)ne_weights[ne_weight_base_idx + j];
if (table_token == eos_token_id && j > 0) {
// Don't let the n-gram context cross an eos boundary. j==0 (the
// current token) is allowed; only break when looking back.
break;
}
const uint64_t term = (uint64_t)table_token * (uint64_t)ne_weights[ne_weight_base_idx + j];
n_gram_id += term % ne_mod;
}
n_gram_id %= ne_mod;
@@ -99,7 +105,8 @@ __global__ void ComputeNGramIdsDecodeKernel(
int max_context_len, // max_context_len
const int64_t* __restrict__ row_indices, // [batch_size]
const int* __restrict__ column_starts, // [batch_size]
int* __restrict__ n_gram_ids // [batch_size, (ne_n-1)*ne_k]
int* __restrict__ n_gram_ids, // [batch_size, (ne_n-1)*ne_k]
int eos_token_id // tokens before an eos are excluded from the n-gram context
) {
const int num_configs = (ne_n - 1) * ne_k;
const int total_outputs = batch_size * num_configs;
@@ -124,6 +131,11 @@ __global__ void ComputeNGramIdsDecodeKernel(
if (token < 0) {
break;
}
if (token == eos_token_id && j > 0) {
// Don't let the n-gram context cross an eos boundary. j==0 (the
// current token) is allowed; only break when looking back.
break;
}
const uint64_t term = static_cast<uint64_t>(token) * static_cast<uint64_t>(ne_weights[weight_offset + j]);
n_gram_id += term % ne_mod;
}
@@ -201,7 +213,8 @@ struct NgramEmbeddingKernel {
const tvm::ffi::TensorView ne_token_table,
const tvm::ffi::TensorView row_indices,
const tvm::ffi::TensorView column_starts,
const tvm::ffi::TensorView n_gram_ids) {
const tvm::ffi::TensorView n_gram_ids,
const int64_t eos_token_id) {
using namespace host;
auto device_ = SymbolicDevice{};
@@ -274,7 +287,8 @@ struct NgramEmbeddingKernel {
max_context_len,
static_cast<const int64_t*>(row_indices.data_ptr()),
static_cast<int*>(column_starts.data_ptr()),
static_cast<int*>(n_gram_ids.data_ptr()));
static_cast<int*>(n_gram_ids.data_ptr()),
static_cast<int>(eos_token_id));
}
static void compute_n_gram_ids_decode(
@@ -286,7 +300,8 @@ struct NgramEmbeddingKernel {
const tvm::ffi::TensorView ne_token_table,
const tvm::ffi::TensorView row_indices,
const tvm::ffi::TensorView column_starts,
const tvm::ffi::TensorView n_gram_ids) {
const tvm::ffi::TensorView n_gram_ids,
const int64_t eos_token_id) {
using namespace host;
auto device_ = SymbolicDevice{};
@@ -354,7 +369,8 @@ struct NgramEmbeddingKernel {
max_context_len,
static_cast<const int64_t*>(row_indices.data_ptr()),
static_cast<const int*>(column_starts.data_ptr()),
static_cast<int*>(n_gram_ids.data_ptr()));
static_cast<int*>(n_gram_ids.data_ptr()),
static_cast<int>(eos_token_id));
}
static void update_token_table(
python/sglang/jit_kernel/ngram_embedding.py
+5
View File
@@ -43,6 +43,7 @@ def compute_n_gram_ids(
row_indices: torch.Tensor,
column_starts: torch.Tensor,
n_gram_ids: torch.Tensor,
eos_token_id: int,
) -> None:
"""
Compute n-gram IDs for embedding.
@@ -59,6 +60,7 @@ def compute_n_gram_ids(
row_indices: row indices for each request
column_starts: column start positions for each request
n_gram_ids: output tensor for n-gram ids
eos_token_id: tokens before an eos are excluded from the n-gram context
"""
module = _jit_ngram_embedding_module()
module.compute_n_gram_ids(
@@ -73,6 +75,7 @@ def compute_n_gram_ids(
row_indices,
column_starts,
n_gram_ids,
eos_token_id,
)
@@ -87,6 +90,7 @@ def compute_n_gram_ids_decode(
row_indices: torch.Tensor,
column_starts: torch.Tensor,
n_gram_ids: torch.Tensor,
eos_token_id: int,
) -> None:
"""
Compute n-gram IDs for decode, where each request contributes one token.
@@ -102,6 +106,7 @@ def compute_n_gram_ids_decode(
row_indices,
column_starts,
n_gram_ids,
eos_token_id,
)
python/sglang/srt/arg_groups/overrides.py
+15
View File
@@ -313,6 +313,8 @@ def _register_for(*architectures: str):
"MistralLarge3ForCausalLM",
"PixtralForConditionalGeneration",
"GlmMoeDsaForCausalLM",
"LongcatFlashForCausalLM",
"LongcatFlashForCausalLMNextN",
)
def _deepseek_family_overrides(server_args: Any, hf_config: Any) -> dict:
"""Order-safe declarations of the DeepSeek/DSA branch. The CP parallel
@@ -1170,6 +1172,8 @@ _DEEPSEEK_FAMILY_ARCHS = frozenset(
"MistralLarge3ForCausalLM",
"PixtralForConditionalGeneration",
"GlmMoeDsaForCausalLM",
"LongcatFlashForCausalLM",
"LongcatFlashForCausalLMNextN",
}
)
@@ -1243,6 +1247,17 @@ def _deepseek_moe_quant_resolution(view: Any) -> dict:
logger.info(
"Use flashinfer_trtllm as MoE runner backend on sm100 for DeepseekV3ForCausalLM"
)
if (
model_arch in ["LongcatFlashForCausalLM", "LongcatFlashForCausalLMNextN"]
and view.fp8_gemm_runner_backend == "auto"
and quantization in ["fp8", "modelopt_fp8"]
and quant_cfg.get("scale_fmt", None) != "ue8m0"
):
overrides["fp8_gemm_runner_backend"] = "flashinfer_trtllm"
logger.info(
"Use flashinfer_trtllm as FP8 GEMM backend on Blackwell for LongCat FP8 "
"checkpoint with non-ue8m0 scales"
)
return overrides
python/sglang/srt/configs/longcat_flash.py
+12
View File
@@ -56,6 +56,9 @@ class LongcatFlashConfig(PretrainedConfig):
ngram_vocab_size_ratio=None,
emb_neighbor_num=None,
emb_split_num=None,
oe_vocab_size_ratio=None,
oe_neighbor_num=None,
oe_split_num=None,
**kwargs,
):
super().__init__(
@@ -105,6 +108,15 @@ class LongcatFlashConfig(PretrainedConfig):
self.zero_expert_type = zero_expert_type
self.routed_scaling_factor = routed_scaling_factor
self.hidden_act = "silu"
if ngram_vocab_size_ratio is None:
ngram_vocab_size_ratio = oe_vocab_size_ratio
if emb_neighbor_num is None:
emb_neighbor_num = oe_neighbor_num
if emb_split_num is None:
emb_split_num = oe_split_num
self.oe_vocab_size_ratio = oe_vocab_size_ratio
self.oe_neighbor_num = oe_neighbor_num
self.oe_split_num = oe_split_num
self.use_ngram_embedding = ngram_vocab_size_ratio is not None
if self.use_ngram_embedding:
self.ngram_embedding_m = int(ngram_vocab_size_ratio * vocab_size)
python/sglang/srt/configs/model_config.py
+2
View File
@@ -110,6 +110,8 @@ def is_deepseek_dsa(config) -> bool:
"MistralLarge3ForCausalLM",
"PixtralForConditionalGeneration",
"GlmMoeDsaForCausalLM",
"LongcatFlashForCausalLM",
"LongcatFlashForCausalLMNextN",
)
and _hf_attr(config, "index_topk") is not None
)
python/sglang/srt/layers/attention/dsa/dsa_indexer.py
+80 -12
View File
@@ -31,7 +31,7 @@ from sglang.srt.layers.attention.dsa.utils import (
is_graph_dsa_split_op_surface,
)
from sglang.srt.layers.dp_attention import attn_tp_all_gather_into_tensor
from sglang.srt.layers.layernorm import LayerNorm
from sglang.srt.layers.layernorm import LayerNorm, RMSNorm
from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype, is_fp8_fnuz
from sglang.srt.layers.utils import MultiPlatformOp
from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
@@ -364,6 +364,7 @@ class Indexer(MultiPlatformOp):
prefix: str = "",
quant_config: Optional[QuantizationConfig] = None,
alt_stream: Optional[torch.cuda.Stream] = None,
config=None,
):
super().__init__()
self.hidden_size = hidden_size
@@ -425,9 +426,15 @@ class Indexer(MultiPlatformOp):
params_dtype=torch.bfloat16,
prefix=add_prefix("weights_proj", prefix),
)
self.k_norm = LayerNorm(
self.head_dim, dtype=torch.bfloat16 if _use_aiter else torch.float32
)
if (
config is not None
and getattr(config, "index_k_norm_type", "layer") == "rms"
):
self.k_norm = RMSNorm(self.head_dim)
else:
self.k_norm = LayerNorm(
self.head_dim, dtype=torch.bfloat16 if _use_aiter else torch.float32
)
self.rotary_emb = get_rope_wrapper(
rope_head_dim,
rotary_dim=rope_head_dim,
@@ -440,6 +447,10 @@ class Indexer(MultiPlatformOp):
self.block_size = block_size
self.scale_fmt = scale_fmt
self.softmax_scale = self.head_dim**-0.5
self.num_init_tokens = self.num_local_tokens = 0
if config is not None:
self.num_init_tokens = getattr(config, "index_init_tokens", 0)
self.num_local_tokens = getattr(config, "index_local_tokens", 0)
self.paged_mqa_logits_backend = DSAPagedMQALogitsBackend.resolve(
get_server_args().dsa_paged_mqa_logits_backend
@@ -786,6 +797,52 @@ class Indexer(MultiPlatformOp):
return
dst.copy_(src)
@staticmethod
def _pad_heads_for_deep_gemm(q_fp8, weights):
"""Pad q and weights to 32 heads when num_heads < 32,
so that block_q = 128/num_heads doesn't exceed seq_len_alignment(4)."""
num_heads = q_fp8.shape[1]
if num_heads >= 32:
return q_fp8, weights, num_heads
target_heads = 32
q_fp8 = torch.nn.functional.pad(q_fp8, (0, 0, 0, target_heads - num_heads))
weights = torch.nn.functional.pad(weights, (0, target_heads - num_heads))
return q_fp8, weights, num_heads
def _mask_init_and_local_tokens(
self,
logits: torch.Tensor,
lengths: torch.Tensor,
row_starts: Optional[torch.Tensor] = None,
):
if self.num_init_tokens == 0 and self.num_local_tokens == 0:
return logits
if row_starts is None:
row_starts = lengths.new_zeros(lengths.shape[0])
num_init_tokens = self.num_init_tokens
num_local_tokens = self.num_local_tokens
if num_init_tokens > 0:
init_idxs = (
torch.arange(
num_init_tokens, dtype=lengths.dtype, device=lengths.device
)[None, :]
+ row_starts[:, None]
)
init_idxs.clamp_max_(logits.shape[-1] - 1)
logits.scatter_(dim=1, index=init_idxs, value=float("inf"))
if num_local_tokens > 0:
local_idxs = (
lengths[:, None]
- 1
+ row_starts[:, None]
- torch.arange(
num_local_tokens, dtype=lengths.dtype, device=lengths.device
)[None, :]
)
local_idxs.clamp_min_(0)
logits.scatter_(dim=1, index=local_idxs, value=float("inf"))
return logits
def _get_topk_paged(
self,
forward_batch: ForwardBatch,
@@ -947,6 +1004,7 @@ class Indexer(MultiPlatformOp):
)
# NOTE(dark): logits should be cleaned in topk_transform
self._mask_init_and_local_tokens(logits, seqlens_32)
topk_result = metadata.topk_transform(logits, self.index_topk)
# Restore possible padding exist in the hidden states.
if not _is_hip and q_offset < q_fp8.shape[0]:
@@ -1121,10 +1179,13 @@ class Indexer(MultiPlatformOp):
clean_logits=False,
)
else:
q_padded, w_padded, _ = self._pad_heads_for_deep_gemm(
q_fp8[:q_offset], weights[:q_offset]
)
logits = deep_gemm.fp8_mqa_logits(
q_fp8[:q_offset],
q_padded,
kv_fp8,
weights[:q_offset],
w_padded,
ks,
ke,
clean_logits=False,
@@ -1132,6 +1193,7 @@ class Indexer(MultiPlatformOp):
assert logits.shape[0] == len(seq_lens_expanded)
assert logits.shape[1] == k_offset
self._mask_init_and_local_tokens(logits, seq_lens_expanded, ks)
raw_topk_result = metadata.topk_transform(logits, self.index_topk, ks=ks)
topk_result[:q_offset] = raw_topk_result
return topk_result
@@ -1173,16 +1235,20 @@ class Indexer(MultiPlatformOp):
clean_logits=False,
)
else:
q_padded, w_padded, _ = self._pad_heads_for_deep_gemm(
q_fp8[start:end], weights[start:end]
)
logits_chunk = deep_gemm.fp8_mqa_logits(
q_fp8[start:end],
q_padded,
kv_fp8,
weights[start:end],
w_padded,
ks[start:end],
ke[start:end],
clean_logits=False,
)
lengths_chunk = seq_lens_expanded[start:end]
self._mask_init_and_local_tokens(logits_chunk, lengths_chunk, ks[start:end])
# RAGGED: use global offset; PAGED: construct local cu_seqlens_q per chunk
if global_topk_offset is not None:
@@ -1379,10 +1445,11 @@ class Indexer(MultiPlatformOp):
ke = ks + ke_offset
actual_seq_q = torch.cat(actual_seq_q_list, dim=0)
with self._with_real_sm_count():
q_padded, w_padded, _ = self._pad_heads_for_deep_gemm(q_fp8, weights)
logits = deep_gemm.fp8_mqa_logits(
q_fp8,
q_padded,
kv_fp8,
weights,
w_padded,
ks,
ke,
clean_logits=False,
@@ -1425,10 +1492,11 @@ class Indexer(MultiPlatformOp):
ke = ks + ke_offset
with self._with_real_sm_count():
q_padded, w_padded, _ = self._pad_heads_for_deep_gemm(q_fp8, weights)
logits = deep_gemm.fp8_mqa_logits(
q_fp8,
q_padded,
kv_fp8,
weights,
w_padded,
ks,
ke,
clean_logits=False,
python/sglang/srt/layers/n_gram_embedding.py
+28 -40
View File
@@ -2,10 +2,7 @@ import torch
from torch import nn
from torch.nn import Parameter
from sglang.jit_kernel.ngram_embedding import (
compute_n_gram_ids,
compute_n_gram_ids_decode,
)
from sglang.jit_kernel.ngram_embedding import compute_n_gram_ids
from sglang.srt.layers.dp_attention import is_dp_attention_enabled
from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@@ -20,6 +17,7 @@ class NgramEmbedding(torch.nn.Module):
over_embedding_m: int,
over_embedding_k: int,
over_embedding_n: int,
eos_token_id: int,
):
super().__init__()
assert (
@@ -30,11 +28,13 @@ class NgramEmbedding(torch.nn.Module):
self.over_embedding_m = over_embedding_m
self.over_embedding_k = over_embedding_k
self.over_embedding_n = over_embedding_n
self.eos_token_id = eos_token_id
use_attn_tp_group = is_dp_attention_enabled()
self.word_embeder = VocabParallelEmbedding(
num_embeddings,
embedding_dim,
enable_tp=is_dp_attention_enabled(),
use_attn_tp_group=use_attn_tp_group,
)
self.n_grams = (over_embedding_n - 1) * over_embedding_k
oe_hidden_dim = embedding_dim // (over_embedding_k * (over_embedding_n - 1))
@@ -51,7 +51,7 @@ class NgramEmbedding(torch.nn.Module):
self.oe_embeder = VocabParallelEmbedding(
num_embeddings=self.exclusive_oe_embedder_size_sums[-1],
embedding_dim=oe_hidden_dim,
enable_tp=is_dp_attention_enabled(),
use_attn_tp_group=use_attn_tp_group,
)
self.oe_projection = nn.Parameter(
@@ -138,40 +138,28 @@ class NgramEmbedding(torch.nn.Module):
or forward_batch.forward_mode.is_decode()
):
ngram_embedding_info = forward_batch.ngram_embedding_info
if forward_batch.forward_mode.is_decode():
compute_n_gram_ids_decode(
ne_n=self.over_embedding_n,
ne_k=self.over_embedding_k,
ne_weights=self.oe_weights,
ne_mods=self.oe_mods,
exclusive_ne_embedder_size_sums=self.exclusive_oe_embedder_size_sums,
ne_token_table=ngram_embedding_info.token_table,
row_indices=forward_batch.req_pool_indices,
column_starts=ngram_embedding_info.column_starts,
n_gram_ids=self.oe_n_gram_ids[: len(input_ids)],
)
else:
torch.cumsum(
ngram_embedding_info.req_lens,
dim=0,
dtype=torch.int32,
out=self.exclusive_req_len_sums[1 : 1 + forward_batch.batch_size],
)
compute_n_gram_ids(
ne_n=self.over_embedding_n,
ne_k=self.over_embedding_k,
ne_weights=self.oe_weights,
ne_mods=self.oe_mods,
tokens=input_ids.to(torch.int32),
exclusive_ne_embedder_size_sums=self.exclusive_oe_embedder_size_sums,
exclusive_req_len_sums=self.exclusive_req_len_sums[
: forward_batch.batch_size + 1
],
ne_token_table=ngram_embedding_info.token_table,
row_indices=forward_batch.req_pool_indices,
column_starts=ngram_embedding_info.column_starts,
n_gram_ids=self.oe_n_gram_ids[: len(input_ids)],
)
torch.cumsum(
ngram_embedding_info.req_lens,
dim=0,
dtype=torch.int32,
out=self.exclusive_req_len_sums[1 : 1 + forward_batch.batch_size],
)
compute_n_gram_ids(
ne_n=self.over_embedding_n,
ne_k=self.over_embedding_k,
ne_weights=self.oe_weights,
ne_mods=self.oe_mods,
tokens=input_ids.to(torch.int32),
exclusive_ne_embedder_size_sums=self.exclusive_oe_embedder_size_sums,
exclusive_req_len_sums=self.exclusive_req_len_sums[
: forward_batch.batch_size + 1
],
ne_token_table=ngram_embedding_info.token_table,
row_indices=forward_batch.req_pool_indices,
column_starts=ngram_embedding_info.column_starts,
n_gram_ids=self.oe_n_gram_ids[: len(input_ids)],
eos_token_id=self.eos_token_id,
)
# [13, seq_len, hidden_dim]
all_hidden_states = torch.empty(
python/sglang/srt/managers/schedule_batch.py
+3
View File
@@ -1744,6 +1744,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# Read by ForwardBatch ngram embedding init
ne_token_table: torch.Tensor = None
# Mask marking chunked (not-yet-finished) prefill requests whose sampled
# pseudo next-token must NOT be written into the ngram token table.
ne_skip_token_table_update: torch.Tensor = None
req_pool_indices: torch.Tensor = None # shape: [b], int64
seq_lens: torch.Tensor = None # shape: [b], int64
python/sglang/srt/managers/scheduler.py
+15
View File
@@ -1342,6 +1342,21 @@ class Scheduler(
),
ignore_tokens=None,
)
# Mark the chunked (not-yet-finished) prefill request so sample()
# skips writing its pseudo next-token into the ngram token table.
# Use self.chunked_req identity (not req.is_chunked) to avoid
# overlap-scheduling timing issues.
if self.chunked_req is not None:
skip_token_table_update = [
req is self.chunked_req for req in batch.reqs
]
batch.ne_skip_token_table_update = (
torch.tensor(
skip_token_table_update, dtype=torch.bool, device=device
)
if any(skip_token_table_update)
else None
)
return batch
def init_deterministic_inference_config(self):
python/sglang/srt/model_executor/forward_batch_info.py
+11
View File
@@ -286,6 +286,9 @@ class NgramEmbeddingInfo:
req_lens: torch.Tensor
out_column_starts: torch.Tensor
out_req_lens: torch.Tensor
# Mask marking chunked (not-yet-finished) prefill requests whose sampled
# pseudo next-token must NOT be written into the token table.
skip_token_table_update: Optional[torch.Tensor] = None
@classmethod
def create(
@@ -295,6 +298,7 @@ class NgramEmbeddingInfo:
device: torch.device,
column_starts=None,
req_lens=None,
skip_token_table_update=None,
) -> NgramEmbeddingInfo:
info = cls(
token_table=token_table,
@@ -302,6 +306,7 @@ class NgramEmbeddingInfo:
req_lens=torch.empty(batch_size, dtype=torch.int32, device=device),
out_column_starts=torch.empty(batch_size, dtype=torch.int32, device=device),
out_req_lens=torch.empty(batch_size, dtype=torch.int32, device=device),
skip_token_table_update=skip_token_table_update,
)
if column_starts is not None:
info.column_starts[:] = column_starts
@@ -316,6 +321,11 @@ class NgramEmbeddingInfo:
req_lens=self.req_lens[:bs],
out_column_starts=self.out_column_starts[:bs],
out_req_lens=self.out_req_lens[:bs],
skip_token_table_update=(
self.skip_token_table_update[:bs]
if self.skip_token_table_update is not None
else None
),
)
@@ -1003,6 +1013,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
device,
column_starts=column_starts,
req_lens=req_lens,
skip_token_table_update=batch.ne_skip_token_table_update,
)
def compute_spec_mrope_positions(
python/sglang/srt/model_executor/model_runner.py
+9 -10
View File
@@ -32,7 +32,6 @@ import torch
import torch.distributed as dist
from torch import nn
from sglang.jit_kernel.ngram_embedding import update_token_table_decode
from sglang.srt.configs import (
BailingHybridConfig,
FalconH1Config,
@@ -165,6 +164,9 @@ from sglang.srt.model_executor.hook_manager import register_forward_hooks
from sglang.srt.model_executor.model_runner_kv_cache_mixin import (
ModelRunnerKVCacheMixin,
)
from sglang.srt.model_executor.ngram_token_table import (
update_ngram_token_table_after_sampling,
)
from sglang.srt.model_executor.pool_configurator import MemoryPoolConfig
from sglang.srt.model_executor.runner import (
EagerRunner,
@@ -2606,15 +2608,12 @@ class ModelRunner(ModelRunnerKVCacheMixin):
ngram_embedding_info = forward_batch.ngram_embedding_info
if ngram_embedding_info is None:
return
ngram_embedding_info.out_column_starts[: forward_batch.batch_size] = (
forward_batch.seq_lens
)
ngram_embedding_info.out_req_lens[: forward_batch.batch_size] = 1
update_token_table_decode(
ne_token_table=ngram_embedding_info.token_table,
tokens=next_token_ids.to(torch.int32),
row_indices=forward_batch.req_pool_indices,
column_starts=ngram_embedding_info.out_column_starts,
update_ngram_token_table_after_sampling(
ngram_embedding_info=ngram_embedding_info,
next_token_ids=next_token_ids,
req_pool_indices=forward_batch.req_pool_indices,
seq_lens=forward_batch.seq_lens,
batch_size=forward_batch.batch_size,
)
def init_decode_cuda_graph(self):
python/sglang/srt/model_executor/ngram_token_table.py
+51
View File
@@ -0,0 +1,51 @@
"""Utilities for updating LongCat ngram embedding token tables."""
from __future__ import annotations
import torch
from sglang.jit_kernel.ngram_embedding import update_token_table
def update_ngram_token_table_after_sampling(
*,
ngram_embedding_info,
next_token_ids: torch.Tensor,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
batch_size: int,
) -> bool:
"""Update the ngram token table with sampled tokens.
Returns whether the token table was updated.
"""
skip_token_table_update = ngram_embedding_info.skip_token_table_update
if skip_token_table_update is not None:
# Skip chunked (not-yet-finished) prefill requests: their sampled token
# is a pseudo prediction and must not pollute the token table.
indices = (~skip_token_table_update).nonzero(as_tuple=True)[0]
if indices.numel() == 0:
return False
update_token_table(
ne_token_table=ngram_embedding_info.token_table,
tokens=next_token_ids[indices].to(torch.int32),
row_indices=req_pool_indices[indices],
column_starts=seq_lens[indices].to(torch.int32),
req_lens=torch.ones(
indices.numel(), dtype=torch.int32, device=next_token_ids.device
),
ignore_tokens=None,
)
return True
ngram_embedding_info.out_column_starts[:batch_size] = seq_lens
ngram_embedding_info.out_req_lens[:batch_size] = 1
update_token_table(
ne_token_table=ngram_embedding_info.token_table,
tokens=next_token_ids.to(torch.int32),
row_indices=req_pool_indices,
column_starts=ngram_embedding_info.out_column_starts,
req_lens=ngram_embedding_info.out_req_lens,
ignore_tokens=None,
)
return True
python/sglang/srt/model_executor/runner/base_runner.py
+6
View File
@@ -147,6 +147,7 @@ def _allocate_decode_buffers(
req_lens=torch.ones([max_bs], dtype=torch.int32),
out_column_starts=torch.zeros([max_bs], dtype=torch.int32),
out_req_lens=torch.ones([max_bs], dtype=torch.int32),
skip_token_table_update=torch.zeros([max_bs], dtype=torch.bool),
)
if ne_token_table is not None
else None
@@ -297,6 +298,7 @@ class BaseRunner(ABC):
num_tokens_per_bs=num_tokens_per_bs,
cache_loc_dtype=torch.int64,
enable_mamba_track=False,
ne_token_table=mr.token_table if mr.use_ngram_embedding else None,
hc_hidden_size=getattr(mr.model_config, "hc_hidden_size", None),
pp_proxy_topk_size=mr.get_pp_proxy_topk_size(),
)
@@ -524,6 +526,10 @@ class BaseRunner(ABC):
global_forward_mode=capture_forward_mode,
lora_ids=lora_ids,
)
if buffers.ngram_embedding_info is not None:
forward_batch.ngram_embedding_info = buffers.ngram_embedding_info.slice(
batch_size
)
if lora_ids is not None:
mr.lora_manager.prepare_lora_batch(forward_batch)
python/sglang/srt/model_executor/runner_utils/buffers.py
+1
View File
@@ -169,6 +169,7 @@ class DecodeInputBuffers(ForwardInputBuffers):
req_lens=torch.ones([max_bs], dtype=torch.int32),
out_column_starts=torch.zeros([max_bs], dtype=torch.int32),
out_req_lens=torch.ones([max_bs], dtype=torch.int32),
skip_token_table_update=torch.zeros([max_bs], dtype=torch.bool),
)
if ne_token_table is not None
else None
python/sglang/srt/models/deepseek_v2.py
+8 -2
View File
@@ -1654,6 +1654,7 @@ class DeepseekV2AttentionMLA(
quant_config=quant_config,
layer_id=layer_id,
alt_stream=alt_stream,
config=config,
)
# Refer: https://arxiv.org/abs/2603.12201 for more details.
# skip_topk: when True, this layer will skip computation and reuse previous layer's topk indices.
@@ -1662,8 +1663,13 @@ class DeepseekV2AttentionMLA(
self.skip_topk = True
self.next_skip_topk = True
else:
self.skip_topk = dsa_layer_skips_topk(config, layer_id)
self.next_skip_topk = dsa_layer_skips_topk(config, layer_id + 1)
index_cli_factor = getattr(config, "cli_factor", 1)
if index_cli_factor > 1:
self.skip_topk = layer_id % index_cli_factor != 0
self.next_skip_topk = (layer_id + 1) % index_cli_factor != 0
else:
self.skip_topk = dsa_layer_skips_topk(config, layer_id)
self.next_skip_topk = dsa_layer_skips_topk(config, layer_id + 1)
self.kv_b_proj = ColumnParallelLinear(
self.kv_lora_rank,
python/sglang/srt/models/longcat_flash.py
+40 -11
View File
@@ -326,8 +326,8 @@ class LongcatFlashDecoderLayer(nn.Module):
v_head_dim=config.v_head_dim,
q_lora_rank=config.q_lora_rank,
kv_lora_rank=config.kv_lora_rank,
rope_theta=config.rope_parameters["rope_theta"],
rope_scaling=None,
rope_theta=config.rope_theta,
rope_scaling=config.rope_scaling,
max_position_embeddings=config.max_position_embeddings,
quant_config=(
None
@@ -420,18 +420,24 @@ class LongcatFlashDecoderLayer(nn.Module):
forward_batch: ForwardBatch,
residual: Optional[torch.Tensor],
zero_allocator: BumpAllocator,
prev_topk_indices: Optional[torch.Tensor],
) -> torch.Tensor:
# first_attn
hidden_states, residual = self.moe_layer_communicator.prepare_attn(
hidden_states, residual, forward_batch
)
if hidden_states.shape[0] != 0:
hidden_states = self.self_attn[0](
attn_out = self.self_attn[0](
positions=positions,
hidden_states=hidden_states,
forward_batch=forward_batch,
zero_allocator=zero_allocator,
prev_topk_indices=prev_topk_indices,
)
if isinstance(attn_out, tuple):
hidden_states, prev_topk_indices = attn_out
else:
hidden_states = attn_out
# moe
hidden_states, residual = self.moe_layer_communicator.prepare_mlp(
@@ -444,15 +450,26 @@ class LongcatFlashDecoderLayer(nn.Module):
moe_hidden_states, moe_residual, forward_batch
)
hidden_states, residual = self.forward_mlp(
hidden_states, positions, residual, forward_batch, zero_allocator
hidden_states, residual, prev_topk_indices = self.forward_mlp(
hidden_states,
positions,
residual,
forward_batch,
zero_allocator,
prev_topk_indices,
)
hidden_states = moe_hidden_states + hidden_states
return hidden_states, residual
return hidden_states, residual, prev_topk_indices
def forward_mlp(
self, hidden_states, positions, residual, forward_batch, zero_allocator
self,
hidden_states,
positions,
residual,
forward_batch,
zero_allocator,
prev_topk_indices,
):
# first_mlp
hidden_states = self.mlps[0](hidden_states)
@@ -464,12 +481,17 @@ class LongcatFlashDecoderLayer(nn.Module):
hidden_states, residual, forward_batch
)
if hidden_states.shape[0] != 0:
hidden_states = self.self_attn[1](
attn_out = self.self_attn[1](
positions=positions,
hidden_states=hidden_states,
forward_batch=forward_batch,
zero_allocator=zero_allocator,
prev_topk_indices=prev_topk_indices,
)
if isinstance(attn_out, tuple):
hidden_states, prev_topk_indices = attn_out
else:
hidden_states = attn_out
# second_mlp
hidden_states, residual = self.mlp_layer_communicator[1].prepare_mlp(
@@ -483,7 +505,7 @@ class LongcatFlashDecoderLayer(nn.Module):
hidden_states, residual, forward_batch
)
return hidden_states, residual
return hidden_states, residual, prev_topk_indices
class LongcatFlashModel(nn.Module):
@@ -506,6 +528,7 @@ class LongcatFlashModel(nn.Module):
over_embedding_m=config.ngram_embedding_m,
over_embedding_k=config.ngram_embedding_k,
over_embedding_n=config.ngram_embedding_n,
eos_token_id=config.eos_token_id,
)
else:
self.use_ngram_embedding = False
@@ -559,13 +582,19 @@ class LongcatFlashModel(nn.Module):
residual = None
aux_hidden_states = []
topk_indices = None
for i in range(total_num_layers):
if i in self.layers_to_capture:
aux_hidden_states.append(hidden_states + residual)
with get_global_expert_distribution_recorder().with_current_layer(i):
layer = self.layers[i]
hidden_states, residual = layer(
positions, hidden_states, forward_batch, residual, zero_allocator
hidden_states, residual, topk_indices = layer(
positions,
hidden_states,
forward_batch,
residual,
zero_allocator,
topk_indices,
)
if hidden_states.shape[0] != 0:
python/sglang/srt/runtime_context.py
+1
View File
@@ -331,6 +331,7 @@ class Flags(_StaticFlags):
sampling_backend: str | None = None
page_size: int | None = None
quantization: str | None = None
fp8_gemm_runner_backend: str = "auto"
disable_overlap_schedule: bool = False
uses_mamba_radix_cache: bool = False
mamba_radix_cache_strategy: str = "auto"
python/sglang/srt/server_args.py
+2
View File
@@ -1340,6 +1340,7 @@ class ServerArgs:
help="Choose the runner backend for Blockwise FP8 GEMM operations. Options: 'auto' (default, auto-selects based on hardware), 'deep_gemm' (JIT-compiled; enabled by default on NVIDIA Hopper (SM90) and Blackwell (SM100) when DeepGEMM is installed), 'flashinfer_trtllm' (optimal for Blackwell and low-latency), 'flashinfer_cutlass' (FlashInfer CUTLASS groupwise FP8 GEMM), 'flashinfer_deepgemm' (Hopper SM90 only; uses swapAB optimization for small M dimensions in decoding), 'cutlass' (optimal for Hopper/Blackwell GPUs and high-throughput), 'triton' (fallback, widely compatible), 'aiter' (ROCm only). ",
cli_name="--fp8-gemm-backend",
choices=FP8_GEMM_RUNNER_BACKEND_CHOICES,
resolvable=True,
),
] = "auto"
fp4_gemm_runner_backend: A[
@@ -3879,6 +3880,7 @@ class ServerArgs:
"MistralLarge3ForCausalLM",
"PixtralForConditionalGeneration",
"GlmMoeDsaForCausalLM",
"LongcatFlashForCausalLM",
]:
# Set attention backend for DeepSeek
if is_deepseek_dsa(hf_config): # DeepSeek 3.2/GLM 5
python/sglang/srt/utils/hf_transformers/config.py
+29 -6
View File
@@ -16,6 +16,7 @@
from pathlib import Path
from typing import Optional
from transformers import PretrainedConfig
from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
from sglang.srt.configs.model_config_parser_registry import (
@@ -51,6 +52,26 @@ def _apply_deepseek_ocr_overrides(config, model):
config._name_or_path = model
_LONGCAT_ARCHS = {
"LongcatCausalLM",
"LongcatFlashForCausalLM",
"LongcatFlashNgramForCausalLM",
}
def _try_load_longcat_config(model, revision: Optional[str], **kwargs):
config_dict, _ = PretrainedConfig.get_config_dict(
model, revision=revision, **kwargs
)
architectures = config_dict.get("architectures") or []
if not any(arch in _LONGCAT_ARCHS for arch in architectures):
return None
return _CONFIG_REGISTRY["longcat_flash"].from_pretrained(
model, revision=revision, **kwargs
)
@register_model_config_parser("hf")
class HfModelConfigParser(ModelConfigParserBase):
def parse(
@@ -60,12 +81,14 @@ class HfModelConfigParser(ModelConfigParserBase):
revision: Optional[str] = None,
**kwargs,
):
config = AutoConfig.from_pretrained(
model,
trust_remote_code=trust_remote_code,
revision=revision,
**kwargs,
)
config = _try_load_longcat_config(model, revision, **kwargs)
if config is None:
config = AutoConfig.from_pretrained(
model,
trust_remote_code=trust_remote_code,
revision=revision,
**kwargs,
)
if (
config.architectures is not None
test/registered/jit/benchmark/bench_ngram_compute_decode.py
+3
View File
@@ -20,6 +20,7 @@ register_cuda_ci(
NE_N = 8
NE_K = 2
VOCAB_SIZE = 32000
EOS_TOKEN_ID = VOCAB_SIZE
MAX_CONTEXT_LEN = 1024
BATCH_SIZE_LIST = get_benchmark_range(
full_range=[1, 2, 8, 32, 128, 512, 1024, 2048, 4096],
@@ -99,6 +100,7 @@ def benchmark(batch_size: int, provider: str):
row_indices,
column_starts,
n_gram_ids,
EOS_TOKEN_ID,
)
else:
@@ -114,6 +116,7 @@ def benchmark(batch_size: int, provider: str):
row_indices,
column_starts,
n_gram_ids,
EOS_TOKEN_ID,
)
return run_benchmark_no_cudagraph(fn)
test/registered/jit/test_ngram_embedding.py
+6 -1
View File
@@ -41,6 +41,7 @@ def test_compute_n_gram_ids_decode_matches_general(batch_size: int) -> None:
ne_n = 8
ne_k = 2
vocab_size = 32000
eos_token_id = vocab_size
max_context_len = 1024
max_running_reqs = batch_size + 8
num_configs = (ne_n - 1) * ne_k
@@ -59,7 +60,9 @@ def test_compute_n_gram_ids_decode_matches_general(batch_size: int) -> None:
column_starts = torch.randint(
0, max_context_len, (batch_size,), dtype=torch.int32, device="cuda"
)
tokens = torch.empty(batch_size, dtype=torch.int32, device="cuda")
tokens = torch.randint(
0, vocab_size, (batch_size,), dtype=torch.int32, device="cuda"
)
exclusive_req_len_sums = torch.arange(
batch_size + 1, dtype=torch.int32, device="cuda"
)
@@ -80,6 +83,7 @@ def test_compute_n_gram_ids_decode_matches_general(batch_size: int) -> None:
row_indices=row_indices,
column_starts=column_starts,
n_gram_ids=n_gram_ids_general,
eos_token_id=eos_token_id,
)
compute_n_gram_ids_decode(
ne_n=ne_n,
@@ -91,6 +95,7 @@ def test_compute_n_gram_ids_decode_matches_general(batch_size: int) -> None:
row_indices=row_indices,
column_starts=column_starts,
n_gram_ids=n_gram_ids_decode,
eos_token_id=eos_token_id,
)
torch.testing.assert_close(n_gram_ids_decode, n_gram_ids_general, atol=0, rtol=0)
test/registered/unit/model_executor/test_ngram_token_table.py
+128
View File
@@ -0,0 +1,128 @@
"""Unit tests for LongCat ngram token table updates."""
import unittest
from types import SimpleNamespace
from unittest.mock import patch
import torch
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
maybe_stub_sgl_kernel()
from sglang.srt.model_executor.ngram_token_table import ( # noqa: E402
update_ngram_token_table_after_sampling,
)
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
def _make_ngram_info(batch_size: int, skip_token_table_update=None):
return SimpleNamespace(
token_table=torch.full((8, 16), -1, dtype=torch.int32),
out_column_starts=torch.empty(batch_size, dtype=torch.int32),
out_req_lens=torch.empty(batch_size, dtype=torch.int32),
skip_token_table_update=skip_token_table_update,
)
class TestNgramTokenTableUpdate(CustomTestCase):
def test_chunked_prefill_mask_skips_pseudo_next_token(self):
info = _make_ngram_info(
4, skip_token_table_update=torch.tensor([False, True, False, True])
)
next_token_ids = torch.tensor([101, 202, 303, 404], dtype=torch.int64)
req_pool_indices = torch.tensor([3, 4, 5, 6], dtype=torch.int64)
seq_lens = torch.tensor([11, 22, 33, 44], dtype=torch.int64)
with patch(
"sglang.srt.model_executor.ngram_token_table.update_token_table"
) as update_mock:
updated = update_ngram_token_table_after_sampling(
ngram_embedding_info=info,
next_token_ids=next_token_ids,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
batch_size=4,
)
self.assertTrue(updated)
update_mock.assert_called_once()
kwargs = update_mock.call_args.kwargs
self.assertIs(kwargs["ne_token_table"], info.token_table)
self.assertTrue(
torch.equal(kwargs["tokens"], torch.tensor([101, 303], dtype=torch.int32))
)
self.assertTrue(
torch.equal(kwargs["row_indices"], torch.tensor([3, 5], dtype=torch.int64))
)
self.assertTrue(
torch.equal(
kwargs["column_starts"], torch.tensor([11, 33], dtype=torch.int32)
)
)
self.assertTrue(
torch.equal(kwargs["req_lens"], torch.ones(2, dtype=torch.int32))
)
self.assertIsNone(kwargs["ignore_tokens"])
def test_all_requests_masked_does_not_update_table(self):
info = _make_ngram_info(2, skip_token_table_update=torch.tensor([True, True]))
with patch(
"sglang.srt.model_executor.ngram_token_table.update_token_table"
) as update_mock:
updated = update_ngram_token_table_after_sampling(
ngram_embedding_info=info,
next_token_ids=torch.tensor([101, 202], dtype=torch.int64),
req_pool_indices=torch.tensor([3, 4], dtype=torch.int64),
seq_lens=torch.tensor([11, 22], dtype=torch.int64),
batch_size=2,
)
self.assertFalse(updated)
update_mock.assert_not_called()
def test_unmasked_update_writes_all_sampled_tokens(self):
info = _make_ngram_info(3)
next_token_ids = torch.tensor([101, 202, 303], dtype=torch.int64)
req_pool_indices = torch.tensor([3, 4, 5], dtype=torch.int64)
seq_lens = torch.tensor([11, 22, 33], dtype=torch.int64)
with patch(
"sglang.srt.model_executor.ngram_token_table.update_token_table"
) as update_mock:
updated = update_ngram_token_table_after_sampling(
ngram_embedding_info=info,
next_token_ids=next_token_ids,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
batch_size=3,
)
self.assertTrue(updated)
update_mock.assert_called_once()
kwargs = update_mock.call_args.kwargs
self.assertIs(kwargs["ne_token_table"], info.token_table)
self.assertTrue(
torch.equal(
kwargs["tokens"], torch.tensor([101, 202, 303], dtype=torch.int32)
)
)
self.assertIs(kwargs["row_indices"], req_pool_indices)
self.assertIs(kwargs["column_starts"], info.out_column_starts)
self.assertIs(kwargs["req_lens"], info.out_req_lens)
self.assertTrue(
torch.equal(
info.out_column_starts, torch.tensor([11, 22, 33], dtype=torch.int32)
)
)
self.assertTrue(
torch.equal(info.out_req_lens, torch.ones(3, dtype=torch.int32))
)
self.assertIsNone(kwargs["ignore_tokens"])
if __name__ == "__main__":
unittest.main()
test/registered/unit/test_model_overrides.py
+1
View File
@@ -89,6 +89,7 @@ class TestModelOverridableWhitelist(CustomTestCase):
"prefill_attention_backend",
"decode_attention_backend",
"flashinfer_allreduce_fusion_backend",
"fp8_gemm_runner_backend",
}
),
)