[Fix] DSv4.1: Fix BCG image handling and add multimodal bounded replay coverage

python/pyproject.toml

@@ -188,7 +188,7 @@ test = [
   # Pinned: a bump moves scoring for every eval test at once, so re-baseline
   # MODEL_SCORE_THRESHOLDS in test/registered/eval/test_text_models_gsm8k_eval.py
   # and the mmlu thresholds of run_eval's other callers before changing it.
-  "sgl-eval==0.1.2",
+  "sgl-eval==0.1.3",
   "sglang[fastokens]",
   "tabulate",
 ]

python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py

@@ -332,6 +332,9 @@ def __init__(self, model_runner: ModelRunner):
         # --- model flags ----------------------------------------------
         self.quant_config = getattr(model_runner.model, "quant_config", None)
         self.is_multimodal = model_runner.model_config.is_multimodal
+        self.is_dsv41_model = (
+            model_runner.model_config.hf_config.model_type == "deepseek_v41"
+        )
         self.enable_lora = model_runner.lora_manager is not None
         # Classification/reward forwards branch on return_pooled_hidden_states;
         # capture must use the same flag value as replay for those models.
@@ -2115,6 +2118,12 @@ def _execute_body_capture(
         full_path = self._is_full_backend
 
         def replay_layer_forward(*args, **layer_kwargs):
+            prepared_ids = layer_kwargs.get("input_ids", args[0] if args else None)
+            # TODO: Add general multimodal replay tests before enabling for other models.
+            if self.is_dsv41_model and prepared_ids is not None:
+                static_forward_batch.input_ids[: prepared_ids.shape[0]].copy_(
+                    prepared_ids
+                )
             # The captured body graph reads activations from the static
             # input_embeds slot. The outer model.forward (run eagerly)
             # passes the live embeddings into layer_model.forward as the

python/sglang/srt/models/deepseek_v4.py

@@ -3659,12 +3659,9 @@ def _forward_layers_hc_pre_from_prev(
                     forward_batch,
                     cp_all_tokens=cp_extend,
                 )
-                # Only multimodal placeholders become image_token_id, so a text-only
-                # batch skips this full-residual where.
                 if (
                     self.config.model_type == "deepseek_v41"
                     and self.config.vision_n_layers > 0
-                    and forward_batch.contains_mm_inputs()
                 ):
                     hidden_states = torch.where(
                         (input_ids == self.config.image_token_id)[:, None, None],

test/registered/e2e/models/test_deepseek_v41_flash_gb300.py

@@ -1,7 +1,8 @@
 """DeepSeek-V4.1-Flash accuracy with DSPARK on four GB300 GPUs.
 
-Run GSM8K and MMLU in TP4+EP4 and DP-attention4 + EP4 configurations,
-plus TRT-LLM attention with decoder bounded replay and breakable prefill.
+Run GSM8K and MMMU-Pro in TP4+EP4 and DP-attention4 + EP4 configurations,
+plus decoder SWA bounded replay inside breakable prefill graphs with the
+default and TRT-LLM attention backends.
 The DP+EP recipe is intentionally exercised rather than skipped: startup or
 accuracy failures must be visible in CI while this path is being validated.
 """
@@ -20,7 +21,7 @@
     try_cached_model,
 )
 
-register_cuda_ci(est_time=1650, stage="base-c", runner_config="4-gpu-gb300")
+register_cuda_ci(est_time=2400, stage="base-c", runner_config="4-gpu-gb300")
 
 MODEL = "deepseek-ai/DeepSeek-V4.1-Flash"
 
@@ -32,11 +33,16 @@ class DSV41FlashAccuracyMixin(GSM8KMixin):
     gsm8k_num_examples = 200
     gsm8k_num_threads = 64
     gsm8k_max_tokens = 4096
-    mmlu_score_threshold = 0.90
-    mmlu_num_examples = 500
-    mmlu_num_threads = 64
-    mmlu_max_tokens = 4096
-
+    mmmu_pro_score_threshold = 0.75
+    mmmu_pro_num_examples = 200
+    mmmu_pro_num_threads = 64
+    # Bound long reasoning responses consistently across all three CI recipes.
+    mmmu_pro_max_tokens = 32768
+    mmmu_pro_load_preset_from_model_id = MODEL
+
+    mem_fraction_static = "0.8"
+    max_running_requests = "64"
+    graph_args = []
     parallel_args = []
     server_env = {}
 
@@ -57,16 +63,17 @@ def setUpClass(cls):
                 "--speculative-dspark-block-size",
                 "5",
                 "--mem-fraction-static",
-                "0.8",
+                cls.mem_fraction_static,
                 "--cuda-graph-max-bs-decode",
                 "64",
                 "--max-running-requests",
-                "64",
+                cls.max_running_requests,
                 "--reasoning-parser",
                 "auto",
                 "--tool-call-parser",
                 "auto",
                 *cls.parallel_args,
+                *cls.graph_args,
             ],
             env=cls.server_env,
         )
