From 5cbea10e2fcd269809c26237c175a470719fc9fe Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=8E=8B=E9=B9=A4=E7=94=B7?= Date: Thu, 16 Jul 2026 15:42:51 +0800 Subject: [PATCH] Fix LongCat n-gram embedding in PD-disaggregated scheduler loops (#31134) Co-authored-by: whn09 Co-authored-by: Claude Opus 4.8 --- python/sglang/srt/disaggregation/decode.py | 6 ++++++ python/sglang/srt/disaggregation/prefill.py | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 664af1f59..ec4f47707 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -2001,6 +2001,9 @@ class SchedulerDisaggregationDecodeMixin: ) self.running_batch = plan.running_batch batch = plan.batch_to_run + batch = self.ngram_embedding_manager.prepare_for_forward( + batch, chunked_req=self.chunked_req + ) self.cur_batch_for_debug = batch # Launch the current batch @@ -2039,6 +2042,9 @@ class SchedulerDisaggregationDecodeMixin: ) self.running_batch = plan.running_batch batch = plan.batch_to_run + batch = self.ngram_embedding_manager.prepare_for_forward( + batch, chunked_req=self.chunked_req + ) self.cur_batch_for_debug = batch # overlap + spec + grammar is unsupported (would desync DP ranks). disable_overlap_for_batch = self.is_disable_overlap_for_batch( diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index a4c14d784..43efa2db9 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -521,6 +521,9 @@ class SchedulerDisaggregationPrefillMixin: ) self.running_batch = plan.running_batch batch = plan.batch_to_run + batch = self.ngram_embedding_manager.prepare_for_forward( + batch, chunked_req=self.chunked_req + ) self.cur_batch_for_debug = batch # Launch the current batch @@ -559,6 +562,9 @@ class SchedulerDisaggregationPrefillMixin: ) self.running_batch = plan.running_batch batch = plan.batch_to_run + batch = self.ngram_embedding_manager.prepare_for_forward( + batch, chunked_req=self.chunked_req + ) self.cur_batch_for_debug = batch # Launch the current batch