[Model Runner V2] Support penalties using bin counts (#29703)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

Author: WoosukKwon

vllm/v1/worker/gpu/input_batch.py

@@ -341,6 +341,8 @@ def _post_update_kernel(
     idx_mapping_ptr,
     num_computed_tokens_ptr,
     last_sampled_tokens_ptr,
+    output_bin_counts_ptr,
+    output_bin_counts_stride,
     sampled_tokens_ptr,
     sampled_tokens_stride,
     num_sampled_ptr,
@@ -357,6 +359,15 @@ def _post_update_kernel(
         )
         tl.store(last_sampled_tokens_ptr + req_state_idx, token_id)
 
+    for i in range(num_sampled):
+        token_id = tl.load(sampled_tokens_ptr + req_id * sampled_tokens_stride + i)
+        token_ptr = (
+            output_bin_counts_ptr + req_state_idx * output_bin_counts_stride + token_id
+        )
+        count = tl.load(token_ptr)
+        count += 1
+        tl.store(token_ptr, count)
+
     query_start = tl.load(query_start_loc_ptr + req_id)
     query_end = tl.load(query_start_loc_ptr + req_id + 1)
     query_len = query_end - query_start
@@ -374,6 +385,8 @@ def post_update(
     num_computed_tokens: torch.Tensor,
     # [max_num_reqs]
     last_sampled_tokens: torch.Tensor,
+    # [max_num_reqs, vocab_size]
+    output_bin_counts: torch.Tensor,
     # [num_reqs, num_speculative_steps + 1]
     sampled_tokens: torch.Tensor,
     # [num_reqs]
@@ -388,6 +401,8 @@ def post_update(
         idx_mapping,
         num_computed_tokens,
         last_sampled_tokens,
+        output_bin_counts,
+        output_bin_counts.stride(0),
         sampled_tokens,
         sampled_tokens.stride(0),
         num_sampled,

vllm/v1/worker/gpu/model_runner.py

@@ -512,7 +512,7 @@ def prepare_inputs(
             idx_mapping_np,
             num_scheduled_tokens,
             query_start_loc_np,
-            self.req_states.prefill_token_ids,
+            self.req_states.prefill_token_ids.np,
             self.req_states.num_computed_prefill_tokens,
             self.input_buffers.input_ids.np,
         )
@@ -681,7 +681,7 @@ def compute_prompt_logprobs(
         # Handle chunked prompts.
         pos_after_step = computed_prefill + input_batch.num_scheduled_tokens
         is_prompt_chunked = pos_after_step < prompt_lens
-        prefill_token_ids = self.req_states.prefill_token_ids
+        prefill_token_ids = self.req_states.prefill_token_ids.np
         query_start_loc = self.input_buffers.query_start_loc.np
         for i, req_id in enumerate(input_batch.req_ids):
             if not needs_prompt_logprobs[i]:
@@ -756,6 +756,7 @@ def postprocess(
             input_batch.idx_mapping,
             self.req_states.num_computed_tokens,
             self.req_states.last_sampled_tokens,
+            self.req_states.output_bin_counts,
             sampled_tokens,
             num_sampled,
             num_rejected,
@@ -785,7 +786,7 @@ def propose_draft(
         idx_mapping_np = input_batch.idx_mapping_np
         with async_barrier(self.spec_decode_event):
             self.input_buffers.next_prefill_tokens.np[:num_reqs] = (
-                self.req_states.prefill_token_ids[
+                self.req_states.prefill_token_ids.np[
                     idx_mapping_np,
                     self.req_states.num_computed_prefill_tokens[idx_mapping_np],
                 ]
@@ -896,7 +897,7 @@ def execute_model(
                 # barrier to avoid race conditions.
                 pos = input_batch.positions[input_batch.logits_indices]
                 sampling_metadata = self.req_states.make_sampling_metadata(
-                    input_batch.idx_mapping_np, pos
+                    input_batch.idx_mapping, input_batch.idx_mapping_np, pos
                 )
                 if input_batch.num_draft_tokens > 0:
                     sampling_metadata = self.req_states.expand_sampling_metadata(

vllm/v1/worker/gpu/penalties.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+
+from vllm.triton_utils import tl, triton
+from vllm.v1.worker.gpu.states import SamplingMetadata
+
+
+@triton.jit
+def _penalties_kernel(
+    logits_ptr,
+    logits_stride,
+    repetition_penalty_ptr,
+    frequency_penalty_ptr,
+    presence_penalty_ptr,
+    idx_mapping_ptr,
+    prompt_bin_counts_ptr,
+    prompt_bin_counts_stride,
+    output_bin_counts_ptr,
+    output_bin_counts_stride,
+    vocab_size,
+    BLOCK_SIZE: tl.constexpr,
+):
+    batch_idx = tl.program_id(0)
+    rep_penalty = tl.load(repetition_penalty_ptr + batch_idx)
+    freq_penalty = tl.load(frequency_penalty_ptr + batch_idx)
+    pres_penalty = tl.load(presence_penalty_ptr + batch_idx)
+
+    use_rep_penalty = rep_penalty != 1.0
+    use_freq_penalty = freq_penalty != 0.0
+    use_pres_penalty = pres_penalty != 0.0
+    if not (use_rep_penalty or use_freq_penalty or use_pres_penalty):
+        # No penalties to apply. Early return.
+        return
+
+    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 = logits.to(tl.float32)
+
+    req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
+    output_bin_counts = tl.load(
+        output_bin_counts_ptr + req_state_idx * output_bin_counts_stride + block,
+        mask=mask,
+    )
+
+    # Apply repetition penalties.
+    if use_rep_penalty:
+        prompt_bin_counts = tl.load(
+            prompt_bin_counts_ptr + req_state_idx * prompt_bin_counts_stride + block,
+            mask=mask,
+        )
+        # If token appears in prompt or output, apply, otherwise use 1.0 for no-op.
+        scale = tl.where((prompt_bin_counts + output_bin_counts) > 0, rep_penalty, 1.0)
+        # If logits are positive, divide by penalty, otherwise multiply by penalty.
+        scale = tl.where(logits > 0, 1.0 / scale, scale)
+        logits *= scale
+
+    # Apply frequency penalties.
+    logits -= freq_penalty * output_bin_counts
+    # Apply presence penalties.
+    logits -= pres_penalty * (output_bin_counts > 0)
+    # Store back to logits.
+    tl.store(logits_ptr + batch_idx * logits_stride + block, logits, mask=mask)
+
+
+def apply_penalties(logits: torch.Tensor, sampling_metadata: SamplingMetadata) -> None:
+    num_reqs, vocab_size = logits.shape
+    BLOCK_SIZE = 8192
+    num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
+    _penalties_kernel[(num_reqs, num_blocks)](
+        logits,
+        logits.stride(0),
+        sampling_metadata.repetition_penalty,
+        sampling_metadata.frequency_penalty,
+        sampling_metadata.presence_penalty,
+        sampling_metadata.idx_mapping,
+        sampling_metadata.prompt_bin_counts,
+        sampling_metadata.prompt_bin_counts.stride(0),
+        sampling_metadata.output_bin_counts,
+        sampling_metadata.output_bin_counts.stride(0),
+        vocab_size,
+        BLOCK_SIZE=BLOCK_SIZE,
+    )

vllm/v1/worker/gpu/sampler.py

@@ -8,6 +8,7 @@
 from vllm.triton_utils import tl, triton
 from vllm.v1.outputs import LogprobsTensors, SamplerOutput
 from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
+from vllm.v1.worker.gpu.penalties import apply_penalties
 from vllm.v1.worker.gpu.states import SamplingMetadata
 
 
@@ -65,6 +66,8 @@ def sample(
         logits = apply_top_k_top_p(
             logits, sampling_metadata.top_k, sampling_metadata.top_p
         )
+        # Apply penalties in place.
+        apply_penalties(logits, sampling_metadata)
 
         sampled = gumbel_sample(
             logits,