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, 18:30 UTC
[CI/Build] Fix illegal memory access and unsupported test in kernels/attention/test_cache.py (#29118)
Signed-off-by: Randall Smith <ransmith@amd.com> Co-authored-by: Randall Smith <ransmith@amd.com>
rasmith committed
Nov 21, 2025 at 09:58 UTC
711241c13cf9c1e543a1948bb25a40623f3da78c
1 file changed
+9
tests/kernels/attention/test_cache.py
+9
@@ -68,6 +68,7 @@ def test_copy_blocks(
68
pytest.skip()
69
current_platform.seed_everything(seed)
70
torch.set_default_device(device)
71
+ torch.cuda.set_device(device)
72
# Generate random block mappings where each source block is mapped to two
73
# destination blocks.
74
assert 2 * num_mappings <= num_blocks
@@ -152,6 +153,7 @@ def test_reshape_and_cache(
153
pytest.skip()
154
current_platform.seed_everything(seed)
155
torch.set_default_device(device)
156
+ torch.cuda.set_device(device)
157
# Create a random slot mapping.
158
num_slots = block_size * num_blocks
159
slot_mapping_lst = random.sample(range(num_slots), num_tokens)
@@ -272,6 +274,7 @@ def test_reshape_and_cache_flash(
274
) -> None:
275
current_platform.seed_everything(seed)
276
torch.set_default_device(device)
277
+ torch.cuda.set_device(device)
278
assert implementation in ["cuda", "triton"]
279
if implementation == "triton" and kv_cache_layout == "HND":
280
pytest.skip("Triton implementation only supports NHD layout.")
@@ -593,6 +596,7 @@ def test_concat_and_cache_mla(
596
) -> None:
597
current_platform.seed_everything(seed)
598
torch.set_default_device(device)
599
+ torch.cuda.set_device(device)
600
601
total_slots = num_blocks * block_size
602
slot_mapping_lst = random.sample(range(total_slots), num_tokens)
@@ -662,11 +666,14 @@ def test_concat_and_cache_ds_mla(
666
seed: int,
667
device: str,
668
) -> None:
669
+ if current_platform.is_rocm():
670
+ pytest.skip("concat_and_cache_mla doesn't support fp8_ds_mla on ROCm")
671
if dtype.itemsize != 2:
672
pytest.skip("ds_mla only supports 16-bit input")
673
kv_cache_dtype = "fp8_ds_mla"
674
current_platform.seed_everything(seed)
675
torch.set_default_device(device)
676
+ torch.cuda.set_device(device)
677
678
total_slots = num_blocks * block_size
679
slot_mapping_lst = random.sample(range(total_slots), num_tokens)
@@ -779,6 +786,7 @@ def test_copy_blocks_mla(
786
) -> None:
787
current_platform.seed_everything(seed)
788
torch.set_default_device(device)
789
+ torch.cuda.set_device(device)
790
791
entry_size = kv_lora_rank + qk_rope_head_dim
792
@@ -843,6 +851,7 @@ def test_swap_blocks_mla(
851
) -> None:
852
current_platform.seed_everything(seed)
853
torch.set_default_device(device)
854
+ torch.cuda.set_device(device)
855
856
entry_size = kv_lora_rank + qk_rope_head_dim
857