Replace Qwen3.5 FLA paths with Attention Gym

.ci/docker/requirements-vlm.txt

@@ -2,4 +2,3 @@ av
 einops
 pillow
 torchvision
-flash-linear-attention

.ci/docker/requirements.txt

@@ -8,4 +8,4 @@ safetensors
 einops
 pillow
 spmd_types==0.2.5
-attn-gym[linear]==0.0.5
+attn-gym[linear]==0.0.8

.github/workflows/integration_test_8gpu_rl.yaml

@@ -80,14 +80,10 @@ jobs:
           --extra-index-url https://download.pytorch.org/whl/nightly/cu130 \
           --index-strategy unsafe-best-match
 
-        # 4. Install the pinned GDN kernels without replacing nightly PyTorch.
-        uv pip install --no-deps \
-          "git+https://github.com/fla-org/flash-linear-attention.git@v0.5.2"
-
-        # 5. Make the checkout importable for subprocesses spawned by the test.
+        # 4. Make the checkout importable for subprocesses spawned by the test.
         export PYTHONPATH="$PWD:${PYTHONPATH:-}"
 
-        # 6. Download HF model checkpoint for tests
+        # 5. Download HF model checkpoint for tests
         MODEL_PATH=$(python -c "from huggingface_hub import snapshot_download; print(snapshot_download('Qwen/Qwen3-0.6B'))")
 
         sudo mkdir -p "$RUNNER_TEMP/artifacts-to-be-uploaded"

pyproject.toml

@@ -28,7 +28,7 @@ dependencies = [
     "einops",
     "pillow",
     "spmd_types==0.2.5",
-    "attn-gym[linear]==0.0.5",
+    "attn-gym[linear]==0.0.8",
 ]
 dynamic = ["version"]
 
@@ -77,5 +77,5 @@ markers = [
 [tool.pyrefly]
 python-version = "3.11"
 project-excludes = ["torchtitan/experiments", "**/tests/**"]
-replace-imports-with-any = ["torchao.*", "torchft", "torchvision.*", "deep_ep.*", "jinja2.*", "fla.*", "helion", "helion.*", "batch_invariant_ops", "torchcomms"]  # optional dependencies
+replace-imports-with-any = ["torchao.*", "torchft", "torchvision.*", "deep_ep.*", "jinja2.*", "helion", "helion.*", "batch_invariant_ops", "torchcomms"]  # optional dependencies
 search-path = ["../pytorch"]  # local built pytorch

scripts/ci/pytorch_ci_test_runner.sh

@@ -41,9 +41,7 @@ case "$COMMAND" in
     model_tests)
         # qwen3_5_fsdp+tp+varlen_attn+per_op_sac: varlen attention needs
         # flash_attn_interface/FA3, which the PyTorch CI image does not install
-        # and which is unavailable on its A10G (sm86) runners anyway. Excluded
-        # here rather than disabled in models.py so torchtitan's own CI, whose
-        # image ships flash-linear-attention, keeps running it.
+        # and which is unavailable on its A10G (sm86) runners anyway.
         python -m tests.integration_tests.run_tests \
             --test_suite models \
             --exclude "qwen3_5_moe_fsdp+tp+ep+pp,qwen3_5_fsdp+tp+varlen_attn+per_op_sac" \

tests/unit_tests/cpu/test_qwen3_8.py

@@ -10,7 +10,7 @@
 import pytest
 import torch
 
-pytest.importorskip("fla")
+pytest.importorskip("attn_gym")
 
 from torchtitan.models.qwen3_5 import Qwen35Model, Qwen35StateDictAdapter
 from torchtitan.models.qwen3_5.sharding import set_qwen35_sharding_config

tests/unit_tests/cpu/test_varlen_attention.py

@@ -29,17 +29,13 @@
 class TestPackedVarlenMetadata(unittest.TestCase):
     def test_document_boundaries(self):
         positions_T = torch.tensor([0, 1, 2, 0, 1, 0, 1, 2, 3])
-        metadata = create_varlen_metadata_for_document(
-            positions_T,
-            include_host_offsets=True,
-        )
+        metadata = create_varlen_metadata_for_document(positions_T)
 
         expected_cu_seq = torch.tensor([0, 3, 5, 9], dtype=torch.int32)
         torch.testing.assert_close(metadata.cu_seq_q, expected_cu_seq)
         torch.testing.assert_close(metadata.cu_seq_k, expected_cu_seq)
         self.assertEqual(metadata.max_q, 4)
         self.assertEqual(metadata.max_k, 4)
-        self.assertEqual(metadata.cu_seq_q_host, (0, 3, 5, 9))
 
 
 class TestPackedVarlenAttention(unittest.TestCase):