minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[SPEC][5/N] feat: batchsize-aware support for adaptive speculative_num_steps (#24055)

Browse Source 
Co-authored-by: 坤钧 <maoyuhan.myh@antgroup.co>
Co-authored-by: alphabetc1 <alphabetc1@users.noreply.github.com>
Co-authored-by: EanWang211123 <wangyiheng@sangfor.com.cn>
Co-authored-by: shuwenn <47200617+alphabetc1@users.noreply.github.com>
Co-authored-by: shuwenn <2508695655@qq.com>
This commit is contained in:
maoyuhan
2026-06-05 15:43:02 -07:00
committed by GitHub
co-authored by 坤钧 alphabetc1 EanWang211123 shuwenn shuwenn
parent c9f582a272
commit 6b180959a8
11 changed files with 657 additions and 152 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs_new/docs/advanced_features/adaptive_speculative_decoding.mdx
+100 -28
View File
@@ -9,7 +9,7 @@ It is designed for workloads whose accept length changes over time, where one st
## Current support
- Only `--speculative-algorithm EAGLE`
- Only `--speculative-algorithm EAGLE` or `EAGLE3`
- Only `--speculative-eagle-topk 1`
- If either condition is not met, SGLang falls back to static speculative settings
@@ -19,8 +19,9 @@ It is designed for workloads whose accept length changes over time, where one st
- If `num_steps` is too small, the draft model could have produced more accepted tokens, but the round stops too early.
- If `num_steps` is too large, the draft model produces many candidate tokens that the target model rejects, so extra draft work is wasted.
- At **high batch sizes**, the cost of each wasted draft step is multiplied across all sequences in the batch, so the optimal step count is often lower than at low batch sizes.
Real traffic often moves between high-acceptance and low-acceptance phases, so one fixed step count is usually a compromise. Adaptive mode tries to follow the workload instead of hard-coding a single global `num_steps`.
Real traffic often moves between high-acceptance and low-acceptance phases, and batch sizes vary continuously. Adaptive mode follows both signals at runtime instead of hard-coding a single global `num_steps`.
## Design overview
@@ -28,9 +29,13 @@ The adaptive mechanism has three pieces:
- `AdaptiveSpeculativeParams`: the EMA-based policy
- `SpecRuntimeState`: the per-tier runtime state bundle
- `AdaptiveController`: the coordinator that chooses a tier and activates the matching runtime state
- `AdaptiveController`: the coordinator that queries the policy for the current batch size and activates the matching runtime state
At startup, SGLang pre-builds one runtime state per candidate tier. By default, the candidate tiers are `candidate_steps = [1, 3, 7]`.
### Per-batch-size independent tracking
The controller maintains **independent EMA trackers for each batch size range**, so observations at small BS don't pollute the large BS signal. Each BS range can have its own candidate steps, hysteresis thresholds, and ceiling coefficient.
BS ranges are defined as lower bounds in the config file (e.g., keys `"1"` and `"8"` mean BS 1–7 uses one slot, BS 8+ uses another). `SpecRuntimeState` objects are shared across BS ranges with the same step count — each state owns CUDA graphs captured for the reachable padded batch sizes of that step.
```mermaid
---
@@ -61,32 +66,37 @@ This matters because `CudaGraphRunner` is shape-dependent. Each candidate tier o
## Runtime flow
The adaptive update happens after verify and affects the next round, not the current one:
The adaptive update happens in two places:
1. **Pre-draft**: query the optimal step for the current batch size and activate if different
2. **Post-verify**: update the matching BS slot's EMA with observed accept lengths
```mermaid
---
title: "EAGLEWorker.forward_batch_generation() — decode path"
---
flowchart TD
Z["⓪ activate_step_by_batch(batch_size)<br/>query optimal step for current BS range, activate if different"]
A["① draft(batch)<br/>draft model multi-step generation with current tier"]
B["② verify(batch, spec_info)<br/>target model tree verification → produces accept_length_per_req"]
B["② verify(batch, spec_info)<br/>target model tree verification → produces num_correct_drafts_per_req"]
C["③ forward_draft_extend_after_decode(batch)<br/>draft model KV-cache catch-up"]
D["④ adaptive_controller.on_verify_complete(accept_lengths)<br/>update EMA, apply warmup / interval / hysteresis gates<br/>if tier changed, select a pre-built state from pool"]
D["④ adaptive_controller.on_verify_complete(num_correct_drafts_per_req, batch_size)<br/>update EMA for matching BS slot, apply warmup / interval / hysteresis gates<br/>if tier changed, select a pre-built state from pool"]
E["worker.apply_runtime_state(state)"]
A --> B --> C --> D --> E
Z --> A --> B --> C --> D --> E
```
> Tier switch happens after the current round completes. Backends and CUDA graphs are never swapped mid-round.
## How the policy decides
After each verify pass, SGLang reads the accepted draft length per request, computes the batch average, smooths it with an exponential moving average (EMA), and switches among the pre-built candidate tiers `[1, 3, 7]` by default.
After each verify pass, SGLang reads the accepted draft length per request, computes the batch average, smooths it with an exponential moving average (EMA), and switches among the candidate tiers for the matching BS slot.
The decision logic is intentionally conservative:
- `warmup_batches` skips the first few batches
- `update_interval` avoids switching every batch
- `down_hysteresis` and `up_hysteresis` reduce oscillation
- `ceiling_coeff` — an optional EMA ceiling rule can cap `num_steps` proportionally to observed draft quality, preventing over-speculation at high BS
Conceptually, the policy probes one step beyond the observed acceptance:
@@ -117,16 +127,21 @@ Example config:
```json
{
"candidate_steps": [1, 3, 7],
"ema_alpha": 0.2,
"warmup_batches": 10,
"update_interval": 5
"update_interval": 5,
"1": {"candidate_steps": [1, 3, 7], "up_hysteresis": 0.0, "down_hysteresis": -0.25, "ceiling_coeff": 0},
"8": {"candidate_steps": [1], "up_hysteresis": 0.0, "down_hysteresis": 0.0, "ceiling_coeff": 0}
}
```
Non-integer keys (`ema_alpha`, `warmup_batches`, `update_interval`) are global overrides applied to every BS slot. Integer keys (`"1"`, `"8"`) define per-BS slots.
## Config file reference
The config file is optional. Any omitted keys use defaults.
The config file is optional. When provided, each integer BS-slot key must specify `candidate_steps`; all other keys fall back to defaults.
### Per-BS slot parameters
<table style={{width: "100%", borderCollapse: "collapse", tableLayout: "fixed"}}>
<colgroup>
@@ -144,9 +159,43 @@ The config file is optional. Any omitted keys use defaults.
<tbody>
<tr>
<td><code>candidate_steps</code></td>
<td><code>[1, 3, 7]</code></td>
<td>Discrete <code>speculative_num_steps</code> tiers that adaptive mode can switch between</td>
<td><em>required</em></td>
<td>Candidate <code>speculative_num_steps</code> tiers for this BS range. Must be a non-empty list of positive ints; a slot that omits it raises a config error</td>
</tr>
<tr>
<td><code>down_hysteresis</code></td>
<td><code>-0.25</code></td>
<td>Extra margin before moving to a smaller step</td>
</tr>
<tr>
<td><code>up_hysteresis</code></td>
<td><code>0.0</code></td>
<td>Extra margin before moving to a larger step</td>
</tr>
<tr>
<td><code>ceiling_coeff</code></td>
<td><code>0</code> (disabled)</td>
<td>EMA ceiling coefficient; set &gt; 0 to cap steps proportionally to draft quality</td>
</tr>
</tbody>
</table>
### Global parameters
<table style={{width: "100%", borderCollapse: "collapse", tableLayout: "fixed"}}>
<colgroup>
<col style={{width: "33.33%"}} />
<col style={{width: "33.33%"}} />
<col style={{width: "33.33%"}} />
</colgroup>
<thead>
<tr>
<th>Key</th>
<th>Default</th>
<th>Meaning</th>
</tr>
</thead>
<tbody>
<tr>
<td><code>ema_alpha</code></td>
<td><code>0.2</code></td>
@@ -162,21 +211,9 @@ The config file is optional. Any omitted keys use defaults.
<td><code>10</code></td>
<td>Number of verify batches to observe before switching</td>
</tr>
<tr>
<td><code>down_hysteresis</code></td>
<td><code>-0.25</code></td>
<td>Extra margin before moving to a smaller step</td>
</tr>
<tr>
<td><code>up_hysteresis</code></td>
<td><code>0.0</code></td>
<td>Extra margin before moving to a larger step</td>
</tr>
</tbody>
</table>
The initial `--speculative-num-steps` is snapped to the nearest value in `candidate_steps`.
## Monitoring
You can inspect the active tier and acceptance metric via `/server_info`:
@@ -190,8 +227,43 @@ curl -s http://127.0.0.1:30000/server_info | jq '.internal_states[0] | {speculat
## Tuning tips
- Start with the default candidate tiers `[1, 3, 7]`
- Use fewer tiers if you want lower startup and graph-memory overhead
- Start with the built-in default (conservative) — it is safe for all draft model qualities
- For strong draft models, use the aggressive config with ceiling rule
- Use fewer candidate steps if you want lower startup GPU memory overhead
- Increase `ema_alpha` to react faster, or lower it for more stability
- Increase `warmup_batches` or `update_interval` if tier switching is too noisy
- At high batch sizes, narrower ladders (e.g., `[1, 2]` or `[1]`) often outperform wide ones
- If your workload is already stable and one static setting is well tuned, adaptive mode may not help much
## Recommended configs
The built-in default is conservative — safe for all draft models but may under-speculate for strong ones. Save one of these as a JSON file and pass via `--speculative-adaptive-config`.
### Conservative (default) — for weak draft models
This is the built-in default: BS 8–31 allows `[1, 3]`, and BS≥32 locks to `step=1` to avoid wasted compute. Best for models like MiniMax-M2.5, DSV4.
```json
{
"1": {"candidate_steps": [1, 3, 7], "up_hysteresis": 0.0, "down_hysteresis": -0.25, "ceiling_coeff": 0},
"8": {"candidate_steps": [1, 3], "up_hysteresis": 0.0, "down_hysteresis": 0.0, "ceiling_coeff": 0},
"32": {"candidate_steps": [1], "up_hysteresis": 0.0, "down_hysteresis": 0.0, "ceiling_coeff": 0}
}
```
### Aggressive — for strong or high-variance draft models
Uses wider ladders with ceiling rule to cap speculation at high BS. Best for models like GLM-4.7-FP8.
```json
{
"1": {"candidate_steps": [1, 3, 7], "up_hysteresis": 0.0, "down_hysteresis": -0.25, "ceiling_coeff": 0},
"8": {"candidate_steps": [1, 3, 7], "up_hysteresis": 0.0, "down_hysteresis": -0.25, "ceiling_coeff": 3.0},
"64": {"candidate_steps": [1, 3], "up_hysteresis": 0.0, "down_hysteresis": -0.25, "ceiling_coeff": 1.67},
"128": {"candidate_steps": [1, 3], "up_hysteresis": 0.0, "down_hysteresis": -0.25, "ceiling_coeff": 1.2}
}
```
### Custom per-model config
For the best performance, benchmark your specific model across batch sizes with different static `num_steps` values, then build a per-BS config that matches each range's optimal step. A well-tuned per-model config might outperform the generic presets above.
python/sglang/srt/arg_groups/speculative_hook.py
+9
View File
@@ -119,6 +119,15 @@ def handle_speculative_decoding(server_args: "ServerArgs") -> None:
if server_args.speculative_adaptive:
_maybe_disable_adaptive(server_args)
if server_args.speculative_adaptive:
from sglang.srt.speculative.adaptive_spec_params import (
validate_adaptive_initial_steps,
)
validate_adaptive_initial_steps(
server_args.speculative_num_steps,
server_args.speculative_adaptive_config,
)
def _handle_dflash(server_args: "ServerArgs") -> None:
python/sglang/srt/managers/scheduler_components/batch_result_processor.py
+3 -1
View File
@@ -541,7 +541,9 @@ class SchedulerBatchResultProcessor:
# Feed the adaptive controller now that accept_lens is on CPU,
# instead of doing a synchronous GPU→CPU copy in the worker hot path.
# BaseSpecWorker provides a no-op default for non-adaptive workers.
self.model_worker.on_verify_complete_cpu(result.num_correct_drafts_per_req_cpu)
self.model_worker.on_verify_complete_cpu(
result.num_correct_drafts_per_req_cpu, batch_size=len(batch.reqs)
)
predict_tokens = []
# In adaptive spec-v2, the worker state may already have switched when this
python/sglang/srt/server_args.py
+1 -2
View File
@@ -5956,7 +5956,7 @@ class ServerArgs:
parser.add_argument(
"--speculative-adaptive-config",
type=str,
help="Path to a JSON config file for adaptive speculative decoding tuning knobs ",
help="Path to a JSON config file for adaptive speculative decoding tuning knobs.",
default=ServerArgs.speculative_adaptive_config,
)
parser.add_argument(
@@ -7310,7 +7310,6 @@ class ServerArgs:
)
candidate_steps = resolve_candidate_steps_from_config(
initial_steps=self.speculative_num_steps,
cfg_path=self.speculative_adaptive_config,
)
# TODO: adaptive spec currently requires topk=1, so each runtime state
python/sglang/srt/speculative/adaptive_runtime_state.py
+30 -13
View File
@@ -1,4 +1,3 @@
import logging
from dataclasses import dataclass
from typing import TYPE_CHECKING, Protocol
@@ -15,8 +14,6 @@ if TYPE_CHECKING:
EAGLEDraftExtendCudaGraphRunner,
)
logger = logging.getLogger(__name__)
@dataclass
class SpecRuntimeState:
@@ -52,7 +49,10 @@ class AdaptiveSpecWorker(Protocol):
speculative_num_steps: int
def build_adaptive_runtime_state(
self, speculative_num_steps: int, speculative_num_draft_tokens: int
self,
speculative_num_steps: int,
speculative_num_draft_tokens: int,
cuda_graph_bs: list[int] | None = None,
) -> SpecRuntimeState: ...
def apply_runtime_state(self, state: SpecRuntimeState) -> None: ...
@@ -62,13 +62,13 @@ class AdaptiveController:
"""Facade that owns adaptive decision-making and runtime state switching.
Works with any worker that implements ``AdaptiveSpecWorker`` protocol:
- ``build_adaptive_runtime_state(steps, draft_tokens)`` → runtime state
- ``apply_runtime_state(state)`` → apply it to the worker
- ``build_adaptive_runtime_state()`` → runtime state
- ``apply_runtime_state()`` → apply it to the worker
The worker only needs to:
1. Call ``register()`` for the initial state, then ``init_states()``
once during startup.
2. Call ``on_verify_complete(num_correct_drafts_per_req)`` after each decode verify.
2. Call ``on_verify_complete()`` after each decode verify.
"""
def __init__(self, worker: AdaptiveSpecWorker, config_path: str | None = None):
@@ -91,22 +91,39 @@ class AdaptiveController:
key = steps if steps is not None else state.speculative_num_steps
self._states[key] = state
def init_states(self) -> None:
def init_states(self, cuda_graph_bs: list[int] | None = None) -> None:
"""Build and register runtime states for all candidate steps."""
for steps in self.params.candidate_steps:
self.params.set_cuda_graph_bs(cuda_graph_bs)
for steps in self.candidate_steps:
if steps in self._states:
continue
pruned_bs = self.params.cuda_graph_bs_for_step(steps)
state = self.worker.build_adaptive_runtime_state(
speculative_num_steps=steps,
speculative_num_draft_tokens=steps + 1,
cuda_graph_bs=pruned_bs,
)
self._states[steps] = state
self._activate(self.params.current_steps)
def on_verify_complete(self, num_correct_drafts_per_req: list[int]) -> None:
# Start on the initial step.
self._activate(self.worker.speculative_num_steps)
def activate_step_by_batch(self, batch_size: int) -> None:
target = self.params.get_steps_for_batch(batch_size)
if target != self.worker.speculative_num_steps:
self._activate(target)
def on_verify_complete(
self, num_correct_drafts_per_req: list[int], batch_size: int
) -> None:
"""Feed verify results; switch runtime state if EMA warrants it."""
if self.params.update(num_correct_drafts_per_req):
self._activate(self.params.current_steps)
new_step = self.params.on_verify_complete(
num_correct_drafts_per_req, batch_size
)
if new_step is not None:
self._activate(new_step)
def _activate(self, speculative_num_steps: int) -> None:
state = self._states.get(speculative_num_steps)
python/sglang/srt/speculative/adaptive_spec_params.py
+182 -59
View File
@@ -5,8 +5,11 @@ Adjusts speculative_num_steps at runtime based on observed acceptance lengths.
from __future__ import annotations
import bisect
import json
import logging
import math
from functools import cached_property
from typing import TYPE_CHECKING
from sglang.srt.utils import log_info_on_rank0
@@ -16,6 +19,28 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
# TODO: add step=0 (nospec fallback) for BS>=8 once supported.
DEFAULT_ADAPTIVE_CONFIG: dict[str, dict] = {
"1": {
"candidate_steps": [1, 3, 7],
"up_hysteresis": 0.0,
"down_hysteresis": -0.25,
"ceiling_coeff": 0,
},
"8": {
"candidate_steps": [1, 3],
"up_hysteresis": 0.0,
"down_hysteresis": 0.0,
"ceiling_coeff": 0,
},
"32": {
"candidate_steps": [1],
"up_hysteresis": 0.0,
"down_hysteresis": 0.0,
"ceiling_coeff": 0,
},
}
def adaptive_unsupported_reason(server_args: ServerArgs) -> str | None:
"""Return why adaptive spec cannot run under the given server args, or None if supported."""
@@ -52,57 +77,69 @@ def adaptive_unsupported_reason(server_args: ServerArgs) -> str | None:
return None
def load_adaptive_config(path: str | None) -> dict[str, object]:
"""Load adaptive speculative config from a JSON file.
def _load_adaptive_config(
cfg_path: str | None,
) -> tuple[dict, dict[int, dict]]:
"""Load and validate adaptive config.
The file may contain any subset of the following keys:
ema_alpha, update_interval, warmup_batches,
down_hysteresis, up_hysteresis, candidate_steps
Returns an empty dict when *path* is ``None``.
Uses ``DEFAULT_ADAPTIVE_CONFIG`` when *cfg_path* is ``None``.
"""
if path is None:
return {}
with open(path) as f:
cfg = json.load(f)
if not isinstance(cfg, dict):
if cfg_path is not None:
with open(cfg_path) as f:
cfg = json.load(f)
else:
cfg = DEFAULT_ADAPTIVE_CONFIG
bs_entries: dict[int, dict] = {}
for key, entry in cfg.items():
if not key.isdigit():
continue
steps = entry.get("candidate_steps")
if (
not isinstance(steps, list)
or not steps
or not all(isinstance(s, int) and s > 0 for s in steps)
):
raise ValueError(
f"BS {key}: candidate_steps must be a list of positive ints, got {steps!r}"
)
bs_entries[int(key)] = entry
if not bs_entries:
raise ValueError(
"speculative_adaptive_config must be a JSON object, "
f"got {type(cfg).__name__}"
"speculative_adaptive_config must contain at least one integer-string "
'BS key, e.g. {"1": {"candidate_steps": [1,3,7]}}. '
f"Got keys: {list(cfg.keys())}"
)
return cfg
def _resolve_candidate_steps(initial_steps: int, cfg: dict[str, object]) -> list[int]:
"""Return sorted, deduplicated candidate steps; inserts *initial_steps* when missing."""
raw = cfg.get("candidate_steps") or (1, 3, 7)
candidates: set[int] = set(raw)
# Ensure the worker's initial speculative_num_steps is itself a candidate.
# Otherwise AdaptiveController.register() would store the worker's pre-built
# runtime state under a key that _activate() never queries, leaking that
# state's draft attn backend and cuda graph buffers for the process lifetime.
if initial_steps not in candidates:
log_info_on_rank0(
logger,
f"Adding initial speculative_num_steps={initial_steps} to "
f"candidate_steps={sorted(candidates)} so the pre-built "
f"runtime state is reused.",
)
candidates.add(initial_steps)
return sorted(candidates)
return cfg, bs_entries
def resolve_candidate_steps_from_config(
initial_steps: int, cfg_path: str | None
cfg_path: str | None = None,
) -> list[int]:
"""Load adaptive config and resolve candidate steps."""
cfg = load_adaptive_config(cfg_path)
return _resolve_candidate_steps(initial_steps, cfg)
"""Union of every BS slot's candidate steps; sizes the runtime buffers."""
_, bs_entries = _load_adaptive_config(cfg_path)
all_steps: set[int] = set()
for entry in bs_entries.values():
all_steps.update(entry["candidate_steps"])
return sorted(all_steps)
class AdaptiveSpeculativeParams:
def validate_adaptive_initial_steps(
initial_steps: int,
cfg_path: str | None = None,
) -> None:
"""Require the initial step to be a candidate of some BS slot."""
candidate_steps = resolve_candidate_steps_from_config(cfg_path)
if initial_steps not in candidate_steps:
raise ValueError(
f"--speculative-num-steps={initial_steps} is not in the adaptive "
f"config candidate_steps {candidate_steps}. Pass one of those values."
)
class AdaptiveStepSlot:
"""Tracks acceptance rate via EMA and adapts num_steps accordingly.
The core idea: if drafts are consistently accepted, try more steps;
@@ -112,38 +149,30 @@ class AdaptiveSpeculativeParams:
- Probes one step beyond observed acceptance
- EMA smoothing prevents oscillation
- Only updates every `update_interval` batches for stability
- num_steps can be selected from different candidate sets on different batch_sizes
"""
def __init__(
self,
initial_steps: int,
cfg_path: str | None = None,
):
cfg = load_adaptive_config(cfg_path)
# TODO: Wider range of candidate_steps (once lazy init is supported).
self.candidate_steps = _resolve_candidate_steps(initial_steps, cfg)
assert (
len(self.candidate_steps) >= 2
), "candidate_steps must have at least 2 distinct values"
def __init__(self, initial_steps: int, cfg: dict):
candidates = sorted(set(cfg["candidate_steps"]))
assert len(candidates) >= 1, "candidate_steps must have at least 1 value"
self.candidate_steps = candidates
self.ema_alpha = cfg.get("ema_alpha", 0.2)
self.update_interval = cfg.get("update_interval", 5)
self.warmup_batches = cfg.get("warmup_batches", 10)
self.down_hysteresis = cfg.get("down_hysteresis", -0.25)
self.up_hysteresis = cfg.get("up_hysteresis", 0.0)
self.ceiling_coeff = cfg.get("ceiling_coeff", 0)
self.current_steps = initial_steps
if initial_steps in self.candidate_steps:
self.current_steps = initial_steps
else:
self.current_steps = self.candidate_steps[len(self.candidate_steps) // 2]
# Initialize EMA at current steps - 1 (neutral starting point)
self.ema_accept_len = float(self.current_steps - 1)
self._batch_count = 0
log_info_on_rank0(
logger,
f"AdaptiveSpeculativeParams initialized: "
f"steps={self.current_steps}, candidate_steps={self.candidate_steps}",
)
def update(self, num_correct_drafts_per_req: list[int]) -> bool:
"""Update EMA with observed accept lengths. Returns True if params changed.
@@ -190,6 +219,14 @@ class AdaptiveSpeculativeParams:
break
target = self.candidate_steps[current_idx]
# EMA ceiling: only caps downward — never blocks step-ups, so the
# system can explore higher steps and let the EMA catch up.
if self.ceiling_coeff > 0:
ceiling = max(1, math.ceil(self.ema_accept_len * self.ceiling_coeff))
if target > ceiling and target <= old_steps:
while current_idx > 0 and self.candidate_steps[current_idx] > ceiling:
current_idx -= 1
target = self.candidate_steps[current_idx]
if target != old_steps:
self.current_steps = target
@@ -200,3 +237,89 @@ class AdaptiveSpeculativeParams:
)
return True
return False
class AdaptiveSpeculativeParams:
"""Routes ``batch_size`` to the correct per-BS slot.
A slot is a per-BS configuration of adaptive step selection.
"""
def __init__(
self,
initial_steps: int,
cfg_path: str | None = None,
):
cfg, bs_entries = _load_adaptive_config(cfg_path)
self._bs_list: list[int] = sorted(bs_entries)
self._slots: dict[int, AdaptiveStepSlot] = {}
self._cuda_graph_bs: list[int] | None = None
for bs, entry in sorted(bs_entries.items()):
self._slots[bs] = AdaptiveStepSlot(
initial_steps=initial_steps,
cfg={**cfg, **entry},
)
first_slot = self._slots[self._bs_list[0]]
log_info_on_rank0(
logger,
f"AdaptiveSpeculativeParams initialized: "
f"steps={first_slot.current_steps}, "
f"candidate_steps={first_slot.candidate_steps}",
)
@cached_property
def candidate_steps(self) -> list[int]:
"""Union of all BS slots' candidate steps."""
return sorted({s for p in self._slots.values() for s in p.candidate_steps})
def set_cuda_graph_bs(self, cuda_graph_bs: list[int] | None) -> None:
self._cuda_graph_bs = sorted(cuda_graph_bs) if cuda_graph_bs else None
def get_steps_for_batch(self, batch_size: int) -> int:
return self._route(batch_size).current_steps
def on_verify_complete(
self, num_correct_drafts_per_req: list[int], batch_size: int
) -> int | None:
"""Feed verify results to the matching BS slot's EMA.
Returns the new step if a switch is warranted, else ``None``.
"""
params = self._route(batch_size)
if params.update(num_correct_drafts_per_req):
return params.current_steps
return None
def cuda_graph_bs_for_step(self, step: int) -> list[int] | None:
"""Return cuda_graph_bs values that can reach *step* at runtime.
Returns ``None`` when CUDA graphs are disabled (``set_cuda_graph_bs``
was never called or was called with ``None``).
"""
if self._cuda_graph_bs is None:
return None
return [
v
for v in self._cuda_graph_bs
if step in self._slots[self._find_closest_bs(v)].candidate_steps
]
def _route(self, batch_size: int) -> AdaptiveStepSlot:
"""Map *batch_size* → pad to CUDA-graph BS → closest slot."""
return self._slots[
self._find_closest_bs(self._pad_to_cuda_graph_bs(batch_size))
]
def _pad_to_cuda_graph_bs(self, batch_size: int) -> int:
if self._cuda_graph_bs is None:
return batch_size
idx = bisect.bisect_left(self._cuda_graph_bs, batch_size)
return (
self._cuda_graph_bs[idx] if idx < len(self._cuda_graph_bs) else batch_size
)
def _find_closest_bs(self, target: int) -> int:
idx = bisect.bisect_right(self._bs_list, target) - 1
return self._bs_list[max(0, idx)]
python/sglang/srt/speculative/base_spec_worker.py
+11 -1
View File
@@ -39,10 +39,20 @@ class BaseSpecWorker(ABC):
# TODO: move this abstract method to BaseTpWorker and call through self.model_runner
pass
def on_verify_complete_cpu(self, num_correct_drafts_per_req: list[int]) -> None:
def on_verify_complete_cpu(
self, num_correct_drafts_per_req: list[int], batch_size: int = 0
) -> None:
"""Hook called after verify finishes and accept counts are on CPU.
Default no-op. Adaptive-aware workers override this to feed the
controller without forcing a GPU→CPU sync in the worker hot path.
"""
pass
def activate_step_by_batch(self, batch_size: int) -> None:
"""Activate the optimal adaptive step for the current batch size.
Default no-op. Adaptive-aware workers override this to switch
the runtime state before each draft round.
"""
pass
python/sglang/srt/speculative/eagle_worker.py
+37 -6
View File
@@ -125,7 +125,8 @@ class EAGLEWorker(TpModelWorker):
self.adaptive_controller: Optional[AdaptiveController] = None
if server_args.speculative_adaptive:
self.adaptive_controller = AdaptiveController(
self, config_path=server_args.speculative_adaptive_config
self,
config_path=server_args.speculative_adaptive_config,
)
# Override the context length of the draft model to be the same as the target model.
@@ -244,7 +245,13 @@ class EAGLEWorker(TpModelWorker):
cuda_graph_runner_for_draft_extend=self.cuda_graph_runner_for_draft_extend,
)
)
self.adaptive_controller.init_states()
self.adaptive_controller.init_states(
cuda_graph_bs=(
None
if self.server_args.disable_cuda_graph
else self.server_args.cuda_graph_bs
),
)
# Some dummy tensors
self.num_new_pages_per_topk = torch.empty(
@@ -352,14 +359,19 @@ class EAGLEWorker(TpModelWorker):
)
def build_adaptive_runtime_state(
self, speculative_num_steps: int, speculative_num_draft_tokens: int
self,
speculative_num_steps: int,
speculative_num_draft_tokens: int,
cuda_graph_bs: list[int] | None = None,
) -> SpecRuntimeState:
"""Build a SpecRuntimeState for the given step configuration."""
tic = time.perf_counter()
before_mem = get_available_gpu_memory(self.device, self.gpu_id)
with self._override_worker_state(
speculative_num_steps, speculative_num_draft_tokens
speculative_num_steps,
speculative_num_draft_tokens,
cuda_graph_bs=cuda_graph_bs,
):
# Reuse existing init methods for draft attention backend and cuda graphs
self.init_attention_backend()
@@ -411,7 +423,10 @@ class EAGLEWorker(TpModelWorker):
@contextmanager
def _override_worker_state(
self, speculative_num_steps: int, speculative_num_draft_tokens: int
self,
speculative_num_steps: int,
speculative_num_draft_tokens: int,
cuda_graph_bs: list[int] | None = None,
):
"""Temporarily override server_args and worker attributes for graph capture."""
sa = self.server_args
@@ -425,11 +440,21 @@ class EAGLEWorker(TpModelWorker):
self.cuda_graph_runner_for_draft_extend,
sa.speculative_num_steps,
sa.speculative_num_draft_tokens,
sa.cuda_graph_bs,
sa.disable_cuda_graph,
)
self.speculative_num_steps = speculative_num_steps
self.speculative_num_draft_tokens = speculative_num_draft_tokens
sa.speculative_num_steps = speculative_num_steps
sa.speculative_num_draft_tokens = speculative_num_draft_tokens
if cuda_graph_bs is not None:
sa.cuda_graph_bs = cuda_graph_bs
# BS-aware adaptive spec may prune cuda_graph_bs to an empty list
# for steps that no BS range uses (e.g. step=1). Disable graph
# capture for those steps; restore in finally so subsequent steps
# are not affected.
if not cuda_graph_bs:
sa.disable_cuda_graph = True
try:
yield
finally:
@@ -443,6 +468,8 @@ class EAGLEWorker(TpModelWorker):
self.cuda_graph_runner_for_draft_extend,
sa.speculative_num_steps,
sa.speculative_num_draft_tokens,
sa.cuda_graph_bs,
sa.disable_cuda_graph,
) = backup
@property
@@ -487,6 +514,9 @@ class EAGLEWorker(TpModelWorker):
can_run_cuda_graph=can_run_cuda_graph,
)
else:
if self.adaptive_controller is not None:
self.adaptive_controller.activate_step_by_batch(batch.batch_size())
set_time_batch(batch.reqs, "set_spec_draft_start_time", trace_only=True)
with (
@@ -547,7 +577,8 @@ class EAGLEWorker(TpModelWorker):
if self.adaptive_controller is not None:
self.adaptive_controller.on_verify_complete(
verify_output.num_correct_drafts_per_req_cpu
verify_output.num_correct_drafts_per_req_cpu,
batch_size=batch.batch_size(),
)
return GenerationBatchResult(
python/sglang/srt/speculative/eagle_worker_v2.py
+44 -7
View File
@@ -851,7 +851,8 @@ class EAGLEWorkerV2(BaseSpecWorker):
self.adaptive_controller: Optional[AdaptiveController] = None
if server_args.speculative_adaptive:
self.adaptive_controller = AdaptiveController(
self, config_path=server_args.speculative_adaptive_config
self,
config_path=server_args.speculative_adaptive_config,
)
# Some dummy tensors
@@ -883,7 +884,13 @@ class EAGLEWorkerV2(BaseSpecWorker):
cuda_graph_runner_for_draft_extend=self._draft_worker.cuda_graph_runner_for_draft_extend,
)
)
self.adaptive_controller.init_states()
self.adaptive_controller.init_states(
cuda_graph_bs=(
None
if self.server_args.disable_cuda_graph
else self.server_args.cuda_graph_bs
),
)
@property
def spec_v2_attn_backends(self) -> tuple:
@@ -943,6 +950,8 @@ class EAGLEWorkerV2(BaseSpecWorker):
)
return batch_output
else:
self.activate_step_by_batch(batch.seq_lens.shape[0])
if batch.spec_info is None:
capture_mode = (
CaptureHiddenMode.NULL
@@ -981,21 +990,34 @@ class EAGLEWorkerV2(BaseSpecWorker):
return batch_output
def on_verify_complete_cpu(self, num_correct_drafts_per_req: list[int]) -> None:
def on_verify_complete_cpu(
self, num_correct_drafts_per_req: list[int], batch_size: int = 0
) -> None:
if self.adaptive_controller is not None:
self.adaptive_controller.on_verify_complete(num_correct_drafts_per_req)
self.adaptive_controller.on_verify_complete(
num_correct_drafts_per_req, batch_size=batch_size
)
def activate_step_by_batch(self, batch_size: int) -> None:
if self.adaptive_controller is not None:
self.adaptive_controller.activate_step_by_batch(batch_size)
# -- Adaptive speculative decoding protocol --
def build_adaptive_runtime_state(
self, speculative_num_steps: int, speculative_num_draft_tokens: int
self,
speculative_num_steps: int,
speculative_num_draft_tokens: int,
cuda_graph_bs=None,
) -> SpecRuntimeState:
"""Build a SpecRuntimeState for the given step configuration."""
tic = time.perf_counter()
before_mem = get_available_gpu_memory(self.device, self.gpu_id)
with self._override_worker_state(
speculative_num_steps, speculative_num_draft_tokens
speculative_num_steps,
speculative_num_draft_tokens,
cuda_graph_bs=cuda_graph_bs,
):
self._draft_worker.init_attention_backend()
self._draft_worker.init_cuda_graphs()
@@ -1081,7 +1103,10 @@ class EAGLEWorkerV2(BaseSpecWorker):
@contextlib.contextmanager
def _override_worker_state(
self, speculative_num_steps: int, speculative_num_draft_tokens: int
self,
speculative_num_steps: int,
speculative_num_draft_tokens: int,
cuda_graph_bs: list[int] | None = None,
):
"""Temporarily override server_args and worker attributes for graph capture."""
sa = self.server_args
@@ -1098,6 +1123,8 @@ class EAGLEWorkerV2(BaseSpecWorker):
dw.cuda_graph_runner_for_draft_extend,
sa.speculative_num_steps,
sa.speculative_num_draft_tokens,
sa.cuda_graph_bs,
sa.disable_cuda_graph,
)
self.speculative_num_steps = speculative_num_steps
@@ -1106,6 +1133,14 @@ class EAGLEWorkerV2(BaseSpecWorker):
dw.speculative_num_draft_tokens = speculative_num_draft_tokens
sa.speculative_num_steps = speculative_num_steps
sa.speculative_num_draft_tokens = speculative_num_draft_tokens
if cuda_graph_bs is not None:
sa.cuda_graph_bs = cuda_graph_bs
# BS-aware adaptive spec may prune cuda_graph_bs to an empty list
# for steps that no BS range uses (e.g. step=1). Disable graph
# capture for those steps; restore in finally so subsequent steps
# are not affected.
if not cuda_graph_bs:
sa.disable_cuda_graph = True
dw._rebuild_topk1_chain_buffers()
try:
@@ -1123,6 +1158,8 @@ class EAGLEWorkerV2(BaseSpecWorker):
dw.cuda_graph_runner_for_draft_extend,
sa.speculative_num_steps,
sa.speculative_num_draft_tokens,
sa.cuda_graph_bs,
sa.disable_cuda_graph,
) = backup
dw._rebuild_topk1_chain_buffers()
test/registered/spec/eagle/test_adaptive_speculative.py
+7 -5
View File
@@ -47,11 +47,13 @@ class TestAdaptiveSpeculativeServer(CustomTestCase):
with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False) as f:
json.dump(
{
"candidate_steps": [1, 3],
"ema_alpha": 1.0,
"warmup_batches": 1,
"update_interval": 1,
"up_hysteresis": 0.0,
"1": {
"candidate_steps": [1, 3],
"ema_alpha": 1.0,
"warmup_batches": 1,
"update_interval": 1,
"up_hysteresis": 0.0,
},
},
f,
)
test/registered/unit/spec/test_adaptive_spec_params.py
+233 -30
View File
@@ -4,6 +4,9 @@ import unittest
from sglang.srt.speculative.adaptive_spec_params import (
AdaptiveSpeculativeParams,
AdaptiveStepSlot,
resolve_candidate_steps_from_config,
validate_adaptive_initial_steps,
)
from sglang.test.ci.ci_register import register_cpu_ci, register_xpu_ci
@@ -11,39 +14,16 @@ register_cpu_ci(est_time=6, suite="base-a-test-cpu")
register_xpu_ci(est_time=10, suite="stage-a-test-1-gpu-xpu")
class TestAdaptiveSpeculativeParams(unittest.TestCase):
def _make_params_from_config(self, initial_steps: int, config: dict[str, object]):
with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
json.dump(config, f)
f.flush()
return AdaptiveSpeculativeParams(
initial_steps=initial_steps, cfg_path=f.name
)
class TestAdaptiveStepSlot(unittest.TestCase):
def _make_params_from_config(self, initial_steps: int, config: dict):
return AdaptiveStepSlot(initial_steps=initial_steps, cfg=config)
def test_params_loads_config_path(self):
with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
json.dump(
{
"candidate_steps": [1, 5],
"ema_alpha": 0.75,
"warmup_batches": 2,
},
f,
)
f.flush()
params = AdaptiveSpeculativeParams(initial_steps=3, cfg_path=f.name)
self.assertEqual(params.candidate_steps, [1, 3, 5])
self.assertEqual(params.ema_alpha, 0.75)
self.assertEqual(params.warmup_batches, 2)
def test_initial_steps_added_to_candidates_when_missing(self):
def test_initial_steps_snaps_to_middle_when_missing(self):
params = self._make_params_from_config(2, {"candidate_steps": [1, 3, 7]})
self.assertEqual(params.candidate_steps, [1, 2, 3, 7])
self.assertEqual(params.current_steps, 2)
self.assertEqual(params.ema_accept_len, 1.0)
self.assertEqual(params.candidate_steps, [1, 3, 7])
self.assertEqual(params.current_steps, 3)
self.assertEqual(params.ema_accept_len, 2.0)
def test_update_respects_warmup_and_interval(self):
params = self._make_params_from_config(
@@ -219,6 +199,229 @@ class TestAdaptiveSpeculativeParams(unittest.TestCase):
self.assertEqual(params.current_steps, 1)
self.assertEqual(params.ema_accept_len, 0.375)
def test_ceiling_coeff_caps_steps(self):
params = self._make_params_from_config(
7,
{
"candidate_steps": [1, 3, 7],
"ema_alpha": 1.0,
"warmup_batches": 0,
"update_interval": 1,
"ceiling_coeff": 1.0,
},
)
# Force low ema to trigger ceiling
params.ema_accept_len = 1.0
self.assertTrue(params.update([1, 1]))
# ceiling = ceil(1.0 * 1.0) = 1, target capped to 1
self.assertEqual(params.current_steps, 1)
def test_ceiling_disabled_by_default(self):
params = self._make_params_from_config(3, {"candidate_steps": [1, 3, 7]})
self.assertEqual(params.ceiling_coeff, 0)
class TestAdaptiveSpeculativeParams(unittest.TestCase):
def test_default_config_loads(self):
params = AdaptiveSpeculativeParams(initial_steps=3)
self.assertEqual(params._bs_list, [1, 8, 32])
self.assertEqual(params._slots[1].candidate_steps, [1, 3, 7])
self.assertEqual(params._slots[8].candidate_steps, [1, 3])
self.assertEqual(params._slots[32].candidate_steps, [1])
def test_config_file(self):
with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
json.dump(
{
"1": {"candidate_steps": [1, 5], "up_hysteresis": 0.3},
"32": {"candidate_steps": [1, 2]},
},
f,
)
f.flush()
params = AdaptiveSpeculativeParams(initial_steps=5, cfg_path=f.name)
self.assertEqual(params._bs_list, [1, 32])
# Slots are built straight from the config; the launch flag never pollutes
# them. initial_steps just selects the smallest slot's starting step.
self.assertEqual(params._slots[1].candidate_steps, [1, 5])
self.assertEqual(params._slots[1].current_steps, 5)
self.assertEqual(params._slots[1].up_hysteresis, 0.3)
self.assertEqual(params._slots[32].candidate_steps, [1, 2])
def test_launch_flag_not_injected_into_slots(self):
# initial_steps lives only in a larger slot. It must NOT be merged into
# any other slot's candidates: slots come straight from the config.
with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
json.dump(
{
"1": {"candidate_steps": [1, 5]},
"8": {"candidate_steps": [1, 3, 7]},
},
f,
)
f.flush()
params = AdaptiveSpeculativeParams(initial_steps=7, cfg_path=f.name)
self.assertEqual(params._slots[1].candidate_steps, [1, 5])
self.assertEqual(params._slots[8].candidate_steps, [1, 3, 7])
# The slot that does not own initial_steps starts at its own median.
self.assertEqual(params._slots[1].current_steps, 5)
self.assertEqual(params._slots[8].current_steps, 7)
def test_invalid_config_raises(self):
with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
json.dump({"not_a_bs": "bad"}, f)
f.flush()
with self.assertRaises(ValueError):
AdaptiveSpeculativeParams(initial_steps=3, cfg_path=f.name)
def test_invalid_steps_raises(self):
with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
json.dump({"1": {"candidate_steps": "bad"}}, f)
f.flush()
with self.assertRaises(ValueError):
AdaptiveSpeculativeParams(initial_steps=3, cfg_path=f.name)
def test_empty_steps_raises(self):
with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
json.dump({"1": {"candidate_steps": []}}, f)
f.flush()
with self.assertRaises(ValueError):
AdaptiveSpeculativeParams(initial_steps=3, cfg_path=f.name)
def test_zero_steps_raises(self):
with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
json.dump({"1": {"candidate_steps": [0]}}, f)
f.flush()
with self.assertRaises(ValueError):
AdaptiveSpeculativeParams(initial_steps=3, cfg_path=f.name)
def test_global_hysteresis_inherited(self):
with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
json.dump(
{
"up_hysteresis": 0.5,
"1": {"candidate_steps": [1, 3]},
},
f,
)
f.flush()
params = AdaptiveSpeculativeParams(initial_steps=3, cfg_path=f.name)
self.assertEqual(params._slots[1].up_hysteresis, 0.5)
def test_entry_hysteresis_overrides_global(self):
with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
json.dump(
{
"up_hysteresis": 0.5,
"1": {"candidate_steps": [1, 3], "up_hysteresis": 0.1},
},
f,
)
f.flush()
params = AdaptiveSpeculativeParams(initial_steps=3, cfg_path=f.name)
self.assertEqual(params._slots[1].up_hysteresis, 0.1)
class TestBatchSizeRouting(unittest.TestCase):
"""BS-aware routing: batch size selects the slot, CUDA-graph BS pads first."""
def _params(self):
# Slots: bs=1 -> [1,3,7], bs=8 -> [1,3], bs=32 -> [1].
return AdaptiveSpeculativeParams(initial_steps=3)
def test_routes_to_floor_slot_without_cuda_graph(self):
params = self._params()
# A batch maps to the largest slot BS <= batch (floor), capped at the top slot.
self.assertEqual(params._route(1).candidate_steps, [1, 3, 7])
self.assertEqual(params._route(7).candidate_steps, [1, 3, 7])
self.assertEqual(params._route(8).candidate_steps, [1, 3])
self.assertEqual(params._route(31).candidate_steps, [1, 3])
self.assertEqual(params._route(32).candidate_steps, [1])
self.assertEqual(params._route(1000).candidate_steps, [1])
def test_cuda_graph_bs_pads_batch_up_before_routing(self):
params = self._params()
params.set_cuda_graph_bs([4, 8, 16, 32])
# bs=5 pads up to the captured graph BS 8 -> slot bs=8.
self.assertEqual(params._route(5).candidate_steps, [1, 3])
# bs=17 pads up to 32 -> slot bs=32.
self.assertEqual(params._route(17).candidate_steps, [1])
# A batch larger than every captured BS keeps its own value -> top slot.
self.assertEqual(params._route(100).candidate_steps, [1])
def test_cuda_graph_bs_for_step_prunes_unreachable_graphs(self):
params = self._params()
params.set_cuda_graph_bs([4, 8, 16, 32])
# step=1 is reachable from every slot.
self.assertEqual(params.cuda_graph_bs_for_step(1), [4, 8, 16, 32])
# step=3 lives in the bs=1 and bs=8 slots: graphs 4,8,16 floor into them.
self.assertEqual(params.cuda_graph_bs_for_step(3), [4, 8, 16])
# step=7 lives only in the bs=1 slot: only graph BS 4 floors into it.
self.assertEqual(params.cuda_graph_bs_for_step(7), [4])
def test_cuda_graph_bs_for_step_returns_none_when_disabled(self):
params = self._params()
self.assertIsNone(params.cuda_graph_bs_for_step(7))
params.set_cuda_graph_bs(None)
self.assertIsNone(params.cuda_graph_bs_for_step(7))
def test_observe_verify_feeds_the_routed_slot(self):
params = self._params()
# Drive the bs=1 slot up with perfect acceptance; the bs=32 slot is
# untouched and stays at its single candidate step.
for _ in range(40):
params.on_verify_complete([7, 7, 7], batch_size=1)
self.assertGreater(params.get_steps_for_batch(1), 1)
self.assertEqual(params.get_steps_for_batch(32), 1)
class TestResolveCandidateSteps(unittest.TestCase):
def test_default_config(self):
steps = resolve_candidate_steps_from_config()
self.assertIn(1, steps)
self.assertIn(3, steps)
self.assertIn(7, steps)
def test_config_file(self):
with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
json.dump({"1": {"candidate_steps": [2, 4]}}, f)
f.flush()
steps = resolve_candidate_steps_from_config(cfg_path=f.name)
self.assertEqual(steps, [2, 4])
def test_unions_and_dedups_across_slots(self):
with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
json.dump(
{
"1": {"candidate_steps": [1, 5]},
"8": {"candidate_steps": [3, 5, 7]},
},
f,
)
f.flush()
steps = resolve_candidate_steps_from_config(cfg_path=f.name)
self.assertEqual(steps, [1, 3, 5, 7])
class TestValidateAdaptiveInitialSteps(unittest.TestCase):
def test_accepts_value_from_any_slot(self):
with tempfile.NamedTemporaryFile("w", suffix=".json") as f:
json.dump(
{
"1": {"candidate_steps": [1, 5]},
"8": {"candidate_steps": [1, 3, 7]},
},
f,
)
f.flush()
# Membership in any slot is enough: 5 lives in the smallest slot,
# 7 only in a larger slot -- both accepted.
validate_adaptive_initial_steps(5, cfg_path=f.name)
validate_adaptive_initial_steps(7, cfg_path=f.name)
# 9 is in no slot -> rejected.
with self.assertRaises(ValueError):
validate_adaptive_initial_steps(9, cfg_path=f.name)
if __name__ == "__main__":
unittest.main()