@tduncan / Nvidia-Optimization-Stack-VLLM / commits / d5cadcee86

[Feature][DSpark]: Logprobs adaptive verification (#52242)

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com> Co-authored-by: Benjamin Chislett <bchislett@nvidia.com>

Naveenraj Kamalakannan committed Aug 25, 2026 at 10:22 UTC d5cadcee8641a9fcec15facb5a9157d157daa207

10 files changed +61 -25

docs/features/speculative_decoding/adaptive_verification.md

+1 -1

@@ -36,7 +36,7 @@ Set `enable_adaptive_verification: false` to verify the full block for every req

37 - The attention backend must tolerate device-decided query lengths, since the CPU lengths only bound them from above. Backends that plan off the CPU lengths are excluded by the attention selector, and rejected at startup for models that hard-wire their backend.

38 - Full cudagraphs are required: step costs are profiled from captured graphs, so `--enforce-eager` is rejected at startup.

39 -- Not supported with LoRA (the per-token LoRA mapping is built from CPU-side boundaries), pipeline parallelism (cost curves and confidences exist only on the last rank), or output logprobs (to be fixed).

39 +- Not supported with LoRA (the per-token LoRA mapping is built from CPU-side boundaries) or pipeline parallelism (cost curves and confidences exist only on the last rank).

41 ## Tuning the cost profile

tests/v1/engine/test_output_processor.py

+1 -1

@@ -411,7 +411,7 @@ def _validate_logprobs(

411 ref_prompt_logprob_toks,

412 ref_prompt_logprob_vals,

413 ref_prompt_token_ranks,

414 - _,

414 + *_,

415 ) = ref_prompt_logprobs

416 for idx, (prompt_token, pos_logprob_dict) in enumerate(

417 zip(prompt_token_ids[1:], prompt_logprobs[1:])

tests/v1/test_outputs.py

+19

@@ -42,6 +42,25 @@ def test_logprobs_tensors_cat():

42 assert LogprobsTensors.cat([first]) is first

45 +def test_logprobs_tensors_tolists_with_tensor_boundaries():

46 + """Adaptive verification hands over the request boundaries as a tensor

47 + (they only exist on device); tolists() must materialize it as a plain

48 + list so slice_request splits requests correctly."""

49 + tensors = LogprobsTensors(

50 + torch.tensor([[1, 2], [3, 4], [5, 6]]),

51 + torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]),

52 + torch.tensor([1, 2, 3]),

53 + cu_num_generated_tokens_tensor=torch.tensor([0, 2, 3], dtype=torch.int32),

54 + )

55 +

56 + lists = tensors.to_cpu_nonblocking().tolists()

57 +

58 + assert lists.cu_num_generated_tokens == [0, 2, 3]

59 + sliced = lists.slice_request(1, 1)

60 + assert sliced.logprob_token_ids.tolist() == [[5, 6]]

61 + assert sliced.sampled_token_ranks.tolist() == [3]

62 +

63 +

64 def test_sampling_mask_tensors_tolist():

65 tensors = SamplingMaskTensors(

66 packed_mask=torch.tensor(

tests/v1/worker/test_gpu_rejection_sampler_chunking.py

@@ -51,6 +51,7 @@ def test_chunked_scores_match_full_batch(logprobs_mode: str):

51 rejection_sampler = object.__new__(RejectionSampler)

52 rejection_sampler.sampler = SimpleNamespace(logprobs_mode=logprobs_mode)

53 rejection_sampler.num_speculative_steps = 3

54 + rejection_sampler.enable_adaptive_verification = False

56 def fake_verify(

57 self,

vllm/sampling_params.py

-14

@@ -974,20 +974,6 @@ class SamplingParams(

974 if speculative_config is None:

975 return

976

977 - # Adaptive verification compacts logits after the forward pass, while

978 - # compute_topk_scores uses the scheduled layout in cu_num_logits_np.

979 - # TODO(lucas): lift this restriction. The true boundaries exist on device as

980 - # cu_num_logits; cu_num_generated_tokens is only read host-side after

981 - # the logprobs D2H, so it could ride along on that copy.

982 - if (

983 - speculative_config.enable_adaptive_verification

984 - and self.num_logprobs is not None

985 - ):

986 - raise ValueError(

987 - "Output logprobs are not supported with DSpark confidence-based "

988 - "verification."

989 - )

990 -

977 # Some sampling parameters are not yet compatible with spec decoding.

978 if self.min_p > _SAMPLING_EPS or self.logit_bias:

979 raise VLLMValidationError(

vllm/v1/engine/logprobs.py

+1 -1

@@ -134,7 +134,7 @@ class LogprobsProcessor:

134 assert self.num_prompt_logprobs is not None

135 assert self.prompt_logprobs is not None

136

137 - token_ids, logprobs, ranks, _ = prompt_logprobs_tensors

137 + token_ids, logprobs, ranks, *_ = prompt_logprobs_tensors

138

139 # Recover shapes.

140 num_prompt_tokens, num_logprobs = logprobs.shape

vllm/v1/outputs.py

+21 -4

@@ -98,27 +98,38 @@ class LogprobsTensors(NamedTuple):

98 logprobs: torch.Tensor

99 # [num_reqs x num_generated_tokens]

100 selected_token_ranks: torch.Tensor

101 - # [num_reqs]

101 + # [num_reqs + 1]

102 cu_num_generated_tokens: list[int] | None = None

103 + # [num_reqs + 1]. Set instead of cu_num_generated_tokens when the

104 + # boundaries only exist on device (adaptive verification); rides along on

105 + # the async D2H copy.

106 + cu_num_generated_tokens_tensor: torch.Tensor | None = None

107

108 def tolists(self, cu_num_generated_tokens: list[int] | None = None):

109 + if cu_num_generated_tokens is None:

110 + if self.cu_num_generated_tokens_tensor is not None:

111 + cu_num_generated_tokens = self.cu_num_generated_tokens_tensor.tolist()

112 + else:

113 + cu_num_generated_tokens = self.cu_num_generated_tokens

114 return LogprobsLists(

115 self.logprob_token_ids.cpu().numpy(),

116 self.logprobs.cpu().numpy(),

117 self.selected_token_ranks.cpu().numpy(),

109 - cu_num_generated_tokens

110 - if cu_num_generated_tokens is not None

111 - else self.cu_num_generated_tokens,

118 + cu_num_generated_tokens,

119 )

120

121 def to_cpu_nonblocking(self) -> "LogprobsTensors":

122 if self.logprob_token_ids.device.type == "cpu":

123 return self

124 + cu_tensor = self.cu_num_generated_tokens_tensor

125 + if cu_tensor is not None:

126 + cu_tensor = cu_tensor.to("cpu", non_blocking=True)

127 return LogprobsTensors(

128 self.logprob_token_ids.to("cpu", non_blocking=True),

129 self.logprobs.to("cpu", non_blocking=True),

130 self.selected_token_ranks.to("cpu", non_blocking=True),

131 self.cu_num_generated_tokens,

132 + cu_tensor,

133 )

134

135 def filter(self, mask: torch.Tensor) -> "LogprobsTensors":

@@ -126,6 +137,9 @@ class LogprobsTensors(NamedTuple):

137 assert self.cu_num_generated_tokens is None, (

138 "filter can't be used with cu_num_generated_tokens"

139 )

140 + assert self.cu_num_generated_tokens_tensor is None, (

141 + "filter can't be used with cu_num_generated_tokens_tensor"

142 + )

143 return LogprobsTensors(

144 self.logprob_token_ids[mask],

145 self.logprobs[mask],

@@ -147,6 +161,9 @@ class LogprobsTensors(NamedTuple):

161 if cu_num_generated_tokens is None:

162 return tensor

163 return tensor._replace(cu_num_generated_tokens=cu_num_generated_tokens)

164 + # The multi-chunk path rebuilds boundaries from the CPU layout, which

165 + # the device-only boundaries of adaptive verification never use.

166 + assert all(tensor.cu_num_generated_tokens_tensor is None for tensor in tensors)

167 return LogprobsTensors(

168 logprob_token_ids=torch.cat(

169 [tensor.logprob_token_ids for tensor in tensors]

vllm/v1/worker/gpu/sample/logprob.py

+4 -2

@@ -110,7 +110,7 @@ def compute_topk_scores(

110 logits: torch.Tensor,

111 num_logprobs: int,

112 sampled_token_ids: torch.Tensor,

113 - cu_num_logits: list[int] | None = None,

113 + cu_num_logits: list[int] | torch.Tensor | None = None,

114 logprob_token_ids_state: "LogprobTokenIdsState | None" = None,

115 expanded_idx_mapping: torch.Tensor | None = None,

116 max_per_req_token_ids: int = 0,

@@ -177,11 +177,13 @@ def compute_topk_scores(

177 vocab_size,

178 BLOCK_SIZE=8192, # type: ignore

179 )

180 + is_tensor = isinstance(cu_num_logits, torch.Tensor)

181 return LogprobsTensors(

182 logprob_token_ids=logprob_token_ids,

183 logprobs=scores,

184 selected_token_ranks=token_ranks,

184 - cu_num_generated_tokens=cu_num_logits,

185 + cu_num_generated_tokens=None if is_tensor else cu_num_logits,

186 + cu_num_generated_tokens_tensor=cu_num_logits if is_tensor else None,

187 )

188

189

vllm/v1/worker/gpu/spec_decode/rejection_sampler.py

+12 -1

@@ -81,6 +81,7 @@ class RejectionSampler:

81 ):

82 self.sampler = sampler

83 self.num_speculative_steps = spec_config.num_speculative_tokens

84 + self.enable_adaptive_verification = spec_config.enable_adaptive_verification

85 rejection_sample_method = spec_config.rejection_sample_method

86 self.use_block_verification: bool = False

87 self.synthetic_conditional_rates: torch.Tensor | None = None

@@ -122,11 +123,20 @@ class RejectionSampler:

123 num_warps=1,

124 )

125 expanded_logits = num_logits != num_reqs

126 + cu_num_generated_tokens: list[int] | torch.Tensor | None = None

127 + if expanded_logits:

128 + if self.enable_adaptive_verification:

129 + # Adaptive verification keeps the true per-request boundaries

130 + # on device only; cu_num_logits_np holds the pre-compacted

131 + # layout.

132 + cu_num_generated_tokens = cu_num_logits.clone()

133 + else:

134 + cu_num_generated_tokens = cu_num_logits_np.tolist()

135 return compute_topk_scores(

136 logits,

137 max_num_logprobs,

138 flat_sampled,

129 - cu_num_logits_np.tolist() if expanded_logits else None,

139 + cu_num_generated_tokens,

140 logits_mode=self.sampler.logprobs_mode

141 in ("raw_logits", "processed_logits"),

142 )

@@ -191,6 +201,7 @@ class RejectionSampler:

201 # guarantees the compacted batch always lands here.

202 request_chunks: Iterable[tuple[int, int]] = ((0, num_reqs),)

203 else:

204 + assert not self.enable_adaptive_verification

205 request_chunks = _iter_request_chunks(cu_num_logits_np, max_chunk_logits)

206

207 sampled_chunks: list[torch.Tensor] = []

vllm/v1/worker/gpu_model_runner.py

+1 -1

@@ -5803,7 +5803,7 @@ class GPUModelRunner(

5803 scores = logits.to(torch.float32)

5804 else:

5805 scores = self.sampler.compute_logprobs(logits)

5806 - token_ids, logprobs, ranks, _ = self.sampler.gather_logprobs(

5806 + token_ids, logprobs, ranks, *_ = self.sampler.gather_logprobs(

5807 scores, num_prompt_logprobs, tgt_token_ids

5808 )

5809