From 16786da7357327f44f3a8f23d17e3c84235d2952 Mon Sep 17 00:00:00 2001 From: zhrrr <43847754+izhuhaoran@users.noreply.github.com> Date: Sat, 7 Feb 2026 02:56:48 +0800 Subject: [PATCH] [Model Runner V2] support apply penalty for spec decode (#33251) Signed-off-by: zhuhaoran --- vllm/v1/worker/gpu/input_batch.py | 14 ++++++-- vllm/v1/worker/gpu/model_runner.py | 25 +++++++++++++-- vllm/v1/worker/gpu/sample/penalties.py | 44 ++++++++++++++++++++++---- vllm/v1/worker/gpu/sample/sampler.py | 22 +++++++++++-- 4 files changed, 91 insertions(+), 14 deletions(-) diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index 1b18768b3..d90b0dc01 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -40,6 +40,8 @@ class InputBatch: idx_mapping_np: np.ndarray # Identical to idx_mapping except for spec decoding. expanded_idx_mapping: torch.Tensor + # [total_num_logits] position within request for each logit + expanded_local_pos: torch.Tensor # [num_reqs] # batch_idx -> num_scheduled_tokens @@ -91,6 +93,7 @@ class InputBatch: idx_mapping_np = np.arange(num_reqs, dtype=np.int32) idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device) expanded_idx_mapping = idx_mapping + expanded_local_pos = torch.zeros(num_reqs, dtype=torch.int32, device=device) num_scheduled_tokens = np.full(num_reqs, num_tokens // num_reqs, dtype=np.int32) num_scheduled_tokens[-1] += num_tokens % num_reqs assert int(num_scheduled_tokens.sum()) == num_tokens @@ -126,6 +129,7 @@ class InputBatch: idx_mapping=idx_mapping, idx_mapping_np=idx_mapping_np, expanded_idx_mapping=expanded_idx_mapping, + expanded_local_pos=expanded_local_pos, num_scheduled_tokens=num_scheduled_tokens, num_tokens=num_tokens, num_tokens_after_padding=num_tokens, @@ -487,6 +491,7 @@ def post_update( def _expand_idx_mapping_kernel( idx_mapping_ptr, expanded_idx_mapping_ptr, + expanded_local_pos_ptr, cu_num_logits_ptr, BLOCK_SIZE: tl.constexpr, ): @@ -499,6 +504,7 @@ def _expand_idx_mapping_kernel( mask = block < num_tokens req_state_idx = tl.load(idx_mapping_ptr + req_idx) tl.store(expanded_idx_mapping_ptr + start_idx + block, req_state_idx, mask=mask) + tl.store(expanded_local_pos_ptr + start_idx + block, block, mask=mask) def expand_idx_mapping( @@ -506,13 +512,17 @@ def expand_idx_mapping( total_num_logits: int, cu_num_logits: torch.Tensor, max_expand_len: int, -) -> torch.Tensor: +) -> tuple[torch.Tensor, torch.Tensor]: num_reqs = idx_mapping.shape[0] expanded_idx_mapping = idx_mapping.new_empty(total_num_logits) + expanded_local_pos = torch.empty( + total_num_logits, dtype=torch.int32, device=idx_mapping.device + ) _expand_idx_mapping_kernel[(num_reqs,)]( idx_mapping, expanded_idx_mapping, + expanded_local_pos, cu_num_logits, BLOCK_SIZE=triton.next_power_of_2(max_expand_len), ) - return expanded_idx_mapping + return expanded_idx_mapping, expanded_local_pos diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 43e26f3c9..416eaa011 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -152,6 +152,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): vocab_size=self.vocab_size, device=self.device, logprobs_mode=self.model_config.logprobs_mode, + num_speculative_tokens=self.num_speculative_steps + 1, ) self.prompt_logprobs_worker = PromptLogprobsWorker(self.max_num_reqs) @@ -318,10 +319,22 @@ class GPUModelRunner(LoRAModelRunnerMixin): idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=self.device) idx_mapping_np = np.arange(num_reqs, dtype=np.int32) pos = torch.zeros(num_reqs, dtype=torch.int64, device=self.device) + dummy_input_ids = torch.zeros(num_reqs, dtype=torch.int32, device=self.device) + expanded_local_pos = torch.zeros( + num_reqs, dtype=torch.int32, device=self.device + ) # NOTE(woosuk): During the initial memory profiling, the sampler may skip # top_k, top_p, and logprobs, using less GPU memory than what is possible # during actual execution. - self.sampler(logits, idx_mapping, idx_mapping_np, idx_mapping_np, pos) + self.sampler( + logits, + idx_mapping, + idx_mapping_np, + idx_mapping_np, + pos, + dummy_input_ids, + expanded_local_pos, + ) @torch.inference_mode() def profile_run(self) -> None: @@ -511,6 +524,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): num_reqs + 1, device=self.device, dtype=torch.int32 ) expanded_idx_mapping = idx_mapping + expanded_local_pos = torch.zeros( + num_reqs, dtype=torch.int32, device=self.device + ) else: num_draft_tokens = np.array( [len(draft_tokens.get(req_id, ())) for req_id in req_ids], @@ -526,7 +542,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): cu_num_logits = async_copy_to_gpu(cu_num_logits_np, device=self.device) max_expand_len = self.num_speculative_steps + 1 - expanded_idx_mapping = expand_idx_mapping( + expanded_idx_mapping, expanded_local_pos = expand_idx_mapping( idx_mapping, total_num_logits, cu_num_logits, max_expand_len ) @@ -627,6 +643,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): idx_mapping=idx_mapping, idx_mapping_np=idx_mapping_np, expanded_idx_mapping=expanded_idx_mapping, + expanded_local_pos=expanded_local_pos, num_scheduled_tokens=num_scheduled_tokens, num_tokens=num_tokens, num_tokens_after_padding=num_tokens_after_padding, @@ -674,6 +691,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): ) -> tuple[SamplerOutput, torch.Tensor, torch.Tensor]: sample_hidden_states = hidden_states[input_batch.logits_indices] sample_pos = input_batch.positions[input_batch.logits_indices] + input_ids = input_batch.input_ids[input_batch.logits_indices] logits = self.model.compute_logits(sample_hidden_states) if grammar_output is not None: # Apply grammar bitmask to the logits in-place. @@ -691,6 +709,8 @@ class GPUModelRunner(LoRAModelRunnerMixin): input_batch.idx_mapping_np, input_batch.cu_num_logits_np, sample_pos, + input_ids, + input_batch.expanded_local_pos, ) if input_batch.num_draft_tokens == 0: @@ -700,7 +720,6 @@ class GPUModelRunner(LoRAModelRunnerMixin): ) else: # Rejection sampling for spec decoding. - input_ids = input_batch.input_ids[input_batch.logits_indices] sampled_tokens, num_sampled = rejection_sample( sampler_output.sampled_token_ids, input_ids, diff --git a/vllm/v1/worker/gpu/sample/penalties.py b/vllm/v1/worker/gpu/sample/penalties.py index 2e6194df5..24928fd10 100644 --- a/vllm/v1/worker/gpu/sample/penalties.py +++ b/vllm/v1/worker/gpu/sample/penalties.py @@ -75,6 +75,9 @@ class PenaltiesState: logits: torch.Tensor, idx_mapping: torch.Tensor, idx_mapping_np: np.ndarray, + input_ids: torch.Tensor, + expanded_local_pos: torch.Tensor, + num_speculative_tokens: int, ) -> None: if not np.any(self.use_penalty[idx_mapping_np]): # No request uses penalties. Skip the kernel launch. @@ -83,11 +86,14 @@ class PenaltiesState: apply_penalties( logits, idx_mapping, + input_ids, + expanded_local_pos, self.repetition_penalty.gpu, self.frequency_penalty.gpu, self.presence_penalty.gpu, self.prompt_bin_mask, self.output_bin_counts, + num_speculative_tokens, ) @@ -96,6 +102,8 @@ def _penalties_kernel( logits_ptr, logits_stride, idx_mapping_ptr, + token_ids_ptr, + expanded_local_pos_ptr, repetition_penalty_ptr, frequency_penalty_ptr, presence_penalty_ptr, @@ -105,9 +113,10 @@ def _penalties_kernel( output_bin_counts_stride, vocab_size, BLOCK_SIZE: tl.constexpr, + MAX_SPEC_LEN: tl.constexpr, ): - batch_idx = tl.program_id(0) - req_state_idx = tl.load(idx_mapping_ptr + batch_idx) + token_idx = tl.program_id(0) + req_state_idx = tl.load(idx_mapping_ptr + token_idx) rep_penalty = tl.load(repetition_penalty_ptr + req_state_idx) freq_penalty = tl.load(frequency_penalty_ptr + req_state_idx) pres_penalty = tl.load(presence_penalty_ptr + req_state_idx) @@ -123,13 +132,27 @@ def _penalties_kernel( block_idx = tl.program_id(1) block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) mask = block < vocab_size - logits = tl.load(logits_ptr + batch_idx * logits_stride + block, mask=mask) + logits = tl.load(logits_ptr + token_idx * logits_stride + block, mask=mask) logits = logits.to(tl.float32) - output_bin_counts = tl.load( + base_output_counts = tl.load( output_bin_counts_ptr + req_state_idx * output_bin_counts_stride + block, mask=mask, + other=0, ) + + # Compute cumulative draft_counts from previous positions in this request + pos = tl.load(expanded_local_pos_ptr + token_idx) + start_idx = token_idx - pos + draft_counts = tl.zeros((BLOCK_SIZE,), dtype=tl.int32) + for prev_pos in tl.static_range(MAX_SPEC_LEN): + if prev_pos < pos: + prev_token = tl.load(token_ids_ptr + start_idx + prev_pos + 1) + token_match = block == prev_token + draft_counts = draft_counts + token_match.to(tl.int32) + + # Total counts = base output counts + cumulative draft counts + output_bin_counts = base_output_counts + draft_counts output_bin_mask = output_bin_counts > 0 # Apply repetition penalties. @@ -138,6 +161,7 @@ def _penalties_kernel( packed_mask = tl.load( prompt_bin_mask_ptr + req_state_idx * prompt_bin_mask_stride + packed_block, mask=packed_block < tl.cdiv(vocab_size, 32), + other=0, ) prompt_bin_mask = (packed_mask[:, None] >> (tl.arange(0, 32)[None, :])) & 1 prompt_bin_mask = prompt_bin_mask.to(tl.int1) @@ -153,25 +177,30 @@ def _penalties_kernel( # Apply presence penalties. logits -= pres_penalty * output_bin_mask # Store back to logits. - tl.store(logits_ptr + batch_idx * logits_stride + block, logits, mask=mask) + tl.store(logits_ptr + token_idx * logits_stride + block, logits, mask=mask) def apply_penalties( logits: torch.Tensor, idx_mapping: torch.Tensor, + token_ids: torch.Tensor, + expanded_local_pos: torch.Tensor, repetition_penalty: torch.Tensor, frequency_penalty: torch.Tensor, presence_penalty: torch.Tensor, prompt_bin_mask: torch.Tensor, output_bin_counts: torch.Tensor, + num_speculative_tokens: int, ) -> None: - num_reqs, vocab_size = logits.shape + num_tokens, vocab_size = logits.shape BLOCK_SIZE = 8192 num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE) - _penalties_kernel[(num_reqs, num_blocks)]( + _penalties_kernel[(num_tokens, num_blocks)]( logits, logits.stride(0), idx_mapping, + token_ids, + expanded_local_pos, repetition_penalty, frequency_penalty, presence_penalty, @@ -181,6 +210,7 @@ def apply_penalties( output_bin_counts.stride(0), vocab_size, BLOCK_SIZE=BLOCK_SIZE, + MAX_SPEC_LEN=num_speculative_tokens, ) diff --git a/vllm/v1/worker/gpu/sample/sampler.py b/vllm/v1/worker/gpu/sample/sampler.py index 25904c076..094fffacf 100644 --- a/vllm/v1/worker/gpu/sample/sampler.py +++ b/vllm/v1/worker/gpu/sample/sampler.py @@ -25,6 +25,7 @@ class Sampler: vocab_size: int, device: torch.device, logprobs_mode: LogprobsMode = "raw_logprobs", + num_speculative_tokens: int = 1, ): if logprobs_mode not in ("processed_logprobs", "raw_logprobs"): raise NotImplementedError(f"Unsupported logprobs_mode: {logprobs_mode}") @@ -34,6 +35,7 @@ class Sampler: self.sampling_states = SamplingStates(max_num_reqs, vocab_size) self.penalties_state = PenaltiesState(max_num_reqs, vocab_size, device) self.logit_bias_state = LogitBiasState(max_num_reqs, device) + self.num_speculative_tokens = num_speculative_tokens def add_request( self, req_idx: int, prompt_len: int, sampling_params: SamplingParams @@ -61,12 +63,19 @@ class Sampler: idx_mapping_np: np.ndarray, cu_num_logits_np: np.ndarray, pos: torch.Tensor, + input_ids: torch.Tensor, + expanded_local_pos: torch.Tensor, ) -> SamplerOutput: # NOTE(woosuk): We intentionally compute num_nans before sampling to make clear # that num_nans is computed before applying penalties and temperature. num_nans = get_num_nans(logits) if self.compute_nans else None sampled, processed_logits = self.sample( - logits, idx_mapping, idx_mapping_np, pos + logits, + idx_mapping, + idx_mapping_np, + pos, + input_ids, + expanded_local_pos, ) max_num_logprobs = self.sampling_states.max_num_logprobs(idx_mapping_np) @@ -98,6 +107,8 @@ class Sampler: idx_mapping: torch.Tensor, idx_mapping_np: np.ndarray, pos: torch.Tensor, + input_ids: torch.Tensor, + expanded_local_pos: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: # Copy logits to a new FP32 tensor. logits = torch.empty_like(logits, dtype=torch.float32).copy_(logits) @@ -106,7 +117,14 @@ class Sampler: self.logit_bias_state.apply_logit_bias(logits, idx_mapping, idx_mapping_np, pos) # Apply penalties in place. - self.penalties_state.apply_penalties(logits, idx_mapping, idx_mapping_np) + self.penalties_state.apply_penalties( + logits, + idx_mapping, + idx_mapping_np, + input_ids, + expanded_local_pos, + self.num_speculative_tokens, + ) # Apply temperature in place. apply_temperature(logits, idx_mapping, self.sampling_states.temperature.gpu)