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, 11:00 UTC
[TPU] Fix the test_sampler (#17820)
Jevin Jiang committed
May 8, 2025 at 02:51 UTC
a463555dee7cd98f88ce1282a730d0f857ad3216
2 files changed
+2
-2
tests/v1/tpu/test_sampler.py
+1
-1
@@ -26,7 +26,7 @@ def test_sampler_different(model_name: str):
26
enforce_eager=False,
27
max_num_seqs=1,
28
max_model_len=512,
29
- max_num_batched_tokens=512)
29
+ max_num_batched_tokens=256)
30
prompts = [
31
"Write a short story about a robot that dreams for the first time."
32
]
vllm/v1/attention/backends/pallas.py
+1
-1
@@ -95,7 +95,7 @@ class PallasMetadata:
95
block_tables: torch.Tensor
96
context_lens: torch.Tensor
97
query_start_loc: torch.Tensor
98
- num_seqs: int
98
+ num_seqs: torch.Tensor
99
100
101
class PallasAttentionBackendImpl(AttentionImpl):