From 22587fb15ce90a31c55e684495dfd16a9e5bf8e7 Mon Sep 17 00:00:00 2001
From: Cheng Wan <54331508+ch-wan@users.noreply.github.com>
Date: Mon, 21 Sep 2026 16:46:22 -0700
Subject: [PATCH] [Fix] Run KV canary hooks for context-parallel prefill
 (#40642)

---
 python/sglang/srt/kv_canary/api.py        |  8 ++++
 test/registered/mock_model/test_e2e_cp.py | 47 +++++++++++++++++++++++
 2 files changed, 55 insertions(+)
 create mode 100644 test/registered/mock_model/test_e2e_cp.py

diff --git a/python/sglang/srt/kv_canary/api.py b/python/sglang/srt/kv_canary/api.py
index 702518720..55d8102a1 100644
--- a/python/sglang/srt/kv_canary/api.py
+++ b/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]:
diff --git a/test/registered/mock_model/test_e2e_cp.py b/test/registered/mock_model/test_e2e_cp.py
new file mode 100644
index 000000000..700d22b3f
--- /dev/null
+++ b/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()