Handle all-empty varlen batches

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

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<int64_t>();
+    } 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<int64_t>();
-    } 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

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;

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

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)
+