diff --git a/README.md b/README.md index 57d7fc2..b1f2cce 100644 --- a/README.md +++ b/README.md @@ -106,6 +106,7 @@ flash_kda.fwd(q, k, v, g, beta, scale, out, A_log, dt_bias, lower_bound, - Currently requires `K = V = 128`. - `initial_state` / `final_state` accept `None` (stateless), bf16, or fp32 tensors. When both are provided, their dtypes must match. - When `cu_seqlens` is provided, `B` must be 1, `T` is the total length across all sequences, and `initial_state` / `final_state` have shape `[N, H, V, K]`. +- `cu_seqlens` may contain repeated offsets for zero-length sequences. Such sequences have no output tokens; their `final_state` equals `initial_state` when provided, or zero otherwise. - When `cu_seqlens` is `None`, each batch element is treated as an independent sequence, and the state shape is `[B, H, V, K]`. ## Development diff --git a/csrc/flash_kda.cpp b/csrc/flash_kda.cpp index 81f5483..b23bc5f 100644 --- a/csrc/flash_kda.cpp +++ b/csrc/flash_kda.cpp @@ -109,6 +109,52 @@ void fwd( TORCH_CHECK(D == 128, "currently only supports D == 128"); + // Determine cu_seqlens and N before constructing any TMA descriptors. An + // all-empty varlen batch has T_total == 0, for which zero-extent TMA + // descriptors are invalid even though its recurrent-state semantics are + // well defined. + bool is_varlen = cu_seqlens.has_value(); + int64_t N_val; + int64_t const* cu_seqlens_dev = nullptr; + + if (is_varlen) { + 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.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; + TORCH_CHECK(N_val > 0, "cu_seqlens must have at least 2 elements"); + cu_seqlens_dev = cu_seqlens_t.data_ptr(); + } else { + N_val = B; + } + + // Validate state shapes: always [N, H, D, D] + if (has_state_in) { + auto& is = initial_state.value(); + TORCH_CHECK(is.dim() == 4, "initial_state must be [N, H, D, D]"); + TORCH_CHECK(is.size(0) == N_val && is.size(1) == H && is.size(2) == D && is.size(3) == D, + "initial_state must be [N, H, D, D]"); + } + if (has_state_out) { + auto& fs = final_state.value(); + TORCH_CHECK(fs.dim() == 4, "final_state must be [N, H, D, D]"); + TORCH_CHECK(fs.size(0) == N_val && fs.size(1) == H && fs.size(2) == D && fs.size(3) == D, + "final_state must be [N, H, D, D]"); + } + + if (T_total == 0) { + if (has_state_out) { + if (has_state_in) { + final_state->copy_(initial_state.value()); + } else { + final_state->zero_(); + } + } + return; + } + // Flatten [B, T, H, D] -> [B*T, H, D] (contiguous, same data pointer) auto q_3d = q.reshape({T_total, H, D}); auto k_3d = k.reshape({T_total, H, D}); @@ -141,38 +187,6 @@ void fwd( void const* initial_state_raw = has_state_in ? initial_state->data_ptr() : nullptr; void* final_state_raw = has_state_out ? final_state->data_ptr() : nullptr; - // Determine cu_seqlens and N - bool is_varlen = cu_seqlens.has_value(); - int64_t N_val; - int64_t const* cu_seqlens_dev = nullptr; - - if (is_varlen) { - 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.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; - TORCH_CHECK(N_val > 0, "cu_seqlens must have at least 2 elements"); - cu_seqlens_dev = cu_seqlens_t.data_ptr(); - } else { - N_val = B; - } - - // Validate state shapes: always [N, H, D, D] - if (has_state_in) { - auto& is = initial_state.value(); - TORCH_CHECK(is.dim() == 4, "initial_state must be [N, H, D, D]"); - TORCH_CHECK(is.size(0) == N_val && is.size(1) == H && is.size(2) == D && is.size(3) == D, - "initial_state must be [N, H, D, D]"); - } - if (has_state_out) { - auto& fs = final_state.value(); - TORCH_CHECK(fs.dim() == 4, "final_state must be [N, H, D, D]"); - TORCH_CHECK(fs.size(0) == N_val && fs.size(1) == H && fs.size(2) == D && fs.size(3) == D, - "final_state must be [N, H, D, D]"); - } - int total_tiles; if (is_varlen) { total_tiles = int((T_total + CHUNK - 1) / CHUNK + N_val); // upper bound for varlen diff --git a/csrc/smxx/fwd_kernel2.cuh b/csrc/smxx/fwd_kernel2.cuh index 26f73fe..dc8d144 100644 --- a/csrc/smxx/fwd_kernel2.cuh +++ b/csrc/smxx/fwd_kernel2.cuh @@ -147,6 +147,7 @@ __global__ void __launch_bounds__(NumThreads) _flash_kda_fwd_recurrence( int H, int N, int64_t const* cu_seqlens, + int const* tile_prefix, int total_tiles ) { using BF16 = cutlass::bfloat16_t; @@ -221,11 +222,7 @@ __global__ void __launch_bounds__(NumThreads) _flash_kda_fwd_recurrence( if constexpr (IsVarlen) { bos = cu_seqlens[seq_idx]; eos = cu_seqlens[seq_idx + 1]; - // Compute tile_base via linear scan (no host-precomputed table) - tile_base = 0; - for (int i = 0; i < seq_idx; i++) { - tile_base += (int(cu_seqlens[i + 1] - cu_seqlens[i]) + CHUNK - 1) / CHUNK; - } + tile_base = tile_prefix[seq_idx]; } else { int T_seq = T_total / N; bos = seq_idx * T_seq; diff --git a/csrc/smxx/fwd_launch.cu b/csrc/smxx/fwd_launch.cu index 91a67a7..fbfed7f 100644 --- a/csrc/smxx/fwd_launch.cu +++ b/csrc/smxx/fwd_launch.cu @@ -210,7 +210,7 @@ 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, T_total, H, N, cu_seqlens_ptr, ws_tile_prefix, total_tiles ); } #endif diff --git a/tests/test_empty_varlen.py b/tests/test_empty_varlen.py new file mode 100644 index 0000000..5fdf86b --- /dev/null +++ b/tests/test_empty_varlen.py @@ -0,0 +1,114 @@ +"""Regression tests for zero-length sequences in variable-length batches.""" + +import math + +import pytest +import torch +import torch.nn.functional as F + +import flash_kda +from torch_ref import torch_ref + + +D = 128 +LOWER_BOUND = -5.0 + + +def _make_inputs(total_tokens, heads): + torch.manual_seed(42) + q = F.normalize( + torch.randn((1, total_tokens, heads, D), dtype=torch.float32, device="cuda"), + p=2, + dim=-1, + ).to(torch.bfloat16) + k = F.normalize( + torch.randn((1, total_tokens, heads, D), dtype=torch.float32, device="cuda"), + p=2, + dim=-1, + ).to(torch.bfloat16) + v = torch.randn((1, total_tokens, heads, D), dtype=torch.bfloat16, device="cuda") + g = torch.randn((1, total_tokens, heads, D), dtype=torch.bfloat16, device="cuda") + beta = torch.randn((1, total_tokens, heads), dtype=torch.bfloat16, device="cuda") + a_log = torch.rand(heads, dtype=torch.float32, device="cuda") + dt_bias = torch.rand(heads, D, dtype=torch.float32, device="cuda") + return q, k, v, g, beta, a_log, dt_bias, 1.0 / math.sqrt(D) + + +def _make_state(shape, dtype): + return torch.arange( + math.prod(shape), dtype=torch.float32, device="cuda" + ).reshape(shape).to(torch.bfloat16).to(dtype) + + +@pytest.mark.parametrize( + "seq_lens", + [ + [0, 17], + [17, 0, 33], + [0, 0, 17, 0], + [0, 0], + ], + ids=["leading_empty", "middle_empty", "multiple_empty", "all_empty"], +) +@pytest.mark.parametrize("heads", [1, 4]) +@pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("has_initial_state", [False, True]) +def test_fwd_varlen_with_empty_sequences( + seq_lens, heads, state_dtype, has_initial_state +): + """Empty sequences have no output tokens and preserve their input state.""" + total_tokens = sum(seq_lens) + sequence_count = len(seq_lens) + cu_seqlens = torch.tensor( + [0, *torch.tensor(seq_lens).cumsum(0).tolist()], + dtype=torch.long, + device="cuda", + ) + q, k, v, g, beta, a_log, dt_bias, scale = _make_inputs(total_tokens, heads) + + initial_kernel = ( + _make_state((sequence_count, heads, D, D), state_dtype) if has_initial_state else None + ) + initial_reference = initial_kernel.clone() if initial_kernel is not None else None + final_kernel = torch.zeros( + sequence_count, heads, D, D, dtype=state_dtype, device="cuda" + ) + final_reference = torch.zeros_like(final_kernel) + output_kernel = torch.empty_like(q) + output_reference = torch.empty_like(q) + + flash_kda.fwd( + q, + k, + v, + g, + beta, + scale, + output_kernel, + A_log=a_log, + dt_bias=dt_bias, + lower_bound=LOWER_BOUND, + initial_state=initial_kernel, + final_state=final_kernel, + cu_seqlens=cu_seqlens, + ) + torch.cuda.synchronize() + torch_ref( + q, + k, + v, + g, + beta, + scale, + output_reference, + A_log=a_log, + dt_bias=dt_bias, + lower_bound=LOWER_BOUND, + initial_state=initial_reference, + final_state=final_reference, + cu_seqlens=cu_seqlens, + ) + + assert torch.equal(output_kernel, output_reference) + assert torch.equal(final_kernel, final_reference) +