Autotune flashinfer extend buckets at warmup (#32556)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

python/sglang/srt/environ.py

@@ -1342,6 +1342,10 @@ class Envs:
     # Sglang Cache Dir
     SGLANG_CACHE_DIR = EnvStr(os.path.expanduser("~/.cache/sglang"))
     SGLANG_FLASHINFER_AUTOTUNE_CACHE = EnvBool(True)
+    # Also autotune one EXTEND-shaped dummy at max_prefill_tokens during
+    # warmup. Opt-in: the extra forward needs transient activation headroom
+    # that small-VRAM or tightly-packed configs may not have.
+    SGLANG_FLASHINFER_AUTOTUNE_EXTEND = EnvBool(False)
     SGLANG_ENABLE_MOE_DEFERRED_FINALIZE = EnvBool(True)
 
     # Plugin system

python/sglang/srt/layers/attention/aiter_backend.py

@@ -128,6 +128,11 @@ _AITER_PARTITION_SIZE_ROCM = 256
 
 
 class AiterAttnBackend(AttentionBackend):
+
+    # kv_indptr/qo_indptr are preallocated at (req pool + 1); an extend batch
+    # can never carry more seqs than the pool.
+    extend_dummy_seqs_capped_by_req_pool: bool = True
+
     def __init__(
         self,
         model_runner: ModelRunner,

python/sglang/srt/layers/attention/base_attn_backend.py

@@ -101,6 +101,11 @@ class AttentionBackend(ABC):
     # Opt out only when this backend never reads seq_lens_cpu / seq_lens_sum.
     needs_cpu_seq_lens: bool = True
 
+    # True for backends that preallocate per-seq extend metadata at req-pool
+    # size (e.g. triton's kv_indptr): dummy extend batches must then keep
+    # batch_size <= req_to_token_pool.size.
+    extend_dummy_seqs_capped_by_req_pool: bool = False
+
     # Most attention backends can rebuild and replace forward metadata before
     # every forward. BCG capture is different: some backends expose metadata
     # tensors to kernels across graph breaks, so the captured graph depends on

python/sglang/srt/layers/attention/dsa_backend.py

@@ -285,6 +285,10 @@ _DSA_IMPL_T: TypeAlias = Literal[
 class DeepseekSparseAttnBackend(
     DeepseekSparseAttnBackendMTPPrecomputeMixin, AttentionBackend
 ):
+
+    # kv_indptr/qo_indptr are preallocated at (req pool + 1); an extend batch
+    # can never carry more seqs than the pool.
+    extend_dummy_seqs_capped_by_req_pool: bool = True
     # Decode/verify/draft graph replay rebuilds metadata from static buffers
     # (page-table width) and never reads seq_lens_cpu / seq_lens_sum; opt out of
     # the D2H sync. The eager fallback derives lengths from GPU seq_lens.

python/sglang/srt/layers/attention/flashinfer_backend.py

@@ -288,6 +288,10 @@ def fast_prefill_plan(
 class FlashInferAttnBackend(AttentionBackend):
     """Flashinfer attention kernels."""
 
+    # kv_indptr/qo_indptr are preallocated at (req pool + 1); an extend batch
+    # can never carry more seqs than the pool.
+    extend_dummy_seqs_capped_by_req_pool: bool = True
+
     def __init__(
         self,
         model_runner: ModelRunner,

python/sglang/srt/layers/attention/flashinfer_mla_backend.py

@@ -208,6 +208,10 @@ class FlashInferMhaChunkKVRunner:
 class FlashInferMLAAttnBackend(AttentionBackend):
     """Flashinfer attention kernels."""
 
+    # kv_indptr/qo_indptr are preallocated at (req pool + 1); an extend batch
+    # can never carry more seqs than the pool.
+    extend_dummy_seqs_capped_by_req_pool: bool = True
+
     # Verify metadata is ragged-layout aware via generate_attn_arg_prefill;
     # graphs key their wrappers by token tier (_verify_graph_key).
     supports_ragged_verify_graph: bool = True

python/sglang/srt/layers/attention/hybrid_attn_backend.py

@@ -44,6 +44,10 @@ class HybridAttnBackend(AttentionBackend):
             self.spec_attn_is_prefill and prefill_backend.needs_cpu_seq_lens
         )
         self.max_context_len = model_runner.model_config.context_len
+        # _select_backend routes EXTEND to prefill_backend unconditionally.
+        self.extend_dummy_seqs_capped_by_req_pool = getattr(
+            prefill_backend, "extend_dummy_seqs_capped_by_req_pool", False
+        )
 
     @property
     def supports_ragged_verify_graph(self) -> bool:

python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py

@@ -962,6 +962,9 @@ class HybridLinearAttnBackend(AttentionBackend):
             full_attn_backend.needs_cpu_seq_lens
             or linear_attn_backend.needs_cpu_seq_lens
         )
+        self.extend_dummy_seqs_capped_by_req_pool = getattr(
+            full_attn_backend, "extend_dummy_seqs_capped_by_req_pool", False
+        ) or getattr(linear_attn_backend, "extend_dummy_seqs_capped_by_req_pool", False)
 
     @property
     def data_type(self):

python/sglang/srt/layers/attention/minimax_sparse_backend.py

@@ -565,6 +565,9 @@ class MiniMaxHybridAttnBackend(AttentionBackend):
         self.sparse = sparse_backend
         self.sparse_layer_ids = sparse_layer_ids
         self.sparse.dense_backend = dense_backend
+        self.extend_dummy_seqs_capped_by_req_pool = getattr(
+            dense_backend, "extend_dummy_seqs_capped_by_req_pool", False
+        ) or getattr(sparse_backend, "extend_dummy_seqs_capped_by_req_pool", False)
 
     def init_forward_metadata(self, forward_batch: ForwardBatch):
         self.sparse.init_forward_metadata(forward_batch)

python/sglang/srt/layers/attention/tbo_backend.py

@@ -20,6 +20,9 @@ class TboAttnBackend(AttentionBackend):
         # reads through TboAttnBackend resolve to the underlying pool.
         self.token_to_kv_pool = primary.token_to_kv_pool
         self.req_to_token_pool = primary.req_to_token_pool
+        self.extend_dummy_seqs_capped_by_req_pool = getattr(
+            primary, "extend_dummy_seqs_capped_by_req_pool", False
+        )
 
     @classmethod
     def init_new(cls, creator: Callable[[], AttentionBackend]):

python/sglang/srt/layers/attention/triton_backend.py

@@ -118,6 +118,10 @@ class TritonAttnBackend(AttentionBackend):
     # buffers; it never reads seq_lens_cpu / seq_lens_sum.
     needs_cpu_seq_lens: bool = False
 
+    # kv_indptr/qo_indptr are preallocated at (req pool + 1); an extend batch
+    # can never carry more seqs than the pool.
+    extend_dummy_seqs_capped_by_req_pool: bool = True
+
     def __init__(
         self,
         model_runner: ModelRunner,

python/sglang/srt/layers/attention/wave_backend.py

@@ -38,6 +38,11 @@ class ForwardMetadata:
 
 
 class WaveAttnBackend(AttentionBackend):
+
+    # kv_indptr/qo_indptr are preallocated at (req pool + 1); an extend batch
+    # can never carry more seqs than the pool.
+    extend_dummy_seqs_capped_by_req_pool: bool = True
+
     def __init__(
         self,
         model_runner: ModelRunner,

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

@@ -42,6 +42,7 @@ from sglang.srt.model_executor.forward_batch_info import (
 )
 from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
 from sglang.srt.model_executor.runner.flashinfer_autotune import (
+    maybe_flashinfer_autotune_extend,
     run_flashinfer_autotune_forward,
     should_run_flashinfer_autotune,
 )
@@ -82,6 +83,7 @@ def _allocate_decode_buffers(
     hc_hidden_size: Optional[int] = None,
     pp_proxy_topk_size: Optional[int] = None,
     pp_proxy_residual_num_blocks: Optional[int] = None,
+    allocate_logits_buffer: bool = True,
 ) -> SimpleNamespace:
     """Allocate the FB-shared decode buffers."""
     with torch.device(device):
@@ -97,9 +99,15 @@ def _allocate_decode_buffers(
             (max_bs * seq_len_fill_value + max_num_token) * num_tokens_per_req,
             dtype=torch.bool,
         )
-        next_token_logits_buffer = torch.zeros(
+        # (max_num_token, vocab) fp32 is large (>10GB at 16k tokens); callers
-            (max_num_token, vocab_size),
+        # whose dummy runs never touch logits (skip_logits autotune) opt out.
-            dtype=torch.float,
+        next_token_logits_buffer = (
+            torch.zeros(
+                (max_num_token, vocab_size),
+                dtype=torch.float,
+            )
+            if allocate_logits_buffer
+            else None
         )
         mamba_track_indices = (
             torch.zeros((max_bs,), dtype=torch.int64) if enable_mamba_track else None
@@ -113,7 +121,7 @@ def _allocate_decode_buffers(
             is_mhc = hc_hidden_size is not None
             hs = hc_hidden_size if is_mhc else hidden_size
             pp_proxy_tensors = {
-                "hidden_states": torch.zeros((max_bs, hs), dtype=dtype),
+                "hidden_states": torch.zeros((max_num_token, hs), dtype=dtype),
             }
             if not is_mhc:
                 # Only Kimi K3 supplies num_blocks: its PP bank is token-major
@@ -237,6 +245,7 @@ class BaseRunner(ABC):
                 buffers is not None
             ), "_autotune_buffers() must return a reusable buffer set for autotune"
             self._flashinfer_autotune(buffers=buffers, batch_size=batch_size)
+            maybe_flashinfer_autotune_extend(self, decode_num_tokens=batch_size)
 
         if (
             envs.SGLANG_PP_PARALLEL_DEEPGEMM_WARMUP.get()
@@ -305,7 +314,13 @@ class BaseRunner(ABC):
 
         run_flashinfer_autotune_forward(self.model_runner, forward_fn, skip_logits=True)
 
-    def _alloc_dummy_decode_buffers(self, max_bs: int, *, num_tokens_per_req: int = 1):
+    def _alloc_dummy_decode_buffers(
+        self,
+        max_bs: int,
+        *,
+        num_tokens_per_req: int = 1,
+        allocate_logits_buffer: bool = True,
+    ):
         """Allocate one static decode-buffer set for a dummy forward, sized to
         (max_bs, max_bs * num_tokens_per_req).
 
@@ -345,6 +360,7 @@ class BaseRunner(ABC):
             hc_hidden_size=getattr(mr.model_config, "hc_hidden_size", None),
             pp_proxy_topk_size=mr.get_pp_proxy_topk_size(),
             pp_proxy_residual_num_blocks=mr.get_pp_proxy_residual_num_blocks(),
+            allocate_logits_buffer=allocate_logits_buffer,
         )
 
     def _dummy_run(
@@ -354,6 +370,7 @@ class BaseRunner(ABC):
         forward_mode_override: Optional[ForwardMode] = None,
         *,
         buffers,
+        extend_num_tokens_per_req: Optional[int] = None,
     ):
         """Run a dummy forward pass for warmup/profiling.
 
@@ -392,6 +409,12 @@ class BaseRunner(ABC):
                 ), "This should not happen"
             capture_forward_mode = ForwardMode.TARGET_VERIFY
             num_tokens_per_req = mr.decode_num_tokens_per_req()
+        if extend_num_tokens_per_req is not None:
+            assert (
+                capture_forward_mode == ForwardMode.EXTEND
+                and not mr.spec_algorithm.is_speculative()
+            ), "extend_num_tokens_per_req requires a non-speculative EXTEND dummy"
+            num_tokens_per_req = extend_num_tokens_per_req
 
         num_tokens = batch_size * num_tokens_per_req
 
@@ -455,12 +478,17 @@ class BaseRunner(ABC):
 
         # For extend mode
         if capture_forward_mode == ForwardMode.EXTEND:
-            seq_len_fill_value = mr.attn_backend.get_cuda_graph_seq_len_fill_value()
+            if extend_num_tokens_per_req is None:
+                per_req_extend_len = mr.attn_backend.get_cuda_graph_seq_len_fill_value()
+            else:
+                per_req_extend_len = extend_num_tokens_per_req
+                seq_lens.fill_(per_req_extend_len)
+                seq_lens_cpu.fill_(per_req_extend_len)
             extend_prefix_lens_cpu = [0] * batch_size
-            extend_seq_lens_cpu = [seq_len_fill_value] * batch_size
+            extend_seq_lens_cpu = [per_req_extend_len] * batch_size
             extend_num_tokens = num_tokens
             extend_seq_lens = torch.full(
-                (batch_size,), seq_len_fill_value, dtype=torch.int32, device=mr.device
+                (batch_size,), per_req_extend_len, dtype=torch.int32, device=mr.device
             )
             extend_prefix_lens = torch.zeros(
                 (batch_size,), dtype=torch.int32, device=mr.device

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

@@ -15,6 +15,7 @@ from __future__ import annotations
 
 import contextlib
 import datetime
+import functools
 import hashlib
 import logging
 from pathlib import Path
@@ -23,6 +24,8 @@ from typing import TYPE_CHECKING, Callable, Optional
 import torch
 
 from sglang.srt.environ import envs
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.srt.utils import empty_context, log_info_on_rank0
 
 if TYPE_CHECKING:
     from sglang.srt.model_executor.model_runner import ModelRunner
@@ -238,3 +241,85 @@ def maybe_flashinfer_autotune_speculative_draft(
 
     run_flashinfer_autotune_forward(mr, run_and_reset, skip_logits=skip_logits)
     tuned_phases.add(phase_key)
+
+
+def maybe_flashinfer_autotune_extend(
+    runner: BaseRunner, *, decode_num_tokens: int
+) -> None:
+    """Also autotune one EXTEND-shaped dummy forward.
+
+    The decode-shaped autotune only covers token counts up to the decode
+    batch size, so larger prefill/extend batches fall outside the tuned
+    buckets and run flashinfer's default heuristic — which can be far
+    slower than the tuned tactic (e.g. trtllm-gen fp4 MoE is ~30% slower
+    untuned at >=8k tokens on sm100). One extra forward at the largest
+    per-rank extend token count tunes all buckets up to it.
+    """
+    if not envs.SGLANG_FLASHINFER_AUTOTUNE_EXTEND.get():
+        return
+    mr = runner.model_runner
+    # max_prefill_tokens is a per-scheduler (per dp-rank) budget, and warmup
+    # runs on all dp ranks at once, so the gathered dummy already reaches the
+    # worst-case serving gather. Do not divide by dp_size.
+    num_tokens = mr.server_args.max_prefill_tokens
+    if num_tokens <= (decode_num_tokens or 0):
+        return  # decode-shaped autotune already covered these buckets
+    if not mr.is_generation or mr.spec_algorithm.is_speculative():
+        # _dummy_run forces TARGET_VERIFY shapes for speculative runners;
+        # extend-bucket autotune for spec configs is a follow-up.
+        return
+    if mr.model_config.is_multimodal:
+        # The dummy runs mm_inputs=None, which multimodal prefill paths iterate.
+        return
+
+    if mr.attn_backend.extend_dummy_seqs_capped_by_req_pool:
+        pool_size = mr.req_to_token_pool.size
+        num_tokens_per_req = (num_tokens + pool_size - 1) // pool_size
+    else:
+        # Packed dummies tune measurably worse tactics for the same token
+        # bucket, so pack only where the backend would otherwise crash. None
+        # (not 1) keeps the backend's own seq_len_fill_value in _dummy_run.
+        num_tokens_per_req = None
+    per_req = num_tokens_per_req or 1
+    batch_size = (num_tokens + per_req - 1) // per_req
+    num_tokens = batch_size * per_req
+
+    buffers = runner._alloc_dummy_decode_buffers(
+        batch_size,
+        num_tokens_per_req=per_req,
+        allocate_logits_buffer=False,
+    )
+    canary_run_ctx = (
+        c.with_active_single_forward_manager(0)
+        if (c := mr.canary_manager) is not None
+        else empty_context()
+    )
+
+    forward_fn = functools.partial(
+        runner._dummy_run,
+        batch_size=batch_size,
+        buffers=buffers,
+        run_ctx=canary_run_ctx,
+        forward_mode_override=ForwardMode.EXTEND,
+        extend_num_tokens_per_req=num_tokens_per_req,
+    )
+
+    log_info_on_rank0(
+        logger,
+        f"FlashInfer autotune: extra EXTEND pass at {num_tokens} tokens "
+        f"({batch_size} seqs x {per_req} tokens).",
+    )
+    try:
+        run_flashinfer_autotune_forward(mr, forward_fn, skip_logits=True)
+    except torch.OutOfMemoryError:
+        # The pass is an optimization; without headroom for the extend-shaped
+        # forward, fall back to untuned extend buckets instead of failing.
+        log_info_on_rank0(
+            logger,
+            "FlashInfer extend autotune skipped: not enough free memory "
+            f"for a {num_tokens}-token dummy forward.",
+        )
+    finally:
+        # release dummy buffers before capture measures free memory
+        del forward_fn, buffers
+        torch.cuda.empty_cache()

test/registered/unit/layers/attention/test_verify_mask.py

@@ -110,6 +110,7 @@ class TestVerifyMaskGate(CustomTestCase):
 class _FakeAttnBackend:
     def __init__(self, verify_mask):
         self.needs_cpu_seq_lens = False
+        self.extend_dummy_seqs_capped_by_req_pool = False
         self.verify_mask = verify_mask
 
 

test/registered/unit/spec/test_dflash_overlap_hostsync.py

@@ -224,7 +224,10 @@ class TestHybridNeedsCpuSeqLens(CustomTestCase):
         from sglang.srt.layers.attention.hybrid_attn_backend import HybridAttnBackend
 
         def backend(flag):
-            return SimpleNamespace(needs_cpu_seq_lens=flag)
+            return SimpleNamespace(
+                needs_cpu_seq_lens=flag,
+                extend_dummy_seqs_capped_by_req_pool=False,
+            )
 
         runner = SimpleNamespace(
             server_args=SimpleNamespace(speculative_attention_mode=spec_mode),