@@ -88,14 +95,16 @@ def test_gsm8k(self):
         super().test_gsm8k()
         self.assert_dspark_accepts_draft_tokens()
 
-    def test_mmlu(self):
+    def test_mmmu_pro(self):
         _run_accuracy_eval(
             self,
-            eval_name="mmlu",
-            score_threshold=self.mmlu_score_threshold,
-            num_examples=self.mmlu_num_examples,
-            num_threads=self.mmlu_num_threads,
-            max_tokens=self.mmlu_max_tokens,
+            eval_name="mmmu_pro",
+            score_threshold=self.mmmu_pro_score_threshold,
+            num_examples=self.mmmu_pro_num_examples,
+            num_threads=self.mmmu_pro_num_threads,
+            model=None,
+            load_preset_from_model_id=self.mmmu_pro_load_preset_from_model_id,
+            max_tokens=self.mmmu_pro_max_tokens,
         )
         self.assert_dspark_accepts_draft_tokens()
 
@@ -123,6 +132,24 @@ class TestDSV41FlashDP4EP4DSpark(DSV41FlashAccuracyMixin, CustomTestCase):
     ]
 
 
+class TestDSV41FlashTP4EP4DSparkBoundedReplayBCG(TestDSV41FlashTP4EP4DSpark):
+    """Exercise bounded late-layer replay inside the breakable prefill graph."""
+
+    # Leave memory for the prefill graph pool. With 32 requests, the tail
+    # capacity is 4096 rows, below the largest 8192-token capture shape.
+    mem_fraction_static = "0.7"
+    max_running_requests = "32"
+    graph_args = [
+        "--enable-decoder-swa-bounded-replay",
+        "--cuda-graph-backend-prefill",
+        "breakable",
+        "--cuda-graph-max-bs-prefill",
+        "8192",
+        "--chunked-prefill-size",
+        "8192",
+    ]
+
+
 class TestDSV41FlashTrtllmBoundedBCG(DSV41FlashAccuracyMixin, CustomTestCase):
     """TRT-LLM varlen attention through the full and bounded-tail graph breaks."""
 

test/registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py

