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

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>

python/sglang/srt/arg_groups/speculative_hook.py

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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()