[DSV4][BCG] Optimize the heavy memory use of C4 Indexer when BCG is enabled (#36534)

Co-authored-by: Yuwei An <ayw.sirius19@gmail.com>

test/registered/attention/unittests/dense/test_tbo.py

@@ -191,6 +191,7 @@ class TestTboAttnDenseAttentionBackendCorrectness(CustomTestCase):
             encoder_lens=None,
             out_cache_loc=batch.out_cache_loc,
             spec_info=batch.spec_info,
+            max_seq_len_override=700,
         )
 
         # Pure mocks (no `wraps=...`) so the dispatcher's slicing/contract is
@@ -215,6 +216,8 @@ class TestTboAttnDenseAttentionBackendCorrectness(CustomTestCase):
         self.assertEqual(
             child_fbs[1].req_pool_indices.shape[0], capture_bs - split_seq_index
         )
+        self.assertEqual(child_fbs[0].max_seq_len_override, 700)
+        self.assertEqual(child_fbs[1].max_seq_len_override, 700)
 
 
 if __name__ == "__main__":

test/registered/attention/unittests/dsv4/test_deepseek_v4.py

@@ -579,8 +579,8 @@ class TestDSV4BreakableCudaGraphMetadataContract(CustomTestCase):
             return replay_metadata
 
         backend._build_forward_metadata = fake_build_forward_metadata
-        forward_batch = SimpleNamespace(name="live")
-        static_forward_batch = SimpleNamespace(name="static")
+        forward_batch = SimpleNamespace(name="live", max_seq_len_override=None)
+        static_forward_batch = SimpleNamespace(name="static", max_seq_len_override=None)
 
         backend.prepare_forward_metadata_for_breakable_cuda_graph_replay(
             capture_metadata,

test/registered/cp/test_cp_strategy_unit.py

@@ -135,6 +135,7 @@ class TestPrefillCPBCGReplay(CustomTestCase):
         runner.has_mha_companion_layers = False
         runner.capture_hidden_mode = CaptureHiddenMode.NULL
         runner.capture_num_tokens = [2048, 2304]
+        runner.max_context_size = None
         runner.max_num_tokens = 2304
         runner.enable_cp_bcg_capture = True
         return runner

test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py

@@ -43,6 +43,7 @@ class TestMultimodalPiecewiseCudaGraph(CustomTestCase):
         runner.has_mha_companion_layers = backend == Backend.BREAKABLE
         runner.capture_hidden_mode = CaptureHiddenMode.NULL
         runner.capture_num_tokens = [4, 16]
+        runner.max_context_size = None
         runner.max_num_tokens = 16
         return runner
 

test/registered/unit/managers/scheduler_components/test_dp_attn.py

@@ -84,6 +84,7 @@ class TestDecodeToExtendConversionVote(CustomTestCase):
     def _vote(self, *, beam):
         runner = Mock(spec=dp_attn.PrefillCudaGraphRunner)
         runner.enable_lora = False
+        runner.max_context_size = None
         runner.can_replay_locally.return_value = True
         batch = SimpleNamespace(
             forward_mode=ForwardMode.DECODE,

test/registered/unit/model_executor/runner/test_hidden_state_graph_recapture.py

@@ -73,6 +73,7 @@ class TestHiddenStateGraphRecapture(CustomTestCase):
         runner._capture_chunked_prefix = False
         runner.capture_hidden_mode = capture_hidden_mode
         runner.capture_num_tokens = [4]
+        runner.max_context_size = None
         runner.max_num_tokens = 4
         return runner
 

test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py

@@ -2,6 +2,9 @@ import unittest
 from types import SimpleNamespace
 from unittest import mock
 
+import torch
+
+import sglang.srt.model_executor.runner.prefill_cuda_graph_runner as runner_module
 from sglang.srt.layers.moe.utils import MoeA2ABackend
 from sglang.srt.model_executor import forward_batch_info
 from sglang.srt.model_executor.cuda_graph_config import Backend
@@ -13,6 +16,7 @@ from sglang.srt.model_executor.forward_batch_info import (
 from sglang.srt.model_executor.runner.prefill_cuda_graph_runner import (
     PrefillCudaGraphRunner,
 )
+from sglang.srt.model_executor.runner.shape_key import ShapeKey
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
@@ -29,6 +33,7 @@ class TestPrefillCudaGraphPadding(CustomTestCase):
         runner.has_mha_companion_layers = False
         runner.capture_hidden_mode = CaptureHiddenMode.NULL
         runner.capture_num_tokens = [4, 16]
+        runner.max_context_size = None
         runner.max_num_tokens = 16
         return runner
 
@@ -44,6 +49,8 @@ class TestPrefillCudaGraphPadding(CustomTestCase):
             return_logprob=False,
             input_ids=list(range(num_tokens)),
             extend_prefix_lens_cpu=[0],
+            seq_lens_cpu=torch.tensor([num_tokens], dtype=torch.int64),
+            seq_lens=torch.tensor([num_tokens], dtype=torch.int64),
         )
 
     def test_rejects_more_than_two_x_token_padding(self):
@@ -67,7 +74,7 @@ class TestPrefillCudaGraphPadding(CustomTestCase):
         runner._prepare_forward_metadata_for_replay(
             forward_batch,
             static_forward_batch,
-            num_tokens=16,
+            shape_key=ShapeKey(size=16),
         )
 
         attn_backend.init_forward_metadata.assert_called_once_with(forward_batch)
@@ -137,6 +144,28 @@ class TestPrefillCudaGraphPadding(CustomTestCase):
                 )
             )
 
