fix: align dict keys with framework naming and remove invalid lru_cache

- Rename build_vsa_metadata() keys to tile_partition_indices /
  reverse_tile_partition_indices to match VideoSparseAttentionMetadata
- Remove @lru_cache from get_non_pad_index (torch.Tensor is unhashable)
- Update tests accordingly

Author: freemty

fastvideo-kernel/python/fastvideo_kernel/vsa_utils.py

@@ -85,7 +85,6 @@ def _sizes(dim_len: int, tile: int, n: int) -> torch.LongTensor:
     ).reshape(-1)
 
 
-@functools.lru_cache(maxsize=10)
 def get_non_pad_index(
     variable_block_sizes: torch.LongTensor,
     max_block_size: int,
@@ -116,7 +115,7 @@ def build_vsa_metadata(
         device: Target device for index tensors.
 
     Returns:
-        Dict with keys: tile_indices, reverse_tile_indices,
+        Dict with keys: tile_partition_indices, reverse_tile_partition_indices,
         variable_block_sizes, non_pad_index, num_tiles, max_block_size.
     """
     if isinstance(device, str):
@@ -137,8 +136,8 @@ def build_vsa_metadata(
     npi = get_non_pad_index(vbs, max_block_size)
 
     return {
-        "tile_indices": tile_indices,
-        "reverse_tile_indices": reverse_tile_indices,
+        "tile_partition_indices": tile_indices,
+        "reverse_tile_partition_indices": reverse_tile_indices,
         "variable_block_sizes": vbs,
         "non_pad_index": npi,
         "num_tiles": num_tiles,

fastvideo-kernel/tests/test_vsa_utils.py

@@ -160,7 +160,7 @@ def test_all_keys_present(self):
         """build_vsa_metadata returns all expected keys."""
         meta = build_vsa_metadata((8, 16, 16), device="cpu")
         expected_keys = {
-            "tile_indices", "reverse_tile_indices",
+            "tile_partition_indices", "reverse_tile_partition_indices",
             "variable_block_sizes", "non_pad_index",
             "num_tiles", "max_block_size",
         }
@@ -169,8 +169,8 @@ def test_all_keys_present(self):
     def test_types(self):
         """Return types are correct."""
         meta = build_vsa_metadata((8, 16, 16), device="cpu")
-        assert isinstance(meta["tile_indices"], torch.Tensor)
-        assert isinstance(meta["reverse_tile_indices"], torch.Tensor)
+        assert isinstance(meta["tile_partition_indices"], torch.Tensor)
+        assert isinstance(meta["reverse_tile_partition_indices"], torch.Tensor)
         assert isinstance(meta["variable_block_sizes"], torch.Tensor)
         assert isinstance(meta["non_pad_index"], torch.Tensor)
         assert isinstance(meta["num_tiles"], tuple)
@@ -191,8 +191,8 @@ def test_consistency(self):
         shape = (8, 16, 16)
         meta = build_vsa_metadata(shape, device="cpu")
         n = math.prod(shape)
-        assert meta["tile_indices"].shape == (n,)
-        assert meta["reverse_tile_indices"].shape == (n,)
+        assert meta["tile_partition_indices"].shape == (n,)
+        assert meta["reverse_tile_partition_indices"].shape == (n,)
         assert meta["variable_block_sizes"].sum().item() == n
         assert meta["non_pad_index"].shape[0] == n