sgl-project--sglang/python/sglang/srt/speculative/dspark_components/dspark_observability.py

from __future__ import annotations

import logging
import math
import statistics
import time
from collections import deque
from contextlib import contextmanager, nullcontext
from enum import Enum
from typing import Callable, ContextManager, Iterator, Optional, Union

import msgspec
import torch

from sglang.srt.environ import envs
from sglang.srt.kv_canary.runner.future_tensor import FutureTensors
from sglang.srt.runtime_context import get_parallel
from sglang.srt.sampling.sampling_params import TOP_K_ALL
from sglang.srt.speculative.dflash_utils import compute_dflash_correct_drafts_and_bonus
from sglang.srt.speculative.dspark_components.dspark_block_accept_estimator import (
    create_block_accept_estimate_recorder,
)
from sglang.srt.speculative.dspark_components.dspark_sts import StsDataRecorder
from sglang.srt.speculative.dspark_components.dspark_verify import (
    verify_logits_adjustments_are_noop,
)

logger = logging.getLogger(__name__)

_NULL_SEGMENT = nullcontext()

ALL_COMPONENTS_TOKEN = "all"


class InfoComponent(str, Enum):
    CORE = "core"
    STEP_CPU_TIME = "step_cpu_time"
    STEP_GPU_TIME = "step_gpu_time"
    DRAFT_GPU_TIME = "draft_gpu_time"
    TARGET_VERIFY_GPU_TIME = "target_verify_gpu_time"
    REQS = "reqs"


class InfoSegment(str, Enum):
    STEP = "step"
    DRAFT = "draft"
    TARGET_VERIFY = "target_verify"


INFO_DUMP_MAX_RECORDS = 200_000
INFO_DUMP_MAX_STEP_CPU_SECONDS = 1.0


def resolve_enabled_components() -> set[InfoComponent]:
    """Components enabled via env: SGLANG_DSPARK_DEBUG_DUMP tokens, plus the
    published SPS-profiling switch SGLANG_DSPARK_ENABLE_SPS_RECORD=1, which is
    an alias for the core,step_cpu_time components the SPS table fit needs."""
    components = resolve_components(envs.SGLANG_DSPARK_DEBUG_DUMP.get())
    if envs.SGLANG_DSPARK_ENABLE_SPS_RECORD.get():
        components |= {InfoComponent.CORE, InfoComponent.STEP_CPU_TIME}
    return components


def resolve_components(raw: tuple[str, ...]) -> set[InfoComponent]:
    tokens = {token.strip() for token in raw if token.strip()}
    if not tokens:
        return set()
    if ALL_COMPONENTS_TOKEN in tokens:
        return set(InfoComponent)
    try:
        return {InfoComponent(token) for token in tokens}
    except ValueError as exc:
        valid = [component.value for component in InfoComponent]
        raise ValueError(
            f"Invalid SGLANG_DSPARK_DEBUG_DUMP token in {sorted(tokens)}; "
            f"valid: {valid} or '{ALL_COMPONENTS_TOKEN}'."
        ) from exc


class ReqDetail(msgspec.Struct, omit_defaults=True):
    req_pool_index: int
    prefix_len: int
    verify_len: int
    acc_len: int
    correct_drafts: int
    cap_trim: int
    bonus_token: int
    draft_tokens: list[int]
    rid: Optional[str] = None
    confidence: Optional[list[float]] = None
    survival: Optional[list[float]] = None


class DecodeStepRecord(msgspec.Struct, omit_defaults=True):
    forward_ct: int
    bs: int = -1
    mode: str = ""
    budget: Optional[int] = None
    lag_steps: Optional[int] = None
    num_running_reqs: int = -1
    num_verify_tokens: int = -1
    verify_tokens_local: int = -1
    verify_tokens_dp_synced: int = -1
    verify_tokens_graph_key: int = -1
    predicted_step_ms: Optional[float] = None
    predicted_theta: Optional[float] = None
    step_cpu_ms: Optional[float] = None
    step_gpu_ms: Optional[float] = None
    draft_gpu_ms: Optional[float] = None
    target_verify_gpu_ms: Optional[float] = None
    reqs: Optional[list[ReqDetail]] = None


