From 3ed46f599f08b83bedd7be6872ab7d0daf237179 Mon Sep 17 00:00:00 2001 From: Jonny Kong Date: Fri, 19 Jun 2026 15:12:07 -0700 Subject: [PATCH] [core] Don't force seq_lens_cpu publication under piecewise CUDA graph (#28633) Co-authored-by: jonnykong Co-authored-by: Liangsheng Yin Co-authored-by: hnyls2002 --- python/sglang/srt/managers/overlap_utils.py | 17 +++++++---------- .../unit/server_args/test_server_args.py | 13 ------------- 2 files changed, 7 insertions(+), 23 deletions(-) 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"