@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 16786da735

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 8, 17:45 UTC

[Model Runner V2] support apply penalty for spec decode (#33251)

Signed-off-by: zhuhaoran <zhuhaoran.zhr@alibaba-inc.com>

zhrrr committed Feb 7, 2026 at 02:56 UTC 16786da7357327f44f3a8f23d17e3c84235d2952

4 files changed +91 -14

vllm/v1/worker/gpu/input_batch.py

+12 -2

@@ -40,6 +40,8 @@ class InputBatch:

40 idx_mapping_np: np.ndarray

41 # Identical to idx_mapping except for spec decoding.

42 expanded_idx_mapping: torch.Tensor

43 + # [total_num_logits] position within request for each logit

44 + expanded_local_pos: torch.Tensor

46 # [num_reqs]

47 # batch_idx -> num_scheduled_tokens

@@ -91,6 +93,7 @@ class InputBatch:

93 idx_mapping_np = np.arange(num_reqs, dtype=np.int32)

94 idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device)

95 expanded_idx_mapping = idx_mapping

96 + expanded_local_pos = torch.zeros(num_reqs, dtype=torch.int32, device=device)

97 num_scheduled_tokens = np.full(num_reqs, num_tokens // num_reqs, dtype=np.int32)

98 num_scheduled_tokens[-1] += num_tokens % num_reqs

99 assert int(num_scheduled_tokens.sum()) == num_tokens

@@ -126,6 +129,7 @@ class InputBatch:

129 idx_mapping=idx_mapping,

130 idx_mapping_np=idx_mapping_np,

131 expanded_idx_mapping=expanded_idx_mapping,

132 + expanded_local_pos=expanded_local_pos,

133 num_scheduled_tokens=num_scheduled_tokens,

134 num_tokens=num_tokens,

135 num_tokens_after_padding=num_tokens,

@@ -487,6 +491,7 @@ def post_update(

491 def _expand_idx_mapping_kernel(

492 idx_mapping_ptr,

493 expanded_idx_mapping_ptr,

494 + expanded_local_pos_ptr,

495 cu_num_logits_ptr,

496 BLOCK_SIZE: tl.constexpr,

497 ):

@@ -499,6 +504,7 @@ def _expand_idx_mapping_kernel(

504 mask = block < num_tokens

505 req_state_idx = tl.load(idx_mapping_ptr + req_idx)

506 tl.store(expanded_idx_mapping_ptr + start_idx + block, req_state_idx, mask=mask)

507 + tl.store(expanded_local_pos_ptr + start_idx + block, block, mask=mask)

508

509

510 def expand_idx_mapping(

@@ -506,13 +512,17 @@ def expand_idx_mapping(

512 total_num_logits: int,

513 cu_num_logits: torch.Tensor,

514 max_expand_len: int,

509 -) -> torch.Tensor:

515 +) -> tuple[torch.Tensor, torch.Tensor]:

516 num_reqs = idx_mapping.shape[0]

517 expanded_idx_mapping = idx_mapping.new_empty(total_num_logits)

518 + expanded_local_pos = torch.empty(

519 + total_num_logits, dtype=torch.int32, device=idx_mapping.device

520 + )

521 _expand_idx_mapping_kernel[(num_reqs,)](

522 idx_mapping,

523 expanded_idx_mapping,

524 + expanded_local_pos,

525 cu_num_logits,

526 BLOCK_SIZE=triton.next_power_of_2(max_expand_len),

527 )

518 - return expanded_idx_mapping

528 + return expanded_idx_mapping, expanded_local_pos

vllm/v1/worker/gpu/model_runner.py

+22 -3

@@ -152,6 +152,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

152 vocab_size=self.vocab_size,

153 device=self.device,

154 logprobs_mode=self.model_config.logprobs_mode,

155 + num_speculative_tokens=self.num_speculative_steps + 1,

156 )

157 self.prompt_logprobs_worker = PromptLogprobsWorker(self.max_num_reqs)

158

@@ -318,10 +319,22 @@ class GPUModelRunner(LoRAModelRunnerMixin):

319 idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=self.device)

320 idx_mapping_np = np.arange(num_reqs, dtype=np.int32)

321 pos = torch.zeros(num_reqs, dtype=torch.int64, device=self.device)

322 + dummy_input_ids = torch.zeros(num_reqs, dtype=torch.int32, device=self.device)

323 + expanded_local_pos = torch.zeros(

324 + num_reqs, dtype=torch.int32, device=self.device

325 + )

326 # NOTE(woosuk): During the initial memory profiling, the sampler may skip

327 # top_k, top_p, and logprobs, using less GPU memory than what is possible

328 # during actual execution.

324 - self.sampler(logits, idx_mapping, idx_mapping_np, idx_mapping_np, pos)

329 + self.sampler(

330 + logits,

331 + idx_mapping,

332 + idx_mapping_np,

333 + idx_mapping_np,

334 + pos,

335 + dummy_input_ids,

336 + expanded_local_pos,

337 + )

338

339 @torch.inference_mode()

340 def profile_run(self) -> None:

@@ -511,6 +524,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):

524 num_reqs + 1, device=self.device, dtype=torch.int32

525 )

526 expanded_idx_mapping = idx_mapping

527 + expanded_local_pos = torch.zeros(

528 + num_reqs, dtype=torch.int32, device=self.device

529 + )

530 else:

531 num_draft_tokens = np.array(

532 [len(draft_tokens.get(req_id, ())) for req_id in req_ids],

@@ -526,7 +542,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

542 cu_num_logits = async_copy_to_gpu(cu_num_logits_np, device=self.device)

543

544 max_expand_len = self.num_speculative_steps + 1

529 - expanded_idx_mapping = expand_idx_mapping(

545 + expanded_idx_mapping, expanded_local_pos = expand_idx_mapping(

546 idx_mapping, total_num_logits, cu_num_logits, max_expand_len

547 )

548

@@ -627,6 +643,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

643 idx_mapping=idx_mapping,

644 idx_mapping_np=idx_mapping_np,

645 expanded_idx_mapping=expanded_idx_mapping,

646 + expanded_local_pos=expanded_local_pos,

647 num_scheduled_tokens=num_scheduled_tokens,

648 num_tokens=num_tokens,

649 num_tokens_after_padding=num_tokens_after_padding,

@@ -674,6 +691,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

691 ) -> tuple[SamplerOutput, torch.Tensor, torch.Tensor]:

692 sample_hidden_states = hidden_states[input_batch.logits_indices]

693 sample_pos = input_batch.positions[input_batch.logits_indices]

694 + input_ids = input_batch.input_ids[input_batch.logits_indices]

695 logits = self.model.compute_logits(sample_hidden_states)

696 if grammar_output is not None:

697 # Apply grammar bitmask to the logits in-place.

@@ -691,6 +709,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):

709 input_batch.idx_mapping_np,

710 input_batch.cu_num_logits_np,

711 sample_pos,

712 + input_ids,

713 + input_batch.expanded_local_pos,

714 )

715

716 if input_batch.num_draft_tokens == 0:

@@ -700,7 +720,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):

720 )

721 else:

722 # Rejection sampling for spec decoding.

703 - input_ids = input_batch.input_ids[input_batch.logits_indices]

723 sampled_tokens, num_sampled = rejection_sample(

724 sampler_output.sampled_token_ids,

725 input_ids,

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

+37 -7

@@ -75,6 +75,9 @@ class PenaltiesState:

75 logits: torch.Tensor,

76 idx_mapping: torch.Tensor,

77 idx_mapping_np: np.ndarray,

78 + input_ids: torch.Tensor,

79 + expanded_local_pos: torch.Tensor,

80 + num_speculative_tokens: int,

81 ) -> None:

82 if not np.any(self.use_penalty[idx_mapping_np]):

83 # No request uses penalties. Skip the kernel launch.

@@ -83,11 +86,14 @@ class PenaltiesState:

