fix: enable Kimi multimodal breakable prefill cuda graph replay (#31391)

test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py

@@ -12,6 +12,13 @@ from sglang.srt.model_executor.cuda_graph_config import (
     CudaGraphConfig,
     PhaseConfig,
 )
+from sglang.srt.model_executor.forward_batch_info import (
+    CaptureHiddenMode,
+    ForwardMode,
+)
+from sglang.srt.model_executor.runner.prefill_cuda_graph_runner import (
+    PrefillCudaGraphRunner,
+)
 from sglang.srt.server_args import ServerArgs
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
@@ -20,6 +27,29 @@ register_cpu_ci(est_time=1, suite="base-a-test-cpu")
 
 
 class TestMultimodalPiecewiseCudaGraph(CustomTestCase):
+    def _make_prefill_runner(self, backend):
+        runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
+        runner._is_full_backend = False
+        runner.prefill_backend_name = backend
+        runner.has_mha_companion_layers = backend == Backend.BREAKABLE
+        runner.capture_hidden_mode = CaptureHiddenMode.NULL
+        runner.max_num_tokens = 16
+        return runner
+
+    def _make_multimodal_forward_batch(self):
+        return SimpleNamespace(
+            batch_size=1,
+            input_embeds=None,
+            replace_embeds=None,
+            mm_inputs=[object()],
+            forward_mode=ForwardMode.EXTEND,
+            capture_hidden_mode=CaptureHiddenMode.NULL,
+            global_num_tokens_cpu=None,
+            return_logprob=False,
+            input_ids=[1, 2, 3, 4],
+            extend_prefix_lens_cpu=[0],
+        )
+
     def test_kimi_k25_lm_prefill_is_opted_in(self):
         self.assertTrue(
             is_multimodal_piecewise_cuda_graph_supported(
@@ -52,6 +82,23 @@ class TestMultimodalPiecewiseCudaGraph(CustomTestCase):
         self.assertEqual(args.cuda_graph_config.prefill.backend, Backend.TC_PIECEWISE)
         disable_if_incompatible.assert_called_once()
 
+    def test_multimodal_inputs_keep_tc_piecewise_prefill_enabled(self):
+        runner = self._make_prefill_runner(Backend.TC_PIECEWISE)
+
+        self.assertTrue(runner.can_run_graph(self._make_multimodal_forward_batch()))
+
+    def test_multimodal_inputs_keep_breakable_prefill_enabled(self):
+        runner = self._make_prefill_runner(Backend.BREAKABLE)
+
+        self.assertTrue(runner.can_run_graph(self._make_multimodal_forward_batch()))
+
+    def test_breakable_prefill_rejects_nonzero_prefix(self):
+        runner = self._make_prefill_runner(Backend.BREAKABLE)
+        forward_batch = self._make_multimodal_forward_batch()
+        forward_batch.extend_prefix_lens_cpu = [1]
+
+        self.assertFalse(runner.can_run_graph(forward_batch))
+
 
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/model_executor/model_runner_components/test_layer_setup.py

@@ -0,0 +1,36 @@
+"""Unit tests for model-runner layer discovery."""
+
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.model_executor.model_runner_components.layer_setup import (
+    compute_attention_and_moe_layers,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+class TestComputeAttentionAndMoeLayers(unittest.TestCase):
+    def test_deepseek_mla_registers_mha_companion(self):
+        attn_mqa = SimpleNamespace()
+        attn_mha = SimpleNamespace()
+        layer_model = SimpleNamespace(
+            layers=[
+                SimpleNamespace(
+                    self_attn=SimpleNamespace(attn_mqa=attn_mqa, attn_mha=attn_mha)
+                )
+            ]
+        )
+
+        attention_layers, _, _, _, mha_companion_layers = (
+            compute_attention_and_moe_layers(layer_model)
+        )
+
+        self.assertEqual(attention_layers, [attn_mqa])
+        self.assertEqual(mha_companion_layers, [attn_mha])
+        self.assertNotIn("_pcg_mha_companion", vars(attn_mqa))
+
+
+if __name__ == "__main__":
+    unittest.main()