diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 5edd4aca8..983da2fbd 100644 --- a/python/sglang/srt/environ.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index 3eaa2e1be..dfcba404f 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/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, diff --git a/python/sglang/srt/layers/attention/base_attn_backend.py b/python/sglang/srt/layers/attention/base_attn_backend.py index bf8bad2e9..1987bf5a6 100644 --- a/python/sglang/srt/layers/attention/base_attn_backend.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index 37fcbe404..ecde3b946 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/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. diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index b524a9afc..63cd3d970 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/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, diff --git a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py index 96518ded3..c87c12c0e 100644 --- a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/hybrid_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_attn_backend.py index 984a0edd3..53b54d949 100644 --- a/python/sglang/srt/layers/attention/hybrid_attn_backend.py +++ b/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: diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index f62abfe89..a3d8f1fae 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/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): diff --git a/python/sglang/srt/layers/attention/minimax_sparse_backend.py b/python/sglang/srt/layers/attention/minimax_sparse_backend.py index 702721ff2..c09e7da12 100644 --- a/python/sglang/srt/layers/attention/minimax_sparse_backend.py +++ b/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) diff --git a/python/sglang/srt/layers/attention/tbo_backend.py b/python/sglang/srt/layers/attention/tbo_backend.py index 27867541d..649d40a63 100644 --- a/python/sglang/srt/layers/attention/tbo_backend.py +++ b/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]): diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index 2daae2f82..37c2b5c04 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/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, diff --git a/python/sglang/srt/layers/attention/wave_backend.py b/python/sglang/srt/layers/attention/wave_backend.py index 9764b7e41..155e215a3 100644 --- a/python/sglang/srt/layers/attention/wave_backend.py +++ b/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, diff --git a/python/sglang/srt/model_executor/runner/base_runner.py b/python/sglang/srt/model_executor/runner/base_runner.py index 288b1e274..f94061db1 100644 --- a/python/sglang/srt/model_executor/runner/base_runner.py +++ b/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_size), - dtype=torch.float, + # (max_num_token, vocab) fp32 is large (>10GB at 16k tokens); callers + # whose dummy runs never touch logits (skip_logits autotune) opt out. + 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 diff --git a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py index b32751716..6790070a6 100644 --- a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py +++ b/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() diff --git a/test/registered/unit/layers/attention/test_verify_mask.py b/test/registered/unit/layers/attention/test_verify_mask.py index 3386c697c..e5c743452 100644 --- a/test/registered/unit/layers/attention/test_verify_mask.py +++ b/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 diff --git a/test/registered/unit/spec/test_dflash_overlap_hostsync.py b/test/registered/unit/spec/test_dflash_overlap_hostsync.py index d3a574211..a78e2e59e 100644 --- a/test/registered/unit/spec/test_dflash_overlap_hostsync.py +++ b/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),