From 4e498b5e5c07480cfb8c046128f0ef8d9a60d8ef Mon Sep 17 00:00:00 2001
From: Lucas Kabela <lucaskabela@meta.com>
Date: Mon, 11 May 2026 22:47:20 -0700
Subject: [PATCH] [Bugfix][Performance Improvement] Improve penalties triton
 kernel performance (#40657)

Signed-off-by: Lucas Kabela <lucaskabela@meta.com>
---
 vllm/v1/worker/gpu/sample/penalties.py | 23 ++++++++---------------
 vllm/v1/worker/gpu/sample/sampler.py   |  1 -
 2 files changed, 8 insertions(+), 16 deletions(-)

diff --git a/vllm/v1/worker/gpu/sample/penalties.py b/vllm/v1/worker/gpu/sample/penalties.py
index 2cfa8b0aec..b2ce2fb812 100644
--- a/vllm/v1/worker/gpu/sample/penalties.py
+++ b/vllm/v1/worker/gpu/sample/penalties.py
@@ -85,7 +85,6 @@ class PenaltiesState:
         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.
@@ -101,7 +100,6 @@ class PenaltiesState:
             self.presence_penalty.gpu,
             self.prompt_bin_mask,
             self.output_bin_counts,
-            num_speculative_tokens,
         )
 
 
@@ -121,7 +119,6 @@ def _penalties_kernel(
     output_bin_counts_stride,
     vocab_size,
     BLOCK_SIZE: tl.constexpr,
-    MAX_SPEC_LEN: tl.constexpr,
 ):
     token_idx = tl.program_id(0)
     req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)
@@ -149,18 +146,16 @@ def _penalties_kernel(
         other=0,
     )
 
-    # Compute cumulative draft_counts from previous positions in this request
+    # Accumulate draft token counts from previous positions directly into
+    # output_bin_counts (preserves its native tensor layout, avoiding an
+    # expensive shared-memory layout conversion after the loop).
     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_counts = base_output_counts
+    for prev_pos in tl.range(pos):
+        prev_token = tl.load(token_ids_ptr + start_idx + prev_pos + 1)
+        token_match = block == prev_token
+        output_bin_counts = output_bin_counts + token_match.to(tl.int32)
     output_bin_mask = output_bin_counts > 0
 
     # Apply repetition penalties.
@@ -198,7 +193,6 @@ def apply_penalties(
     presence_penalty: torch.Tensor,
     prompt_bin_mask: torch.Tensor,
     output_bin_counts: torch.Tensor,
-    num_speculative_tokens: int,
 ) -> None:
     num_tokens, vocab_size = logits.shape
     BLOCK_SIZE = 8192
@@ -218,7 +212,6 @@ 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 5d91d5b2f0..7b4b01666f 100644
--- a/vllm/v1/worker/gpu/sample/sampler.py
+++ b/vllm/v1/worker/gpu/sample/sampler.py
@@ -142,7 +142,6 @@ class Sampler:
             idx_mapping_np,
             input_ids,
             expanded_local_pos,
-            self.num_speculative_tokens,
         )
 
         # Apply bad words masking in place.