diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index 3d087cd94..be995d078 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/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. diff --git a/python/sglang/srt/model_executor/runner/base_runner.py b/python/sglang/srt/model_executor/runner/base_runner.py index 1597344ca..c421cc419 100644 --- a/python/sglang/srt/model_executor/runner/base_runner.py +++ b/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, diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index 230884885..b9c1eae94 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/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, diff --git a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py index f01256fd4..85818efa7 100644 --- a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py +++ b/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. diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index 13371f432..5722c40f8 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/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, 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 1d75e303d..02b0cf42c 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 @@ -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, diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py index df7fab33c..d0d9acf7e 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py +++ b/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, diff --git a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py index 23e8cdee0..97c98ad52 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py +++ b/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,