@@ -386,6 +386,7 @@ def test_body_replay_uses_plural_embeds_and_skips_embedding_on_later_pp_stage(se
                 with self.subTest(first_rank=first_rank, positional=positional):
                     runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
                     runner._is_full_backend = False
+                    runner.is_dsv41_model = False
                     runner._qwen_bcg_hc_sidechannel = False
                     runner._qwen_bcg_mtp_draft = False
                     runner._input_embeds_arg_idx = 3
@@ -505,6 +506,7 @@ def _check_bcg_eager_tail(self, live_embeds):
 
         runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
         runner._is_full_backend = False
+        runner.is_dsv41_model = False
         runner._qwen_bcg_hc_sidechannel = False
         runner._qwen_bcg_mtp_draft = True
         runner._input_embeds_arg_idx = None
@@ -531,6 +533,7 @@ def _check_bcg_eager_tail(self, live_embeds):
     def test_bcg_mtp_draft_replay_fills_hc_width_input_embeds_slot(self):
         runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
         runner._is_full_backend = False
+        runner.is_dsv41_model = False
         runner._qwen_bcg_hc_sidechannel = False
         runner._qwen_bcg_mtp_draft = True
         runner._input_embeds_arg_idx = 3

test/registered/unit/models/test_deepseek_v41_bcg_multimodal.py

@@ -0,0 +1,121 @@
+"""Regression tests for multimodal IDs and Engram masking during graph replay."""
+
+import unittest
+from contextlib import nullcontext
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+import sglang.srt.models.deepseek_v4 as dsv4
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.srt.model_executor.runner.prefill_cuda_graph_runner import (
+    PrefillCudaGraphRunner,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="base-a-test-cpu")
+
+
+class TestDeepseekV41BCGMultimodal(CustomTestCase):
+    def test_replay_uses_prepared_image_ids_without_mutating_scheduler_ids(self):
+        for keyword in (False, True):
+            with self.subTest(keyword=keyword):
+                scheduler_ids = torch.tensor([7, dsv4.MM_PAD_SHIFT_VALUE + 1, 9, 0])
+                static_ids = scheduler_ids.clone()
+                image_id = 42
+                layer = SimpleNamespace(forward=None)
+                runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
+                runner._is_full_backend = False
+                runner.is_dsv41_model = True
+                runner._qwen_bcg_hc_sidechannel = False
+                runner._use_draft_input_embeds = False
+                runner.layer_model = layer
+                runner.buffer_registry = SimpleNamespace(has_slot=lambda name: False)
+                runner._prefill_forward_context = lambda *a, **kw: nullcontext()
+                runner.backend = SimpleNamespace(
+                    replay=lambda *a, **kw: static_ids.clone()
+                )
+
+                def outer_forward(ids, positions, batch):
+                    model = SimpleNamespace(
+                        vision=object(), config=SimpleNamespace(image_token_id=image_id)
+                    )
+                    mapped, _ = (
+                        dsv4.DeepseekV4ForCausalLM.prepare_language_model_inputs(
+                            model, ids, batch
+                        )
+                    )
+                    if keyword:
+                        return layer.forward(
+                            input_ids=mapped, positions=positions, forward_batch=batch
+                        )
+                    return layer.forward(mapped, positions, batch)
+
+                runner.model_runner = SimpleNamespace(
+                    pp_group=SimpleNamespace(is_first_rank=True),
+                    model=SimpleNamespace(forward=outer_forward),
+                )
+                live = SimpleNamespace(input_ids=scheduler_ids, mm_input_embeds=None)
+                static = SimpleNamespace(
+                    input_ids=static_ids,
+                    positions=torch.arange(4),
+                    mm_input_embeds=None,
+                    mm_inputs=None,
+                    forward_mode=ForwardMode.EXTEND,
+                )
+                actual = runner._execute_body_capture(live, static, 4, 3, None)
+                torch.testing.assert_close(actual, torch.tensor([7, image_id, 9, 0]))
+                self.assertEqual(scheduler_ids[1].item(), dsv4.MM_PAD_SHIFT_VALUE + 1)
+
+    def test_image_engram_mask_does_not_depend_on_capture_batch_mm_metadata(self):
+        image_id = 42
+        ids = torch.tensor([7, image_id, 9])
+        hidden = torch.arange(6, dtype=torch.float32).reshape(3, 1, 2)
+
+        class Engram:
+            layer_hash_index = 0
+
+            def __call__(self, residual, *args, **kwargs):
+                return residual + 10
+
+        layer = SimpleNamespace(
+            engram=Engram(),
+            hc_cfg=None,
+            forward_hc_pre_from_prev=lambda **kw: kw["state"],
+        )
+        model = SimpleNamespace(
+            pp_group=SimpleNamespace(world_size=1),
+            engram_hasher=lambda *a: torch.zeros((3, 1), dtype=torch.long),
+            late_layer_start=None,
+            start_layer=0,
+            end_layer=1,
+            layers=[layer],
+            config=SimpleNamespace(
+                model_type="deepseek_v41", vision_n_layers=1, image_token_id=image_id
+            ),
+        )
+        expected = hidden + 10
+        expected[1] = hidden[1]
+        for contains_mm in (False, True):
+            batch = SimpleNamespace(
+                forward_mode=ForwardMode.EXTEND, contains_mm_inputs=lambda: contains_mm
+            )
+            with (
+                patch.object(dsv4, "is_cp_active", return_value=False),
+                patch.object(dsv4, "is_in_breakable_cuda_graph", return_value=False),
+                patch.object(
+                    dsv4, "get_forward", return_value=SimpleNamespace(sp_active=False)
+                ),
+                patch.object(dsv4, "check_cuda_graph_backend", return_value=True),
+                patch.object(dsv4, "_is_hip", False),
+            ):
+                actual, _, _ = dsv4.DeepseekV4Model._forward_layers_hc_pre_from_prev(
+                    model, torch.arange(3), hidden, batch, ids, ids, False, []
+                )
+            torch.testing.assert_close(actual, expected)
+
+
+if __name__ == "__main__":
+    unittest.main()