class DecodeStepObservation(msgspec.Struct):
    forward_ct: int
    bs: int
    mode: str
    budget: Optional[int]
    lag_steps: Optional[int]
    num_verify_tokens: int
    verify_tokens_local: int
    verify_tokens_dp_synced: int
    verify_tokens_graph_key: int
    predicted_step_ms: Optional[float]
    predicted_theta: Optional[float]
    verify_lens: Optional[torch.Tensor]
    confidence: Optional[torch.Tensor]
    req_pool_indices: torch.Tensor
    prefix_lens: torch.Tensor
    draft_tokens: torch.Tensor
    bonus_tokens: torch.Tensor
    correct_len: torch.Tensor
    cap_trim_lens: torch.Tensor
    commit_lens: torch.Tensor
    rids: Optional[list[str]]


class _PendingStep(msgspec.Struct):
    forward_ct: int
    bs: int
    mode: str
    budget: Optional[int]
    lag_steps: Optional[int]
    num_verify_tokens: int
    verify_tokens_local: int
    verify_tokens_dp_synced: int
    verify_tokens_graph_key: int
    predicted_step_ms: Optional[float]
    predicted_theta: Optional[float]
    step_cpu_ms: Optional[float]
    rids: Optional[list[str]]
    future: Optional[FutureTensors]
    segment_events: dict[InfoSegment, tuple[torch.cuda.Event, torch.cuda.Event]]


