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 12, 22:15 UTC
[CI/Build] Fix test_prefix_prefill for AMD (#28905)
Signed-off-by: Ryan Rock <ryan.rock@amd.com>
Ryan Rock committed
Nov 19, 2025 at 15:04 UTC
68d7231991cc307d6865eac5bfca551c06f67465
1 file changed
+6
-6
tests/kernels/attention/test_prefix_prefill.py
+6
-6
@@ -174,11 +174,11 @@ def test_contexted_kv_attention(
174
block_table = values[: BS * max_block_per_request].view(BS, max_block_per_request)
175
b_seq_len = torch.tensor(seq_lens, dtype=torch.int32)
176
b_ctx_len = torch.tensor(ctx_lens, dtype=torch.int32)
177
- b_start_loc = torch.cumsum(torch.tensor([0] + query_lens, dtype=torch.int32), dim=0)
177
+ b_start_loc = torch.cumsum(torch.tensor([0] + query_lens), dim=0).to(torch.int32)
178
max_input_len = MAX_SEQ_LEN
179
# copy kv to cache
180
- b_seq_start_loc = torch.cumsum(
181
- torch.tensor([0] + seq_lens[:-1], dtype=torch.int32), dim=0
180
+ b_seq_start_loc = torch.cumsum(torch.tensor([0] + seq_lens[:-1]), dim=0).to(
181
+ torch.int32
182
)
183
for i in range(BS):
184
for j in range(query_lens[i]):
@@ -417,11 +417,11 @@ def test_contexted_kv_attention_alibi(
417
block_table = values[: BS * max_block_per_request].view(BS, max_block_per_request)
418
b_seq_len = torch.tensor(seq_lens, dtype=torch.int32)
419
b_ctx_len = torch.tensor(ctx_lens, dtype=torch.int32)
420
- b_start_loc = torch.cumsum(torch.tensor([0] + query_lens, dtype=torch.int32), dim=0)
420
+ b_start_loc = torch.cumsum(torch.tensor([0] + query_lens), dim=0).to(torch.int32)
421
max_input_len = MAX_SEQ_LEN
422
# copy kv to cache
423
- b_seq_start_loc = torch.cumsum(
424
- torch.tensor([0] + seq_lens[:-1], dtype=torch.int32), dim=0
423
+ b_seq_start_loc = torch.cumsum(torch.tensor([0] + seq_lens[:-1]), dim=0).to(
424
+ torch.int32
425
)
426
for i in range(BS):
427
for j in range(query_lens[i]):