[kernel] Extract VSA utility functions into fastvideo_kernel

[freemty]

Summary

• Extract 4 tile-partition/variable-block-size utility functions from fastvideo/attention/backends/video_sparse_attn.py into standalone fastvideo_kernel/vsa_utils.py
• Add build_vsa_metadata() convenience entry point for one-call metadata construction
• Enable external users to call VSA kernels without depending on the full fastvideo framework (pip install fastvideo-kernel is now self-contained)

Motivation

The VSA kernel lives in fastvideo_kernel (independently installable), but calling it requires helper functions buried in the fastvideo framework layer. This forces external users to install the entire framework just for metadata construction. See #782.

Changes

File	Description
fastvideo_kernel/vsa_utils.py	New module: 5 functions + VSA_TILE_SIZE constant (~147 lines)
fastvideo_kernel/__init__.py	Re-export new utilities
tests/test_vsa_utils.py	20 CPU-only tests covering all functions + consistency with framework

Design decisions

• Parameterized tile_size: construct_variable_block_sizes accepts tile_size as argument instead of hardcoding VSA_TILE_SIZE, enabling non-default tile configurations
• lru_cache on pure functions: index computations are deterministic for a given shape — caching avoids redundant recomputation
• CPU-only tests: utilities are pure PyTorch index math, no GPU/kernel required
• Consistency test: verifies output matches the framework-level functions exactly (skipped if fastvideo is not installed)

Test plan

• pytest fastvideo-kernel/tests/test_vsa_utils.py -v — 20 tests, all pass (CPU-only)
• CI pre-commit checks
• Verify framework-level import still works after downstream refactor (future PR)

Closes #782

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI

	Protection	Waiting on
🔴	PR merge requirements	👀 reviews and 🤖 CI

🔴 PR merge requirements

Waiting for

• #approved-reviews-by>=1
• check-success=fastcheck-passed
• check-success=full-suite-passed

This rule is failing.

• #approved-reviews-by>=1
• check-success=fastcheck-passed
• check-success=full-suite-passed
• check-success~=pre-commit
• title~=(?i)^\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model|skill|skills|infra)\]

[gemini-code-assist]

Code Review

This pull request introduces standalone VSA metadata utilities in vsa_utils.py along with comprehensive unit tests in test_vsa_utils.py and exports them in the package initialization. A critical issue was identified in vsa_utils.py where @functools.lru_cache is applied to get_non_pad_index. Since this function accepts a torch.Tensor which is unhashable, this will raise a TypeError at runtime and must be removed.

[gemini-code-assist]

Applying @functools.lru_cache to get_non_pad_index will raise a TypeError: unhashable type: 'Tensor' at runtime because variable_block_sizes is a torch.Tensor, which is mutable and not hashable in PyTorch.

Since get_non_pad_index is a lightweight operation and is only called once per metadata construction, caching it is unnecessary and causes immediate failures. Please remove the @functools.lru_cache decorator from this function.

def get_non_pad_index(
    variable_block_sizes: torch.LongTensor,
    max_block_size: int,
) -> torch.LongTensor:

[alexzms]

Quick question on the dict keys returned by build_vsa_metadata() in fastvideo-kernel/python/fastvideo_kernel/vsa_utils.py:

• The dict uses tile_indices and reverse_tile_indices.
• The framework VideoSparseAttentionMetadata dataclass in fastvideo/attention/backends/video_sparse_attn.py (lines 138–139) uses the longer names tile_partition_indices and reverse_tile_partition_indices.

Since the follow-up PR will swap the framework over to importing from here, was there a reason to shorten the names? If not, aligning them now (using tile_partition_indices / reverse_tile_partition_indices in the returned dict too) would let the follow-up be a clean import ... as ... swap rather than touching every call site that destructures the metadata.

Happy to be wrong if you meant this as a cleaner external-facing API — just want to flag it before merge.

[freemty]

Good catch! You're right — aligning the dict keys now avoids unnecessary churn in the follow-up PR.

Updated build_vsa_metadata() to return tile_partition_indices / reverse_tile_partition_indices (matching VideoSparseAttentionMetadata). Also removed the @lru_cache on get_non_pad_index since it accepts a torch.Tensor which is unhashable (would raise TypeError at runtime).

Tests updated and passing.

[mergify]

This PR has merge conflicts with the base branch. Please rebase:

git fetch origin main
git rebase origin/main
# Resolve any conflicts, then:
git push --force-with-lease