86 apply_penalties(

87 logits,

88 idx_mapping,

89 + input_ids,

90 + expanded_local_pos,

91 self.repetition_penalty.gpu,

92 self.frequency_penalty.gpu,

93 self.presence_penalty.gpu,

94 self.prompt_bin_mask,

95 self.output_bin_counts,

96 + num_speculative_tokens,

97 )

@@ -96,6 +102,8 @@ def _penalties_kernel(

102 logits_ptr,

103 logits_stride,

104 idx_mapping_ptr,

105 + token_ids_ptr,

106 + expanded_local_pos_ptr,

107 repetition_penalty_ptr,

108 frequency_penalty_ptr,

109 presence_penalty_ptr,

@@ -105,9 +113,10 @@ def _penalties_kernel(

113 output_bin_counts_stride,

114 vocab_size,

115 BLOCK_SIZE: tl.constexpr,

116 + MAX_SPEC_LEN: tl.constexpr,

117 ):

109 - batch_idx = tl.program_id(0)

110 - req_state_idx = tl.load(idx_mapping_ptr + batch_idx)

118 + token_idx = tl.program_id(0)

119 + req_state_idx = tl.load(idx_mapping_ptr + token_idx)

120 rep_penalty = tl.load(repetition_penalty_ptr + req_state_idx)

121 freq_penalty = tl.load(frequency_penalty_ptr + req_state_idx)

122 pres_penalty = tl.load(presence_penalty_ptr + req_state_idx)

@@ -123,13 +132,27 @@ def _penalties_kernel(

132 block_idx = tl.program_id(1)

133 block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)

134 mask = block < vocab_size

126 - logits = tl.load(logits_ptr + batch_idx * logits_stride + block, mask=mask)

135 + logits = tl.load(logits_ptr + token_idx * logits_stride + block, mask=mask)

136 logits = logits.to(tl.float32)

137

129 - output_bin_counts = tl.load(

138 + base_output_counts = tl.load(

139 output_bin_counts_ptr + req_state_idx * output_bin_counts_stride + block,

140 mask=mask,

141 + other=0,

142 )

143 +

144 + # Compute cumulative draft_counts from previous positions in this request

145 + pos = tl.load(expanded_local_pos_ptr + token_idx)

146 + start_idx = token_idx - pos

147 + draft_counts = tl.zeros((BLOCK_SIZE,), dtype=tl.int32)

148 + for prev_pos in tl.static_range(MAX_SPEC_LEN):

149 + if prev_pos < pos:

150 + prev_token = tl.load(token_ids_ptr + start_idx + prev_pos + 1)

151 + token_match = block == prev_token

152 + draft_counts = draft_counts + token_match.to(tl.int32)

153 +

154 + # Total counts = base output counts + cumulative draft counts

155 + output_bin_counts = base_output_counts + draft_counts

156 output_bin_mask = output_bin_counts > 0

157

158 # Apply repetition penalties.

@@ -138,6 +161,7 @@ def _penalties_kernel(

161 packed_mask = tl.load(

162 prompt_bin_mask_ptr + req_state_idx * prompt_bin_mask_stride + packed_block,

163 mask=packed_block < tl.cdiv(vocab_size, 32),

164 + other=0,

165 )

166 prompt_bin_mask = (packed_mask[:, None] >> (tl.arange(0, 32)[None, :])) & 1

167 prompt_bin_mask = prompt_bin_mask.to(tl.int1)

@@ -153,25 +177,30 @@ def _penalties_kernel(

177 # Apply presence penalties.

178 logits -= pres_penalty * output_bin_mask

179 # Store back to logits.

156 - tl.store(logits_ptr + batch_idx * logits_stride + block, logits, mask=mask)

180 + tl.store(logits_ptr + token_idx * logits_stride + block, logits, mask=mask)

181

182

183 def apply_penalties(

184 logits: torch.Tensor,

185 idx_mapping: torch.Tensor,

186 + token_ids: torch.Tensor,

187 + expanded_local_pos: torch.Tensor,

188 repetition_penalty: torch.Tensor,

189 frequency_penalty: torch.Tensor,

190 presence_penalty: torch.Tensor,

191 prompt_bin_mask: torch.Tensor,

192 output_bin_counts: torch.Tensor,

193 + num_speculative_tokens: int,

194 ) -> None:

168 - num_reqs, vocab_size = logits.shape

195 + num_tokens, vocab_size = logits.shape

196 BLOCK_SIZE = 8192

197 num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)

171 - _penalties_kernel[(num_reqs, num_blocks)](

198 + _penalties_kernel[(num_tokens, num_blocks)](

199 logits,

200 logits.stride(0),

201 idx_mapping,

202 + token_ids,

203 + expanded_local_pos,

204 repetition_penalty,

205 frequency_penalty,

206 presence_penalty,

@@ -181,6 +210,7 @@ def apply_penalties(

210 output_bin_counts.stride(0),

211 vocab_size,

212 BLOCK_SIZE=BLOCK_SIZE,

213 + MAX_SPEC_LEN=num_speculative_tokens,

214 )

215

216

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

+20 -2

@@ -25,6 +25,7 @@ class Sampler:

25 vocab_size: int,

26 device: torch.device,

27 logprobs_mode: LogprobsMode = "raw_logprobs",

28 + num_speculative_tokens: int = 1,

29 ):

30 if logprobs_mode not in ("processed_logprobs", "raw_logprobs"):

31 raise NotImplementedError(f"Unsupported logprobs_mode: {logprobs_mode}")

@@ -34,6 +35,7 @@ class Sampler:

35 self.sampling_states = SamplingStates(max_num_reqs, vocab_size)

36 self.penalties_state = PenaltiesState(max_num_reqs, vocab_size, device)

37 self.logit_bias_state = LogitBiasState(max_num_reqs, device)

38 + self.num_speculative_tokens = num_speculative_tokens

40 def add_request(

41 self, req_idx: int, prompt_len: int, sampling_params: SamplingParams

@@ -61,12 +63,19 @@ class Sampler:

63 idx_mapping_np: np.ndarray,

64 cu_num_logits_np: np.ndarray,

65 pos: torch.Tensor,

66 + input_ids: torch.Tensor,

67 + expanded_local_pos: torch.Tensor,

68 ) -> SamplerOutput:

69 # NOTE(woosuk): We intentionally compute num_nans before sampling to make clear

70 # that num_nans is computed before applying penalties and temperature.

71 num_nans = get_num_nans(logits) if self.compute_nans else None

72 sampled, processed_logits = self.sample(

69 - logits, idx_mapping, idx_mapping_np, pos

73 + logits,

74 + idx_mapping,

75 + idx_mapping_np,

76 + pos,

77 + input_ids,

78 + expanded_local_pos,

79 )

81 max_num_logprobs = self.sampling_states.max_num_logprobs(idx_mapping_np)

@@ -98,6 +107,8 @@ class Sampler:

107 idx_mapping: torch.Tensor,

108 idx_mapping_np: np.ndarray,

109 pos: torch.Tensor,

110 + input_ids: torch.Tensor,

111 + expanded_local_pos: torch.Tensor,

112 ) -> tuple[torch.Tensor, torch.Tensor]:

113 # Copy logits to a new FP32 tensor.

114 logits = torch.empty_like(logits, dtype=torch.float32).copy_(logits)

@@ -106,7 +117,14 @@ class Sampler:

117 self.logit_bias_state.apply_logit_bias(logits, idx_mapping, idx_mapping_np, pos)

118

119 # Apply penalties in place.

109 - self.penalties_state.apply_penalties(logits, idx_mapping, idx_mapping_np)

120 + self.penalties_state.apply_penalties(

121 + logits,

122 + idx_mapping,

123 + idx_mapping_np,

124 + input_ids,

125 + expanded_local_pos,

126 + self.num_speculative_tokens,

127 + )

128

129 # Apply temperature in place.

130 apply_temperature(logits, idx_mapping, self.sampling_states.temperature.gpu)