+    def test_rejects_context_above_fixed_maximum(self):
+        runner = self._make_runner()
+        runner.max_context_size = 700
+
+        much_shorter = self._make_forward_batch(4)
+        much_shorter.seq_lens_cpu.fill_(200)
+        self.assertTrue(runner.can_run_graph(much_shorter))
+
+        uncovered = self._make_forward_batch(4)
+        uncovered.seq_lens_cpu.fill_(701)
+        self.assertFalse(runner.can_run_graph(uncovered))
+
+    def test_unsupported_path_ignores_max_context_size(self):
+        runner = self._make_runner()
+        runner.max_context_size = 1024
+
+        with self.assertLogs(runner_module.logger, level="WARNING") as logs:
+            runner._ignore_max_context_size("test path")
+
+        self.assertIsNone(runner.max_context_size)
+        self.assertIn("fixed metadata extent", "\n".join(logs.output))
+
 
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py

@@ -244,6 +244,8 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
     def test_static_batch_preserves_consumed_multimodal_embeddings(self):
         runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
         runner.capture_num_tokens = [4]
+        runner.max_context_size = None
+        runner._capture_chunked_prefix = False
         runner.buffer_registry = _FakeBatchRegistry()
         runner.model_runner = SimpleNamespace(attn_tp_sequence_sharded=lambda _: False)
         runner.enable_cp_bcg_capture = False
@@ -477,6 +479,7 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
         runner.capture_hidden_mode = CaptureHiddenMode.NULL
         runner.max_num_tokens = 32
         runner.capture_num_tokens = [4]
+        runner.max_context_size = None
         runner.backend = SimpleNamespace()
         runner.prefill_backend_name = Backend.FULL
         runner.has_mha_companion_layers = False

test/registered/unit/server_args/test_server_args.py

@@ -21,6 +21,7 @@ from sglang.srt.arg_groups.attention_hook import (
 from sglang.srt.arg_groups.cuda_graph_hook import (
     apply_cuda_graph_compatibility,
     disable_tc_piecewise_cudagraph_if_incompatible,
+    finalize_cuda_graph_prefill_max_context,
     handle_cuda_graph_config,
 )
 from sglang.srt.arg_groups.hicache_hook import (
@@ -2114,6 +2115,35 @@ class TestCudaGraphConfigDataclassAccess(CustomTestCase):
         self.assertEqual(config.compiler, "eager")
 
 
+class TestCudaGraphPrefillMaxContextResolution(CustomTestCase):
+    @staticmethod
+    def _make_args(max_context_size, model_context_len=4096, page_size=64):
+        args = ServerArgs(
+            model_path="dummy",
+            page_size=page_size,
+            cuda_graph_config=CudaGraphConfig(
+                prefill=PhaseConfig(
+                    backend=Backend.BREAKABLE,
+                    max_context_size=max_context_size,
+                )
+            ),
+        )
+        args._model_config = SimpleNamespace(context_len=model_context_len)
+        return args
+
+    def test_rejects_invalid_values_during_resolution(self):
+        cases = (
+            (0, "positive integer"),
+            (-1, "positive integer"),
+            (4097, "model context length"),
+        )
+        for max_context_size, expected_error in cases:
+            with self.subTest(max_context_size=max_context_size):
+                args = self._make_args(max_context_size)
+                with self.assertRaisesRegex(ValueError, expected_error):
+                    finalize_cuda_graph_prefill_max_context(args)
+
+
 class TestPipelineParallelPrefillCudaGraphPolicy(CustomTestCase):
     def test_pp_prefill_graph_is_opt_in(self):
         cases = (

test/registered/unit/test_server_args_cli_metadata.py

@@ -45,6 +45,11 @@ class TestServerArgsMigratedCliMetadata(CustomTestCase):
             self.actions_by_option["--prefill-delayer-forward-passes-buckets"].nargs,
             "+",
         )
+        self.assertIs(
+            self.actions_by_option["--cuda-graph-prefill-max-context"].type,
+            human_readable_int,
+        )
+        self.assertIsNone(self.actions_by_option["--context-bucket"].nargs)
         self.assertEqual(
             self.actions_by_option["--schedule-policy"].choices,
             [
@@ -76,6 +81,19 @@ class TestServerArgsMigratedCliMetadata(CustomTestCase):
                 self.assertEqual(args.dp_size, 3)
                 self.assertEqual(ServerArgs.from_cli_args(args).dp_size, 3)
 
+    def test_prefill_max_context_accepts_human_readable_values(self):
+        for option in (
+            "--cuda-graph-prefill-max-context",
+            "--context-bucket",
+        ):
+            with self.subTest(option=option):
+                args = self.parser.parse_args(["--model", "dummy", option, "200k"])
+
+                self.assertEqual(
+                    ServerArgs.from_cli_args(args).cuda_graph_prefill_max_context,
+                    200_000,
+                )
+
     def test_migrated_and_manual_options_parse_together(self):
         args = self.parser.parse_args(
             [