Add a nullcontext placeholder in the forward path for KV-canary (#26801)

python/sglang/srt/model_executor/cuda_graph_runner.py

@@ -1122,6 +1122,8 @@ class CudaGraphRunner:
 
             self.deepep_adapter.capture(is_extend_in_batch=False)
 
+            canary_ctx = contextlib.nullcontext()
+            with canary_ctx:
                 for _ in range(2):
                     self.device_module.synchronize()
                     self.model_runner.tp_group.barrier()

python/sglang/srt/model_executor/model_runner.py

@@ -3203,7 +3203,11 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             if torch.autograd._profiler_enabled()
             else contextlib.nullcontext()
         )
+
+        canary_ctx = contextlib.nullcontext()
+
         with (
+            canary_ctx,
             step_span_ctx,
             get_global_expert_distribution_recorder().with_forward_pass(
                 self.forward_pass_id,

python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py

@@ -433,7 +433,11 @@ class EAGLEDraftExtendCudaGraphRunner:
                 spec_info=spec_info,
             )
             self.deepep_adapter.capture(is_extend_in_batch=True)
+
+            canary_ctx = contextlib.nullcontext()
+            with canary_ctx:
                 self._capture_init(run_once)
+
                 out = self._capture_graph(
                     graph, get_global_graph_memory_pool(), stream, run_once
                 )

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -363,10 +363,13 @@ class EagleDraftWorker(BaseDraftWorker):
             self.speculative_num_steps,
         )
 
+        canary_outside_ctx = contextlib.nullcontext()
+
+        with canary_outside_ctx:
             # Run draft
             if can_cuda_graph:
-            parent_list, top_scores_index, draft_tokens = self.cuda_graph_runner.replay(
+                parent_list, top_scores_index, draft_tokens = (
-                forward_batch,
+                    self.cuda_graph_runner.replay(forward_batch)
                 )
             else:
                 if (
@@ -500,9 +503,10 @@ class EagleDraftWorker(BaseDraftWorker):
             # Run forward under a per-step ForwardContext so the model layer
             # reads attn_backends[i] for the i-th draft step. ``_forward_raw``
             # honors the outer context and does not override.
+            canary_index_ctx = contextlib.nullcontext()
             with forward_context(
                 ForwardContext(attn_backend=self.draft_attn_backend.attn_backends[i])
-            ):
+            ), canary_index_ctx:
                 logits_output = self.draft_runner.forward(
                     forward_batch, skip_attn_backend_init=True
                 ).logits_output
@@ -614,6 +618,9 @@ class EagleDraftWorker(BaseDraftWorker):
         forward_batch.return_logprob = False
         if mm_input_embeds is not None:
             forward_batch.mm_input_embeds = mm_input_embeds
+
+        canary_ctx = contextlib.nullcontext()
+        with canary_ctx:
             logits_output = self.draft_runner.forward(forward_batch).logits_output
         maybe_detect_nan(logits_output.next_token_logits, "draft_extend_for_prefill")
         maybe_detect_inf(logits_output.next_token_logits, "draft_extend_for_prefill")
@@ -668,6 +675,9 @@ class EagleDraftWorker(BaseDraftWorker):
             self.cuda_graph_runner_for_draft_extend
             and self.cuda_graph_runner_for_draft_extend.can_run(forward_batch)
         )
+
+        canary_ctx = contextlib.nullcontext()
+        with canary_ctx:
             if can_cuda_graph:
                 draft_logits_output = self.cuda_graph_runner_for_draft_extend.replay(
                     forward_batch