class DsparkInfoDumper:
    def __init__(
        self,
        *,
        components: set[Union[InfoComponent, str]],
        gamma: int,
        verify_num_draft_tokens: int,
        attn_tp_rank: int,
        device: torch.device,
        mode_value: str,
        sps_report_interval: int = 0,
        max_records: int = INFO_DUMP_MAX_RECORDS,
        max_step_cpu_seconds: float = INFO_DUMP_MAX_STEP_CPU_SECONDS,
        clock: Callable[[], float] = time.monotonic,
    ) -> None:
        self.gamma = int(gamma)
        self.verify_num_draft_tokens = int(verify_num_draft_tokens)
        self.attn_tp_rank = int(attn_tp_rank)
        self.device = device
        self.mode_value = mode_value
        self._clock = clock
        self._max_step_cpu_seconds = max_step_cpu_seconds

        self._components: set[InfoComponent] = {
            InfoComponent(component) for component in components
        }
        self._sps_report_interval = int(sps_report_interval)
        if self._sps_report_interval > 0:
            self._components.add(InfoComponent.STEP_GPU_TIME)
        # Dedup within an attention-TP group only: records describe the
        # DP-rank-local batch, so under dp-attention every DP rank must keep
        # dumping (the SPS profiler reads one payload per DP rank).
        self.enabled = bool(self._components) and self.attn_tp_rank == 0
        self._sps_window: list[tuple[float, float]] = []
        self._sps_mismatched = 0

        self._records: deque[DecodeStepRecord] = deque(maxlen=max_records)
        self._pending: Optional[_PendingStep] = None
        self._prev_stamp: Optional[float] = None

        self._d2h_stream: Optional[torch.cuda.Stream] = None
        if self.enabled and InfoComponent.REQS in self._components:
            self._d2h_stream = torch.cuda.Stream(device=device)

        self._current_segments: dict[
            InfoSegment, tuple[torch.cuda.Event, torch.cuda.Event]
        ] = {}
        self._open_segments: dict[InfoSegment, torch.cuda.Event] = {}

    def begin_step(self) -> None:
        if not self.enabled:
            return
        self._current_segments = {}
        self._open_segments = {}
        if InfoComponent.STEP_GPU_TIME in self._components:
            self._open_segment(InfoSegment.STEP)

    def segment(self, name: Union[InfoSegment, str]) -> ContextManager[None]:
        if not self.enabled:
            return _NULL_SEGMENT
        segment = InfoSegment(name)
        if not self._segment_enabled(segment):
            return _NULL_SEGMENT
        return self._active_segment(segment)

    @contextmanager
    def _active_segment(self, segment: InfoSegment) -> Iterator[None]:
        self._open_segment(segment)
        try:
            yield
        finally:
            self._close_segment(segment)

    def observe_decode_step(self, obs: DecodeStepObservation) -> None:
        if not self.enabled:
            return
        if InfoComponent.STEP_GPU_TIME in self._components:
            self._close_segment(InfoSegment.STEP)

        now = self._clock()
        step_cpu_ms = self._step_cpu_ms(now=now)
        self._drain_pending()

        future = (
            self._stage_reqs(obs) if InfoComponent.REQS in self._components else None
        )
        self._pending = _PendingStep(
            forward_ct=int(obs.forward_ct),
            bs=int(obs.bs),
            mode=obs.mode,
            budget=None if obs.budget is None else int(obs.budget),
            lag_steps=None if obs.lag_steps is None else int(obs.lag_steps),
            num_verify_tokens=int(obs.num_verify_tokens),
            verify_tokens_local=int(obs.verify_tokens_local),
            verify_tokens_dp_synced=int(obs.verify_tokens_dp_synced),
            verify_tokens_graph_key=int(obs.verify_tokens_graph_key),
            predicted_step_ms=obs.predicted_step_ms,
            predicted_theta=obs.predicted_theta,
            step_cpu_ms=step_cpu_ms,
            rids=obs.rids,
            future=future,
            segment_events=self._current_segments,
        )
        self._current_segments = {}
        self._prev_stamp = now

    def note_non_decode_step(self) -> None:
        if not self.enabled:
            return
        self._drain_pending()
        self._prev_stamp = None
        self._current_segments = {}
        self._open_segments = {}

    def flush(self) -> None:
        if not self.enabled:
            return
        self._drain_pending()

    def clear(self) -> None:
        self._records.clear()
        self._pending = None
        self._prev_stamp = None
        self._current_segments = {}
        self._open_segments = {}
        self._sps_window = []
        self._sps_mismatched = 0

    def dump(self) -> Optional[dict]:
        if not self.enabled:
            return None
        self.flush()
        return {
            "mode": self.mode_value,
            "gamma": self.gamma,
            "verify_num_draft_tokens": self.verify_num_draft_tokens,
            "components": sorted(component.value for component in self._components),
            "records": [msgspec.to_builtins(record) for record in self._records],
        }

    def _segment_enabled(self, segment: InfoSegment) -> bool:
        if segment is InfoSegment.STEP:
            return InfoComponent.STEP_GPU_TIME in self._components
        if segment is InfoSegment.DRAFT:
            return InfoComponent.DRAFT_GPU_TIME in self._components
        if segment is InfoSegment.TARGET_VERIFY:
            return InfoComponent.TARGET_VERIFY_GPU_TIME in self._components
        return False

    def _open_segment(self, segment: InfoSegment) -> None:
        start = torch.cuda.Event(enable_timing=True)
        start.record()
        self._open_segments[segment] = start

    def _close_segment(self, segment: InfoSegment) -> None:
        start = self._open_segments.pop(segment, None)
        if start is None:
            return
        end = torch.cuda.Event(enable_timing=True)
        end.record()
        self._current_segments[segment] = (start, end)

    def _stage_reqs(self, obs: DecodeStepObservation) -> Optional[FutureTensors]:
        tensors: dict[str, torch.Tensor] = {
            "req_pool_indices": obs.req_pool_indices,
            "prefix_lens": obs.prefix_lens,
            "draft_tokens": obs.draft_tokens,
            "bonus_tokens": obs.bonus_tokens,
            "correct_len": obs.correct_len,
            "cap_trim_lens": obs.cap_trim_lens,
            "commit_lens": obs.commit_lens,
        }
        if obs.verify_lens is not None:
            tensors["verify_lens"] = obs.verify_lens
        if obs.confidence is not None:
            tensors["confidence"] = obs.confidence
        return FutureTensors.device_to_host(tensors, d2h_stream=self._d2h_stream)

    def _drain_pending(self) -> None:
        pending = self._pending
        self._pending = None
        if pending is None:
            return

        record = DecodeStepRecord(forward_ct=pending.forward_ct)
        if InfoComponent.CORE in self._components:
            record.bs = pending.bs
            record.mode = pending.mode
            record.budget = pending.budget
            record.lag_steps = pending.lag_steps
            record.num_running_reqs = pending.bs
            record.num_verify_tokens = pending.num_verify_tokens
            record.verify_tokens_local = pending.verify_tokens_local
            record.verify_tokens_dp_synced = pending.verify_tokens_dp_synced
            record.verify_tokens_graph_key = pending.verify_tokens_graph_key
            record.predicted_step_ms = pending.predicted_step_ms
            record.predicted_theta = pending.predicted_theta
        if InfoComponent.STEP_CPU_TIME in self._components:
            record.step_cpu_ms = pending.step_cpu_ms
        if InfoComponent.STEP_GPU_TIME in self._components:
            record.step_gpu_ms = self._segment_ms(pending, InfoSegment.STEP)
        if InfoComponent.DRAFT_GPU_TIME in self._components:
            record.draft_gpu_ms = self._segment_ms(pending, InfoSegment.DRAFT)
        if InfoComponent.TARGET_VERIFY_GPU_TIME in self._components:
            record.target_verify_gpu_ms = self._segment_ms(
                pending, InfoSegment.TARGET_VERIFY
            )
        if InfoComponent.REQS in self._components and pending.future is not None:
            record.reqs = self._build_reqs(
                host=pending.future.wait(), bs=pending.bs, rids=pending.rids
            )
        elif pending.future is not None:
            pending.future.wait()

        self._records.append(record)
        if self._sps_report_interval > 0:
            self._report_sps_prediction(pending=pending, step_gpu_ms=record.step_gpu_ms)

    def _report_sps_prediction(
        self, *, pending: _PendingStep, step_gpu_ms: Optional[float]
    ) -> None:
        predicted = pending.predicted_step_ms
        if predicted is None or step_gpu_ms is None:
            return
        matched = (
            pending.budget is not None
            and pending.bs + pending.budget == pending.num_verify_tokens
        )
        if not matched:
            self._sps_mismatched += 1
            return
        self._sps_window.append((predicted, step_gpu_ms))
        if len(self._sps_window) < self._sps_report_interval:
            return

        predictions = [p for p, _ in self._sps_window]
        actuals = [a for _, a in self._sps_window]
        abs_err = [abs(p - a) for p, a in self._sps_window]
        rel_err = [abs(p - a) / a * 100 for p, a in self._sps_window if a > 0]
        total = len(self._sps_window) + self._sps_mismatched
        logger.info(
            "DSpark SPS prediction: n=%d  mean predicted=%.3fms  mean actual=%.3fms  "
            "MAE=%.3fms  median rel-err=%.1f%%  mean bias(pred-actual)=%+.3fms  "
            "M_mismatch_rate=%.1f%% (%d/%d)",
            len(self._sps_window),
            statistics.fmean(predictions),
            statistics.fmean(actuals),
            statistics.fmean(abs_err),
            statistics.median(rel_err) if rel_err else float("nan"),
            statistics.fmean([p - a for p, a in self._sps_window]),
            self._sps_mismatched / total * 100 if total else 0.0,
            self._sps_mismatched,
            total,
        )
        self._sps_window = []
        self._sps_mismatched = 0

    def _step_cpu_ms(self, *, now: float) -> Optional[float]:
        prev = self._prev_stamp
        if prev is None:
            return None
        step_cpu = now - prev
        if not (0.0 < step_cpu <= self._max_step_cpu_seconds):
            return None
        return round(step_cpu * 1000.0, 4)

    def _segment_ms(
        self, pending: _PendingStep, segment: InfoSegment
    ) -> Optional[float]:
        events = pending.segment_events.get(segment)
        if events is None:
            return None
        start, end = events
        end.synchronize()
        elapsed_ms = start.elapsed_time(end)
        if elapsed_ms > self._max_step_cpu_seconds * 1000.0:
            return None
        return round(elapsed_ms, 4)

    def _build_reqs(
        self, *, host: dict, bs: int, rids: Optional[list[str]]
    ) -> list[ReqDetail]:
        req_ids = host["req_pool_indices"].tolist()
        prefixes = host["prefix_lens"].tolist()
        draft_rows = host["draft_tokens"].tolist()
        bonus = host["bonus_tokens"].tolist()
        correct = host["correct_len"].tolist()
        cap_trim = host["cap_trim_lens"].tolist()
        commit = host["commit_lens"].tolist()
        verify_lens = host["verify_lens"].tolist() if "verify_lens" in host else None
        if "confidence" in host:
            conf_host = host["confidence"].float()
            conf_rows = conf_host.tolist()
            survival_rows = torch.cumprod(conf_host, dim=1).tolist()
        else:
            conf_rows = None
            survival_rows = None

        reqs: list[ReqDetail] = []
        for row in range(bs):
            verify_len = (
                self.verify_num_draft_tokens
                if verify_lens is None
                else int(verify_lens[row])
            )
            reqs.append(
                ReqDetail(
                    rid=None if rids is None else rids[row],
                    req_pool_index=int(req_ids[row]),
                    prefix_len=int(prefixes[row]),
                    verify_len=verify_len,
                    acc_len=int(commit[row]),
                    correct_drafts=int(correct[row]),
                    cap_trim=int(cap_trim[row]),
                    bonus_token=int(bonus[row]),
                    draft_tokens=[int(t) for t in draft_rows[row]],
                    confidence=(
                        None
                        if conf_rows is None
                        else [round(float(p), 4) for p in conf_rows[row]]
                    ),
                    survival=(
                        None
                        if survival_rows is None
                        else [round(float(p), 4) for p in survival_rows[row]]
                    ),
                )
            )
        return reqs


EPS_PROB = 1e-8


def _format_float(value: float, digits: int = 4) -> str:
    value = float(value)
    if math.isnan(value):
        return "nan"
    return f"{value:.{digits}f}"


class PerPositionConfidenceMetrics:

    def __init__(
        self,
        *,
        gamma: int,
        device: torch.device,
        num_coarse_bins: int = 15,
        num_fine_bins: int = 1024,
    ) -> None:
        self.gamma = int(gamma)
        self.num_coarse_bins = int(num_coarse_bins)
        self.num_fine_bins = int(num_fine_bins)
        self.coarse_count = torch.zeros(
            (self.gamma, self.num_coarse_bins), dtype=torch.float64, device=device
        )
        self.coarse_pred = torch.zeros_like(self.coarse_count)
        self.coarse_target = torch.zeros_like(self.coarse_count)
        self.fine_pos = torch.zeros(
            (self.gamma, self.num_fine_bins), dtype=torch.float64, device=device
        )
        self.fine_neg = torch.zeros_like(self.fine_pos)
        self.brier_num = torch.zeros(self.gamma, dtype=torch.float64, device=device)

    def update(self, *, survival: torch.Tensor, prefix_mask: torch.Tensor) -> None:
        assert survival.shape == prefix_mask.shape
        assert survival.dim() == 2 and survival.shape[1] == self.gamma

        probs = survival.to(torch.float64).clamp(EPS_PROB, 1.0 - EPS_PROB)
        targets = prefix_mask.to(torch.float64)
        bs = probs.shape[0]

        probs_flat = probs.reshape(-1)
        targets_flat = targets.reshape(-1)
        weights = torch.ones_like(probs_flat)
        pos_idx = (
            torch.arange(self.gamma, device=probs.device)
            .view(1, -1)
            .expand(bs, self.gamma)
            .reshape(-1)
        )

        coarse_idx = (
            (probs_flat * self.num_coarse_bins)
            .long()
            .clamp_(0, self.num_coarse_bins - 1)
        )
        flat_coarse = pos_idx * self.num_coarse_bins + coarse_idx
        self.coarse_count.view(-1).scatter_add_(0, flat_coarse, weights)
        self.coarse_pred.view(-1).scatter_add_(0, flat_coarse, probs_flat)
        self.coarse_target.view(-1).scatter_add_(0, flat_coarse, targets_flat)

        fine_idx = (
            (probs_flat * self.num_fine_bins).long().clamp_(0, self.num_fine_bins - 1)
        )
        flat_fine = pos_idx * self.num_fine_bins + fine_idx
        self.fine_pos.view(-1).scatter_add_(0, flat_fine, targets_flat)
        self.fine_neg.view(-1).scatter_add_(0, flat_fine, 1.0 - targets_flat)

        self.brier_num.add_((probs - targets).pow(2).sum(dim=0))

    @staticmethod
    def _auroc_from_hist(pos_hist: torch.Tensor, neg_hist: torch.Tensor) -> float:
        total_pos = float(pos_hist.sum())
        total_neg = float(neg_hist.sum())
        if total_pos <= 0.0 or total_neg <= 0.0:
            return float("nan")
        cum_neg = torch.cumsum(neg_hist, dim=0)
        cum_neg_before = cum_neg - neg_hist
        pair = (pos_hist * cum_neg_before).sum() + 0.5 * (pos_hist * neg_hist).sum()
        return float(pair) / (total_pos * total_neg)

    def compute(self) -> list[dict]:
        coarse_count = self.coarse_count.cpu()
        coarse_pred = self.coarse_pred.cpu()
        coarse_target = self.coarse_target.cpu()
        fine_pos = self.fine_pos.cpu()
        fine_neg = self.fine_neg.cpu()
        brier_num = self.brier_num.cpu()

        out: list[dict] = []
        for pos in range(self.gamma):
            weights = coarse_count[pos]
            total = float(weights.sum())
            if total <= 1e-12:
                out.append(
                    {
                        "position": pos,
                        "total_weight": 0.0,
                        "ece": float("nan"),
                        "auc": float("nan"),
                        "brier": float("nan"),
                        "pred_mean": float("nan"),
                        "target_mean": float("nan"),
                        "reliability": [],
                    }
                )
                continue

            denom = weights.clamp_min(1e-12)
            avg_pred = coarse_pred[pos] / denom
            avg_target = coarse_target[pos] / denom
            bin_err = (avg_pred - avg_target).abs()
            ece = float((bin_err * weights).sum()) / total
            auc = self._auroc_from_hist(fine_pos[pos], fine_neg[pos])
            brier = float(brier_num[pos]) / total
            reliability = []
            for bin_idx in range(self.num_coarse_bins):
                weight = float(weights[bin_idx])
                if weight <= 0.0:
                    continue
                reliability.append(
                    {
                        "bin": bin_idx,
                        "range": [
                            bin_idx / self.num_coarse_bins,
                            (bin_idx + 1) / self.num_coarse_bins,
                        ],
                        "avg_pred": float(avg_pred[bin_idx]),
                        "avg_target": float(avg_target[bin_idx]),
                        "weight": weight,
                    }
                )
            out.append(
                {
                    "position": pos,
                    "total_weight": total,
                    "ece": ece,
                    "auc": auc,
                    "brier": brier,
                    "pred_mean": float(coarse_pred[pos].sum()) / total,
                    "target_mean": float(coarse_target[pos].sum()) / total,
                    "reliability": reliability,
                }
            )
        return out

    def format_table(self) -> str:
        rows = self.compute()
        header = (
            f"{'pos':>3} {'count':>12} {'pred':>8} {'target':>8} "
            f"{'ece':>8} {'auc':>8} {'brier':>8}"
        )
        lines = [
            "DSpark confidence-head per-position calibration "
            "(cumprod survival vs leading-correct-prefix)",
            header,
        ]
        for row in rows:
            lines.append(
                f"{row['position']:>3} {row['total_weight']:>12.0f} "
                f"{_format_float(row['pred_mean']):>8} "
                f"{_format_float(row['target_mean']):>8} "
                f"{_format_float(row['ece']):>8} "
                f"{_format_float(row['auc']):>8} "
                f"{_format_float(row['brier']):>8}"
            )
        return "\n".join(lines)


