chore: import upstream snapshot with attribution

python/sglang/srt/speculative/adaptive_runtime_state.py

@@ -0,0 +1,134 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Protocol
+
+from sglang.srt.speculative.adaptive_spec_params import AdaptiveSpeculativeParams
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+    from sglang.srt.model_executor.cpu_graph_runner import CPUGraphRunner
+    from sglang.srt.model_executor.runner import DecodeCudaGraphRunner
+    from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
+        EAGLEDraftCudaGraphRunner,
+    )
+    from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import (
+        EAGLEDraftExtendCudaGraphRunner,
+    )
+
+
+@dataclass
+class SpecRuntimeState:
+    """A complete set of runtime resources bound to a specific speculative
+    decoding configuration.
+
+    Each decode round runs three stages — draft, verify, extend — and every
+    stage has shape-dependent resources (attention backends and CUDA graphs)
+    that must match the current configuration.  Switching adaptive steps
+    means swapping the entire state atomically.
+    """
+
+    # -- Configuration (determines shapes for all stages) --
+    speculative_num_steps: int
+    speculative_num_draft_tokens: int
+
+    # -- Draft stage: draft model multi-step autoregressive generation --
+    draft_attn_backend: "AttentionBackend | None"
+    cuda_graph_runner: "EAGLEDraftCudaGraphRunner | None"
+
+    # -- Verify stage: target model one-pass tree verification --
+    target_attn_backend: "AttentionBackend"
+    target_graph_runner: "DecodeCudaGraphRunner | CPUGraphRunner | None"
+
+    # -- Extend stage: draft model KV cache catch-up after verify --
+    draft_extend_attn_backend: "AttentionBackend | None"
+    cuda_graph_runner_for_draft_extend: "EAGLEDraftExtendCudaGraphRunner | None"
+
+
+class AdaptiveSpecWorker(Protocol):
+    """Protocol that a worker must implement to use AdaptiveController."""
+
+    speculative_num_steps: int
+
+    def build_adaptive_runtime_state(
+        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: ...
+
+
+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
+
+    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.
+    """
+
+    def __init__(self, worker: AdaptiveSpecWorker, config_path: str | None = None):
+        self.worker = worker
+        self.params = AdaptiveSpeculativeParams(
+            initial_steps=worker.speculative_num_steps,
+            cfg_path=config_path,
+        )
+        self._states: dict[int, SpecRuntimeState] = {}
+
+    @property
+    def candidate_steps(self) -> list[int]:
+        return self.params.candidate_steps
+
+    def register(self, state: SpecRuntimeState, steps: int | None = None) -> None:
+        """Register a pre-built runtime state.
+
+        *steps* defaults to state.speculative_num_steps when not given.
+        """
+        key = steps if steps is not None else state.speculative_num_steps
+        self._states[key] = state
+
+    def init_states(self, cuda_graph_bs: list[int] | None = None) -> None:
+        """Build and register runtime states for all 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
+
+        # 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."""
+        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)
+        if state is None:
+            raise ValueError(
+                f"Missing adaptive runtime state for steps={speculative_num_steps}"
+            )
+        self.worker.apply_runtime_state(state)