diff --git a/fla/__init__.py b/fla/__init__.py index 7812db8ece..c57600e9a0 100644 --- a/fla/__init__.py +++ b/fla/__init__.py @@ -5,40 +5,11 @@ # For a list of all contributors, visit: # https://github.com/fla-org/flash-linear-attention/graphs/contributors -import importlib from pkgutil import extend_path __path__ = extend_path(__path__, __name__) __version__ = "0.5.2" -__all__: list[str] = [] +from fla import modules, ops # noqa: E402 - -def _import_optional_public_module(module_name: str): - try: - return importlib.import_module(module_name) - except ModuleNotFoundError as exc: - missing = exc.name - # The extension package is optional. Treat its absence, or the absence - # of an external runtime dependency, as the extension being unavailable. - if missing == module_name or (missing is not None and missing.split('.', 1)[0] != 'fla'): - return None - raise - - -def _export_public_api(module) -> None: - globals()[module.__name__.rsplit('.', maxsplit=1)[-1]] = module - for name in module.__all__: - if name.endswith('Config'): - continue - globals()[name] = getattr(module, name) - __all__.append(name) - - -_layers = _import_optional_public_module('fla.layers') -_models = _import_optional_public_module('fla.models') -if _layers is not None and _models is not None: - _export_public_api(_layers) - _export_public_api(_models) - -del _import_optional_public_module, _export_public_api, _layers, _models +__all__ = ["modules", "ops"] diff --git a/fla/modules/__init__.py b/fla/modules/__init__.py index 86129dcbe4..d60ae817b4 100644 --- a/fla/modules/__init__.py +++ b/fla/modules/__init__.py @@ -5,48 +5,7 @@ # For a list of all contributors, visit: # https://github.com/fla-org/flash-linear-attention/graphs/contributors -from fla.modules.convolution import ImplicitLongConvolution, LongConvolution, ShortConvolution -from fla.modules.fused_bitlinear import BitLinear, FusedBitLinear -from fla.modules.fused_cross_entropy import FusedCrossEntropyLoss -from fla.modules.fused_kl_div import FusedKLDivLoss -from fla.modules.fused_linear_cross_entropy import FusedLinearCrossEntropyLoss -from fla.modules.fused_norm_gate import ( - FusedLayerNormGated, - FusedLayerNormSwishGate, - FusedLayerNormSwishGateLinear, - FusedRMSNormGated, - FusedRMSNormSwishGate, - FusedRMSNormSwishGateLinear, -) -from fla.modules.l2norm import L2Norm -from fla.modules.layernorm import GroupNorm, GroupNormLinear, LayerNorm, LayerNormLinear, RMSNorm, RMSNormLinear -from fla.modules.mlp import GatedMLP -from fla.modules.rotary import RotaryEmbedding -from fla.modules.token_shift import TokenShift +from fla.modules.conv.short_conv import ShortConvolution +from fla.modules.fused_norm_gate import FusedRMSNormGated -__all__ = [ - 'BitLinear', - 'FusedBitLinear', - 'FusedCrossEntropyLoss', - 'FusedKLDivLoss', - 'FusedLayerNormGated', - 'FusedLayerNormSwishGate', - 'FusedLayerNormSwishGateLinear', - 'FusedLinearCrossEntropyLoss', - 'FusedRMSNormGated', - 'FusedRMSNormSwishGate', - 'FusedRMSNormSwishGateLinear', - 'GatedMLP', - 'GroupNorm', - 'GroupNormLinear', - 'ImplicitLongConvolution', - 'L2Norm', - 'LayerNorm', - 'LayerNormLinear', - 'LongConvolution', - 'RMSNorm', - 'RMSNormLinear', - 'RotaryEmbedding', - 'ShortConvolution', - 'TokenShift', -] +__all__ = ["FusedRMSNormGated", "ShortConvolution"] diff --git a/fla/modules/backends/__init__.py b/fla/modules/backends/__init__.py index 6dca8b39ba..d3ddae502f 100644 --- a/fla/modules/backends/__init__.py +++ b/fla/modules/backends/__init__.py @@ -7,11 +7,8 @@ """Module-level backends for FLA components such as rotary and cross-entropy.""" -from fla.modules.backends.triton_ascend import TritonAscendBackend -from fla.ops.backends import BackendRegistry, dispatch +from fla.ops.backends import dispatch -modules_registry = BackendRegistry("modules") - -modules_registry.register(TritonAscendBackend()) +modules_registry = None __all__ = ['dispatch', 'modules_registry'] diff --git a/fla/modules/conv/causal_conv1d.py b/fla/modules/conv/causal_conv1d.py index 52e087310b..fb0741bfb8 100644 --- a/fla/modules/conv/causal_conv1d.py +++ b/fla/modules/conv/causal_conv1d.py @@ -9,6 +9,8 @@ import torch +from fla.modules.conv.cp import causal_conv1d_cp +from fla.modules.conv.triton import CausalConv1dFunction from fla.ops.cp import FLACPContext from fla.utils import input_guard @@ -65,11 +67,6 @@ def causal_conv1d( Tuple of (output, final_state). If `output_final_state` is `False`, the final state is `None`. """ - # Import here to avoid circular dependencies - from fla.modules.conv.cp import causal_conv1d_cp - from fla.modules.conv.cuda import causal_conv1d_cuda, fast_causal_conv1d_fn - from fla.modules.conv.triton import CausalConv1dFunction - if cp_context is not None: assert initial_state is None, "Initial state is not supported for CP" assert output_final_state is False, "Output final state is not supported for CP" @@ -98,6 +95,8 @@ def causal_conv1d( ) return y, final_state elif backend == 'mix': + from fla.modules.conv.cuda import fast_causal_conv1d_fn + seq_idx = kwargs.get('seq_idx') return fast_causal_conv1d_fn( x, @@ -113,6 +112,8 @@ def causal_conv1d( seq_idx=seq_idx, ) elif backend == 'cuda': + from fla.modules.conv.cuda import causal_conv1d_cuda + return causal_conv1d_cuda( x, weight, diff --git a/fla/modules/conv/short_conv.py b/fla/modules/conv/short_conv.py index 3a4757271e..2509e29189 100644 --- a/fla/modules/conv/short_conv.py +++ b/fla/modules/conv/short_conv.py @@ -7,6 +7,8 @@ """Short convolution implementation for efficient causal convolutions.""" +from __future__ import annotations + import warnings import torch diff --git a/fla/modules/conv/triton/ops.py b/fla/modules/conv/triton/ops.py index 9f32953d6c..4ac21e439f 100644 --- a/fla/modules/conv/triton/ops.py +++ b/fla/modules/conv/triton/ops.py @@ -61,7 +61,7 @@ def causal_conv1d_fwd( NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT) NB = triton.cdiv(B*T, 1024) - y = torch.empty_like(x, memory_format=torch.contiguous_format) + y = torch.empty_like(x) def grid(meta): return (triton.cdiv(D, meta['BD']), NT, B) causal_conv1d_fwd_kernel[grid]( @@ -378,6 +378,12 @@ def forward( chunk_indices: torch.LongTensor | None = None, chunk_size: int = 64, ): + ctx.has_bias = bias is not None + ctx.has_residual = residual is not None + ctx.has_initial_state = initial_state is not None + ctx.has_cu_seqlens = cu_seqlens is not None + ctx.has_cu_seqlens_cpu = cu_seqlens_cpu is not None + ctx.has_chunk_indices = chunk_indices is not None BT = chunk_size if cu_seqlens is not None and chunk_indices is None: chunk_indices = prepare_chunk_indices(cu_seqlens, BT, cu_seqlens_cpu=cu_seqlens_cpu) @@ -421,4 +427,12 @@ def backward(ctx, dy: torch.Tensor, dht: torch.Tensor | None = None): chunk_indices=ctx.chunk_indices, layout_fallback=ctx.layout_fallback, ) - return dx, dw, db, dr, dh0, None, None, None, None, None, None + return ( + (dx, dw) + + ((db,) if ctx.has_bias else ()) + + ((dr,) if ctx.has_residual else ()) + + ((dh0,) if ctx.has_initial_state else ()) + + ((None,) if ctx.has_cu_seqlens else ()) + + ((None,) if ctx.has_cu_seqlens_cpu else ()) + + ((None,) if ctx.has_chunk_indices else ()) + ) diff --git a/fla/modules/fused_norm_gate.py b/fla/modules/fused_norm_gate.py index 58c4deeec9..55bf75f182 100644 --- a/fla/modules/fused_norm_gate.py +++ b/fla/modules/fused_norm_gate.py @@ -656,6 +656,7 @@ def forward( residual_in_fp32: bool = False, is_rms_norm: bool = False, ): + ctx.has_bias = bias is not None x_shape_og = x.shape g_shape_og = g.shape # reshape input data into 2D tensor @@ -716,16 +717,9 @@ def backward(ctx, dy, *args): x_dtype=ctx.x_dtype, ) return ( - dx.reshape(ctx.x_shape_og), - dg.reshape(ctx.g_shape_og), - dw, - db, - None, - dres_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - None, - None, - None, - None, + (dx.reshape(ctx.x_shape_og), dg.reshape(ctx.g_shape_og), dw) + + ((db,) if ctx.has_bias else ()) + + ((dres_in.reshape(ctx.x_shape_og),) if ctx.has_residual else ()) ) diff --git a/fla/ops/__init__.py b/fla/ops/__init__.py index 3f962e5dfd..991f210967 100644 --- a/fla/ops/__init__.py +++ b/fla/ops/__init__.py @@ -5,87 +5,6 @@ # For a list of all contributors, visit: # https://github.com/fla-org/flash-linear-attention/graphs/contributors -from .abc import chunk_abc -from .attn import parallel_attn -from .attnres import fused_attnres -from .based import fused_chunk_based, parallel_based -from .comba import chunk_comba, fused_recurrent_comba -from .delta_rule import chunk_delta_rule, fused_chunk_delta_rule, fused_recurrent_delta_rule -from .forgetting_attn import parallel_forgetting_attn -from .gated_delta_rule import chunk_gated_delta_rule, chunk_gdn, fused_recurrent_gated_delta_rule, fused_recurrent_gdn -from .generalized_delta_rule import ( - chunk_dplr_delta_rule, - chunk_iplr_delta_rule, - fused_recurrent_dplr_delta_rule, - fused_recurrent_iplr_delta_rule, -) -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 -from .mesa_net import chunk_mesa_net -from .nsa import parallel_nsa -from .parallax import parallel_parallax -from .path_attn import parallel_path_attn -from .retention import chunk_retention, fused_chunk_retention, fused_recurrent_retention, parallel_retention -from .rwkv6 import chunk_rwkv6, fused_recurrent_rwkv6 -from .rwkv7 import chunk_rwkv7, fused_recurrent_rwkv7 -from .simple_gla import chunk_simple_gla, fused_chunk_simple_gla, fused_recurrent_simple_gla, parallel_simple_gla -from .wall_attn import parallel_wall_attn, parallel_wall_attn_decode +from fla.ops import cp, kda, utils -__all__ = [ - 'chunk_abc', - 'chunk_comba', - 'chunk_delta_rule', - 'chunk_dplr_delta_rule', - 'chunk_gated_delta_rule', - 'chunk_gdn', - 'chunk_gla', - 'chunk_gsa', - 'chunk_iplr_delta_rule', - 'chunk_kda', - 'chunk_lightning_attn', - 'chunk_linear_attn', - 'chunk_log_linear_attn', - 'chunk_mesa_net', - 'chunk_retention', - 'chunk_rwkv6', - 'chunk_rwkv7', - 'chunk_simple_gla', - 'fused_attnres', - 'fused_chunk_based', - 'fused_chunk_delta_rule', - 'fused_chunk_gla', - 'fused_chunk_linear_attn', - 'fused_chunk_retention', - 'fused_chunk_simple_gla', - 'fused_recurrent_comba', - 'fused_recurrent_delta_rule', - 'fused_recurrent_dplr_delta_rule', - 'fused_recurrent_gated_delta_rule', - 'fused_recurrent_gdn', - 'fused_recurrent_gla', - 'fused_recurrent_gsa', - 'fused_recurrent_hgrn', - 'fused_recurrent_iplr_delta_rule', - 'fused_recurrent_kda', - 'fused_recurrent_lightning_attn', - 'fused_recurrent_linear_attn', - 'fused_recurrent_retention', - 'fused_recurrent_rwkv6', - 'fused_recurrent_rwkv7', - 'fused_recurrent_simple_gla', - 'parallel_attn', - 'parallel_based', - 'parallel_forgetting_attn', - 'parallel_nsa', - 'parallel_parallax', - 'parallel_path_attn', - 'parallel_retention', - 'parallel_simple_gla', - 'parallel_wall_attn', - 'parallel_wall_attn_decode', -] +__all__ = ["cp", "kda", "utils"] diff --git a/fla/ops/backends/__init__.py b/fla/ops/backends/__init__.py index 4fd2881bef..fa07795184 100644 --- a/fla/ops/backends/__init__.py +++ b/fla/ops/backends/__init__.py @@ -164,6 +164,8 @@ def dispatch(operation: str): that passes the verifier for the given function call. """ def decorator(func: F) -> F: + return func + if _DISPATCH_DISABLED: return func func_name = func.__name__ diff --git a/fla/ops/common/gate.py b/fla/ops/common/gate.py index 06c8486686..b76c8bf4b7 100644 --- a/fla/ops/common/gate.py +++ b/fla/ops/common/gate.py @@ -57,7 +57,7 @@ def fused_beta_sigmoid_bwd_kernel( @dispatch('common') def fused_beta_sigmoid_fwd(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: y = torch.empty_like(x, dtype=torch.float32) - n_elements = x.numel() + n_elements = x.shape.numel() grid = (triton.cdiv(n_elements, _BETA_SIGMOID_BLOCK_SIZE),) fused_beta_sigmoid_fwd_kernel[grid]( x, @@ -73,7 +73,7 @@ def fused_beta_sigmoid_fwd(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: @dispatch('common') def fused_beta_sigmoid_bwd(x: torch.Tensor, dy: torch.Tensor, scale: float = 1.0) -> torch.Tensor: dx = torch.empty_like(x) - n_elements = x.numel() + n_elements = x.shape.numel() grid = (triton.cdiv(n_elements, _BETA_SIGMOID_BLOCK_SIZE),) fused_beta_sigmoid_bwd_kernel[grid]( x, @@ -103,7 +103,7 @@ def forward(ctx, x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: def backward(ctx, dy: torch.Tensor): (x,) = ctx.saved_tensors dx = fused_beta_sigmoid_bwd(x, dy, ctx.scale) - return dx.type_as(x), None + return (dx.type_as(x),) def fused_beta_sigmoid(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: diff --git a/fla/ops/cp/comm.py b/fla/ops/cp/comm.py index 89a1cac000..46c584122a 100644 --- a/fla/ops/cp/comm.py +++ b/fla/ops/cp/comm.py @@ -34,11 +34,16 @@ def all_gather_into_tensor( Returns: Tuple of (output tensor, handle if async_op else None) """ - world_size = dist.get_world_size(group=group) + if async_op: + raise NotImplementedError("KDA context parallel currently supports synchronous all-gather only") + gathered = [] + dist.all_gather(gathered, inp, group=group, sync_op=True) + gathered_tensor = torch.stack(gathered, dim=0) if out is None: - out = torch.empty(world_size, *inp.shape, device=inp.device, dtype=inp.dtype) - handle = dist.all_gather_into_tensor(out, inp, group=group, async_op=async_op) - return out, handle + out = gathered_tensor + else: + out.copy_(gathered_tensor) + return out, None def all_reduce_sum( @@ -57,7 +62,7 @@ def all_reduce_sum( Returns: Tuple of (reduced tensor, handle if async_op else None) """ - handle = dist.all_reduce(inp, op=dist.ReduceOp.SUM, group=group, async_op=async_op) + handle = dist.all_reduce(inp, op=dist.ReduceOp.SUM, group=group, sync_op=not async_op) return inp, handle diff --git a/fla/ops/cp/context.py b/fla/ops/cp/context.py index 0a0b8113ac..77246cc825 100644 --- a/fla/ops/cp/context.py +++ b/fla/ops/cp/context.py @@ -10,6 +10,7 @@ from dataclasses import dataclass from typing import TYPE_CHECKING +import paddle import torch import torch.distributed as dist @@ -90,11 +91,14 @@ def get_cp_cu_seqlens( # Optimization: cu_seqlens is sorted, use searchsorted to quickly locate boundaries # Find first sequence whose end > rank_start # cu_seqlens_cpu[1:] contains all sequence end points - start_seq_idx = torch.searchsorted(cu_seqlens_cpu[1:], rank_start, side='right') + # TODO: avoid scalar tensor allocations once paddle.searchsorted accepts Python scalar values. + rank_start_tensor = paddle.to_tensor(rank_start, dtype=cu_seqlens_cpu.dtype, place=cu_seqlens_cpu.place) + rank_end_tensor = paddle.to_tensor(rank_end, dtype=cu_seqlens_cpu.dtype, place=cu_seqlens_cpu.place) + start_seq_idx = paddle.searchsorted(cu_seqlens_cpu[1:], rank_start_tensor, side='right').item() # Find first sequence whose start >= rank_end, sequences before this may overlap # cu_seqlens_cpu[:-1] contains all sequence start points - end_seq_idx = torch.searchsorted(cu_seqlens_cpu[:-1], rank_end, side='left') + end_seq_idx = paddle.searchsorted(cu_seqlens_cpu[:-1], rank_end_tensor, side='left').item() # Slice cu_seqlens_cpu[start_seq_idx : end_seq_idx + 1] to get relevant global cu_seqlens nodes # +1 because end_seq_idx is an open boundary, and cu_seqlens length is num_seqs + 1 diff --git a/fla/ops/gla/chunk.py b/fla/ops/gla/chunk.py index f1a30a2738..484d898dc5 100644 --- a/fla/ops/gla/chunk.py +++ b/fla/ops/gla/chunk.py @@ -1416,7 +1416,6 @@ def backward(ctx, do, dht): return dq.to(q), dk.to(k), dv.to(v), dg, None, dh0, None, None, None, None -@torch.compiler.disable def chunk_gla( q: torch.Tensor, k: torch.Tensor, diff --git a/fla/ops/kda/__init__.py b/fla/ops/kda/__init__.py index f896d2b9bc..4bf1f66ded 100644 --- a/fla/ops/kda/__init__.py +++ b/fla/ops/kda/__init__.py @@ -6,9 +6,5 @@ # https://github.com/fla-org/flash-linear-attention/graphs/contributors from .chunk import chunk_kda -from .fused_recurrent import fused_recurrent_kda -__all__ = [ - "chunk_kda", - "fused_recurrent_kda", -] +__all__ = ["chunk_kda"] diff --git a/fla/ops/kda/chunk.py b/fla/ops/kda/chunk.py index 39d2b4ffb2..c67a356770 100644 --- a/fla/ops/kda/chunk.py +++ b/fla/ops/kda/chunk.py @@ -9,6 +9,7 @@ import warnings +import paddle import torch from fla.modules.l2norm import l2norm_bwd, l2norm_fwd @@ -51,6 +52,11 @@ def forward( return_intermediate_states: bool = False, cp_context: FLACPContext | None = None, ): + ctx.has_A_log = A_log is not None + ctx.has_dt_bias = dt_bias is not None + ctx.has_initial_state = initial_state is not None + ctx.has_cu_seqlens = cu_seqlens is not None + ctx.has_cu_seqlens_cpu = cu_seqlens_cpu is not None # Apply l2norm q_rstd, k_rstd = None, None if use_qk_l2norm_in_kernel: @@ -96,7 +102,7 @@ def forward( ) if return_intermediate_states: - assert torch.is_inference_mode_enabled(), "return_intermediate_states is only allowed in inference mode" + assert not paddle.is_grad_enabled(), "return_intermediate_states is only allowed in inference mode" assert disable_recompute is False, "return_intermediate_states must be used with disable_recompute=False" return o.type_as(q), final_state, h @@ -124,7 +130,7 @@ def forward( def backward( ctx, do: torch.Tensor, - dht: torch.Tensor, + dht: torch.Tensor | None = None, ): (q, q_rstd, k, k_rstd, v, g_cumsum, g_input, beta_raw, beta, A_log, dt_bias, Aqk, Akk, w, u, qg, kg, v_new, h, @@ -169,12 +175,17 @@ def backward( if ctx.use_beta_sigmoid_in_kernel: db = fused_beta_sigmoid_bwd(beta_raw, db, scale=2.0 if ctx.allow_neg_eigval else 1.0) - return (dq.to(q), dk.to(k), dv.to(v), dg.to(g_input), db.to(beta_raw), dA, dbias, None, dh0, - None, None, None, None, None, None, None, None, None, None, None, None, None, None) + return ( + (dq.to(q), dk.to(k), dv.to(v), dg.to(g_input), db.to(beta_raw)) + + ((dA,) if ctx.has_A_log else ()) + + ((dbias,) if ctx.has_dt_bias else ()) + + ((dh0,) if ctx.has_initial_state else ()) + + ((None,) if ctx.has_cu_seqlens else ()) + + ((None,) if ctx.has_cu_seqlens_cpu else ()) + ) @dispatch('kda') -@torch.compiler.disable def chunk_kda( q: torch.Tensor, k: torch.Tensor, diff --git a/fla/ops/kda/gate.py b/fla/ops/kda/gate.py index ac8c54121f..451a48c363 100644 --- a/fla/ops/kda/gate.py +++ b/fla/ops/kda/gate.py @@ -225,7 +225,7 @@ def kda_gate_fwd( output_dtype: torch.dtype = torch.float32, ) -> torch.Tensor: H, K = g.shape[-2:] - T = g.numel() // (H * K) + T = g.shape.numel() // (H * K) yg = torch.empty_like(g, dtype=output_dtype) @@ -257,7 +257,7 @@ def kda_gate_bwd( lower_bound: float | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: H, K = g.shape[-2:] - T = g.numel() // (H * K) + T = g.shape.numel() // (H * K) BT = 32 NT = triton.cdiv(T, BT) @@ -302,6 +302,7 @@ def forward( lower_bound: float | None = None, output_dtype: torch.dtype = torch.float32, ) -> torch.Tensor: + ctx.has_dt_bias = dt_bias is not None yg = kda_gate_fwd( g=g, A_log=A_log, @@ -325,11 +326,10 @@ def backward(ctx, dyg: torch.Tensor): dyg=dyg, lower_bound=ctx.lower_bound ) - return dg, dA, dbias, None, None + return (dg, dA) + ((dbias,) if ctx.has_dt_bias else ()) @dispatch('kda') -@torch.compiler.disable def fused_kda_gate( g: torch.Tensor, A_log: torch.Tensor, diff --git a/fla/ops/utils/__init__.py b/fla/ops/utils/__init__.py index 88acd8b9cb..f5efd5ad99 100644 --- a/fla/ops/utils/__init__.py +++ b/fla/ops/utils/__init__.py @@ -5,61 +5,11 @@ # For a list of all contributors, visit: # https://github.com/fla-org/flash-linear-attention/graphs/contributors -from .csr import prepare_block_csr -from .cumsum import ( - chunk_global_cumsum, - chunk_global_cumsum_scalar, - chunk_global_cumsum_vector, - chunk_local_cumsum, - chunk_local_cumsum_scalar, - chunk_local_cumsum_vector, -) -from .index import ( - get_max_num_splits, - prepare_chunk_indices, - prepare_chunk_offsets, - prepare_cu_seqlens_from_lens, - prepare_cu_seqlens_from_mask, - prepare_lens, - prepare_lens_from_mask, - prepare_position_ids, - prepare_sequence_ids, - prepare_token_indices, -) -from .logsumexp import logsumexp_fwd -from .matmul import addmm, matmul -from .pack import pack_sequence, unpack_sequence -from .pooling import mean_pooling -from .softmax import softmax_bwd, softmax_fwd -from .softplus import softplus -from .solve_tril import solve_tril +from .cumsum import chunk_local_cumsum +from .index import prepare_chunk_indices, prepare_chunk_offsets __all__ = [ - "addmm", - "chunk_global_cumsum", - "chunk_global_cumsum_scalar", - "chunk_global_cumsum_vector", "chunk_local_cumsum", - "chunk_local_cumsum_scalar", - "chunk_local_cumsum_vector", - "get_max_num_splits", - "logsumexp_fwd", - "matmul", - "mean_pooling", - "pack_sequence", - "prepare_block_csr", "prepare_chunk_indices", "prepare_chunk_offsets", - "prepare_cu_seqlens_from_lens", - "prepare_cu_seqlens_from_mask", - "prepare_lens", - "prepare_lens_from_mask", - "prepare_position_ids", - "prepare_sequence_ids", - "prepare_token_indices", - "softmax_bwd", - "softmax_fwd", - "softplus", - "solve_tril", - "unpack_sequence", ] diff --git a/fla/ops/utils/index.py b/fla/ops/utils/index.py index 843bf25aa9..7b5597378d 100644 --- a/fla/ops/utils/index.py +++ b/fla/ops/utils/index.py @@ -5,6 +5,8 @@ # For a list of all contributors, visit: # https://github.com/fla-org/flash-linear-attention/graphs/contributors +from __future__ import annotations + import torch import torch.nn.functional as F import triton @@ -126,10 +128,10 @@ def _segmented_arange(counts: torch.LongTensor) -> tuple[torch.LongTensor, torch host (one device sync). Pass host-side counts to avoid it. """ seg_id = torch.repeat_interleave( - torch.arange(counts.numel(), device=counts.device, dtype=counts.dtype), + torch.arange(counts.shape.numel(), device=counts.device, dtype=counts.dtype), counts, - ) - seg_start = F.pad(counts.cumsum(0), (1, 0))[:-1] + ).to(counts.dtype) + seg_start = F.pad(counts.cumsum(0), (1, 0))[:-1].to(counts.dtype) intra_idx = torch.arange(seg_id.shape[0], device=counts.device, dtype=counts.dtype) - seg_start[seg_id] return seg_id, intra_idx @@ -159,7 +161,7 @@ def prepare_chunk_indices( cu_seqlens_cpu: torch.LongTensor | None = None, ) -> torch.LongTensor: src = cu_seqlens_cpu if cu_seqlens_cpu is not None else cu_seqlens - chunk_counts = (prepare_lens(src) + (chunk_size - 1)).div(chunk_size, rounding_mode='floor') + chunk_counts = (prepare_lens(src) + (chunk_size - 1)) // chunk_size seg_id, intra_chunk_idx = _segmented_arange(chunk_counts) return torch.stack([seg_id, intra_chunk_idx], 1).to(cu_seqlens) diff --git a/fla/utils/_device.py b/fla/utils/_device.py index 80796edf83..dad319ba8b 100644 --- a/fla/utils/_device.py +++ b/fla/utils/_device.py @@ -6,7 +6,6 @@ # https://github.com/fla-org/flash-linear-attention/graphs/contributors import contextlib -import functools import logging import os import platform @@ -202,29 +201,18 @@ def check_shared_mem(arch: str = "none", tensor_idx: int = 0) -> bool: return False -if check_pytorch_version('2.4'): - if device == 'cpu': - device = 'cuda' - device_torch_lib = getattr(torch, device) - autocast_custom_fwd = functools.partial(torch.amp.custom_fwd, device_type=device) - autocast_custom_bwd = functools.partial(torch.amp.custom_bwd, device_type=device) +def autocast_custom_fwd(fn): + return fn - def custom_device_ctx(index: int): - if index is None: - return contextlib.nullcontext() - try: - return device_torch_lib.device(index) - except (AttributeError, AssertionError, RuntimeError): - return contextlib.nullcontext() -else: - assert device == 'cuda', 'Only cuda device is supported for PyTorch version < 2.4.0.' - autocast_custom_fwd = device_torch_lib.amp.custom_fwd - autocast_custom_bwd = device_torch_lib.amp.custom_bwd - - def custom_device_ctx(index: int): - if index is None: - return contextlib.nullcontext() - try: - return torch.cuda.device(index) - except (AttributeError, AssertionError, RuntimeError): - return contextlib.nullcontext() + +def autocast_custom_bwd(fn): + return fn + + +def custom_device_ctx(index: int): + if index is None: + return contextlib.nullcontext() + try: + return device_torch_lib.device(index) + except (AttributeError, AssertionError, RuntimeError): + return contextlib.nullcontext() diff --git a/tests/paddle/kda_eager_reference.py b/tests/paddle/kda_eager_reference.py new file mode 100644 index 0000000000..bd4b3a12ee --- /dev/null +++ b/tests/paddle/kda_eager_reference.py @@ -0,0 +1,114 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +from __future__ import annotations + +from collections.abc import Mapping + +import numpy as np +import paddle + + +def make_leaf(array: np.ndarray, dtype: str, stop_gradient: bool = False) -> paddle.Tensor: + tensor = paddle.to_tensor(array).astype(dtype) + tensor.stop_gradient = stop_gradient + return tensor + + +def clone_inputs(arrays: Mapping[str, np.ndarray], dtypes: Mapping[str, str]) -> dict[str, paddle.Tensor]: + return {name: make_leaf(array, dtypes[name]) for name, array in arrays.items()} + + +def normalize(x: paddle.Tensor) -> paddle.Tensor: + x = x.astype("float32") + return x * paddle.rsqrt((x.square()).sum(axis=-1, keepdim=True) + 1e-6) + + +def apply_gate( + g: paddle.Tensor, + A_log: paddle.Tensor, + dt_bias: paddle.Tensor | None, + lower_bound: float | None, +) -> paddle.Tensor: + H, _ = g.shape[-2:] + g = g.astype("float32") + if dt_bias is not None: + g = g + dt_bias.reshape([H, -1]) + rate = paddle.exp(A_log.reshape([H, 1]).astype("float32")) + if lower_bound is None: + return -rate * paddle.nn.functional.softplus(g) + return lower_bound * paddle.nn.functional.sigmoid(rate * g) + + +def recurrent_kda( + q: paddle.Tensor, + k: paddle.Tensor, + v: paddle.Tensor, + g: paddle.Tensor, + beta: paddle.Tensor, + scale: float | None = None, + initial_state: paddle.Tensor | None = None, + output_final_state: bool = False, +) -> tuple[paddle.Tensor, paddle.Tensor | None]: + dtype = v.dtype + B, T, H, K = q.shape + HV, V = v.shape[2:] + value_heads_per_qk_head = HV // H + scale = K ** -0.5 if scale is None else scale + + q = paddle.repeat_interleave(q.astype("float32"), value_heads_per_qk_head, axis=2) * scale + k = paddle.repeat_interleave(k.astype("float32"), value_heads_per_qk_head, axis=2) + v, g, beta = (value.astype("float32") for value in (v, g, beta)) + + state = paddle.zeros([B, HV, K, V], dtype="float32") + if initial_state is not None: + state = state + initial_state.astype("float32") + + outputs = [] + for token_idx in range(T): + q_i, k_i, v_i, g_i, beta_i = ( + q[:, token_idx], + k[:, token_idx], + v[:, token_idx], + g[:, token_idx], + beta[:, token_idx], + ) + state = state * paddle.exp(g_i).unsqueeze(-1) + prediction = (k_i.unsqueeze(-1) * state).sum(axis=-2) + update = beta_i.unsqueeze(-1).unsqueeze(-1) * k_i.unsqueeze(-1) * (v_i - prediction).unsqueeze(-2) + state = state + update + outputs.append((q_i.unsqueeze(-1) * state).sum(axis=-2)) + + output = paddle.stack(outputs, axis=1).astype(dtype) + return output, state if output_final_state else None + + +def tensor_error(ref: paddle.Tensor, actual: paddle.Tensor) -> tuple[float, float]: + ref = ref.detach().astype("float32") + actual = actual.detach().astype("float32") + diff = ref - actual + max_abs = diff.abs().max().item() + rms_error = diff.square().mean().sqrt().item() + rms_ref = ref.square().mean().sqrt().item() + return max_abs, rms_error / (rms_ref + 1e-8) + + +def assert_close( + name: str, + ref: paddle.Tensor, + actual: paddle.Tensor, + ratio: float, + abs_atol: float = 1e-6, +) -> None: + assert list(ref.shape) == list(actual.shape), f"{name}: shape {list(actual.shape)} != {list(ref.shape)}" + assert ref.dtype == actual.dtype, f"{name}: dtype {actual.dtype} != {ref.dtype}" + assert bool(paddle.isfinite(ref).all().item()), f"{name}: non-finite reference" + assert bool(paddle.isfinite(actual).all().item()), f"{name}: non-finite result" + max_abs, error_ratio = tensor_error(ref, actual) + if max_abs <= abs_atol: + return + assert error_ratio < ratio, f"{name}: max_abs={max_abs:.6f}, error_ratio={error_ratio:.6f}, tolerance={ratio:.6f}" diff --git a/tests/paddle/test_cp_kda.py b/tests/paddle/test_cp_kda.py new file mode 100644 index 0000000000..18e5f22298 --- /dev/null +++ b/tests/paddle/test_cp_kda.py @@ -0,0 +1,144 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +import os + +import numpy as np +import paddle +import paddle.distributed as dist +import pytest +from kda_eager_reference import apply_gate, assert_close, clone_inputs, normalize, recurrent_kda + + +def _cp_worker(world_size: int) -> None: + paddle.enable_compat(scope={"fla", "triton"}, silent=True) + dist.init_parallel_env() + rank = dist.get_rank() + group = dist.new_group(ranks=list(range(world_size))) + + from fla.ops.cp import build_cp_context + from fla.ops.kda import chunk_kda + + try: + lengths = [20, 44] + B, T, H, HV, K, V = 1, sum(lengths), 1, 1, 64, 32 + local_tokens = T // world_size + local_start = rank * local_tokens + local_end = local_start + local_tokens + rng = np.random.default_rng(20260729) + arrays = { + "q": rng.standard_normal([B, T, H, K], dtype=np.float32), + "k": rng.standard_normal([B, T, H, K], dtype=np.float32), + "v": rng.standard_normal([B, T, HV, V], dtype=np.float32), + "g": rng.standard_normal([B, T, HV, K], dtype=np.float32), + "beta": rng.standard_normal([B, T, HV], dtype=np.float32), + "A_log": rng.uniform(-0.5, 0.5, [HV]).astype(np.float32), + "dt_bias": rng.uniform(-1.0, 1.0, [HV * K]).astype(np.float32), + "do": rng.standard_normal([B, T, HV, V], dtype=np.float32), + } + dtypes = { + "q": "bfloat16", + "k": "bfloat16", + "v": "bfloat16", + "g": "bfloat16", + "beta": "bfloat16", + "A_log": "float32", + "dt_bias": "float32", + "do": "float32", + } + local_arrays = { + name: ( + array[:, local_start:local_end] + if name in {"q", "k", "v", "g", "beta", "do"} + else array + ) + for name, array in arrays.items() + } + tri_inputs = clone_inputs(local_arrays, dtypes) + ref_inputs = clone_inputs(arrays, dtypes) + + cu_values = np.cumsum([0, *lengths]).astype(np.int32) + cu_seqlens = paddle.to_tensor(cu_values) + cu_seqlens_cpu = paddle.to_tensor(cu_values, place=paddle.CPUPlace()) + cp_context = build_cp_context(cu_seqlens, group=group, cu_seqlens_cpu=cu_seqlens_cpu) + tri, final_state = chunk_kda( + q=tri_inputs["q"], + k=tri_inputs["k"], + v=tri_inputs["v"], + g=tri_inputs["g"], + beta=tri_inputs["beta"], + A_log=tri_inputs["A_log"], + dt_bias=tri_inputs["dt_bias"], + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + safe_gate=True, + lower_bound=-5.0, + cp_context=cp_context, + ) + assert final_state is None + (tri.astype("float32") * tri_inputs["do"]).sum().backward() + + ref_gate = apply_gate(ref_inputs["g"], ref_inputs["A_log"], ref_inputs["dt_bias"], lower_bound=-5.0) + ref_beta = paddle.nn.functional.sigmoid(ref_inputs["beta"]) + ref_outputs = [] + start = 0 + for length in lengths: + end = start + length + ref_output, _ = recurrent_kda( + q=normalize(ref_inputs["q"][:, start:end]), + k=normalize(ref_inputs["k"][:, start:end]), + v=ref_inputs["v"][:, start:end], + g=ref_gate[:, start:end], + beta=ref_beta[:, start:end], + ) + ref_outputs.append(ref_output) + start = end + ref = paddle.concat(ref_outputs, axis=1) + (ref.astype("float32") * ref_inputs["do"]).sum().backward() + + assert_close("o", ref[:, local_start:local_end], tri, 0.008) + for name, tolerance in { + "q": 0.01, + "k": 0.01, + "v": 0.01, + "g": 0.025, + "beta": 0.025, + }.items(): + assert_close( + f"d{name}", + ref_inputs[name].grad[:, local_start:local_end], + tri_inputs[name].grad, + tolerance, + ) + + for name, tolerance, abs_atol in (("A_log", 0.025, 1e-6), ("dt_bias", 0.01, 1e-3)): + dist.all_reduce(tri_inputs[name].grad, group=group) + assert_close(f"d{name}", ref_inputs[name].grad, tri_inputs[name].grad, tolerance, abs_atol=abs_atol) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif(paddle.device.cuda.device_count() < 2, reason="Paddle KDA CP requires at least two GPUs") +def test_chunk_kda_context_parallel(monkeypatch: pytest.MonkeyPatch): + cloud_prefixes = ( + "PADDLE_CLUSTER_", + "PADDLE_CURRENT_ENDPOINT", + "PADDLE_IS_LOCAL", + "PADDLE_NUM_GRADIENT_SERVERS", + "PADDLE_TRAINER", + "PADDLE_TRAINERS", + "PADDLE_TRAINING_ROLE", + "PADDLE_WORKERS_IP_PORT_LIST", + "POD_", + "TRAINER_", + "TRAINERS", + ) + for name in tuple(os.environ): + if name.startswith(cloud_prefixes): + monkeypatch.delenv(name) + dist.spawn(_cp_worker, args=(2,), nprocs=2) diff --git a/tests/paddle/test_kda.py b/tests/paddle/test_kda.py new file mode 100644 index 0000000000..e8617b861f --- /dev/null +++ b/tests/paddle/test_kda.py @@ -0,0 +1,277 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +import sys + +import numpy as np +import paddle +import pytest + +paddle.enable_compat(scope={"fla", "triton"}, silent=True) + +from kda_eager_reference import apply_gate, assert_close, clone_inputs, normalize, recurrent_kda # noqa: E402 + +from fla.ops.kda import chunk_fwd as kda_chunk_fwd # noqa: E402 +from fla.ops.kda import chunk_kda # noqa: E402 +from fla.ops.kda.gate import fused_kda_gate # noqa: E402 + + +def _random_arrays( + *, + B: int, + T: int, + H: int, + HV: int, + K: int, + V: int, + seed: int, +) -> dict[str, np.ndarray]: + rng = np.random.default_rng(seed) + return { + "q": rng.standard_normal([B, T, H, K], dtype=np.float32), + "k": rng.standard_normal([B, T, H, K], dtype=np.float32), + "v": rng.standard_normal([B, T, HV, V], dtype=np.float32), + "g": rng.standard_normal([B, T, HV, K], dtype=np.float32), + "beta": rng.standard_normal([B, T, HV], dtype=np.float32), + "A_log": rng.uniform(-0.5, 0.5, [HV]).astype(np.float32), + "dt_bias": rng.uniform(-1.0, 1.0, [HV * K]).astype(np.float32), + "h0": rng.standard_normal([B, HV, K, V], dtype=np.float32), + "do": rng.standard_normal([B, T, HV, V], dtype=np.float32), + "dht": rng.standard_normal([B, HV, K, V], dtype=np.float32), + } + + +def _run_dense( + inputs: dict[str, paddle.Tensor], + *, + use_gate_in_kernel: bool, + use_beta_sigmoid_in_kernel: bool, + safe_gate: bool, + disable_recompute: bool, +) -> tuple[paddle.Tensor, paddle.Tensor]: + lower_bound = -5.0 if safe_gate else None + output, final_state = chunk_kda( + q=inputs["q"], + k=inputs["k"], + v=inputs["v"], + g=inputs["g"], + beta=inputs["beta"], + A_log=inputs["A_log"] if use_gate_in_kernel else None, + dt_bias=inputs["dt_bias"] if use_gate_in_kernel else None, + initial_state=inputs["h0"], + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=use_gate_in_kernel, + use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel, + safe_gate=safe_gate, + lower_bound=lower_bound, + disable_recompute=disable_recompute, + ) + loss = ( + output.astype("float32") * inputs["do"] + ).sum() + ( + final_state * inputs["dht"] + ).sum() + loss.backward() + return output, final_state + + +def _run_dense_reference( + inputs: dict[str, paddle.Tensor], + *, + use_gate_in_kernel: bool, + use_beta_sigmoid_in_kernel: bool, + safe_gate: bool, +) -> tuple[paddle.Tensor, paddle.Tensor]: + lower_bound = -5.0 if safe_gate else None + gate = ( + apply_gate(inputs["g"], inputs["A_log"], inputs["dt_bias"], lower_bound) + if use_gate_in_kernel + else inputs["g"] + ) + beta = paddle.nn.functional.sigmoid(inputs["beta"]) if use_beta_sigmoid_in_kernel else inputs["beta"] + output, final_state = recurrent_kda( + q=normalize(inputs["q"]), + k=normalize(inputs["k"]), + v=inputs["v"], + g=gate, + beta=beta, + initial_state=inputs["h0"], + output_final_state=True, + ) + loss = ( + output.astype("float32") * inputs["do"] + ).sum() + ( + final_state * inputs["dht"] + ).sum() + loss.backward() + return output, final_state + + +def test_kda_dependency_closure(): + loaded = {name for name in sys.modules if name == "fla" or name.startswith("fla.")} + allowed_op_roots = {"backends", "common", "cp", "gla", "kda", "utils"} + loaded_op_roots = { + name.split(".")[2] + for name in loaded + if name.startswith("fla.ops.") and len(name.split(".")) > 2 + } + + assert "fla._paddle" not in loaded + assert "fla._framework" not in loaded + assert not any(name.startswith("fla.layers") for name in loaded) + assert not any(name.startswith("fla.models") for name in loaded) + assert "fla.ops.gla.chunk" in loaded + assert kda_chunk_fwd.chunk_gla_fwd_o_gk.__module__ == "fla.ops.gla.chunk" + assert loaded_op_roots <= allowed_op_roots + + +@pytest.mark.parametrize("disable_recompute", [False, True]) +def test_chunk_kda_dense_forward_backward(disable_recompute: bool): + arrays = _random_arrays(B=1, T=64, H=1, HV=2, K=64, V=48, seed=42) + dtypes = { + "q": "bfloat16", + "k": "bfloat16", + "v": "bfloat16", + "g": "bfloat16", + "beta": "bfloat16", + "A_log": "float32", + "dt_bias": "float32", + "h0": "float32", + "do": "float32", + "dht": "float32", + } + tri_inputs = clone_inputs(arrays, dtypes) + ref_inputs = clone_inputs(arrays, dtypes) + + tri, tri_ht = _run_dense( + tri_inputs, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + safe_gate=True, + disable_recompute=disable_recompute, + ) + ref, ref_ht = _run_dense_reference( + ref_inputs, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + safe_gate=True, + ) + + assert_close("o", ref, tri, 0.005) + assert_close("ht", ref_ht, tri_ht, 0.005) + for name, tolerance in { + "q": 0.008, + "k": 0.008, + "v": 0.008, + "g": 0.02, + "beta": 0.02, + "A_log": 0.02, + "dt_bias": 0.008, + "h0": 0.008, + }.items(): + assert_close(f"d{name}", ref_inputs[name].grad, tri_inputs[name].grad, tolerance) + + +@pytest.mark.parametrize("has_bias", [False, True]) +def test_fused_kda_gate_backward(has_bias: bool): + arrays = _random_arrays(B=1, T=37, H=2, HV=2, K=64, V=32, seed=123) + dtypes = { + "g": "float32", + "A_log": "float32", + "dt_bias": "float32", + } + tri_inputs = clone_inputs({name: arrays[name] for name in dtypes}, dtypes) + ref_inputs = clone_inputs({name: arrays[name] for name in dtypes}, dtypes) + tri_bias = tri_inputs["dt_bias"] if has_bias else None + ref_bias = ref_inputs["dt_bias"] if has_bias else None + + tri = fused_kda_gate(tri_inputs["g"], tri_inputs["A_log"], tri_bias, lower_bound=-5.0) + ref = apply_gate(ref_inputs["g"], ref_inputs["A_log"], ref_bias, lower_bound=-5.0) + dy = paddle.to_tensor(arrays["g"]) + (tri * dy).sum().backward() + (ref * dy).sum().backward() + + assert_close("gate", ref, tri, 1e-4) + assert_close("dg", ref_inputs["g"].grad, tri_inputs["g"].grad, 1e-4) + assert_close("dA", ref_inputs["A_log"].grad, tri_inputs["A_log"].grad, 1e-4) + if has_bias: + assert_close("dbias", ref_inputs["dt_bias"].grad, tri_inputs["dt_bias"].grad, 1e-4) + + +def test_chunk_kda_varlen_backward(): + lengths = [13, 19, 32] + B, T, H, HV, K, V = 1, sum(lengths), 1, 1, 64, 40 + arrays = _random_arrays(B=B, T=T, H=H, HV=HV, K=K, V=V, seed=2026) + arrays["g"] = np.minimum(arrays["g"], -0.01).astype(np.float32) + arrays["beta"] = (1.0 / (1.0 + np.exp(-arrays["beta"]))).astype(np.float32) + arrays["h0"] = np.random.default_rng(7).standard_normal([len(lengths), HV, K, V], dtype=np.float32) + arrays["dht"] = np.random.default_rng(8).standard_normal([len(lengths), HV, K, V], dtype=np.float32) + dtypes = { + "q": "float16", + "k": "float16", + "v": "float16", + "g": "float32", + "beta": "float16", + "A_log": "float32", + "dt_bias": "float32", + "h0": "float32", + "do": "float32", + "dht": "float32", + } + tri_inputs = clone_inputs(arrays, dtypes) + ref_inputs = clone_inputs(arrays, dtypes) + cu_values = np.cumsum([0, *lengths]).astype(np.int32) + cu_seqlens = paddle.to_tensor(cu_values) + cu_seqlens_cpu = paddle.to_tensor(cu_values, place=paddle.CPUPlace()) + + tri, tri_ht = chunk_kda( + q=normalize(tri_inputs["q"]).astype(tri_inputs["q"].dtype), + k=normalize(tri_inputs["k"]).astype(tri_inputs["k"].dtype), + v=tri_inputs["v"], + g=tri_inputs["g"], + beta=tri_inputs["beta"], + initial_state=tri_inputs["h0"], + output_final_state=True, + cu_seqlens=cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, + disable_recompute=True, + ) + ((tri.astype("float32") * tri_inputs["do"]).sum() + (tri_ht * tri_inputs["dht"]).sum()).backward() + + ref_outputs = [] + ref_states = [] + start = 0 + for sequence_idx, length in enumerate(lengths): + end = start + length + ref_output, ref_state = recurrent_kda( + q=normalize(ref_inputs["q"][:, start:end]), + k=normalize(ref_inputs["k"][:, start:end]), + v=ref_inputs["v"][:, start:end], + g=ref_inputs["g"][:, start:end], + beta=ref_inputs["beta"][:, start:end], + initial_state=ref_inputs["h0"][sequence_idx:sequence_idx + 1], + output_final_state=True, + ) + ref_outputs.append(ref_output) + ref_states.append(ref_state) + start = end + ref = paddle.concat(ref_outputs, axis=1) + ref_ht = paddle.concat(ref_states, axis=0) + ((ref.astype("float32") * ref_inputs["do"]).sum() + (ref_ht * ref_inputs["dht"]).sum()).backward() + + assert_close("o", ref, tri, 0.005) + assert_close("ht", ref_ht, tri_ht, 0.005) + for name, tolerance in { + "q": 0.008, + "k": 0.008, + "v": 0.008, + "g": 0.02, + "beta": 0.02, + "h0": 0.008, + }.items(): + assert_close(f"d{name}", ref_inputs[name].grad, tri_inputs[name].grad, tolerance) diff --git a/tests/paddle/test_kimi_delta_attention.py b/tests/paddle/test_kimi_delta_attention.py new file mode 100644 index 0000000000..0017a02a72 --- /dev/null +++ b/tests/paddle/test_kimi_delta_attention.py @@ -0,0 +1,281 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +import subprocess +import sys +from types import SimpleNamespace + +import paddle +import pytest + +paddle.enable_compat(scope={"fla", "triton"}, silent=True) + +import fla # noqa: E402 + +paddle.disable_compat() + +from fla.modules import FusedRMSNormGated, ShortConvolution # noqa: E402 +from fla.ops.kda import chunk_kda # noqa: E402 +from fla.ops.utils.index import prepare_cu_seqlens_from_mask, prepare_lens_from_mask # noqa: E402 +from fla.utils import tensor_cache # noqa: E402 + + +@tensor_cache +def _get_unpad_data(attention_mask: paddle.Tensor) -> tuple[paddle.Tensor, paddle.Tensor]: + indices = paddle.nonzero(attention_mask.flatten()).flatten() + return indices, prepare_cu_seqlens_from_mask(attention_mask) + + +def _unpad(x: paddle.Tensor, indices: paddle.Tensor) -> paddle.Tensor: + return x.reshape([-1, x.shape[-1]])[indices].unsqueeze(0) + + +def _pad(x: paddle.Tensor, indices: paddle.Tensor, batch_size: int, seq_len: int) -> paddle.Tensor: + return paddle.scatter_nd(indices.unsqueeze(-1), x.squeeze(0), [batch_size * seq_len, x.shape[-1]]).reshape( + [batch_size, seq_len, x.shape[-1]] + ) + + +class _KimiDeltaAttentionTrainingHarness(paddle.nn.Layer): + """Training-only mirror of the Hugging Face KimiDeltaAttention data flow.""" + + def __init__(self, config: SimpleNamespace): + super().__init__() + self.hidden_size = config.hidden_size + self.head_dim = config.linear_attn_config["head_dim"] + self.num_heads = config.linear_attn_config["num_heads"] + projection_size = self.head_dim * self.num_heads + + self.q_proj = paddle.nn.Linear(self.hidden_size, projection_size, bias_attr=False) + self.k_proj = paddle.nn.Linear(self.hidden_size, projection_size, bias_attr=False) + self.v_proj = paddle.nn.Linear(self.hidden_size, projection_size, bias_attr=False) + self.q_conv1d = ShortConvolution(projection_size, config.linear_attn_config["short_conv_kernel_size"]) + self.k_conv1d = ShortConvolution(projection_size, config.linear_attn_config["short_conv_kernel_size"]) + self.v_conv1d = ShortConvolution(projection_size, config.linear_attn_config["short_conv_kernel_size"]) + + self.A_log = self.create_parameter( + shape=[self.num_heads], + dtype="float32", + default_initializer=paddle.nn.initializer.Assign(paddle.log(paddle.uniform([self.num_heads], min=1, max=16))), + ) + self.f_a_proj = paddle.nn.Linear(self.hidden_size, self.head_dim, bias_attr=False) + self.f_b_proj = paddle.nn.Linear(self.head_dim, projection_size, bias_attr=False) + self.dt_bias = self.create_parameter( + shape=[projection_size], + dtype="float32", + default_initializer=paddle.nn.initializer.Uniform(-1, 1), + ) + self.b_proj = paddle.nn.Linear(self.hidden_size, self.num_heads, bias_attr=False) + self.g_a_proj = paddle.nn.Linear(self.hidden_size, self.head_dim, bias_attr=False) + self.g_b_proj = paddle.nn.Linear(self.head_dim, projection_size, bias_attr=False) + self.o_norm = FusedRMSNormGated(self.head_dim, eps=config.rms_norm_eps, activation="sigmoid") + self.o_proj = paddle.nn.Linear(projection_size, self.hidden_size, bias_attr=False) + self.gate_lower_bound = config.linear_attn_config["gate_lower_bound"] + + def forward( + self, + hidden_states: paddle.Tensor, + attention_mask: paddle.Tensor | None = None, + ) -> paddle.Tensor: + batch_size, seq_len, _ = hidden_states.shape + indices = None + cu_seqlens = None + if attention_mask is not None: + indices, cu_seqlens = _get_unpad_data(attention_mask) + hidden_states = _unpad(hidden_states, indices) + + q, _ = self.q_conv1d(self.q_proj(hidden_states), cu_seqlens=cu_seqlens) + k, _ = self.k_conv1d(self.k_proj(hidden_states), cu_seqlens=cu_seqlens) + v, _ = self.v_conv1d(self.v_proj(hidden_states), cu_seqlens=cu_seqlens) + g = self.f_b_proj(self.f_a_proj(hidden_states)).reshape([*hidden_states.shape[:-1], self.num_heads, self.head_dim]) + beta = self.b_proj(hidden_states).astype("float32") + q = q.reshape([*q.shape[:-1], self.num_heads, self.head_dim]) + k = k.reshape([*k.shape[:-1], self.num_heads, self.head_dim]) + v = v.reshape([*v.shape[:-1], self.num_heads, self.head_dim]) + + o, _ = chunk_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=self.A_log, + dt_bias=self.dt_bias, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + safe_gate=True, + lower_bound=self.gate_lower_bound, + transpose_state_layout=True, + cu_seqlens=cu_seqlens, + ) + gate = self.g_b_proj(self.g_a_proj(hidden_states)).reshape( + [*hidden_states.shape[:-1], self.num_heads, self.head_dim] + ) + o = self.o_norm(o, gate).reshape([*hidden_states.shape[:-1], self.num_heads * self.head_dim]) + o = self.o_proj(o) + return o if indices is None else _pad(o, indices, batch_size, seq_len) + + +def _assert_finite_gradients(layer: paddle.nn.Layer, x: paddle.Tensor) -> None: + assert x.grad is not None + assert paddle.isfinite(x.grad).all().item() + for name, parameter in layer.named_parameters(): + assert parameter.grad is not None, name + assert paddle.isfinite(parameter.grad).all().item(), name + + +def test_huggingface_kda_training_import_contract(): + result = subprocess.run( + [ + sys.executable, + "-c", + """ +import sys + +import paddle + +paddle.enable_compat(scope={"fla", "triton"}, silent=True) + +import fla + +paddle.disable_compat() + +from fla.modules import FusedRMSNormGated, ShortConvolution +from fla.modules.backends import dispatch as modules_dispatch +from fla.ops.backends import dispatch as ops_dispatch +from fla.ops.kda import chunk_kda +from fla.ops.utils.index import prepare_cu_seqlens_from_mask, prepare_lens_from_mask +from fla.utils import tensor_cache + +assert FusedRMSNormGated is not None +assert ShortConvolution is not None +assert modules_dispatch is ops_dispatch +assert callable(chunk_kda) +assert callable(prepare_cu_seqlens_from_mask) +assert callable(prepare_lens_from_mask) +assert callable(tensor_cache) +assert fla.modules is sys.modules["fla.modules"] +assert fla.modules.conv.cp is sys.modules["fla.modules.conv.cp"] +assert fla.modules.conv.triton is sys.modules["fla.modules.conv.triton"] +assert fla.ops is sys.modules["fla.ops"] +assert fla.ops.cp is sys.modules["fla.ops.cp"] +assert fla.ops.kda is sys.modules["fla.ops.kda"] +assert fla.ops.kda.chunk_kda is chunk_kda +assert fla.ops.utils is sys.modules["fla.ops.utils"] +assert "fla.ops.kda" in sys.modules +""", + ], + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stderr + assert fla.modules is sys.modules["fla.modules"] + assert fla.modules.conv.cp is sys.modules["fla.modules.conv.cp"] + assert fla.modules.conv.triton is sys.modules["fla.modules.conv.triton"] + assert fla.ops is sys.modules["fla.ops"] + assert fla.ops.cp is sys.modules["fla.ops.cp"] + assert fla.ops.kda is sys.modules["fla.ops.kda"] + assert fla.ops.kda.chunk_kda is chunk_kda + assert fla.ops.utils is sys.modules["fla.ops.utils"] + assert ShortConvolution.__module__ == "fla.modules.conv.short_conv" + assert FusedRMSNormGated.__module__ == "fla.modules.fused_norm_gate" + assert callable(chunk_kda) + assert callable(prepare_lens_from_mask) + assert callable(prepare_cu_seqlens_from_mask) + assert callable(tensor_cache) + + +def test_short_convolution_varlen_matches_independent_sequences(): + paddle.seed(42) + conv = ShortConvolution(64, 4, activation="silu") + x = paddle.randn([1, 37, 64], dtype="float32") + x.stop_gradient = False + x_ref = x.detach().clone() + x_ref.stop_gradient = False + weight_ref = conv.weight.detach().clone() + weight_ref.stop_gradient = False + cu_seqlens = paddle.to_tensor([0, 13, 37], dtype="int32") + + actual, _ = conv(x, cu_seqlens=cu_seqlens) + expected = [] + for bos, eos in [(0, 13), (13, 37)]: + y = paddle.nn.functional.conv1d( + x_ref[:, bos:eos].transpose([0, 2, 1]), + weight_ref, + padding=3, + groups=64, + ) + expected.append(paddle.nn.functional.silu(y[:, :, :eos - bos].transpose([0, 2, 1]))) + expected = paddle.concat(expected, axis=1) + + paddle.testing.assert_close(actual, expected, rtol=1e-6, atol=1e-6) + output_gradient = paddle.randn(actual.shape, dtype="float32") + (actual * output_gradient).sum().backward() + (expected * output_gradient).sum().backward() + paddle.testing.assert_close(x.grad, x_ref.grad, rtol=1e-5, atol=1e-5) + paddle.testing.assert_close(conv.weight.grad, weight_ref.grad, rtol=1e-5, atol=1e-5) + + +def test_fused_rms_norm_gated_forward_backward(): + paddle.seed(42) + x = paddle.randn([2, 17, 2, 64], dtype="float32") + g = paddle.randn([2, 17, 2, 64], dtype="float32") + x.stop_gradient = False + g.stop_gradient = False + layer = FusedRMSNormGated(64, eps=1e-6, activation="sigmoid") + x_ref = x.detach().clone() + x_ref.stop_gradient = False + g_ref = g.detach().clone() + g_ref.stop_gradient = False + weight_ref = layer.weight.detach().clone() + weight_ref.stop_gradient = False + + actual = layer(x, g) + expected = x_ref * paddle.rsqrt(x_ref.square().mean(axis=-1, keepdim=True) + 1e-6) + expected = expected * weight_ref * paddle.nn.functional.sigmoid(g_ref) + paddle.testing.assert_close(actual, expected, rtol=1e-6, atol=1e-6) + + output_gradient = paddle.randn(actual.shape, dtype="float32") + (actual * output_gradient).sum().backward() + (expected * output_gradient).sum().backward() + paddle.testing.assert_close(x.grad, x_ref.grad, rtol=1e-5, atol=1e-5) + paddle.testing.assert_close(g.grad, g_ref.grad, rtol=1e-5, atol=1e-5) + paddle.testing.assert_close(layer.weight.grad, weight_ref.grad, rtol=1e-5, atol=1e-5) + + +@pytest.mark.parametrize("use_padding_mask", [False, True]) +def test_kimi_delta_attention_training_forward_backward(use_padding_mask: bool): + paddle.seed(42) + config = SimpleNamespace( + hidden_size=128, + rms_norm_eps=1e-6, + linear_attn_config={ + "short_conv_kernel_size": 4, + "head_dim": 128, + "num_heads": 1, + "gate_lower_bound": -5.0, + }, + ) + layer = _KimiDeltaAttentionTrainingHarness(config) + hidden_states = paddle.randn([2 if use_padding_mask else 1, 64, 128], dtype="float32") + hidden_states.stop_gradient = False + attention_mask = None + if use_padding_mask: + attention_mask = paddle.to_tensor([[1] * 53 + [0] * 11, [1] * 64], dtype="bool") + + with paddle.amp.auto_cast(enable=True, dtype="bfloat16"): + output = layer(hidden_states, attention_mask) + loss = output.astype("float32").square().mean() + loss.backward() + + assert output.shape == hidden_states.shape + assert paddle.isfinite(output).all().item() + if attention_mask is not None: + assert (output[0, 53:] == 0).all().item() + _assert_finite_gradients(layer, hidden_states)