diff --git a/.ci/docker/requirements.txt b/.ci/docker/requirements.txt index 8df75d8144..a2cec99127 100644 --- a/.ci/docker/requirements.txt +++ b/.ci/docker/requirements.txt @@ -8,3 +8,4 @@ safetensors einops pillow spmd_types==0.2.5 +attn-gym[linear]==0.0.5 diff --git a/pyproject.toml b/pyproject.toml index 70d7fbb446..8adfbd9fe2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,7 @@ dependencies = [ "einops", "pillow", "spmd_types==0.2.5", + "attn-gym[linear]==0.0.5", ] dynamic = ["version"] diff --git a/tests/integration_tests/__init__.py b/tests/integration_tests/__init__.py index cab9cf1c32..12d26bfba4 100644 --- a/tests/integration_tests/__init__.py +++ b/tests/integration_tests/__init__.py @@ -35,6 +35,8 @@ class OverrideDefinitions: ngpu: int = 4 disabled: bool = False skip_rocm_test: bool = False + required_cuda_capabilities: Sequence[tuple[int, int]] = () + """CUDA compute capabilities on which the test can run.""" timeout: int | None = None golden_numerics_path: str | None = None """Run through loss_compare.py using this mode-specific golden path.""" diff --git a/tests/integration_tests/models.py b/tests/integration_tests/models.py index e10ec4431e..cfc6f53b40 100755 --- a/tests/integration_tests/models.py +++ b/tests/integration_tests/models.py @@ -206,5 +206,6 @@ def build_model_tests_list() -> list[OverrideDefinitions]: test_descr="Kimi K3 multimodal FSDP", test_name="kimi_k3_mm_fsdp", ngpu=2, + required_cuda_capabilities=((10, 0), (10, 3)), ), ] diff --git a/tests/integration_tests/run_tests.py b/tests/integration_tests/run_tests.py index 8d1209116c..450d42e0bf 100644 --- a/tests/integration_tests/run_tests.py +++ b/tests/integration_tests/run_tests.py @@ -14,6 +14,8 @@ from concurrent.futures import Future, ThreadPoolExecutor from pathlib import Path +import torch + from torchtitan.tools.logging import logger from torchtitan.trainer import Trainer @@ -332,10 +334,14 @@ def run_single_test( def _filter_tests( args, test_list: list[OverrideDefinitions] -) -> tuple[list[OverrideDefinitions], list[OverrideDefinitions]]: +) -> tuple[ + list[OverrideDefinitions], + list[OverrideDefinitions], + list[OverrideDefinitions], +]: """Filter tests by name, scope, disabled state, architecture, and GPU count. - Returns (runnable, skipped_due_to_ngpu). + Returns (runnable, skipped_due_to_ngpu, skipped_due_to_cuda_capability). """ exclude_set = set() if hasattr(args, "exclude") and args.exclude: @@ -343,6 +349,13 @@ def _filter_tests( runnable: list[OverrideDefinitions] = [] skipped_ngpu: list[OverrideDefinitions] = [] + skipped_cuda_capability: list[OverrideDefinitions] = [] + cuda_capability = ( + torch.cuda.get_device_capability() + if getattr(args, "gpu_arch_type", "cuda") == "cuda" + and torch.cuda.is_available() + else None + ) for test_flavor in test_list: if args.test_name != "all" and test_flavor.test_name != args.test_name: continue @@ -361,11 +374,17 @@ def _filter_tests( and test_flavor.skip_rocm_test ): continue + if ( + test_flavor.required_cuda_capabilities + and cuda_capability not in test_flavor.required_cuda_capabilities + ): + skipped_cuda_capability.append(test_flavor) + continue if execution_mode != "fake_pg" and args.ngpu < test_flavor.ngpu: skipped_ngpu.append(test_flavor) continue runnable.append(test_flavor) - return runnable, skipped_ngpu + return runnable, skipped_ngpu, skipped_cuda_capability def run_tests( @@ -374,12 +393,18 @@ def run_tests( parallel: bool = True, ): """Run all integration tests to test the core features of TorchTitan.""" - runnable, skipped_ngpu = _filter_tests(args, test_list) + runnable, skipped_ngpu, skipped_cuda_capability = _filter_tests(args, test_list) for test_flavor in skipped_ngpu: logger.info( f"Skipping test {test_flavor.test_name} that requires {test_flavor.ngpu} gpus," f" because --ngpu arg is {args.ngpu}" ) + for test_flavor in skipped_cuda_capability: + logger.info( + f"Skipping test {test_flavor.test_name} because its required CUDA " + "capability is unavailable; supported capabilities are " + f"{tuple(test_flavor.required_cuda_capabilities)}" + ) failed_tests: list[tuple[str, str]] = [] execution_mode = getattr(args, "execution_mode", "real_pg") diff --git a/tests/unit_tests/cpu/test_integration_test_definitions.py b/tests/unit_tests/cpu/test_integration_test_definitions.py index acd122e46a..7223ce38c1 100644 --- a/tests/unit_tests/cpu/test_integration_test_definitions.py +++ b/tests/unit_tests/cpu/test_integration_test_definitions.py @@ -18,7 +18,11 @@ from tests.integration_tests.flux import build_flux_test_list from tests.integration_tests.h100 import build_h100_tests_list from tests.integration_tests.models import build_model_tests_list -from tests.integration_tests.run_tests import _parse_test_suites, run_single_test +from tests.integration_tests.run_tests import ( + _filter_tests, + _parse_test_suites, + run_single_test, +) def test_hf_checkpoint_load_path_comes_from_test_config(monkeypatch) -> None: @@ -72,6 +76,41 @@ def test_parse_multiple_integration_test_suites() -> None: ) +def test_filter_tests_skips_unsupported_cuda_capability(monkeypatch) -> None: + monkeypatch.setattr( + "tests.integration_tests.run_tests.torch.cuda.is_available", lambda: True + ) + monkeypatch.setattr( + "tests.integration_tests.run_tests.torch.cuda.get_device_capability", + lambda: (8, 6), + ) + supported = OverrideDefinitions(test_name="supported") + blackwell_only = OverrideDefinitions( + test_name="blackwell_only", + required_cuda_capabilities=((10, 0), (10, 3)), + ) + args = type( + "Args", + (), + { + "test_name": "all", + "execution_mode": "real_pg", + "test_scope": "all", + "gpu_arch_type": "cuda", + "ngpu": 8, + "exclude": None, + }, + )() + + runnable, skipped_ngpu, skipped_cuda_capability = _filter_tests( + args, [supported, blackwell_only] + ) + + assert runnable == [supported] + assert not skipped_ngpu + assert skipped_cuda_capability == [blackwell_only] + + def test_h100_tests_are_registered_in_separate_suite() -> None: assert {test.test_name for test in build_h100_tests_list()} == { "2d_asynctp_compile", diff --git a/tests/unit_tests/test_kda_attention.py b/tests/unit_tests/test_kda_attention.py new file mode 100644 index 0000000000..015871915b --- /dev/null +++ b/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() diff --git a/tests/unit_tests/test_kimi_k3.py b/tests/unit_tests/test_kimi_k3.py index 9d02f1e010..00d6fa6430 100644 --- a/tests/unit_tests/test_kimi_k3.py +++ b/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) diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index e559f1469c..6ea10930b6 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/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, diff --git a/torchtitan/models/kimi_k3/kda.py b/torchtitan/models/kimi_k3/kda.py index f31a8eb9b1..ac53adc104 100644 --- a/torchtitan/models/kimi_k3/kda.py +++ b/torchtitan/models/kimi_k3/kda.py @@ -4,23 +4,22 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""Kimi Delta Attention modules for Kimi K3.""" +"""Kimi Delta Attention using Attention Gym kernels.""" from dataclasses import dataclass import torch import torch.nn.functional as F -from fla.ops.kda import chunk_kda +from attn_gym.linear.kda import bound_gate, chunk_kda +from attn_gym.linear.kda.fwd.triton.l2norm_fwd import l2norm +from attn_gym.linear.kda.short_conv import causal_conv1d from torch import nn -from torchtitan.models.common import Conv1d, Linear -from torchtitan.models.common.attention import AttentionMasksType +from torchtitan.models.common.attention import AttentionMasksType, VarlenMetadata +from torchtitan.models.common.linear import Linear +from torchtitan.models.common.nn_modules import Conv1d from torchtitan.protocols.module import Module -# Shape suffixes: -# T = packed tokens, D = model dimension, H = heads, -# K = key head dimension, V = value head dimension, C = projection channels. - class KimiRMSNormGated(Module): """Per-head RMSNorm followed by a sigmoid output gate.""" @@ -35,59 +34,150 @@ def __init__(self, config: Config): self.eps = config.eps self.weight = nn.Parameter(torch.empty(config.dim)) - def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: - input_dtype = x.dtype - x_float = x.float() - variance = x_float.pow(2).mean(dim=-1, keepdim=True) - x_float = x_float * torch.rsqrt(variance + self.eps) - x_float = self.weight.float() * x_float - return (x_float * torch.sigmoid(gate.float())).to(input_dtype) + def forward(self, x_TNV: torch.Tensor, gate_TNV: torch.Tensor) -> torch.Tensor: + input_dtype = x_TNV.dtype + normalized_TNV = F.rms_norm( + x_TNV.float(), + (x_TNV.shape[-1],), + self.weight.float(), + self.eps, + ) + return (normalized_TNV * gate_TNV.float().sigmoid()).to(input_dtype) -class KimiKDAKernel(Module): - """Stateless dispatch to FLA's chunked KDA kernel.""" +class KDAKernel(Module): + """Apply KDA preprocessing and the Attention Gym kernel.""" @dataclass(kw_only=True, slots=True) class Config(Module.Config): - lower_bound: float | None = -5.0 + lower_bound: float = -5.0 + + def __post_init__(self): + if not -5.0 <= self.lower_bound < 0.0: + raise ValueError( + "KDA lower_bound must be in the safe range [-5, 0), " + f"got {self.lower_bound}." + ) def __init__(self, config: Config): super().__init__() self.lower_bound = config.lower_bound - if self.lower_bound is not None and not (-5.0 <= self.lower_bound < 0.0): - raise ValueError("KDA lower_bound must be in the safe range [-5, 0).") def forward( self, - q_BLHK: torch.Tensor, - k_BLHK: torch.Tensor, - v_BLHV: torch.Tensor, - gate_BLHK: torch.Tensor, - beta_BLH: torch.Tensor, - A_log_H: torch.Tensor, - dt_bias_HK: torch.Tensor, + q_BTNK: torch.Tensor, + k_BTNK: torch.Tensor, + v_BTNV: torch.Tensor, + raw_gate_BTNK: torch.Tensor, + raw_beta_BTN: torch.Tensor, + A_log_N: torch.Tensor, + dt_bias_NK: torch.Tensor, + *, + cu_seqlens: torch.Tensor | None = None, ) -> torch.Tensor: - out_BLHV, _ = chunk_kda( - q_BLHK, - k_BLHK, - v_BLHV, - gate_BLHK, - beta_BLH, - A_log=A_log_H, - dt_bias=dt_bias_HK.reshape(-1), - use_qk_l2norm_in_kernel=True, - use_gate_in_kernel=True, - use_beta_sigmoid_in_kernel=True, - safe_gate=self.lower_bound is not None, + if not q_BTNK.is_cuda: + raise RuntimeError("Attention Gym KDA requires CUDA tensors.") + capability = torch.cuda.get_device_capability(q_BTNK.device) + if capability not in {(10, 0), (10, 3)}: + raise RuntimeError( + "Attention Gym KDA requires Blackwell SM100/SM103; " + f"got CUDA capability {capability}." + ) + + gate_BTNK = bound_gate( + raw_gate_BTNK, + # TODO: The long-term solution is to specify mixed precision per FQN + # instead of per layer. https://github.com/pytorch/pytorch/issues/156784 + A_log_N.float(), + dt_bias_NK.float(), lower_bound=self.lower_bound, + impl="fused", + ) + output_BTNV, _ = chunk_kda( + l2norm(q_BTNK), + l2norm(k_BTNK), + v_BTNV, + gate_BTNK, + raw_beta_BTN.float().sigmoid(), + cu_seqlens=cu_seqlens, ) - return out_BLHV + return output_BTNV -class KimiDeltaAttention(Module): +class InnerKDA(Module): + """Run short convolution and KDA behind the vLLM replacement boundary.""" + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + head_dim: int + kernel: KDAKernel.Config + + def __post_init__(self): + if self.head_dim != 128: + raise ValueError( + "Attention Gym KDA requires head_dim=128, " f"got {self.head_dim}." + ) + + def __init__(self, config: Config): + super().__init__() + self.head_dim = config.head_dim + self.kernel = config.kernel.build() + + def forward( + self, + query_TC: torch.Tensor, + key_TC: torch.Tensor, + value_TC: torch.Tensor, + raw_gate_TNK: torch.Tensor, + raw_beta_TN: torch.Tensor, + conv_q_weight_C1W: torch.Tensor, + conv_k_weight_C1W: torch.Tensor, + conv_v_weight_C1W: torch.Tensor, + A_log_N: torch.Tensor, + dt_bias_NK: torch.Tensor, + cu_seqlens: torch.Tensor | None, + ) -> torch.Tensor: + raw_gate_BTNK = raw_gate_TNK.unsqueeze(0) + raw_beta_BTN = raw_beta_TN.unsqueeze(0) + mixed_qkv_BTC = torch.cat( + (query_TC, key_TC, value_TC), + dim=-1, + ).unsqueeze(0) + conv_weight_C1W = torch.cat( + (conv_q_weight_C1W, conv_k_weight_C1W, conv_v_weight_C1W), + dim=0, + ) + conv_output_BTC = causal_conv1d( + mixed_qkv_BTC, + conv_weight_C1W[:, 0], + activation="silu", + cu_seqlens=cu_seqlens, + ) + assert isinstance(conv_output_BTC, torch.Tensor) + + q_BTC, k_BTC, v_BTC = conv_output_BTC.chunk(3, dim=-1) + q_BTNK, k_BTNK, v_BTNV = ( + tensor.unflatten(-1, (-1, self.head_dim)) + for tensor in (q_BTC, k_BTC, v_BTC) + ) + output_BTNV = self.kernel( + q_BTNK, + k_BTNK, + v_BTNV, + raw_gate_BTNK, + raw_beta_BTN, + A_log_N, + dt_bias_NK, + cu_seqlens=cu_seqlens, + ) + return output_BTNV.squeeze(0) + + +class KDA(Module): + """Kimi Delta Attention with checkpoint-compatible Kimi K3 parameters.""" + @dataclass(kw_only=True, slots=True) class Config(Module.Config): - dim: int num_heads: int head_dim: int conv_kernel_size: int @@ -101,15 +191,26 @@ class Config(Module.Config): forget_b: Linear.Config beta: Linear.Config output_gate: Linear.Config - kernel: Module.Config + inner_kda: Module.Config output_norm: KimiRMSNormGated.Config output_proj: Linear.Config + def __post_init__(self): + if self.num_heads < 1: + raise ValueError(f"num_heads must be positive, got {self.num_heads}") + if self.head_dim != 128: + raise ValueError( + "Attention Gym KDA requires head_dim=128, " f"got {self.head_dim}." + ) + if self.conv_kernel_size < 1: + raise ValueError( + f"conv_kernel_size must be positive, got {self.conv_kernel_size}" + ) + def __init__(self, config: Config): super().__init__() self.num_heads = config.num_heads self.head_dim = config.head_dim - self.conv_kernel_size = config.conv_kernel_size self.q_proj = config.q_proj.build() self.k_proj = config.k_proj.build() @@ -121,17 +222,13 @@ def __init__(self, config: Config): self.forget_b = config.forget_b.build() self.beta = config.beta.build() self.output_gate = config.output_gate.build() - self.kernel = config.kernel.build() + self.inner_kda = config.inner_kda.build() self.output_norm = config.output_norm.build() self.output_proj = config.output_proj.build() self.A_log = nn.Parameter(torch.empty(config.num_heads)) self.dt_bias = nn.Parameter(torch.empty(config.num_heads, config.head_dim)) - def _causal_conv(self, x_TC: torch.Tensor, conv: Conv1d) -> torch.Tensor: - x_1CT = F.pad(x_TC.T.unsqueeze(0), (self.conv_kernel_size - 1, 0)) - return F.silu(conv(x_1CT)).squeeze(0).T - def forward( self, x_TD: torch.Tensor, @@ -139,37 +236,38 @@ def forward( positions: torch.Tensor | None = None, ) -> torch.Tensor: del positions - if attention_masks is not None: - raise NotImplementedError( - "Kimi K3 reference KDA does not support packed-document masks." + if x_TD.ndim != 2: + raise ValueError( + f"KDA input must have shape [T, D], got {tuple(x_TD.shape)}." ) + if attention_masks is None: + cu_seqlens = None + elif isinstance(attention_masks, VarlenMetadata): + cu_seqlens = attention_masks.cu_seq_q + else: + raise ValueError( + "KDA attention_masks must be VarlenMetadata or None, " + f"got {type(attention_masks).__name__}." + ) num_tokens = x_TD.shape[0] - q_THK = self._causal_conv(self.q_proj(x_TD), self.q_conv).view( - num_tokens, self.num_heads, self.head_dim - ) - k_THK = self._causal_conv(self.k_proj(x_TD), self.k_conv).view( - num_tokens, self.num_heads, self.head_dim - ) - v_THV = self._causal_conv(self.v_proj(x_TD), self.v_conv).view( - num_tokens, self.num_heads, self.head_dim - ) - forget_THK = self.forget_b(self.forget_a(x_TD)).view( + raw_gate_TNK = self.forget_b(self.forget_a(x_TD)).reshape( num_tokens, self.num_heads, self.head_dim ) - beta_TH = self.beta(x_TD).float() - - out_THV = self.kernel( - q_THK.unsqueeze(0), - k_THK.unsqueeze(0), - v_THV.unsqueeze(0), - forget_THK.unsqueeze(0), - beta_TH.unsqueeze(0), + raw_beta_TN = self.beta(x_TD).reshape(num_tokens, self.num_heads) + out_TNV = self.inner_kda( + self.q_proj(x_TD), + self.k_proj(x_TD), + self.v_proj(x_TD), + raw_gate_TNK, + raw_beta_TN, + self.q_conv.weight, + self.k_conv.weight, + self.v_conv.weight, self.A_log, self.dt_bias, - ).squeeze(0) - output_gate_THV = self.output_gate(x_TD).view( - num_tokens, self.num_heads, self.head_dim + cu_seqlens, ) - out_THV = self.output_norm(out_THV, output_gate_THV) - return self.output_proj(out_THV.reshape(num_tokens, -1)) + + output_gate_TNV = self.output_gate(x_TD).view_as(out_TNV) + return self.output_proj(self.output_norm(out_TNV, output_gate_TNV).flatten(-2)) diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index 8314b7f999..a052387ed1 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -30,7 +30,7 @@ ) from torchtitan.protocols.module import Module -from .kda import KimiDeltaAttention +from .kda import KDA from .moe import KimiFeedForward, KimiLatentMoE from .vision_encoder import KimiK3VisionEncoder @@ -163,7 +163,7 @@ class Config(Module.Config): layer_id: int attn_res_block_size: int attention: KimiMLAAttention.Config | None - delta_attention: KimiDeltaAttention.Config | None + delta_attention: KDA.Config | None feed_forward: KimiFeedForward.Config | None moe: KimiLatentMoE.Config | None attention_norm: RMSNorm.Config @@ -299,7 +299,7 @@ def get_nparams_and_flops( v_head_dim=attention.v_head_dim, seq_len=seq_len, ) - elif isinstance(layer.delta_attention, KimiDeltaAttention.Config): + elif isinstance(layer.delta_attention, KDA.Config): delta_attention = layer.delta_attention attention_op_flops += delta_rule_flops_per_token( num_heads=delta_attention.num_heads,