kda in torchtitan (training)

.ci/docker/requirements-vlm.txt

@@ -3,3 +3,4 @@ einops
 pillow
 torchvision
 flash-linear-attention
+attention-gym[linear]

pyproject.toml

@@ -73,5 +73,5 @@ testpaths = ["tests"]
 [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.*", "fla.*", "attn_gym.*", "helion", "helion.*", "batch_invariant_ops", "torchcomms"]  # optional dependencies
 search-path = ["../pytorch"]  # local built pytorch

tests/unit_tests/test_kda_attention.py

@@ -0,0 +1,173 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Unit tests for the KDA linear-attention layer."""
+
+import importlib.util
+import unittest
+
+import spmd_types as spmd
+import torch
+
+from torchtitan.distributed.parallel_dims import MeshAxisName
+from torchtitan.models.common import Conv1d, Linear, RMSNorm
+from torchtitan.models.common.attention import create_varlen_metadata_for_document
+from torchtitan.models.common.decoder_sharding import (
+    dense_sequence_parallel_placement,
+    set_kda_sharding,
+)
+from torchtitan.models.kimi_k3.kda import InnerKDA, KDA, KDAKernel
+
+_HAS_BLACKWELL = (
+    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)}
+)
+
+
+def _kda_config() -> KDA.Config:
+    def linear(in_features: int, out_features: int) -> Linear.Config:
+        return Linear.Config(
+            in_features=in_features,
+            out_features=out_features,
+            bias=False,
+        )
+
+    projection_dim = 256
+
+    def conv() -> Conv1d.Config:
+        return Conv1d.Config(
+            in_channels=projection_dim,
+            out_channels=projection_dim,
+            kernel_size=4,
+            groups=projection_dim,
+            bias=False,
+        )
+
+    return KDA.Config(
+        num_heads=2,
+        head_dim=128,
+        conv_kernel_size=4,
+        q_proj=linear(32, projection_dim),
+        k_proj=linear(32, projection_dim),
+        v_proj=linear(32, projection_dim),
+        q_conv=conv(),
+        k_conv=conv(),
+        v_conv=conv(),
+        forget_a=linear(32, 128),
+        forget_b=linear(128, projection_dim),
+        beta=linear(32, 2),
+        output_gate=linear(32, projection_dim),
+        inner_kda=InnerKDA.Config(
+            head_dim=128,
+            kernel=KDAKernel.Config(),
+        ),
+        output_norm=RMSNorm.Config(normalized_shape=128),
+        output_proj=linear(projection_dim, 32),
+    )
+
+
+class TestKDASharding(unittest.TestCase):
+    def test_folded_token_tp_contracts(self):
+        config = _kda_config()
+        input_layout = dense_sequence_parallel_placement()
+        set_kda_sharding(
+            config,
+            attention_input_layout=input_layout,
+            enable_sp=True,
+            cp_enabled=False,
+        )
+
+        tp_axis = MeshAxisName.TP
+        q_proj_sharding = config.q_proj.sharding_config
+        output_proj_sharding = config.output_proj.sharding_config
+        inner_kda_sharding = config.inner_kda.sharding_config
+        kda_sharding = config.sharding_config
+        self.assertEqual(
+            q_proj_sharding.state_shardings["weight"].axis_types[tp_axis],
+            spmd.S(0),
+        )
+        self.assertEqual(
+            output_proj_sharding.state_shardings["weight"].axis_types[tp_axis],
+            spmd.S(1),
+        )
+        self.assertEqual(
+            set(inner_kda_sharding.in_src_shardings),
+            {
+                "query_TC",
+                "key_TC",
+                "value_TC",
+                "raw_gate_TNK",
+                "raw_beta_TN",
+                "conv_q_weight_C1W",
+                "conv_k_weight_C1W",
+                "conv_v_weight_C1W",
+                "A_log_N",
+                "dt_bias_NK",
+                "cu_seqlens",
+            },
+        )
+        head_layout = inner_kda_sharding.in_src_shardings["raw_gate_TNK"]
+        self.assertEqual(head_layout.per_axis_spmd_types()[tp_axis], spmd.S(1))
+        self.assertEqual(
+            kda_sharding.in_src_shardings,
+            {"x_TD": input_layout},
+        )
+        self.assertEqual(
+            len(inner_kda_sharding.local_map.in_grad_placements),
+            11,
+        )
+
+
+@unittest.skipUnless(
+    _HAS_BLACKWELL, "KDA requires Attention Gym on CUDA capability 10.0 or 10.3"
+)
+class TestKDA(unittest.TestCase):
+    def _make_kda(self):
+        model = _kda_config().build()
+        model = model.to(device="cuda", dtype=torch.bfloat16)
+        torch.manual_seed(1)
+        with torch.no_grad():
+            for param in model.parameters():
+                param.normal_(mean=0.0, std=0.02)
+            model.A_log.uniform_(1.0, 16.0).log_()
+            model.dt_bias.zero_()
+            model.output_norm.weight.fill_(1.0)
+        return model
+
+    def _inputs(self, seed: int, tokens: int = 128) -> torch.Tensor:
+        torch.manual_seed(seed)
+        return torch.randn(tokens, 32, device="cuda", dtype=torch.bfloat16)
+
+    def test_varlen_matches_independent_document_forwards(self):
+        lengths = (37, 64, 91)
+        x_TD = self._inputs(seed=2, tokens=sum(lengths))
+        positions_T = torch.tensor(
+            [index for length in lengths for index in range(length)],
+            device="cuda",
+            dtype=torch.int32,
+        )
+        masks = create_varlen_metadata_for_document(
+            positions_T,
+            include_host_offsets=True,
+        )
+        self.assertEqual(masks.cu_seq_q_host, (0, 37, 101, 192))
+
+        model = self._make_kda()
+        packed_TD = model(x_TD, masks)
+        independent_TD = torch.cat(
+            [model(document_TD, None) for document_TD in x_TD.split(lengths)]
+        )
+        torch.testing.assert_close(
+            packed_TD.float(),
+            independent_TD.float(),
+            rtol=2e-2,
+            atol=2e-2,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit_tests/test_kimi_k3.py

@@ -7,11 +7,10 @@
 import unittest
 
 import torch
-import torch.nn.functional as F
 from torch.nn.attention.flex_attention import BlockMask
 
 from torchtitan.models.kimi_k3 import _kimi_k3_config, _vision_encoder_config
-from torchtitan.models.kimi_k3.kda import KimiKDAKernel
+from torchtitan.models.kimi_k3.kda import KDAKernel
 from torchtitan.models.kimi_k3.model import KimiK3Model
 from torchtitan.models.kimi_k3.state_dict_adapter import KimiK3StateDictAdapter
 
@@ -31,7 +30,7 @@ def _small_model_config() -> KimiK3Model.Config:
         qk_nope_head_dim=16,
         qk_rope_head_dim=16,
         v_head_dim=16,
-        kda_head_dim=64,
+        kda_head_dim=128,
         conv_kernel_size=3,
         dense_hidden_dim=128,
         latent_dim=32,
@@ -65,29 +64,19 @@ def _kda_recurrent_reference(
     A_log_H: torch.Tensor,
     dt_bias_HK: torch.Tensor,
     *,
-    lower_bound: float | None,
+    lower_bound: float,
 ) -> torch.Tensor:
-    """Explicit KDA recurrence in FP32, matching the released Kimi K3 math.
-
-    ``lower_bound`` selects the same two gate activations FLA exposes through
-    ``safe_gate``: the bounded ``lower_bound * sigmoid(...)`` form when set,
-    and ``-exp(A_log) * softplus(...)`` when ``None``.
-    """
+    """Explicit bounded KDA recurrence in FP32."""
     input_dtype = q_BLHK.dtype
     q_BLHK = q_BLHK.float()
     k_BLHK = k_BLHK.float()
     q_BLHK = q_BLHK * torch.rsqrt(q_BLHK.square().sum(dim=-1, keepdim=True) + 1e-6)
     k_BLHK = k_BLHK * torch.rsqrt(k_BLHK.square().sum(dim=-1, keepdim=True) + 1e-6)
     v_BLHV = v_BLHV.float()
-    if lower_bound is None:
-        log_decay_BLHK = -torch.exp(A_log_H.float()).view(1, 1, -1, 1) * F.softplus(
-            gate_BLHK.float() + dt_bias_HK.float()
-        )
-    else:
-        log_decay_BLHK = lower_bound * torch.sigmoid(
-            torch.exp(A_log_H.float()).view(1, 1, -1, 1)
-            * (gate_BLHK.float() + dt_bias_HK.float())
-        )
+    log_decay_BLHK = lower_bound * torch.sigmoid(
+        torch.exp(A_log_H.float()).view(1, 1, -1, 1)
+        * (gate_BLHK.float() + dt_bias_HK.float())
+    )
     decay_BLHK = torch.exp(log_decay_BLHK)
     beta_BLH = torch.sigmoid(beta_BLH.float())
 
@@ -125,10 +114,14 @@ def test_flex_attention_mask(self):
         attention_masks = model.get_attention_masks(positions)
         self.assertIsInstance(attention_masks, BlockMask)
 
-    @unittest.skipIf(not torch.cuda.is_available(), "FLA KDA kernel requires CUDA.")
-    def test_fla_kda_kernel_matches_recurrent_reference(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.",
+    )
+    def test_attention_gym_kda_kernel_matches_recurrent_reference(self):
         torch.manual_seed(1)
-        head_dim = 64
+        head_dim = 128
         num_heads = 3
 
         def parameter(*shape: int) -> torch.Tensor:
@@ -139,61 +132,64 @@ def parameter(*shape: int) -> torch.Tensor:
                 requires_grad=True,
             )
 
-        for lower_bound in (-5.0, None):
-            with self.subTest(lower_bound=lower_bound):
-                A_log_H = torch.rand(num_heads, device="cuda")
-                A_log_H = A_log_H.uniform_(1.0, 16.0).log().requires_grad_()
-                actual_inputs = (
-                    parameter(2, 64, num_heads, head_dim),
-                    parameter(2, 64, num_heads, head_dim),
-                    parameter(2, 64, num_heads, head_dim),
-                    parameter(2, 64, num_heads, head_dim),
-                    parameter(2, 64, num_heads),
-                    A_log_H,
-                    parameter(num_heads, head_dim),
-                )
-                expected_inputs = tuple(
-                    tensor.detach().clone().requires_grad_() for tensor in actual_inputs
-                )
-
-                kernel = KimiKDAKernel.Config(lower_bound=lower_bound).build()
-                actual_BLHV = kernel(*actual_inputs)
-                expected_BLHV = _kda_recurrent_reference(
-                    *expected_inputs,
-                    lower_bound=lower_bound,
-                )
-
-                # The chunked kernel accumulates over chunk boundaries and uses
-                # reduced-precision matmuls internally, so it does not reproduce
-                # the sequential FP32 recurrence bit for bit.
-                torch.testing.assert_close(
-                    actual_BLHV,
-                    expected_BLHV,
-                    atol=2e-3,
-                    rtol=2e-3,
-                )
-                output_grad_BLHV = torch.randn_like(actual_BLHV)
-                actual_grads = torch.autograd.grad(
-                    actual_BLHV,
-                    actual_inputs,
-                    grad_outputs=output_grad_BLHV,
-                )
-                expected_grads = torch.autograd.grad(
-                    expected_BLHV,
-                    expected_inputs,
-                    grad_outputs=output_grad_BLHV,
-                )
-                for actual_grad, expected_grad in zip(
-                    actual_grads,
-                    expected_grads,
-                    strict=True,
-                ):
-                    torch.testing.assert_close(
-                        actual_grad,
-                        expected_grad,
-                        atol=2e-2,
-                        rtol=2e-2,
-                    )
+        lower_bound = -5.0
+        A_log_H = (
+            torch.empty(num_heads, device="cuda")
+            .uniform_(1.0, 16.0)
+            .log_()
+            .requires_grad_()
+        )
+        actual_inputs = (
+            parameter(2, 64, num_heads, head_dim),
+            parameter(2, 64, num_heads, head_dim),
+            parameter(2, 64, num_heads, head_dim),
+            parameter(2, 64, num_heads, head_dim),
+            parameter(2, 64, num_heads),
+            A_log_H,
+            parameter(num_heads, head_dim),
+        )
+        expected_inputs = tuple(
+            tensor.detach().clone().requires_grad_() for tensor in actual_inputs
+        )
+
+        kernel = KDAKernel.Config(lower_bound=lower_bound).build()
+        actual_BLHV = kernel(*actual_inputs)
+        expected_BLHV = _kda_recurrent_reference(
+            *expected_inputs,
+            lower_bound=lower_bound,
+        )
+
+        # The chunked kernel accumulates over chunk boundaries and uses
+        # reduced-precision matmuls internally, so it does not reproduce
+        # the sequential FP32 recurrence bit for bit.
+        torch.testing.assert_close(
+            actual_BLHV,
+            expected_BLHV,
+            atol=2e-3,
+            rtol=2e-3,
+        )
+        output_grad_BLHV = torch.randn_like(actual_BLHV)
+        actual_grads = torch.autograd.grad(
+            actual_BLHV,
+            actual_inputs,
+            grad_outputs=output_grad_BLHV,
+        )
+        expected_grads = torch.autograd.grad(
+            expected_BLHV,
+            expected_inputs,
+            grad_outputs=output_grad_BLHV,
+        )
+        for actual_grad, expected_grad in zip(
+            actual_grads,
+            expected_grads,
+            strict=True,
+        ):
+            torch.testing.assert_close(
+                actual_grad,
+                expected_grad,
+                atol=2e-2,
+                rtol=2e-2,
+            )
 
     def test_state_dict_round_trips_through_hf_adapter(self):
         torch.manual_seed(2)