diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index 05b753f55..3af719f9d 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/python/sglang/srt/managers/overlap_utils.py @@ -5,7 +5,7 @@ from typing import TYPE_CHECKING, Sequence, Union import torch from sglang.srt.environ import envs -from sglang.srt.model_executor.cuda_graph_config import Backend +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_utils import spec_need_hidden_states from sglang.srt.speculative.triton_ops.gather_spec_extras import gather_spec_extras from sglang.srt.utils import is_cuda, is_hip, is_npu @@ -16,7 +16,6 @@ if TYPE_CHECKING: from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.eagle_info import EagleDraftInput - from sglang.srt.speculative.spec_info import SpeculativeAlgorithm def decide_needs_cpu_seq_lens( @@ -25,18 +24,16 @@ def decide_needs_cpu_seq_lens( ) -> bool: """Whether FutureMap must publish seq_lens_cpu / sum. - OR over per-backend needs_cpu_seq_lens; force True under TBO / piecewise CG - (they read the CPU mirror outside the backend layer). + OR over per-backend needs_cpu_seq_lens; force True under TBO (it reads the + CPU mirror outside the backend layer to split the batch) or ngram (its + USE_FULL_MASK verify path reads the host mirror regardless of backend). """ if server_args.enable_two_batch_overlap: # FIXME: support TBO without seq lens cpu value return True - cuda_graph_config = server_args.cuda_graph_config - if ( - cuda_graph_config is not None - and cuda_graph_config.prefill.backend == Backend.TC_PIECEWISE - ): - # FIXME: support PCG without seq lens cpu value + if SpeculativeAlgorithm.from_string(server_args.speculative_algorithm).is_ngram(): + # ngram's USE_FULL_MASK verify path reads seq_lens_cpu per req to size + # the tree mask, regardless of the attn backend (e.g. Triton opts out). return True # Skip unset slots (e.g. draft_extend_attn_backend on some spec configs); # missing flag -> True so undeclared backends stay on the legacy path. diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 5672946e1..be2215977 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -1009,19 +1009,6 @@ class TestPrefillOnlyDisableKvCache(unittest.TestCase): class TestCudaGraphConfigDataclassAccess(CustomTestCase): - def test_overlap_force_cpu_seq_lens_with_tc_piecewise_prefill(self): - from sglang.srt.managers.overlap_utils import decide_needs_cpu_seq_lens - - server_args = SimpleNamespace( - enable_two_batch_overlap=False, - cuda_graph_config=CudaGraphConfig( - prefill=PhaseConfig(backend=Backend.TC_PIECEWISE) - ), - ) - attn_backend = SimpleNamespace(needs_cpu_seq_lens=False) - - self.assertTrue(decide_needs_cpu_seq_lens(server_args, [attn_backend])) - @patch( "sglang.srt.model_executor.runner_backend." "tc_piecewise_cuda_graph_backend.get_moe_a2a_backend"