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 12, 08:30 UTC
[Model Runner V2] Support penalties using bin counts (#29703)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Nov 28, 2025 at 17:53 UTC
1dcafb3dea62f556011be4df6f71769aa7260561
5 files changed
+280
-14
vllm/v1/worker/gpu/input_batch.py
+15
@@ -341,6 +341,8 @@ def _post_update_kernel(
341
idx_mapping_ptr,
342
num_computed_tokens_ptr,
343
last_sampled_tokens_ptr,
344
+ output_bin_counts_ptr,
345
+ output_bin_counts_stride,
346
sampled_tokens_ptr,
347
sampled_tokens_stride,
348
num_sampled_ptr,
@@ -357,6 +359,15 @@ def _post_update_kernel(
359
)
360
tl.store(last_sampled_tokens_ptr + req_state_idx, token_id)
361
362
+ for i in range(num_sampled):
363
+ token_id = tl.load(sampled_tokens_ptr + req_id * sampled_tokens_stride + i)
364
+ token_ptr = (
365
+ output_bin_counts_ptr + req_state_idx * output_bin_counts_stride + token_id
366
+ )
367
+ count = tl.load(token_ptr)
368
+ count += 1
369
+ tl.store(token_ptr, count)
370
+
371
query_start = tl.load(query_start_loc_ptr + req_id)
372
query_end = tl.load(query_start_loc_ptr + req_id + 1)
373
query_len = query_end - query_start
@@ -374,6 +385,8 @@ def post_update(
385
num_computed_tokens: torch.Tensor,
386
# [max_num_reqs]
387
last_sampled_tokens: torch.Tensor,
388
+ # [max_num_reqs, vocab_size]
389
+ output_bin_counts: torch.Tensor,
390
# [num_reqs, num_speculative_steps + 1]
391
sampled_tokens: torch.Tensor,
392
# [num_reqs]
@@ -388,6 +401,8 @@ def post_update(
401
idx_mapping,
402
num_computed_tokens,
403
last_sampled_tokens,
404
+ output_bin_counts,
405
+ output_bin_counts.stride(0),
406
sampled_tokens,
407
sampled_tokens.stride(0),
408
num_sampled,
vllm/v1/worker/gpu/model_runner.py
+5
-4
@@ -512,7 +512,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
512
idx_mapping_np,
513
num_scheduled_tokens,
514
query_start_loc_np,
515
- self.req_states.prefill_token_ids,
515
+ self.req_states.prefill_token_ids.np,
516
self.req_states.num_computed_prefill_tokens,
517
self.input_buffers.input_ids.np,
518
)
@@ -681,7 +681,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
681
# Handle chunked prompts.
682
pos_after_step = computed_prefill + input_batch.num_scheduled_tokens
683
is_prompt_chunked = pos_after_step < prompt_lens
684
- prefill_token_ids = self.req_states.prefill_token_ids
684
+ prefill_token_ids = self.req_states.prefill_token_ids.np
685
query_start_loc = self.input_buffers.query_start_loc.np
686
for i, req_id in enumerate(input_batch.req_ids):
687
if not needs_prompt_logprobs[i]:
@@ -756,6 +756,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
756
input_batch.idx_mapping,
757
self.req_states.num_computed_tokens,
758
self.req_states.last_sampled_tokens,
759
+ self.req_states.output_bin_counts,
760
sampled_tokens,
761
num_sampled,
762
num_rejected,
@@ -785,7 +786,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
786
idx_mapping_np = input_batch.idx_mapping_np
787
with async_barrier(self.spec_decode_event):
788
self.input_buffers.next_prefill_tokens.np[:num_reqs] = (
788
- self.req_states.prefill_token_ids[
789
+ self.req_states.prefill_token_ids.np[
790
idx_mapping_np,
791
self.req_states.num_computed_prefill_tokens[idx_mapping_np],
792
]
@@ -896,7 +897,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
897
# barrier to avoid race conditions.
898
pos = input_batch.positions[input_batch.logits_indices]
899
sampling_metadata = self.req_states.make_sampling_metadata(
899
- input_batch.idx_mapping_np, pos
900
+ input_batch.idx_mapping, input_batch.idx_mapping_np, pos
901
)
902
if input_batch.num_draft_tokens > 0:
903
sampling_metadata = self.req_states.expand_sampling_metadata(
vllm/v1/worker/gpu/penalties.py
new
+85
@@ -0,0 +1,85 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+
5
+from vllm.triton_utils import tl, triton
6
+from vllm.v1.worker.gpu.states import SamplingMetadata
7
+
8
+
9
+@triton.jit
10
+def _penalties_kernel(
11
+ logits_ptr,
12
+ logits_stride,
13
+ repetition_penalty_ptr,
14
+ frequency_penalty_ptr,
15
+ presence_penalty_ptr,
16
+ idx_mapping_ptr,
17
+ prompt_bin_counts_ptr,
18
+ prompt_bin_counts_stride,
19
+ output_bin_counts_ptr,
20
+ output_bin_counts_stride,
21
+ vocab_size,
22
+ BLOCK_SIZE: tl.constexpr,
23
+):
24
+ batch_idx = tl.program_id(0)
25
+ rep_penalty = tl.load(repetition_penalty_ptr + batch_idx)
26
+ freq_penalty = tl.load(frequency_penalty_ptr + batch_idx)
27
+ pres_penalty = tl.load(presence_penalty_ptr + batch_idx)
28
+
29
+ use_rep_penalty = rep_penalty != 1.0
30
+ use_freq_penalty = freq_penalty != 0.0
31
+ use_pres_penalty = pres_penalty != 0.0
32
+ if not (use_rep_penalty or use_freq_penalty or use_pres_penalty):
33
+ # No penalties to apply. Early return.
34
+ return
35
+
36
+ block_idx = tl.program_id(1)
37
+ block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
38
+ mask = block < vocab_size
39
+ logits = tl.load(logits_ptr + batch_idx * logits_stride + block, mask=mask)
40
+ logits = logits.to(tl.float32)
41
+
42
+ req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
43
+ output_bin_counts = tl.load(
44
+ output_bin_counts_ptr + req_state_idx * output_bin_counts_stride + block,
45
+ mask=mask,
46
+ )
47
+
48
+ # Apply repetition penalties.
49
+ if use_rep_penalty:
50
+ prompt_bin_counts = tl.load(
51
+ prompt_bin_counts_ptr + req_state_idx * prompt_bin_counts_stride + block,
52
+ mask=mask,
53
+ )
54
+ # If token appears in prompt or output, apply, otherwise use 1.0 for no-op.
55
+ scale = tl.where((prompt_bin_counts + output_bin_counts) > 0, rep_penalty, 1.0)
56
+ # If logits are positive, divide by penalty, otherwise multiply by penalty.
57
+ scale = tl.where(logits > 0, 1.0 / scale, scale)
58
+ logits *= scale
59
+
60
+ # Apply frequency penalties.
61
+ logits -= freq_penalty * output_bin_counts
62
+ # Apply presence penalties.
63
+ logits -= pres_penalty * (output_bin_counts > 0)
64
+ # Store back to logits.
65
+ tl.store(logits_ptr + batch_idx * logits_stride + block, logits, mask=mask)
66
+
67
+
68
+def apply_penalties(logits: torch.Tensor, sampling_metadata: SamplingMetadata) -> None:
69
+ num_reqs, vocab_size = logits.shape
70
+ BLOCK_SIZE = 8192
71
+ num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
72
+ _penalties_kernel[(num_reqs, num_blocks)](
73
+ logits,
74
+ logits.stride(0),
75
+ sampling_metadata.repetition_penalty,
76
+ sampling_metadata.frequency_penalty,
77
+ sampling_metadata.presence_penalty,
78
+ sampling_metadata.idx_mapping,
79
+ sampling_metadata.prompt_bin_counts,
80
+ sampling_metadata.prompt_bin_counts.stride(0),
81
+ sampling_metadata.output_bin_counts,
82
+ sampling_metadata.output_bin_counts.stride(0),
83
+ vocab_size,
84
+ BLOCK_SIZE=BLOCK_SIZE,
85
+ )
vllm/v1/worker/gpu/sampler.py
+3
@@ -8,6 +8,7 @@ from vllm.config.model import LogprobsMode
8
from vllm.triton_utils import tl, triton
9
from vllm.v1.outputs import LogprobsTensors, SamplerOutput
10
from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
11
+from vllm.v1.worker.gpu.penalties import apply_penalties
12
from vllm.v1.worker.gpu.states import SamplingMetadata
13
14
@@ -65,6 +66,8 @@ class Sampler:
66
logits = apply_top_k_top_p(
67
logits, sampling_metadata.top_k, sampling_metadata.top_p
68
)
69
+ # Apply penalties in place.
70
+ apply_penalties(logits, sampling_metadata)
71
72
sampled = gumbel_sample(
73
logits,
vllm/v1/worker/gpu/states.py
+172
-10
@@ -8,6 +8,8 @@ import torch
8
from vllm.lora.request import LoRARequest
9
from vllm.sampling_params import SamplingParams
10
from vllm.triton_utils import tl, triton
11
+from vllm.utils.platform_utils import is_uva_available
12
+from vllm.utils.torch_utils import get_cuda_view_from_cpu_tensor
13
from vllm.v1.outputs import LogprobsTensors
14
from vllm.v1.utils import CpuGpuBuffer
15
@@ -23,12 +25,21 @@ class SamplingMetadata:
25
top_p: torch.Tensor | None
26
top_k: torch.Tensor | None
27
28
+ repetition_penalty: torch.Tensor
29
+ frequency_penalty: torch.Tensor
30
+ presence_penalty: torch.Tensor
31
+
32
seeds: torch.Tensor
33
pos: torch.Tensor
34
35
# None means no logprobs, 0 means sampled token logprobs only
36
max_num_logprobs: int | None
37
38
+ # For penalties
39
+ idx_mapping: torch.Tensor
40
+ prompt_bin_counts: torch.Tensor
41
+ output_bin_counts: torch.Tensor
42
+
43
@classmethod
44
def make_dummy(
45
cls,
@@ -44,17 +55,35 @@ class SamplingMetadata:
55
# top_k = torch.full((num_reqs,), 20, dtype=torch.int32, device=device)
56
top_p = None
57
top_k = None
58
+ # NOTE(woosuk): We must set penalties to their default values to make sure
59
+ # the penalties kernel does not touch the placeholder bin_counts tensors.
60
+ repetition_penalty = torch.ones(num_reqs, dtype=torch.float32, device=device)
61
+ frequency_penalty = torch.zeros(num_reqs, dtype=torch.float32, device=device)
62
+ presence_penalty = torch.zeros(num_reqs, dtype=torch.float32, device=device)
63
seeds = torch.zeros(num_reqs, dtype=torch.int64, device=device)
64
pos = torch.zeros(num_reqs, dtype=torch.int64, device=device)
65
max_num_logprobs = 20
66
67
+ idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device)
68
+ # NOTE(woosuk): These are placeholder tensors to avoid None checks in the
69
+ # penalties kernel. We use 2 instead of 1 as vocab_size to avoid Triton
70
+ # specialization and re-compilation at runtime.
71
+ prompt_bin_counts = torch.zeros(num_reqs, 2, dtype=torch.int32, device=device)
72
+ output_bin_counts = torch.zeros(num_reqs, 2, dtype=torch.int32, device=device)
73
+
74
return cls(
75
temperature=temperature,
76
top_p=top_p,
77
top_k=top_k,
78
+ repetition_penalty=repetition_penalty,
79
+ frequency_penalty=frequency_penalty,
80
+ presence_penalty=presence_penalty,
81
seeds=seeds,
82
pos=pos,
83
max_num_logprobs=max_num_logprobs,
84
+ idx_mapping=idx_mapping,
85
+ prompt_bin_counts=prompt_bin_counts,
86
+ output_bin_counts=output_bin_counts,
87
)
88
89
@@ -83,9 +112,10 @@ class RequestState:
112
self.extra_data: dict[str, ExtraData] = {}
113
114
self.prompt_len = np.zeros(self.max_num_reqs, dtype=np.int32)
86
- self.prefill_token_ids = np.zeros(
87
- (self.max_num_reqs, self.max_model_len),
88
- dtype=np.int32,
115
+ # NOTE(woosuk): This tensor can be extremely large (e.g., several GBs)
116
+ # depending on the configured max_num_reqs and max_model_len.
117
+ self.prefill_token_ids = UvaBuffer(
118
+ self.max_num_reqs, self.max_model_len, dtype=torch.int32
119
)
120
self.prefill_len = self._make_buffer(self.max_num_reqs, dtype=torch.int32)
121
@@ -119,6 +149,9 @@ class RequestState:
149
self.temperature = self._make_param(self.max_num_reqs, torch.float32)
150
self.top_p = self._make_param(self.max_num_reqs, torch.float32)
151
self.top_k = self._make_param(self.max_num_reqs, torch.int32)
152
+ self.repetition_penalty = self._make_param(self.max_num_reqs, torch.float32)
153
+ self.frequency_penalty = self._make_param(self.max_num_reqs, torch.float32)
154
+ self.presence_penalty = self._make_param(self.max_num_reqs, torch.float32)
155
self.seeds = self._make_param(self.max_num_reqs, torch.int64)
156
157
self.num_logprobs = np.empty(self.max_num_reqs, dtype=np.int32)
@@ -126,6 +159,16 @@ class RequestState:
159
self.num_logprobs.fill(-1)
160
self.needs_prompt_logprobs = np.zeros(self.max_num_reqs, dtype=bool)
161
162
+ # Statistics for penalties.
163
+ # TODO(woosuk): These tensors are rarely used but can be extremely large.
164
+ # Optimize the memory usage.
165
+ self.prompt_bin_counts = torch.zeros(
166
+ self.max_num_reqs, self.vocab_size, dtype=torch.int32, device=self.device
167
+ )
168
+ self.output_bin_counts = torch.zeros(
169
+ self.max_num_reqs, self.vocab_size, dtype=torch.int32, device=self.device
170
+ )
171
+
172
def _make_param(self, size: int, dtype: torch.dtype) -> "Param":
173
return Param(size, dtype=dtype, device=self.device, pin_memory=self.pin_memory)
174
@@ -159,7 +202,7 @@ class RequestState:
202
f"prefill_len {prefill_len} < prompt_len {prompt_len}"
203
)
204
self.prefill_len.np[req_idx] = prefill_len
162
- self.prefill_token_ids[req_idx, :prefill_len] = prefill_token_ids
205
+ self.prefill_token_ids.np[req_idx, :prefill_len] = prefill_token_ids
206
207
self.num_computed_prefill_tokens[req_idx] = num_computed_tokens
208
# FIXME(woosuk): This triggers a GPU operation whenever adding a new request.
@@ -178,6 +221,18 @@ class RequestState:
221
else:
222
top_k = self.vocab_size
223
self.top_k.np[req_idx] = top_k
224
+ self.repetition_penalty.np[req_idx] = sampling_params.repetition_penalty
225
+ self.frequency_penalty.np[req_idx] = sampling_params.frequency_penalty
226
+ self.presence_penalty.np[req_idx] = sampling_params.presence_penalty
227
+
228
+ if use_penalty(sampling_params):
229
+ bincount(
230
+ self.prefill_token_ids.gpu[req_idx],
231
+ prefill_len,
232
+ prompt_len,
233
+ self.prompt_bin_counts[req_idx],
234
+ self.output_bin_counts[req_idx],
235
+ )
236
237
if sampling_params.seed is not None:
238
seed = sampling_params.seed
@@ -206,24 +261,32 @@ class RequestState:
261
262
def make_sampling_metadata(
263
self,
209
- idx_mapping: np.ndarray,
264
+ idx_mapping: torch.Tensor,
265
+ idx_mapping_np: np.ndarray,
266
pos: torch.Tensor,
267
) -> SamplingMetadata:
212
- temperature = self.temperature.np[idx_mapping]
268
+ temperature = self.temperature.np[idx_mapping_np]
269
temperature = self.temperature.copy_np_to_gpu(temperature)
270
215
- top_p = self.top_p.np[idx_mapping]
271
+ top_p = self.top_p.np[idx_mapping_np]
272
no_top_p = np.all(top_p == 1.0)
273
top_p = self.top_p.copy_np_to_gpu(top_p) if not no_top_p else None
274
219
- top_k = self.top_k.np[idx_mapping]
275
+ top_k = self.top_k.np[idx_mapping_np]
276
no_top_k = np.all(top_k == self.vocab_size)
277
top_k = self.top_k.copy_np_to_gpu(top_k) if not no_top_k else None
278
223
- seeds = self.seeds.np[idx_mapping]
279
+ rep_penalty = self.repetition_penalty.np[idx_mapping_np]
280
+ rep_penalty = self.repetition_penalty.copy_np_to_gpu(rep_penalty)
281
+ freq_penalty = self.frequency_penalty.np[idx_mapping_np]
282
+ freq_penalty = self.frequency_penalty.copy_np_to_gpu(freq_penalty)
283
+ pres_penalty = self.presence_penalty.np[idx_mapping_np]
284
+ pres_penalty = self.presence_penalty.copy_np_to_gpu(pres_penalty)
285
+
286
+ seeds = self.seeds.np[idx_mapping_np]
287
seeds = self.seeds.copy_np_to_gpu(seeds)
288
226
- num_logprobs = self.num_logprobs[idx_mapping]
289
+ num_logprobs = self.num_logprobs[idx_mapping_np]
290
max_num_logprobs: int | None = int(np.max(num_logprobs))
291
if max_num_logprobs == -1:
292
max_num_logprobs = None
@@ -232,9 +295,15 @@ class RequestState:
295
temperature=temperature,
296
top_p=top_p,
297
top_k=top_k,
298
+ repetition_penalty=rep_penalty,
299
+ frequency_penalty=freq_penalty,
300
+ presence_penalty=pres_penalty,
301
seeds=seeds,
302
pos=pos,
303
max_num_logprobs=max_num_logprobs,
304
+ idx_mapping=idx_mapping,
305
+ prompt_bin_counts=self.prompt_bin_counts,
306
+ output_bin_counts=self.output_bin_counts,
307
)
308
309
def expand_sampling_metadata(
@@ -294,6 +363,14 @@ class ExtraData:
363
in_progress_prompt_logprobs: list[LogprobsTensors] = field(default_factory=list)
364
365
366
+class UvaBuffer:
367
+ def __init__(self, *size: int | torch.SymInt, dtype: torch.dtype):
368
+ assert is_uva_available()
369
+ self.cpu = torch.zeros(*size, dtype=dtype, device="cpu", pin_memory=True)
370
+ self.np = self.cpu.numpy()
371
+ self.gpu = get_cuda_view_from_cpu_tensor(self.cpu)
372
+
373
+
374
# NOTE(woosuk): Re-compilation can happen at runtime since top_p and top_k can be None.
375
@triton.jit
376
def _expand_sampling_metadata_kernel(
@@ -304,6 +381,12 @@ def _expand_sampling_metadata_kernel(
381
top_k_ptr,
382
expanded_top_k_ptr,
383
seeds_ptr,
384
+ rep_penalty_ptr,
385
+ expanded_rep_penalty_ptr,
386
+ freq_penalty_ptr,
387
+ expanded_freq_penalty_ptr,
388
+ pres_penalty_ptr,
389
+ expanded_pres_penalty_ptr,
390
expanded_seeds_ptr,
391
cu_num_logits_ptr,
392
BLOCK_SIZE: tl.constexpr,
@@ -327,6 +410,15 @@ def _expand_sampling_metadata_kernel(
410
top_k = tl.load(top_k_ptr + req_idx)
411
tl.store(expanded_top_k_ptr + start_idx + block, top_k, mask=mask)
412
413
+ rep_penalty = tl.load(rep_penalty_ptr + req_idx)
414
+ tl.store(expanded_rep_penalty_ptr + start_idx + block, rep_penalty, mask=mask)
415
+
416
+ freq_penalty = tl.load(freq_penalty_ptr + req_idx)
417
+ tl.store(expanded_freq_penalty_ptr + start_idx + block, freq_penalty, mask=mask)
418
+
419
+ pres_penalty = tl.load(pres_penalty_ptr + req_idx)
420
+ tl.store(expanded_pres_penalty_ptr + start_idx + block, pres_penalty, mask=mask)
421
+
422
seed = tl.load(seeds_ptr + req_idx)
423
tl.store(expanded_seeds_ptr + start_idx + block, seed, mask=mask)
424
@@ -341,6 +433,9 @@ def expand_sampling_metadata(
433
expanded_temp = create_empty(sampling_metadata.temperature)
434
expanded_top_p = create_empty(sampling_metadata.top_p)
435
expanded_top_k = create_empty(sampling_metadata.top_k)
436
+ expanded_repetition_penalty = create_empty(sampling_metadata.repetition_penalty)
437
+ expanded_frequency_penalty = create_empty(sampling_metadata.frequency_penalty)
438
+ expanded_presence_penalty = create_empty(sampling_metadata.presence_penalty)
439
expanded_seeds = create_empty(sampling_metadata.seeds)
440
441
num_reqs = cu_num_logits.shape[0] - 1
@@ -351,6 +446,12 @@ def expand_sampling_metadata(
446
expanded_top_p,
447
sampling_metadata.top_k,
448
expanded_top_k,
449
+ sampling_metadata.repetition_penalty,
450
+ expanded_repetition_penalty,
451
+ sampling_metadata.frequency_penalty,
452
+ expanded_frequency_penalty,
453
+ sampling_metadata.presence_penalty,
454
+ expanded_presence_penalty,
455
sampling_metadata.seeds,
456
expanded_seeds,
457
cu_num_logits,
@@ -361,6 +462,67 @@ def expand_sampling_metadata(
462
top_p=expanded_top_p,
463
top_k=expanded_top_k,
464
seeds=expanded_seeds,
465
+ repetition_penalty=expanded_repetition_penalty,
466
+ frequency_penalty=expanded_frequency_penalty,
467
+ presence_penalty=expanded_presence_penalty,
468
pos=sampling_metadata.pos,
469
max_num_logprobs=sampling_metadata.max_num_logprobs,
470
+ # TODO(woosuk): Support penalties with spec decoding.
471
+ idx_mapping=sampling_metadata.idx_mapping,
472
+ prompt_bin_counts=sampling_metadata.prompt_bin_counts,
473
+ output_bin_counts=sampling_metadata.output_bin_counts,
474
+ )
475
+
476
+
477
+def use_penalty(sampling_params: SamplingParams) -> bool:
478
+ return (
479
+ sampling_params.repetition_penalty != 1.0
480
+ or sampling_params.frequency_penalty != 0.0
481
+ or sampling_params.presence_penalty != 0.0
482
+ )
483
+
484
+
485
+@triton.jit(do_not_specialize=["prefill_len", "prompt_len"])
486
+def _bincount_kernel(
487
+ prefill_token_ids_ptr,
488
+ prefill_len,
489
+ prompt_len,
490
+ prompt_bin_counts_ptr,
491
+ output_bin_counts_ptr,
492
+ BLOCK_SIZE: tl.constexpr,
493
+):
494
+ block_idx = tl.program_id(0)
495
+ if block_idx * BLOCK_SIZE >= prefill_len:
496
+ return
497
+
498
+ block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
499
+ if block_idx * BLOCK_SIZE < prompt_len:
500
+ mask = block < prompt_len
501
+ prefill_tokens = tl.load(prefill_token_ids_ptr + block, mask=mask)
502
+ tl.atomic_add(prompt_bin_counts_ptr + prefill_tokens, 1, mask=mask)
503
+ if (block_idx + 1) * BLOCK_SIZE >= prompt_len:
504
+ mask = block < prefill_len
505
+ mask &= block >= prompt_len
506
+ prefill_tokens = tl.load(prefill_token_ids_ptr + block, mask=mask)
507
+ tl.atomic_add(output_bin_counts_ptr + prefill_tokens, 1, mask=mask)
508
+
509
+
510
+def bincount(
511
+ prefill_token_ids: torch.Tensor,
512
+ prefill_len: int,
513
+ prompt_len: int,
514
+ prompt_bin_counts: torch.Tensor,
515
+ output_bin_counts: torch.Tensor,
516
+) -> None:
517
+ prompt_bin_counts.zero_()
518
+ output_bin_counts.zero_()
519
+ BLOCK_SIZE = 1024
520
+ num_blocks = triton.cdiv(prefill_len, BLOCK_SIZE)
521
+ _bincount_kernel[(num_blocks,)](
522
+ prefill_token_ids,
523
+ prefill_len,
524
+ prompt_len,
525
+ prompt_bin_counts,
526
+ output_bin_counts,
527
+ BLOCK_SIZE=BLOCK_SIZE,
528
)