diff --git a/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py b/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py index 6f0a7b75f..4f2486cef 100644 --- a/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py +++ b/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py @@ -199,6 +199,8 @@ class MlxModelRunnerStub(ModelRunner): self.graph_mem_usage = 0 self.attn_backend = None + self.init_ngram_embedding_manager() + logger.info( f"MLX stub: initialized minimal pools " f"(max_total_num_tokens={self.max_total_num_tokens}, " diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 9f8fc168d..846a0cc6f 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -38,7 +38,6 @@ import torch.distributed from torch.cuda import Stream as CudaStream from torch.distributed import barrier -from sglang.jit_kernel.ngram_embedding import update_token_table from sglang.srt.configs.model_config import ModelConfig, ModelImpl, is_minimax_sparse from sglang.srt.constrained.grammar_manager import GrammarManager from sglang.srt.debug_utils.pr_fix_toggle import maybe_revert_pr_fix @@ -225,7 +224,7 @@ from sglang.srt.managers.utils import ( ) from sglang.srt.mem_cache import kv_cache_builder from sglang.srt.mem_cache.common import maybe_cache_unfinished_req, release_kv_cache -from sglang.srt.model_executor.forward_batch_info import ForwardMode, PPProxyTensors +from sglang.srt.model_executor.forward_batch_info import PPProxyTensors from sglang.srt.model_loader.utils import get_resolved_model_impl from sglang.srt.multiplex.multiplexing_mixin import SchedulerMultiplexMixin from sglang.srt.observability.metrics_collector import SchedulerMetricsCollector @@ -1303,71 +1302,16 @@ class Scheduler( self.batch_record_ct = 0 def maybe_init_ngram_embedding(self): + self.ngram_embedding_manager = ( + self.tp_worker.model_runner.ngram_embedding_manager + ) self.use_ngram_embedding = self.tp_worker.model_config.use_ngram_embedding if self.use_ngram_embedding: - self.token_table = self.tp_worker.model_runner.token_table + self.token_table = self.tp_worker.model_runner.ngram_embedding_manager.table hf_config = self.tp_worker.model_config.hf_config self.ngram_embedding_n = hf_config.ngram_embedding_n self.ngram_embedding_k = hf_config.ngram_embedding_k - def _maybe_prepare_ngram_embedding( - self, batch: Optional[ScheduleBatch] - ) -> Optional[ScheduleBatch]: - """Fill the token table for ngram embedding before a forward pass.""" - if batch is None or not self.use_ngram_embedding: - return batch - batch.ne_token_table = self.token_table - if batch.forward_mode == ForwardMode.EXTEND: - all_tokens = [] - column_starts = [] - request_lengths = [] - for req in batch.reqs: - start = len(req.prefix_indices) - end = start + req.extend_range.length - fill_ids = req.origin_input_ids + req.output_ids - if start == 0: - tokens = fill_ids[start:end] - column_starts.append(0) - elif start < self.ngram_embedding_n: - tokens = fill_ids[0:end] - column_starts.append(0) - else: - # Prepend n-1 tokens before prefix_len for n-gram context - tokens = fill_ids[start - self.ngram_embedding_n + 1 : end] - column_starts.append(start - self.ngram_embedding_n + 1) - all_tokens.extend(tokens) - request_lengths.append(len(tokens)) - dtype = self.token_table.dtype - device = self.token_table.device - update_token_table( - ne_token_table=self.token_table, - tokens=torch.tensor(all_tokens, dtype=dtype, device=device), - row_indices=batch.req_pool_indices, - column_starts=torch.tensor( - column_starts, dtype=torch.int32, device=device - ), - req_lens=torch.tensor( - request_lengths, dtype=torch.int32, device=device - ), - 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): """Initialize deterministic inference configuration for different attention backends.""" if not self.server_args.enable_deterministic_inference: @@ -2806,7 +2750,9 @@ class Scheduler( ) # Handle ngram embedding - ret = self._maybe_prepare_ngram_embedding(ret) + ret = self.ngram_embedding_manager.prepare_for_forward( + ret, chunked_req=self.chunked_req + ) if ret: set_schedule_time_batch(ret) diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 3104b7fcf..988956e2f 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -848,7 +848,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): ret.positions = positions ret.extend_logprob_start_lens_cpu = extend_logprob_start_lens - if model_runner.use_ngram_embedding: + if model_runner.ngram_embedding_manager.enabled: ret._init_ngram_embedding_info(batch, device) if model_runner.model_config.model_is_mrope: diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 126c9851c..16627622b 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -129,6 +129,9 @@ from sglang.srt.model_executor.forward_context import ( ) from sglang.srt.model_executor.graph_shared_output import GraphSharedOutput from sglang.srt.model_executor.hook_manager import register_forward_hooks +from sglang.srt.model_executor.model_runner_components.ngram_embedding_manager import ( + NgramEmbeddingManager, +) from sglang.srt.model_executor.model_runner_components.remote_instance_weight_transporter import ( RemoteInstanceWeightTransporter, ) @@ -141,9 +144,6 @@ from sglang.srt.model_executor.model_runner_components.weight_updater import ( 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, @@ -558,6 +558,16 @@ class ModelRunner(ModelRunnerKVCacheMixin): gpu_id=self.gpu_id, ) + def init_ngram_embedding_manager(self): + self.ngram_embedding_manager = NgramEmbeddingManager.from_model( + model=self.model, + model_config=self.model_config, + req_to_token_pool=self.req_to_token_pool, + server_args=self.server_args, + max_running_requests=self.max_running_requests, + device=self.device, + ) + def init_msprobe(self): # Init the msprobe try: @@ -773,7 +783,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): ) # Init ngram embedding token table - self.maybe_init_ngram_embedding() + self.init_ngram_embedding_manager() if self.enable_hisparse: from sglang.srt.managers.hisparse_coordinator import HiSparseCoordinator @@ -1627,45 +1637,6 @@ class ModelRunner(ModelRunnerKVCacheMixin): full_attention_backend = ATTENTION_BACKENDS[backend_str](self) return attn_backend_wrapper(self, full_attention_backend) - def maybe_init_ngram_embedding(self): - self.use_ngram_embedding = self.model_config.use_ngram_embedding - if self.use_ngram_embedding: - from sglang.srt.layers.n_gram_embedding import NgramEmbedding - - # Sized to mirror req_to_token (indexed by req_pool_idx). - self.token_table = torch.empty( - self.req_to_token_pool.req_to_token.shape[0], - self.model_config.context_len, - dtype=torch.int32, - device=self.device, - ) - chunked_prefill_size = self.server_args.chunked_prefill_size - assert ( - chunked_prefill_size is not None and chunked_prefill_size > 0 - ), "Ngram embedding requires chunked prefill to be enabled (chunked_prefill_size > 0)" - for module in self.model.modules(): - if isinstance(module, NgramEmbedding): - module.init_buffers( - self.max_running_requests, chunked_prefill_size, self.device - ) - - def maybe_update_ngram_token_table( - self, - next_token_ids: torch.Tensor, - forward_batch: ForwardBatch, - ): - """Update the ngram embedding token table after sampling.""" - ngram_embedding_info = forward_batch.ngram_embedding_info - if ngram_embedding_info is None: - return - 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): """Capture device graphs.""" self.decode_cuda_graph_runner = None @@ -2316,7 +2287,10 @@ class ModelRunner(ModelRunnerKVCacheMixin): else forward_batch.seq_lens - 1 ), ) - self.maybe_update_ngram_token_table(next_token_ids, forward_batch) + self.ngram_embedding_manager.update_after_decode( + next_token_ids=next_token_ids, + forward_batch=forward_batch, + ) return next_token_ids def compute_logprobs_only( diff --git a/python/sglang/srt/model_executor/model_runner_components/ngram_embedding_manager.py b/python/sglang/srt/model_executor/model_runner_components/ngram_embedding_manager.py new file mode 100644 index 000000000..71e8a3322 --- /dev/null +++ b/python/sglang/srt/model_executor/model_runner_components/ngram_embedding_manager.py @@ -0,0 +1,190 @@ +"""Utilities for updating LongCat ngram embedding token tables.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.jit_kernel.ngram_embedding import update_token_table +from sglang.srt.configs.model_config import ModelConfig +from sglang.srt.managers.schedule_batch import ForwardMode +from sglang.srt.mem_cache.memory_pool import ReqToTokenPool +from sglang.srt.server_args import ServerArgs + +if TYPE_CHECKING: + from sglang.srt.managers.schedule_batch import Req, ScheduleBatch + from sglang.srt.model_executor.forward_batch_info import ForwardBatch + + +@dataclass(frozen=True, slots=True, kw_only=True) +class NgramEmbeddingManager: + enabled: bool + table: Optional[torch.Tensor] + n: int + k: int + + @classmethod + def from_model( + cls, + *, + model: torch.nn.Module, + model_config: ModelConfig, + req_to_token_pool: ReqToTokenPool, + server_args: ServerArgs, + max_running_requests: int, + device: str, + ): + token_table = None + ngram_embedding_n = 0 + ngram_embedding_k = 0 + use_ngram_embedding = model_config.use_ngram_embedding + if use_ngram_embedding: + from sglang.srt.layers.n_gram_embedding import NgramEmbedding + + # Sized to mirror req_to_token (indexed by req_pool_idx). + token_table = torch.empty( + req_to_token_pool.req_to_token.shape[0], + model_config.context_len, + dtype=torch.int32, + device=device, + ) + chunked_prefill_size = server_args.chunked_prefill_size + assert ( + chunked_prefill_size is not None and chunked_prefill_size > 0 + ), "Ngram embedding requires chunked prefill to be enabled (chunked_prefill_size > 0)" + for module in model.modules(): + if isinstance(module, NgramEmbedding): + module.init_buffers( + max_running_requests, chunked_prefill_size, device + ) + hf_config = model_config.hf_config + ngram_embedding_n = hf_config.ngram_embedding_n + ngram_embedding_k = hf_config.ngram_embedding_k + return cls( + enabled=use_ngram_embedding, + table=token_table, + n=ngram_embedding_n, + k=ngram_embedding_k, + ) + + def update_after_decode( + self, + next_token_ids: torch.Tensor, + forward_batch: ForwardBatch, + ): + """Update the ngram embedding token table after sampling.""" + ngram_embedding_info = forward_batch.ngram_embedding_info + if ngram_embedding_info is None: + return + 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 prepare_for_forward( + self, + batch: Optional[ScheduleBatch], + *, + chunked_req: Optional[Req], + ) -> Optional[ScheduleBatch]: + """Fill the token table for ngram embedding before a forward pass.""" + if batch is None or not self.enabled: + return batch + batch.ne_token_table = self.table + if batch.forward_mode == ForwardMode.EXTEND: + all_tokens = [] + column_starts = [] + request_lengths = [] + for req in batch.reqs: + start = len(req.prefix_indices) + end = start + req.extend_range.length + fill_ids = req.origin_input_ids + req.output_ids + if start == 0: + tokens = fill_ids[start:end] + column_starts.append(0) + elif start < self.n: + tokens = fill_ids[0:end] + column_starts.append(0) + else: + # Prepend n-1 tokens before prefix_len for n-gram context + tokens = fill_ids[start - self.n + 1 : end] + column_starts.append(start - self.n + 1) + all_tokens.extend(tokens) + request_lengths.append(len(tokens)) + dtype = self.table.dtype + device = self.table.device + update_token_table( + ne_token_table=self.table, + tokens=torch.tensor(all_tokens, dtype=dtype, device=device), + row_indices=batch.req_pool_indices, + column_starts=torch.tensor( + column_starts, dtype=torch.int32, device=device + ), + req_lens=torch.tensor( + request_lengths, dtype=torch.int32, device=device + ), + 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 chunked_req is not None: + skip_token_table_update = [req is 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 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 diff --git a/python/sglang/srt/model_executor/ngram_token_table.py b/python/sglang/srt/model_executor/ngram_token_table.py deleted file mode 100644 index adbb0a99b..000000000 --- a/python/sglang/srt/model_executor/ngram_token_table.py +++ /dev/null @@ -1,51 +0,0 @@ -"""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 diff --git a/python/sglang/srt/model_executor/runner/base_runner.py b/python/sglang/srt/model_executor/runner/base_runner.py index 85388a30f..9958f7ec2 100644 --- a/python/sglang/srt/model_executor/runner/base_runner.py +++ b/python/sglang/srt/model_executor/runner/base_runner.py @@ -312,7 +312,11 @@ class BaseRunner(ABC): num_tokens_per_req=num_tokens_per_req, cache_loc_dtype=torch.int64, enable_mamba_track=False, - ne_token_table=mr.token_table if mr.use_ngram_embedding else None, + ne_token_table=( + mr.ngram_embedding_manager.table + if mr.ngram_embedding_manager.enabled + else None + ), hc_hidden_size=getattr(mr.model_config, "hc_hidden_size", None), pp_proxy_topk_size=mr.get_pp_proxy_topk_size(), ) diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index c4b7e5424..983b8648c 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -213,7 +213,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): self.enable_two_batch_overlap = ( model_runner.server_args.enable_two_batch_overlap ) - self.use_ngram_embedding = model_runner.use_ngram_embedding + self.use_ngram_embedding = model_runner.ngram_embedding_manager.enabled if self.use_ngram_embedding: hf_config = model_runner.model_config.hf_config self.ngram_embedding_n = hf_config.ngram_embedding_n @@ -362,7 +362,9 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): cache_loc_dtype=self._cache_loc_dtype(), enable_mamba_track=enable_mamba_track, ne_token_table=( - model_runner.token_table if self.use_ngram_embedding else None + model_runner.ngram_embedding_manager.table + if self.use_ngram_embedding + else None ), hc_hidden_size=getattr( self.model_runner.model_config, "hc_hidden_size", None diff --git a/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py b/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py index 687e29a8c..f8fdd0bae 100644 --- a/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py +++ b/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py @@ -37,6 +37,7 @@ def make_runner( return SimpleNamespace( server_args=args, + model_config=SimpleNamespace(), hybrid_gdn_config=SimpleNamespace( linear_key_head_dim=key_dim, linear_value_head_dim=value_dim, @@ -60,6 +61,11 @@ class TestFlashInferGDNPrefillBackendPolicy(unittest.TestCase): flashinfer_available=True, ): with ( + patch.object( + gdn_backend, + "hybrid_gdn_config", + return_value=runner.hybrid_gdn_config, + ), patch.object(gdn_backend, "is_cuda", return_value=cuda), patch.object(torch.cuda, "get_device_capability", return_value=capability), patch.object(torch.version, "cuda", cuda_version), diff --git a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py b/test/registered/unit/managers/test_scheduler_chunked_req_gate.py index 9f4860fd4..14ed036ec 100644 --- a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py +++ b/test/registered/unit/managers/test_scheduler_chunked_req_gate.py @@ -97,7 +97,10 @@ def _scheduler_for_get_next_batch(*, tree_cache, chunked_req) -> Scheduler: s.dp_attn_adapter.maybe_prepare_mlp_sync_batch = MagicMock( side_effect=lambda batch, **_: batch ) - s._maybe_prepare_ngram_embedding = MagicMock(side_effect=lambda batch: batch) + s.ngram_embedding_manager = MagicMock() + s.ngram_embedding_manager.prepare_for_forward = MagicMock( + side_effect=lambda batch, **_: batch + ) s.update_running_batch = MagicMock(side_effect=lambda batch: batch) s.tree_cache = tree_cache s.chunked_req = chunked_req diff --git a/test/registered/unit/model_executor/test_ngram_token_table.py b/test/registered/unit/model_executor/model_runner_components/test_ngram_embedding_manager.py similarity index 91% rename from test/registered/unit/model_executor/test_ngram_token_table.py rename to test/registered/unit/model_executor/model_runner_components/test_ngram_embedding_manager.py index 56ec6ad81..3010d86db 100644 --- a/test/registered/unit/model_executor/test_ngram_token_table.py +++ b/test/registered/unit/model_executor/model_runner_components/test_ngram_embedding_manager.py @@ -11,7 +11,7 @@ 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 +from sglang.srt.model_executor.model_runner_components.ngram_embedding_manager import ( # noqa: E402 update_ngram_token_table_after_sampling, ) @@ -37,7 +37,7 @@ class TestNgramTokenTableUpdate(CustomTestCase): seq_lens = torch.tensor([11, 22, 33, 44], dtype=torch.int64) with patch( - "sglang.srt.model_executor.ngram_token_table.update_token_table" + "sglang.srt.model_executor.model_runner_components.ngram_embedding_manager.update_token_table" ) as update_mock: updated = update_ngram_token_table_after_sampling( ngram_embedding_info=info, @@ -71,7 +71,7 @@ class TestNgramTokenTableUpdate(CustomTestCase): 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" + "sglang.srt.model_executor.model_runner_components.ngram_embedding_manager.update_token_table" ) as update_mock: updated = update_ngram_token_table_after_sampling( ngram_embedding_info=info, @@ -91,7 +91,7 @@ class TestNgramTokenTableUpdate(CustomTestCase): seq_lens = torch.tensor([11, 22, 33], dtype=torch.int64) with patch( - "sglang.srt.model_executor.ngram_token_table.update_token_table" + "sglang.srt.model_executor.model_runner_components.ngram_embedding_manager.update_token_table" ) as update_mock: updated = update_ngram_token_table_after_sampling( ngram_embedding_info=info,