diff --git a/python/pyproject.toml b/python/pyproject.toml index c1cff46ed9c4..31d96b5f0dd5 100755 --- a/python/pyproject.toml +++ b/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", ] diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index 7209bd003b19..119997e5bf88 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/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 diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index ef6df00ba22f..fa27c7ed904b 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/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], diff --git a/test/registered/e2e/models/test_deepseek_v41_flash_gb300.py b/test/registered/e2e/models/test_deepseek_v41_flash_gb300.py index 2173ae71d59b..d12bbb3a18ee 100644 --- a/test/registered/e2e/models/test_deepseek_v41_flash_gb300.py +++ b/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.""" diff --git a/test/registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py b/test/registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py index 7f554e313f5a..37b07ba2cb0a 100644 --- a/test/registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py +++ b/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 diff --git a/test/registered/unit/models/test_deepseek_v41_bcg_multimodal.py b/test/registered/unit/models/test_deepseek_v41_bcg_multimodal.py new file mode 100644 index 000000000000..9a0ab7c7d7f4 --- /dev/null +++ b/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()