[ModelRunner V2] Bound memory for large logprobs requests (#46746)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Jun 25, 2026 at 18:04 UTC
32bb3195f0b93f6971781479591f7a6ee666e7dc
1 file changed
+18
-10
vllm/v1/worker/gpu/sample/logprob.py
+18
-10
@@ -9,6 +9,9 @@ from vllm.triton_utils import tl, triton
9
from vllm.v1.outputs import LogprobsTensors
10
from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor
11
12
+# Upper bound on the topk kernel's per-iteration gather width.
13
+_MAX_TOPK_BLOCK = 1024
14
+
15
16
@triton.jit
17
def _topk_log_softmax_kernel(
@@ -19,7 +22,7 @@ def _topk_log_softmax_kernel(
22
topk,
23
vocab_size,
24
BLOCK_SIZE: tl.constexpr,
22
- PADDED_TOPK: tl.constexpr,
25
+ TOPK_BLOCK_SIZE: tl.constexpr,
26
):
27
req_idx = tl.program_id(0).to(tl.int64)
28
row_ptr = logits_ptr + req_idx * logits_stride
@@ -42,14 +45,16 @@ def _topk_log_softmax_kernel(
45
se += tl.sum(e)
46
lse = tl.log(se)
47
45
- k_offset = tl.arange(0, PADDED_TOPK)
46
- k_mask = k_offset < topk
47
- topk_ids = tl.load(topk_ids_ptr + req_idx * topk + k_offset, mask=k_mask, other=0)
48
-
49
- logits = tl.load(row_ptr + topk_ids, mask=k_mask)
50
- logits = logits.to(tl.float32)
51
- o = logits - max_val - lse
52
- tl.store(output_ptr + req_idx * topk + k_offset, o, mask=k_mask)
48
+ for j in range(0, topk, TOPK_BLOCK_SIZE):
49
+ k_offset = j + tl.arange(0, TOPK_BLOCK_SIZE)
50
+ k_mask = k_offset < topk
51
+ topk_ids = tl.load(
52
+ topk_ids_ptr + req_idx * topk + k_offset, mask=k_mask, other=0
53
+ )
54
+ logits = tl.load(row_ptr + topk_ids, mask=k_mask)
55
+ logits = logits.to(tl.float32)
56
+ o = logits - max_val - lse
57
+ tl.store(output_ptr + req_idx * topk + k_offset, o, mask=k_mask)
58
59
60
@triton.jit
@@ -85,6 +90,9 @@ def compute_token_logprobs(
90
token_ids = token_ids.to(torch.int64)
91
num_logprobs = token_ids.shape[1]
92
logprobs = logits.new_empty((batch_size, num_logprobs), dtype=torch.float32)
93
+ # Cap the kernel's per-iteration width so very large num_logprobs requests
94
+ # stream the gather in bounded-size chunks, avoiding excessive mem use.
95
+ topk_block_size = min(triton.next_power_of_2(num_logprobs), _MAX_TOPK_BLOCK)
96
_topk_log_softmax_kernel[(batch_size,)](
97
logprobs,
98
logits,
@@ -93,7 +101,7 @@ def compute_token_logprobs(
101
num_logprobs,
102
vocab_size,
103
BLOCK_SIZE=1024, # type: ignore
96
- PADDED_TOPK=triton.next_power_of_2(num_logprobs),
104
+ TOPK_BLOCK_SIZE=topk_block_size,
105
)
106
return logprobs
107