kda in torchtitan (training)

.ci/docker/requirements-vlm.txt

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

.ci/docker/requirements.txt

@@ -8,3 +8,5 @@ safetensors
 einops
 pillow
 spmd_types==0.2.3
+flash-linear-attention
+attn-gym[linear]==0.0.5

pyproject.toml

@@ -76,5 +76,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.*", "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,133 @@
+# 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 torch
+
+from torchtitan.models.common import Conv1d, Linear
+from torchtitan.models.common.attention import create_varlen_metadata_for_document
+from torchtitan.models.kimi_k3.kda import InnerKDA, KDA, KDAKernel, KimiRMSNormGated
+
+_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=KimiRMSNormGated.Config(dim=128),
+        output_proj=linear(projection_dim, 32),
+    )
+
+
+@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_documents(self):
+        lengths = (37, 64, 91)
+        x_TD = self._inputs(seed=2, tokens=sum(lengths)).requires_grad_()
+        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,
+        )
+        output_grad_TD = torch.randn_like(packed_TD)
+        parameters = tuple(model.parameters())
+        packed_grads = torch.autograd.grad(
+            packed_TD,
+            (x_TD, *parameters),
+            output_grad_TD,
+        )
+        independent_grads = torch.autograd.grad(
+            independent_TD,
+            (x_TD, *parameters),
+            output_grad_TD,
+        )
+        torch.testing.assert_close(
+            packed_grads,
+            independent_grads,
+            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,18 @@ 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).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 +113,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 +131,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)

torchtitan/models/kimi_k3/__init__.py

@@ -26,7 +26,7 @@
 from torchtitan.protocols.model import ModelConfigConverter
 from torchtitan.protocols.model_spec import ModelSpec
 
-from .kda import KimiDeltaAttention, KimiKDAKernel, KimiRMSNormGated
+from .kda import InnerKDA, KDA, KDAKernel, KimiRMSNormGated
 from .model import KimiK3Model, KimiK3TransformerBlock, KimiMLAAttention
 from .moe import KimiFeedForward, KimiGroupedExperts, KimiLatentMoE
 from .parallelize import parallelize_kimi_k3
@@ -170,7 +170,7 @@ def _kda_config(
     num_heads: int,
     head_dim: int,
     conv_kernel_size: int,
-) -> KimiDeltaAttention.Config:
+) -> KDA.Config:
     projection_dim = num_heads * head_dim
 
     def conv() -> Conv1d.Config:
@@ -183,8 +183,7 @@ def conv() -> Conv1d.Config:
             param_init=_CONV_INIT,
         )
 
-    return KimiDeltaAttention.Config(
-        dim=dim,
+    return KDA.Config(
         num_heads=num_heads,
         head_dim=head_dim,
         conv_kernel_size=conv_kernel_size,
@@ -198,7 +197,10 @@ def conv() -> Conv1d.Config:
         forget_b=_linear(head_dim, projection_dim),
         beta=_linear(dim, num_heads),
         output_gate=_linear(dim, projection_dim),
-        kernel=KimiKDAKernel.Config(lower_bound=-5.0),
+        inner_kda=InnerKDA.Config(
+            head_dim=head_dim,
+            kernel=KDAKernel.Config(),
+        ),
         output_norm=KimiRMSNormGated.Config(
             dim=head_dim,
             eps=1e-5,
@@ -466,7 +468,7 @@ def _debugmodel(attn_backend: str) -> KimiK3Model.Config:
         qk_nope_head_dim=64,
         qk_rope_head_dim=32,
         v_head_dim=64,
-        kda_head_dim=64,
+        kda_head_dim=128,
         conv_kernel_size=4,
         dense_hidden_dim=4096,
         latent_dim=512,