class AdaptiveVerificationManager:
def __init__(
self,
req_states: "RequestState",
query_start_loc: torch.Tensor,
num_bonus_tokens: int,
max_total_logits: int,
):
self.req_states = req_states
self.num_speculative_steps = req_states.num_speculative_steps
device = req_states.device
self._copy_stream = torch.cuda.Stream(device)
self.num_bonus_tokens = num_bonus_tokens
# Rejection sampling verifies logits in one contiguous chunk; the
# chunked path indexes by scheduled (untrimmed) offsets and cannot
# address the compacted layout, so the budget must fit one chunk.
self._max_total_logits = max_total_logits
self.query_start_loc = query_start_loc
self.cost_tables: tuple[np.ndarray, np.ndarray] | None = None
# Largest cudagraph-captured token count; above it nothing pads.
self._cudagraph_limit = 0
self._batch_budget: tuple[dict[str, int], dict[str, int], int] | None = None
max_num_reqs = req_states.max_num_reqs
# Current per-slot confidences
self._confidence_probs = torch.empty(
(max_num_reqs, self.num_speculative_steps),
dtype=torch.float32,
device=device,
)
self._batch_draft_capacity = torch.empty(
max_num_reqs, dtype=torch.int32, device=device
)
self._num_non_draft_tokens = torch.empty_like(query_start_loc[:-1])
self._cu_num_logits = torch.empty_like(query_start_loc)
# Two D2H slots preserve stale inputs for budget selection.
self._stale_confidences = [
CpuGpuBuffer(
max_num_reqs,
self.num_speculative_steps,
dtype=torch.float32,
device=device,
)
for _ in range(2)
]
self._copy_events = [torch.cuda.Event(blocking=True) for _ in range(2)]
self._pending_resets: list[int] = []
self._stale_idx = 0
for slot in self._stale_confidences:
slot.np.fill(1.0)
def add_request(self, req_idx: int) -> None:
self._stale_confidences[self._stale_idx].np[req_idx].fill(1.0)
self._pending_resets.append(req_idx)
self._confidence_probs[req_idx].fill_(1.0)
def batches_to_profile(self, capture_sizes: list[int]) -> Iterator[dict[str, int]]:
"""Dummy-run kwargs whose step timings seed the cost tables.
Run these inside StepTimingCollector.collect(), then hand the block's
timings to set_initial_cost_curves."""
max_num_tokens = self.req_states.max_num_batched_tokens
size = self._cudagraph_limit = capture_sizes[-1] if capture_sizes else 0
# Also profile beyond the capture limit: real steps run there
# (piecewise/eager) and linear extrapolation badly underestimates
# them. These runs double as JIT warmup for the piecewise shapes.
tail_sizes: set[int] = set()
if size:
tail_sizes.add(min(size + size // 2, max_num_tokens))
while size < max_num_tokens:
size = min(size * 2, max_num_tokens)
tail_sizes.add(size)
tail_sizes -= set(capture_sizes)
for num_tokens in capture_sizes + sorted(tail_sizes):
for _ in range(_PROFILE_REPLAYS):
yield {
"num_tokens": num_tokens,
"context_len": envs.VLLM_ADAPTIVE_VERIFICATION_PROFILE_CONTEXT_LEN,
}
def set_initial_cost_curves(self, samples: list[StepTimingSample]) -> None:
def median_curve(
points: Iterable[tuple[int, float]],
) -> list[tuple[int, float]]:
grouped: defaultdict[int, list[float]] = defaultdict(list)
for key, value in points:
grouped[key].append(value)
return [(k, float(np.median(v))) for k, v in sorted(grouped.items())]
# Draft curve: eager-target steps inflate drafter timings (the CPU is
# still launching kernels, opening gaps between the drafter's events),
# and request counts — unlike token counts — collide across execution
# modes, so only graph-replay samples may price the draft curve.
draft_curve = median_curve(
(s.num_reqs, s.drafter_ms) for s in samples if s.full_cudagraph
)
verify_curve = median_curve(
(s.num_target_tokens, s.forward_ms) for s in samples
)
self.set_cost_curves(draft_curve, verify_curve)
def set_cost_curves(
self,
draft_curve: list[tuple[int, float]],
verify_curve: list[tuple[int, float]],
) -> None:
draft_curve, verify_curve = get_tp_group().broadcast_object(
(draft_curve, verify_curve), src=0
)
if not draft_curve or not verify_curve:
raise RuntimeError(
"Adaptive verification could not profile step costs. Pass "
"`enable_adaptive_verification=false` in the speculative config to "
"verify a fixed number of drafts instead."
)
self.cost_tables = build_cost_tables_from_curves(
draft_curve,
verify_curve,
self.req_states.max_num_reqs,
self.req_states.max_num_batched_tokens,
self._cudagraph_limit,
)
logger.debug("DSpark cost tables: %s", self.cost_tables)
def record_confidences(
self,
confidence_probs: torch.Tensor,
input_batch: "InputBatch",
) -> None:
"""Publish this step's raw confidences for the ranking kernel and start
copying them to the CPU, where a later step's budget reads them."""
num_reqs = input_batch.num_reqs
ready_idx = self._stale_idx ^ 1
with gpu_sync_allowed():
self._copy_events[ready_idx].synchronize()
if self._pending_resets:
self._stale_confidences[ready_idx].np[self._pending_resets] = 1.0
self._pending_resets.clear()
# Last step's copy has landed: budgets read it, this step overwrites the
# slot they were reading before.
self._stale_idx, write_idx = ready_idx, self._stale_idx
self._confidence_probs[input_batch.idx_mapping] = confidence_probs[:num_reqs]
write_slot = self._stale_confidences[write_idx]
write_slot.gpu.copy_(self._confidence_probs)
current_stream = torch.cuda.current_stream(self.req_states.device)
self._copy_stream.wait_stream(current_stream)
with stream(self._copy_stream, current_stream):
write_slot.copy_to_cpu()
self._copy_events[write_idx].record()
def get_num_tokens(
self,
num_tokens_per_req: dict[str, int],
draft_tokens: dict[str, list[int]],
) -> int:
"""Token count once the draft budget is trimmed to fit.
Stashes the chosen budget in ``_batch_budget`` for the compaction and
reallocation that follow in the same step.
"""
assert self.cost_tables is not None
req_ids = list(num_tokens_per_req)
num_reqs = len(req_ids)
scheduled_tokens = np.fromiter(
num_tokens_per_req.values(), dtype=np.int32, count=num_reqs
)
scheduled_drafts = np.fromiter(
(len(draft_tokens.get(req_id, ())) for req_id in req_ids),
dtype=np.int32,
count=num_reqs,
)
num_non_draft_tokens = scheduled_tokens - scheduled_drafts
slots = np.fromiter(
(self.req_states.req_id_to_index[req_id] for req_id in req_ids),
dtype=np.int32,
count=len(req_ids),
)
stale_confidences = self._stale_confidences[self._stale_idx].np[slots]
survival_probability = np.cumprod(stale_confidences.astype(np.float64), axis=1)
steps = np.arange(self.num_speculative_steps)
valid = steps[None, :] < scheduled_drafts[:, None]
scores = np.sort(survival_probability[valid])[::-1]
num_non_draft_tokens_total = int(num_non_draft_tokens.sum())
max_draft_budget = min(
int(scheduled_drafts.sum()),
max(0, self._max_total_logits - num_reqs * self.num_bonus_tokens),
)
scores = scores[:max_draft_budget]
draft_cost_ms, verify_cost_ms = self.cost_tables
num_sampling_requests = np.count_nonzero(
self.req_states.num_computed_tokens_np[slots] + num_non_draft_tokens
>= self.req_states.prefill_len.np[slots]
)
num_tokens_to_estimated_accepted_tokens = np.concatenate(
([num_sampling_requests], num_sampling_requests + np.cumsum(scores))
)
costs = (
draft_cost_ms[len(req_ids)]
+ verify_cost_ms[
num_non_draft_tokens_total : num_non_draft_tokens_total
+ max_draft_budget
+ 1
]
)
num_drafts_per_req = {
req_id: int(num_drafts)
for req_id, num_drafts in zip(req_ids, scheduled_drafts, strict=True)
}
num_non_draft_tokens_per_req = {
req_id: int(num_tokens)
for req_id, num_tokens in zip(req_ids, num_non_draft_tokens, strict=True)
}
draft_budget = int(np.argmax(num_tokens_to_estimated_accepted_tokens / costs))
self._batch_budget = (
num_drafts_per_req,
num_non_draft_tokens_per_req,
draft_budget,
)
return sum(num_non_draft_tokens_per_req.values()) + draft_budget
def compact_batch(
self,
num_draft_tokens_per_req: np.ndarray,
num_scheduled_tokens: np.ndarray,
cu_num_logits_np: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
"""Compact the CPU batch to the chosen draft budget.
Returns the compacted per-request token counts and the CPU cu_num_logits_np.
If the draft budget is 0, we can know cu_num_logits_np exactly, otherwise
its unchanged/an-upper-bound.
"""
batch_budget = self._batch_budget
assert batch_budget is not None
_, _, draft_budget = batch_budget
num_drafts = int(num_draft_tokens_per_req.sum())
if draft_budget == num_drafts:
return num_scheduled_tokens, cu_num_logits_np
num_non_draft_tokens = num_scheduled_tokens - num_draft_tokens_per_req
if draft_budget == 0:
# The draft budget is 0, so we can know cu_num_logits_np exactly. This helps
# when we would exceed the sampler logit chunk size.
num_reqs = num_scheduled_tokens.shape[0]
cu_num_logits_np = (
np.arange(num_reqs + 1, dtype=cu_num_logits_np.dtype)
* self.num_bonus_tokens
)
return num_non_draft_tokens, cu_num_logits_np
is_verification_request = num_draft_tokens_per_req > 0
num_verification_reqs = int(is_verification_request.sum())
# sort_batch_req_ids keeps verification requests at the front.
assert np.all(is_verification_request[:num_verification_reqs])
# for the CPU side buffer we distribute draft tokens evenly
draft_lens_cpu = np.zeros_like(num_non_draft_tokens)
draft_lens_cpu[:num_verification_reqs] = draft_budget // num_verification_reqs
draft_lens_cpu[: draft_budget % num_verification_reqs] += 1
return num_non_draft_tokens + draft_lens_cpu, cu_num_logits_np
def reallocate_drafts(
self, req_ids: list[str], idx_mapping: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, int]:
batch_budget, self._batch_budget = self._batch_budget, None
assert batch_budget is not None
num_drafts_per_req, num_non_draft_tokens_per_req, draft_budget = batch_budget
num_reqs = idx_mapping.shape[0]
scheduled_drafts = np.fromiter(
(num_drafts_per_req[req_id] for req_id in req_ids),
dtype=np.int32,
count=num_reqs,
)
num_non_draft_tokens = np.fromiter(
(num_non_draft_tokens_per_req[req_id] for req_id in req_ids),
dtype=np.int32,
count=num_reqs,
)
num_tokens = int(num_non_draft_tokens.sum()) + draft_budget
# Rank draft slots by survival probability and admit the best prefix.
# capacities enters holding each request's valid draft count (the kernel
# uses it to bound eligible slots) and leaves holding the admitted count.
capacities = self._batch_draft_capacity[:num_reqs]
if draft_budget == 0:
capacities.zero_()
else:
async_copy_to_gpu(scheduled_drafts, out=capacities)
if draft_budget < int(scheduled_drafts.sum()):
_assign_draft_token_budget_compiled(
self._confidence_probs,
idx_mapping,
capacities,
draft_budget,
self.num_speculative_steps,
)
num_non_draft_tokens_gpu = self._num_non_draft_tokens[:num_reqs]
async_copy_to_gpu(
num_non_draft_tokens,
out=num_non_draft_tokens_gpu,
)
self._cu_num_logits[:1].zero_()
torch.cumsum(
capacities + self.num_bonus_tokens,
dim=0,
out=self._cu_num_logits[1 : num_reqs + 1],
)
self.query_start_loc[:1].zero_()
torch.cumsum(
capacities + num_non_draft_tokens_gpu,
dim=0,
out=self.query_start_loc[1 : num_reqs + 1],
)
self.query_start_loc[num_reqs + 1 :].fill_(num_tokens)
return (
self._cu_num_logits[: num_reqs + 1],
self.query_start_loc,
draft_budget,
)