# Varlen Flash Attention Implementation Plan > **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. **Goal:** Add variable-length (varlen) packed-sequence support to `rtx-flash-attention` so mixed-length batches can be processed without padding waste, using a `cu_seqlens` prefix-sum descriptor. **Architecture:** A new Rust file `flash_varlen_forward.rs` provides a pure-Rust CPU simulation (`varlen_attention_cpu`) that is always compiled, plus a `#[cfg(feature = "cuda")]`-gated `FlashVarlenKernel` struct that loads a new CUDA kernel via NVRTC. A new `SdpaBackend::VarLen` variant is added to `backend_selector.rs`. The CUDA kernel maps each grid block to a (q-tile, head) pair using a linear scan over `cu_seqlens_q`, then runs tiled online-softmax attention with BLOCK_Q=64 / BLOCK_K=64 tiles. **Tech Stack:** Rust 2021, cudarc 0.16 (CUDA feature), half::bf16, existing FlashError/FlashResult error types. ## Global Constraints - Only modify files inside `crates/training/rtx-flash-attention/` - `cargo check -p rtx-flash-attention` must produce zero errors/warnings after each task - `cargo test -p rtx-flash-attention --lib` must pass (42 existing + 8 new = 50 tests total after Task 3) - No new Cargo dependencies - No GPU hardware required to run tests (all 8 new tests use `varlen_attention_cpu`) - Clippy-clean: no `#[allow(clippy::...)]` suppressions without a comment explaining why --- ## File Map | Status | Path | Role | |--------|------|------| | **CREATE** | `src/kernels/cuda/flash_attention_varlen.cu` | CUDA kernel — varlen forward pass | | **CREATE** | `src/kernels/flash_varlen_forward.rs` | Rust wrapper + CPU simulation | | **MODIFY** | `src/kernels/mod.rs` | Export the new module and re-exports | | **MODIFY** | `src/backend_selector.rs` | Add `SdpaBackend::VarLen` variant | --- ## Task 1: CUDA Kernel Source — `flash_attention_varlen.cu` **Files:** - Create: `crates/training/rtx-flash-attention/src/kernels/cuda/flash_attention_varlen.cu` **Interfaces:** - Produces: CUDA `__global__` function `flash_attention_varlen_forward` with the signature shown in Step 1. - [ ] **Step 1: Write the CUDA kernel file** Create the file with the following content exactly. The key design decisions are: - `BLOCK_Q = 64`, `BLOCK_K = 64` — tile dimensions (match spec) - Grid: `(ceil(max_seqlen_q / BLOCK_Q), num_heads, 1)` — blocks identified by `(q_tile_idx, head_idx)` - Sequence identification: linear scan over `cu_seqlens_q` in the block prologue - Accumulation in `f32`; output stored as `__nv_bfloat16` - Online softmax (Dao FA2 algorithm): running `m` (row max) and `l` (denominator) ```c // Variable-length FlashAttention forward kernel (packed sequences) // // Preconditions (caller-enforced): // q, k, v are packed [total_tokens, num_heads, head_dim] in row-major, dtype __nv_bfloat16 // out is same shape, pre-allocated // lse is [total_tokens, num_heads] in f32 // cu_seqlens_q[0] == 0, cu_seqlens_q[batch_size] == total_tokens_q // cu_seqlens_k[0] == 0, cu_seqlens_k[batch_size] == total_tokens_k // max_seqlen_q >= any individual sequence q length // 0 < head_dim <= 256, head_dim % 8 == 0 // // Grid/block assignment: // blockDim = (128, 1, 1) — 4 warps // gridDim = (ceil(max_seqlen_q / BLOCK_Q), num_heads, 1) #include #include #include #include #define BLOCK_Q 64 #define BLOCK_K 64 // Per-token element offset in packed layout [total_tokens, num_heads, head_dim] // token_idx is the absolute position in the packed buffer. __device__ __forceinline__ int elem_offset(int token_idx, int head_idx, int d, int num_heads, int head_dim) { return (token_idx * num_heads + head_idx) * head_dim + d; } // Find sequence index s such that cu_seqlens[s] <= pos < cu_seqlens[s+1]. // Linear scan is correct for all batch sizes; constant-time for batch <= 64. __device__ __forceinline__ int find_sequence(const int* cu_seqlens, int batch_size, int pos) { int s = 0; while (s < batch_size - 1 && cu_seqlens[s + 1] <= pos) { s++; } return s; } extern "C" __global__ void flash_attention_varlen_forward( const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ k, const __nv_bfloat16* __restrict__ v, __nv_bfloat16* __restrict__ out, float* __restrict__ lse, const int* __restrict__ cu_seqlens_q, const int* __restrict__ cu_seqlens_k, int max_seqlen_q, int batch_size, int num_heads, int head_dim, float softmax_scale, int causal ) { // ------------------------------------------------------------------------- // Identify which q-tile and head this block handles // ------------------------------------------------------------------------- const int q_tile_idx = blockIdx.x; // which BLOCK_Q tile within max_seqlen_q const int head_idx = blockIdx.y; const int tid = threadIdx.x; // 0..127 const int q_tile_start_global = q_tile_idx * BLOCK_Q; // ------------------------------------------------------------------------- // Identify which sequence owns this q-tile (linear scan over cu_seqlens_q) // ------------------------------------------------------------------------- int s = find_sequence(cu_seqlens_q, batch_size, q_tile_start_global); const int seq_start_q = cu_seqlens_q[s]; const int seq_end_q = cu_seqlens_q[s + 1]; const int seq_len_q = seq_end_q - seq_start_q; const int seq_start_k = cu_seqlens_k[s]; const int seq_end_k = cu_seqlens_k[s + 1]; const int seq_len_k = seq_end_k - seq_start_k; // Local q-tile start within this sequence const int local_q_start = q_tile_start_global - seq_start_q; // Early exit: this block is completely past the end of sequence s if (local_q_start >= seq_len_q) return; const int q_rows_this_tile = min(BLOCK_Q, seq_len_q - local_q_start); // ------------------------------------------------------------------------- // Shared memory layout: // [0 .. BLOCK_Q * head_dim) : Q tile (bf16) // [BLOCK_Q*D*2 .. BLOCK_Q*D*2 + BLOCK_K*D*2) : K tile (bf16) // [next .. next + BLOCK_K*D*2) : V tile (bf16) // We access shared mem via float-aligned pointers for vectorised loads. // ------------------------------------------------------------------------- extern __shared__ char smem_raw[]; __nv_bfloat16* smem_q = (__nv_bfloat16*)smem_raw; __nv_bfloat16* smem_k = smem_q + BLOCK_Q * head_dim; __nv_bfloat16* smem_v = smem_k + BLOCK_K * head_dim; // ------------------------------------------------------------------------- // Load Q tile into shared memory // Each of 128 threads loads elements strided across [q_row, d]. // ------------------------------------------------------------------------- for (int i = tid; i < q_rows_this_tile * head_dim; i += blockDim.x) { int row = i / head_dim; int col = i % head_dim; int global_token = seq_start_q + local_q_start + row; smem_q[row * head_dim + col] = q[elem_offset(global_token, head_idx, col, num_heads, head_dim)]; } // Zero-pad rows that don't exist in this tile for (int i = q_rows_this_tile * head_dim + tid; i < BLOCK_Q * head_dim; i += blockDim.x) { smem_q[i] = __float2bfloat16(0.0f); } __syncthreads(); // ------------------------------------------------------------------------- // Per-row accumulators in registers (one row per thread — split 64 rows // across 128 threads, so each thread "owns" rows tid/2 with half the threads // handling even/odd columns via the inner d-loop). // For simplicity we assign each thread one q-row cyclically. // ------------------------------------------------------------------------- // We unroll over q rows: thread tid handles q-row (tid) if tid < BLOCK_Q. // For 128 threads and BLOCK_Q=64, threads 0..63 process one row each. // Threads 64..127 do nothing except participate in K/V loads. // Accumulators: acc[d] for output, running_m and running_l for online softmax float acc[256]; // max supported head_dim float running_m = -FLT_MAX; float running_l = 0.0f; for (int d = 0; d < head_dim; d++) acc[d] = 0.0f; const int my_q_row = tid; // my row within the q tile (tid 0..63 active, 64..127 idle) const int my_active = (my_q_row < q_rows_this_tile) ? 1 : 0; // ------------------------------------------------------------------------- // Outer loop over K/V tiles // ------------------------------------------------------------------------- const int num_k_tiles = (seq_len_k + BLOCK_K - 1) / BLOCK_K; for (int k_tile = 0; k_tile < num_k_tiles; k_tile++) { const int k_tile_start = k_tile * BLOCK_K; const int k_rows_this_tile = min(BLOCK_K, seq_len_k - k_tile_start); // ------------------------------------------------------------------ // Load K tile into shared memory // ------------------------------------------------------------------ for (int i = tid; i < k_rows_this_tile * head_dim; i += blockDim.x) { int row = i / head_dim; int col = i % head_dim; int global_token = seq_start_k + k_tile_start + row; smem_k[row * head_dim + col] = k[elem_offset(global_token, head_idx, col, num_heads, head_dim)]; } for (int i = k_rows_this_tile * head_dim + tid; i < BLOCK_K * head_dim; i += blockDim.x) { smem_k[i] = __float2bfloat16(0.0f); } // ------------------------------------------------------------------ // Load V tile into shared memory // ------------------------------------------------------------------ for (int i = tid; i < k_rows_this_tile * head_dim; i += blockDim.x) { int row = i / head_dim; int col = i % head_dim; int global_token = seq_start_k + k_tile_start + row; smem_v[row * head_dim + col] = v[elem_offset(global_token, head_idx, col, num_heads, head_dim)]; } for (int i = k_rows_this_tile * head_dim + tid; i < BLOCK_K * head_dim; i += blockDim.x) { smem_v[i] = __float2bfloat16(0.0f); } __syncthreads(); if (!my_active) { __syncthreads(); continue; } // ------------------------------------------------------------------ // Compute attention scores for my_q_row x all k_rows: S[j] = Q[my_q_row] . K[j] // ------------------------------------------------------------------ float S[BLOCK_K]; for (int j = 0; j < k_rows_this_tile; j++) { float dot = 0.0f; for (int d = 0; d < head_dim; d++) { dot += __bfloat162float(smem_q[my_q_row * head_dim + d]) * __bfloat162float(smem_k[j * head_dim + d]); } S[j] = dot * softmax_scale; // Causal mask: token (seq_start_q + local_q_start + my_q_row) cannot // attend to token (seq_start_k + k_tile_start + j) if j > my_q_row+local_q_start. if (causal) { int q_pos = local_q_start + my_q_row; int k_pos = k_tile_start + j; if (k_pos > q_pos) S[j] = -FLT_MAX; } // Pad out-of-sequence k positions if (j >= k_rows_this_tile) S[j] = -FLT_MAX; } for (int j = k_rows_this_tile; j < BLOCK_K; j++) S[j] = -FLT_MAX; // ------------------------------------------------------------------ // Online softmax update (FA2 algorithm 1) // ------------------------------------------------------------------ float tile_m = -FLT_MAX; for (int j = 0; j < k_rows_this_tile; j++) tile_m = fmaxf(tile_m, S[j]); float m_new = fmaxf(running_m, tile_m); // Rescale existing accumulator float rescale = expf(running_m - m_new); for (int d = 0; d < head_dim; d++) acc[d] *= rescale; float l_rescale = running_l * rescale; // Accumulate weighted V float tile_l = 0.0f; for (int j = 0; j < k_rows_this_tile; j++) { float p = expf(S[j] - m_new); tile_l += p; for (int d = 0; d < head_dim; d++) { acc[d] += p * __bfloat162float(smem_v[j * head_dim + d]); } } running_m = m_new; running_l = l_rescale + tile_l; __syncthreads(); } // ------------------------------------------------------------------------- // Write output and LSE // ------------------------------------------------------------------------- if (!my_active) return; float inv_l = (running_l > 0.0f) ? (1.0f / running_l) : 0.0f; int global_token_out = seq_start_q + local_q_start + my_q_row; for (int d = 0; d < head_dim; d++) { out[elem_offset(global_token_out, head_idx, d, num_heads, head_dim)] = __float2bfloat16(acc[d] * inv_l); } // lse[token, head] = log(l) + m (log-sum-exp in standard form) float lse_val = (running_l > 0.0f) ? (logf(running_l) + running_m) : -FLT_MAX; lse[global_token_out * num_heads + head_idx] = lse_val; } ``` - [ ] **Step 2: Verify file exists** ```bash ls -la /slab/projects/rustyverse/rustytorch/crates/training/rtx-flash-attention/src/kernels/cuda/flash_attention_varlen.cu ``` Expected: file listed with non-zero size. --- ## Task 2: Rust Wrapper + CPU Simulation — `flash_varlen_forward.rs` **Files:** - Create: `crates/training/rtx-flash-attention/src/kernels/flash_varlen_forward.rs` **Interfaces:** - Consumes: `crate::error::{FlashError, FlashResult}` (always); `cudarc::driver::{CudaContext, CudaSlice, CudaStream, LaunchConfig, PushKernelArg}`, `cudarc::nvrtc::compile_ptx`, `half::bf16` (cuda feature only). - Produces: - `pub fn varlen_attention_cpu(q: &[f32], k: &[f32], v: &[f32], cu_seqlens: &[usize], num_heads: usize, head_dim: usize, softmax_scale: f32, causal: bool) -> Vec` — always available. - `#[cfg(feature = "cuda")] pub struct FlashVarlenKernel` with: - `pub fn new(ctx: &Arc) -> FlashResult` - `pub fn forward(&self, q: &CudaSlice, k: &CudaSlice, v: &CudaSlice, out: &mut CudaSlice, lse: &mut CudaSlice, cu_seqlens_q: &CudaSlice, cu_seqlens_k: &CudaSlice, max_seqlen_q: usize, batch_size: usize, num_heads: usize, head_dim: usize, softmax_scale: f32, causal: bool, stream: &Arc) -> FlashResult<()>` - [ ] **Step 1: Write the Rust file** Create `/slab/projects/rustyverse/rustytorch/crates/training/rtx-flash-attention/src/kernels/flash_varlen_forward.rs` with the following content: ```rust //! Variable-length (varlen) Flash Attention forward pass. //! //! This module provides: //! - [`varlen_attention_cpu`] — pure-Rust O(n²) reference implementation for //! testing and CPU fallback. No GPU required. //! - [`FlashVarlenKernel`] — CUDA kernel wrapper (cuda feature only). //! //! # Packed layout //! //! All tensors use **packed** (also called "varlen" or "jagged") layout: //! instead of `[batch, heads, seq_len, head_dim]` with padding, sequences are //! concatenated along the token axis: `[total_tokens, heads, head_dim]`. //! //! A `cu_seqlens` prefix-sum array of length `batch_size + 1` describes the //! boundaries: //! - `cu_seqlens[0] = 0` //! - `cu_seqlens[b + 1] = cu_seqlens[b] + seqlen_b` //! - `cu_seqlens[batch_size] = total_tokens` //! //! Sequence `b` occupies token indices `cu_seqlens[b] .. cu_seqlens[b+1]`. use crate::error::FlashResult; #[cfg(feature = "cuda")] use crate::error::FlashError; #[cfg(feature = "cuda")] use cudarc::driver::{CudaContext, CudaSlice, CudaStream, LaunchConfig, PushKernelArg}; #[cfg(feature = "cuda")] use cudarc::nvrtc::compile_ptx; #[cfg(feature = "cuda")] use half::bf16; #[cfg(feature = "cuda")] use std::sync::Arc; #[cfg(feature = "cuda")] const VARLEN_CUDA_SOURCE: &str = include_str!("cuda/flash_attention_varlen.cu"); // ============================================================================= // CPU reference implementation // ============================================================================= /// Variable-length attention on the CPU using naive O(n²) per-sequence dot-product. /// /// # Layout /// /// `q`, `k`, `v` are packed flat slices in row-major order: /// `[total_tokens * num_heads * head_dim]` where the logical index for /// token `t`, head `h`, dimension `d` is `(t * num_heads + h) * head_dim + d`. /// /// `cu_seqlens` has length `batch_size + 1`; `cu_seqlens[b..b+1]` is a half-open /// range of token indices for sequence `b`. Sequences with zero length /// (i.e. `cu_seqlens[b] == cu_seqlens[b+1]`) are silently skipped. /// /// # Returns /// /// A `Vec` of length `total_tokens * num_heads * head_dim` — the attention /// output in the same packed layout. /// /// # Panics /// /// Panics if `cu_seqlens` is empty (must have at least one entry: `[0]`). /// /// # Example /// /// ``` /// use rtx_flash_attention::kernels::varlen_attention_cpu; /// /// let head_dim = 4; /// let num_heads = 1; /// // Two sequences: seq0 = 2 tokens, seq1 = 3 tokens → 5 total tokens /// let cu_seqlens = vec![0usize, 2, 5]; /// let total = cu_seqlens[cu_seqlens.len() - 1]; /// let scale = 1.0 / (head_dim as f32).sqrt(); /// /// let q: Vec = (0..total * num_heads * head_dim).map(|i| i as f32 * 0.01).collect(); /// let k = q.clone(); /// let v = q.clone(); /// /// let out = varlen_attention_cpu(&q, &k, &v, &cu_seqlens, num_heads, head_dim, scale, false); /// assert_eq!(out.len(), total * num_heads * head_dim); /// ``` pub fn varlen_attention_cpu( q: &[f32], k: &[f32], v: &[f32], cu_seqlens: &[usize], num_heads: usize, head_dim: usize, softmax_scale: f32, causal: bool, ) -> Vec { assert!(!cu_seqlens.is_empty(), "cu_seqlens must have at least one entry"); let batch_size = cu_seqlens.len() - 1; let total_tokens = cu_seqlens[batch_size]; let mut output = vec![0.0f32; total_tokens * num_heads * head_dim]; // Helper: flat index for packed layout [total_tokens, num_heads, head_dim] let idx = |token: usize, head: usize, d: usize| -> usize { (token * num_heads + head) * head_dim + d }; for b in 0..batch_size { let seq_start = cu_seqlens[b]; let seq_end = cu_seqlens[b + 1]; let seq_len = seq_end - seq_start; // Empty sequence — skip without touching output if seq_len == 0 { continue; } for h in 0..num_heads { // Compute attention weights for every (query, key) pair in this sequence. // scores[qi * seq_len + ki] = Q[qi] . K[ki] * scale let mut scores = vec![0.0f32; seq_len * seq_len]; for qi in 0..seq_len { for ki in 0..seq_len { // Causal: query at position qi cannot attend to key at ki > qi if causal && ki > qi { scores[qi * seq_len + ki] = f32::NEG_INFINITY; continue; } let mut dot = 0.0f32; for d in 0..head_dim { dot += q[idx(seq_start + qi, h, d)] * k[idx(seq_start + ki, h, d)]; } scores[qi * seq_len + ki] = dot * softmax_scale; } } // Softmax row-by-row (online — subtract row max for numerical stability) let mut weights = vec![0.0f32; seq_len * seq_len]; for qi in 0..seq_len { let row = &scores[qi * seq_len..(qi + 1) * seq_len]; let row_max = row .iter() .copied() .fold(f32::NEG_INFINITY, f32::max); let exps: Vec = row.iter().map(|&s| (s - row_max).exp()).collect(); let sum: f32 = exps.iter().sum(); let inv_sum = if sum > 0.0 { 1.0 / sum } else { 0.0 }; for ki in 0..seq_len { weights[qi * seq_len + ki] = exps[ki] * inv_sum; } } // Output: O[qi, d] = Σ_ki weights[qi, ki] * V[ki, d] for qi in 0..seq_len { for d in 0..head_dim { let mut acc = 0.0f32; for ki in 0..seq_len { acc += weights[qi * seq_len + ki] * v[idx(seq_start + ki, h, d)]; } output[idx(seq_start + qi, h, d)] = acc; } } } } output } // ============================================================================= // CUDA kernel wrapper // ============================================================================= /// Compiled FlashAttention varlen forward kernel for a single CUDA device. /// /// Create once per device; reuse across calls. The compiled PTX is cached /// inside the struct. /// /// # Feature gate /// /// Only available when the `cuda` feature is enabled. #[cfg(feature = "cuda")] pub struct FlashVarlenKernel { module: Arc, _ctx: Arc, } #[cfg(feature = "cuda")] impl FlashVarlenKernel { /// Compile and load the varlen CUDA kernel into `ctx`. /// /// # Errors /// /// Returns [`FlashError::Cuda`] if NVRTC compilation fails or the PTX /// cannot be loaded into the device context. pub fn new(ctx: &Arc) -> FlashResult { let ptx = compile_ptx(VARLEN_CUDA_SOURCE) .map_err(|e| FlashError::cuda(format!("varlen NVRTC compilation failed: {e:?}")))?; let module = ctx .load_module(ptx) .map_err(|e| FlashError::cuda(format!("varlen PTX load failed: {e:?}")))?; Ok(Self { module, _ctx: Arc::clone(ctx), }) } /// Launch the varlen forward kernel on `stream`. /// /// # Arguments /// /// - `q`, `k`, `v` — packed `[total_tokens, num_heads, head_dim]` in BF16. /// - `out` — pre-allocated output buffer, same shape as `q`. /// - `lse` — pre-allocated log-sum-exp buffer `[total_tokens, num_heads]` in f32. /// - `cu_seqlens_q` / `cu_seqlens_k` — device buffers of length `batch_size + 1`. /// - `max_seqlen_q` — maximum query sequence length across the batch. /// - `batch_size` — number of sequences in the batch. /// /// # Safety invariants (enforced by caller) /// /// - All device buffers must be allocated on the same device as `ctx`. /// - Buffer lengths: `q.len() >= total_tokens_q * num_heads * head_dim`. /// - `cu_seqlens_q[0] == 0`, `cu_seqlens_q[batch_size] == total_tokens_q`. /// - `head_dim <= 256` and `head_dim % 8 == 0`. /// /// # Errors /// /// Returns [`FlashError::Cuda`] if the kernel symbol is not found or the /// driver rejects the launch configuration. #[allow(clippy::too_many_arguments)] // kernel interface requires all parameters pub fn forward( &self, q: &CudaSlice, k: &CudaSlice, v: &CudaSlice, out: &mut CudaSlice, lse: &mut CudaSlice, cu_seqlens_q: &CudaSlice, cu_seqlens_k: &CudaSlice, max_seqlen_q: usize, batch_size: usize, num_heads: usize, head_dim: usize, softmax_scale: f32, causal: bool, stream: &Arc, ) -> FlashResult<()> { let kernel = self .module .load_function("flash_attention_varlen_forward") .map_err(|e| { FlashError::cuda(format!("varlen kernel symbol not found: {e:?}")) })?; // Grid: (ceil(max_seqlen_q / BLOCK_Q), num_heads, 1) const BLOCK_Q: usize = 64; let grid_x = max_seqlen_q.div_ceil(BLOCK_Q) as u32; let grid_y = num_heads as u32; // Shared memory: 3 tiles × BLOCK_K × head_dim × sizeof(bf16) const BLOCK_K: usize = 64; let smem_bytes = (3 * BLOCK_K * head_dim * 2) as u32; // bf16 = 2 bytes let launch_cfg = LaunchConfig { block_dim: (128, 1, 1), grid_dim: (grid_x, grid_y, 1), shared_mem_bytes: smem_bytes, }; let causal_int = causal as i32; let batch_size_i32 = batch_size as i32; let num_heads_i32 = num_heads as i32; let head_dim_i32 = head_dim as i32; let max_seqlen_q_i32 = max_seqlen_q as i32; // Safety: // - All CudaSlice buffers are device-resident on the same device as `_ctx`. // - `out` and `lse` are mutable unique references, so no aliasing. // - The kernel reads cu_seqlens_q/k as read-only int arrays; they have // length `batch_size + 1` ≥ 2 (enforced by caller). // - head_dim ≤ 256 ensures the `acc[256]` register array is not overflowed. // - The CUDA kernel performs a bounds check per q-row (early return when // the tile is past the end of its sequence). unsafe { let mut builder = stream.launch_builder(&kernel); builder.arg(q); builder.arg(k); builder.arg(v); builder.arg(out); builder.arg(lse); builder.arg(cu_seqlens_q); builder.arg(cu_seqlens_k); builder.arg(&max_seqlen_q_i32); builder.arg(&batch_size_i32); builder.arg(&num_heads_i32); builder.arg(&head_dim_i32); builder.arg(&softmax_scale); builder.arg(&causal_int); builder.launch(launch_cfg) .map_err(|e| FlashError::cuda(format!("varlen kernel launch failed: {e:?}")))?; } Ok(()) } } // ============================================================================= // Tests — all CPU, no GPU required // ============================================================================= #[cfg(test)] mod tests { use super::varlen_attention_cpu; // Helper: flat index for packed [total_tokens, num_heads, head_dim] fn idx(token: usize, head: usize, d: usize, num_heads: usize, head_dim: usize) -> usize { (token * num_heads + head) * head_dim + d } // Build a simple packed buffer filled with a constant value per token. fn constant_qkv( cu_seqlens: &[usize], num_heads: usize, head_dim: usize, val: f32, ) -> Vec { let total = *cu_seqlens.last().unwrap(); vec![val; total * num_heads * head_dim] } // ------------------------------------------------------------------------- // Test 1: Single sequence — must match standard (non-varlen) attention // ------------------------------------------------------------------------- #[test] fn test_varlen_cpu_single_sequence_matches_regular() { let num_heads = 2; let head_dim = 8; let seq_len = 4; let scale = 1.0 / (head_dim as f32).sqrt(); // cu_seqlens for a single sequence of length 4 let cu_seqlens = vec![0usize, seq_len]; let total = seq_len; // Random-ish deterministic data let q: Vec = (0..total * num_heads * head_dim) .map(|i| (i as f32) * 0.1) .collect(); let k = q.clone(); let v = q.clone(); let out = varlen_attention_cpu(&q, &k, &v, &cu_seqlens, num_heads, head_dim, scale, false); // Output has the right length assert_eq!(out.len(), total * num_heads * head_dim); // For a single sequence, varlen output must be identical to the same // computation run as a single entry — we verify by calling again with // the same data (idempotency check + non-zero output sanity). let out2 = varlen_attention_cpu(&q, &k, &v, &cu_seqlens, num_heads, head_dim, scale, false); for i in 0..out.len() { assert!( (out[i] - out2[i]).abs() < 1e-6, "output not deterministic at index {i}: {} vs {}", out[i], out2[i] ); } // All values must be finite (no NaN/Inf for well-conditioned inputs) for (i, &v) in out.iter().enumerate() { assert!(v.is_finite(), "non-finite output at index {i}: {v}"); } } // ------------------------------------------------------------------------- // Test 2: Two sequences are independent — tokens from seq0 do not appear // in seq1's attention output // ------------------------------------------------------------------------- #[test] fn test_varlen_cpu_two_sequences_independent() { let num_heads = 1; let head_dim = 4; let scale = 1.0 / (head_dim as f32).sqrt(); // Sequence 0: 3 tokens, all-ones // Sequence 1: 4 tokens, all-twos let cu_seqlens = vec![0usize, 3, 7]; let total = 7; let mut q = vec![0.0f32; total * num_heads * head_dim]; let mut k = vec![0.0f32; total * num_heads * head_dim]; let mut v = vec![0.0f32; total * num_heads * head_dim]; // seq0 tokens (0..3) → value 1.0 for t in 0..3 { for d in 0..head_dim { q[idx(t, 0, d, num_heads, head_dim)] = 1.0; k[idx(t, 0, d, num_heads, head_dim)] = 1.0; v[idx(t, 0, d, num_heads, head_dim)] = 1.0; } } // seq1 tokens (3..7) → value 100.0 (very different magnitude) for t in 3..7 { for d in 0..head_dim { q[idx(t, 0, d, num_heads, head_dim)] = 100.0; k[idx(t, 0, d, num_heads, head_dim)] = 100.0; v[idx(t, 0, d, num_heads, head_dim)] = 100.0; } } let out = varlen_attention_cpu(&q, &k, &v, &cu_seqlens, num_heads, head_dim, scale, false); // seq0 output must be ~1.0 (attending only to seq0 tokens, all equal) for t in 0..3 { for d in 0..head_dim { let val = out[idx(t, 0, d, num_heads, head_dim)]; assert!( (val - 1.0).abs() < 1e-5, "seq0 token {t} dim {d}: expected ~1.0, got {val}" ); } } // seq1 output must be ~100.0 for t in 3..7 { for d in 0..head_dim { let val = out[idx(t, 0, d, num_heads, head_dim)]; assert!( (val - 100.0).abs() < 1e-5, "seq1 token {t} dim {d}: expected ~100.0, got {val}" ); } } } // ------------------------------------------------------------------------- // Test 3: Output shape is [total_tokens, heads, head_dim] // ------------------------------------------------------------------------- #[test] fn test_varlen_cpu_output_shape() { let num_heads = 4; let head_dim = 16; let cu_seqlens = vec![0usize, 5, 12, 15]; // 3 seqs, lengths 5, 7, 3 let total = 15; let scale = 1.0 / (head_dim as f32).sqrt(); let q = vec![1.0f32; total * num_heads * head_dim]; let k = q.clone(); let v = q.clone(); let out = varlen_attention_cpu(&q, &k, &v, &cu_seqlens, num_heads, head_dim, scale, false); assert_eq!( out.len(), total * num_heads * head_dim, "output length must equal total_tokens * num_heads * head_dim" ); } // ------------------------------------------------------------------------- // Test 4: cu_seqlens semantics — [0, 3, 7] means seq0=3 tokens, seq1=4 tokens // ------------------------------------------------------------------------- #[test] fn test_varlen_cu_seqlens_correct() { let num_heads = 1; let head_dim = 4; let scale = 1.0 / (head_dim as f32).sqrt(); let cu_seqlens = vec![0usize, 3, 7]; // Verify: batch_size = cu_seqlens.len() - 1 = 2 let batch_size = cu_seqlens.len() - 1; assert_eq!(batch_size, 2); // seq0 length = cu_seqlens[1] - cu_seqlens[0] = 3 assert_eq!(cu_seqlens[1] - cu_seqlens[0], 3); // seq1 length = cu_seqlens[2] - cu_seqlens[1] = 4 assert_eq!(cu_seqlens[2] - cu_seqlens[1], 4); let total = cu_seqlens[batch_size]; let q = vec![0.5f32; total * num_heads * head_dim]; let k = q.clone(); let v = q.clone(); let out = varlen_attention_cpu(&q, &k, &v, &cu_seqlens, num_heads, head_dim, scale, false); assert_eq!(out.len(), total * num_heads * head_dim); } // ------------------------------------------------------------------------- // Test 5: Causal mask — token i cannot attend to token j > i (within sequence) // ------------------------------------------------------------------------- #[test] fn test_varlen_causal_mask() { let num_heads = 1; let head_dim = 4; let seq_len = 4; let scale = 1.0 / (head_dim as f32).sqrt(); let cu_seqlens = vec![0usize, seq_len]; // Use distinct V values per token so any cross-contamination is detectable. // V[token t, head 0, all dims] = (t+1) as f32 let total = seq_len; let mut q = vec![1.0f32; total * num_heads * head_dim]; let mut k = vec![1.0f32; total * num_heads * head_dim]; let mut v = vec![0.0f32; total * num_heads * head_dim]; for t in 0..seq_len { for d in 0..head_dim { v[idx(t, 0, d, num_heads, head_dim)] = (t + 1) as f32; } } let out_causal = varlen_attention_cpu(&q, &k, &v, &cu_seqlens, num_heads, head_dim, scale, true); let out_non_causal = varlen_attention_cpu(&q, &k, &v, &cu_seqlens, num_heads, head_dim, scale, false); // Token 0 in causal mode: can only attend to itself (k=0). // V[0,*] = 1.0, so output for token 0 must be 1.0 in causal mode. for d in 0..head_dim { let val = out_causal[idx(0, 0, d, num_heads, head_dim)]; assert!( (val - 1.0).abs() < 1e-5, "causal: token 0 dim {d} expected 1.0, got {val}" ); } // In non-causal mode, token 0 attends to all tokens uniformly (Q=K=1.0), // so output = mean(V) = (1+2+3+4)/4 = 2.5 for d in 0..head_dim { let val = out_non_causal[idx(0, 0, d, num_heads, head_dim)]; assert!( (val - 2.5).abs() < 1e-5, "non-causal: token 0 dim {d} expected 2.5, got {val}" ); } // Last token in causal mode attends to all tokens (same as non-causal for last token) for d in 0..head_dim { let causal_last = out_causal[idx(seq_len - 1, 0, d, num_heads, head_dim)]; let noncausal_last = out_non_causal[idx(seq_len - 1, 0, d, num_heads, head_dim)]; assert!( (causal_last - noncausal_last).abs() < 1e-5, "causal last token should match non-causal: got {causal_last} vs {noncausal_last}" ); } } // ------------------------------------------------------------------------- // Test 6: Attention weights sum to 1.0 per row within 1e-5 // ------------------------------------------------------------------------- #[test] fn test_varlen_softmax_sums_to_one() { let num_heads = 2; let head_dim = 8; let scale = 1.0 / (head_dim as f32).sqrt(); // Two sequences: 3 and 5 tokens let cu_seqlens = vec![0usize, 3, 8]; let total = 8; // V is the identity: each token t, dim d = (t == d) ? 1.0 : 0.0 // This lets us recover the attention weight for each (q, k) pair // as the output value at dimension k. Works only when head_dim >= seq_len, // but here head_dim=8 >= max(3,5)=5, so the trick works. let q = vec![1.0f32; total * num_heads * head_dim]; let k = vec![1.0f32; total * num_heads * head_dim]; // V = identity: V[t, h, d] = if d == t % head_dim { 1.0 } else { 0.0 } // (use global token index t, so different tokens map to different dims) let mut v = vec![0.0f32; total * num_heads * head_dim]; for t in 0..total { for h in 0..num_heads { let d = t % head_dim; // unique dim per token (modulo to stay in range) v[idx(t, h, d, num_heads, head_dim)] = 1.0; } } let out = varlen_attention_cpu(&q, &k, &v, &cu_seqlens, num_heads, head_dim, scale, false); // For each sequence, for each q-row, sum weights extracted via V=identity. // Since q=k=const, the attention is uniform → weights = 1/seq_len. // The output at dim d = weight for token whose unique dim is d. // Sum of all output dims within a sequence = sum of all weights = 1.0. for seq_idx in 0..2 { let seq_start = cu_seqlens[seq_idx]; let seq_end = cu_seqlens[seq_idx + 1]; let seq_len = seq_end - seq_start; for qi in 0..seq_len { let global_token = seq_start + qi; for h in 0..num_heads { // Sum over dims that correspond to tokens in *this* sequence let mut weight_sum = 0.0f32; for ki in 0..seq_len { let kt = seq_start + ki; let d = kt % head_dim; weight_sum += out[idx(global_token, h, d, num_heads, head_dim)]; } assert!( (weight_sum - 1.0).abs() < 1e-5, "seq{seq_idx} q={qi} h={h}: weights sum to {weight_sum}, expected 1.0" ); } } } } // ------------------------------------------------------------------------- // Test 7: Empty sequence — cu_seqlens = [0, 0, 4] → skip seq0, process seq1 // ------------------------------------------------------------------------- #[test] fn test_varlen_empty_sequence_handled() { let num_heads = 1; let head_dim = 4; let scale = 1.0 / (head_dim as f32).sqrt(); // seq0 is empty (0 tokens), seq1 has 4 tokens let cu_seqlens = vec![0usize, 0, 4]; let total = 4; // only seq1 tokens exist let q = vec![1.0f32; total * num_heads * head_dim]; let k = q.clone(); let v = q.clone(); // Must not panic — empty seq0 is silently skipped let out = varlen_attention_cpu(&q, &k, &v, &cu_seqlens, num_heads, head_dim, scale, false); assert_eq!(out.len(), total * num_heads * head_dim); // seq1 tokens (indices 0..4 in the packed buffer) should have finite values for i in 0..out.len() { assert!(out[i].is_finite(), "non-finite at index {i}: {}", out[i]); } // For uniform q/k/v=1.0, output should be 1.0 for i in 0..out.len() { assert!( (out[i] - 1.0).abs() < 1e-5, "expected 1.0 at index {i}, got {}", out[i] ); } } // ------------------------------------------------------------------------- // Test 8: Standard head_dim=64 (production default) // ------------------------------------------------------------------------- #[test] fn test_varlen_head_dim_64() { let num_heads = 8; let head_dim = 64; let scale = 1.0 / (head_dim as f32).sqrt(); let cu_seqlens = vec![0usize, 7, 15]; // seq0=7, seq1=8 let total = 15; let q: Vec = (0..total * num_heads * head_dim) .map(|i| (i as f32) * 0.001) .collect(); let k = q.clone(); let v: Vec = (0..total * num_heads * head_dim) .map(|i| (i as f32) * 0.002) .collect(); let out = varlen_attention_cpu(&q, &k, &v, &cu_seqlens, num_heads, head_dim, scale, false); assert_eq!(out.len(), total * num_heads * head_dim); for (i, &val) in out.iter().enumerate() { assert!(val.is_finite(), "non-finite at index {i}: {val}"); } } } ``` - [ ] **Step 2: Verify the file exists** ```bash ls -la /slab/projects/rustyverse/rustytorch/crates/training/rtx-flash-attention/src/kernels/flash_varlen_forward.rs ``` Expected: file with non-zero size. --- ## Task 3: Wire Up Exports and Add `SdpaBackend::VarLen` **Files:** - Modify: `crates/training/rtx-flash-attention/src/kernels/mod.rs:1-57` - Modify: `crates/training/rtx-flash-attention/src/backend_selector.rs` (lines 38–68 for the enum, lines 205–216 for `supports_backend`, lines 55–68 for `Display`) **Interfaces:** - Consumes: `flash_varlen_forward::varlen_attention_cpu` and `flash_varlen_forward::FlashVarlenKernel` from Task 2. - Produces: `pub use kernels::varlen_attention_cpu` available at crate root (via `lib.rs` re-export chain); `SdpaBackend::VarLen` in the public enum. - [ ] **Step 1: Add module declaration and re-exports to `mod.rs`** Edit `src/kernels/mod.rs`. Add after the existing `#[cfg(feature = "cuda")] pub mod ptx;` block (line 14) and before the `#[cfg(feature = "metal")] pub mod metal;` block (line 18): ```rust // varlen module is always compiled (contains CPU simulation unconditionally) pub mod flash_varlen_forward; pub use flash_varlen_forward::varlen_attention_cpu; #[cfg(feature = "cuda")] pub use flash_varlen_forward::FlashVarlenKernel; ``` The resulting block after the edit looks like: ```rust //! GPU kernels for Flash Attention implementation #[cfg(feature = "cuda")] pub mod flash_forward; #[cfg(feature = "cuda")] pub mod flash_backward; #[cfg(feature = "cuda")] pub mod flash_v3_forward; #[cfg(feature = "cuda")] pub mod utils; #[cfg(feature = "cuda")] pub mod simple; #[cfg(feature = "cuda")] pub mod ptx; #[cfg(feature = "cuda")] pub mod manager; // varlen module is always compiled (contains CPU simulation unconditionally) pub mod flash_varlen_forward; pub use flash_varlen_forward::varlen_attention_cpu; #[cfg(feature = "cuda")] pub use flash_varlen_forward::FlashVarlenKernel; #[cfg(feature = "metal")] pub mod metal; #[cfg(all(test, feature = "cuda"))] mod manager_test; #[cfg(all(test, feature = "metal"))] mod metal_test; #[cfg(all(test, feature = "cuda"))] mod test_kernel_launch; #[cfg(all(test, feature = "cuda"))] mod cudarc_api_test; // Re-export the implementation based on feature #[cfg(feature = "cuda")] pub use simple::{FlashCudaKernels, KernelResult}; #[cfg(feature = "cuda")] pub use flash_forward::{ FlashForwardKernel, KernelExecutionResult, KernelPerformanceInfo, KernelPerformanceCache, FP8Format, FP8KernelConfig, }; #[cfg(feature = "cuda")] pub use flash_v3_forward::FlashV3ForwardKernel; #[cfg(feature = "metal")] pub use metal::{FlashMetalKernels, MetalKernelResult}; // The complex CUDA kernel implementation has been replaced with a simplified version // for compilation purposes. This allows the advanced crates to compile while maintaining // the required API surface. In production, this would be replaced with optimized CUDA kernels. ``` - [ ] **Step 2: Add `SdpaBackend::VarLen` to `backend_selector.rs`** In `src/backend_selector.rs`, locate the `SdpaBackend` enum (lines 38–54) and add `VarLen` as a new variant. Add it after `Cpu`: ```rust /// Available SDPA backends #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum SdpaBackend { /// FlashAttention v2 - optimal for long sequences FlashAttention, /// FlashAttention v3 - WGMMA + TMA + warp specialization (Hopper/Blackwell) FlashAttentionV3, /// Standard mathematical attention - simple, debuggable Math, /// Memory-efficient chunked attention MemoryEfficient, /// NVIDIA cuDNN attention (Ampere+) CuDnn, /// Custom Metal implementation for Apple Silicon Metal, /// Fallback CPU implementation Cpu, /// Variable-length packed-sequence attention (no padding) VarLen, } ``` - [ ] **Step 3: Add `Display` arm for `VarLen`** In the `impl std::fmt::Display for SdpaBackend` block (lines 56–68), add: ```rust SdpaBackend::VarLen => write!(f, "VarLen"), ``` after the `SdpaBackend::Cpu` arm. - [ ] **Step 4: Add `supports_backend` arm for `VarLen`** In `HardwareCapabilities::supports_backend` (lines 206–217), add: ```rust SdpaBackend::VarLen => true, // CPU simulation always available; CUDA variant when feature is on ``` after the `SdpaBackend::Cpu => true,` arm. - [ ] **Step 5: Run `cargo check` and confirm clean** ```bash ~/.cargo/bin/cargo check -p rtx-flash-attention 2>&1 ``` Expected: zero errors, zero warnings (or only pre-existing warnings from other files). - [ ] **Step 6: Run `cargo test --lib` and confirm 50 tests pass** ```bash ~/.cargo/bin/cargo test -p rtx-flash-attention --lib 2>&1 | tail -30 ``` Expected output contains: ``` test result: ok. 50 passed; 0 failed; ... ``` (42 existing + 8 new varlen tests = 50) - [ ] **Step 7: Commit** ```bash git -C /slab/projects/rustyverse/rustytorch add \ crates/training/rtx-flash-attention/src/kernels/cuda/flash_attention_varlen.cu \ crates/training/rtx-flash-attention/src/kernels/flash_varlen_forward.rs \ crates/training/rtx-flash-attention/src/kernels/mod.rs \ crates/training/rtx-flash-attention/src/backend_selector.rs git -C /slab/projects/rustyverse/rustytorch commit -m "$(cat <<'EOF' feat(flash-attention): add varlen packed-sequence support Implements variable-length (varlen) FlashAttention that processes mixed-length batches without padding waste: - New CUDA kernel flash_attention_varlen_forward with BLOCK_Q=64 / BLOCK_K=64 tiling; grid=(ceil(max_seqlen_q/64), num_heads, 1). Each block uses a linear scan over cu_seqlens_q to identify its owning sequence and exits early when past sequence end. - New Rust module flash_varlen_forward: always-compiled CPU simulation (varlen_attention_cpu) for testing + #[cfg(cuda)] FlashVarlenKernel. - SdpaBackend::VarLen variant added to backend_selector. - 8 new CPU-only tests; total test count: 50. Co-Authored-By: Claude Sonnet 4.6 EOF )" ``` --- ## Self-Review Checklist ### Spec Coverage | Spec requirement | Task covering it | |---|---| | New CUDA kernel `flash_attention_varlen.cu` | Task 1 | | Grid mapping: `(ceil(max_seqlen_q/BLOCK_Q), num_heads, 1)` | Task 1 Step 1 | | Sequence identification by linear scan | Task 1 Step 1 | | Early exit when tile is past sequence end | Task 1 Step 1 | | BLOCK_Q=64, BLOCK_K=64 | Task 1 Step 1 | | `FlashVarlenKernel` Rust struct with `new` + `forward` | Task 2 Step 1 | | `varlen_attention_cpu` with correct signature | Task 2 Step 1 | | Export from `src/kernels/mod.rs` | Task 3 Step 1 | | `SdpaBackend::VarLen` | Task 3 Steps 2–4 | | 8 tests — all CPU | Task 2 Step 1 (tests block) | | `cargo check` clean | Task 3 Step 5 | | `cargo test --lib` passes (50 total) | Task 3 Step 6 | | No new dependencies | Verified — only existing crate features used | ### Placeholder Scan No "TBD", "TODO", "implement later", or "similar to Task N" placeholders — all code blocks are complete and self-contained. ### Type Consistency - `varlen_attention_cpu(q: &[f32], k: &[f32], v: &[f32], cu_seqlens: &[usize], num_heads: usize, head_dim: usize, softmax_scale: f32, causal: bool) -> Vec` — used consistently in both the implementation (Task 2) and all 8 test calls. - `FlashVarlenKernel::new(ctx: &Arc) -> FlashResult` — used in Task 2. - `FlashVarlenKernel::forward(...)` — full 15-argument signature defined once in Task 2, re-exported in Task 3. - `SdpaBackend::VarLen` — added in Task 3 Step 2, Display arm in Step 3, `supports_backend` arm in Step 4 — three places, all consistent spelling.