diff --git a/benchmarks/ops/benchmark_kda.py b/benchmarks/ops/benchmark_kda.py new file mode 100644 index 0000000000..1c75fc39cc --- /dev/null +++ b/benchmarks/ops/benchmark_kda.py @@ -0,0 +1,147 @@ +# -*- coding: utf-8 -*- + +import os + +import torch +import triton +from flash_attn import flash_attn_func +from torch.nn import functional as F + +from fla.ops.comba import chunk_comba +from fla.ops.gated_delta_rule import chunk_gated_delta_rule +from fla.ops.generalized_delta_rule import chunk_dplr_delta_rule +from fla.ops.kda import chunk_kda + + +@triton.testing.perf_report( + triton.testing.Benchmark( + # argument names to use as an x-axis for the plot + x_names=['T'], + # different possible values for `x_name` + x_vals=[256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536], + # argument name whose value corresponds to a different line in the plot + line_arg='provider', + # possible values for `line_arg`` + line_vals=['gdn', 'comba', 'kda', 'dplr', 'attn'], + # label name for the lines + line_names=['gdn', 'comba', 'kda', 'dplr', 'attn'], + # line styles + styles=[('blue', '-'), ('red', '-.'), ('green', '-'), ('orange', '-.'), + ('purple', '-'), ('brown', '-.'), ('pink', '-'), ('gray', '-.')], + ylabel="Execution Time (ms)", # label name for the y-axis + # name for the plot. Used also as a file name for saving the plot. + plot_name="Performance", + args={}, + ) +) +def benchmark(T, provider): + from fla.utils import device + dtype = torch.bfloat16 + B, H, D = 1, 16, 128 + + # Set TMA environment variable based on provider + original_tma_env = os.environ.get('FLA_USE_TMA', '0') + + if provider.endswith('_no_tma'): + os.environ['FLA_USE_TMA'] = '0' + provider_base = provider.replace('_no_tma', '') + else: + os.environ['FLA_USE_TMA'] = '1' + provider_base = provider + + quantiles = [0.5, 0.2, 0.8] + results = 0, 0, 0 + + do = torch.randn(B, T, H, D, dtype=dtype, device=device) + if provider_base == 'gdn': + q = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) + k = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) + v = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) + g = F.logsigmoid(torch.randn(B, T, H, dtype=dtype, device=device)).requires_grad_(True) + beta = torch.randn(B, T, H, dtype=dtype, device=device).sigmoid().requires_grad_(True) + results = triton.testing.do_bench( + lambda: chunk_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + use_qk_l2norm_in_kernel=True, + )[0].backward(do), + quantiles=quantiles + ) + elif provider_base == 'attn': + q = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) + k = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) + v = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) + results = triton.testing.do_bench( + lambda: flash_attn_func( + q=q, + k=k, + v=v, + ).backward(do), + quantiles=quantiles + ) + elif provider_base == 'comba': + q = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) + k = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) + p = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) + v = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) + g = F.logsigmoid(torch.randn(B, T, H, dtype=torch.float, device=device)).requires_grad_(True) + beta = torch.randn(B, T, H, dtype=dtype, device=device).sigmoid().requires_grad_(True) + results = triton.testing.do_bench( + lambda: chunk_comba( + q=q, + k=k, + p=p, + v=v, + g=g, + beta=beta, + use_qk_l2norm_in_kernel=True, + )[0].backward(do), + quantiles=quantiles + ) + elif provider_base == 'kda': + q = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) + k = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) + v = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) + g = F.logsigmoid(torch.randn(B, T, H, D, dtype=dtype, device=device)).requires_grad_(True) + beta = torch.randn(B, T, H, dtype=dtype, device=device).sigmoid().requires_grad_(True) + results = triton.testing.do_bench( + lambda: chunk_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + use_qk_l2norm_in_kernel=True, + )[0].backward(do), + quantiles=quantiles + ) + elif provider_base == 'dplr': + q = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) + k = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) + a = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) + b = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) + v = torch.randn(B, T, H, D, dtype=dtype, device=device).requires_grad_(True) + g = F.logsigmoid(torch.randn(B, T, H, D, dtype=dtype, device=device)).requires_grad_(True) + beta = torch.randn(B, T, H, dtype=dtype, device=device).sigmoid().requires_grad_(True) + results = triton.testing.do_bench( + lambda: chunk_dplr_delta_rule( + q=q, + k=k, + v=v, + a=a, + b=b, + gk=g, + )[0].backward(do), + quantiles=quantiles + ) + + # Restore original TMA environment variable + os.environ['FLA_USE_TMA'] = original_tma_env + return results + + +if __name__ == '__main__': + benchmark.run(print_data=True) diff --git a/fla/layers/__init__.py b/fla/layers/__init__.py index 3527a387ef..8044097256 100644 --- a/fla/layers/__init__.py +++ b/fla/layers/__init__.py @@ -15,6 +15,7 @@ from .gsa import GatedSlotAttention from .hgrn import HGRNAttention from .hgrn2 import HGRN2Attention +from .kda import KimiDeltaAttention from .lightnet import LightNetAttention from .linear_attn import LinearAttention from .log_linear_mamba2 import LogLinearMamba2 @@ -45,6 +46,7 @@ 'GatedSlotAttention', 'HGRNAttention', 'HGRN2Attention', + 'KimiDeltaAttention', 'LightNetAttention', 'LinearAttention', 'LogLinearMamba2', diff --git a/fla/layers/gated_deltanet.py b/fla/layers/gated_deltanet.py index 5d2ce34611..000731f8a4 100644 --- a/fla/layers/gated_deltanet.py +++ b/fla/layers/gated_deltanet.py @@ -281,7 +281,7 @@ def forward( initial_state=recurrent_state, output_final_state=use_cache, cu_seqlens=cu_seqlens, - use_qk_l2norm_in_kernel=True + use_qk_l2norm_in_kernel=True, ) elif mode == 'fused_recurrent': o, recurrent_state = fused_recurrent_gated_delta_rule( @@ -293,7 +293,7 @@ def forward( initial_state=recurrent_state, output_final_state=use_cache, cu_seqlens=cu_seqlens, - use_qk_l2norm_in_kernel=True + use_qk_l2norm_in_kernel=True, ) else: raise NotImplementedError(f"Not supported mode `{mode}`.") diff --git a/fla/layers/kda.py b/fla/layers/kda.py new file mode 100644 index 0000000000..54c32515ca --- /dev/null +++ b/fla/layers/kda.py @@ -0,0 +1,280 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING, Dict, Optional, Tuple + +import torch +import torch.nn as nn +from einops import rearrange, repeat +from torch.nn import functional as F + +from fla.layers.utils import get_unpad_data, index_first_axis, pad_input +from fla.modules import FusedRMSNormGated, ShortConvolution +from fla.ops.kda import chunk_kda, fused_recurrent_kda +from fla.ops.kda.gate import fused_kda_gate + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + from fla.models.utils import Cache + + +class KimiDeltaAttention(nn.Module): + """ + Kimi Delta Attention (KDA) layer implementation. + + Each layer contains approximately 6*hidden_size*hidden_size parameters. + + Parameter allocation: + - q_proj: hidden_size * key_dim (where key_dim = num_heads * head_dim) + - k_proj: hidden_size * key_dim + - v_proj: hidden_size * value_dim (where value_dim = num_v_heads * head_dim * expand_v) + - o_proj: value_dim * hidden_size + - f_proj: hidden_size * head_v_dim + head_v_dim * key_dim (with bias) + - b_proj: hidden_size * num_heads + - g_proj: hidden_size * head_v_dim + head_v_dim * value_dim (with bias) + - A: num_heads parameters + - Plus convolution layers when use_short_conv=True + + Args: + hidden_size (int, Optional): + The hidden size of the input. Default: 2048. + expand_v (float, Optional): + The expansion ratio for the value dimension. Default: 1.0. + head_dim (int, Optional): + The dimension of each head. Default: 128. + num_heads (int, Optional): + The number of heads. Default: 16. + num_v_heads (int, Optional): + The number of heads for the value projection, equal to `num_heads` if `None`. + GVA (Grouped Value Attention) is applied if `num_v_heads` > `num_heads`. Default: `None`. + mode (str, Optional): + Which Gated DeltaNet kernel to use. + Currently available: `chunk` and `fused_recurrent`. + Default: `chunk`. + use_short_conv (bool, Optional): + Whether to use short convolutions. Default: `True`. + allow_neg_eigval (bool, Optional): + Allow negative eigenvalues. Default: `False`. If set to `True`, the beta will be multiplied by 2. + See reference: + [Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues](https://arxiv.org/abs/2411.12537) + conv_size (int, Optional): + The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4. + conv_bias (bool, Optional): + Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`. + layer_idx (int, Optional): + The index of the layer. Default: None. + norm_eps (float, Optional): + The epsilon value for the normalization layer. Default: 1e-5. + """ + + def __init__( + self, + hidden_size: int = 2048, + expand_v: float = 1, + head_dim: int = 128, + num_heads: int = 16, + num_v_heads: int = None, + mode: str = 'chunk', + use_short_conv: bool = True, + allow_neg_eigval: bool = False, + conv_size: int = 4, + conv_bias: bool = False, + layer_idx: int = None, + norm_eps: float = 1e-5, + **kwargs + ) -> KimiDeltaAttention: + super().__init__() + + self.mode = mode + self.allow_neg_eigval = allow_neg_eigval + self.hidden_size = hidden_size + self.expand_v = expand_v + + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + + self.head_dim = head_dim + self.num_heads = num_heads + self.num_v_heads = num_v_heads if num_v_heads is not None else num_heads + + self.head_k_dim = head_dim + self.head_v_dim = int(self.head_dim * self.expand_v) + self.key_dim = int(self.num_heads * self.head_k_dim) + self.value_dim = int(self.num_v_heads * self.head_v_dim) + self.layer_idx = layer_idx + + # Consistency check: Ensure expand_v produces integer values + if not math.isclose(self.num_v_heads * self.head_dim * expand_v, self.value_dim, rel_tol=1e-5): + raise ValueError( + f"expand_v={expand_v} does not produce an integer value when multiplied by key_dim={self.key_dim}. " + f"Resulting value_dim would be {self.num_v_heads * self.head_dim * expand_v}, which is invalid for nn.Linear." + ) + if self.num_v_heads > self.num_heads and self.num_v_heads % self.num_heads != 0: + raise ValueError( + f"num_v_heads={self.num_v_heads} must be divisible by num_heads={self.num_heads}." + ) + + if not math.isclose(head_dim * expand_v, self.head_v_dim, rel_tol=1e-5): + raise ValueError( + f"expand_v={expand_v} does not produce an integer value when multiplied by head_dim={head_dim}. " + f"Resulting head_v_dim would be {head_dim * expand_v}, which is invalid for FusedRMSNormGated." + ) + assert mode in ['chunk', 'fused_recurrent'], f"Not supported mode `{mode}`." + + self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + + if use_short_conv: + self.q_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation='silu' + ) + self.k_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation='silu' + ) + self.v_conv1d = ShortConvolution( + hidden_size=self.value_dim, + kernel_size=conv_size, + bias=conv_bias, + activation='silu' + ) + + self.A = nn.Parameter(torch.log(torch.empty(self.num_heads, dtype=torch.float32).uniform_(1, 16))) + self.f_proj = nn.Sequential( + nn.Linear(hidden_size, self.head_v_dim, bias=False), + nn.Linear(self.head_v_dim, self.key_dim, bias=True) + ) + self.b_proj = nn.Linear(hidden_size, self.num_heads, bias=False) + + self.g_proj = nn.Sequential( + nn.Linear(hidden_size, self.head_v_dim, bias=False), + nn.Linear(self.head_v_dim, self.value_dim, bias=True) + ) + self.o_norm = FusedRMSNormGated(self.head_v_dim, activation='sigmoid', eps=norm_eps) + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[Cache] = None, + use_cache: Optional[bool] = False, + output_attentions: Optional[bool] = False, + **kwargs: Unpack[Dict] + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + batch_size, q_len, _ = hidden_states.shape + # change to inference mode. + mode = 'fused_recurrent' if q_len <= 64 else self.mode + if self.training: + assert mode == 'chunk', "Only chunk mode is supported in training." + + last_state = None + if past_key_values is not None and len(past_key_values) > self.layer_idx: + last_state = past_key_values[self.layer_idx] + + cu_seqlens = kwargs.get('cu_seqlens', None) + if attention_mask is not None: + indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:]) + hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0) + + if self.use_short_conv: + conv_state_q, conv_state_k, conv_state_v = None, None, None + if last_state is not None: + conv_state_q, conv_state_k, conv_state_v = last_state['conv_state'] + q, conv_state_q = self.q_conv1d( + x=self.q_proj(hidden_states), + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + k, conv_state_k = self.k_conv1d( + x=self.k_proj(hidden_states), + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + v, conv_state_v = self.v_conv1d( + x=self.v_proj(hidden_states), + cache=conv_state_v, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + else: + q = F.silu(self.q_proj(hidden_states)) + k = F.silu(self.k_proj(hidden_states)) + v = F.silu(self.v_proj(hidden_states)) + + g = self.f_proj(hidden_states) + g = fused_kda_gate(g, self.A, self.head_k_dim) + beta = self.b_proj(hidden_states).sigmoid() + + q, k = map(lambda x: rearrange(x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k)) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim) + + if self.num_v_heads > self.num_heads: + q, k = map(lambda x: repeat(x, '... h d -> ... (h g) d', g=self.num_v_heads // self.num_heads), (q, k)) + + if self.allow_neg_eigval: + beta = beta * 2. + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + if mode == 'chunk': + o, recurrent_state = chunk_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=recurrent_state, + output_final_state=use_cache, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens, + ) + elif mode == 'fused_recurrent': + o, recurrent_state = fused_recurrent_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=recurrent_state, + output_final_state=use_cache, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens, + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + if past_key_values is not None: + past_key_values.update( + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None, + layer_idx=self.layer_idx, + offset=q_len + ) + + o = self.o_norm(o, rearrange(self.g_proj(hidden_states), '... (h d) -> ... h d', d=self.head_v_dim)) + o = rearrange(o, 'b t h d -> b t (h d)') + o = self.o_proj(o) + if attention_mask is not None: + o = pad_input(o.squeeze(0), indices, batch_size, q_len) + + return o, None, past_key_values diff --git a/fla/models/__init__.py b/fla/models/__init__.py index 902ccbf810..c6d24b93f8 100644 --- a/fla/models/__init__.py +++ b/fla/models/__init__.py @@ -16,6 +16,7 @@ from fla.models.gsa import GSAConfig, GSAForCausalLM, GSAModel from fla.models.hgrn import HGRNConfig, HGRNForCausalLM, HGRNModel from fla.models.hgrn2 import HGRN2Config, HGRN2ForCausalLM, HGRN2Model +from fla.models.kda import KDAConfig, KDAForCausalLM, KDAModel from fla.models.lightnet import LightNetConfig, LightNetForCausalLM, LightNetModel from fla.models.linear_attn import LinearAttentionConfig, LinearAttentionForCausalLM, LinearAttentionModel from fla.models.log_linear_mamba2 import LogLinearMamba2Config, LogLinearMamba2ForCausalLM, LogLinearMamba2Model @@ -46,6 +47,7 @@ 'GSAConfig', 'GSAForCausalLM', 'GSAModel', 'HGRNConfig', 'HGRNForCausalLM', 'HGRNModel', 'HGRN2Config', 'HGRN2ForCausalLM', 'HGRN2Model', + 'KDAConfig', 'KDAForCausalLM', 'KDAModel', 'LightNetConfig', 'LightNetForCausalLM', 'LightNetModel', 'LinearAttentionConfig', 'LinearAttentionForCausalLM', 'LinearAttentionModel', 'LogLinearMamba2Config', 'LogLinearMamba2ForCausalLM', 'LogLinearMamba2Model', diff --git a/fla/models/kda/__init__.py b/fla/models/kda/__init__.py new file mode 100644 index 0000000000..b3f4190fa4 --- /dev/null +++ b/fla/models/kda/__init__.py @@ -0,0 +1,12 @@ +# -*- coding: utf-8 -*- + +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.kda.configuration_kda import KDAConfig +from fla.models.kda.modeling_kda import KDAForCausalLM, KDAModel + +AutoConfig.register(KDAConfig.model_type, KDAConfig, exist_ok=True) +AutoModel.register(KDAConfig, KDAModel, exist_ok=True) +AutoModelForCausalLM.register(KDAConfig, KDAForCausalLM, exist_ok=True) + +__all__ = ['KDAConfig', 'KDAForCausalLM', 'KDAModel'] diff --git a/fla/models/kda/configuration_kda.py b/fla/models/kda/configuration_kda.py new file mode 100644 index 0000000000..8cffab901f --- /dev/null +++ b/fla/models/kda/configuration_kda.py @@ -0,0 +1,87 @@ +# -*- coding: utf-8 -*- + +from typing import Dict, Optional + +from transformers.configuration_utils import PretrainedConfig + + +class KDAConfig(PretrainedConfig): + model_type = 'kda' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + attn_mode: str = "chunk", + hidden_size: int = 2048, + expand_v: float = 1.0, + use_short_conv: bool = True, + allow_neg_eigval: bool = False, + conv_size: int = 4, + head_dim: int = 128, + num_heads: int = 16, + num_v_heads: Optional[int] = None, + max_position_embeddings: int = 2048, + hidden_ratio: Optional[int] = 4, + intermediate_size: Optional[int] = None, + hidden_act: str = "swish", + num_hidden_layers: int = 24, + norm_eps: float = 1e-6, + attn: Optional[Dict] = None, + use_cache: bool = True, + pad_token_id: Optional[int] = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + initializer_range: float = 0.02, + fuse_norm: bool = True, + fuse_swiglu: bool = True, + fuse_cross_entropy: bool = True, + use_l2warp: bool = False, + vocab_size: int = 32000, + **kwargs + ): + self.attn_mode = attn_mode + self.hidden_size = hidden_size + self.expand_v = expand_v + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.head_dim = head_dim + self.num_heads = num_heads + self.num_v_heads = num_v_heads + self.max_position_embeddings = max_position_embeddings + + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.hidden_act = hidden_act + self.num_hidden_layers = num_hidden_layers + self.norm_eps = norm_eps + self.attn = attn + self.use_cache = use_cache + self.initializer_range = initializer_range + + self.fuse_norm = fuse_norm + self.fuse_swiglu = fuse_swiglu + self.fuse_cross_entropy = fuse_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + self.allow_neg_eigval = allow_neg_eigval + + if attn is not None: + if not isinstance(attn, Dict): + raise ValueError("attn must be a dictionary") + if 'layers' not in attn: + raise ValueError("Layer indices must be provided to initialize hybrid attention layers") + if 'num_heads' not in attn: + raise ValueError("Number of heads must be provided to initialize hybrid attention layers") + attn['num_kv_heads'] = attn.get('num_kv_heads', attn['num_heads']) + attn['qkv_bias'] = attn.get('qkv_bias', False) + attn['window_size'] = attn.get('window_size', None) + attn['rope_theta'] = attn.get('rope_theta', 10000.) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/kda/modeling_kda.py b/fla/models/kda/modeling_kda.py new file mode 100644 index 0000000000..0f66974807 --- /dev/null +++ b/fla/models/kda/modeling_kda.py @@ -0,0 +1,394 @@ +# -*- coding: utf-8 -*- + +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union + +import torch +import torch.nn as nn +from transformers.generation import GenerationMixin +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.attn import Attention +from fla.layers.kda import KimiDeltaAttention +from fla.models.kda.configuration_kda import KDAConfig +from fla.models.utils import Cache +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss +from fla.modules import GatedMLP as KDAMLP +from fla.modules import RMSNorm +from fla.modules.l2warp import l2_warp + +try: + from torch.distributed.tensor import DTensor +except (ImportError, AttributeError): + DTensor = None + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +logger = logging.get_logger(__name__) + + +class KDABlock(nn.Module): + def __init__(self, config: KDAConfig, layer_idx: int): + super().__init__() + + self.config = config + self.layer_idx = layer_idx + + self.attn_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + if config.attn is not None and layer_idx in config.attn['layers']: + self.attn = Attention( + hidden_size=config.hidden_size, + num_heads=config.attn['num_heads'], + num_kv_heads=config.attn['num_kv_heads'], + qkv_bias=config.attn['qkv_bias'], + window_size=config.attn['window_size'], + rope_theta=config.attn['rope_theta'], + max_position_embeddings=config.max_position_embeddings, + layer_idx=layer_idx + ) + else: + self.attn = KimiDeltaAttention( + mode=config.attn_mode, + hidden_size=config.hidden_size, + expand_v=config.expand_v, + head_dim=config.head_dim, + num_heads=config.num_heads, + num_v_heads=config.num_v_heads, + use_short_conv=config.use_short_conv, + allow_neg_eigval=config.allow_neg_eigval, + conv_size=config.conv_size, + norm_eps=config.norm_eps, + layer_idx=layer_idx, + ) + self.mlp_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + self.mlp = KDAMLP( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + fuse_swiglu=config.fuse_swiglu + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + use_cache: Optional[bool] = False, + output_attentions: Optional[bool] = False, + **kwargs: Unpack[Dict] + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + residual = hidden_states + hidden_states = self.attn_norm(hidden_states) + hidden_states, attentions, past_key_values = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs + ) + if self.config.fuse_norm: + hidden_states, residual = self.mlp_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = self.mlp(hidden_states, **kwargs) + hidden_states = residual + hidden_states + + outputs = (hidden_states, attentions, past_key_values) + + return outputs + + +class KDAPreTrainedModel(PreTrainedModel): + + config_class = KDAConfig + base_model_prefix = 'model' + supports_gradient_checkpointing = True + _no_split_modules = ['KDABlock'] + _supports_cache_class = True + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights( + self, + module: nn.Module, + prenorm_residual_strategy: Optional[str] = None, + num_residuals_per_layer: int = 2, + ): + if isinstance(module, KimiDeltaAttention) and next(module.parameters()).device.type != 'meta': + with torch.no_grad(): + module.A.copy_(nn.init.uniform_(module.A, a=1, b=16).log()) + dt = torch.exp( + nn.init.uniform_(module.f_proj[1].bias) * (math.log(0.1) - math.log(0.001)) + math.log(0.001) + ).clamp(min=1e-4) + inv_dt = dt + torch.log(-torch.expm1(-dt)) + module.f_proj[1].bias.copy_(inv_dt) + module.f_proj[1].bias._is_hf_initialized = True + if isinstance(module, (nn.Linear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None and not getattr(module.bias, '_is_hf_initialized', False): + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + if prenorm_residual_strategy is not None: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, 'o_proj'): + p = module.o_proj.weight + elif hasattr(module, 'down_proj'): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + if prenorm_residual_strategy == 'rescale': + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + elif prenorm_residual_strategy == 'zero': + nn.init.zeros_(p) + else: + raise ValueError(f"Invalid prenorm_residual_strategy: {prenorm_residual_strategy}") + + +class KDAModel(KDAPreTrainedModel): + + def __init__(self, config: KDAConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([KDABlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + self.norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, # noqa + inputs_embeds: Optional[torch.FloatTensor] = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + **kwargs: Unpack[Dict] + ) -> Union[Tuple, BaseModelOutputWithPast]: + if output_attentions: + warnings.warn("`KDAModel` does not `output_attentions` now, setting it to `False`.") + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + hidden_states = inputs_embeds + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + if self.gradient_checkpointing and self.training and use_cache: + logger.warning_once("`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...") + use_cache = False + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + for layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if self.gradient_checkpointing and self.training: + hidden_states, attentions, past_key_values = self._gradient_checkpointing_func( + layer.__call__, + hidden_states, + attention_mask, + past_key_values, + use_cache, + output_attentions, + **kwargs + ) + else: + hidden_states, attentions, past_key_values = layer( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs + ) + + if output_attentions: + all_attns += (attentions,) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns] if i is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns + ) + + +class KDAForCausalLM(KDAPreTrainedModel, GenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = KDAModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embeddings + + def set_input_embeddings(self, value): + self.model.embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def generate(self, *args, **kwargs): + try: + return super().generate(*args, **kwargs) + except AttributeError as exception: + if 'past_key_values' in str(exception): + raise AttributeError( + f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, " + f"which is not supported for {self.__class__.__name__}. " + f"Try another generation strategy instead. " + f"For the available generation strategies, check this doc: " + f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies" + ) + else: + raise exception + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + logits_to_keep: Optional[int] = 0, + **kwargs: Unpack[Dict] + ) -> Union[Tuple, CausalLMOutputWithPast]: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs + ) + + hidden_states = outputs[0] + fuse_linear_and_cross_entropy = self.config.fuse_cross_entropy and self.training and labels is not None + + loss, logits = None, None + if not fuse_linear_and_cross_entropy or labels is None: + logits = self.lm_head(hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:]) + if labels is not None: + if getattr(self, 'criterion', None) is None: + if fuse_linear_and_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if fuse_linear_and_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/ops/__init__.py b/fla/ops/__init__.py index e45884b6ff..af9413acc5 100644 --- a/fla/ops/__init__.py +++ b/fla/ops/__init__.py @@ -16,6 +16,7 @@ from .gla import chunk_gla, fused_chunk_gla, fused_recurrent_gla from .gsa import chunk_gsa, fused_recurrent_gsa from .hgrn import fused_recurrent_hgrn +from .kda import chunk_kda, fused_recurrent_kda from .lightning_attn import chunk_lightning_attn, fused_recurrent_lightning_attn from .linear_attn import chunk_linear_attn, fused_chunk_linear_attn, fused_recurrent_linear_attn from .log_linear_attn import chunk_log_linear_attn @@ -37,6 +38,7 @@ 'chunk_comba', 'fused_recurrent_comba', 'chunk_dplr_delta_rule', 'chunk_iplr_delta_rule', 'fused_recurrent_dplr_delta_rule', 'fused_recurrent_iplr_delta_rule', + 'chunk_kda', 'fused_recurrent_kda', 'chunk_gla', 'fused_chunk_gla', 'fused_recurrent_gla', 'chunk_gsa', 'fused_recurrent_gsa', 'fused_recurrent_hgrn', diff --git a/fla/ops/gla/chunk.py b/fla/ops/gla/chunk.py index f51451f59d..e345f9267d 100644 --- a/fla/ops/gla/chunk.py +++ b/fla/ops/gla/chunk.py @@ -97,8 +97,9 @@ def chunk_gla_fwd_A_kernel_intra_sub_inter( }) @triton.autotune( configs=[ - triton.Config({}, num_warps=num_warps) + triton.Config({}, num_warps=num_warps, num_stages=num_stages) for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3] ], key=["BK", "BT"], **autotune_cache_kwargs @@ -283,10 +284,11 @@ def chunk_gla_fwd_A_kernel_intra_sub_intra_merge( }) @triton.autotune( configs=[ - triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps) + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) for BK in [32, 64] for BV in [64, 128] for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] ], key=['BT'], **autotune_cache_kwargs @@ -362,8 +364,9 @@ def chunk_gla_fwd_kernel_o( }) @triton.autotune( configs=[ - triton.Config({}, num_warps=num_warps) + triton.Config({}, num_warps=num_warps, num_stages=num_stages) for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] ], key=['BK', 'NC', 'BT'], **autotune_cache_kwargs @@ -497,8 +500,9 @@ def chunk_gla_bwd_kernel_intra( }) @triton.autotune( configs=[ - triton.Config({}, num_warps=num_warps) + triton.Config({}, num_warps=num_warps, num_stages=num_stages) for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] ], key=['BV', 'BT'], **autotune_cache_kwargs @@ -547,10 +551,11 @@ def chunk_gla_bwd_kernel_dA( }) @triton.autotune( configs=[ - triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps) + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) for BK in BK_LIST for BV in BV_LIST for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] ], key=['BT'], **autotune_cache_kwargs @@ -1184,8 +1189,7 @@ def forward( output_final_state, cu_seqlens, ): - T = q.shape[1] - chunk_size = min(64, max(16, triton.next_power_of_2(T))) + chunk_size = min(64, max(16, triton.next_power_of_2(q.shape[1]))) g_cumsum, A, _, ht, o = chunk_gla_fwd( q=q, @@ -1294,14 +1298,12 @@ def chunk_gla( >>> q, k, v, g = map(lambda x: rearrange(x, 'b t h d -> 1 (b t) h d'), (q, k, v, g)) # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) - >>> o_var, ht_var = chunk_gla( + >>> o, ht = chunk_gla( q, k, v, g, initial_state=h0, output_final_state=True, cu_seqlens=cu_seqlens ) - >>> assert o.allclose(o_var.view(o.shape)) - >>> assert ht.allclose(ht_var) """ if cu_seqlens is not None: if q.shape[0] != 1: diff --git a/fla/ops/kda/__init__.py b/fla/ops/kda/__init__.py new file mode 100644 index 0000000000..5bc7120181 --- /dev/null +++ b/fla/ops/kda/__init__.py @@ -0,0 +1,7 @@ +from .chunk import chunk_kda +from .fused_recurrent import fused_recurrent_kda + +__all__ = [ + "chunk_kda", + "fused_recurrent_kda" +] diff --git a/fla/ops/kda/chunk.py b/fla/ops/kda/chunk.py new file mode 100644 index 0000000000..cbee39a3a4 --- /dev/null +++ b/fla/ops/kda/chunk.py @@ -0,0 +1,351 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from typing import Optional + +import torch + +from fla.modules.l2norm import l2norm_bwd, l2norm_fwd +from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h +from fla.ops.common.chunk_o import chunk_bwd_dv_local +from fla.ops.gla.chunk import chunk_gla_bwd_dA, chunk_gla_fwd_o_gk +from fla.ops.kda.chunk_inter import chunk_kda_bwd_dqkwg +from fla.ops.kda.chunk_intra import chunk_kda_bwd_intra, chunk_kda_fwd_intra +from fla.ops.kda.wy_fast import prepare_wy_repr_bwd, recompute_w_u_fwd +from fla.ops.utils import chunk_local_cumsum +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + + +def chunk_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: Optional[torch.LongTensor] = None +): + chunk_size = 64 + g = chunk_local_cumsum(g, chunk_size=chunk_size, cu_seqlens=cu_seqlens) + # the intra Aqk is kept in fp32 + # the computation has very marginal effect on the entire throughput + Aqk, Akk = chunk_kda_fwd_intra( + q=q, + k=k, + gk=g, + beta=beta, + scale=scale, + cu_seqlens=cu_seqlens, + output_dtype=torch.float32 + ) + w, u, _, kg = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=Akk, + gk=g, + cu_seqlens=cu_seqlens, + ) + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + k=kg, + w=w, + u=u, + gk=g, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + + o = chunk_gla_fwd_o_gk( + q=q, + v=v_new, + g=g, + A=Aqk, + h=h, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size + ) + return g, o, Aqk, Akk, final_state + + +def chunk_kda_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + Aqk: torch.Tensor, + Akk: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + do: torch.Tensor, + dht: torch.Tensor, + cu_seqlens: Optional[torch.LongTensor] = None, +): + chunk_size = 64 + w, u, qg, kg = recompute_w_u_fwd( + q=q, + k=k, + v=v, + beta=beta, + A=Akk, + gk=g, + cu_seqlens=cu_seqlens, + ) + h, v_new, _ = chunk_gated_delta_rule_fwd_h( + k=kg, + w=w, + u=u, + gk=g, + initial_state=initial_state, + output_final_state=False, + cu_seqlens=cu_seqlens, + ) + dv = chunk_bwd_dv_local( + q=q, + k=k, + do=do, + A=Aqk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size + ) + + dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( + q=qg, + k=kg, + w=w, + gk=g, + h0=initial_state, + dht=dht, + do=do, + dv=dv, + scale=scale, + cu_seqlens=cu_seqlens, + ) + + # dq dk in fp32 + dAqk = chunk_gla_bwd_dA( + v=v_new, + do=do, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size + ) + dq, dk, dw, dg = chunk_kda_bwd_dqkwg( + q=q, + k=k, + v=v_new, + w=w, + g=g, + h=h, + dv=dv, + do=do, + dh=dh, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + ) + dk2, dv, db, dg2, dAkk = prepare_wy_repr_bwd( + k=k, + v=v, + beta=beta, + gk=g, + A=Akk, + dw=dw, + du=dv, + cu_seqlens=cu_seqlens, + ) + dq, dk2, db, dg2 = chunk_kda_bwd_intra( + q=q, + k=k, + g=g, + beta=beta, + dAqk=dAqk, + dAkk=dAkk, + dq=dq, + dk=dk2, + db=db, + dg=dg2, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size + ) + dk.add_(dk2) + dg.add_(dg2) + return dq, dk, dv, db, dg, dh0 + + +class ChunkKDAFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: Optional[torch.LongTensor] = None, + ): + q_rstd, k_rstd = None, None + if use_qk_l2norm_in_kernel: + q, q_rstd = l2norm_fwd(q) + k, k_rstd = l2norm_fwd(k) + + g, o, Aqk, Akk, final_state = chunk_kda_fwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + ctx.save_for_backward(q, q_rstd, k, k_rstd, v, g, beta, Aqk, Akk, initial_state, cu_seqlens) + ctx.scale = scale + ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel + return o.to(q.dtype), final_state + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward( + ctx, + do: torch.Tensor, + dht: torch.Tensor + ): + q, q_rstd, k, k_rstd, v, g, beta, Aqk, Akk, initial_state, cu_seqlens = ctx.saved_tensors + dq, dk, dv, db, dg, dh0 = chunk_kda_bwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + Aqk=Aqk, + Akk=Akk, + scale=ctx.scale, + initial_state=initial_state, + do=do, + dht=dht, + cu_seqlens=cu_seqlens, + ) + if ctx.use_qk_l2norm_in_kernel: + dq = l2norm_bwd(q, q_rstd, dq) + dk = l2norm_bwd(k, k_rstd, dk) + return dq.to(q), dk.to(k), dv.to(v), dg.to(g), db.to(beta), None, dh0, None, None, None + + +@torch.compiler.disable +def chunk_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: Optional[torch.LongTensor] = None, + **kwargs +): + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + g (torch.Tensor): + (forget) gating tensor (in log space!) of shape `[B, T, H]`. + beta (torch.Tensor): + betas of shape `[B, T, H]`. + scale (Optional[float]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + use_qk_l2norm_in_kernel (bool): + Whether to apply L2norm to the q,k tensor internally. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.kda import chunk_kda + # inputs with equal lengths + >>> B, T, H, K, V = 4, 2048, 4, 512, 512 + >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') + >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1) + >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda') + >>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda').sigmoid() + >>> g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda')) + >>> h0 = torch.randn(B, H, K, V, dtype=torch.bfloat16, device='cuda') + >>> o, ht = chunk_kda( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o, ht = chunk_kda( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing." + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}." + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + o, final_state = ChunkKDAFunction.apply( + q, + k, + v, + g, + beta, + scale, + initial_state, + output_final_state, + use_qk_l2norm_in_kernel, + cu_seqlens, + ) + return o, final_state diff --git a/fla/ops/kda/chunk_inter.py b/fla/ops/kda/chunk_inter.py new file mode 100644 index 0000000000..105b9d3336 --- /dev/null +++ b/fla/ops/kda/chunk_inter.py @@ -0,0 +1,188 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.op import exp +from fla.utils import autotune_cache_kwargs, check_shared_mem + +BK_LIST = [32, 64] if check_shared_mem() else [16, 32] +BV_LIST = [64, 128] if check_shared_mem('ampere') else [16, 32] + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for BK in BK_LIST + for BV in BV_LIST + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['BT'], + **autotune_cache_kwargs +) +@triton.jit(do_not_specialize=['T']) +def chunk_kda_bwd_kernel_inter( + q, + k, + v, + g, + h, + do, + dh, + dq, + dk, + dv, + dw, + dg, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr +): + i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + v += (bos * H + i_h) * V + g += (bos * H + i_h) * K + h += (i_tg * H + i_h) * K*V + do += (bos * H + i_h) * V + dh += (i_tg * H + i_h) * K*V + dq += (bos * H + i_h) * K + dk += (bos * H + i_h) * K + dw += (bos * H + i_h) * K + dv += (bos * H + i_h) * V + dg += (bos * H + i_h) * K + + p_g = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_g = tl.load(p_g, boundary_check=(0, 1)) + p_gn = g + (min(T, i_t * BT + BT) - 1) * H*K + o_k + b_gn = tl.load(p_gn, mask=m_k, other=0) + b_dq = tl.zeros([BT, BK], dtype=tl.float32) + b_dk = tl.zeros([BT, BK], dtype=tl.float32) + b_dw = tl.zeros([BT, BK], dtype=tl.float32) + b_dgk = tl.zeros([BK,], dtype=tl.float32) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [BV, BK] + b_h = tl.load(p_h, boundary_check=(0, 1)) + b_dh = tl.load(p_dh, boundary_check=(0, 1)) + + # [BK] + b_dgk += tl.sum(b_h * b_dh, axis=0) + # [BT, BK] + b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) + b_dk += tl.dot(b_v, b_dh.to(b_v.dtype)) + + p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_dv = tl.load(p_dv, boundary_check=(0, 1)) + b_dw += tl.dot(b_dv.to(b_v.dtype), b_h.to(b_v.dtype)) + + p_dw = tl.make_block_ptr(dw, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + tl.store(p_dw, -b_dw.to(p_dw.dtype.element_ty), boundary_check=(0, 1)) + + b_dgk *= exp(b_gn) + b_dq *= scale + b_dq = b_dq * exp(b_g) + b_dk = b_dk * exp(b_gn[None, :] - b_g) + + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dg = tl.make_block_ptr(dg, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_dgk += tl.sum(b_dk * b_k, axis=0) + b_dg = b_q * b_dq - b_k * b_dk + b_dg = b_dg - tl.cumsum(b_dg, axis=0) + tl.sum(b_dg, axis=0)[None, :] + b_dgk[None, :] + + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_kda_bwd_dqkwg( + q: torch.Tensor, + k: torch.Tensor, + w: torch.Tensor, + v: torch.Tensor, + h: torch.Tensor, + g: torch.Tensor, + do: torch.Tensor, + dh: torch.Tensor, + dv: torch.Tensor, + scale: Optional[float] = None, + cu_seqlens: Optional[torch.LongTensor] = None, + chunk_size: int = 64 +): + B, T, H, K, V = *k.shape, v.shape[-1] + BT = min(chunk_size, max(16, triton.next_power_of_2(T))) + + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + dq = torch.empty_like(q, dtype=torch.float) + dk = torch.empty_like(k, dtype=torch.float) + dw = torch.empty_like(w) + dg = torch.empty_like(g) + def grid(meta): return (triton.cdiv(K, meta['BK']), NT, B * H) + chunk_kda_bwd_kernel_inter[grid]( + q=q, + k=k, + v=v, + g=g, + h=h, + do=do, + dh=dh, + dq=dq, + dk=dk, + dv=dv, + dw=dw, + dg=dg, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + ) + return dq, dk, dw, dg diff --git a/fla/ops/kda/chunk_intra.py b/fla/ops/kda/chunk_intra.py new file mode 100644 index 0000000000..59214dc1dc --- /dev/null +++ b/fla/ops/kda/chunk_intra.py @@ -0,0 +1,540 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import chunk_local_cumsum, prepare_chunk_indices, solve_tril +from fla.ops.utils.op import exp +from fla.utils import autotune_cache_kwargs + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64] + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["BC"], + **autotune_cache_kwargs +) +@triton.jit(do_not_specialize=['T']) +def chunk_kda_fwd_kernel_intra_sub_inter( + q, + k, + g, + beta, + Aqk, + Akk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + NC: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + i_i, i_j = i_c // NC, i_c % NC + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT + i_i * BC >= T: + return + if i_i <= i_j: + return + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + g += (bos * H + i_h) * K + Aqk += (bos * H + i_h) * BT + Akk += (bos * H + i_h) * BT + + p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)) + + b_Aqk = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk = tl.zeros([BC, BC], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_g = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + b_kt = tl.make_block_ptr(k, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) + p_gk = tl.make_block_ptr(g, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) + + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + # [BK,] + b_gn = tl.load(g + (i_t * BT + i_i * BC) * H*K + o_k, mask=m_k, other=0) + # [BC, BK] + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) * exp(b_g - b_gn[None, :]) + # [BK, BC] + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kt = tl.load(b_kt, boundary_check=(0, 1)) + # [BC, BC] + b_ktg = b_kt * exp(b_gn[:, None] - b_gk) + b_Akk += tl.dot(b_k, b_ktg) + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_qg = b_q * exp(b_g - b_gn[None, :]) * scale + b_Aqk += tl.dot(b_qg, b_ktg) + + b_Akk *= b_b[:, None] + + p_Akk = tl.make_block_ptr(Akk, (T, BT), (H*BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) + tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + p_Aqk = tl.make_block_ptr(Aqk, (T, BT), (H*BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4, 8] + ], + key=["BK", "BT"], + **autotune_cache_kwargs +) +@triton.jit(do_not_specialize=['T']) +def chunk_kda_fwd_kernel_intra_sub_intra( + q, + k, + g, + beta, + Aqk, + Akk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT + i_i * BC >= T: + return + + o_i = tl.arange(0, BC) + o_k = tl.arange(0, BK) + m_k = o_k < K + m_A = (i_t * BT + i_i * BC + o_i) < T + o_A = (bos + i_t * BT + i_i * BC + o_i) * H*BT + i_h * BT + i_i * BC + + p_q = tl.make_block_ptr(q + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), (1, 0)) + p_g = tl.make_block_ptr(g + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)) + + p_b = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h + b_k = b_k * tl.load(p_b, mask=m_A, other=0)[:, None] + + p_kt = k + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k + p_gk = g + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k + + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + b_kt = tl.load(p_kt, mask=m_k, other=0).to(tl.float32) + b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32) + b_ktg = b_kt[None, :] * exp(b_g - b_gk[None, :]) + b_Aqk = tl.sum(b_q * b_ktg, 1) + b_Aqk = tl.where(o_i >= j, b_Aqk * scale, 0.) + b_Akk = tl.sum(b_k * b_ktg, 1) + b_Akk = tl.where(o_i > j, b_Akk, 0.) + tl.store(Aqk + o_A + j, b_Aqk, mask=m_A) + tl.store(Akk + o_A + j, b_Akk, mask=m_A) + p_kt += H*K + p_gk += H*K + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['BK', 'NC', 'BT'], + **autotune_cache_kwargs +) +@triton.jit(do_not_specialize=['B', 'T']) +def chunk_kda_bwd_kernel_intra( + q, + k, + g, + beta, + dAqk, + dAkk, + dq, + dq2, + dk, + dk2, + dg, + db, + cu_seqlens, + chunk_indices, + B, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + NC: tl.constexpr, + IS_VARLEN: tl.constexpr +): + i_kc, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + i_k, i_i = i_kc // NC, i_kc % NC + + all = B * T + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + else: + bos, eos = i_b * T, i_b * T + T + T = eos - bos + if i_t * BT + i_i * BC >= T: + return + + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + g += (bos * H + i_h) * K + beta += bos * H + i_h + + dAqk += (bos * H + i_h) * BT + dAkk += (bos * H + i_h) * BT + dq += (bos * H + i_h) * K + dq2 += (bos * H + i_h) * K + dk += (bos * H + i_h) * K + dk2 += (bos * H + i_h) * K + dg += (bos * H + i_h) * K + db += (i_k * all + bos) * H + i_h + + p_g = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + b_g = tl.load(p_g, boundary_check=(0, 1)) + + p_b = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)) + + b_dq2 = tl.zeros([BC, BK], dtype=tl.float32) + b_dk2 = tl.zeros([BC, BK], dtype=tl.float32) + if i_i > 0: + p_gn = g + (i_t * BT + i_i * BC) * H*K + o_k + # [BK,] + b_gn = tl.load(p_gn, mask=m_k, other=0) + for i_j in range(0, i_i): + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) + p_gk = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) + p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (H*BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) + p_dAkk = tl.make_block_ptr(dAkk, (T, BT), (H*BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) + # [BC, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kg = b_k * exp(b_gn[None, :] - b_gk) + # [BC, BC] + b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) + b_dAkk = tl.load(p_dAkk, boundary_check=(0, 1)) + # [BC, BK] + b_dq2 += tl.dot(b_dAqk, b_kg) + b_dk2 += tl.dot(b_dAkk, b_kg) + b_dq2 *= exp(b_g - b_gn[None, :]) + b_dk2 *= exp(b_g - b_gn[None, :]) + + o_i = tl.arange(0, BC) + m_dA = (i_t * BT + i_i * BC + o_i) < T + o_dA = (i_t * BT + i_i * BC + o_i) * H*BT + i_i * BC + p_kj = k + (i_t * BT + i_i * BC) * H*K + o_k + p_gkj = g + (i_t * BT + i_i * BC) * H*K + o_k + + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + # [BC] + b_dAqk = tl.load(dAqk + o_dA + j, mask=m_dA, other=0) + b_dAkk = tl.load(dAkk + o_dA + j, mask=m_dA, other=0) + # [BK] + b_kj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) + b_gkj = tl.load(p_gkj, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + m_i = o_i[:, None] >= j + # [BC, BK] + b_dq2 += tl.where(m_i, b_dAqk[:, None] * b_kj[None, :] * exp(b_g - b_gkj[None, :]), 0.) + b_dk2 += tl.where(m_i, b_dAkk[:, None] * b_kj[None, :] * exp(b_g - b_gkj[None, :]), 0.) + + p_kj += H*K + p_gkj += H*K + b_db = tl.sum(b_dk2 * b_k, 1) + b_dk2 *= b_b[:, None] + + p_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_dq2 = tl.make_block_ptr(dq2, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_db = tl.make_block_ptr(db, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,)) + + b_dg = b_q * b_dq2 + b_dq2 = b_dq2 + tl.load(p_dq, boundary_check=(0, 1)) + tl.store(p_dq2, b_dq2.to(p_dq2.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + + tl.debug_barrier() + b_dkt = tl.zeros([BC, BK], dtype=tl.float32) + + NC = min(NC, tl.cdiv(T - i_t * BT, BC)) + if i_i < NC - 1: + p_gn = g + (min(i_t * BT + i_i * BC + BC, T) - 1) * H*K + o_k + # [BK,] + b_gn = tl.load(p_gn, mask=m_k, other=0) + for i_j in range(i_i + 1, NC): + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k*BK), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) + p_gk = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k*BK), (BC, BK), (1, 0)) + p_b = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT + i_j * BC,), (BC,), (0,)) + p_dAqk = tl.make_block_ptr(dAqk, (BT, T), (1, H*BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) + p_dAkk = tl.make_block_ptr(dAkk, (BT, T), (1, H*BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) + # [BC] + b_b = tl.load(p_b, boundary_check=(0,)) + # [BC, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_kb = tl.load(p_k, boundary_check=(0, 1)) * b_b[:, None] + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + # [BC, BC] + b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) + b_dAkk = tl.load(p_dAkk, boundary_check=(0, 1)) + + o_j = i_t * BT + i_j * BC + o_i + m_j = o_j < T + # [BC, BK] + b_qg = b_q * tl.where(m_j[:, None], exp(b_gk - b_gn[None, :]), 0) + b_kbg = b_kb * tl.where(m_j[:, None], exp(b_gk - b_gn[None, :]), 0) + # [BC, BK] + # (SY 09/17) important to not use bf16 here to have a good precision. + b_dkt += tl.dot(b_dAqk, b_qg) + b_dkt += tl.dot(b_dAkk, b_kbg) + b_dkt *= exp(b_gn[None, :] - b_g) + o_dA = (i_t * BT + i_i * BC) * H*BT + i_i * BC + o_i + p_qj = q + (i_t * BT + i_i * BC) * H*K + o_k + p_kj = k + (i_t * BT + i_i * BC) * H*K + o_k + p_gkj = g + (i_t * BT + i_i * BC) * H*K + o_k + p_bj = beta + (i_t * BT + i_i * BC) * H + + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + # [BC,] + b_dAqk = tl.load(dAqk + o_dA + j * H*BT) + b_dAkk = tl.load(dAkk + o_dA + j * H*BT) + # [BK,] + b_qj = tl.load(p_qj, mask=m_k, other=0).to(tl.float32) + b_kbj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) * tl.load(p_bj) + b_gkj = tl.load(p_gkj, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + m_i = o_i[:, None] <= j + b_dkt += tl.where(m_i, b_dAqk[:, None] * b_qj[None, :] * exp(b_gkj[None, :] - b_g), 0.) + b_dkt += tl.where(m_i, b_dAkk[:, None] * b_kbj[None, :] * exp(b_gkj[None, :] - b_g), 0.) + + p_qj += H*K + p_kj += H*K + p_gkj += H*K + p_bj += H + p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_dk2 = tl.make_block_ptr(dk2, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_dg = tl.make_block_ptr(dg, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + + b_dg += (b_dk2 - b_dkt) * b_k + b_dk2 += tl.load(p_dk, boundary_check=(0, 1)) + b_dk2 += b_dkt + + tl.store(p_dk2, b_dk2.to(p_dk2.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_kda_fwd_intra( + q: torch.Tensor, + k: torch.Tensor, + gk: Optional[torch.Tensor] = None, + beta: Optional[torch.Tensor] = None, + scale: Optional[float] = None, + cu_seqlens: Optional[torch.LongTensor] = None, + chunk_size: int = 64, + output_dtype: torch.dtype = torch.float32 +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + The query tensor of shape `[B, T, H, K]`. + k (torch.Tensor): + The key tensor of shape `[B, T, H, K]`. + gk (torch.Tensor): + The cumulative sum of the gate tensor of shape `[B, T, H, K]` applied to the key tensor. Default: `None`. + beta (torch.Tensor): + The beta tensor of shape `[B, T, H]`. Default: `None`. + scale (Optional[float]): + The scale factor. Default: `None`. + cu_seqlens (torch.LongTensor): + The cumulative sequence lengths of the input tensor. + Default: None + chunk_size (int): + The chunk size. Default: 64. + output_dtype (torch.dtype): + The dtype of the output tensor. Default: `torch.float32` + + Returns: + Aqk (torch.Tensor): + The intra Aqk tensor of shape `[B, T, H, BT]` where `BT` is the chunk size. + Akk (torch.Tensor): + The intra Akk tensor of shape `[B, T, H, BT]` where `BT` is the chunk size. + """ + B, T, H, K = k.shape + assert K <= 256 + BT = chunk_size + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + BC = min(16, BT) + NC = triton.cdiv(BT, BC) + BK = max(triton.next_power_of_2(K), 16) + + Aqk = torch.zeros(B, T, H, BT, device=k.device, dtype=output_dtype) + Akk = torch.zeros(B, T, H, BT, device=k.device, dtype=output_dtype) + grid = (NT, NC * NC, B * H) + chunk_kda_fwd_kernel_intra_sub_inter[grid]( + q=q, + k=k, + g=gk, + beta=beta, + Aqk=Aqk, + Akk=Akk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + NC=NC, + ) + + grid = (NT, NC, B * H) + chunk_kda_fwd_kernel_intra_sub_intra[grid]( + q=q, + k=k, + g=gk, + beta=beta, + Aqk=Aqk, + Akk=Akk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + BK=BK, + ) + Akk = solve_tril( + A=Akk, + cu_seqlens=cu_seqlens, + output_dtype=k.dtype + ) + return Aqk, Akk + + +def chunk_kda_bwd_intra( + q: torch.Tensor, + k: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + dAqk: torch.Tensor, + dAkk: torch.Tensor, + dq: torch.Tensor, + dk: torch.Tensor, + db: torch.Tensor, + dg: torch.Tensor, + cu_seqlens: Optional[torch.LongTensor] = None, + chunk_size: int = 64 +): + B, T, H, K = k.shape + BT = chunk_size + BC = min(16, BT) + BK = min(64, triton.next_power_of_2(K)) + + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NC = triton.cdiv(BT, BC) + NK = triton.cdiv(K, BK) + + dq2 = torch.empty_like(q) + dk2 = torch.empty_like(k) + db2 = beta.new_empty(NK, *beta.shape, dtype=torch.float) + dg2 = torch.empty_like(dg, dtype=torch.float) + grid = (NK * NC, NT, B * H) + chunk_kda_bwd_kernel_intra[grid]( + q=q, + k=k, + g=g, + beta=beta, + dAqk=dAqk, + dAkk=dAkk, + dq=dq, + dq2=dq2, + dk=dk, + dk2=dk2, + dg=dg2, + db=db2, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + B=B, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + BK=BK, + NC=NC, + ) + dq = dq2 + dk = dk2 + db = db2.sum(0).add_(db) + dg = chunk_local_cumsum(dg2.add_(dg), chunk_size=chunk_size, reverse=True, cu_seqlens=cu_seqlens) + + return dq, dk2, db, dg diff --git a/fla/ops/kda/fused_recurrent.py b/fla/ops/kda/fused_recurrent.py new file mode 100644 index 0000000000..b57c4476b6 --- /dev/null +++ b/fla/ops/kda/fused_recurrent.py @@ -0,0 +1,114 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from typing import Optional, Tuple + +import torch + +from fla.ops.gated_delta_rule.fused_recurrent import fused_recurrent_gated_delta_rule + + +def fused_recurrent_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor = None, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: Optional[torch.LongTensor] = None, + **kwargs +) -> Tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, HV, V]`. + GVA is applied if `HV > H`. + g (torch.Tensor): + g (decays) of shape `[B, T, HV]`. + beta (torch.Tensor): + betas of shape `[B, T, HV]`. + scale (Optional[float]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, HV, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`. + use_qk_l2norm_in_kernel (Optional[bool]): + Whether to use L2 normalization in the kernel. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, HV, V]`. + final_state (torch.Tensor): + Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.kda import fused_recurrent_kda + # inputs with equal lengths + >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512 + >>> q = torch.randn(B, T, H, K, device='cuda') + >>> k = F.normalize(torch.randn(B, T, H, K, device='cuda'), p=2, dim=-1) + >>> v = torch.randn(B, T, HV, V, device='cuda') + >>> g = F.logsigmoid(torch.rand(B, T, HV, K, device='cuda')) + >>> beta = torch.rand(B, T, HV, device='cuda').sigmoid() + >>> h0 = torch.randn(B, HV, K, V, device='cuda') + >>> o, ht = fused_recurrent_kda( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, g, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, g, beta)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o_var, ht_var = fused_recurrent_kda( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing." + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}." + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + + o, final_state = fused_recurrent_gated_delta_rule( + q=q, + k=k, + v=v, + gk=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + cu_seqlens=cu_seqlens, + ) + return o, final_state diff --git a/fla/ops/kda/gate.py b/fla/ops/kda/gate.py new file mode 100644 index 0000000000..9e17e3ba33 --- /dev/null +++ b/fla/ops/kda/gate.py @@ -0,0 +1,349 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl +from einops import rearrange + +from fla.ops.utils.op import log +from fla.utils import autotune_cache_kwargs, input_guard, is_amd + +BT_LIST_AUTOTUNE = [32, 64, 128] +NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if is_amd else [4, 8, 16, 32] + + +def kda_gate_ref( + g: torch.Tensor, + A: torch.Tensor, + head_k_dim: int, + g_bias: torch.Tensor | None = None, + beta=1.0, threshold=20.0 +) -> torch.Tensor: + """ + Torch reference implementation for KDA gate computation. + + Computes: g = -A.exp().unsqueeze(-1) * softplus(rearrange(g, '... (h d) -> ... h d', d=head_k_dim)) + + Supports both formats: + - Standard: [batch_size, seq_len, num_heads * head_k_dim] + - vLLM: [num_tokens, num_heads * head_k_dim] + + Args: + g: Input tensor of shape [..., num_heads * head_k_dim] + A: Parameter tensor of shape [num_heads] or [1, 1, num_heads, 1] + g_bias : Optional bias tensor added to g before activation, shape [num_heads * head_k_dim] + head_k_dim: Dimension of each head + + Returns: + Output tensor of shape [..., num_heads, head_k_dim] + """ + # Rearrange g to separate heads: [..., H*D] -> [..., H, D] + A = A.view(-1) # Flatten A to [num_heads] to handle any input shape + if g_bias is not None: + g = g + g_bias + g = rearrange(g, '... (h d) -> ... h d', d=head_k_dim) + + # Apply the gate computation: -A.exp().unsqueeze(-1) * softplus(g) + # A: [H] -> [H, 1] for broadcasting + A_exp = -A.float().exp().unsqueeze(-1) # [H, 1] + g_softplus = F.softplus(g.float(), beta, threshold) # [..., H, D] + + return A_exp * g_softplus + + +@triton.autotune( + configs=[ + triton.Config({'BT': bt}, num_warps=nw, num_stages=ns) + for bt in BT_LIST_AUTOTUNE + for nw in NUM_WARPS_AUTOTUNE + for ns in [2, 3] + ], + key=['H', 'D'], + **autotune_cache_kwargs +) +@triton.jit +def kda_gate_fwd_kernel( + g, A, y, + g_bias, + beta: tl.constexpr, + threshold: tl.constexpr, + T, + H, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + HAS_BIAS: tl.constexpr +): + i_t, i_h = tl.program_id(0), tl.program_id(1) + n_t = i_t * BT + + b_a = tl.load(A + i_h).to(tl.float32) + b_a = -tl.exp(b_a) + + stride_row = H * D + stride_col = 1 + + g_ptr = tl.make_block_ptr( + base=g + i_h * D, + shape=(T, D), + strides=(stride_row, stride_col), + offsets=(n_t, 0), + block_shape=(BT, BD), + order=(1, 0), + ) + + y_ptr = tl.make_block_ptr( + base=y + i_h * D, + shape=(T, D), + strides=(stride_row, stride_col), + offsets=(n_t, 0), + block_shape=(BT, BD), + order=(1, 0), + ) + + b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32) + + if HAS_BIAS: + n_d = tl.arange(0, BD) + bias_mask = n_d < D + b_bias = tl.load(g_bias + i_h * D + n_d, mask=bias_mask, other=0.0).to(tl.float32) + b_g = b_g + b_bias[None, :] + + # softplus(x, beta) = (1/beta) * log(1 + exp(beta * x)) + # When beta * x > threshold, use linear approximation x + # Use threshold to switch to linear when beta*x > threshold + g_scaled = b_g * beta + use_linear = g_scaled > threshold + sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) + b_y = b_a * sp + + tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=nw, num_stages=ns) + for nw in NUM_WARPS_AUTOTUNE + for ns in [2, 3] + ], + key=['H', 'D'], + **autotune_cache_kwargs +) +@triton.jit +def kda_gate_bwd_kernel( + g, + A, + dy, + dg, + dA, + g_bias, + beta: tl.constexpr, + threshold: tl.constexpr, + T, + H: tl.constexpr, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + i_t, i_h = tl.program_id(0), tl.program_id(1) + n_t = i_t * BT + + a_h = tl.load(A + i_h).to(tl.float32) + neg_exp_a = -tl.exp(a_h) + + stride_row = H * D + stride_col = 1 + + g_ptr = tl.make_block_ptr( + base=g + i_h * D, + shape=(T, D), + strides=(stride_row, stride_col), + offsets=(n_t, 0), + block_shape=(BT, BD), + order=(1, 0), + ) + dy_ptr = tl.make_block_ptr( + base=dy + i_h * D, + shape=(T, D), + strides=(stride_row, stride_col), + offsets=(n_t, 0), + block_shape=(BT, BD), + order=(1, 0), + ) + dg_ptr = tl.make_block_ptr( + base=dg + i_h * D, + shape=(T, D), + strides=(stride_row, stride_col), + offsets=(n_t, 0), + block_shape=(BT, BD), + order=(1, 0), + ) + + b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32) # [BT, BD] + b_dy = tl.load(dy_ptr, boundary_check=(0, 1)).to(tl.float32) # [BT, BD] + + if HAS_BIAS: + n_d = tl.arange(0, BD) + bias_mask = n_d < D + b_bias = tl.load(g_bias + i_h * D + n_d, mask=bias_mask, other=0.0).to(tl.float32) + b_g = b_g + b_bias[None, :] + + # softplus(g + bias) + g_scaled = b_g * beta + use_linear = g_scaled > threshold + sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) + + sig = tl.sigmoid(g_scaled) + + # grad_g = dy * (-exp(A)) * sigmoid(beta*g) + b_dg = b_dy * (neg_exp_a * sig) + tl.store(dg_ptr, b_dg.to(dg_ptr.dtype.element_ty), boundary_check=(0, 1)) + + contrib = b_dy * (neg_exp_a * sp) + tile_sum = tl.sum(tl.sum(contrib, axis=1), axis=0) + + out_off = i_t * H + i_h + tl.store(dA + out_off, tile_sum) + + +def kda_gate_fwd( + g: torch.Tensor, + A: torch.Tensor, + head_k_dim: int, + g_bias: torch.Tensor | None = None, + beta: float = 1.0, + threshold: float = 20.0 +) -> torch.Tensor: + """ + Forward pass for KDA gate: + input g: [..., H*D] + param A: [H] or [1, 1, H, 1] + beta: softplus beta parameter + threshold: softplus threshold parameter + return : [..., H, D] + """ + orig_shape = g.shape[:-1] + + g = g.view(-1, g.shape[-1]) + T = g.shape[0] + HD = g.shape[1] + H = A.numel() + assert HD == H * head_k_dim + + y = torch.empty_like(g, dtype=torch.float32) + + def grid(meta): return (triton.cdiv(T, meta['BT']), H) + + kda_gate_fwd_kernel[grid]( + g, A, y, g_bias, + beta, threshold, + T, H, head_k_dim, + BD=triton.next_power_of_2(head_k_dim), + HAS_BIAS=g_bias is not None + ) + + y = y.view(*orig_shape, H, head_k_dim) + return y + + +def kda_gate_bwd( + grad_output: torch.Tensor, # [..., H, D] + g: torch.Tensor, # [..., H*D] + A: torch.Tensor, # [H] + head_k_dim: int, + g_bias: torch.Tensor | None = None, + beta: float = 1.0, + threshold: float = 20.0, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + + g_flat = g.view(-1, g.shape[-1]) + T = g_flat.shape[0] + A_ori_shape = A.shape + + H = A.numel() + D = head_k_dim + + dy = grad_output.view(T, H * D) + dg = torch.empty_like(g_flat, dtype=torch.float32) + + BT = 32 + NT = triton.cdiv(T, BT) + dA = torch.empty((NT, H), dtype=torch.float32, device=g.device) + + grid = (triton.cdiv(T, BT), H) + kda_gate_bwd_kernel[grid]( + g_flat, A, dy, dg, dA, g_bias, + beta, threshold, + T, H, D, + BT=BT, + BD=triton.next_power_of_2(D), + HAS_BIAS=g_bias is not None + ) + + dA = dA.sum(0).view(A_ori_shape).type_as(A) + dgbias = dg.sum(0).type_as(g_bias) if g_bias is not None else None + dg = dg.view(g.shape).type_as(g) + return dg, dA, dgbias + + +class KDAGateFunction(torch.autograd.Function): + """ + Autograd function for KDA gate computation. + + Supports both formats: + - Standard: [batch_size, seq_len, num_heads * head_k_dim] + - vLLM: [num_tokens, num_heads * head_k_dim] + """ + + @input_guard + @staticmethod + def forward(ctx, g: torch.Tensor, A: torch.Tensor, head_k_dim: int, + g_bias: torch.Tensor | None = None, + beta: float = 1.0, + threshold: float = 20.0) -> torch.Tensor: + ctx.save_for_backward(g, A) + ctx.g_bias = g_bias + ctx.head_k_dim = head_k_dim + ctx.beta = beta + ctx.threshold = threshold + + return kda_gate_fwd(g, A, head_k_dim, g_bias, beta, threshold) + + @input_guard + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, None, None, None]: + g, A = ctx.saved_tensors + head_k_dim = ctx.head_k_dim + beta = ctx.beta + threshold = ctx.threshold + g_bias = ctx.g_bias + + grad_g, grad_A, grad_gbias = kda_gate_bwd(grad_output, g, A, head_k_dim, g_bias, beta, threshold) + return grad_g, grad_A, None, grad_gbias, None, None + + +def fused_kda_gate(g: torch.Tensor, A: torch.Tensor, head_k_dim: int, + g_bias: torch.Tensor | None = None, + beta: float = 1.0, threshold: float = 20.0) -> torch.Tensor: + """ + Fused KDA gate computation with autograd support. + + Supports both formats: + - Standard: [batch_size, seq_len, num_heads * head_k_dim] + - vLLM: [num_tokens, num_heads * head_k_dim] + + Args: + g: Input tensor of shape [..., num_heads * head_k_dim] + A: Parameter tensor of shape [num_heads] or [1, 1, num_heads, 1] + head_k_dim: Dimension of each head + beta: softplus beta parameter + threshold: softplus threshold parameter + + Returns: + Output tensor of shape [..., num_heads, head_k_dim] + """ + return KDAGateFunction.apply(g, A, head_k_dim, g_bias, beta, threshold) diff --git a/fla/ops/kda/naive.py b/fla/ops/kda/naive.py new file mode 100644 index 0000000000..d5e856d4b3 --- /dev/null +++ b/fla/ops/kda/naive.py @@ -0,0 +1,102 @@ +# -*- coding: utf-8 -*- + +from typing import Optional + +import torch +from einops import rearrange + + +def naive_recurrent_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: Optional[float] = None, + initial_state: Optional[torch.Tensor] = None, + output_final_state: bool = False, +): + dtype = v.dtype + B, T, H, K, V = *q.shape, v.shape[-1] + if scale is None: + scale = K ** -0.5 + + q, k, v, g, beta = map(lambda x: x.to(torch.float), [q, k, v, g, beta]) + q = q * scale + + S = k.new_zeros(B, H, K, V).to(q) + if initial_state is not None: + S += initial_state + o = torch.zeros_like(v) + for i in range(0, T): + q_i, k_i, v_i, g_i, b_i = q[:, i], k[:, i], v[:, i], g[:, i], beta[:, i] + S = S * g_i[..., None].exp() + S = S + torch.einsum('b h k, b h v -> b h k v', b_i[..., None] * k_i, v_i - (k_i[..., None] * S).sum(-2)) + o[:, i] = torch.einsum('b h k, b h k v -> b h v', q_i, S) + if not output_final_state: + S = None + return o.to(dtype), S + + +def naive_chunk_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: Optional[float] = None, + initial_state: Optional[torch.Tensor] = None, + output_final_state: bool = False, + chunk_size: int = 64 +): + dtype = v.dtype + B, T, H, K, V = *q.shape, v.shape[-1] + BT = chunk_size + NT = T // BT + if scale is None: + scale = K ** -0.5 + assert T % BT == 0 + + q, k, v, g, beta = map(lambda x: rearrange(x, 'b (n c) h ... -> b h n c ...', c=BT).to(torch.float), [q, k, v, g, beta]) + q = q * scale + g = g.cumsum(-2) + + # note that diagonal is masked. + mask = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=0) + + A = torch.zeros(*q.shape[:-1], BT, dtype=torch.float, device=q.device) + for i in range(BT): + k_i = k[..., i, :] + g_i = g[..., i:i+1, :] + A[..., i] = torch.einsum('... c d, ... d -> ... c', k * (g - g_i).exp(), k_i) + A = A * beta[..., None] + + A = -A.masked_fill(mask, 0) + for i in range(1, BT): + A[..., i, :i] = A[..., i, :i].clone() + (A[..., i, :, None].clone() * A[..., :, :i].clone()).sum(-2) + A = (A + torch.eye(BT, dtype=torch.float, device=q.device)) * beta[..., None, :] + + w = A @ (g.exp() * k) + u = A @ v + + S = k.new_zeros(B, H, K, V).to(q) + if initial_state is not None: + S += initial_state + o = torch.zeros_like(v) + mask = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=1) + for i in range(0, NT): + # [B, H, BT, ...] + q_i, k_i, u_i, g_i, w_i = q[:, :, i], k[:, :, i], u[:, :, i], g[:, :, i], w[:, :, i] + A = torch.zeros(B, H, BT, BT, dtype=torch.float, device=q.device) + for j in range(BT): + k_j = k[:, :, i, j] + g_j = g[:, :, i, j:j+1, :] + A[..., j] = torch.einsum('... c d, ... d -> ... c', q_i * (g_i - g_j).exp(), k_j) + A = A.masked_fill(mask, 0) + v_i = u_i - w_i @ S + o[:, :, i] = (q_i * g_i.exp()) @ S + A @ v_i + S = S * rearrange(g_i[:, :, -1].exp(), 'b h k -> b h k 1') + S += rearrange((g_i[:, :, -1:] - g_i).exp() * k_i, 'b h c k -> b h k c') @ v_i + if not output_final_state: + S = None + return rearrange(o, 'b h n c d -> b (n c) h d').to(dtype), S diff --git a/fla/ops/kda/wy_fast.py b/fla/ops/kda/wy_fast.py new file mode 100644 index 0000000000..a55bbd600d --- /dev/null +++ b/fla/ops/kda/wy_fast.py @@ -0,0 +1,304 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl + +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.op import exp +from fla.utils import autotune_cache_kwargs, check_shared_mem, is_tf32_supported + + +@triton.heuristics({ + 'STORE_QG': lambda args: args['qg'] is not None, + 'STORE_KG': lambda args: args['kg'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None +}) +@triton.autotune( + configs=[ + triton.Config({'DOT_PRECISION': DOT_PRECISION}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + for DOT_PRECISION in (["tf32x3", "ieee"] if is_tf32_supported else ["ieee"]) + ], + key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], + **autotune_cache_kwargs +) +@triton.jit(do_not_specialize=['T']) +def recompute_w_u_fwd_kernel( + q, + k, + qg, + kg, + v, + beta, + w, + u, + A, + gk, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + STORE_QG: tl.constexpr, + STORE_KG: tl.constexpr, + IS_VARLEN: tl.constexpr, + DOT_PRECISION: tl.constexpr +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + p_b = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)) + + p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + b_A = tl.load(p_A, boundary_check=(0, 1)) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_u = tl.make_block_ptr(u + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_vb = (b_v * b_b[:, None]).to(b_v.dtype) + b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + + for i_k in range(tl.cdiv(K, BK)): + p_w = tl.make_block_ptr(w + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kb = b_k * b_b[:, None] + + p_gk = tl.make_block_ptr(gk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kb *= exp(b_gk) + if STORE_QG: + p_q = tl.make_block_ptr(q + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_qg = tl.make_block_ptr(qg + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_qg = b_q * exp(b_gk) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + if STORE_KG: + last_idx = min(i_t * BT + BT, T) - 1 + + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + b_gn = tl.load(gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.) + b_kg = b_k * exp(b_gn - b_gk) + + p_kg = tl.make_block_ptr(kg + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + + b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], + **autotune_cache_kwargs +) +@triton.jit(do_not_specialize=['T']) +def prepare_wy_repr_bwd_kernel( + k, + v, + beta, + gk, + A, + dA, + dw, + du, + dk, + dv, + db, + dg, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + p_b = tl.make_block_ptr(beta + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_db = tl.make_block_ptr(db + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) + + b_b = tl.load(p_b, boundary_check=(0,)) + b_db = tl.zeros([BT], dtype=tl.float32) + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_dA = tl.zeros([BT, BT], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dw = tl.make_block_ptr(dw + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + # [BT, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + + p_gk = tl.make_block_ptr(gk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + b_gk_exp = exp(tl.load(p_gk, boundary_check=(0, 1))) + b_kbg = b_k * b_b[:, None] * b_gk_exp + b_dw = tl.load(p_dw, boundary_check=(0, 1)) + + b_dA += tl.dot(b_dw, tl.trans(b_kbg).to(b_dw.dtype)) + b_dkbg = tl.dot(b_A, b_dw) + b_dk = b_dkbg * b_gk_exp * b_b[:, None] + b_db += tl.sum(b_dkbg * b_k * b_gk_exp, 1) + b_dg = b_kbg * b_dkbg + + p_dg = tl.make_block_ptr(dg + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_du = tl.make_block_ptr(du + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_vb = (b_v * b_b[:, None]).to(b_v.dtype) + b_du = tl.load(p_du, boundary_check=(0, 1)) + b_dA += tl.dot(b_du, tl.trans(b_vb)) + b_dvb = tl.dot(b_A, b_du) + b_dv = b_dvb * b_b[:, None] + b_db += tl.sum(b_dvb * b_v, 1) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) + b_dA = tl.where(m_A, b_dA, 0) + b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) + b_dA = tl.dot(b_A, b_dA.to(b_A.dtype)) + + b_dA = tl.where(m_A, -b_dA, 0) + + # if using gk, save dA first and handle dk in another kernel + p_dA = tl.make_block_ptr(dA + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + + +def recompute_w_u_fwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + q: Optional[torch.Tensor] = None, + gk: Optional[torch.Tensor] = None, + cu_seqlens: Optional[torch.LongTensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = A.shape[-1] + BK = 64 + BV = 64 + + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + w = torch.empty_like(k) + u = torch.empty_like(v) + qg = torch.empty_like(q) if q is not None else None + kg = torch.empty_like(k) if gk is not None else None + recompute_w_u_fwd_kernel[(NT, B*H)]( + q=q, + k=k, + qg=qg, + kg=kg, + v=v, + beta=beta, + w=w, + u=u, + A=A, + gk=gk, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return w, u, qg, kg + + +def prepare_wy_repr_bwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + gk: torch.Tensor, + A: torch.Tensor, + dw: torch.Tensor, + du: torch.Tensor, + cu_seqlens: Optional[torch.LongTensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = 64 + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + CONST_TILING = 64 if check_shared_mem() else 32 + BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) + BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) + + dk = torch.empty_like(k, dtype=torch.float) + dv = torch.empty_like(v) + dg = torch.empty_like(gk, dtype=torch.float) + dA = torch.empty_like(A, dtype=torch.float) + db = torch.empty_like(beta, dtype=torch.float) + prepare_wy_repr_bwd_kernel[(NT, B * H)]( + k=k, + v=v, + beta=beta, + gk=gk, + A=A, + dA=dA, + dw=dw, + du=du, + dk=dk, + dv=dv, + db=db, + dg=dg, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return dk, dv, db, dg, dA diff --git a/tests/models/test_modeling_kda.py b/tests/models/test_modeling_kda.py new file mode 100644 index 0000000000..808115d10a --- /dev/null +++ b/tests/models/test_modeling_kda.py @@ -0,0 +1,56 @@ +# -*- coding: utf-8 -*- + +import pytest +import torch + +from fla.models import KDAConfig + +from .test_modeling_base import run_test_generation, run_test_model_forward_backward + + +# =================================================================================== +# Test for Modeling (Forward/Backward Pass) +# =================================================================================== +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + ] + ] +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, KDAConfig, use_l2warp=use_l2warp, dtype=dtype) + + +# =================================================================================== +# Test for Generation +# =================================================================================== +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 4, 2000, 8, 64, torch.float16), + ] + ] +) +def test_generation( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): + run_test_generation(L, B, T, H, D, KDAConfig, dtype) diff --git a/tests/ops/test_kda.py b/tests/ops/test_kda.py new file mode 100644 index 0000000000..3b0ec0671e --- /dev/null +++ b/tests/ops/test_kda.py @@ -0,0 +1,381 @@ +# -*- coding: utf-8 -*- + +import os +from typing import List + +import pytest +import torch +import torch.nn.functional as F + +from fla.ops.kda import chunk_kda, fused_recurrent_kda +from fla.ops.kda.gate import fused_kda_gate, kda_gate_ref +from fla.ops.kda.naive import naive_chunk_kda, naive_recurrent_kda +from fla.utils import assert_close, device, is_intel_alchemist + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'scale', 'gate_logit_normalizer', 'dtype'), + [ + pytest.param( + *test, + id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-{}".format(*test) + ) + for test in [ + (1, 64, 1, 64, 1, 1, torch.float), + (2, 512, 3, 60, 1, 1, torch.float), + (4, 1024, 4, 128, 0.1, 1, torch.float), + (4, 1024, 4, 128, 1, 10, torch.float), + ] + ] +) +def test_naive_chunk( + B: int, + T: int, + H: int, + D: int, + scale: float, + gate_logit_normalizer: float, + dtype: torch.dtype, +): + torch.manual_seed(42) + if is_intel_alchemist and D > 128: + pytest.skip(reason='chunk_gated_delta_rule is not supported on alchemist for D>128') + + q = torch.rand(B, T, H, D, dtype=dtype) + k = torch.rand(B, T, H, D, dtype=dtype) + v = torch.rand(B, T, H, D, dtype=dtype) + g = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float)) / gate_logit_normalizer + beta = torch.randn(B, T, H, dtype=dtype).sigmoid() + h0 = torch.randn(B, H, D, D, dtype=torch.float32) + q, k, v, g, beta, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, g, beta, h0)) + + ref, ref_ht = naive_recurrent_kda( + q=F.normalize(q.clone(), p=2, dim=-1), + k=F.normalize(k.clone(), p=2, dim=-1), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + ) + + tri, tri_ht = naive_chunk_kda( + q=F.normalize(q.clone(), p=2, dim=-1), + k=F.normalize(k.clone(), p=2, dim=-1), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + ) + assert_close('o', ref, tri, 0.005) + assert_close('ht', ref_ht, tri_ht, 0.005) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'scale', 'gate_logit_normalizer', 'use_qk_l2norm_in_kernel', 'dtype'), + [ + pytest.param( + *test, + id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-use_qk_l2norm_in_kernel{}-{}".format(*test) + ) + for test in [ + (1, 64, 1, 64, 1, 1, False, torch.float), + (2, 512, 3, 60, 1, 1, False, torch.float), + (3, 1000, 4, 100, 0.1, 1, True, torch.float), + (4, 1024, 4, 128, 0.1, 1, False, torch.float), + ] + ] +) +def test_fused_recurrent( + B: int, + T: int, + H: int, + D: int, + scale: float, + gate_logit_normalizer: float, + use_qk_l2norm_in_kernel: bool, + dtype: torch.dtype, +): + torch.manual_seed(42) + if is_intel_alchemist and D > 128: + pytest.skip(reason='chunk_gated_delta_rule is not supported on alchemist for D>128') + + q = torch.rand(B, T, H, D, dtype=dtype) + k = torch.rand(B, T, H, D, dtype=dtype) + v = torch.rand(B, T, H, D, dtype=dtype) + g = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float)) / gate_logit_normalizer + beta = torch.randn(B, T, H, dtype=dtype).sigmoid() + h0 = torch.randn(B, H, D, D, dtype=torch.float32) + q, k, v, g, beta, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, g, beta, h0)) + + ref, ref_ht = naive_recurrent_kda( + q=F.normalize(q.clone(), p=2, dim=-1), + k=F.normalize(k.clone(), p=2, dim=-1), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + ) + + tri, tri_ht = fused_recurrent_kda( + q=F.normalize(q.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else q.clone(), + k=F.normalize(k.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else k.clone(), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + ) + assert_close('o', ref, tri, 0.005) + assert_close('ht', ref_ht, tri_ht, 0.005) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'scale', 'gate_logit_normalizer', 'mask_p', 'use_qk_l2norm_in_kernel', 'dtype', 'tma'), + [ + pytest.param( + *test, + id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-mask_p{}-use_qk_l2norm_in_kernel{}-{}-tma{}".format(*test) + ) + for test in [ + (1, 63, 1, 64, 1, 1, 0, False, torch.float16, True), + (2, 500, 3, 60, 1, 1, 0, False, torch.float16, True), + (2, 1000, 3, 64, 0.1, 1, 0.5, False, torch.float16, False), + (3, 1024, 4, 100, 1, 0.1, 0, False, torch.float16, False), + (4, 1024, 4, 128, 0.1, 1, 0, False, torch.float16, True), + (4, 1024, 4, 128, 0.1, 1, 0, True, torch.float16, True), + (2, 1500, 4, 128, 0.1, 10, 0, False, torch.float16, False), + (4, 2048, 8, 64, 0.1, 1, 0, False, torch.float16, True), + ] + ] +) +def test_chunk( + B: int, + T: int, + H: int, + D: int, + scale: float, + gate_logit_normalizer: float, + mask_p: float, + use_qk_l2norm_in_kernel: bool, + dtype: torch.dtype, + tma: bool +): + torch.manual_seed(42) + if not tma: + os.environ['FLA_USE_TMA'] = '0' + else: + os.environ['FLA_USE_TMA'] = '1' + q = torch.rand(B, T, H, D, dtype=dtype) + k = torch.rand(B, T, H, D, dtype=dtype) + v = torch.rand(B, T, H, D, dtype=dtype) + g = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float)) / gate_logit_normalizer + g = g * (torch.rand_like(g) > mask_p) + beta = torch.randn(B, T, H, dtype=dtype).sigmoid() + h0 = torch.randn(B, H, D, D, dtype=torch.float32) + q, k, v, g, beta, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, g, beta, h0)) + do = torch.randn_like(v) + dht = torch.randn_like(h0) + + ref, ref_ht = naive_recurrent_kda( + q=F.normalize(q.clone(), p=2, dim=-1), + k=F.normalize(k.clone(), p=2, dim=-1), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + ) + ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True) + ref_dq, ref_dk, ref_dv, ref_dg, ref_db, ref_dh0 = q.grad, k.grad, v.grad, g.grad, beta.grad, h0.grad + q.grad = k.grad = v.grad = g.grad = beta.grad = h0.grad = None + + tri, tri_ht = chunk_kda( + q=F.normalize(q.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else q.clone(), + k=F.normalize(k.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else k.clone(), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + ) + ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True) + tri_dq, tri_dk, tri_dv, tri_dg, tri_db, tri_dh0 = q.grad, k.grad, v.grad, g.grad, beta.grad, h0.grad + q.grad = k.grad = v.grad = g.grad = beta.grad = h0.grad = None + + assert_close('o', ref, tri, 0.005) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.008) + assert_close('dk', ref_dk, tri_dk, 0.008) + assert_close('dv', ref_dv, tri_dv, 0.008) + assert_close('dg', ref_dg, tri_dg, 0.02) + assert_close('db', ref_db, tri_db, 0.02) + assert_close('dh0', ref_dh0, tri_dh0, 0.008) + + +@pytest.mark.parametrize( + ('H', 'D', 'mask_p', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="H{}-D{}-mask_p{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (4, 60, 0, [0, 15], torch.float16), + (4, 64, 0, [0, 256, 500, 1000], torch.float16), + (4, 128, 0.5, [0, 256, 500, 1000], torch.float16), + (4, 100, 0, [0, 15, 100, 300, 1200, 2000], torch.float16), + (4, 256, 0, [0, 15, 100, 300, 1200, 4096], torch.float16), + ] + ] +) +@pytest.mark.skipif( + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set' +) +def test_chunk_varlen( + H: int, + D: int, + mask_p: float, + cu_seqlens: List[int], + dtype: torch.dtype, +): + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + # randomly split the sequence into N segments + cu_seqlens = torch.LongTensor(cu_seqlens).to(device) + T = cu_seqlens[-1] + N = len(cu_seqlens) - 1 + + # seq-first required for inputs with variable lengths + q = torch.randn((1, T, H, D), dtype=dtype) + k = F.normalize(torch.randn(1, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype) + v = torch.randn((1, T, H, D), dtype=dtype) + g = F.logsigmoid(torch.randn(1, T, H, D, dtype=torch.float)) + g = g * (torch.rand_like(g) > mask_p) + beta = torch.rand(1, T, H, dtype=dtype).sigmoid() + h0 = torch.randn((N, H, D, D), dtype=dtype) + + q, k, v, g, beta, h0 = map(lambda x: x.to(device).requires_grad_(), (q, k, v, g, beta, h0)) + do = torch.randn_like(v) + dht = torch.rand_like(h0) + + tri, tri_ht = chunk_kda( + q=q.clone(), + k=k.clone(), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + initial_state=h0.clone(), + output_final_state=True, + cu_seqlens=cu_seqlens, + ) + ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True) + tri_dq, tri_dk, tri_dv, tri_dg, tri_db, tri_dh0 = q.grad, k.grad, v.grad, g.grad, beta.grad, h0.grad + q.grad = k.grad = v.grad = g.grad = beta.grad = h0.grad = None + + ref = [] + ref_ht = [] + for i in range(N): + ref_i, ref_ht_i = naive_recurrent_kda( + q=q[:, cu_seqlens[i]:cu_seqlens[i+1]], + k=k[:, cu_seqlens[i]:cu_seqlens[i+1]], + v=v[:, cu_seqlens[i]:cu_seqlens[i+1]], + beta=beta[:, cu_seqlens[i]:cu_seqlens[i+1]], + g=g[:, cu_seqlens[i]:cu_seqlens[i+1]], + initial_state=h0[i], + output_final_state=True, + ) + ref.append(ref_i) + ref_ht.append(ref_ht_i) + ref = torch.cat(ref, 1) + ref_ht = torch.cat(ref_ht, 0) + + ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True) + ref_dq, ref_dk, ref_dv, ref_dg, ref_db, ref_dh0 = q.grad, k.grad, v.grad, g.grad, beta.grad, h0.grad + + assert_close('o', ref, tri, 0.005) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.007) + assert_close('dk', ref_dk, tri_dk, 0.008) + assert_close('dv', ref_dv, tri_dv, 0.007) + assert_close('dg', ref_dg, tri_dg, 0.015) + assert_close('db', ref_db, tri_db, 0.015) + assert_close('dh0', ref_dh0, tri_dh0, 0.007) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'use_bias'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-bias{}".format(*test)) + for test in [ + (1, 2, 2, 12, False), + (1, 32, 2, 16, False), + (2, 64, 4, 32, False), + (4, 128, 8, 64, False), + (4, 128, 8, 128, False), + # Add bias tests + (1, 2, 2, 12, True), + (1, 32, 2, 16, True), + (2, 64, 4, 32, True), + (4, 128, 8, 64, True), + (4, 128, 8, 128, True), + ] + ] +) +def test_kda_gate( + B: int, + T: int, + H: int, + D: int, + use_bias: bool, +): + """Test kda gate forward and backward pass - reference vs Triton implementation""" + torch.manual_seed(42) + + g = torch.randn(B, T, H * D, dtype=torch.float32) + # Ensure some values are > 20 to test the threshold logic in softplus + g = g * 30 # Scale up to get values > 20 + A = torch.log(torch.randn(1, 1, H, 1, dtype=torch.float32).uniform_(1, 16)) + g_bias = torch.randn(H * D, dtype=torch.float32) if use_bias else None + + # Move to device and set requires_grad + g, A = map(lambda x: x.to(device).requires_grad_(True), (g, A)) + if g_bias is not None: + g_bias = g_bias.to(device).requires_grad_(True) + + # Create gradient output + do = torch.randn_like(g).view(B, T, H, D) + + # Reference implementation + ref = kda_gate_ref(g.clone(), A.clone(), D, g_bias.clone() if g_bias is not None else None) + # Triton implementation + tri = fused_kda_gate(g.clone(), A.clone(), D, g_bias.clone() if g_bias is not None else None) + + # Backward pass + ((ref * do).sum()).backward(retain_graph=True) + ref_dg, ref_dA = g.grad, A.grad + ref_dgbias = g_bias.grad if g_bias is not None else None + g.grad = A.grad = None + if g_bias is not None: + g_bias.grad = None + + ((tri * do).sum()).backward(retain_graph=True) + tri_dg, tri_dA = g.grad, A.grad + tri_dgbias = g_bias.grad if g_bias is not None else None + g.grad = A.grad = None + if g_bias is not None: + g_bias.grad = None + + assert_close('o', ref, tri, 1e-4) + assert_close('dg', ref_dg, tri_dg, 1e-4) + assert_close('dA', ref_dA, tri_dA, 1e-4) + if use_bias: + assert_close('dgbias', ref_dgbias, tri_dgbias, 1e-4)