[kimi k3] Allow Attention Gym KDA on CUDA capability 9.0+

tests/unit_tests/gpu/test_kda_attention.py

@@ -15,10 +15,10 @@
 from torchtitan.models.common.attention import create_varlen_metadata_for_document
 from torchtitan.models.kimi_k3.kda import InnerKDA, KDA, KDAKernel, KimiRMSNormGated
 
-_HAS_BLACKWELL = (
+_HAS_ATTENTION_GYM_KDA = (
     importlib.util.find_spec("attn_gym") is not None
     and torch.cuda.is_available()
-    and torch.cuda.get_device_capability() in {(10, 0), (10, 3)}
+    and torch.cuda.get_device_capability() >= (9, 0)
 )
 
 
@@ -65,7 +65,7 @@ def conv() -> Conv1d.Config:
 
 
 @unittest.skipUnless(
-    _HAS_BLACKWELL, "KDA requires Attention Gym on CUDA capability 10.0 or 10.3"
+    _HAS_ATTENTION_GYM_KDA, "KDA requires Attention Gym on CUDA capability 9.0+"
 )
 class TestKDA(unittest.TestCase):
     def _make_kda(self):

tests/unit_tests/gpu/test_kimi_k3.py

@@ -132,9 +132,8 @@ def test_padded_tail_is_one_kda_segment(self):
         )
 
     @unittest.skipIf(
-        not torch.cuda.is_available()
-        or torch.cuda.get_device_capability() not in {(10, 0), (10, 3)},
-        "Attention Gym KDA requires CUDA capability 10.0 or 10.3.",
+        not torch.cuda.is_available() or torch.cuda.get_device_capability() < (9, 0),
+        "Attention Gym KDA requires CUDA capability 9.0 or newer.",
     )
     def test_attention_gym_kda_kernel_matches_recurrent_reference(self):
         torch.manual_seed(1)

torchtitan/models/kimi_k3/README.md

@@ -11,6 +11,10 @@ Install the additional dependencies:
 pip install -r .ci/docker/requirements-vlm.txt
 ```
 
+KDA uses Attention Gym kernels, which need an NVIDIA GPU with CUDA capability
+9.0 or newer. Attention Gym runs its CuTe kernels on SM100/SM103 and its Triton
+kernels on other GPUs.
+
 ## Architecture
 
 Kimi K3 is built on Kimi Delta Attention (KDA) and Attention Residuals

torchtitan/models/kimi_k3/kda.py

@@ -90,9 +90,11 @@ def forward(
         if not q_1THK.is_cuda:
             raise RuntimeError("Attention Gym KDA requires CUDA tensors.")
         capability = torch.cuda.get_device_capability(q_1THK.device)
-        if capability not in {(10, 0), (10, 3)}:
+        # The fused bounded gate needs TMA (SM90+). chunk_kda runs its CuTe
+        # kernels on SM100/SM103 and its Triton kernels on other NVIDIA GPUs.
+        if capability < (9, 0):
             raise RuntimeError(
-                "Attention Gym KDA requires Blackwell SM100/SM103; "
+                "Attention Gym KDA requires CUDA capability 9.0 or newer; "
                 f"got CUDA capability {capability}."
             )