diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py index 2f4ddb5c2..5f2006edb 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/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 diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index ddcf2c620..11084782b 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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, diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index e5a73ffe3..d3cb9acab 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/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 diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 311e7bd93..11a1ba1b6 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/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,