class ConfidenceMetricsProbe:

    def __init__(
        self,
        *,
        gamma: int,
        verify_num_draft_tokens: int,
        tp_rank: int,
        print_every: int = 256,
    ) -> None:
        self.gamma = int(gamma)
        self.verify_num_draft_tokens = int(verify_num_draft_tokens)
        self.tp_rank = int(tp_rank)
        self.print_every = int(print_every)
        self._metrics: Optional[PerPositionConfidenceMetrics] = None
        self._step_ct: int = 0
        self._compact_warned: bool = False

    def maybe_observe(
        self,
        *,
        carries_confidence: bool,
        is_compact_mode: bool,
        confidence_raw: Optional[torch.Tensor],
        verify_ids_2d: torch.Tensor,
        target_logits: torch.Tensor,
        bs: int,
    ) -> None:
        if not envs.SGLANG_DSPARK_DEBUG_CONFIDENCE_METRICS.get():
            return
        if self.tp_rank != 0:
            return
        if not carries_confidence:
            return
        if is_compact_mode:
            if not self._compact_warned:
                logger.warning(
                    "SGLANG_DSPARK_DEBUG_CONFIDENCE_METRICS is ignored under "
                    "SGLANG_RAGGED_VERIFY_MODE=compact (padded verify rows corrupt the "
                    "per-position prefix label); run cap-accept or static to measure it."
                )
                self._compact_warned = True
            return
        if confidence_raw is None:
            return

        target_predict = torch.argmax(target_logits, dim=-1).view(
            bs, self.verify_num_draft_tokens
        )
        num_correct_drafts, _ = compute_dflash_correct_drafts_and_bonus(
            candidates=verify_ids_2d,
            target_predict=target_predict,
        )
        positions = torch.arange(self.gamma, device=confidence_raw.device).view(1, -1)
        prefix_mask = (positions < num_correct_drafts.view(-1, 1)).to(torch.float32)
        survival = torch.cumprod(torch.sigmoid(confidence_raw.float()), dim=1)

        if self._metrics is None:
            self._metrics = PerPositionConfidenceMetrics(
                gamma=self.gamma, device=confidence_raw.device
            )
        self._metrics.update(survival=survival, prefix_mask=prefix_mask)
        self._step_ct += 1
        if self._step_ct % self.print_every == 0:
            logger.info("%s", self._metrics.format_table())


