[bugfix] fix pyproject install and VSA precision test (#726)

.buildkite/pipeline.yml

@@ -163,6 +163,7 @@ steps:
             - path:
                 - "csrc/attn/vsa/**"
                 - "csrc/attn/tk/**"
+                - "csrc/attn/tests/test_vsa.py"
                 - "csrc/attn/setup_vsa.py"
                 - "csrc/attn/config_vsa.py"
                 - "csrc/attn/vsa.cpp"

csrc/attn/README.md

@@ -84,7 +84,7 @@ out = sliding_tile_attention(q, k, v, window_size, 0, False)
 ### Test
 ```bash
 python tests/test_sta.py # test STA
-python tests/test_block_sparse.py # test VSA
+python tests/test_vsa.py # test VSA
 ```
 ### Benchmark
 ```bash

fastvideo/pipelines/pipeline_batch_info.py

@@ -9,11 +9,13 @@ in a functional manner, reducing the need for explicit parameter passing.
 
 import pprint
 from dataclasses import asdict, dataclass, field
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import PIL.Image
 import torch
-from torchcodec.decoders import VideoDecoder
+
+if TYPE_CHECKING:
+    from torchcodec.decoders import VideoDecoder
 
 from fastvideo.attention import AttentionMetadata
 from fastvideo.configs.sample.teacache import TeaCacheParams, WanTeaCacheParams
@@ -198,5 +200,5 @@ class TrainingBatch:
 
 @dataclass
 class PreprocessBatch(ForwardBatch):
-    video_loader: list[VideoDecoder] = field(default_factory=list)
+    video_loader: list["VideoDecoder"] = field(default_factory=list)
     video_file_name: list[str] = field(default_factory=list)

fastvideo/tests/modal/pr_test.py

@@ -52,7 +52,7 @@ def run_test(pytest_command: str):
     git clone {git_repo} /FastVideo &&
     cd /FastVideo &&
     {checkout_command} &&
-    uv pip install -e .[test] --index-strategy unsafe-best-match &&
+    uv pip install -e .[test] &&
     {pytest_command}
     """
     
@@ -100,7 +100,7 @@ def run_precision_tests_STA():
 
 @app.function(gpu="H100:1", image=image, timeout=900)
 def run_precision_tests_VSA():
-    run_test("python csrc/attn/tests/test_block_sparse.py")
+    run_test("python csrc/attn/tests/test_vsa.py")
 
 @app.function(gpu="L40S:1", image=image, timeout=3600)
 def run_inference_lora_tests():

pyproject.toml

@@ -53,7 +53,27 @@ dependencies = [
 ]
 
 [tool.uv]
-extra-index-url = ["https://download.pytorch.org/whl/cu128"]
+prerelease = "allow"
+
+[tool.uv.sources]
+torch = [
+  { index = "pytorch-cpu", marker = "sys_platform != 'linux'" },
+  { index = "pytorch-cu128", marker = "sys_platform == 'linux'" },
+]
+torchvision = [
+  { index = "pytorch-cpu", marker = "sys_platform != 'linux'" },
+  { index = "pytorch-cu128", marker = "sys_platform == 'linux'" },
+]
+
+[[tool.uv.index]]
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true
+
+[[tool.uv.index]]
+name = "pytorch-cu128"
+url = "https://download.pytorch.org/whl/cu128"
+explicit = true
 
 [project.optional-dependencies]