[ModelRunner V2] Bound memory for large logprobs requests (#46746)

Signed-off-by: Nick Hill <nickhill123@gmail.com>

vllm/v1/worker/gpu/sample/logprob.py

@@ -9,6 +9,9 @@ from vllm.triton_utils import tl, triton
 from vllm.v1.outputs import LogprobsTensors
 from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor
 
+# Upper bound on the topk kernel's per-iteration gather width.
+_MAX_TOPK_BLOCK = 1024
+
 
 @triton.jit
 def _topk_log_softmax_kernel(
@@ -19,7 +22,7 @@ def _topk_log_softmax_kernel(
     topk,
     vocab_size,
     BLOCK_SIZE: tl.constexpr,
-    PADDED_TOPK: tl.constexpr,
+    TOPK_BLOCK_SIZE: tl.constexpr,
 ):
     req_idx = tl.program_id(0).to(tl.int64)
     row_ptr = logits_ptr + req_idx * logits_stride
@@ -42,14 +45,16 @@ def _topk_log_softmax_kernel(
         se += tl.sum(e)
     lse = tl.log(se)
 
-    k_offset = tl.arange(0, PADDED_TOPK)
-    k_mask = k_offset < topk
-    topk_ids = tl.load(topk_ids_ptr + req_idx * topk + k_offset, mask=k_mask, other=0)
-
-    logits = tl.load(row_ptr + topk_ids, mask=k_mask)
-    logits = logits.to(tl.float32)
-    o = logits - max_val - lse
-    tl.store(output_ptr + req_idx * topk + k_offset, o, mask=k_mask)
+    for j in range(0, topk, TOPK_BLOCK_SIZE):
+        k_offset = j + tl.arange(0, TOPK_BLOCK_SIZE)
+        k_mask = k_offset < topk
+        topk_ids = tl.load(
+            topk_ids_ptr + req_idx * topk + k_offset, mask=k_mask, other=0
+        )
+        logits = tl.load(row_ptr + topk_ids, mask=k_mask)
+        logits = logits.to(tl.float32)
+        o = logits - max_val - lse
+        tl.store(output_ptr + req_idx * topk + k_offset, o, mask=k_mask)
 
 
 @triton.jit
@@ -85,6 +90,9 @@ def compute_token_logprobs(
     token_ids = token_ids.to(torch.int64)
     num_logprobs = token_ids.shape[1]
     logprobs = logits.new_empty((batch_size, num_logprobs), dtype=torch.float32)
+    # Cap the kernel's per-iteration width so very large num_logprobs requests
+    # stream the gather in bounded-size chunks, avoiding excessive mem use.
+    topk_block_size = min(triton.next_power_of_2(num_logprobs), _MAX_TOPK_BLOCK)
     _topk_log_softmax_kernel[(batch_size,)](
         logprobs,
         logits,
@@ -93,7 +101,7 @@ def compute_token_logprobs(
         num_logprobs,
         vocab_size,
         BLOCK_SIZE=1024,  # type: ignore
-        PADDED_TOPK=triton.next_power_of_2(num_logprobs),
+        TOPK_BLOCK_SIZE=topk_block_size,
     )
     return logprobs