diff --git a/vllm/v1/worker/gpu/sample/logprob.py b/vllm/v1/worker/gpu/sample/logprob.py index 0028e8c3a9..cb2cf1a590 100644 --- a/vllm/v1/worker/gpu/sample/logprob.py +++ b/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