[core] Don't force seq_lens_cpu publication under piecewise CUDA graph (#28633)

Co-authored-by: jonnykong <jonnykong@fb.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
Co-authored-by: hnyls2002 <lsyincs@gmail.com>

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
+    OR over per-backend needs_cpu_seq_lens; force True under TBO (it reads the
-    (they read the CPU mirror outside the backend layer).
+    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 SpeculativeAlgorithm.from_string(server_args.speculative_algorithm).is_ngram():
-    if (
+        # ngram's USE_FULL_MASK verify path reads seq_lens_cpu per req to size
-        cuda_graph_config is not None
+        # the tree mask, regardless of the attn backend (e.g. Triton opts out).
-        and cuda_graph_config.prefill.backend == Backend.TC_PIECEWISE
-    ):
-        # FIXME: support PCG without seq lens cpu value
         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.

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"