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

python/sglang/srt/model_executor/cuda_graph_runner.py

@@ -1122,20 +1122,22 @@ class CudaGraphRunner:
 
             self.deepep_adapter.capture(is_extend_in_batch=False)
 
-            for _ in range(2):
-                self.device_module.synchronize()
-                self.model_runner.tp_group.barrier()
-                run_once()
-                attn_backend.on_after_cuda_graph_warmup()
+            canary_ctx = contextlib.nullcontext()
+            with canary_ctx:
+                for _ in range(2):
+                    self.device_module.synchronize()
+                    self.model_runner.tp_group.barrier()
+                    run_once()
+                    attn_backend.on_after_cuda_graph_warmup()
 
-            if get_global_graph_memory_pool() is None:
-                set_global_graph_memory_pool(self.device_module.graph_pool_handle())
-            # Set graph pool id globally to be able to use symmetric memory
-            set_graph_pool_id(get_global_graph_memory_pool())
+                if get_global_graph_memory_pool() is None:
+                    set_global_graph_memory_pool(self.device_module.graph_pool_handle())
+                # Set graph pool id globally to be able to use symmetric memory
+                set_graph_pool_id(get_global_graph_memory_pool())
 
-            out = self._capture_graph(
-                graph, get_global_graph_memory_pool(), stream, run_once
-            )
+                out = self._capture_graph(
+                    graph, get_global_graph_memory_pool(), stream, run_once
+                )
 
         return graph, out
 

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,10 +433,14 @@ class EAGLEDraftExtendCudaGraphRunner:
                 spec_info=spec_info,
             )
             self.deepep_adapter.capture(is_extend_in_batch=True)
-            self._capture_init(run_once)
-            out = self._capture_graph(
-                graph, get_global_graph_memory_pool(), stream, run_once
-            )
+
+            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
+                )
 
         set_global_graph_memory_pool(graph.pool())
         return graph, out

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -363,22 +363,25 @@ class EagleDraftWorker(BaseDraftWorker):
             self.speculative_num_steps,
         )
 
-        # Run draft
-        if can_cuda_graph:
-            parent_list, top_scores_index, draft_tokens = self.cuda_graph_runner.replay(
-                forward_batch,
-            )
-        else:
-            if (
-                not forward_batch.forward_mode.is_idle()
-                and self.speculative_num_steps > 1
-            ):
-                # Skip attention backend init for 1-step draft,
-                # `draft_forward` only does sample in this case.
-                self.draft_attn_backend.init_forward_metadata(forward_batch)
-            parent_list, top_scores_index, draft_tokens = self.draft_forward(
-                forward_batch
-            )
+        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(forward_batch)
+                )
+            else:
+                if (
+                    not forward_batch.forward_mode.is_idle()
+                    and self.speculative_num_steps > 1
+                ):
+                    # Skip attention backend init for 1-step draft,
+                    # `draft_forward` only does sample in this case.
+                    self.draft_attn_backend.init_forward_metadata(forward_batch)
+                parent_list, top_scores_index, draft_tokens = self.draft_forward(
+                    forward_batch
+                )
 
         if batch.forward_mode.is_idle():
             return EagleVerifyInput.create_idle_input(
@@ -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,7 +618,10 @@ class EagleDraftWorker(BaseDraftWorker):
         forward_batch.return_logprob = False
         if mm_input_embeds is not None:
             forward_batch.mm_input_embeds = mm_input_embeds
-        logits_output = self.draft_runner.forward(forward_batch).logits_output
+
+        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,14 +675,17 @@ class EagleDraftWorker(BaseDraftWorker):
             self.cuda_graph_runner_for_draft_extend
             and self.cuda_graph_runner_for_draft_extend.can_run(forward_batch)
         )
-        if can_cuda_graph:
-            draft_logits_output = self.cuda_graph_runner_for_draft_extend.replay(
-                forward_batch
-            )
-        else:
-            draft_logits_output = self.draft_runner.forward(
-                forward_batch, skip_attn_backend_init=True
-            ).logits_output
+
+        canary_ctx = contextlib.nullcontext()
+        with canary_ctx:
+            if can_cuda_graph:
+                draft_logits_output = self.cuda_graph_runner_for_draft_extend.replay(
+                    forward_batch
+                )
+            else:
+                draft_logits_output = self.draft_runner.forward(
+                    forward_batch, skip_attn_backend_init=True
+                ).logits_output
 
         maybe_detect_nan(
             draft_logits_output.next_token_logits,