[Refactor] Split the FlashInfer autotune dummy-run flag from the LM-head policy (#34336)

python/sglang/srt/layers/logits_processor.py

@@ -70,26 +70,26 @@ _UNQUANTIZED_LM_HEAD_METHODS = {
     "PackWeightMethod",
 }
 
-# When set, LogitsProcessor.forward returns an empty output and skips the
-# LM head + tensor-parallel all-gather. FlashInfer autotune only profiles
-# attention/MoE/GEMM kernels, so the LM-head all-gather is wasted work --
-# and its [batch * dp_size, vocab] output OOMs under DP attention with a
-# tight mem_fraction_static.
-_in_autotune_dummy_run = False
+# None outside a FlashInfer autotune pass; inside one, whether that pass runs the
+# LM head. Not-None means the forward's output is discarded -- attention backends
+# read that via get_in_autotune_dummy_run() to skip a cross-node exchange.
+# Skipping the LM head skips its [batch * dp_size, vocab] all-gather, which OOMs
+# under DP attention with a tight mem_fraction_static.
+_autotune_run_lm_head: Optional[bool] = None
 
 
 def get_in_autotune_dummy_run() -> bool:
-    return _in_autotune_dummy_run
+    return _autotune_run_lm_head is not None
 
 
 @contextmanager
-def autotune_dummy_run_mode():
-    global _in_autotune_dummy_run
-    _in_autotune_dummy_run = True
+def autotune_dummy_run_mode(*, run_lm_head: bool):
+    global _autotune_run_lm_head
+    _autotune_run_lm_head = run_lm_head
     try:
         yield
     finally:
-        _in_autotune_dummy_run = False
+        _autotune_run_lm_head = None
 
 
 @dataclasses.dataclass
@@ -344,10 +344,10 @@ class LogitsProcessor(nn.Module):
             multi_item_delimiter_indices = logits_metadata.multi_item_delimiter_indices
             logits_metadata = LogitsMetadata.from_forward_batch(logits_metadata)
 
-        # Autotune dummy run discards this output; see _in_autotune_dummy_run.
-        # Placed before the MIS / DLLM / common dispatch so all three LM-head
-        # paths are skipped.
-        if _in_autotune_dummy_run:
+        # Autotune dummy run discards this output. `is False` not `not`: None
+        # means no autotune pass, which must not skip. Placed before the MIS /
+        # DLLM / common dispatch so all three LM-head paths are skipped.
+        if _autotune_run_lm_head is False:
             return LogitsProcessorOutput(next_token_logits=None)
 
         # Multi-item scoring only for prefill-only requests with pre-computed indices.

python/sglang/srt/model_executor/runner/base_runner.py

@@ -100,7 +100,7 @@ def _allocate_decode_buffers(
             dtype=torch.bool,
         )
         # (max_num_token, vocab) fp32 is large (>10GB at 16k tokens); callers
-        # whose dummy runs never touch logits (skip_logits autotune) opt out.
+        # whose dummy runs never touch logits (run_lm_head=False autotune) opt out.
         next_token_logits_buffer = (
             torch.zeros(
                 (max_num_token, vocab_size),
@@ -314,7 +314,9 @@ class BaseRunner(ABC):
                 run_ctx=canary_run_ctx,
             )
 
-        run_flashinfer_autotune_forward(self.model_runner, forward_fn, skip_logits=True)
+        run_flashinfer_autotune_forward(
+            self.model_runner, forward_fn, run_lm_head=False
+        )
 
     def _alloc_dummy_decode_buffers(
         self,

python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py

@@ -1133,7 +1133,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
                     self,
                     run_once,
                     post_warmup_hook=post_warmup_hook,
-                    skip_logits=False,
+                    run_lm_head=True,
                 )
                 self.backend.capture_one(
                     shape_key,

python/sglang/srt/model_executor/runner/flashinfer_autotune.py

@@ -170,7 +170,7 @@ def flashinfer_autotune_cache_path(model_runner: ModelRunner) -> Path:
 
 
 @contextlib.contextmanager
-def flashinfer_autotune_context(model_runner: ModelRunner, *, skip_logits: bool):
+def flashinfer_autotune_context(model_runner: ModelRunner, *, run_lm_head: bool):
     from flashinfer.autotuner import autotune
 
     mr = model_runner
@@ -193,27 +193,24 @@ def flashinfer_autotune_context(model_runner: ModelRunner, *, skip_logits: bool)
     # calls on default stream (unsupported by CUDA) when --enable-symm-mem is used.
     mr.forward_stream.wait_stream(torch.cuda.current_stream())
     with torch.get_device_module(mr.device).stream(mr.forward_stream):
-        maybe_skip_logits = contextlib.nullcontext()
-        if skip_logits:
-            from sglang.srt.layers.logits_processor import autotune_dummy_run_mode
+        from sglang.srt.layers.logits_processor import autotune_dummy_run_mode
 
-            maybe_skip_logits = autotune_dummy_run_mode()
         skip_ops = get_flashinfer_autotune_skip_ops(mr)
         with autotune(
             True,
             cache=str(autotune_cache),
             skip_ops=skip_ops,
-        ), maybe_skip_logits:
+        ), autotune_dummy_run_mode(run_lm_head=run_lm_head):
             yield
     torch.cuda.current_stream().wait_stream(mr.forward_stream)
     logger.info("FlashInfer autotune completed.")
 
 
 def run_flashinfer_autotune_forward(
-    model_runner: ModelRunner, forward_fn: Callable[[], None], *, skip_logits: bool
+    model_runner: ModelRunner, forward_fn: Callable[[], None], *, run_lm_head: bool
 ) -> None:
     """Run flashinfer autotune forward."""
-    with flashinfer_autotune_context(model_runner, skip_logits=skip_logits):
+    with flashinfer_autotune_context(model_runner, run_lm_head=run_lm_head):
         forward_fn()
 
 
@@ -222,7 +219,7 @@ def maybe_flashinfer_autotune_speculative_draft(
     forward_fn: Callable[[], None],
     *,
     post_warmup_hook: Optional[Callable[[], None]] = None,
-    skip_logits: bool = False,
+    run_lm_head: bool = True,
 ) -> None:
     """Run speculative draft flashinfer autotune."""
     mr = runner.model_runner
@@ -245,7 +242,7 @@ def maybe_flashinfer_autotune_speculative_draft(
         if post_warmup_hook is not None:
             post_warmup_hook()
 
-    run_flashinfer_autotune_forward(mr, run_and_reset, skip_logits=skip_logits)
+    run_flashinfer_autotune_forward(mr, run_and_reset, run_lm_head=run_lm_head)
     tuned_phases.add(phase_key)
 
 
@@ -316,7 +313,7 @@ def maybe_flashinfer_autotune_extend(
         f"({batch_size} seqs x {per_req} tokens).",
     )
     try:
-        run_flashinfer_autotune_forward(mr, forward_fn, skip_logits=True)
+        run_flashinfer_autotune_forward(mr, forward_fn, run_lm_head=False)
     except torch.OutOfMemoryError:
         # The pass is an optimization; without headroom for the extend-shaped
         # forward, fall back to untuned extend buckets instead of failing.

python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -479,7 +479,7 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
                 self,
                 run_once,
                 post_warmup_hook=post_warmup_hook,
-                skip_logits=False,
+                run_lm_head=True,
             )
             self.backend.capture_one(
                 shape_key,

python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py

@@ -456,7 +456,7 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
                     self,
                     run_once,
                     post_warmup_hook=post_warmup_hook,
-                    skip_logits=False,
+                    run_lm_head=True,
                 )
                 self.backend.capture_one(
                     shape_key,

python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py

@@ -363,7 +363,7 @@ class FrozenKVMTPCudaGraphRunner(DecodeCudaGraphRunner):
                     self,
                     run_once,
                     post_warmup_hook=post_warmup_hook,
-                    skip_logits=False,
+                    run_lm_head=True,
                 )
                 self.backend.capture_one(
                     shape_key,

python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py

@@ -451,7 +451,7 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
                 self,
                 run_once,
                 post_warmup_hook=post_warmup_hook,
-                skip_logits=False,
+                run_lm_head=True,
             )
             self.backend.capture_one(
                 shape_key,