_STS_COLLECT_FLUSH_EVERY: int = 256


class DsparkStepObservers:
    """Facade over the per-step observability sinks (info dumper, confidence
    probe, STS collection, block-accept estimator). The worker's decode path
    makes one call per step; all sink gating and field derivation live here
    so the hot path stays free of observer plumbing."""

    def __init__(
        self,
        *,
        planner,
        gamma: int,
        verify_num_draft_tokens: int,
        tp_rank: int,
        device,
        simulate_acc_len: float,
    ) -> None:
        self._planner = planner
        self._gamma = int(gamma)
        self._verify_num_draft_tokens = int(verify_num_draft_tokens)
        self._simulate_acc_len = float(simulate_acc_len)

        self._confidence_probe = ConfidenceMetricsProbe(
            gamma=gamma,
            verify_num_draft_tokens=verify_num_draft_tokens,
            tp_rank=tp_rank,
        )
        self._info_dumper = DsparkInfoDumper(
            components=resolve_enabled_components(),
            gamma=gamma,
            verify_num_draft_tokens=verify_num_draft_tokens,
            attn_tp_rank=get_parallel().attn_tp_rank,
            device=device,
            mode_value=planner.mode_value,
            sps_report_interval=envs.SGLANG_DSPARK_LOG_SPS_PRED_INTERVAL.get(),
        )
        self._block_accept_recorder = create_block_accept_estimate_recorder(
            gamma=gamma, device=device, tp_rank=tp_rank
        )
        if self._simulate_acc_len > 0 and self._block_accept_recorder is not None:
            raise ValueError(
                "SGLANG_DSPARK_BLOCK_ACCEPT_ESTIMATE_PATH cannot be combined with "
                "SGLANG_SIMULATE_ACC_LEN (simulated correct_len breaks the "
                "accept-probability bookkeeping of the estimator)."
            )
        self._sts_collect_path = envs.SGLANG_DSPARK_STS_COLLECT_PATH.get()
        self._sts_recorder: Optional[StsDataRecorder] = None

    # --- step lifecycle -------------------------------------------------

    def begin_step(self) -> None:
        self._info_dumper.begin_step()

    def segment(self, name: Union[InfoSegment, str]) -> ContextManager[None]:
        return self._info_dumper.segment(name)

    def note_prefill_step(self) -> None:
        self._info_dumper.note_non_decode_step()
        if self._block_accept_recorder is not None:
            self._block_accept_recorder.flush()

    def note_idle_decode_step(self) -> None:
        self._info_dumper.note_non_decode_step()

    # --- scheduler-facing hooks ------------------------------------------

    def dump_info_records(self) -> Optional[dict]:
        dumped = self._info_dumper.dump()
        if dumped is None:
            return None
        dumped["simulate_acc_len"] = (
            self._simulate_acc_len if self._simulate_acc_len > 0 else None
        )
        return dumped

    def clear_info_records(self) -> None:
        self._info_dumper.clear()

    def block_accept_estimate_log_suffix(self) -> Optional[str]:
        if self._block_accept_recorder is None:
            return None
        return self._block_accept_recorder.estimate_log_suffix()

    def note_request_finished(self, *, rid: str, natural_stop: bool) -> None:
        if self._block_accept_recorder is None:
            return
        self._block_accept_recorder.note_request_finished(
            rid=rid, natural_stop=natural_stop
        )

    # --- per-step observation --------------------------------------------

    def observe_verify_step(
        self,
        *,
        forward_ct: int,
        reqs,
        bs: int,
        proposal_folded: bool,
        verify_ids_2d: torch.Tensor,
        target_logits: Optional[torch.Tensor],
        layout,
        confidence: Optional[torch.Tensor],
        prefix_lens: torch.Tensor,
        draft_tokens: torch.Tensor,
        draft_block,
        sampling_info,
        correct_len: torch.Tensor,
        cap_trim_lens: torch.Tensor,
        bonus: torch.Tensor,
        commit_lens: torch.Tensor,
        verify_token_budget: Optional[int],
        req_pool_indices: torch.Tensor,
        verify_tier_num_tokens: int,
        dp_tier_num_tokens: Optional[int],
    ) -> None:
        planner = self._planner
        if not proposal_folded:
            self._maybe_record_sts_collect(
                verify_ids_2d=verify_ids_2d,
                target_logits=target_logits,
                bs=bs,
            )
            self._confidence_probe.maybe_observe(
                carries_confidence=planner.carries_confidence,
                is_compact_mode=planner.is_compact_mode,
                confidence_raw=planner.last_confidence_raw,
                verify_ids_2d=verify_ids_2d,
                target_logits=target_logits,
                bs=bs,
            )
        if self._block_accept_recorder is not None and not proposal_folded:
            self._block_accept_recorder.observe_verify_step(
                forward_ct=forward_ct,
                rids=[req.rid for req in reqs],
                draft_tokens=draft_tokens,
                corrected_logits=draft_block.corrected_logits,
                draft_temperatures=draft_block.temperatures,
                greedy_mask=draft_block.greedy_mask,
                target_logits=target_logits,
                target_temperatures=(
                    sampling_info.temperatures
                    if sampling_info is not None
                    else draft_block.temperatures
                ),
                truncated_sampling_mask=(
                    (sampling_info.top_ks != TOP_K_ALL)
                    | (sampling_info.top_ps != 1.0)
                    | (sampling_info.min_ps > 0)
                    if sampling_info is not None
                    else None
                ),
                logits_adjustments_are_noop=verify_logits_adjustments_are_noop(
                    sampling_info
                ),
                correct_len=correct_len,
                cap_trim_lens=cap_trim_lens,
                bonus=bonus,
                prefix_lens=prefix_lens,
                layout=layout,
            )
        if self._info_dumper.enabled:
            budget_decision = planner.take_budget_decision()
            predicted_step_ms = (
                None
                if budget_decision is None
                or budget_decision.predicted_step_seconds is None
                else budget_decision.predicted_step_seconds * 1e3
            )
            predicted_theta = (
                None if budget_decision is None else budget_decision.predicted_theta
            )
            num_verify_tokens = (
                layout.graph_num_tokens
                if layout is not None
                else int(verify_ids_2d.numel())
            )
            self._info_dumper.observe_decode_step(
                DecodeStepObservation(
                    forward_ct=forward_ct,
                    bs=bs,
                    mode=planner.mode_value,
                    budget=verify_token_budget,
                    lag_steps=planner.lag_steps,
                    num_verify_tokens=num_verify_tokens,
                    verify_tokens_local=verify_tier_num_tokens,
                    verify_tokens_dp_synced=(
                        -1 if dp_tier_num_tokens is None else int(dp_tier_num_tokens)
                    ),
                    verify_tokens_graph_key=num_verify_tokens,
                    predicted_step_ms=predicted_step_ms,
                    predicted_theta=predicted_theta,
                    verify_lens=layout.verify_lens if layout is not None else None,
                    confidence=confidence,
                    req_pool_indices=req_pool_indices,
                    prefix_lens=prefix_lens,
                    draft_tokens=draft_tokens,
                    bonus_tokens=bonus,
                    correct_len=correct_len,
                    cap_trim_lens=cap_trim_lens,
                    commit_lens=commit_lens,
                    rids=[req.rid for req in reqs],
                )
            )

    def _maybe_record_sts_collect(
        self,
        *,
        verify_ids_2d: torch.Tensor,
        target_logits: Optional[torch.Tensor],
        bs: int,
    ) -> None:
        if not self._sts_collect_path:
            return
        if not self._planner.carries_confidence:
            return
        confidence_raw = self._planner.last_confidence_raw
        if confidence_raw is None:
            return
        if self._sts_recorder is None:
            self._sts_recorder = StsDataRecorder(
                path_stem=self._sts_collect_path,
                gamma=self._gamma,
                flush_every=_STS_COLLECT_FLUSH_EVERY,
            )
        target_predict = torch.argmax(target_logits, dim=-1).view(
            bs, self._verify_num_draft_tokens
        )
        num_correct_drafts, _ = compute_dflash_correct_drafts_and_bonus(
            candidates=verify_ids_2d,
            target_predict=target_predict,
        )
        self._sts_recorder.record(
            confidence_raw=confidence_raw,
            num_correct_drafts=num_correct_drafts,
        )