[Fix] Run KV canary hooks for context-parallel prefill (#40642)

python/sglang/srt/kv_canary/api.py

@@ -20,6 +20,7 @@ from sglang.srt.model_executor.cuda_graph_config import (
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.runtime_context import (
     get_disagg,
+    get_parallel,
     get_spec,
 )
 
@@ -123,6 +124,13 @@ def _patch_model_forward(*, model_runner: ModelRunner, manager: CanaryManager) -
             return output
 
     wrap_method(model_runner.model, "forward", wrapper=_with_canary_bracketing)
+    if get_parallel().enable_prefill_cp:
+        # CP prefill calls the transformer body directly, bypassing the outer
+        # model.forward. Decode still enters through the outer model; the shared
+        # bracket scope prevents the body from running a second pair of hooks.
+        wrap_method(
+            model_runner.model.model, "forward", wrapper=_with_canary_bracketing
+        )
 
 
 def _extract_forward_batch(args, kwargs) -> Optional[ForwardBatch]:

test/registered/mock_model/test_e2e_cp.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import unittest
+
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.mock_model.utils import run_mock_model_bench_serving
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=50, stage="extra-a", runner_config="2-gpu-large")
+
+
+class TestE2EContextParallel(CustomTestCase):
+    def test_cp_prefill_then_decode_no_canary_violation(self) -> None:
+        # CP prefill enters the transformer body directly, while decode calls
+        # the outer model.forward. Both need exactly one canary bracket.
+        run_mock_model_bench_serving(
+            extra_server_args=[
+                "--tp",
+                "2",
+                "--attn-cp-size",
+                "2",
+                "--enable-prefill-cp",
+                "--cp-strategy",
+                "zigzag",
+                "--attention-backend",
+                "fa3",
+                "--kv-canary-real-data",
+                "all",
+                "--mem-fraction-static",
+                "0.2",
+                "--max-total-tokens",
+                "4096",
+                "--max-running-requests",
+                "8",
+                "--context-length",
+                "256",
+                "--cuda-graph-max-bs-decode",
+                "4",
+            ],
+            num_prompts=4,
+            random_input_len=32,
+            random_output_len=8,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()