[CI] Batch the swap_blocks verification instead of copying block by block (#54558)
Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>
stefankoncarevic committed
Sep 2, 2026 at 10:12 UTC
62588e0592ad5af8d3d4a536ace01afb44ebaed5
1 file changed
+11
-7
tests/kernels/attention/test_cache.py
+11
-7
@@ -646,13 +646,17 @@ def test_swap_blocks(
646
block_mapping_tensor,
647
)
648
649
- for src, dst in block_mapping:
650
- torch.testing.assert_close(
651
- src_key_caches_clone[src].cpu(), dist_key_caches[0][dst].cpu()
652
- )
653
- torch.testing.assert_close(
654
- src_value_caches_clone[src].cpu(), dist_value_caches[0][dst].cpu()
655
- )
649
+ src_indices = block_mapping_tensor[:, 0].to(src_key_caches_clone.device)
650
+ dst_indices = block_mapping_tensor[:, 1].to(dist_key_caches[0].device)
651
+
652
+ torch.testing.assert_close(
653
+ src_key_caches_clone.index_select(0, src_indices).cpu(),
654
+ dist_key_caches[0].index_select(0, dst_indices).cpu(),
655
+ )
656
+ torch.testing.assert_close(
657
+ src_value_caches_clone.index_select(0, src_indices).cpu(),
658
+ dist_value_caches[0].index_select(0, dst_indices).cpu(),
659
+ )
660
661
662
@pytest.mark.parametrize("num_heads", NUM_HEADS)