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

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()