Skip to content

GitLab

Unverified Commit a463555d authored by Jevin Jiang's avatar Jevin Jiang Committed by GitHub
Browse files

[TPU] Fix the test_sampler (#17820)

parent ca04b97c
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
tests/v1/tpu/test_sampler.py
View file @ a463555d
...@@ -26,7 +26,7 @@ def test_sampler_different(model_name: str): ...@@ -26,7 +26,7 @@ def test_sampler_different(model_name: str):
enforce_eager=False, enforce_eager=False,
max_num_seqs=1, max_num_seqs=1,
max_model_len=512, max_model_len=512,
max_num_batched_tokens=512) max_num_batched_tokens=256)
prompts = [ prompts = [
"Write a short story about a robot that dreams for the first time." "Write a short story about a robot that dreams for the first time."
] ]
......
vllm/v1/attention/backends/pallas.py
View file @ a463555d
...@@ -95,7 +95,7 @@ class PallasMetadata: ...@@ -95,7 +95,7 @@ class PallasMetadata:
block_tables: torch.Tensor block_tables: torch.Tensor
context_lens: torch.Tensor context_lens: torch.Tensor
query_start_loc: torch.Tensor query_start_loc: torch.Tensor
num_seqs: int num_seqs: torch.Tensor
class PallasAttentionBackendImpl(AttentionImpl): class PallasAttentionBackendImpl(AttentionImpl):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment