refactor(runner): reuse a prepared static buffer for every dummy run (#28740)

python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py

@@ -1,4 +1,5 @@
 import logging
+import math
 import os
 import time
 from contextlib import contextmanager, nullcontext
@@ -434,14 +435,28 @@ def pp_parallel_deep_gemm_warmup(runner) -> None:
     # in-seq-split). _dummy_run does not pad q/hidden like the real flow, so
     # an unaligned bs makes DSA's padded num_splits longer than the q tokens
     # and trips FlashMLA's "num_splits must have shape (b+1)" check.
+    from sglang.srt.layers.dp_attention import get_attention_tp_size
     from sglang.srt.layers.utils.cp_utils import get_cp_padding_align_size
+    from sglang.srt.utils.common import require_mlp_sync
 
     n_sms = torch.cuda.get_device_properties(model_runner.device).multi_processor_count
     block_m = 64
     cp = max(get_cp_padding_align_size(), 1)
+
+    attn_tp_size = get_attention_tp_size()
+    mlp_sync = require_mlp_sync(model_runner.server_args)
+
+    def _align(bs: int) -> int:
+        # Align to lcm(cp, attn_tp_size) so the CP multiple isn't undone by a
+        # later attn_tp align (e.g. cp=2, attn_tp=3: 128 -> 128 -> 129).
+        align = cp
+        if mlp_sync and attn_tp_size > 1:
+            align = math.lcm(cp, attn_tp_size)
+        return ceil_align(bs, align)
+
     batch_sizes = sorted(
         {
-            ceil_align(bs, cp)
+            _align(bs)
             for bs in (
                 1,
                 2 * block_m,
@@ -467,16 +482,24 @@ def pp_parallel_deep_gemm_warmup(runner) -> None:
         disagg_mode,
     )
 
+    # One buffer set sized to the largest shape, reused across the sweep
+    # (the decode runner's max_bs is too small for n_sms*block_m).
+    dummy_buffers = runner._alloc_dummy_decode_buffers(max(batch_sizes))
+
     t0 = time.perf_counter()
     with torch.inference_mode():
         for bs in batch_sizes:
             if run_decode:
                 runner._dummy_run(
-                    batch_size=bs, forward_mode_override=ForwardMode.DECODE
+                    batch_size=bs,
+                    forward_mode_override=ForwardMode.DECODE,
+                    buffers=dummy_buffers,
                 )
             if run_extend:
                 runner._dummy_run(
-                    batch_size=bs, forward_mode_override=ForwardMode.EXTEND
+                    batch_size=bs,
+                    forward_mode_override=ForwardMode.EXTEND,
+                    buffers=dummy_buffers,
                 )
 
     logger.info(

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

@@ -22,7 +22,7 @@ import logging
 from abc import ABC, abstractmethod
 from pathlib import Path
 from types import SimpleNamespace
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, Optional, Tuple
 
 import torch
 
@@ -52,7 +52,6 @@ from sglang.srt.utils import (
     require_gathered_buffer,
     require_mlp_tp_gather,
 )
-from sglang.srt.utils.common import ceil_align, require_mlp_sync
 
 if TYPE_CHECKING:
     from sglang.srt.model_executor.model_runner import ModelRunner
@@ -210,7 +209,11 @@ class BaseRunner(ABC):
         self._pre_initialize_flashinfer_allreduce_workspace()
 
         if self._should_run_flashinfer_autotune():
-            self._flashinfer_autotune()
+            buffers, batch_size = self._autotune_buffers()
+            assert (
+                buffers is not None
+            ), "_autotune_buffers() must return a reusable buffer set for autotune"
+            self._flashinfer_autotune(buffers=buffers, batch_size=batch_size)
 
         if (
             envs.SGLANG_PP_PARALLEL_DEEPGEMM_WARMUP.get()
@@ -313,8 +316,14 @@ class BaseRunner(ABC):
 
         return True
 
-    def _flashinfer_autotune(self):
+    def _flashinfer_autotune(self, *, buffers, batch_size):
-        """Run flashinfer autotune."""
+        """Run flashinfer autotune.
+
+        buffers / batch_size: a prepared static decode-buffer set and its bs,
+        reused for the dummy forward instead of allocating a throwaway set.
+        Supplied by warmup() (the decode runner's captured buffers when a graph
+        runner exists; a freshly-allocated dummy set in the eager path).
+        """
         from flashinfer.autotuner import autotune
 
         from sglang.srt.layers.logits_processor import autotune_dummy_run_mode
@@ -346,7 +355,7 @@ class BaseRunner(ABC):
                 autotune(True, cache=str(autotune_cache)),
                 autotune_dummy_run_mode(),
             ):
-                self._dummy_run(batch_size=mr.req_to_token_pool.size)
+                self._dummy_run(batch_size=batch_size, buffers=buffers)
         torch.cuda.current_stream().wait_stream(mr.forward_stream)
         logger.info("FlashInfer autotune completed.")
 
@@ -386,16 +395,66 @@ class BaseRunner(ABC):
             cache_dir / f"rank_tp{mr.tp_rank}_pp{mr.pp_rank}_dp{mr.dp_rank or 0}.json"
         )
 
+    def _alloc_dummy_decode_buffers(self, max_bs: int, *, num_tokens_per_bs: int = 1):
+        """Allocate one static decode-buffer set for a dummy forward, sized to
+        (max_bs, max_bs * num_tokens_per_bs).
+
+        The PP-parallel DeepGEMM warmup sweeps batch sizes far larger than any
+        runner's max_bs (up to ~n_sms*block_m), so no pre-allocated runner buffer
+        set fits; it builds one here and hands it to _dummy_run (reused across the
+        sweep; _dummy_run slices it per shape). The flashinfer autotune does NOT
+        use this -- it reuses an existing runner's buffers via _autotune_buffers
+        (the eager input registry, or the decode cuda-graph runner's captured
+        buffers).
+        """
+        mr = self.model_runner
+        return _allocate_decode_buffers(
+            device=mr.device,
+            max_bs=max_bs,
+            max_num_token=max_bs * num_tokens_per_bs,
+            hidden_size=mr.model_config.hidden_size,
+            vocab_size=mr.model_config.vocab_size,
+            dtype=mr.model_config.dtype,
+            dp_size=mr.server_args.dp_size,
+            pp_size=mr.server_args.pp_size,
+            is_encoder_decoder=mr.model_config.is_encoder_decoder,
+            require_mlp_tp_gather=require_mlp_tp_gather(mr.server_args),
+            seq_len_fill_value=mr.attn_backend.get_cuda_graph_seq_len_fill_value(),
+            encoder_len_fill_value=(
+                getattr(mr.model_config.hf_config, "max_source_positions", 0)
+                if mr.model_config.is_encoder_decoder
+                else 0
+            ),
+            num_tokens_per_bs=num_tokens_per_bs,
+            cache_loc_dtype=torch.int64,
+            enable_mamba_track=False,
+            hc_hidden_size=getattr(mr.model_config, "hc_hidden_size", None),
+        )
+
     def _dummy_run(
         self,
         batch_size: int,
         run_ctx=None,
         forward_mode_override: Optional[ForwardMode] = None,
+        *,
+        buffers,
     ):
         """Run a dummy forward pass for warmup/profiling.
 
         forward_mode_override forces EXTEND/DECODE regardless of
         is_generation (used by the PP-parallel DeepGEMM warmup).
+
+        buffers: a prepared static buffer set (or lightweight adapter exposing
+        the same fields), sized >= this dummy shape, which _dummy_run slices to
+        (batch_size, num_tokens). The caller owns the shape and the allocation --
+        the flashinfer autotune reuses an existing runner's buffers via
+        _autotune_buffers (the eager input registry, or the decode cuda-graph
+        runner's captured buffers); the PP-DeepGEMM warmup builds one via
+        _alloc_dummy_decode_buffers. _dummy_run never allocates and never re-pads
+        (autotune must run at the reused shape; the PP warmup pre-pads and sizes
+        its buffer to match). next_token_logits_buffer is optional -- a live
+        autotune forward returns logits fresh, so the eager-reuse path passes
+        None (only the PP warmup set still carries one).
         """
         mr = self.model_runner
         if forward_mode_override is not None:
@@ -422,12 +481,22 @@ class BaseRunner(ABC):
 
         num_tokens = batch_size * num_tokens_per_bs
 
-        # Keep warmup aligned with scheduler MLP-sync padding.
+        # Caller owns the shape: passes a static buffer >= the dummy shape; no
-        if require_mlp_sync(mr.server_args):
+        # allocation, no re-padding (would overflow the reused buffers).
-            attn_tp_size = get_parallel().attn_tp_size
+        assert (
-            if attn_tp_size > 1 and num_tokens % attn_tp_size != 0:
+            buffers is not None
-                num_tokens = ceil_align(num_tokens, attn_tp_size)
+            and num_tokens <= buffers.input_ids.shape[0]
-                batch_size = num_tokens // num_tokens_per_bs
+            and batch_size <= buffers.seq_lens.shape[0]
+        ), (
+            f"_dummy_run needs a static buffer >= (num_tokens={num_tokens}, "
+            f"batch_size={batch_size}); got "
+            + (
+                "None"
+                if buffers is None
+                else f"(input_ids={buffers.input_ids.shape[0]}, "
+                f"seq_lens={buffers.seq_lens.shape[0]})"
+            )
+        )
 
         seq_len_fill_value = mr.attn_backend.get_cuda_graph_seq_len_fill_value()
 
@@ -451,28 +520,25 @@ class BaseRunner(ABC):
         if require_gathered_buffer(mr.server_args):
             assert require_mlp_tp_gather_ or require_attn_tp_gather(mr.server_args)
 
-        buffers = _allocate_decode_buffers(
+        input_ids = buffers.input_ids[:num_tokens]
-            device=mr.device,
+        positions = buffers.positions[:num_tokens]
-            max_bs=batch_size,
+        out_cache_loc = buffers.out_cache_loc[:num_tokens]
-            max_num_token=num_tokens,
+        # Eager-reuse drops the logits buffer; only buffer sets that carry one slice it.
-            hidden_size=mr.model_config.hidden_size,
+        next_token_logits_buffer = (
-            vocab_size=mr.model_config.vocab_size,
+            buffers.next_token_logits_buffer[:num_tokens]
-            dtype=mr.model_config.dtype,
+            if buffers.next_token_logits_buffer is not None
-            dp_size=mr.server_args.dp_size,
+            else None
-            pp_size=mr.server_args.pp_size,
-            is_encoder_decoder=mr.model_config.is_encoder_decoder,
-            require_mlp_tp_gather=require_mlp_tp_gather_,
-            seq_len_fill_value=seq_len_fill_value,
-            encoder_len_fill_value=(
-                getattr(mr.model_config.hf_config, "max_source_positions", 0)
-                if mr.model_config.is_encoder_decoder
-                else 0
-            ),
-            num_tokens_per_bs=num_tokens_per_bs,
-            cache_loc_dtype=torch.int64,
-            enable_mamba_track=False,
-            hc_hidden_size=getattr(mr.model_config, "hc_hidden_size", None),
         )
+        mrope_positions = buffers.mrope_positions[:, :num_tokens]
+        req_pool_indices = buffers.req_pool_indices[:batch_size]
+        seq_lens = buffers.seq_lens[:batch_size]
+        seq_lens_cpu = buffers.seq_lens_cpu[:batch_size]
+        encoder_lens = (
+            buffers.encoder_lens[:batch_size]
+            if buffers.encoder_lens is not None
+            else None
+        )
+
         buffers.num_token_non_padded[...] = num_tokens
 
         # For extend mode
@@ -558,17 +624,17 @@ class BaseRunner(ABC):
         forward_batch = ForwardBatch(
             forward_mode=capture_forward_mode,
             batch_size=batch_size,
-            input_ids=buffers.input_ids,
+            input_ids=input_ids,
-            req_pool_indices=buffers.req_pool_indices,
+            req_pool_indices=req_pool_indices,
-            seq_lens=buffers.seq_lens,
+            seq_lens=seq_lens,
-            seq_lens_cpu=buffers.seq_lens_cpu,
+            seq_lens_cpu=seq_lens_cpu,
-            next_token_logits_buffer=buffers.next_token_logits_buffer,
+            next_token_logits_buffer=next_token_logits_buffer,
-            orig_seq_lens=buffers.seq_lens,
+            orig_seq_lens=seq_lens,
-            out_cache_loc=buffers.out_cache_loc,
+            out_cache_loc=out_cache_loc,
-            seq_lens_sum=buffers.seq_lens.sum().item(),
+            seq_lens_sum=seq_lens.sum().item(),
-            encoder_lens=buffers.encoder_lens,
+            encoder_lens=encoder_lens,
             return_logprob=False,
-            positions=buffers.positions,
+            positions=positions,
             extend_num_tokens=extend_num_tokens,
             extend_seq_lens=extend_seq_lens,
             extend_prefix_lens=extend_prefix_lens,
@@ -580,7 +646,7 @@ class BaseRunner(ABC):
             global_num_tokens_for_logprob_gpu=buffers.global_num_tokens_for_logprob_gpu,
             dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(),
             global_dp_buffer_len=global_dp_buffer_len,
-            mrope_positions=buffers.mrope_positions,
+            mrope_positions=mrope_positions,
             spec_algorithm=mr.spec_algorithm,
             spec_info=spec_info,
             capture_hidden_mode=capture_hidden_mode,
@@ -616,7 +682,7 @@ class BaseRunner(ABC):
                 kwargs["get_embedding"] = True
 
             logits_output_or_pp_proxy_tensors = mr.model.forward(
-                buffers.input_ids,
+                input_ids,
                 forward_batch.positions,
                 forward_batch,
                 **kwargs,
@@ -629,6 +695,11 @@ class BaseRunner(ABC):
             with torch.inference_mode(), run_ctx or empty_context():
                 run_once()
 
+    def _autotune_buffers(self) -> Tuple[Optional[Any], Optional[int]]:
+        """Return (buffers, bs) for the autotune dummy forward to reuse; the
+        EagerRunner and DecodeCudaGraphRunner override this."""
+        return None, None
+
     @abstractmethod
     def can_run_graph(self, forward_batch: ForwardBatch) -> bool: ...
 

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

@@ -358,6 +358,20 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
                 f"Capture cuda graph failed: {e}\n" f"{CUDA_GRAPH_CAPTURE_FAILED_MSG}"
             )
 
+    def _autotune_buffers(self):
+        """Reuse these static decode buffers (sized to max_bs) for the warmup
+        flashinfer-autotune dummy forward instead of allocating a throwaway set
+        — see BaseRunner._autotune_buffers / BaseRunner._dummy_run.
+
+        The dummy forward derives its shape from max_bs and must match these
+        buffers exactly; _dummy_run asserts that. Every autotune-reachable
+        decode shape (plain decode, spec target-verify) matches. DLLM would not
+        (its buffers hold block_size tokens/bs while the dummy run derives 1),
+        but DLLM does not use a flashinfer MoE backend, so autotune never runs
+        for it and this is never reached there.
+        """
+        return self.buffers, self.max_bs
+
     def maybe_init_pdmux(self):
         if self.enable_pdmux:
             self.stream_groups = get_stream_groups()

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

@@ -18,7 +18,8 @@ from __future__ import annotations
 import contextlib
 import logging
 from dataclasses import replace
-from typing import TYPE_CHECKING, Any, Tuple, Union
+from types import SimpleNamespace
+from typing import TYPE_CHECKING, Any, Optional, Tuple, Union
 
 import torch
 
@@ -41,7 +42,7 @@ from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph impo
     enable_tc_piecewise_cuda_graph,
     set_tc_piecewise_forward_context,
 )
-from sglang.srt.utils import is_hip
+from sglang.srt.utils import is_hip, require_mlp_tp_gather
 
 logger = logging.getLogger(__name__)
 
@@ -99,6 +100,8 @@ class EagerRunner(BaseRunner):
             else mr.max_total_num_tokens
         )
         max_num_token = max(prefill_ceiling, max_bs * num_tokens_per_bs)
+        self._eager_max_bs = max_bs
+        self._eager_num_tokens_per_bs = num_tokens_per_bs
         is_encoder_decoder = mr.model_config.is_encoder_decoder
         self._eager_registry = build_eager_registry(
             device=mr.device,
@@ -119,6 +122,83 @@ class EagerRunner(BaseRunner):
         # Eager has no capture step, so warm up here (run-once via mr._kernel_warmed_up).
         self.warmup()
 
+    def _autotune_buffers(self) -> Tuple[Any, int]:
+        """Adapter over the eager registry for the autotune dummy forward; fills
+        in fields the registry omits (logits buffer, pp_proxy, custom_mask)."""
+        mr = self.model_runner
+        reg = self._eager_registry
+        max_bs = self._eager_max_bs
+
+        def _slot(name):
+            return reg.get_slot(name).buffer if reg.has_slot(name) else None
+
+        # num_token_non_padded / global_num_tokens_* are not registered on the
+        # eager registry (build_eager_registry passes enable_num_token_non_padded
+        # =False, register_global_num_tokens=False); _dummy_run writes + reads
+        # them unconditionally, so supply tiny fresh tensors here.
+        num_token_non_padded = torch.zeros((1,), dtype=torch.int32, device=mr.device)
+        global_dim = (
+            mr.server_args.dp_size if require_mlp_tp_gather(mr.server_args) else 1
+        )
+        global_num_tokens_gpu = torch.zeros(
+            (global_dim,), dtype=torch.int32, device=mr.device
+        )
+        global_num_tokens_for_logprob_gpu = torch.zeros(
+            (global_dim,), dtype=torch.int32, device=mr.device
+        )
+
+        # custom_mask: only consumed by create_dummy_verify_input (spec). Size it
+        # like the decode path's custom_mask for a spec target worker.
+        custom_mask: Optional[torch.Tensor] = None
+        if mr.spec_algorithm.is_speculative():
+            num_tokens_per_bs = self._eager_num_tokens_per_bs
+            max_num_token = reg.max_num_tokens
+            seq_len_fill_value = mr.attn_backend.get_cuda_graph_seq_len_fill_value()
+            custom_mask = torch.ones(
+                (max_bs * seq_len_fill_value + max_num_token) * num_tokens_per_bs,
+                dtype=torch.bool,
+                device=mr.device,
+            )
+
+        # pp_proxy_tensors: only read when pp_size>1. _dummy_run slices each value
+        # [:pp_hidden_tokens] (pp_hidden_tokens <= num_tokens), so size the first
+        # dim to the registry's token ceiling. Mirror _allocate_decode_buffers'
+        # keys/dtypes (mHC flattens residual into hidden_states of hc_hidden_size).
+        pp_proxy_tensors = None
+        if mr.server_args.pp_size > 1:
+            hidden_size = mr.model_config.hidden_size
+            hc_hidden_size = getattr(mr.model_config, "hc_hidden_size", None)
+            is_mhc = hc_hidden_size is not None
+            hs = hc_hidden_size if is_mhc else hidden_size
+            rows = reg.max_num_tokens
+            pp_proxy_tensors = {
+                "hidden_states": torch.zeros(
+                    (rows, hs), dtype=mr.dtype, device=mr.device
+                ),
+            }
+            if not is_mhc:
+                pp_proxy_tensors["residual"] = torch.zeros(
+                    (rows, hidden_size), dtype=mr.dtype, device=mr.device
+                )
+
+        adapter = SimpleNamespace(
+            input_ids=_slot("input_ids"),
+            positions=_slot("positions"),
+            out_cache_loc=_slot("out_cache_loc"),
+            req_pool_indices=_slot("req_pool_indices"),
+            seq_lens=_slot("seq_lens"),
+            seq_lens_cpu=_slot("seq_lens_cpu"),
+            mrope_positions=_slot("mrope_positions"),
+            encoder_lens=_slot("encoder_lens"),
+            next_token_logits_buffer=None,
+            num_token_non_padded=num_token_non_padded,
+            global_num_tokens_gpu=global_num_tokens_gpu,
+            global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu,
+            custom_mask=custom_mask,
+            pp_proxy_tensors=pp_proxy_tensors,
+        )
+        return adapter, max_bs
+
     def can_run_graph(self, forward_batch: ForwardBatch) -> bool:
         # Eager never runs a cuda graph; callers dispatch on isinstance(...,
         # EagerRunner) and must not route an eager batch into a replay branch.