diff --git a/.gitignore b/.gitignore index 66b35c4..43a0a58 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,7 @@ build/** .cache/** *.egg-info/** *.so +*.pyd *.ncu-rep *.nsys-rep .vscode/** diff --git a/README.md b/README.md index 57d7fc2..06c8192 100644 --- a/README.md +++ b/README.md @@ -7,9 +7,9 @@ FlashKDA: Flash Kimi Delta Attention — high-performance KDA kernels built on C - **2026-04-22** — Deep-Dive Blog: the design decisions behind FlashKDA v1, read it [here](docs/20260420-flashkda-v1-deep-dive.md). ## Requirements -- SM90 and above -- CUDA 12.9 and above - PyTorch 2.4 and above +- CPU backend: any PyTorch-supported CPU +- CUDA backend: SM90 and above, CUDA 12.9 and above ## Installation ```bash @@ -27,6 +27,64 @@ FLASH_KDA_CUDA_ARCHS=all pip install -v --no-build-isolation . Supported values are `auto` (default), `all`, or a comma-separated arch list such as `90a,100a`. +For a CPU-only installation, skip the CUDA extension build: + +```bash +FLASH_KDA_BUILD_CUDA=0 pip install -v --no-build-isolation . +``` + +Set `FLASH_KDA_BUILD_CUDA=1` to require a CUDA build and fail immediately when +the CUDA toolkit is unavailable. The default `auto` mode builds the extension +when a CUDA device is visible or explicit architectures are requested. + +The C++ CPU extension is built by default. Set both build flags to `0` for a +compiler-free, pure-PyTorch installation: + +```bash +FLASH_KDA_BUILD_CPU=0 FLASH_KDA_BUILD_CUDA=0 pip install -v --no-build-isolation . +``` + +## PyTorch CPU and training backend + +`flash_kda.torch_kda` implements the recurrent KDA definition for CPU training. +It uses a C++/ATen recurrent forward and analytical backward when the CPU +extension is available, with a native PyTorch implementation as the portable +fallback and numerical oracle. Both paths participate in PyTorch autograd, +support arbitrary key/value dimensions and grouped value heads, and follow +FLA's `chunk_kda` argument conventions. + +```python +from flash_kda import torch_kda + +out, final_state = torch_kda( + q=q, k=k, v=v, g=g, beta=beta, + A_log=A_log, dt_bias=dt_bias, + lower_bound=-5.0, + initial_state=initial_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, +) + +loss = out.square().mean() +loss.backward() +``` + +The CPU backend is primarily intended as an accessible research and prototyping +path for integrating KDA into conventional PyTorch models, validating model +innovations, running correctness checks, and experimenting with small workloads. +The C++ path removes the Python token loop, but it still composes ATen tensor +operations and is not yet a fused AVX/NEON kernel for long-sequence training. +Pass `use_cpp_backend=False` to `torch_kda` to force the PyTorch oracle or to +compute higher-order gradients; the analytical C++ backward supports +first-order training gradients. The existing CUTLASS path remains the +high-performance CUDA inference backend. + +FLA does not currently dispatch CPU `chunk_kda` calls to FlashKDA automatically. +The compatible API added here is intended to support a follow-up FLA backend +integration. + ## Using FlashKDA as an FLA backend Once installed, FlashKDA is auto-dispatched from `flash-linear-attention`'s `chunk_kda`. See [fla-org/flash-linear-attention#852](https://github.com/fla-org/flash-linear-attention/pull/852) for integration details. diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp new file mode 100644 index 0000000..b959624 --- /dev/null +++ b/csrc/cpu/torch_bindings.cpp @@ -0,0 +1,162 @@ +#include + +#include + +namespace { + +void check_recurrent_inputs( + const torch::Tensor& q, + const torch::Tensor& k, + const torch::Tensor& v, + const torch::Tensor& g, + const torch::Tensor& beta, + const torch::Tensor& initial_state +) { + TORCH_CHECK(q.device().is_cpu(), "q must be a CPU tensor"); + TORCH_CHECK(k.device().is_cpu() && v.device().is_cpu() && g.device().is_cpu(), + "k, v, and g must be CPU tensors"); + TORCH_CHECK(beta.device().is_cpu() && initial_state.device().is_cpu(), + "beta and initial_state must be CPU tensors"); + TORCH_CHECK(q.scalar_type() == torch::kFloat32 || q.scalar_type() == torch::kFloat64, + "CPU KDA supports float32 and float64 compute tensors"); + TORCH_CHECK(k.scalar_type() == q.scalar_type() && v.scalar_type() == q.scalar_type() && + g.scalar_type() == q.scalar_type() && beta.scalar_type() == q.scalar_type() && + initial_state.scalar_type() == q.scalar_type(), + "all CPU KDA tensors must use the same dtype"); + TORCH_CHECK(q.dim() == 4 && k.dim() == 4 && v.dim() == 4 && g.dim() == 4, + "q, k, v, and g must be 4D tensors"); + TORCH_CHECK(beta.dim() == 3 && initial_state.dim() == 4, + "beta must be 3D and initial_state must be 4D"); + TORCH_CHECK(q.sizes() == k.sizes() && q.sizes() == g.sizes(), + "q, k, and g must have the same shape"); + + const auto batch = q.size(0); + const auto sequence_length = q.size(1); + const auto heads = q.size(2); + const auto key_dim = q.size(3); + const auto value_dim = v.size(3); + TORCH_CHECK(v.size(0) == batch && v.size(1) == sequence_length && v.size(2) == heads, + "v must match q's batch, sequence, and head dimensions"); + TORCH_CHECK(beta.size(0) == batch && beta.size(1) == sequence_length && beta.size(2) == heads, + "beta must match q's batch, sequence, and head dimensions"); + TORCH_CHECK(initial_state.size(0) == batch && initial_state.size(1) == heads && + initial_state.size(2) == key_dim && initial_state.size(3) == value_dim, + "initial_state must have shape [B, H, K, V]"); +} + +std::vector recurrent_forward( + const torch::Tensor& q, + const torch::Tensor& k, + const torch::Tensor& v, + const torch::Tensor& g, + const torch::Tensor& beta, + const torch::Tensor& initial_state +) { + check_recurrent_inputs(q, k, v, g, beta, initial_state); + TORCH_CHECK(q.size(1) > 0, "the C++ CPU backend requires a non-empty sequence"); + + auto state = initial_state; + std::vector outputs; + outputs.reserve(q.size(1)); + for (int64_t token_index = 0; token_index < q.size(1); ++token_index) { + const auto q_token = q.select(1, token_index); + const auto k_token = k.select(1, token_index); + const auto v_token = v.select(1, token_index); + const auto gate_token = g.select(1, token_index); + const auto beta_token = beta.select(1, token_index); + + state = state * gate_token.exp().unsqueeze(-1); + const auto predicted_value = (k_token.unsqueeze(-1) * state).sum(-2); + const auto value_delta = beta_token.unsqueeze(-1) * (v_token - predicted_value); + state = state + k_token.unsqueeze(-1) * value_delta.unsqueeze(-2); + outputs.push_back((q_token.unsqueeze(-1) * state).sum(-2)); + } + + return {torch::stack(outputs, 1), state}; +} + +std::vector recurrent_backward( + const torch::Tensor& q, + const torch::Tensor& k, + const torch::Tensor& v, + const torch::Tensor& g, + const torch::Tensor& beta, + const torch::Tensor& initial_state, + const torch::Tensor& grad_output, + const torch::Tensor& grad_final_state +) { + check_recurrent_inputs(q, k, v, g, beta, initial_state); + TORCH_CHECK(grad_output.sizes() == v.sizes(), "grad_output must match v's shape"); + TORCH_CHECK(grad_final_state.sizes() == initial_state.sizes(), + "grad_final_state must match initial_state's shape"); + + std::vector states; + states.reserve(q.size(1) + 1); + auto state = initial_state; + states.push_back(state); + for (int64_t token_index = 0; token_index < q.size(1); ++token_index) { + const auto k_token = k.select(1, token_index); + const auto v_token = v.select(1, token_index); + const auto gate_token = g.select(1, token_index); + const auto beta_token = beta.select(1, token_index); + + const auto decayed_state = state * gate_token.exp().unsqueeze(-1); + const auto predicted_value = (k_token.unsqueeze(-1) * decayed_state).sum(-2); + const auto value_delta = beta_token.unsqueeze(-1) * (v_token - predicted_value); + state = decayed_state + k_token.unsqueeze(-1) * value_delta.unsqueeze(-2); + states.push_back(state); + } + + auto dq = torch::zeros_like(q); + auto dk = torch::zeros_like(k); + auto dv = torch::zeros_like(v); + auto dg = torch::zeros_like(g); + auto dbeta = torch::zeros_like(beta); + auto dstate = grad_final_state; + + for (int64_t token_index = q.size(1) - 1; token_index >= 0; --token_index) { + const auto q_token = q.select(1, token_index); + const auto k_token = k.select(1, token_index); + const auto v_token = v.select(1, token_index); + const auto gate_token = g.select(1, token_index); + const auto beta_token = beta.select(1, token_index); + const auto output_gradient = grad_output.select(1, token_index); + const auto previous_state = states[token_index]; + const auto current_state = states[token_index + 1]; + const auto decay = gate_token.exp(); + const auto decayed_state = previous_state * decay.unsqueeze(-1); + const auto predicted_value = (k_token.unsqueeze(-1) * decayed_state).sum(-2); + const auto residual = v_token - predicted_value; + const auto value_delta = beta_token.unsqueeze(-1) * residual; + + const auto dq_token = (current_state * output_gradient.unsqueeze(-2)).sum(-1); + auto total_state_gradient = dstate + q_token.unsqueeze(-1) * output_gradient.unsqueeze(-2); + const auto dk_from_update = (total_state_gradient * value_delta.unsqueeze(-2)).sum(-1); + const auto dvalue_delta = (total_state_gradient * k_token.unsqueeze(-1)).sum(-2); + const auto dbeta_token = (dvalue_delta * residual).sum(-1); + const auto dresidual = dvalue_delta * beta_token.unsqueeze(-1); + const auto dv_token = dresidual; + const auto dpredicted_value = -dresidual; + const auto dk_from_prediction = + (decayed_state * dpredicted_value.unsqueeze(-2)).sum(-1); + const auto ddecayed_state = + total_state_gradient + k_token.unsqueeze(-1) * dpredicted_value.unsqueeze(-2); + const auto dg_token = (ddecayed_state * previous_state).sum(-1) * decay; + dstate = ddecayed_state * decay.unsqueeze(-1); + + dq.select(1, token_index).copy_(dq_token); + dk.select(1, token_index).copy_(dk_from_update + dk_from_prediction); + dv.select(1, token_index).copy_(dv_token); + dg.select(1, token_index).copy_(dg_token); + dbeta.select(1, token_index).copy_(dbeta_token); + } + + return {dq, dk, dv, dg, dbeta, dstate}; +} + +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + module.def("recurrent_forward", &recurrent_forward, "KDA recurrent forward (CPU)"); + module.def("recurrent_backward", &recurrent_backward, "KDA recurrent backward (CPU)"); +} diff --git a/flash_kda/__init__.py b/flash_kda/__init__.py index cc03493..52baf0a 100644 --- a/flash_kda/__init__.py +++ b/flash_kda/__init__.py @@ -1,5 +1,12 @@ import torch -from flash_kda_C import fwd as _fwd_raw, get_workspace_size + +from .torch_backend import has_cpu_extension, torch_kda + +try: + from flash_kda_C import fwd as _fwd_raw, get_workspace_size +except ImportError: + _fwd_raw = None + get_workspace_size = None def fwd(q, k, v, g, beta, scale, out, A_log, dt_bias, lower_bound, initial_state=None, final_state=None, cu_seqlens=None): @@ -28,14 +35,50 @@ def fwd(q, k, v, g, beta, scale, out, A_log, dt_bias, lower_bound, initial_state Notes: * Currently requires ``K = V = 128``. - * All input tensors must be CUDA, contiguous, and have the dtypes - listed above. + * The native extension requires contiguous CUDA tensors with the + dtypes listed above. + * CPU tensors use the differentiable PyTorch backend. Prefer + :func:`torch_kda` for new training code because it returns tensors + instead of writing into caller-provided buffers. """ - B, T_seq, H = q.shape[0], q.shape[1], q.shape[2] - T_total = B * T_seq - N = cu_seqlens.numel() - 1 if cu_seqlens is not None else B + if q.is_cuda and _fwd_raw is not None: + B, T_seq, H = q.shape[0], q.shape[1], q.shape[2] + T_total = B * T_seq + N = cu_seqlens.numel() - 1 if cu_seqlens is not None else B + + workspace = torch.empty(get_workspace_size(T_total, H, N), dtype=torch.uint8, device=q.device) + + _fwd_raw(q, k, v, g, beta, float(scale), out, workspace, A_log, dt_bias, lower_bound, + initial_state=initial_state, final_state=final_state, cu_seqlens=cu_seqlens) + return + + output, state = torch_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=final_state is not None, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + state_v_first=True, + cu_seqlens=cu_seqlens, + safe_gate=True, + lower_bound=lower_bound, + A_log=A_log, + dt_bias=dt_bias, + ) + out.copy_(output) + if final_state is not None: + final_state.copy_(state) + + +def has_cuda_extension(): + """Return whether the compiled FlashKDA CUDA extension is available.""" + return _fwd_raw is not None - workspace = torch.empty(get_workspace_size(T_total, H, N), dtype=torch.uint8, device=q.device) - _fwd_raw(q, k, v, g, beta, float(scale), out, workspace, A_log, dt_bias, lower_bound, - initial_state=initial_state, final_state=final_state, cu_seqlens=cu_seqlens) +__all__ = ["fwd", "has_cpu_extension", "has_cuda_extension", "torch_kda"] diff --git a/flash_kda/torch_backend.py b/flash_kda/torch_backend.py new file mode 100644 index 0000000..a7d6830 --- /dev/null +++ b/flash_kda/torch_backend.py @@ -0,0 +1,316 @@ +from __future__ import annotations + +import torch +import torch.nn.functional as F +from torch.autograd.function import once_differentiable + +try: + from flash_kda_cpu_C import recurrent_backward as _cpu_recurrent_backward + from flash_kda_cpu_C import recurrent_forward as _cpu_recurrent_forward +except ImportError: + _cpu_recurrent_backward = None + _cpu_recurrent_forward = None + + +class _CppRecurrentKDA(torch.autograd.Function): + @staticmethod + def forward(ctx, q, k, v, g, beta, initial_state): + output, final_state = _cpu_recurrent_forward( + q.contiguous(), + k.contiguous(), + v.contiguous(), + g.contiguous(), + beta.contiguous(), + initial_state.contiguous(), + ) + ctx.save_for_backward(q, k, v, g, beta, initial_state) + return output, final_state + + @staticmethod + @once_differentiable + def backward(ctx, grad_output, grad_final_state): + q, k, v, g, beta, initial_state = ctx.saved_tensors + if grad_output is None: + grad_output = torch.zeros_like(v) + if grad_final_state is None: + grad_final_state = torch.zeros_like(initial_state) + return tuple( + _cpu_recurrent_backward( + q.contiguous(), + k.contiguous(), + v.contiguous(), + g.contiguous(), + beta.contiguous(), + initial_state.contiguous(), + grad_output.contiguous(), + grad_final_state.contiguous(), + ) + ) + + +def has_cpu_extension() -> bool: + """Return whether the compiled CPU recurrent extension is available.""" + return _cpu_recurrent_forward is not None + + +def _compute_dtype(dtype: torch.dtype) -> torch.dtype: + if dtype in (torch.float16, torch.bfloat16): + return torch.float32 + if dtype in (torch.float32, torch.float64): + return dtype + raise TypeError(f"KDA inputs must use a floating-point dtype, got {dtype}") + + +def _l2_normalize(value: torch.Tensor, eps: float) -> torch.Tensor: + return value * torch.rsqrt(value.square().sum(dim=-1, keepdim=True) + eps) + + +def _validate_inputs( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, +) -> tuple[int, int, int, int, int, int]: + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4 or g.ndim != 4 or beta.ndim != 3: + raise ValueError("expected q, k, v, g to be 4D and beta to be 3D") + if q.shape != k.shape: + raise ValueError(f"q and k must have the same shape, got {q.shape} and {k.shape}") + + batch, sequence_length, query_heads, key_dim = q.shape + value_heads, value_dim = v.shape[2], v.shape[3] + if v.shape[:2] != (batch, sequence_length): + raise ValueError("v must match q's batch and sequence dimensions") + if value_heads % query_heads != 0: + raise ValueError(f"value heads ({value_heads}) must be divisible by query heads ({query_heads})") + if g.shape != (batch, sequence_length, value_heads, key_dim): + raise ValueError( + f"g must have shape {(batch, sequence_length, value_heads, key_dim)}, got {tuple(g.shape)}" + ) + if beta.shape != (batch, sequence_length, value_heads): + raise ValueError( + f"beta must have shape {(batch, sequence_length, value_heads)}, got {tuple(beta.shape)}" + ) + + devices = {tensor.device for tensor in (q, k, v, g, beta)} + if len(devices) != 1: + raise ValueError("q, k, v, g, and beta must be on the same device") + if not all(tensor.is_floating_point() for tensor in (q, k, v, g, beta)): + raise TypeError("q, k, v, g, and beta must use floating-point dtypes") + + return batch, sequence_length, query_heads, key_dim, value_heads, value_dim + + +def _activate_gate( + g: torch.Tensor, + A_log: torch.Tensor | None, + dt_bias: torch.Tensor | None, + lower_bound: float | None, + value_heads: int, + key_dim: int, +) -> torch.Tensor: + if A_log is None: + raise ValueError("A_log is required when use_gate_in_kernel=True") + if A_log.shape != (value_heads,): + raise ValueError(f"A_log must have shape {(value_heads,)}, got {tuple(A_log.shape)}") + + if dt_bias is None: + biased_gate = g + else: + if dt_bias.numel() != value_heads * key_dim: + raise ValueError(f"dt_bias must contain {value_heads * key_dim} values, got {dt_bias.numel()}") + bias = dt_bias.reshape(1, 1, value_heads, key_dim) + biased_gate = g + bias + gate_scale = A_log.exp().reshape(1, 1, value_heads, 1) + if lower_bound is None: + return -gate_scale * F.softplus(biased_gate) + return float(lower_bound) * torch.sigmoid(gate_scale * biased_gate) + + +def _run_sequence( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + state: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + outputs = [] + for token_index in range(q.shape[1]): + q_token = q[:, token_index] + k_token = k[:, token_index] + v_token = v[:, token_index] + gate_token = g[:, token_index] + beta_token = beta[:, token_index] + + state = state * gate_token.exp().unsqueeze(-1) + predicted_value = torch.einsum("bhk,bhkv->bhv", k_token, state) + value_delta = beta_token.unsqueeze(-1) * (v_token - predicted_value) + state = state + torch.einsum("bhk,bhv->bhkv", k_token, value_delta) + outputs.append(torch.einsum("bhk,bhkv->bhv", q_token, state)) + + if outputs: + output = torch.stack(outputs, dim=1) + else: + output = v.new_empty(v.shape) + return output, state + + +def _run_sequence_dispatch( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + state: torch.Tensor, + use_cpp_backend: bool | None, +) -> tuple[torch.Tensor, torch.Tensor]: + can_use_cpp = has_cpu_extension() and q.device.type == "cpu" and q.shape[1] > 0 + if use_cpp_backend is True and not can_use_cpp: + raise RuntimeError("the compiled FlashKDA CPU extension is not available for these inputs") + if use_cpp_backend is not False and can_use_cpp: + return _CppRecurrentKDA.apply(q, k, v, g, beta, state) + return _run_sequence(q, k, v, g, beta, state) + + +def torch_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + use_gate_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = False, + allow_neg_eigval: bool = False, + state_v_first: bool = False, + cu_seqlens: torch.Tensor | None = None, + cu_seqlens_cpu: torch.Tensor | None = None, + safe_gate: bool = False, + lower_bound: float | None = None, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + l2norm_eps: float = 1e-6, + disable_recompute: bool = False, + return_intermediate_states: bool = False, + cp_context=None, + chunk_size: int = 64, + use_cpp_backend: bool | None = None, + **kwargs, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Run differentiable Kimi Delta Attention with native PyTorch operations. + + This backend is intended for CPU research, model prototyping, correctness + checks, and small workloads. It follows FLA's ``chunk_kda`` argument + conventions while using the recurrent KDA definition, so PyTorch autograd + supplies the backward pass automatically. + """ + if kwargs: + unknown = ", ".join(sorted(kwargs)) + raise TypeError(f"unexpected keyword arguments: {unknown}") + del safe_gate, disable_recompute, chunk_size + if return_intermediate_states: + raise ValueError("return_intermediate_states is not supported by the PyTorch backend") + if cp_context is not None: + raise ValueError("context parallelism is not supported by the PyTorch backend") + + batch, sequence_length, query_heads, key_dim, value_heads, value_dim = _validate_inputs(q, k, v, g, beta) + compute_dtype = _compute_dtype(v.dtype) + tensors = (q, k, v, g, beta) + q_compute, k_compute, v_compute, g_compute, beta_compute = ( + tensor.to(compute_dtype) for tensor in tensors + ) + + if use_qk_l2norm_in_kernel: + q_compute = _l2_normalize(q_compute, l2norm_eps) + k_compute = _l2_normalize(k_compute, l2norm_eps) + + group_size = value_heads // query_heads + q_compute = q_compute.repeat_interleave(group_size, dim=2) + k_compute = k_compute.repeat_interleave(group_size, dim=2) + q_compute = q_compute * (key_dim ** -0.5 if scale is None else float(scale)) + + if use_gate_in_kernel: + if A_log is not None and A_log.device != q.device: + raise ValueError("A_log must be on the same device as q") + if dt_bias is not None and dt_bias.device != q.device: + raise ValueError("dt_bias must be on the same device as q") + g_compute = _activate_gate( + g_compute, + None if A_log is None else A_log.to(compute_dtype), + None if dt_bias is None else dt_bias.to(compute_dtype), + lower_bound, + value_heads, + key_dim, + ) + if use_beta_sigmoid_in_kernel: + beta_compute = torch.sigmoid(beta_compute) + if allow_neg_eigval: + beta_compute = beta_compute * 2.0 + + sequence_lengths = cu_seqlens_cpu if cu_seqlens_cpu is not None else cu_seqlens + if sequence_lengths is not None: + if batch != 1: + raise ValueError("batch size must be 1 when cu_seqlens is provided") + if sequence_lengths.ndim != 1 or sequence_lengths.numel() < 2: + raise ValueError("cu_seqlens must be a 1D tensor with at least two elements") + sequence_offsets = sequence_lengths.detach().to(device="cpu", dtype=torch.long).tolist() + if sequence_offsets[0] != 0 or sequence_offsets[-1] != sequence_length: + raise ValueError("cu_seqlens must start at 0 and end at the total sequence length") + if any(end < start for start, end in zip(sequence_offsets, sequence_offsets[1:])): + raise ValueError("cu_seqlens must be nondecreasing") + state_count = len(sequence_offsets) - 1 + else: + sequence_offsets = None + state_count = batch + + expected_state_shape = ( + (state_count, value_heads, value_dim, key_dim) + if state_v_first + else (state_count, value_heads, key_dim, value_dim) + ) + if initial_state is None: + state = v_compute.new_zeros(state_count, value_heads, key_dim, value_dim) + else: + if tuple(initial_state.shape) != expected_state_shape: + raise ValueError(f"initial_state must have shape {expected_state_shape}, got {tuple(initial_state.shape)}") + if initial_state.device != q.device: + raise ValueError("initial_state must be on the same device as q") + state = initial_state.to(compute_dtype) + if state_v_first: + state = state.transpose(-1, -2) + + if sequence_offsets is None: + output, final_state = _run_sequence_dispatch( + q_compute, + k_compute, + v_compute, + g_compute, + beta_compute, + state, + use_cpp_backend, + ) + else: + output_parts = [] + final_states = [] + for sequence_index, (start, end) in enumerate(zip(sequence_offsets, sequence_offsets[1:])): + sequence_output, sequence_state = _run_sequence_dispatch( + q_compute[:, start:end], + k_compute[:, start:end], + v_compute[:, start:end], + g_compute[:, start:end], + beta_compute[:, start:end], + state[sequence_index:sequence_index + 1], + use_cpp_backend, + ) + output_parts.append(sequence_output) + final_states.append(sequence_state) + output = torch.cat(output_parts, dim=1) + final_state = torch.cat(final_states, dim=0) + + if state_v_first: + final_state = final_state.transpose(-1, -2) + return output.to(v.dtype), final_state if output_final_state else None diff --git a/setup.py b/setup.py index 76e44ed..a6f0aa5 100644 --- a/setup.py +++ b/setup.py @@ -1,10 +1,9 @@ import os import subprocess from setuptools import setup -from torch.utils.cpp_extension import CUDAExtension, BuildExtension, CUDA_HOME +from torch.utils.cpp_extension import CppExtension, CUDAExtension, BuildExtension, CUDA_HOME this_dir = os.path.dirname(os.path.abspath(__file__)) -subprocess.run(["git", "submodule", "update", "--init", "cutlass"]) def is_flag_set(flag: str) -> bool: @@ -52,39 +51,80 @@ def get_arch_flags(): return flags -ext_modules = [ - CUDAExtension( - name='flash_kda_C', - sources=[ - 'csrc/flash_kda.cpp', - 'csrc/smxx/fwd_launch.cu', - ], - include_dirs=[ - os.path.join(this_dir, 'cutlass', 'include'), - os.path.join(this_dir, 'cutlass', 'examples', 'common'), - os.path.join(this_dir, 'cutlass', 'tools', 'util', 'include'), - os.path.join(this_dir, 'csrc'), - ], - extra_compile_args={ - 'cxx': ['-O3', '-Wno-psabi'], - 'nvcc': [ - '-O3', - '-U__CUDA_NO_HALF_OPERATORS__', - '-U__CUDA_NO_HALF_CONVERSIONS__', - '-U__CUDA_NO_HALF2_OPERATORS__', - '-U__CUDA_NO_BFLOAT16_CONVERSIONS__', - '--expt-relaxed-constexpr', - '--expt-extended-lambda', - '--use_fast_math', - '--ptxas-options=-v,--register-usage-level=10,--warn-on-spills', - '-lineinfo', - *get_nvcc_thread_args(), - *get_arch_flags(), +def should_build_cuda(): + requested = os.getenv("FLASH_KDA_BUILD_CUDA", "auto").lower() + if requested in ("0", "false", "no", "off"): + return False + if requested in ("1", "true", "yes", "on"): + if CUDA_HOME is None: + raise RuntimeError("FLASH_KDA_BUILD_CUDA=1 requires a CUDA toolkit") + return True + if requested != "auto": + raise ValueError("FLASH_KDA_BUILD_CUDA must be auto, 0, or 1") + if CUDA_HOME is None: + return False + + import torch + + return torch.cuda.is_available() or os.getenv("FLASH_KDA_CUDA_ARCHS", "auto").lower() != "auto" + + +def should_build_cpu(): + requested = os.getenv("FLASH_KDA_BUILD_CPU", "1").lower() + if requested in ("0", "false", "no", "off"): + return False + if requested in ("1", "true", "yes", "on"): + return True + raise ValueError("FLASH_KDA_BUILD_CPU must be 0 or 1") + + +ext_modules = [] +if should_build_cpu(): + cpu_compile_args = ['/O2'] if os.name == 'nt' else ['-O3'] + ext_modules.append( + CppExtension( + name='flash_kda_cpu_C', + sources=['csrc/cpu/torch_bindings.cpp'], + extra_compile_args=cpu_compile_args, + ) + ) + +if should_build_cuda(): + subprocess.run(["git", "submodule", "update", "--init", "cutlass"], check=True) + ext_modules.append( + CUDAExtension( + name='flash_kda_C', + sources=[ + 'csrc/flash_kda.cpp', + 'csrc/smxx/fwd_launch.cu', ], - }, + include_dirs=[ + os.path.join(this_dir, 'cutlass', 'include'), + os.path.join(this_dir, 'cutlass', 'examples', 'common'), + os.path.join(this_dir, 'cutlass', 'tools', 'util', 'include'), + os.path.join(this_dir, 'csrc'), + ], + extra_compile_args={ + 'cxx': ['-O3', '-Wno-psabi'], + 'nvcc': [ + '-O3', + '-U__CUDA_NO_HALF_OPERATORS__', + '-U__CUDA_NO_HALF_CONVERSIONS__', + '-U__CUDA_NO_HALF2_OPERATORS__', + '-U__CUDA_NO_BFLOAT16_CONVERSIONS__', + '--expt-relaxed-constexpr', + '--expt-extended-lambda', + '--use_fast_math', + '--ptxas-options=-v,--register-usage-level=10,--warn-on-spills', + '-lineinfo', + *get_nvcc_thread_args(), + *get_arch_flags(), + ], + }, + ) ) -] -cmdclass = {"build_ext": BuildExtension} + +cmdclass = {"build_ext": BuildExtension} if ext_modules else {} rev = os.getenv("FLASH_KDA_VERSION_SUFFIX", "") if not rev: diff --git a/tests/test_cpu.py b/tests/test_cpu.py new file mode 100644 index 0000000..31b56ae --- /dev/null +++ b/tests/test_cpu.py @@ -0,0 +1,190 @@ +import torch +import pytest + +import flash_kda + + +def _make_inputs(batch=2, sequence_length=4, query_heads=2, value_heads=2, key_dim=3, value_dim=4): + torch.manual_seed(42) + return ( + torch.randn(batch, sequence_length, query_heads, key_dim, dtype=torch.float64, requires_grad=True), + torch.randn(batch, sequence_length, query_heads, key_dim, dtype=torch.float64, requires_grad=True), + torch.randn(batch, sequence_length, value_heads, value_dim, dtype=torch.float64, requires_grad=True), + torch.randn(batch, sequence_length, value_heads, key_dim, dtype=torch.float64, requires_grad=True), + torch.randn(batch, sequence_length, value_heads, dtype=torch.float64, requires_grad=True), + torch.randn(value_heads, dtype=torch.float64, requires_grad=True), + torch.randn(value_heads, key_dim, dtype=torch.float64, requires_grad=True), + ) + + +def _run(inputs, initial_state=None, **kwargs): + q, k, v, g, beta, A_log, dt_bias = inputs + return flash_kda.torch_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=0.5, + initial_state=initial_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + lower_bound=-5.0, + A_log=A_log, + dt_bias=dt_bias, + **kwargs, + ) + + +def test_torch_kda_cpu_forward_and_backward(): + inputs = _make_inputs() + initial_state = torch.randn(2, 2, 3, 4, dtype=torch.float64, requires_grad=True) + + output, final_state = _run(inputs, initial_state=initial_state) + loss = output.square().mean() + final_state.square().mean() + loss.backward() + + assert output.shape == (2, 4, 2, 4) + assert final_state.shape == (2, 2, 3, 4) + for tensor in (*inputs, initial_state): + assert tensor.grad is not None + assert torch.isfinite(tensor.grad).all() + + +@pytest.mark.parametrize( + "use_cpp_backend", + [ + False, + pytest.param( + True, + marks=pytest.mark.skipif( + not flash_kda.has_cpu_extension(), + reason="compiled CPU extension is unavailable", + ), + ), + ], +) +def test_torch_kda_cpu_gradcheck(use_cpp_backend): + inputs = _make_inputs(batch=1, sequence_length=2, query_heads=1, value_heads=1, key_dim=2, value_dim=2) + initial_state = torch.randn(1, 1, 2, 2, dtype=torch.float64, requires_grad=True) + + def function(*arguments): + return _run(arguments[:-1], initial_state=arguments[-1], use_cpp_backend=use_cpp_backend) + + assert torch.autograd.gradcheck(function, (*inputs, initial_state), fast_mode=True) + + +@pytest.mark.skipif(not flash_kda.has_cpu_extension(), reason="compiled CPU extension is unavailable") +def test_cpp_backend_matches_pytorch_forward_and_gradients(): + inputs = _make_inputs(batch=1, sequence_length=3, query_heads=1, value_heads=2, key_dim=2, value_dim=3) + initial_state = torch.randn(1, 2, 2, 3, dtype=torch.float64, requires_grad=True) + + def evaluate(use_cpp_backend): + cloned_inputs = tuple(tensor.detach().clone().requires_grad_(True) for tensor in inputs) + cloned_state = initial_state.detach().clone().requires_grad_(True) + output, final_state = _run( + cloned_inputs, + initial_state=cloned_state, + use_cpp_backend=use_cpp_backend, + ) + gradients = torch.autograd.grad( + output.square().sum() + final_state.square().sum(), + (*cloned_inputs, cloned_state), + ) + return output, final_state, gradients + + cpp_output, cpp_state, cpp_gradients = evaluate(True) + torch_output, torch_state, torch_gradients = evaluate(False) + + torch.testing.assert_close(cpp_output, torch_output) + torch.testing.assert_close(cpp_state, torch_state) + for cpp_gradient, torch_gradient in zip(cpp_gradients, torch_gradients): + torch.testing.assert_close(cpp_gradient, torch_gradient) + + +def test_torch_kda_cpu_supports_fla_gate_and_grouped_value_heads(): + inputs = _make_inputs(batch=1, query_heads=1, value_heads=2, key_dim=3, value_dim=2) + q, k, v, g, beta, A_log, dt_bias = inputs + + output, final_state = flash_kda.torch_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + output_final_state=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + ) + (output.sum() + final_state.sum()).backward() + + assert output.shape == (1, 4, 2, 2) + assert final_state.shape == (1, 2, 3, 2) + assert A_log.grad is not None + assert dt_bias.grad is not None + + +def test_torch_kda_cpu_varlen_matches_individual_sequences(): + inputs = _make_inputs(batch=1, sequence_length=5, query_heads=1, value_heads=1, key_dim=2, value_dim=3) + initial_state = torch.randn(2, 1, 3, 2, dtype=torch.float64) + cu_seqlens = torch.tensor([0, 2, 5], dtype=torch.long) + + output, final_state = _run( + inputs, + initial_state=initial_state, + state_v_first=True, + cu_seqlens=cu_seqlens, + ) + + expected_outputs = [] + expected_states = [] + for sequence_index, (start, end) in enumerate(zip(cu_seqlens[:-1], cu_seqlens[1:])): + sliced_inputs = tuple( + tensor[:, start:end] if tensor.ndim >= 3 and tensor.shape[0] == 1 else tensor + for tensor in inputs + ) + sequence_output, sequence_state = _run( + sliced_inputs, + initial_state=initial_state[sequence_index:sequence_index + 1], + state_v_first=True, + ) + expected_outputs.append(sequence_output) + expected_states.append(sequence_state) + + assert torch.allclose(output, torch.cat(expected_outputs, dim=1)) + assert torch.allclose(final_state, torch.cat(expected_states, dim=0)) + + +def test_legacy_fwd_uses_cpu_backend(): + inputs = _make_inputs(batch=1, sequence_length=3, query_heads=1, value_heads=1, key_dim=2, value_dim=2) + q, k, v, g, beta, A_log, dt_bias = inputs + initial_state = torch.randn(1, 1, 2, 2, dtype=torch.float64) + output = torch.empty_like(v) + final_state = torch.empty_like(initial_state) + + flash_kda.fwd( + q, + k, + v, + g, + beta, + 0.5, + output, + A_log, + dt_bias, + -5.0, + initial_state=initial_state, + final_state=final_state, + ) + expected_output, expected_state = _run( + inputs, + initial_state=initial_state, + state_v_first=True, + ) + + assert torch.allclose(output, expected_output) + assert torch.allclose(final_state, expected_state)