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 15, 17:15 UTC
[V1] TPU - Fix the chunked prompt bug (#15713)
Signed-off-by: Alexander Matveev <amatveev@redhat.com>
Alexander Matveev committed
Mar 28, 2025 at 16:19 UTC
c3f687ac227f93cfb8a5558da871d9dfa68095ab
2 files changed
+17
-1
tests/v1/tpu/test_basic.py
+4
-1
@@ -48,7 +48,10 @@ def test_models(
48
49
with vllm_runner(
50
model,
51
- max_model_len=8192,
51
+ # Note: max_num_batched_tokens == 1024 is needed here to
52
+ # actually test chunked prompt
53
+ max_num_batched_tokens=1024,
54
+ max_model_len=8196,
55
gpu_memory_utilization=0.7,
56
max_num_seqs=16,
57
tensor_parallel_size=tensor_parallel_size) as vllm_model:
vllm/v1/worker/tpu_model_runner.py
+13
@@ -618,6 +618,7 @@ class TPUModelRunner:
618
# Update the cache state concurrently. Code above will not block until
619
# we use `selected_token_ids`. Add mark_step if post-processing changes
620
request_seq_lens: list[tuple[int, CachedRequestState, int]] = []
621
+ discard_sampled_tokens_req_indices = []
622
for i, req_id in zip(range(num_reqs), self.input_batch.req_ids):
623
assert req_id is not None
624
req_state = self.requests[req_id]
@@ -633,6 +634,10 @@ class TPUModelRunner:
634
# This relies on cuda-specific torch-internal impl details
635
generator.set_offset(generator.get_offset() - 4)
636
637
+ # Record the index of the request that should not be sampled,
638
+ # so that we could clear the sampled tokens before returning.
639
+ discard_sampled_tokens_req_indices.append(i)
640
+
641
assert all(
642
req_id is not None for req_id in
643
self.input_batch.req_ids[:num_reqs]), "req_ids contains None"
@@ -646,11 +651,19 @@ class TPUModelRunner:
651
if max_gen_len == 1:
652
valid_sampled_token_ids = selected_token_ids.tolist()
653
654
+ # Mask out the sampled tokens that should not be sampled.
655
+ # TODO: Keep in sync with gpu_model_runner.py, in particular
656
+ # the "else" case here
657
+ for i in discard_sampled_tokens_req_indices:
658
+ valid_sampled_token_ids[i].clear()
659
+
660
+ # Append sampled tokens
661
for i, req_state, seq_len in request_seq_lens:
662
token_id = valid_sampled_token_ids[i][0]
663
self.input_batch.token_ids_cpu[i, seq_len] = token_id
664
req_state.output_token_ids.append(token_id)
665
self.input_batch.num_tokens[i] += 1
666
+
667
else:
668
valid_mask = selected_token_ids != INVALID_TOKEN_ID
669
gen_lens = valid_mask.sum(dim=1).tolist()