From 93015a9e6b11e584f4285f71d7779c4855f7d1a8 Mon Sep 17 00:00:00 2001 From: YAMY <74099316+YAMY1234@users.noreply.github.com> Date: Tue, 23 Jun 2026 13:31:47 -0700 Subject: [PATCH] fix(runner): autotune flashinfer MoE on a decode-shaped buffer (#29069) --- .../srt/model_executor/runner/eager_runner.py | 99 ++++--------------- 1 file changed, 20 insertions(+), 79 deletions(-) diff --git a/python/sglang/srt/model_executor/runner/eager_runner.py b/python/sglang/srt/model_executor/runner/eager_runner.py index 691fd28e5..0034cceca 100644 --- a/python/sglang/srt/model_executor/runner/eager_runner.py +++ b/python/sglang/srt/model_executor/runner/eager_runner.py @@ -18,8 +18,7 @@ from __future__ import annotations import contextlib import logging from dataclasses import replace -from types import SimpleNamespace -from typing import TYPE_CHECKING, Any, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Tuple, Union import torch @@ -42,7 +41,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, require_mlp_tp_gather +from sglang.srt.utils import is_hip from sglang.srt.utils.common import ceil_align, require_mlp_sync logger = logging.getLogger(__name__) @@ -133,86 +132,28 @@ class EagerRunner(BaseRunner): 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).""" + """Decode-shaped dummy buffers (bs * num_tokens_per_bs) for the warmup + flashinfer-autotune forward. + + flashinfer's MoE autotuner times candidate tactics against the buffer it + is given, so it must match the live decode shape for the cached tactic to + be optimal at decode. The eager input registry spans the prefill token + ceiling; the dummy run only needs the decode-sized slice. + """ 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 + num_tokens_per_bs = 1 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, + num_tokens_per_bs = ( + mr.spec_algorithm.get_num_tokens_per_bs_for_target_verify( + mr.server_args.speculative_num_draft_tokens, mr.is_draft_worker + ) ) - - # 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 - ) - pp_proxy_topk_size = mr.get_pp_proxy_topk_size() - if pp_proxy_topk_size is not None: - pp_proxy_tensors["topk_indices"] = torch.zeros( - (rows, pp_proxy_topk_size), dtype=torch.int32, 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 ( + self._alloc_dummy_decode_buffers( + self._eager_max_bs, num_tokens_per_bs=num_tokens_per_bs + ), + self._eager_max_bs, ) - return adapter, max_bs def can_run_graph(self, forward_batch: ForwardBatch) -> bool: # Eager never runs a cuda graph; callers dispatch on isinstance(...,