config: the alias form of the runner-side instance read

The previous batch counted `self.server_args.X` and called the runner surface
done. It was not: the same read spelled through a local alias --
`server_args = model_runner.server_args` (or `sa = kvc.server_args`, `args = ...`)
followed by `server_args.leaf` -- is the same process-global read wearing a
local name, and the AST census counts **57 of them** across eleven files that
the grep never saw. Census per function, following the alias.

52 were leaves and go to their bag (`spec` 11, `schedule` 9, `memory` 7,
`exec.graph` 5, `exec.moe` 5, `parallel` 4, `disagg` 4, `model` 3,
`exec.mamba` 2, `exec.overlap` 2). Five were not leaves:
three derived members on the eager runner --
`max_speculative_num_draft_tokens` and `enable_mamba_extra_buffer` already had
accessors, and `max_prefill_buffer_tokens` gets one (all its inputs are `schedule`
leaves plus the configured PP size, so it derives from the bags and follows a
post-publish override; `TestDerivedPredicatesAgreeAcrossTiers` pins it against
the member over a 48-case matrix) -- plus `get_attention_backends()`, which the
same commit routes through `attention_backends()`, and a dict that merely shares
the name (`server_args_dict.items`). That dict is the one read left behind.

`build_attention_backends` also stops resolving the pair from the record: it
runs after publish, so it asks `attention_backends()` like every other consumer.
The draft override on the runner still wins first.

`dispatch_event_loop`'s three PP checks read the *configured* PP size, not the
live topology: the MLX runner stub never initializes torch.distributed, so the
live property asserts before the MLX event loop can start (a Codex catch). The
configured leaf answers the same value wherever the live groups exist.

`flashinfer_gdn_prefill_default`'s guard is the one read here that asks what the
*operator* named rather than what the config resolved to, and the bag leaf now
answers exactly that: the per-runner auto-default is stamped on the runner and
deliberately never recorded process-wide, so nothing writes that leaf after
launch and reading it back cannot mistake another runner's default for a flag.

Three test doubles injected a `SimpleNamespace`/`MagicMock` record for exactly
these reads and now publish instead (pool configurator, cache registry, GDN
prefill policy) -- the fixture publishes what the case configures and hands the
published instance to the whole-object contracts that still take one.

