diff --git a/csrc/flash_kda.cpp b/csrc/flash_kda.cpp index 81f5483..5253b19 100644 --- a/csrc/flash_kda.cpp +++ b/csrc/flash_kda.cpp @@ -39,7 +39,8 @@ void fwd( double lower_bound, std::optional initial_state = std::nullopt, std::optional final_state = std::nullopt, - std::optional cu_seqlens = std::nullopt + std::optional cu_seqlens = std::nullopt, + std::optional intermediate_state = std::nullopt ) { TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda() && g.is_cuda() && beta.is_cuda() && out.is_cuda() && workspace.is_cuda(), "all tensors must be on CUDA"); @@ -57,6 +58,7 @@ void fwd( bool has_state_in = initial_state.has_value(); bool has_state_out = final_state.has_value(); bool state_fp32 = false; + bool has_intermediate_state = intermediate_state.has_value(); if (has_state_in) { auto& is = initial_state.value(); @@ -150,6 +152,9 @@ void fwd( TORCH_CHECK(B == 1, "B must be 1 when cu_seqlens is provided"); auto& cu_seqlens_t = cu_seqlens.value(); TORCH_CHECK(cu_seqlens_t.is_cuda(), "cu_seqlens must be on CUDA"); + TORCH_CHECK(cu_seqlens_t.device() == q.device(), + "cu_seqlens must be on the same CUDA device as q"); + TORCH_CHECK(cu_seqlens_t.is_contiguous(), "cu_seqlens must be contiguous"); TORCH_CHECK(cu_seqlens_t.dtype() == torch::kLong, "cu_seqlens must be int64"); TORCH_CHECK(cu_seqlens_t.dim() == 1, "cu_seqlens must be 1D"); N_val = cu_seqlens_t.numel() - 1; @@ -180,36 +185,60 @@ void fwd( total_tiles = int(N_val * ((T_seq + CHUNK - 1) / CHUNK)); // exact for batched } + if (has_intermediate_state) { + auto& ims = intermediate_state.value(); + TORCH_CHECK(ims.is_cuda() && ims.is_contiguous(), + "intermediate_state must be a contiguous CUDA tensor"); + TORCH_CHECK(ims.device() == q.device(), + "intermediate_state must be on the same CUDA device as q"); + TORCH_CHECK(ims.dtype() == torch::kBFloat16, + "intermediate_state must be bfloat16"); + TORCH_CHECK(ims.dim() == 4 && ims.size(0) == H && + ims.size(1) == total_tiles && ims.size(2) == D && ims.size(3) == D, + "intermediate_state must be [H, total_tiles, D, D]"); + } + auto intermediate_state_ptr = has_intermediate_state + ? reinterpret_cast(intermediate_state->data_ptr()) + : nullptr; + // Dispatch based on state configuration and varlen - #define LAUNCH(HI, HO, FP32, VL) \ - launch_fwd<128, HI, HO, FP32, VL>( \ + #define LAUNCH(HI, HO, FP32, HIS, VL) \ + launch_fwd<128, HI, HO, FP32, HIS, VL>( \ q_ptr, k_ptr, v_ptr, g_ptr, beta_t_ptr, \ - initial_state_raw, scale_f, final_state_raw, out_ptr, \ + initial_state_raw, scale_f, final_state_raw, intermediate_state_ptr, out_ptr, \ workspace_ptr, total_tiles, \ int(T_total), int(H), int(N_val), cu_seqlens_dev, \ A_log_ptr, dt_bias_ptr, gate_scale, stream) - #define DISPATCH_STATE(VL) \ + #define DISPATCH_STATE(HIS, VL) \ if (!has_state_in && !has_state_out) { \ - LAUNCH(false, false, false, VL); \ + LAUNCH(false, false, false, HIS, VL); \ } else if (has_state_in && has_state_out && state_fp32) { \ - LAUNCH(true, true, true, VL); \ + LAUNCH(true, true, true, HIS, VL); \ } else if (has_state_in && has_state_out && !state_fp32) { \ - LAUNCH(true, true, false, VL); \ + LAUNCH(true, true, false, HIS, VL); \ } else if (!has_state_in && has_state_out && state_fp32) { \ - LAUNCH(false, true, true, VL); \ + LAUNCH(false, true, true, HIS, VL); \ } else if (!has_state_in && has_state_out && !state_fp32) { \ - LAUNCH(false, true, false, VL); \ + LAUNCH(false, true, false, HIS, VL); \ } else if (has_state_in && !has_state_out && state_fp32) { \ - LAUNCH(true, false, true, VL); \ + LAUNCH(true, false, true, HIS, VL); \ } else { \ - LAUNCH(true, false, false, VL); \ + LAUNCH(true, false, false, HIS, VL); \ } - if (is_varlen) { - DISPATCH_STATE(true); + if (has_intermediate_state) { + if (is_varlen) { + DISPATCH_STATE(true, true); + } else { + DISPATCH_STATE(true, false); + } } else { - DISPATCH_STATE(false); + if (is_varlen) { + DISPATCH_STATE(false, true); + } else { + DISPATCH_STATE(false, false); + } } #undef DISPATCH_STATE @@ -223,7 +252,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("workspace"), py::arg("A_log"), py::arg("dt_bias"), py::arg("lower_bound"), py::arg("initial_state") = py::none(), py::arg("final_state") = py::none(), - py::arg("cu_seqlens") = py::none()); + py::arg("cu_seqlens") = py::none(), + py::arg("intermediate_state") = py::none()); m.def("get_workspace_size", static_cast(&get_workspace_size), "Get workspace size in bytes", diff --git a/csrc/fwd.h b/csrc/fwd.h index deb82b6..783abdc 100644 --- a/csrc/fwd.h +++ b/csrc/fwd.h @@ -3,7 +3,7 @@ #include -template +template void launch_fwd( cutlass::bfloat16_t const* q_ptr, cutlass::bfloat16_t const* k_ptr, @@ -13,6 +13,7 @@ void launch_fwd( void const* initial_state_ptr, float scale, void* final_state_ptr, + cutlass::bfloat16_t* intermediate_state_ptr, cutlass::bfloat16_t* out_ptr, void* workspace_ptr, int total_tiles, diff --git a/csrc/smxx/fwd_kernel2.cuh b/csrc/smxx/fwd_kernel2.cuh index 26f73fe..e070a4b 100644 --- a/csrc/smxx/fwd_kernel2.cuh +++ b/csrc/smxx/fwd_kernel2.cuh @@ -128,6 +128,7 @@ template < bool HasStateIn = true, bool HasStateOut = true, bool StateFP32 = false, + bool HasIntermediateState = false, bool IsVarlen = true > __global__ void __launch_bounds__(NumThreads) _flash_kda_fwd_recurrence( @@ -143,6 +144,7 @@ __global__ void __launch_bounds__(NumThreads) _flash_kda_fwd_recurrence( CUTE_GRID_CONSTANT TmaStoreState const tma_store_final_state, CUTE_GRID_CONSTANT TmaStoreOut const tma_store_out, cutlass::bfloat16_t* out_raw_ptr, + cutlass::bfloat16_t* intermediate_state_ptr, int T_total, int H, int N, @@ -732,6 +734,18 @@ __global__ void __launch_bounds__(NumThreads) _flash_kda_fwd_recurrence( } compute_barrier.arrive_and_wait(); + if constexpr (HasIntermediateState) { + // StateSmemLayout maps the logical [D, D] state to its swizzled + // shared-memory storage. Snapshot after the whole update is + // visible so every chunk is independently reusable by a caller. + int64_t snapshot_offset = + (int64_t(head_idx) * total_tiles + tile_base + t) * D * D; + for (int i = compute_tid; i < D * D; i += kComputeThreads) { + intermediate_state_ptr[snapshot_offset + i] = + s_acc(i / D, i % D); + } + } + #ifndef TMA_DISABLE_ALL cutlass::arch::fence_view_async_shared(); store_pipeline.producer_commit(out_write); diff --git a/csrc/smxx/fwd_launch.cu b/csrc/smxx/fwd_launch.cu index 91a67a7..55ee291 100644 --- a/csrc/smxx/fwd_launch.cu +++ b/csrc/smxx/fwd_launch.cu @@ -3,7 +3,7 @@ #include "fwd_kernel2.cuh" // ==================== launch_fwd ==================== -template +template void launch_fwd( cutlass::bfloat16_t const* q_ptr, cutlass::bfloat16_t const* k_ptr, @@ -13,6 +13,7 @@ void launch_fwd( void const* initial_state_ptr, float scale, void* final_state_ptr, + cutlass::bfloat16_t* intermediate_state_ptr, cutlass::bfloat16_t* out_ptr, void* workspace_ptr, int total_tiles, @@ -195,7 +196,7 @@ void launch_fwd( decltype(tma_store_final_state), decltype(tma_store_out), CHUNK, D, kInputStages, kOutputStages, kK2Threads, - HasStateIn, HasStateOut, StateFP32, IsVarlen + HasStateIn, HasStateOut, StateFP32, HasIntermediateState, IsVarlen >; cudaFuncSetAttribute(kernel2, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size_k2); @@ -210,29 +211,32 @@ void launch_fwd( tma_load_initial_state, tma_store_final_state, tma_store_out, - out_ptr, T_total, H, N, cu_seqlens_ptr, total_tiles + out_ptr, intermediate_state_ptr, + T_total, H, N, cu_seqlens_ptr, total_tiles ); } #endif } // Explicit instantiations -#define INSTANTIATE_LAUNCH_FWD(D, HI, HO, FP32, VL) \ - template void launch_fwd( \ +#define INSTANTIATE_LAUNCH_FWD(D, HI, HO, FP32, HIS, VL) \ + template void launch_fwd( \ cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, \ cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, \ cutlass::bfloat16_t const*, void const*, float, void*, \ - cutlass::bfloat16_t*, void*, int, int, int, int, \ + cutlass::bfloat16_t*, cutlass::bfloat16_t*, void*, int, int, int, int, \ int64_t const*, float const*, float const*, float, cudaStream_t); -#define INSTANTIATE_STATE_VARIANTS(VL) \ - INSTANTIATE_LAUNCH_FWD(128, true, true, false, VL) \ - INSTANTIATE_LAUNCH_FWD(128, true, true, true, VL) \ - INSTANTIATE_LAUNCH_FWD(128, false, false, false, VL) \ - INSTANTIATE_LAUNCH_FWD(128, false, true, false, VL) \ - INSTANTIATE_LAUNCH_FWD(128, true, false, false, VL) \ - INSTANTIATE_LAUNCH_FWD(128, false, true, true, VL) \ - INSTANTIATE_LAUNCH_FWD(128, true, false, true, VL) - -INSTANTIATE_STATE_VARIANTS(true) // varlen -INSTANTIATE_STATE_VARIANTS(false) // non-varlen +#define INSTANTIATE_STATE_VARIANTS(HIS, VL) \ + INSTANTIATE_LAUNCH_FWD(128, true, true, false, HIS, VL) \ + INSTANTIATE_LAUNCH_FWD(128, true, true, true, HIS, VL) \ + INSTANTIATE_LAUNCH_FWD(128, false, false, false, HIS, VL) \ + INSTANTIATE_LAUNCH_FWD(128, false, true, false, HIS, VL) \ + INSTANTIATE_LAUNCH_FWD(128, true, false, false, HIS, VL) \ + INSTANTIATE_LAUNCH_FWD(128, false, true, true, HIS, VL) \ + INSTANTIATE_LAUNCH_FWD(128, true, false, true, HIS, VL) + +INSTANTIATE_STATE_VARIANTS(false, true) // varlen, no chunk-state snapshots +INSTANTIATE_STATE_VARIANTS(false, false) // batched, no chunk-state snapshots +INSTANTIATE_STATE_VARIANTS(true, true) // varlen, chunk-state snapshots +INSTANTIATE_STATE_VARIANTS(true, false) // batched, chunk-state snapshots diff --git a/flash_kda/__init__.py b/flash_kda/__init__.py index cc03493..89832d2 100644 --- a/flash_kda/__init__.py +++ b/flash_kda/__init__.py @@ -2,7 +2,45 @@ from flash_kda_C import fwd as _fwd_raw, get_workspace_size -def fwd(q, k, v, g, beta, scale, out, A_log, dt_bias, lower_bound, initial_state=None, final_state=None, cu_seqlens=None): +def get_intermediate_state_shape(q, cu_seqlens=None): + """Return the required ``[H, tile_capacity, D, D]`` snapshot-buffer shape. + + In ragged mode, ``tile_capacity`` is an upper bound with at most one unused + tile per sequence. Use ``get_intermediate_state_tile_prefix`` to locate the + compact valid chunk states without synchronizing to the host. + """ + B, T_seq, H, D = q.shape + if cu_seqlens is None: + total_tiles = B * ((T_seq + 15) // 16) + else: + total_tiles = (B * T_seq + 15) // 16 + cu_seqlens.numel() - 1 + return H, total_tiles, D, D + + +def get_intermediate_state_tile_prefix(q, cu_seqlens=None): + """Return device-resident offsets for the compact chunk-state layout. + + The state after chunk ``t`` of sequence ``n`` is at + ``intermediate_state[head, tile_prefix[n] + t]``. + ``tile_prefix[-1]`` is the number of valid snapshot slots, so this helper + locates ragged snapshots without host synchronization. Snapshots are useful + for inspection or a serial handoff, but cannot alone implement exact CP + prefix composition: KDA's input-state dependence requires the affine + transition prototype in flash_kda.cp. + """ + B, T_seq = q.shape[:2] + if cu_seqlens is None: + tiles_per_sequence = (T_seq + 15) // 16 + return torch.arange(B + 1, dtype=torch.long, device=q.device) * tiles_per_sequence + if cu_seqlens.device != q.device: + raise ValueError("cu_seqlens must be on the same device as q") + sequence_tiles = torch.div(cu_seqlens[1:] - cu_seqlens[:-1] + 15, 16, + rounding_mode="floor") + return torch.cat((torch.zeros(1, dtype=torch.long, device=cu_seqlens.device), + torch.cumsum(sequence_tiles, dim=0))) + + +def fwd(q, k, v, g, beta, scale, out, A_log, dt_bias, lower_bound, initial_state=None, final_state=None, cu_seqlens=None, intermediate_state=None): """FlashKDA forward (Flash Kimi Delta Attention). Args: @@ -25,6 +63,12 @@ def fwd(q, k, v, g, beta, scale, out, A_log, dt_bias, lower_bound, initial_state recurrent state. Same dtype/shape rules as ``initial_state``. cu_seqlens (torch.Tensor, optional): Cumulative sequence lengths, int64, shape ``[N+1]``. When provided, ``B`` must be 1. + intermediate_state (torch.Tensor, optional): BF16 output buffer for the + recurrent state after each 16-token chunk. Its shape must equal + ``get_intermediate_state_shape(q, cu_seqlens)``. In ragged mode, use + ``get_intermediate_state_tile_prefix(q, cu_seqlens)``: chunk ``t`` + of sequence ``n`` is stored at ``[head, tile_prefix[n] + t]``; + slots from ``tile_prefix[-1]`` onward are left untouched. Notes: * Currently requires ``K = V = 128``. @@ -38,4 +82,5 @@ def fwd(q, k, v, g, beta, scale, out, A_log, dt_bias, lower_bound, initial_state 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) + initial_state=initial_state, final_state=final_state, cu_seqlens=cu_seqlens, + intermediate_state=intermediate_state) diff --git a/flash_kda/cp.py b/flash_kda/cp.py new file mode 100644 index 0000000..8c79abd --- /dev/null +++ b/flash_kda/cp.py @@ -0,0 +1,590 @@ +"""Exact PyTorch transition algebra for chunked KDA context parallelism. + +This correctness prototype uses the kernel's post-preprocessing space: key is +already L2-normalized, log_decay is the A_log/dt_bias gate in log2 units, and +beta is already sigmoid-activated. A chunk maps S [H, V, K] as S @ A + B. +The affine map, unlike a boundary-state snapshot, supports an exact CP scan. + +"Exact" here means the high-precision affine recurrence. Materializing the +state as BF16 after every chunk introduces a nonlinear rounding step, so the +legacy bitwise path requires a serial fallback rather than a fixed (A, B) map. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Sequence + +import torch +import torch.nn.functional as F + +try: + import triton + import triton.language as tl +except ImportError: # pragma: no cover - exercised only by CPU-only PyTorch builds. + triton = None + tl = None + + +KDA_CHUNK_SIZE = 16 + + +if triton is not None: + + @triton.jit + def _kda_segment_accumulate_kernel( + key_decayed_ptr, + right_factor_ptr, + total_decay_ptr, + value_ptr, + matrix_ptr, + bias_ptr, + num_chunks, + HEADS: tl.constexpr, + KEY_DIM: tl.constexpr, + VALUE_DIM: tl.constexpr, + CHUNK: tl.constexpr, + BLOCK_K: tl.constexpr, + ): + """Keep one A/B row resident while applying every low-rank chunk.""" + program = tl.program_id(0) + rows_per_head = KEY_DIM + VALUE_DIM + head = program // rows_per_head + row = program - head * rows_per_head + offsets = tl.arange(0, BLOCK_K) + is_matrix = row < KEY_DIM + state = tl.where(is_matrix & (offsets == row), 1.0, 0.0).to(tl.float32) + value_row = row - KEY_DIM + + for chunk in range(0, num_chunks): + decay_offsets = (chunk * HEADS + head) * KEY_DIM + offsets + decay = tl.load(total_decay_ptr + decay_offsets) + update = tl.zeros((BLOCK_K,), dtype=tl.float32) + for token in tl.static_range(0, CHUNK): + factor_offsets = ( + ((chunk * HEADS + head) * CHUNK + token) * KEY_DIM + offsets + ) + key_decayed = tl.load(key_decayed_ptr + factor_offsets) + coefficient = tl.sum(state * key_decayed, axis=0) + value_offset = ( + ((chunk * HEADS + head) * CHUNK + token) * VALUE_DIM + + value_row + ) + value = tl.load( + value_ptr + value_offset, mask=~is_matrix, other=0.0 + ) + right_factor = tl.load(right_factor_ptr + factor_offsets) + update += (value - coefficient) * right_factor + state = state * decay + update + + matrix_offsets = head * KEY_DIM * KEY_DIM + row * KEY_DIM + offsets + bias_offsets = head * VALUE_DIM * KEY_DIM + value_row * KEY_DIM + offsets + tl.store(matrix_ptr + matrix_offsets, state, mask=is_matrix) + tl.store(bias_ptr + bias_offsets, state, mask=~is_matrix) + + +@dataclass(frozen=True) +class KDATransition: + """Affine state map with shapes [..., H, K, K] and [..., H, V, K].""" + + matrix: torch.Tensor + bias: torch.Tensor + + +def _check_transition(transition: KDATransition) -> None: + if transition.matrix.ndim < 3 or transition.bias.ndim < 3: + raise ValueError("transition tensors must have at least 3 dimensions") + if transition.matrix.shape[:-3] != transition.bias.shape[:-3]: + raise ValueError("transition matrix and bias must have identical batch shapes") + heads, key_dim, key_dim_2 = transition.matrix.shape[-3:] + if key_dim != key_dim_2: + raise ValueError("transition matrix must have shape [..., H, K, K]") + if transition.bias.shape[-3] != heads or transition.bias.shape[-1] != key_dim: + raise ValueError("transition bias must have shape [..., H, V, K]") + if transition.matrix.device != transition.bias.device: + raise ValueError("transition matrix and bias must share a device") + if transition.matrix.dtype != transition.bias.dtype: + raise ValueError("transition matrix and bias must share a dtype") + + +def identity_transition( + heads: int, + value_dim: int, + key_dim: int, + *, + device: torch.device | None = None, + dtype: torch.dtype = torch.float32, +) -> KDATransition: + """Return the identity map for a [H, V, K] recurrent state.""" + if heads <= 0 or value_dim <= 0 or key_dim <= 0: + raise ValueError("heads, value_dim, and key_dim must be positive") + matrix = torch.eye(key_dim, device=device, dtype=dtype).expand(heads, -1, -1).clone() + bias = torch.zeros((heads, value_dim, key_dim), device=device, dtype=dtype) + return KDATransition(matrix=matrix, bias=bias) + + +def apply_transition(state: torch.Tensor, transition: KDATransition) -> torch.Tensor: + """Apply an affine transition, broadcasting leading segment dimensions.""" + _check_transition(transition) + if state.ndim < 3: + raise ValueError("state must have shape [..., H, V, K]") + if state.shape[-3:] != transition.bias.shape[-3:]: + raise ValueError("state shape is incompatible with transition") + if state.device != transition.matrix.device: + raise ValueError("state and transition must share a device") + try: + torch.broadcast_shapes(state.shape[:-3], transition.matrix.shape[:-3]) + except RuntimeError as error: + raise ValueError("state and transition batch shapes are not broadcastable") from error + state = state.to(transition.matrix.dtype) + return state @ transition.matrix + transition.bias + + +def compose_transitions(first: KDATransition, second: KDATransition) -> KDATransition: + """Return the transition for applying first and then second. + + For first(S) = S @ A0 + B0 and second(S) = S @ A1 + B1, this returns + S @ (A0 @ A1) + (B0 @ A1 + B1). + """ + _check_transition(first) + _check_transition(second) + if first.matrix.shape != second.matrix.shape or first.bias.shape != second.bias.shape: + raise ValueError("transitions must have identical shapes") + if first.matrix.device != second.matrix.device or first.matrix.dtype != second.matrix.dtype: + raise ValueError("transitions must share device and dtype") + return KDATransition( + matrix=first.matrix @ second.matrix, + bias=first.bias @ second.matrix + second.bias, + ) + + +def compose_transition_sequence(transitions: Sequence[KDATransition]) -> KDATransition: + """Compose a non-empty ordered sequence of transitions.""" + if not transitions: + raise ValueError("cannot compose an empty transition sequence") + result = transitions[0] + for transition in transitions[1:]: + result = compose_transitions(result, transition) + return result + + +def stack_transitions(transitions: Sequence[KDATransition]) -> KDATransition: + """Stack equal-shaped per-segment transitions along a new leading axis.""" + if not transitions: + raise ValueError("cannot stack an empty transition sequence") + first = transitions[0] + _check_transition(first) + for transition in transitions[1:]: + _check_transition(transition) + if transition.matrix.shape != first.matrix.shape or transition.bias.shape != first.bias.shape: + raise ValueError("transitions must have identical shapes") + if transition.matrix.device != first.matrix.device or transition.matrix.dtype != first.matrix.dtype: + raise ValueError("transitions must share device and dtype") + return KDATransition( + matrix=torch.stack([transition.matrix for transition in transitions]), + bias=torch.stack([transition.bias for transition in transitions]), + ) + + +def exclusive_prefix_transitions( + transitions: Sequence[KDATransition], +) -> tuple[KDATransition, ...]: + """Serial reference for the start transition of every CP segment.""" + if not transitions: + return () + first = transitions[0] + _check_transition(first) + prefix = identity_transition( + first.matrix.shape[-3], + first.bias.shape[-2], + first.matrix.shape[-1], + device=first.matrix.device, + dtype=first.matrix.dtype, + ) + result = [] + for transition in transitions: + _check_transition(transition) + result.append(prefix) + prefix = compose_transitions(prefix, transition) + return tuple(result) + + +def exclusive_prefix_scan(transitions: Sequence[KDATransition]) -> KDATransition: + """GPU-resident O(log P) exclusive scan over P segment transitions. + + This Hillis-Steele reference launches batched matrix multiplications at each + tree level and supports any positive segment count, not only powers of two. + A fused implementation can retain the same composition law. + """ + stacked = stack_transitions(transitions) + matrix, bias = stacked.matrix, stacked.bias + segment_count = matrix.shape[0] + stride = 1 + while stride < segment_count: + right_matrix = matrix[stride:] + matrix = torch.cat( + (matrix[:stride], matrix[:-stride] @ right_matrix), dim=0 + ) + bias = torch.cat( + (bias[:stride], bias[:-stride] @ right_matrix + bias[stride:]), dim=0 + ) + stride *= 2 + + identity = identity_transition( + matrix.shape[-3], + bias.shape[-2], + matrix.shape[-1], + device=matrix.device, + dtype=matrix.dtype, + ) + return KDATransition( + matrix=torch.cat((identity.matrix.unsqueeze(0), matrix[:-1]), dim=0), + bias=torch.cat((identity.bias.unsqueeze(0), bias[:-1]), dim=0), + ) + + +def segment_start_states( + initial_state: torch.Tensor, + transitions: Sequence[KDATransition], +) -> torch.Tensor: + """Return [P, H, V, K] initial states for all P CP segments.""" + return apply_transition(initial_state, exclusive_prefix_scan(transitions)) + + +def distributed_exclusive_prefix_transition( + local_transition: KDATransition, + group=None, +) -> KDATransition: + """Compute this rank's exclusive CP prefix with a NCCL all-gather. + + This correctness baseline communicates one dense segment summary per rank, + then runs the GPU-resident tree scan locally. It has O(P) communication and + memory per rank; a production implementation should replace the all-gather + with a fused distributed tree while retaining the same composition law. + """ + import torch.distributed as dist + + _check_transition(local_transition) + if local_transition.matrix.ndim != 3 or local_transition.bias.ndim != 3: + raise ValueError("each rank must contribute one unbatched transition") + if not local_transition.matrix.is_cuda: + raise ValueError("distributed CP transitions must reside on CUDA") + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError("torch.distributed must be initialized") + + rank = dist.get_rank(group) + world_size = dist.get_world_size(group) + matrix = local_transition.matrix.contiguous() + bias = local_transition.bias.contiguous() + matrix_elements = matrix.numel() + payload = torch.cat((matrix.reshape(-1), bias.reshape(-1))) + gathered = torch.empty( + world_size * payload.numel(), device=payload.device, dtype=payload.dtype + ) + all_gather = getattr(dist, "all_gather_single", dist.all_gather_into_tensor) + all_gather(gathered, payload, group=group) + gathered = gathered.view(world_size, -1) + transitions = [ + KDATransition( + matrix=gathered[index, :matrix_elements].view_as(matrix), + bias=gathered[index, matrix_elements:].view_as(bias), + ) + for index in range(world_size) + ] + prefixes = exclusive_prefix_scan(transitions) + return KDATransition(matrix=prefixes.matrix[rank], bias=prefixes.bias[rank]) + + +def _chunk_factors( + key: torch.Tensor, + log_decay: torch.Tensor, + beta: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Build the chunk-local KDA factors in a high-precision reference space.""" + if key.ndim != 3: + raise ValueError("key and log_decay must have shape [T, H, K]") + if log_decay.shape != key.shape: + raise ValueError("log_decay must have the same shape as key") + if beta.shape != key.shape[:2]: + raise ValueError("beta must have shape [T, H]") + if not key.is_floating_point() or not log_decay.is_floating_point() or not beta.is_floating_point(): + raise ValueError("key, log_decay, and beta must be floating point") + if key.device != log_decay.device or key.device != beta.device: + raise ValueError("key, log_decay, and beta must share a device") + if key.shape[0] == 0 or key.shape[0] > KDA_CHUNK_SIZE: + raise ValueError(f"chunk length must be in [1, {KDA_CHUNK_SIZE}]") + + dtype = torch.promote_types(torch.promote_types(key.dtype, log_decay.dtype), beta.dtype) + if dtype in (torch.float16, torch.bfloat16): + dtype = torch.float32 + key_h = key.transpose(0, 1).to(dtype) + decay_h = log_decay.transpose(0, 1).to(dtype) + beta_h = beta.transpose(0, 1).to(dtype) + cumulative_decay = torch.cumsum(decay_h, dim=1) + decay_prefix = torch.exp2(cumulative_decay) + key_decayed = key_h * decay_prefix + key_inverse = key_h * torch.exp2(-cumulative_decay) + total_decay = decay_prefix[:, -1, :] + key_restored = key_inverse * total_decay.unsqueeze(1) + + gram = key_decayed @ key_inverse.transpose(-1, -2) + lower = torch.tril(gram, diagonal=-1) * beta_h.unsqueeze(-1) + chunk_len = key.shape[0] + eye = torch.eye(chunk_len, dtype=dtype, device=key.device).expand(key.shape[1], -1, -1) + inverse = eye.clone() + power = eye + for _ in range(1, chunk_len): + power = power @ lower + inverse = inverse + power + return key_decayed, key_restored, total_decay, inverse, beta_h + + +def kda_chunk_transition( + key: torch.Tensor, + value: torch.Tensor, + log_decay: torch.Tensor, + beta: torch.Tensor, +) -> KDATransition: + """Return the exact affine state map for one KDA chunk (T <= 16).""" + if value.ndim != 3 or value.shape[:2] != key.shape[:2]: + raise ValueError("value must have shape [T, H, V] matching key") + if value.device != key.device or not value.is_floating_point(): + raise ValueError("value must be floating point and share key's device") + + key_decayed, key_restored, total_decay, inverse, beta_h = _chunk_factors( + key, log_decay, beta + ) + value_h = value.transpose(0, 1).to(inverse.dtype) + # P = (I - L)^-1 diag(beta), which multiplies the chunk residual. + residual_operator = inverse * beta_h.unsqueeze(-2) + right_factor = residual_operator.transpose(-1, -2) @ key_restored + return KDATransition( + matrix=torch.diag_embed(total_decay) - key_decayed.transpose(-1, -2) @ right_factor, + bias=value_h.transpose(-1, -2) @ right_factor, + ) + + +def _segment_factors_batched( + key: torch.Tensor, + value: torch.Tensor, + log_decay: torch.Tensor, + beta: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Build all rank-16 chunk factors with a fixed number of GPU launches.""" + tokens, heads, key_dim = key.shape + value_dim = value.shape[-1] + padding = (-tokens) % KDA_CHUNK_SIZE + if padding: + key = F.pad(key, (0, 0, 0, 0, 0, padding)) + value = F.pad(value, (0, 0, 0, 0, 0, padding)) + log_decay = F.pad(log_decay, (0, 0, 0, 0, 0, padding)) + beta = F.pad(beta, (0, 0, 0, padding)) + + chunks = key.shape[0] // KDA_CHUNK_SIZE + dtype = torch.promote_types( + torch.promote_types(key.dtype, log_decay.dtype), beta.dtype + ) + if dtype in (torch.float16, torch.bfloat16): + dtype = torch.float32 + key_h = ( + key.reshape(chunks, KDA_CHUNK_SIZE, heads, key_dim) + .permute(0, 2, 1, 3) + .to(dtype) + ) + value_h = ( + value.reshape(chunks, KDA_CHUNK_SIZE, heads, value_dim) + .permute(0, 2, 1, 3) + .to(dtype) + ) + decay_h = ( + log_decay.reshape(chunks, KDA_CHUNK_SIZE, heads, key_dim) + .permute(0, 2, 1, 3) + .to(dtype) + ) + beta_h = ( + beta.reshape(chunks, KDA_CHUNK_SIZE, heads) + .permute(0, 2, 1) + .to(dtype) + ) + cumulative_decay = torch.cumsum(decay_h, dim=2) + decay_prefix = torch.exp2(cumulative_decay) + key_decayed = key_h * decay_prefix + key_inverse = key_h * torch.exp2(-cumulative_decay) + total_decay = decay_prefix[:, :, -1, :] + key_restored = key_inverse * total_decay.unsqueeze(2) + + gram = key_decayed @ key_inverse.transpose(-1, -2) + lower = torch.tril(gram, diagonal=-1) * beta_h.unsqueeze(-1) + eye = torch.eye( + KDA_CHUNK_SIZE, dtype=dtype, device=key.device + ).view(1, 1, KDA_CHUNK_SIZE, KDA_CHUNK_SIZE) + inverse = eye.expand(chunks, heads, -1, -1).clone() + power = eye.expand(chunks, heads, -1, -1) + for _ in range(1, KDA_CHUNK_SIZE): + power = power @ lower + inverse = inverse + power + residual_operator = inverse * beta_h.unsqueeze(-2) + right_factor = residual_operator.transpose(-1, -2) @ key_restored + return ( + key_decayed.contiguous(), + right_factor.contiguous(), + total_decay.contiguous(), + value_h.contiguous(), + ) + + +def _validate_segment_inputs( + key: torch.Tensor, + value: torch.Tensor, + log_decay: torch.Tensor, + beta: torch.Tensor, +) -> None: + if key.ndim != 3 or key.shape[0] == 0: + raise ValueError("key must have non-empty shape [T, H, K]") + if value.ndim != 3 or value.shape[:2] != key.shape[:2]: + raise ValueError("value must have shape [T, H, V] matching key") + if log_decay.shape != key.shape or beta.shape != key.shape[:2]: + raise ValueError("log_decay and beta shapes must match key") + if value.device != key.device or log_decay.device != key.device or beta.device != key.device: + raise ValueError("all KDA inputs must share a device") + if not all( + tensor.is_floating_point() for tensor in (key, value, log_decay, beta) + ): + raise ValueError("all KDA inputs must be floating point") + + +def kda_segment_transition_reference( + key: torch.Tensor, + value: torch.Tensor, + log_decay: torch.Tensor, + beta: torch.Tensor, +) -> KDATransition: + """PyTorch reference using each chunk's rank-at-most-16 update. + + Unlike composing a dense transition after every chunk, this costs + O((K + V) * K * C) per chunk for C <= 16 instead of O(K^3 + V*K^2). + The final dense (A, B) pair is the segment summary communicated by CP. + """ + _validate_segment_inputs(key, value, log_decay, beta) + + first_end = min(KDA_CHUNK_SIZE, key.shape[0]) + first_factors = _chunk_factors( + key[:first_end], log_decay[:first_end], beta[:first_end] + ) + dtype = first_factors[3].dtype + transition = identity_transition( + key.shape[1], value.shape[2], key.shape[2], device=key.device, dtype=dtype + ) + matrix, bias = transition.matrix, transition.bias + + for start in range(0, key.shape[0], KDA_CHUNK_SIZE): + end = min(start + KDA_CHUNK_SIZE, key.shape[0]) + key_decayed, key_restored, total_decay, inverse, beta_h = _chunk_factors( + key[start:end], log_decay[start:end], beta[start:end] + ) + value_h = value[start:end].transpose(0, 1).to(dtype) + residual_operator = inverse * beta_h.unsqueeze(-2) + right_factor = residual_operator.transpose(-1, -2) @ key_restored + key_decayed_t = key_decayed.transpose(-1, -2) + + # Post-multiply the accumulated dense transition by + # diag(total_decay) - key_decayed.T @ right_factor without forming the + # chunk's dense KxK matrix first. + matrix = ( + matrix * total_decay.unsqueeze(-2) + - (matrix @ key_decayed_t) @ right_factor + ) + bias = ( + bias * total_decay.unsqueeze(-2) + - (bias @ key_decayed_t) @ right_factor + + value_h.transpose(-1, -2) @ right_factor + ) + return KDATransition(matrix=matrix, bias=bias) + + +def kda_segment_transition_triton( + key: torch.Tensor, + value: torch.Tensor, + log_decay: torch.Tensor, + beta: torch.Tensor, +) -> KDATransition: + """Fused CUDA D=V=128 segment transition with a register-resident row.""" + _validate_segment_inputs(key, value, log_decay, beta) + if triton is None: + raise RuntimeError("Triton is required for the fused segment transition") + if not key.is_cuda: + raise ValueError("the fused segment transition requires CUDA tensors") + if key.shape[-1] != 128 or value.shape[-1] != 128: + raise ValueError("the fused segment transition currently requires K = V = 128") + + key_decayed, right_factor, total_decay, value_h = _segment_factors_batched( + key, value, log_decay, beta + ) + chunks, heads, _, key_dim = key_decayed.shape + value_dim = value_h.shape[-1] + matrix = torch.empty( + (heads, key_dim, key_dim), device=key.device, dtype=torch.float32 + ) + bias = torch.empty( + (heads, value_dim, key_dim), device=key.device, dtype=torch.float32 + ) + _kda_segment_accumulate_kernel[(heads * (key_dim + value_dim),)]( + key_decayed, + right_factor, + total_decay, + value_h, + matrix, + bias, + chunks, + HEADS=heads, + KEY_DIM=key_dim, + VALUE_DIM=value_dim, + CHUNK=KDA_CHUNK_SIZE, + BLOCK_K=triton.next_power_of_2(key_dim), + num_warps=4, + ) + return KDATransition(matrix=matrix, bias=bias) + + +def kda_segment_transition( + key: torch.Tensor, + value: torch.Tensor, + log_decay: torch.Tensor, + beta: torch.Tensor, +) -> KDATransition: + """Return a segment summary, dispatching the fused CUDA D=128 path.""" + if ( + triton is not None + and key.is_cuda + and key.ndim == 3 + and value.ndim == 3 + and key.shape[-1] == 128 + and value.shape[-1] == 128 + ): + return kda_segment_transition_triton(key, value, log_decay, beta) + return kda_segment_transition_reference(key, value, log_decay, beta) + + +def kda_chunk_update( + state: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + log_decay: torch.Tensor, + beta: torch.Tensor, +) -> torch.Tensor: + """Direct recurrence for checking kda_chunk_transition independently.""" + key_decayed, key_restored, total_decay, inverse, beta_h = _chunk_factors( + key, log_decay, beta + ) + if value.ndim != 3 or value.shape[:2] != key.shape[:2]: + raise ValueError("value must have shape [T, H, V] matching key") + if state.shape != (key.shape[1], value.shape[2], key.shape[2]): + raise ValueError("state must have shape [H, V, K] compatible with inputs") + + state_h = state.to(inverse.dtype) + value_h = value.transpose(0, 1).to(inverse.dtype) + residual = value_h - key_decayed @ state_h.transpose(-1, -2) + update = key_restored.transpose(-1, -2) @ ( + inverse @ (residual * beta_h.unsqueeze(-1)) + ) + return (update + total_decay.unsqueeze(-1) * state_h.transpose(-1, -2)).transpose(-1, -2) diff --git a/tests/distributed_cp_smoke.py b/tests/distributed_cp_smoke.py new file mode 100644 index 0000000..d7e62a0 --- /dev/null +++ b/tests/distributed_cp_smoke.py @@ -0,0 +1,196 @@ +"""Two-GPU NCCL and end-to-end FlashKDA smoke tests for CP transitions.""" + +import importlib.util +import math +import os +import sys +from pathlib import Path + +_REPO = Path(__file__).parents[1] +sys.path.insert(0, str(_REPO)) + +import flash_kda +import torch +import torch.distributed as dist +import torch.nn.functional as F +from torch_ref import LOG2E, fp32_ex2_ftz, l2_normalize_kernel_match, sigmoid_ext + + +_CP_PATH = _REPO / "flash_kda" / "cp.py" +_SPEC = importlib.util.spec_from_file_location("flash_kda_cp_distributed", _CP_PATH) +assert _SPEC is not None and _SPEC.loader is not None +cp = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = cp +_SPEC.loader.exec_module(cp) + + +def make_segment(seed: int, tokens: int, device: torch.device): + generator = torch.Generator().manual_seed(seed) + heads, key_dim, value_dim = 2, 128, 128 + key = F.normalize( + torch.randn((tokens, heads, key_dim), generator=generator), p=2, dim=-1 + ).to(device=device, dtype=torch.bfloat16) + value = torch.randn( + (tokens, heads, value_dim), generator=generator + ).to(device=device, dtype=torch.bfloat16) + log_decay = ( + -0.1 * torch.rand((tokens, heads, key_dim), generator=generator) + ).to(device) + beta = ( + 0.1 + 0.8 * torch.rand((tokens, heads), generator=generator) + ).to(device) + return key, value, log_decay, beta + + +def make_raw_inputs(seed: int, tokens: int, heads: int, dim: int, device): + generator = torch.Generator().manual_seed(seed) + shape = (1, tokens, heads, dim) + tensors = ( + torch.randn(shape, generator=generator).to(torch.bfloat16), + torch.randn(shape, generator=generator).to(torch.bfloat16), + torch.randn(shape, generator=generator).to(torch.bfloat16), + torch.randn(shape, generator=generator).to(torch.bfloat16), + torch.randn((1, tokens, heads), generator=generator).to(torch.bfloat16), + torch.rand((heads,), generator=generator), + torch.randn((heads, dim), generator=generator), + ) + return tuple(tensor.to(device) for tensor in tensors) + + +def slice_segment(inputs, start: int, end: int): + q, k, v, g, beta, a_log, dt_bias = inputs + return ( + q[:, start:end].contiguous(), + k[:, start:end].contiguous(), + v[:, start:end].contiguous(), + g[:, start:end].contiguous(), + beta[:, start:end].contiguous(), + a_log, + dt_bias, + ) + + +def transition_from_raw_inputs(inputs, lower_bound: float): + _, key, value, gate, beta, a_log, dt_bias = inputs + tokens, heads, dim = key.shape[1:] + key = l2_normalize_kernel_match(key.reshape(tokens, heads, dim)) + gate = gate.reshape(tokens, heads, dim).float() + dt_bias.unsqueeze(0) + a_log_exp = fp32_ex2_ftz(a_log * LOG2E).view(1, heads, 1) + log_decay = lower_bound * LOG2E * sigmoid_ext.sigmoid_tanh_fp32( + a_log_exp * gate + ) + beta = sigmoid_ext.sigmoid_tanh_fp32(beta.reshape(tokens, heads).float()) + return cp.kda_segment_transition( + key, value.reshape(tokens, heads, dim), log_decay, beta + ) + + +def run_flash_kda(inputs, initial_state, lower_bound: float, snapshots=None): + q, k, v, g, beta, a_log, dt_bias = inputs + output = torch.empty_like(q) + final_state = torch.empty_like(initial_state) + flash_kda.fwd( + q, + k, + v, + g, + beta, + 1.0 / math.sqrt(q.shape[-1]), + output, + a_log, + dt_bias, + lower_bound, + initial_state=initial_state, + final_state=final_state, + intermediate_state=snapshots, + ) + return output, final_state + + +def check_flash_kda_end_to_end(rank: int, device: torch.device) -> None: + """Use the distributed prefix as the real FlashKDA initial state.""" + heads, dim, tokens_per_rank = 2, 128, 128 + lower_bound = -5.0 + global_inputs = make_raw_inputs( + 20260821, 2 * tokens_per_rank, heads, dim, device + ) + start = rank * tokens_per_rank + end = start + tokens_per_rank + local_inputs = slice_segment(global_inputs, start, end) + + local_transition = transition_from_raw_inputs(local_inputs, lower_bound) + prefix = cp.distributed_exclusive_prefix_transition(local_transition) + zero_state = torch.zeros((heads, dim, dim), device=device, dtype=torch.float32) + cp_start = cp.apply_transition(zero_state, prefix) + cp_output, _ = run_flash_kda( + local_inputs, cp_start.unsqueeze(0), lower_bound + ) + + snapshots = torch.empty( + flash_kda.get_intermediate_state_shape(global_inputs[0]), + device=device, + dtype=torch.bfloat16, + ) + serial_output, _ = run_flash_kda( + global_inputs, zero_state.unsqueeze(0), lower_bound, snapshots + ) + torch.cuda.synchronize(device) + expected_output = serial_output[:, start:end] + expected_start = ( + zero_state + if rank == 0 + else snapshots[:, tokens_per_rank // cp.KDA_CHUNK_SIZE - 1].float() + ) + + if rank == 0: + torch.testing.assert_close(cp_start, expected_start, rtol=0, atol=0) + torch.testing.assert_close(cp_output, expected_output, rtol=0, atol=0) + else: + # The affine scan stays in FP32, while the legacy serial kernel rounds + # its recurrent state to BF16 after each chunk. + torch.testing.assert_close( + cp_start, expected_start, rtol=5e-2, atol=3e-2 + ) + torch.testing.assert_close( + cp_output, expected_output, rtol=5e-3, atol=1e-3 + ) + + +def main() -> None: + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + dist.init_process_group("nccl", device_id=device) + rank = dist.get_rank() + world_size = dist.get_world_size() + if world_size != 2: + raise RuntimeError("this smoke test expects exactly two ranks") + + # Build the independent oracle before the first NCCL operation. CUDA work + # after rank-divergent P2P can otherwise be ordered behind a collective on + # the peer and turn a test-only synchronization pattern into a deadlock. + if rank == 0: + expected = cp.identity_transition( + 2, 128, 128, device=device, dtype=torch.float32 + ) + else: + expected = cp.kda_segment_transition(*make_segment(20260820, 32, device)) + local = cp.kda_segment_transition(*make_segment(20260820 + rank, 32, device)) + prefix = cp.distributed_exclusive_prefix_transition(local) + + torch.testing.assert_close(prefix.matrix, expected.matrix, rtol=2e-5, atol=2e-5) + torch.testing.assert_close(prefix.bias, expected.bias, rtol=2e-5, atol=2e-5) + + initial_state = torch.zeros((2, 128, 128), device=device, dtype=torch.bfloat16) + actual_start = cp.apply_transition(initial_state, prefix) + expected_start = cp.apply_transition(initial_state, expected) + torch.testing.assert_close(actual_start, expected_start, rtol=2e-5, atol=2e-5) + check_flash_kda_end_to_end(rank, device) + dist.barrier() + if rank == 0: + print("distributed CP + FlashKDA end-to-end smoke test passed on 2 GPUs") + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/test_cp_transition.py b/tests/test_cp_transition.py new file mode 100644 index 0000000..632cfa6 --- /dev/null +++ b/tests/test_cp_transition.py @@ -0,0 +1,258 @@ +"""Mathematical validation for the exact KDA CP transition prototype.""" + +import importlib.util +import sys +from pathlib import Path + +import torch +import torch.nn.functional as F + + +_CP_PATH = Path(__file__).parents[1] / "flash_kda" / "cp.py" +_SPEC = importlib.util.spec_from_file_location("flash_kda_cp", _CP_PATH) +assert _SPEC is not None and _SPEC.loader is not None +cp = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = cp +_SPEC.loader.exec_module(cp) + + +def _chunk( + generator: torch.Generator, + tokens: int, + heads: int, + key_dim: int, + value_dim: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + key = F.normalize( + torch.randn((tokens, heads, key_dim), generator=generator, dtype=torch.float64), + p=2, + dim=-1, + ) + value = 0.1 * torch.randn( + (tokens, heads, value_dim), generator=generator, dtype=torch.float64 + ) + log_decay = -0.1 * torch.rand( + (tokens, heads, key_dim), generator=generator, dtype=torch.float64 + ) + beta = 0.1 + 0.8 * torch.rand( + (tokens, heads), generator=generator, dtype=torch.float64 + ) + return key, value, log_decay, beta + + +def _concatenate_chunks( + chunks: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]], +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + return tuple(torch.cat([chunk[index] for chunk in chunks], dim=0) for index in range(4)) + + +def _assert_close(actual: torch.Tensor, expected: torch.Tensor) -> None: + torch.testing.assert_close(actual, expected, rtol=1e-11, atol=1e-11) + + +def test_transition_matches_direct_kda_block_recurrence(): + generator = torch.Generator().manual_seed(20260814) + heads, key_dim, value_dim = 2, 9, 7 + state = 0.1 * torch.randn( + (heads, value_dim, key_dim), generator=generator, dtype=torch.float64 + ) + chunk = _chunk(generator, 13, heads, key_dim, value_dim) + + transition = cp.kda_chunk_transition(*chunk) + _assert_close(cp.apply_transition(state, transition), cp.kda_chunk_update(state, *chunk)) + + +def test_transition_composition_is_associative(): + generator = torch.Generator().manual_seed(20260815) + chunks = [_chunk(generator, length, 2, 8, 6) for length in (16, 7, 11)] + transitions = [cp.kda_chunk_transition(*chunk) for chunk in chunks] + + left = cp.compose_transitions( + cp.compose_transitions(transitions[0], transitions[1]), transitions[2] + ) + right = cp.compose_transitions( + transitions[0], cp.compose_transitions(transitions[1], transitions[2]) + ) + + _assert_close(left.matrix, right.matrix) + _assert_close(left.bias, right.bias) + + +def test_segment_low_rank_accumulation_matches_dense_chunk_composition(): + generator = torch.Generator().manual_seed(20260816) + chunks = [_chunk(generator, length, 2, 10, 7) for length in (16, 16, 16, 3)] + dense = cp.compose_transition_sequence( + [cp.kda_chunk_transition(*chunk) for chunk in chunks] + ) + low_rank = cp.kda_segment_transition(*_concatenate_chunks(chunks)) + + _assert_close(low_rank.matrix, dense.matrix) + _assert_close(low_rank.bias, dense.bias) + + +def test_two_rank_cp_exclusive_prefix_matches_serial_recurrence(): + """Two logical CP ranks receive the exact start state for their segment.""" + generator = torch.Generator().manual_seed(20260817) + heads, key_dim, value_dim = 2, 10, 5 + initial_state = 0.1 * torch.randn( + (heads, value_dim, key_dim), generator=generator, dtype=torch.float64 + ) + chunks = [ + _chunk(generator, length, heads, key_dim, value_dim) + for length in (16, 9, 16, 5) + ] + transitions = [cp.kda_chunk_transition(*chunk) for chunk in chunks] + + # Rank 0 owns chunks [0, 1]; rank 1 owns chunks [2, 3]. + local_transitions = [ + cp.compose_transition_sequence(transitions[:2]), + cp.compose_transition_sequence(transitions[2:]), + ] + prefixes = cp.exclusive_prefix_transitions(local_transitions) + rank0_start = cp.apply_transition(initial_state, prefixes[0]) + rank1_start = cp.apply_transition(initial_state, prefixes[1]) + + serial = initial_state + serial_after_rank0 = initial_state + for index, chunk in enumerate(chunks): + serial = cp.kda_chunk_update(serial, *chunk) + if index == 1: + serial_after_rank0 = serial + + rank0_end = cp.apply_transition(rank0_start, local_transitions[0]) + rank1_end = cp.apply_transition(rank1_start, local_transitions[1]) + full_transition = cp.compose_transition_sequence(transitions) + + _assert_close(rank0_end, serial_after_rank0) + _assert_close(rank1_start, serial_after_rank0) + _assert_close(rank1_end, serial) + _assert_close(cp.apply_transition(initial_state, full_transition), serial) + + +def test_tree_prefix_scan_matches_serial_for_non_power_of_two(): + generator = torch.Generator().manual_seed(20260818) + heads, key_dim, value_dim = 2, 9, 6 + initial_state = 0.1 * torch.randn( + (heads, value_dim, key_dim), generator=generator, dtype=torch.float64 + ) + segments = [ + _concatenate_chunks( + [_chunk(generator, length, heads, key_dim, value_dim) for length in lengths] + ) + for lengths in ((16,), (5, 13), (9,), (16, 2), (7,)) + ] + transitions = [cp.kda_segment_transition(*segment) for segment in segments] + serial_prefixes = cp.exclusive_prefix_transitions(transitions) + tree_prefixes = cp.exclusive_prefix_scan(transitions) + + assert tree_prefixes.matrix.shape[0] == 5 + for index, serial_prefix in enumerate(serial_prefixes): + _assert_close(tree_prefixes.matrix[index], serial_prefix.matrix) + _assert_close(tree_prefixes.bias[index], serial_prefix.bias) + + starts = cp.segment_start_states(initial_state, transitions) + for index, serial_prefix in enumerate(serial_prefixes): + _assert_close(starts[index], cp.apply_transition(initial_state, serial_prefix)) + + +def test_cuda_triton_segment_matches_reference_with_tail_chunk(): + if not torch.cuda.is_available(): + return + + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(20260819) + tokens, heads, key_dim, value_dim = 35, 2, 128, 128 + key = F.normalize( + torch.randn( + (tokens, heads, key_dim), + generator=generator, + device=device, + dtype=torch.float32, + ), + p=2, + dim=-1, + ).to(torch.bfloat16) + value = torch.randn( + (tokens, heads, value_dim), + generator=generator, + device=device, + dtype=torch.bfloat16, + ) + log_decay = -0.1 * torch.rand( + (tokens, heads, key_dim), generator=generator, device=device + ) + beta = 0.1 + 0.8 * torch.rand( + (tokens, heads), generator=generator, device=device + ) + + expected = cp.kda_segment_transition_reference(key, value, log_decay, beta) + actual = cp.kda_segment_transition_triton(key, value, log_decay, beta) + torch.testing.assert_close(actual.matrix, expected.matrix, rtol=3e-3, atol=3e-4) + torch.testing.assert_close(actual.bias, expected.bias, rtol=3e-3, atol=3e-4) + + +def test_cuda_bf16_segment_scan_reconstructs_serial_state(): + """Exercise segment accumulation and a 3-rank scan at FlashKDA D=128.""" + if not torch.cuda.is_available(): + return + + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(20260819) + heads, key_dim, value_dim = 4, 128, 128 + initial_state = ( + 0.05 + * torch.randn( + (heads, value_dim, key_dim), + generator=generator, + device=device, + dtype=torch.float32, + ) + ).to(torch.bfloat16) + chunks = [] + for length in (16, 16, 9, 16): + key = F.normalize( + torch.randn( + (length, heads, key_dim), + generator=generator, + device=device, + dtype=torch.float32, + ), + p=2, + dim=-1, + ).to(torch.bfloat16) + value = torch.randn( + (length, heads, value_dim), + generator=generator, + device=device, + dtype=torch.bfloat16, + ) + log_decay = -0.1 * torch.rand( + (length, heads, key_dim), + generator=generator, + device=device, + dtype=torch.float32, + ) + beta = 0.1 + 0.8 * torch.rand( + (length, heads), generator=generator, device=device, dtype=torch.float32 + ) + chunks.append((key, value, log_decay, beta)) + + segment_chunks = (chunks[:1], chunks[1:3], chunks[3:]) + transitions = [ + cp.kda_segment_transition(*_concatenate_chunks(segment)) + for segment in segment_chunks + ] + starts = cp.segment_start_states(initial_state, transitions) + + serial = initial_state + expected_starts = [] + for segment in segment_chunks: + expected_starts.append(serial.to(torch.float32)) + for chunk in segment: + serial = cp.kda_chunk_update(serial, *chunk) + + for actual, expected in zip(starts, expected_starts): + torch.testing.assert_close(actual, expected, rtol=5e-4, atol=5e-4) + + final_from_scan = cp.apply_transition(starts[-1], transitions[-1]) + torch.testing.assert_close(final_from_scan, serial, rtol=5e-4, atol=5e-4) diff --git a/tests/test_fwd.py b/tests/test_fwd.py index 144bedd..e4a3fe9 100644 --- a/tests/test_fwd.py +++ b/tests/test_fwd.py @@ -424,6 +424,140 @@ def test_fwd_varlen_vs_fla(): print("Assert results: Success") +def test_fwd_intermediate_state_varlen(): + """Each ragged chunk snapshot must exactly match the reference recurrence.""" + H, D = 2, 128 + LOWER_BOUND = -5.0 + seq_lens = [17, 33, 65] + T_total = sum(seq_lens) + N = len(seq_lens) + valid_tiles = sum((seq_len + 15) // 16 for seq_len in seq_lens) + cu_seqlens = torch.tensor([0, 17, 50, 115], dtype=torch.long, device="cuda") + + torch.manual_seed(17) + shape = (1, T_total, H, D) + q = F.normalize(torch.randn(shape, dtype=torch.float32, device="cuda"), p=2, dim=-1).to(torch.bfloat16) + k = F.normalize(torch.randn(shape, dtype=torch.float32, device="cuda"), p=2, dim=-1).to(torch.bfloat16) + v = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + g = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + beta = torch.randn((1, T_total, H), dtype=torch.bfloat16, device="cuda") + A_log = torch.rand(H, dtype=torch.float32, device="cuda") + dt_bias = torch.rand(H, D, dtype=torch.float32, device="cuda") + initial_state = torch.randn((N, H, D, D), dtype=torch.bfloat16, device="cuda") + scale = 1.0 / math.sqrt(D) + + snapshot_shape = flash_kda.get_intermediate_state_shape(q, cu_seqlens) + tile_prefix = flash_kda.get_intermediate_state_tile_prefix(q, cu_seqlens) + assert snapshot_shape == (H, 11, D, D) + assert torch.equal(tile_prefix, torch.tensor([0, 2, 5, 10], device="cuda")) + intermediate_state = torch.full(snapshot_shape, float("nan"), dtype=torch.bfloat16, device="cuda") + intermediate_state_ref = torch.full_like(intermediate_state, float("nan")) + out_kernel = torch.zeros_like(q) + out_ref = torch.zeros_like(q) + + flash_kda.fwd(q, k, v, g, beta, scale, out_kernel, + A_log=A_log, dt_bias=dt_bias, lower_bound=LOWER_BOUND, + initial_state=initial_state.clone(), cu_seqlens=cu_seqlens, + intermediate_state=intermediate_state) + torch.cuda.synchronize() + torch_ref(q, k, v, g, beta, scale, out_ref, + A_log=A_log, dt_bias=dt_bias, lower_bound=LOWER_BOUND, + initial_state=initial_state.clone(), cu_seqlens=cu_seqlens, + intermediate_state=intermediate_state_ref) + + # The unmodified SM90 kernel differs from the torch reference by one BF16 + # ULP on this ragged input, so allow that established output tolerance while + # retaining an exact check for the new chunk-state ABI below. + torch.testing.assert_close(out_kernel, out_ref, rtol=1e-2, atol=1e-7) + assert torch.equal(intermediate_state[:, :valid_tiles], + intermediate_state_ref[:, :valid_tiles]), "varlen chunk state mismatch" + assert torch.isnan(intermediate_state[:, valid_tiles:]).all(), "unused ragged capacity was modified" + + # Kernel-side raw indexing must reject strided CUDA views. + strided_source = torch.empty(cu_seqlens.numel() * 2 - 1, dtype=torch.long, device="cuda") + strided_source[::2] = cu_seqlens + strided_source[1::2] = 0 + noncontiguous_cu_seqlens = strided_source[::2] + assert not noncontiguous_cu_seqlens.is_contiguous() + try: + flash_kda.fwd(q, k, v, g, beta, scale, torch.zeros_like(q), + A_log=A_log, dt_bias=dt_bias, lower_bound=LOWER_BOUND, + initial_state=initial_state.clone(), + cu_seqlens=noncontiguous_cu_seqlens, + intermediate_state=torch.empty(snapshot_shape, dtype=torch.bfloat16, device="cuda")) + except RuntimeError as exc: + assert "cu_seqlens must be contiguous" in str(exc) + else: + raise AssertionError("non-contiguous cu_seqlens was accepted") + + +def test_fwd_intermediate_state_batched(): + """Snapshot dispatch also supports batched fixed-length sequences.""" + B, T, H, D = 2, 17, 2, 128 + LOWER_BOUND = -5.0 + torch.manual_seed(23) + shape = (B, T, H, D) + q = F.normalize(torch.randn(shape, dtype=torch.float32, device="cuda"), p=2, dim=-1).to(torch.bfloat16) + k = F.normalize(torch.randn(shape, dtype=torch.float32, device="cuda"), p=2, dim=-1).to(torch.bfloat16) + v = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + g = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + beta = torch.randn((B, T, H), dtype=torch.bfloat16, device="cuda") + A_log = torch.rand(H, dtype=torch.float32, device="cuda") + dt_bias = torch.rand(H, D, dtype=torch.float32, device="cuda") + initial_state = torch.randn((B, H, D, D), dtype=torch.bfloat16, device="cuda") + scale = 1.0 / math.sqrt(D) + + intermediate_state = torch.empty(flash_kda.get_intermediate_state_shape(q), + dtype=torch.bfloat16, device="cuda") + intermediate_state_ref = torch.empty_like(intermediate_state) + out_kernel = torch.zeros_like(q) + out_ref = torch.zeros_like(q) + + flash_kda.fwd(q, k, v, g, beta, scale, out_kernel, + A_log=A_log, dt_bias=dt_bias, lower_bound=LOWER_BOUND, + initial_state=initial_state.clone(), intermediate_state=intermediate_state) + torch.cuda.synchronize() + torch_ref(q, k, v, g, beta, scale, out_ref, + A_log=A_log, dt_bias=dt_bias, lower_bound=LOWER_BOUND, + initial_state=initial_state.clone(), intermediate_state=intermediate_state_ref) + + assert torch.equal(out_kernel, out_ref), "batched output mismatch" + assert torch.equal(intermediate_state, intermediate_state_ref), "batched chunk state mismatch" + + # Exercise the remaining HasIntermediateState dispatches permanently. + state_variants = [ + ("no-state", None, None), + ("bf16-in", torch.randn((B, H, D, D), dtype=torch.bfloat16, device="cuda"), None), + ("bf16-final", None, torch.empty((B, H, D, D), dtype=torch.bfloat16, device="cuda")), + ("bf16-in-out", torch.randn((B, H, D, D), dtype=torch.bfloat16, device="cuda"), + torch.empty((B, H, D, D), dtype=torch.bfloat16, device="cuda")), + ("fp32-in", torch.randn((B, H, D, D), dtype=torch.float32, device="cuda"), None), + ("fp32-final", None, torch.empty((B, H, D, D), dtype=torch.float32, device="cuda")), + ("fp32-in-out", torch.randn((B, H, D, D), dtype=torch.float32, device="cuda"), + torch.empty((B, H, D, D), dtype=torch.float32, device="cuda")), + ] + for label, state_in, state_out in state_variants: + final_kernel = torch.empty_like(state_out) if state_out is not None else None + final_ref = torch.empty_like(state_out) if state_out is not None else None + snapshot_kernel = torch.empty_like(intermediate_state) + snapshot_ref = torch.empty_like(intermediate_state) + out_kernel = torch.zeros_like(q) + out_ref = torch.zeros_like(q) + flash_kda.fwd(q, k, v, g, beta, scale, out_kernel, + A_log=A_log, dt_bias=dt_bias, lower_bound=LOWER_BOUND, + initial_state=state_in, final_state=final_kernel, + intermediate_state=snapshot_kernel) + torch.cuda.synchronize() + torch_ref(q, k, v, g, beta, scale, out_ref, + A_log=A_log, dt_bias=dt_bias, lower_bound=LOWER_BOUND, + initial_state=state_in, final_state=final_ref, + intermediate_state=snapshot_ref) + assert torch.equal(out_kernel, out_ref), f"{label} output mismatch" + assert torch.equal(snapshot_kernel, snapshot_ref), f"{label} chunk state mismatch" + if state_out is not None: + assert torch.equal(final_kernel, final_ref), f"{label} final state mismatch" + + if __name__ == "__main__": test_fwd() test_fwd_varlen() diff --git a/tests/torch_ref.py b/tests/torch_ref.py index 9be8e82..572d422 100644 --- a/tests/torch_ref.py +++ b/tests/torch_ref.py @@ -121,7 +121,7 @@ def l2_normalize_kernel_match(x): # Torch reference implementation # ============================================================ -def torch_ref(q, k, v, g, beta, scale, out, A_log, dt_bias, lower_bound, initial_state=None, final_state=None, cu_seqlens=None): +def torch_ref(q, k, v, g, beta, scale, out, A_log, dt_bias, lower_bound, initial_state=None, final_state=None, cu_seqlens=None, intermediate_state=None): """Torch reference, supports both fixed-length and variable-length sequences. Input: [B, T, H, D] (4D). B must be 1 when cu_seqlens is provided. @@ -177,6 +177,7 @@ def torch_ref(q, k, v, g, beta, scale, out, A_log, dt_bias, lower_bound, initial else: work_state = torch.zeros(N, H, D, D, dtype=torch.bfloat16, device=device) + tile_base = 0 for seq_idx in range(N): bos = cu_seqlens[seq_idx].item() eos = cu_seqlens[seq_idx + 1].item() @@ -243,6 +244,10 @@ def torch_ref(q, k, v, g, beta, scale, out, A_log, dt_bias, lower_bound, initial work_state[seq_idx, h] = fp32_fma(delta_s, state_slice.to(torch.float32).t(), g_total_exp).to(torch.bfloat16).t() out[t0:t0 + actual_len, h] = _out[:actual_len] + if intermediate_state is not None: + intermediate_state[h, tile_base + chunk_idx].copy_(work_state[seq_idx, h]) + + tile_base += n_chunks if final_state is not None: if state_fp32: