[tx] Cuda tile for expert parallelism #880

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft

agolajko wants to merge 27 commits into NovaSky-AI:main from agolajko:feature/cutile-ep

skyrl-tx/pyproject.toml

-Original file line number
+Diff line change
@@ Expand Up / @@ -28,6 +28,8 @@ dependencies = [ @@
     [project.optional-dependencies]
     gpu = [
         "jax[cuda12]>=0.7.2",
+        "torch>=2.0.0",
+        "cuda-tile",  # cuTile Python from PyPI (requires CUDA Toolkit 13.1+)
     ]
     tpu = [
@@ Expand Down @@

skyrl-tx/tests/README.md

-Original file line number
+Diff line change
@@ Expand Up / @@ -5,3 +5,6 @@ uv run --extra dev --extra tinker pytest -v @@
     ## Run Specific tests
     uv run --extra dev --extra tinker pytest -v -s tests/models/test_qwen3_generate.py::test_qwen3_generate_speed
+    ## Run Cutile tests
+    TX_USE_CUTILE_LORA=1 uv run pytest tests/cutile/test_cutile_lora_equivalence.py -v

Provide feedback