The functions this sweep partially converted stop mixing sources (review
catches): the flash-attention constructor's remaining seed reads
(`speculative_eagle_topk`, `speculative_algorithm`, both deterministic gates)
read their bags next to the leaves already converted;
`_should_disable_scheduler_metadata_precompute` reads the parallel config
leaves itself instead of taking the record (its alias binding was the last
use); and the autotune gates (`disable_flashinfer_autotune`, deterministic,
`flashinfer_autotune_skip_ops`) join the moe leaves the same function already
reads from the bags. The pool-configurator fixture drops a parameter nothing
published or read.
This commit is contained in:
Cheng Wan
2026-08-15 00:39:03 -07:00
committed by GitHub
parent 61908870f6
commit f2ab6e306b
18 changed files with 319 additions and 147 deletions
@@ -30,7 +30,14 @@ from sglang.srt.layers.utils.cp_utils import (
from sglang.srt.mem_cache.memory_pool import KVWriteLoc
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.runtime_context import get_schedule, get_spec
from sglang.srt.runtime_context import (
get_exec,
get_memory,
get_model,
get_parallel,
get_schedule,
get_spec,
)
from sglang.srt.speculative.ragged_verify import build_ragged_target_verify_geometry
from sglang.srt.speculative.spec_info import SpecInput, SpeculativeAlgorithm
from sglang.srt.speculative.spec_utils import resolve_num_tokens_per_req
@@ -49,8 +56,8 @@ from sglang.kernels.ops.attention.flash_attention import (
)
def _should_disable_scheduler_metadata_precompute(server_args) -> bool:
return bool(server_args.enable_prefill_cp or server_args.enable_dp_attention)
def _should_disable_scheduler_metadata_precompute() -> bool:
return bool(get_parallel().enable_prefill_cp or get_parallel().enable_dp_attention)
@dataclass
@@ -202,7 +209,7 @@ class FlashAttentionBackend(AttentionBackend):
and model_runner.token_to_kv_pool.swa_layer_nums > 0
)
self.topk = model_runner.server_args.speculative_eagle_topk or 0
self.topk = get_spec().speculative_eagle_topk or 0
self.speculative_num_steps = speculative_num_steps
self.speculative_num_draft_tokens = get_spec().speculative_num_draft_tokens
if (
@@ -214,7 +221,7 @@ class FlashAttentionBackend(AttentionBackend):
phase="target_verify",
server_args=model_runner.server_args,
spec_algorithm=SpeculativeAlgorithm.from_string(
model_runner.server_args.speculative_algorithm
get_spec().speculative_algorithm
),
is_draft_worker=True,
num_draft_tokens=int(self.speculative_num_draft_tokens),
@@ -283,7 +290,7 @@ class FlashAttentionBackend(AttentionBackend):
self._get_fa_runtime_policy = None
self._get_scheduler_metadata = None
if model_runner.server_args.enable_deterministic_inference:
if get_exec().deterministic.enable_deterministic_inference:
# Must precede the first kernel compile.
from sglang.kernels.ops.attention.flash_attn.cute.batch_invariance import (
set_batch_invariant,
@@ -311,7 +318,7 @@ class FlashAttentionBackend(AttentionBackend):
self.has_softcap = _softcapping is not None and _softcapping > 0.0
# num_splits == 0 delegates SplitKV sizing to the selected FA runtime.
deterministic = model_runner.server_args.enable_deterministic_inference
deterministic = get_exec().deterministic.enable_deterministic_inference
if self._get_fa_runtime_policy is None:
self.num_splits = 1 if deterministic else 0
self.decode_num_splits = self.num_splits
@@ -339,11 +346,10 @@ class FlashAttentionBackend(AttentionBackend):
# guard wraps both set_kv_buffer and set_mla_kv_buffer. Without this
# gate, MLA + is_embedding would skip the write but still read stale
# cache via get_key_buffer in the absorbed-MLA path.
server_args = model_runner.server_args
self.fa_skip_kv_cache = (
server_args.is_embedding
and server_args.chunked_prefill_size == -1
and server_args.disable_radix_cache
get_model().is_embedding
and get_schedule().chunked_prefill_size == -1
and get_memory().disable_radix_cache
and not self.use_mla
)
@@ -353,7 +359,7 @@ class FlashAttentionBackend(AttentionBackend):
# combine kernel (flash_fwd_combine_launch_template.h:52). Leaving
# scheduler_metadata unset uses the existing per-layer metadata path.
self._disable_scheduler_metadata_precompute = (
_should_disable_scheduler_metadata_precompute(server_args)
_should_disable_scheduler_metadata_precompute()
)
def _compute_scheduler_metadata(
@@ -18,6 +18,7 @@ from sglang.srt.layers.radix_linear_attention import RadixLinearAttention
from sglang.srt.mem_cache.memory_pool import MambaPool
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.runtime_context import get_exec, get_memory, get_schedule
from sglang.srt.utils import is_cpu, is_cuda, is_hip, is_npu, is_xpu
from sglang.srt.utils.common import rank0_log
@@ -64,23 +65,22 @@ elif is_cpu():
def flashinfer_gdn_prefill_default(model_runner: ModelRunner) -> Optional[str]:
"""FlashInfer for the narrow SM100 GDN prefill domain we validated, else None."""
args = model_runner.server_args
if (
args.linear_attn_prefill_backend is not None
or args.linear_attn_backend != "triton"
or args.enable_page_major_kv_layout
get_exec().mamba.linear_attn_prefill_backend is not None
or get_exec().mamba.linear_attn_backend != "triton"
or get_memory().enable_page_major_kv_layout
or not is_cuda()
or torch.cuda.get_device_capability()[0] != 10
):
return None
cuda_version = torch.version.cuda
chunk_size = args.chunked_prefill_size
chunk_size = get_schedule().chunked_prefill_size
config = hybrid_gdn_config(model_runner.model_config)
if (
cuda_version is None
or int(cuda_version.split(".", 1)[0]) < 13
or args.enable_dynamic_chunking
or get_schedule().enable_dynamic_chunking
or chunk_size is None
or not 1 <= chunk_size <= 8192
or getattr(config, "linear_key_head_dim", None) != 128
+5 -5
View File
@@ -29,6 +29,7 @@ from typing import TYPE_CHECKING, Any, Deque, Dict, List, Optional, Set, Tuple,
from sglang.srt.runtime_context import (
attention_backends,
configured_pp_size,
get_device,
get_disagg,
get_exec,
@@ -4900,13 +4901,12 @@ class Scheduler(
def dispatch_event_loop(scheduler: Scheduler):
# Dispatch to the appropriate event loop based on the disaggregation mode
server_args = scheduler.server_args
# The live PP property asserts before torch.distributed init (MLX stub).
disaggregation_mode: DisaggregationMode = scheduler.disaggregation_mode
if disaggregation_mode == DisaggregationMode.NULL:
if scheduler.enable_pdmux:
scheduler.event_loop_pdmux()
elif server_args.pp_size > 1:
elif configured_pp_size() > 1:
scheduler.event_loop_pp()
elif scheduler.enable_overlap_mlx:
scheduler.event_loop_overlap_mlx()
@@ -4915,14 +4915,14 @@ def dispatch_event_loop(scheduler: Scheduler):
else:
scheduler.event_loop_normal()
elif disaggregation_mode == DisaggregationMode.PREFILL:
if server_args.pp_size > 1:
if configured_pp_size() > 1:
scheduler.event_loop_pp_disagg_prefill()
elif scheduler.enable_overlap:
scheduler.event_loop_overlap_disagg_prefill()
else:
scheduler.event_loop_normal_disagg_prefill()
elif disaggregation_mode == DisaggregationMode.DECODE:
if server_args.pp_size > 1:
if configured_pp_size() > 1:
scheduler.event_loop_pp_disagg_decode()
elif scheduler.enable_overlap:
scheduler.event_loop_overlap_disagg_decode()
@@ -338,15 +338,15 @@ class SchedulerMetricsReporter:
"num_draft_tokens": 0,
}
# Fallback to server_args if draft_worker does not have the attributes.
server_args = self.scheduler.server_args
# Fallback to the published `spec` bag if draft_worker does not have
# the attributes.
num_steps = getattr(
draft_worker, "speculative_num_steps", server_args.speculative_num_steps
draft_worker, "speculative_num_steps", get_spec().speculative_num_steps
)
num_draft_tokens = getattr(
draft_worker,
"speculative_num_draft_tokens",
server_args.speculative_num_draft_tokens,
get_spec().speculative_num_draft_tokens,
)
return {
+5 -4
View File
@@ -17,6 +17,7 @@ from typing import TYPE_CHECKING, Any, Callable, Optional
from sglang.srt.environ import envs
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.runtime_context import get_memory
from sglang.srt.utils.tensor_bridge import use_mlx
if TYPE_CHECKING:
@@ -128,7 +129,7 @@ def default_radix_cache_factory(ctx: TreeCacheBuildContext) -> BasePrefixCache:
)
return cache
if server_args.enable_lmcache:
if get_memory().enable_lmcache:
from sglang.srt.mem_cache.storage.lmcache.lmc_radix_cache import (
LMCRadixCache,
)
@@ -141,7 +142,7 @@ def default_radix_cache_factory(ctx: TreeCacheBuildContext) -> BasePrefixCache:
tp_group=ctx.tp_group,
)
if server_args.enable_flexkv:
if get_memory().enable_flexkv:
# Importing the package side-effect registers the explicit
# ``--radix-cache-backend=flexkv`` factory; we then call the
# factory directly so --enable-flexkv stands on its own.
@@ -151,8 +152,8 @@ def default_radix_cache_factory(ctx: TreeCacheBuildContext) -> BasePrefixCache:
# Honor a CLI --flexkv-config-file by forwarding it via the env
# var that FlexKV's config loader actually reads.
if server_args.flexkv_config_file and not os.environ.get("FLEXKV_CONFIG_PATH"):
os.environ["FLEXKV_CONFIG_PATH"] = server_args.flexkv_config_file
if get_memory().flexkv_config_file and not os.environ.get("FLEXKV_CONFIG_PATH"):
os.environ["FLEXKV_CONFIG_PATH"] = get_memory().flexkv_config_file
return _flexkv_factory(ctx)
from sglang.srt.mem_cache.radix_cache import RadixCache
@@ -39,7 +39,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_utils.capture_mode import model_capture_mode
from sglang.srt.runtime_context import get_flags, get_parallel, get_spec
from sglang.srt.runtime_context import get_exec, get_flags, get_parallel, get_spec
from sglang.srt.utils import (
empty_context,
log_info_on_rank0,
@@ -127,14 +127,13 @@ def set_torch_compile_config():
def get_batch_sizes_to_capture(model_runner: ModelRunner):
# torch compile speeds up decoding by reducing python overhead on CPU
server_args = model_runner.server_args
# Reuse cuda_graph_config[decode].bs here.
# Users can customize the batch sizes supported by cpu_graph, such as:
# --cuda-graph-bs-decode 1 2 4 8 16
capture_bs = server_args.cuda_graph_config.decode.bs
capture_bs = get_exec().graph.cuda_graph_config.decode.bs
assert (
max(capture_bs) <= server_args.torch_compile_max_bs
), f"{capture_bs=}, {server_args.torch_compile_max_bs=}"
max(capture_bs) <= get_exec().graph.torch_compile_max_bs
), f"{capture_bs=}, {get_exec().graph.torch_compile_max_bs=}"
capture_bs = [bs for bs in capture_bs if bs <= model_runner.req_to_token_pool.size]
capture_bs = list(sorted(set(capture_bs)))
assert len(capture_bs) > 0 and capture_bs[0] > 0, f"{capture_bs=}"
@@ -18,6 +18,8 @@ if TYPE_CHECKING:
from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.runtime_context import attention_backends, get_disagg, get_exec
logger = logging.getLogger(__name__)
@@ -66,7 +68,6 @@ def configure_aux_hidden_state_capture(
def build_attention_backends(*, model_runner: ModelRunner) -> AttentionBackends:
"""Init attention kernel backend."""
server_args = model_runner.server_args
# TODO: Refactor device-specific init branches into platform interface (separate PR).
if model_runner.device in ("cuda", "musa"):
@@ -81,7 +82,7 @@ def build_attention_backends(*, model_runner: ModelRunner) -> AttentionBackends:
),
)
if server_args.enable_pdmux:
if get_disagg().enable_pdmux:
attn_backend = _build_resolved_backend(
model_runner=model_runner, resolved=resolved, init_new_workspace=True
)
@@ -91,10 +92,12 @@ def build_attention_backends(*, model_runner: ModelRunner) -> AttentionBackends:
resolved=resolved,
init_new_workspace=False,
)
for _ in range(server_args.sm_group_num)
for _ in range(get_disagg().sm_group_num)
]
decode_attn_backend = decode_attn_backend_group[0]
elif server_args.enable_two_batch_overlap and not model_runner.is_draft_worker:
elif (
get_exec().overlap.enable_two_batch_overlap and not model_runner.is_draft_worker
):
attn_backend = TboAttnBackend.init_new(
lambda: _build_resolved_backend(
model_runner=model_runner,
@@ -161,7 +164,6 @@ def resolve_attention_backend_strs(
target and draft coexist in one process, so it cannot come from the
process-wide config.
"""
server_args = model_runner.server_args
is_draft_worker = model_runner.is_draft_worker
draft_attn_backend = model_runner.draft_attention_backend
if is_draft_worker and draft_attn_backend:
@@ -172,7 +174,7 @@ def resolve_attention_backend_strs(
decode=draft_attn_backend,
is_draft_override=True,
)
prefill, decode = server_args.get_attention_backends()
prefill, decode = attention_backends()
return ResolvedAttentionBackendStr(prefill=prefill, decode=decode)
@@ -36,7 +36,14 @@ from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
get_compress_state_write_pad,
)
from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool
from sglang.srt.runtime_context import get_parallel
from sglang.srt.runtime_context import (
get_disagg,
get_memory,
get_parallel,
get_schedule,
get_spec,
max_speculative_num_draft_tokens,
)
from sglang.srt.utils.common import (
ceil_align,
ceil_div,
@@ -533,10 +540,9 @@ class SWAChunkCapPoolConfigurator(HybridSWAPoolConfigurator):
super().__init__(kvc)
assert self._full_layers_num > 0
sa = kvc.server_args
page_size = kvc.page_size
window = kvc.sliding_window_size
draft_tokens = sa.speculative_num_draft_tokens or 1
draft_tokens = get_spec().speculative_num_draft_tokens or 1
eviction_interval = max(1, envs.SGLANG_SWA_EVICTION_INTERVAL.get())
"""
@@ -544,40 +550,47 @@ class SWAChunkCapPoolConfigurator(HybridSWAPoolConfigurator):
Padding to make sure eviction point is page-aligned.
"""
trailing_tokens = window + eviction_interval * draft_tokens + page_size
if sa.speculative_algorithm is None:
if get_spec().speculative_algorithm is None:
decode_alloc = page_size
elif sa.disable_overlap_schedule:
elif get_schedule().disable_overlap_schedule:
# spec-v1: new_tokens_required_next_decode per request.
decode_alloc = spec_decode_alloc_len_per_request(sa)
decode_alloc = spec_decode_alloc_len_per_request(
page_size=page_size,
speculative_num_steps=get_spec().speculative_num_steps,
speculative_eagle_topk=get_spec().speculative_eagle_topk,
speculative_num_draft_tokens=get_spec().speculative_num_draft_tokens,
)
else:
# spec-v2: the overlap allocator keeps 2 * alloc_len outstanding
# (eagle_utils.eagle_prepare_for_decode: kv_committed_len + 2 * alloc_len).
decode_alloc = 2 * get_alloc_len_per_decode(sa)
decode_alloc = 2 * get_alloc_len_per_decode(
kvc.server_args,
max_draft_tokens=max_speculative_num_draft_tokens(),
)
per_request = trailing_tokens + decode_alloc
num_reqs = sa.max_running_requests // kvc.ps.attn_dp_size
if sa.disaggregation_mode == "decode":
num_reqs = get_schedule().max_running_requests // kvc.ps.attn_dp_size
if get_disagg().disaggregation_mode == "decode":
self._swa_cap = (
per_request * num_reqs
+ (window + page_size) * sa.disaggregation_decode_extra_slots
+ (window + page_size) * get_disagg().disaggregation_decode_extra_slots
)
else:
chunks_in_flight = 1 if sa.disable_overlap_schedule else 2
chunks_in_flight = 1 if get_schedule().disable_overlap_schedule else 2
self._swa_cap = (
per_request * num_reqs
+ chunks_in_flight * sa.chunked_prefill_size
+ chunks_in_flight * get_schedule().chunked_prefill_size
+ page_size
)
@staticmethod
def is_applicable(kvc: KVCacheConfigurator) -> bool:
"""True when SWAChunkCache can be sized from explicit max requests."""
sa = kvc.server_args
if sa.max_running_requests is None:
if get_schedule().max_running_requests is None:
return False
if not sa.disable_radix_cache:
if not get_memory().disable_radix_cache:
return False
if sa.chunked_prefill_size is None:
if get_schedule().chunked_prefill_size is None:
return False
if kvc.sliding_window_size is None:
return False
@@ -23,7 +23,10 @@ from contextlib import contextmanager
from typing import TYPE_CHECKING, Any, List, Sequence, Tuple
from sglang.srt.model_executor.runner.base_runner import BaseRunner
from sglang.srt.runtime_context import get_flags
from sglang.srt.runtime_context import (
get_exec,
get_flags,
)
from sglang.srt.utils import (
get_cuda_graph_batch_size_alignment,
get_cuda_graph_max_batch_size,
@@ -68,14 +71,14 @@ def get_batch_sizes_to_capture(
"""
server_args = model_runner.server_args
capture_bs = list(server_args.cuda_graph_config.decode.bs)
capture_bs = list(get_exec().graph.cuda_graph_config.decode.bs)
num_max_requests = model_runner.req_to_token_pool.size
mul_base = get_cuda_graph_batch_size_alignment(server_args)
# TBO splits each request's rows across two micro-batches, so the
# alignment constraint applies per request rather than per token row.
alignment_width = captured_req_width
if server_args.enable_two_batch_overlap:
if get_exec().overlap.enable_two_batch_overlap:
alignment_width = 1
# pad `num_max_requests` to avoid being filtered out
@@ -92,7 +95,7 @@ def get_batch_sizes_to_capture(
assert len(capture_bs) > 0 and capture_bs[0] > 0, f"{capture_bs=}"
compile_bs = (
[bs for bs in capture_bs if bs <= server_args.torch_compile_max_bs]
[bs for bs in capture_bs if bs <= get_exec().graph.torch_compile_max_bs]
if get_flags().capture.enable_torch_compile
else []
)
@@ -49,6 +49,13 @@ 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.runtime_context import (
get_parallel,
get_spec,
mamba_extra_buffer_enabled,
max_prefill_buffer_tokens,
max_speculative_num_draft_tokens,
)
from sglang.srt.utils import is_hip
from sglang.srt.utils.common import (
ceil_align,
@@ -76,14 +83,14 @@ class EagerRunner(BaseRunner):
num_tokens_per_req = 1
if mr.spec_algorithm.is_speculative():
# speculative_adaptive can grow draft tokens at runtime; size to the max.
num_draft_tokens = sa.max_speculative_num_draft_tokens or 1
num_draft_tokens = max_speculative_num_draft_tokens() or 1
if mr.is_draft_worker:
num_tokens_per_req = max(
sa.speculative_eagle_topk or 1,
get_spec().speculative_eagle_topk or 1,
num_draft_tokens,
(
2 * (sa.speculative_num_steps or 0)
if sa.enable_multi_layer_eagle
2 * (get_spec().speculative_num_steps or 0)
if get_spec().enable_multi_layer_eagle
else 0
),
)
@@ -100,14 +107,14 @@ class EagerRunner(BaseRunner):
if (
mr.is_draft_worker
and mr.spec_algorithm.is_frozen_kv_mtp()
and sa.speculative_eagle_topk > 1
and get_spec().speculative_eagle_topk > 1
):
# Frozen-KV MTP expands the draft batch by topk on the bs axis
# (expand_for_topk_draft) before the eager fallback.
max_bs *= sa.speculative_eagle_topk
max_bs *= get_spec().speculative_eagle_topk
# Mirror prepare_mlp_sync_batch padding so the registry holds what load_batch copies.
max_bs = get_eager_max_batch_size(sa, max_bs)
prefill_ceiling = max(mr.max_total_num_tokens, sa.max_prefill_buffer_tokens())
prefill_ceiling = max(mr.max_total_num_tokens, max_prefill_buffer_tokens())
max_num_token = max(prefill_ceiling, max_bs * num_tokens_per_req)
if require_mlp_sync(sa):
from sglang.srt.layers.cp.padding import get_cp_padding_align_size
@@ -123,7 +130,7 @@ class EagerRunner(BaseRunner):
max_num_token=max_num_token,
cache_loc_dtype=torch.int64,
enable_mamba_track=(
sa.enable_mamba_extra_buffer() and mr.spec_algorithm.is_none()
mamba_extra_buffer_enabled() and mr.spec_algorithm.is_none()
),
is_encoder_decoder=is_encoder_decoder,
encoder_len_fill_value=(
@@ -134,7 +141,7 @@ class EagerRunner(BaseRunner):
encoder_lens_dtype=(
torch.int64 if torch.device(mr.device).type == "cpu" else torch.int32
),
dp_size=sa.dp_size,
dp_size=get_parallel().dp_size,
)
# Eager has no capture step, so warm up here (run-once via mr._kernel_warmed_up).
self.warmup()
@@ -25,6 +25,11 @@ import torch
from sglang.srt.environ import envs
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.runtime_context import (
get_exec,
get_model,
get_spec,
)
from sglang.srt.utils import empty_context, log_info_on_rank0
if TYPE_CHECKING:
@@ -37,7 +42,7 @@ FLASHINFER_AUTOTUNE_WORKAROUND_SKIPS = frozenset()
def get_flashinfer_autotune_skip_ops(model_runner: ModelRunner) -> set[str]:
skip_ops = set(model_runner.server_args.flashinfer_autotune_skip_ops or ())
skip_ops = set(get_exec().kernel.flashinfer_autotune_skip_ops or ())
skip_ops.update(FLASHINFER_AUTOTUNE_WORKAROUND_SKIPS)
return skip_ops
@@ -49,28 +54,29 @@ def should_run_flashinfer_autotune(
mr = model_runner
if mr.device != "cuda":
return False
if mr.server_args.disable_flashinfer_autotune:
if get_exec().kernel.disable_flashinfer_autotune:
return False
if mr.server_args.enable_deterministic_inference:
if get_exec().deterministic.enable_deterministic_inference:
# Tuned configs are per problem shape, so the reduction order would follow
# the batch shape.
return False
server_args = mr.server_args
if for_speculative_draft:
backend_str = (
server_args.speculative_moe_runner_backend or server_args.moe_runner_backend
get_spec().speculative_moe_runner_backend
or get_exec().moe.moe_runner_backend
)
a2a_backend_str = (
server_args.speculative_moe_a2a_backend or server_args.moe_a2a_backend
get_spec().speculative_moe_a2a_backend or get_exec().moe.moe_a2a_backend
)
else:
backend_str = server_args.moe_runner_backend
a2a_backend_str = server_args.moe_a2a_backend
backend_str = get_exec().moe.moe_runner_backend
a2a_backend_str = get_exec().moe.moe_a2a_backend
# Autotune can run before the MoE backend globals are initialized, so read
# the target or draft backend from server_args. CuteDSL v1 bypasses
# MoeRunner, and its dummy dispatch can exceed DeepEP low-latency's token limit.
# the configured backends -- the draft leaves (`get_spec()`) or the target
# leaves (`get_exec().moe`) above. CuteDSL v1 bypasses MoeRunner, and its
# dummy dispatch can exceed DeepEP low-latency's token limit.
if backend_str == "flashinfer_cutedsl" and a2a_backend_str == "deepep":
return False
@@ -130,12 +136,11 @@ def flashinfer_autotune_cache_path(model_runner: ModelRunner) -> Path:
arch = f"sm{major}{minor}"
flashinfer_version = getattr(flashinfer, "__version__", "unknown")
server_args = mr.server_args
model_key_parts = [
str(server_args.model_path),
str(get_model().model_path),
str(mr.dtype),
str(server_args.quantization),
str(server_args.moe_runner_backend),
str(get_model().quantization),
str(get_exec().moe.moe_runner_backend),
str(mr.ps.tp_size),
str(mr.ps.pp_size),
str(mr.ps.attn_dp_size),
+29
View File
@@ -1434,6 +1434,35 @@ def remote_instance_transfer_engine_enabled(load_format: str | None = None) -> b
return remote_instance_transfer_engine_of(get_model(), load_format)
def max_prefill_buffer_tokens() -> int:
"""The prefill-buffer ceiling: ``chunked_prefill_size``, except PP dynamic
chunking can grow chunks toward ``max_prefill_tokens`` and probe at 1.25x.
Every input is a published leaf (``schedule`` plus the configured PP size),
so this derives from the bags and follows a post-publish override;
``ServerArgs.max_prefill_buffer_tokens`` is the pre-publish equivalent and
``TestDerivedPredicatesAgreeAcrossTiers`` pins the two equal.
"""
import math
schedule = get_schedule()
chunked = (
schedule.chunked_prefill_size
if schedule.chunked_prefill_size and schedule.chunked_prefill_size > 0
else 0
)
tokens = chunked
if (
schedule.enable_dynamic_chunking
and _configured_parallel("pp_size") > 1
and chunked
):
tokens = max(
tokens, schedule.max_prefill_tokens or 0, math.ceil(chunked * 1.25)
)
return tokens
def pre_capture_activation_reserve_mb(gpu_mem: float | None) -> float:
"""The activation working-set reserve held back before cuda-graph capture.
+12 -6
View File
@@ -3903,14 +3903,20 @@ def ceil_align(x: int, y: int) -> int:
return ceil_div(x, y) * y
def spec_decode_alloc_len_per_request(server_args) -> int:
def spec_decode_alloc_len_per_request(
*,
page_size,
speculative_num_steps,
speculative_eagle_topk,
speculative_num_draft_tokens,
) -> int:
"""Per-request KV tokens a (spec-v1) decode step allocates: the draft-decode
topk*num_steps peak vs. the verify num_draft_tokens, page-aligned.
topk*num_steps peak vs. the verify num_draft_tokens, page-aligned. A pure
function of the resolved values its one caller reads off the bags.
"""
page_size = server_args.page_size
len_per_topk = server_args.speculative_num_steps or 1
spec_topk = server_args.speculative_eagle_topk or 1
spec_tokens = server_args.speculative_num_draft_tokens or 1
len_per_topk = speculative_num_steps or 1
spec_topk = speculative_eagle_topk or 1
spec_tokens = speculative_num_draft_tokens or 1
if page_size > 1 and spec_topk > 1:
# last partial page and ceil alignment