diff --git a/csrc/cuda/attention/flash_attention_varlen_sm90.cu b/csrc/cuda/attention/flash_attention_varlen_sm90.cu new file mode 100644 index 00000000..5a0a019f --- /dev/null +++ b/csrc/cuda/attention/flash_attention_varlen_sm90.cu @@ -0,0 +1,1134 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Hopper (SM90) forward-only FlashAttention: causal masking, packed +// variable-length (cu_seqlens) input, attention-domain LSE export. +// +// Compute path is TMA (global->shared loads) + PTX `mma.sync.aligned.m16n8k16` +// (not literal `wgmma.mma_async`) -- the same combination every other "SM90" +// kernel in this repo already uses (fused_logp_sm90.cu, +// fused_linear_logp_sm90.cu). Softmax/PV math, tiling constants, and the +// mma.sync fragment layout are adapted from +// csrc/cuda/attention/prefix_shared_attention.cu (a working dense mma.sync +// attention kernel already in this repo); the TMA double-buffering pipeline +// is adapted from csrc/cuda/fused_linear_logp_sm90.cu. This is the +// cross-platform-baseline-matching forward kernel described in +// docs/operators/attention-varlen.md; validated against +// rl_engine/kernels/ops/triton/triton_attn.py's varlen path, which is the +// semantic reference. No backward pass this milestone. + +#include "../../utils/tma_utils.cuh" +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr int BLOCK_Q = 64; +constexpr int BLOCK_KV = 64; +constexpr int DIM = 128; // only head_dim supported this milestone +constexpr int NUM_WARPS = 4; +constexpr int WARP_SIZE = 32; +constexpr int TB_SIZE = NUM_WARPS * WARP_SIZE; +constexpr int STAGES = 2; // K/V double buffering + +constexpr int MMA_M = 16; +constexpr int MMA_N = 8; +constexpr int MMA_K = 16; +constexpr int WARP_Q = BLOCK_Q / NUM_WARPS; + +static_assert(WARP_Q % MMA_M == 0, "WARP_Q must be a multiple of MMA_M"); +static_assert(DIM % MMA_K == 0 && DIM % MMA_N == 0, "DIM must be a multiple of MMA_K/MMA_N"); +static_assert(BLOCK_KV % MMA_N == 0 && BLOCK_KV % MMA_K == 0, + "BLOCK_KV must be a multiple of MMA_N/MMA_K"); + +__device__ __host__ constexpr int cdiv(int a, int b) { return (a + b - 1) / b; } + +// Tensor-core / shared-memory helpers -- same PTX and fragment layout as +// csrc/cuda/attention/prefix_shared_attention.cu, validated on this repo's +// Hopper GPUs. Kept as a local copy (not shared via a header) matching how +// fused_linear_logp_sm90.cu keeps its own copy of the analogous helpers. +__device__ inline void ldmatrix_x4(uint32_t regs[4], uint32_t addr) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];" + : "=r"(regs[0]), "=r"(regs[1]), "=r"(regs[2]), "=r"(regs[3]) + : "r"(addr)); +} + +__device__ inline void ldmatrix_x4_trans(uint32_t regs[4], uint32_t addr) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];" + : "=r"(regs[0]), "=r"(regs[1]), "=r"(regs[2]), "=r"(regs[3]) + : "r"(addr)); +} + +__device__ inline void mma_m16n8k16(uint32_t A[4], uint32_t B[2], float D[4]) { + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, %13};" + : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), + "f"(D[0]), "f"(D[1]), "f"(D[2]), "f"(D[3])); +} + +// 2D bf16 tensor map, swizzle pinned to NONE: this kernel reads its tiles with +// plain row-major ldmatrix addressing (no manual XOR swizzle), so the TMA +// writes must be unswizzled too. Same rationale/pattern as +// fused_linear_logp_sm90.cu's helper of the same name (kept as a local copy, +// not shared, to match that file's own precedent). +inline void init_tensor_map_noswizzle(CUtensorMap *tmap, const nv_bfloat16 *gmem, + uint64_t gmem_height, uint64_t gmem_width, + uint32_t box_height, uint32_t box_width) { + uint64_t size[2] = {gmem_width, gmem_height}; + uint64_t stride[1] = {gmem_width * sizeof(nv_bfloat16)}; + uint32_t box[2] = {box_width, box_height}; + uint32_t elem_stride[2] = {1, 1}; + CUresult res = cuTensorMapEncodeTiled( + tmap, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, (void *)gmem, size, stride, box, elem_stride, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_NONE, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TORCH_CHECK(res == CUDA_SUCCESS, "cuTensorMapEncodeTiled failed for flash_attention_varlen_sm90"); +} + +// Q/K/V are packed [total_tokens, H, DIM] bf16, treated as logical 2D +// [total_tokens, H*DIM] matrices (row stride H*DIM elements -- the tensor's +// natural contiguous stride). Each TMA box spans one head's DIM-wide column +// slice (x = h*DIM) at a given packed-row offset (y = seq_start + tile_off). +// Cross-sequence-boundary reads (a KV tile near a short sequence's tail +// pulling in the next sequence's real K/V rows) are real and are handled by +// explicit seqlen_k/causal masking of the S=QK^T fragments below, not by +// making TMA itself sequence-aware -- the same design choice +// docs/operators/attention-varlen.md documents for the Triton kernel. +__global__ __launch_bounds__(TB_SIZE) void flash_attention_varlen_sm90_kernel( + const __grid_constant__ CUtensorMap q_tmap, + const __grid_constant__ CUtensorMap k_tmap, + const __grid_constant__ CUtensorMap v_tmap, + const int *__restrict__ cu_seqlens_q, + const int *__restrict__ cu_seqlens_k, + nv_bfloat16 *__restrict__ Out, // [total_q, H, DIM] + float *__restrict__ Lse, // [total_q, H] + int H, + float sm_scale, + bool causal) { + const int tid = threadIdx.x; + const int warp_id = tid / WARP_SIZE; + const int lane_id = tid % WARP_SIZE; + + const int q_block_id = blockIdx.x; + const int b = blockIdx.y / H; + const int h = blockIdx.y % H; + + const int q_start = cu_seqlens_q[b]; + const int seqlen_q = cu_seqlens_q[b + 1] - q_start; + if (q_block_id * BLOCK_Q >= seqlen_q) + return; // also makes zero-length sequences a no-op for every CTA assigned to them + + const int k_start = cu_seqlens_k[b]; + const int seqlen_k = cu_seqlens_k[b + 1] - k_start; + const int causal_offset = seqlen_k - seqlen_q; + + const int hi_rows = causal ? min(seqlen_k, (q_block_id + 1) * BLOCK_Q + causal_offset) + : seqlen_k; + const int num_kv_iter = cdiv(max(hi_rows, 0), BLOCK_KV); + // Row indices computed from warp_id/lane_id below are local to this + // Q-tile's shared-memory layout; add this to get the row's position + // within the sequence (what seqlen_q/causal comparisons and the + // Out/Lse global address need). + const int q_tile_row_base = q_block_id * BLOCK_Q; + + extern __shared__ __align__(1024) char smem[]; + nv_bfloat16 *sQ = reinterpret_cast(smem); + nv_bfloat16 *sK = sQ + BLOCK_Q * DIM; + nv_bfloat16 *sV = sK + STAGES * BLOCK_KV * DIM; + int *mbar_base = reinterpret_cast(sV + STAGES * BLOCK_KV * DIM); + + const uint64_t sQ_tma = __cvta_generic_to_shared(sQ); + const uint64_t sK_tma = __cvta_generic_to_shared(sK); + const uint64_t sV_tma = __cvta_generic_to_shared(sV); + const uint32_t sQ_base = static_cast(sQ_tma); + const uint32_t sK_base = static_cast(sK_tma); + const uint32_t sV_base = static_cast(sV_tma); + + const uint64_t mbar_q = __cvta_generic_to_shared(mbar_base); + uint64_t mbar_kv[STAGES]; +#pragma unroll + for (int s = 0; s < STAGES; ++s) + mbar_kv[s] = __cvta_generic_to_shared(mbar_base + 2 * (1 + s)); + + if (tid == 0) { + mbarrier_init(mbar_q, 1); +#pragma unroll + for (int s = 0; s < STAGES; ++s) + mbarrier_init(mbar_kv[s], 1); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + __syncthreads(); + + // Q: single TMA load for the whole CTA lifetime (Q doesn't change across + // the KV loop). + const uint32_t q_tile_bytes = BLOCK_Q * DIM * sizeof(nv_bfloat16); + if (tid == 0) { + mbarrier_arrive_expect_tx(mbar_q, q_tile_bytes); + tma_2d_g2s(sQ_tma, &q_tmap, h * DIM, q_start + q_block_id * BLOCK_Q, mbar_q); + } + mbarrier_wait(mbar_q, 0); + __syncthreads(); + + uint32_t Q_smem_thread; + { + const int row_off = warp_id * WARP_Q + (lane_id % 16); + const int col_off = (lane_id / 16) * 8; + Q_smem_thread = sQ_base + (row_off * DIM + col_off) * sizeof(nv_bfloat16); + } + uint32_t K_smem_thread; + { + const int row_off = lane_id % 8; + const int col_off = (lane_id / 8) * 8; + K_smem_thread = sK_base + (row_off * DIM + col_off) * sizeof(nv_bfloat16); + } + uint32_t V_smem_thread; + { + const int row_off = lane_id % 16; + const int col_off = (lane_id / 16) * 8; + V_smem_thread = sV_base + (row_off * DIM + col_off) * sizeof(nv_bfloat16); + } + + uint32_t Q_rmem[WARP_Q / MMA_M][DIM / MMA_K][4]; +#pragma unroll + for (int mi = 0; mi < WARP_Q / MMA_M; mi++) +#pragma unroll + for (int kd = 0; kd < DIM / MMA_K; kd++) { + uint32_t addr = + Q_smem_thread + mi * MMA_M * DIM * sizeof(nv_bfloat16) + kd * MMA_K * sizeof(nv_bfloat16); + ldmatrix_x4(Q_rmem[mi][kd], addr); + } + __syncthreads(); + + if (num_kv_iter == 0) { + // No valid keys for this Q-tile (e.g. seqlen_k == 0). Write a + // degenerate all-masked result rather than leaving Out/Lse + // uninitialized: out=0, lse=-inf-ish sentinel via log(tiny). +#pragma unroll + for (int mi = 0; mi < WARP_Q / MMA_M; mi++) { + const int row0 = q_tile_row_base + warp_id * WARP_Q + mi * MMA_M + lane_id / 4; + if (lane_id % 4 == 0) { + const float sentinel_lse = -CUDART_INF_F; + if (row0 < seqlen_q) + Lse[(q_start + row0) * H + h] = sentinel_lse; + if (row0 + 8 < seqlen_q) + Lse[(q_start + row0 + 8) * H + h] = sentinel_lse; + } +#pragma unroll + for (int d = 0; d < DIM / MMA_N; d++) { + const int col = d * MMA_N + (lane_id % 4) * 2; + nv_bfloat162 zero2 = __float22bfloat162_rn({0.0f, 0.0f}); + if (row0 < seqlen_q) + *reinterpret_cast(Out + ((q_start + row0) * H + h) * DIM + col) = + zero2; + if (row0 + 8 < seqlen_q) + *reinterpret_cast(Out + ((q_start + row0 + 8) * H + h) * DIM + + col) = zero2; + } + } + return; + } + + const uint32_t kv_tile_bytes = 2 * BLOCK_KV * DIM * sizeof(nv_bfloat16); // K+V combined + auto issue_kv = [&](int kv_id) { + if (kv_id < 0 || kv_id >= num_kv_iter) + return; + const int buf = kv_id % STAGES; + const uint32_t off = buf * BLOCK_KV * DIM * sizeof(nv_bfloat16); + mbarrier_arrive_expect_tx(mbar_kv[buf], kv_tile_bytes); + tma_2d_g2s(sK_tma + off, &k_tmap, h * DIM, k_start + kv_id * BLOCK_KV, mbar_kv[buf]); + tma_2d_g2s(sV_tma + off, &v_tmap, h * DIM, k_start + kv_id * BLOCK_KV, mbar_kv[buf]); + }; + + int phase[STAGES]; +#pragma unroll + for (int s = 0; s < STAGES; ++s) + phase[s] = 0; + if (tid == 0) { +#pragma unroll + for (int s = 0; s < STAGES - 1; ++s) + issue_kv(s); + } + + uint32_t K_rmem[BLOCK_KV / MMA_N][DIM / MMA_K][2]; + uint32_t P_rmem[WARP_Q / MMA_M][BLOCK_KV / MMA_K][4]; + uint32_t V_rmem[BLOCK_KV / MMA_K][DIM / MMA_N][2]; + + float O_rmem[WARP_Q / MMA_M][DIM / MMA_N][4] = {}; + float rowmax[WARP_Q / MMA_M][2]; + float rowsumexp[WARP_Q / MMA_M][2] = {}; +#pragma unroll + for (int mi = 0; mi < WARP_Q / MMA_M; mi++) { + rowmax[mi][0] = -FLT_MAX; + rowmax[mi][1] = -FLT_MAX; + } + + for (int kv_id = 0; kv_id < num_kv_iter; kv_id++) { + const int buf = kv_id % STAGES; + const uint32_t buf_off = buf * BLOCK_KV * DIM * sizeof(nv_bfloat16); + + if (tid == 0) + issue_kv(kv_id + STAGES - 1); + + mbarrier_wait(mbar_kv[buf], phase[buf]); + phase[buf] ^= 1; + __syncthreads(); + + float S_rmem[WARP_Q / MMA_M][BLOCK_KV / MMA_N][4] = {}; + +#pragma unroll + for (int nkv = 0; nkv < BLOCK_KV / MMA_N; nkv++) +#pragma unroll + for (int kd = 0; kd < DIM / MMA_K; kd += 2) { + uint32_t addr = K_smem_thread + buf_off; + addr += nkv * MMA_N * DIM * sizeof(nv_bfloat16); + addr += kd * MMA_K * sizeof(nv_bfloat16); + ldmatrix_x4(K_rmem[nkv][kd], addr); + } + +#pragma unroll + for (int mi = 0; mi < WARP_Q / MMA_M; mi++) +#pragma unroll + for (int nkv = 0; nkv < BLOCK_KV / MMA_N; nkv++) +#pragma unroll + for (int kd = 0; kd < DIM / MMA_K; kd++) + mma_m16n8k16(Q_rmem[mi][kd], K_rmem[nkv][kd], S_rmem[mi][nkv]); + +#pragma unroll + for (int mi = 0; mi < WARP_Q / MMA_M; mi++) { + const int row_a = q_tile_row_base + warp_id * WARP_Q + mi * MMA_M + lane_id / 4; + const int row_b = row_a + 8; + +#pragma unroll + for (int nkv = 0; nkv < BLOCK_KV / MMA_N; nkv++) { + float *regs = S_rmem[mi][nkv]; + const int col_a = kv_id * BLOCK_KV + nkv * MMA_N + (lane_id % 4) * 2; + const int col_b = col_a + 1; + regs[0] *= sm_scale; + regs[1] *= sm_scale; + regs[2] *= sm_scale; + regs[3] *= sm_scale; + if (col_a >= seqlen_k || (causal && row_a < col_a - causal_offset)) + regs[0] = -CUDART_INF_F; + if (col_b >= seqlen_k || (causal && row_a < col_b - causal_offset)) + regs[1] = -CUDART_INF_F; + if (col_a >= seqlen_k || (causal && row_b < col_a - causal_offset)) + regs[2] = -CUDART_INF_F; + if (col_b >= seqlen_k || (causal && row_b < col_b - causal_offset)) + regs[3] = -CUDART_INF_F; + } + + float this_rowmax[2]; +#pragma unroll + for (int nkv = 0; nkv < BLOCK_KV / MMA_N; nkv++) { + float *regs = S_rmem[mi][nkv]; + if (nkv == 0) { + this_rowmax[0] = max(regs[0], regs[1]); + this_rowmax[1] = max(regs[2], regs[3]); + } else { + this_rowmax[0] = max(this_rowmax[0], max(regs[0], regs[1])); + this_rowmax[1] = max(this_rowmax[1], max(regs[2], regs[3])); + } + } + this_rowmax[0] = max(this_rowmax[0], __shfl_xor_sync(0xFFFFFFFFu, this_rowmax[0], 1)); + this_rowmax[0] = max(this_rowmax[0], __shfl_xor_sync(0xFFFFFFFFu, this_rowmax[0], 2)); + this_rowmax[1] = max(this_rowmax[1], __shfl_xor_sync(0xFFFFFFFFu, this_rowmax[1], 1)); + this_rowmax[1] = max(this_rowmax[1], __shfl_xor_sync(0xFFFFFFFFu, this_rowmax[1], 2)); + this_rowmax[0] = max(this_rowmax[0], rowmax[mi][0]); + this_rowmax[1] = max(this_rowmax[1], rowmax[mi][1]); + + float rescale[2]; + rescale[0] = __expf(rowmax[mi][0] - this_rowmax[0]); + rescale[1] = __expf(rowmax[mi][1] - this_rowmax[1]); +#pragma unroll + for (int d = 0; d < DIM / MMA_N; d++) { + O_rmem[mi][d][0] *= rescale[0]; + O_rmem[mi][d][1] *= rescale[0]; + O_rmem[mi][d][2] *= rescale[1]; + O_rmem[mi][d][3] *= rescale[1]; + } + rowmax[mi][0] = this_rowmax[0]; + rowmax[mi][1] = this_rowmax[1]; + + float this_rowsumexp[2]; +#pragma unroll + for (int nkv = 0; nkv < BLOCK_KV / MMA_N; nkv++) { + float *regs = S_rmem[mi][nkv]; + regs[0] = __expf(regs[0] - rowmax[mi][0]); + regs[1] = __expf(regs[1] - rowmax[mi][0]); + regs[2] = __expf(regs[2] - rowmax[mi][1]); + regs[3] = __expf(regs[3] - rowmax[mi][1]); + if (nkv == 0) { + this_rowsumexp[0] = regs[0] + regs[1]; + this_rowsumexp[1] = regs[2] + regs[3]; + } else { + this_rowsumexp[0] += regs[0] + regs[1]; + this_rowsumexp[1] += regs[2] + regs[3]; + } + + nv_bfloat162 *pfrag = reinterpret_cast(P_rmem[mi][nkv / 2]); + pfrag[(nkv % 2) * 2] = __float22bfloat162_rn({regs[0], regs[1]}); + pfrag[(nkv % 2) * 2 + 1] = __float22bfloat162_rn({regs[2], regs[3]}); + } + this_rowsumexp[0] += __shfl_xor_sync(0xFFFFFFFFu, this_rowsumexp[0], 1); + this_rowsumexp[0] += __shfl_xor_sync(0xFFFFFFFFu, this_rowsumexp[0], 2); + this_rowsumexp[1] += __shfl_xor_sync(0xFFFFFFFFu, this_rowsumexp[1], 1); + this_rowsumexp[1] += __shfl_xor_sync(0xFFFFFFFFu, this_rowsumexp[1], 2); + rowsumexp[mi][0] = rowsumexp[mi][0] * rescale[0] + this_rowsumexp[0]; + rowsumexp[mi][1] = rowsumexp[mi][1] * rescale[1] + this_rowsumexp[1]; + } + +#pragma unroll + for (int nkv = 0; nkv < BLOCK_KV / MMA_K; nkv++) +#pragma unroll + for (int d = 0; d < DIM / MMA_N; d += 2) { + uint32_t addr = V_smem_thread + buf_off; + addr += nkv * MMA_K * DIM * sizeof(nv_bfloat16); + addr += d * MMA_N * sizeof(nv_bfloat16); + ldmatrix_x4_trans(V_rmem[nkv][d], addr); + } + +#pragma unroll + for (int mi = 0; mi < WARP_Q / MMA_M; mi++) +#pragma unroll + for (int d = 0; d < DIM / MMA_N; d++) +#pragma unroll + for (int nkv = 0; nkv < BLOCK_KV / MMA_K; nkv++) + mma_m16n8k16(P_rmem[mi][nkv], V_rmem[nkv][d], O_rmem[mi][d]); + + __syncthreads(); + } + +#pragma unroll + for (int mi = 0; mi < WARP_Q / MMA_M; mi++) { + const int row0 = q_tile_row_base + warp_id * WARP_Q + mi * MMA_M + lane_id / 4; + const float inv0 = 1.0f / fmaxf(rowsumexp[mi][0], 1e-30f); + const float inv1 = 1.0f / fmaxf(rowsumexp[mi][1], 1e-30f); + + if (lane_id % 4 == 0) { + const float lse0 = rowmax[mi][0] + logf(fmaxf(rowsumexp[mi][0], 1e-30f)); + const float lse1 = rowmax[mi][1] + logf(fmaxf(rowsumexp[mi][1], 1e-30f)); + if (row0 < seqlen_q) + Lse[(q_start + row0) * H + h] = lse0; + if (row0 + 8 < seqlen_q) + Lse[(q_start + row0 + 8) * H + h] = lse1; + } + +#pragma unroll + for (int d = 0; d < DIM / MMA_N; d++) { + const int col = d * MMA_N + (lane_id % 4) * 2; + float *regs = O_rmem[mi][d]; + regs[0] *= inv0; + regs[1] *= inv0; + regs[2] *= inv1; + regs[3] *= inv1; + + if (row0 < seqlen_q) + *reinterpret_cast(Out + ((q_start + row0) * H + h) * DIM + col) = + __float22bfloat162_rn({regs[0], regs[1]}); + if (row0 + 8 < seqlen_q) + *reinterpret_cast(Out + ((q_start + row0 + 8) * H + h) * DIM + col) = + __float22bfloat162_rn({regs[2], regs[3]}); + } + } +} + +// Delta[row,h] = sum_d(Out[row,h,d] * DO[row,h,d]), the per-query-row term +// the backward recompute needs (`_bwd_preprocess_varlen` in triton_attn.py). +// One block per (row, head); DIM=128 threads, one per element, block reduce. +__global__ void flash_attention_bwd_preprocess_kernel( + const nv_bfloat16 *__restrict__ Out, + const nv_bfloat16 *__restrict__ DO, + float *__restrict__ Delta, + int total_q, + int H) { + const int row = blockIdx.x; + const int h = blockIdx.y; + const int tid = threadIdx.x; + if (row >= total_q) + return; + + const int base = (row * H + h) * DIM; + float val = __bfloat162float(Out[base + tid]) * __bfloat162float(DO[base + tid]); + +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + val += __shfl_xor_sync(0xFFFFFFFFu, val, off); + + __shared__ float warp_sums[DIM / WARP_SIZE]; + const int warp = tid / WARP_SIZE; + const int lane = tid % WARP_SIZE; + if (lane == 0) + warp_sums[warp] = val; + __syncthreads(); + if (tid == 0) { + float total = 0.0f; +#pragma unroll + for (int w = 0; w < DIM / WARP_SIZE; w++) + total += warp_sums[w]; + Delta[row * H + h] = total; + } +} + +// Backward: grid over KV-tiles (the CTA's fixed, exclusively-owned dK/dV +// range), looping over Q-tiles internally -- the transpose of forward's +// grid-over-Q-tiles structure. Mirrors `_bwd_kernel_varlen` in +// triton_attn.py exactly (same lo-bound causal skip, same recompute-based +// algorithm using the saved `Lse` instead of an online rescale). +// +// P/ds only ever need to be read back as an mma.sync "A" operand (never +// "B"), so rather than using `ldmatrix_x4_trans` on their natural +// [query,key] layout (which is the trick this file uses for tensors sourced +// directly from HBM via TMA, e.g. V/dO/Q's second, transposed reads), P/ds +// are written to a small shared-memory scratch tile already +// row/col-swapped (key-major) by the store itself -- 4 independent scalar +// stores per thread, not vectorized, since the transposed target addresses +// aren't adjacent -- then read back with a plain (non-trans) `ldmatrix_x4` +// using the same addressing *shape* as this file's Q-as-A-operand load, +// substituting BLOCK_KV/BLOCK_Q for BLOCK_Q/DIM's roles. +__global__ __launch_bounds__(TB_SIZE) void flash_attention_varlen_sm90_bwd_kernel( + const __grid_constant__ CUtensorMap q_tmap, + const __grid_constant__ CUtensorMap k_tmap, + const __grid_constant__ CUtensorMap v_tmap, + const __grid_constant__ CUtensorMap do_tmap, + const int *__restrict__ cu_seqlens_q, + const int *__restrict__ cu_seqlens_k, + const float *__restrict__ Lse, // [total_q, H] + const float *__restrict__ Delta, // [total_q, H] + nv_bfloat16 *__restrict__ DQ, // [total_q, H, DIM] -- atomic-add target + nv_bfloat16 *__restrict__ DK, // [total_k, H, DIM] + nv_bfloat16 *__restrict__ DV, // [total_k, H, DIM] + int H, + float sm_scale, + bool causal) { + const int tid = threadIdx.x; + const int warp_id = tid / WARP_SIZE; + const int lane_id = tid % WARP_SIZE; + + const int kv_block_id = blockIdx.x; + const int b = blockIdx.y / H; + const int h = blockIdx.y % H; + + const int k_start = cu_seqlens_k[b]; + const int seqlen_k = cu_seqlens_k[b + 1] - k_start; + if (kv_block_id * BLOCK_KV >= seqlen_k) + return; + + const int q_start = cu_seqlens_q[b]; + const int seqlen_q = cu_seqlens_q[b + 1] - q_start; + const int causal_offset = seqlen_k - seqlen_q; + const int kv_tile_row_base = kv_block_id * BLOCK_KV; + + // Query rows before `kv_tile_row_base - causal_offset` never attend to + // this KV tile; the Q-loop below can start there instead of row 0. If + // this pushes `lo >= seqlen_q` the loop below runs zero iterations and + // dk/dv (register-accumulated, initialized to zero) are still written + // unconditionally at the end -- correctly zero, no special case needed. + int lo = causal ? max(0, kv_tile_row_base - causal_offset) : 0; + lo = (lo / BLOCK_Q) * BLOCK_Q; + + extern __shared__ __align__(1024) char smem[]; + nv_bfloat16 *sK = reinterpret_cast(smem); + nv_bfloat16 *sV = sK + BLOCK_KV * DIM; + nv_bfloat16 *sQ = sV + BLOCK_KV * DIM; + nv_bfloat16 *sDO = sQ + STAGES * BLOCK_Q * DIM; + nv_bfloat16 *sPT = sDO + STAGES * BLOCK_Q * DIM; // [BLOCK_KV][BLOCK_Q] transpose scratch + int *mbar_base = reinterpret_cast(sPT + BLOCK_KV * BLOCK_Q); + + const uint64_t sK_tma = __cvta_generic_to_shared(sK); + const uint64_t sV_tma = __cvta_generic_to_shared(sV); + const uint64_t sQ_tma = __cvta_generic_to_shared(sQ); + const uint64_t sDO_tma = __cvta_generic_to_shared(sDO); + const uint32_t sK_base = static_cast(sK_tma); + const uint32_t sV_base = static_cast(sV_tma); + const uint32_t sQ_base = static_cast(sQ_tma); + const uint32_t sDO_base = static_cast(sDO_tma); + const uint32_t sPT_base = static_cast(__cvta_generic_to_shared(sPT)); + + const uint64_t mbar_kv = __cvta_generic_to_shared(mbar_base); + uint64_t mbar_qdo[STAGES]; +#pragma unroll + for (int s = 0; s < STAGES; ++s) + mbar_qdo[s] = __cvta_generic_to_shared(mbar_base + 2 * (1 + s)); + + if (tid == 0) { + mbarrier_init(mbar_kv, 1); +#pragma unroll + for (int s = 0; s < STAGES; ++s) + mbarrier_init(mbar_qdo[s], 1); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + __syncthreads(); + + // K + V: single TMA load for this CTA's fixed KV-tile lifetime. + const uint32_t kv_tile_bytes = 2 * BLOCK_KV * DIM * sizeof(nv_bfloat16); + if (tid == 0) { + mbarrier_arrive_expect_tx(mbar_kv, kv_tile_bytes); + tma_2d_g2s(sK_tma, &k_tmap, h * DIM, k_start + kv_tile_row_base, mbar_kv); + tma_2d_g2s(sV_tma, &v_tmap, h * DIM, k_start + kv_tile_row_base, mbar_kv); + } + mbarrier_wait(mbar_kv, 0); + __syncthreads(); + + // K/V "natural" addressing (this file's forward-K-in-QK^T role: B + // operand, contracts over D, plain ldmatrix_x4). + uint32_t K_smem_thread_nat; + { + const int row_off = lane_id % 8; + const int col_off = (lane_id / 8) * 8; + K_smem_thread_nat = sK_base + (row_off * DIM + col_off) * sizeof(nv_bfloat16); + } + uint32_t V_smem_thread_nat; + { + const int row_off = lane_id % 8; + const int col_off = (lane_id / 8) * 8; + V_smem_thread_nat = sV_base + (row_off * DIM + col_off) * sizeof(nv_bfloat16); + } + // K "transposed" addressing (this file's forward-V-in-P@V role: B + // operand, contracts over the KV/seq dimension, ldmatrix_x4_trans) -- + // used for dQ = ds @ K. + uint32_t K_smem_thread_trans; + { + const int row_off = lane_id % 16; + const int col_off = (lane_id / 16) * 8; + K_smem_thread_trans = sK_base + (row_off * DIM + col_off) * sizeof(nv_bfloat16); + } + // Q/DO "natural" addressing (forward-Q-in-QK^T role: A operand). + uint32_t Q_smem_thread_nat; + { + const int row_off = warp_id * WARP_Q + (lane_id % 16); + const int col_off = (lane_id / 16) * 8; + Q_smem_thread_nat = sQ_base + (row_off * DIM + col_off) * sizeof(nv_bfloat16); + } + uint32_t DO_smem_thread_nat; + { + const int row_off = warp_id * WARP_Q + (lane_id % 16); + const int col_off = (lane_id / 16) * 8; + DO_smem_thread_nat = sDO_base + (row_off * DIM + col_off) * sizeof(nv_bfloat16); + } + // Q/DO "transposed" addressing (forward-V-in-P@V role: B operand, + // contracts over the query dimension) -- Q for dK = ds^T @ Q, DO for + // dV = P^T @ dO. + uint32_t Q_smem_thread_trans; + { + const int row_off = lane_id % 16; + const int col_off = (lane_id / 16) * 8; + Q_smem_thread_trans = sQ_base + (row_off * DIM + col_off) * sizeof(nv_bfloat16); + } + uint32_t DO_smem_thread_trans; + { + const int row_off = lane_id % 16; + const int col_off = (lane_id / 16) * 8; + DO_smem_thread_trans = sDO_base + (row_off * DIM + col_off) * sizeof(nv_bfloat16); + } + // sPT[BLOCK_KV][BLOCK_Q] read-as-A-operand addressing: same shape as + // this file's Q-as-A-operand read, with BLOCK_Q substituted for DIM as + // the row stride (sPT's rows only span BLOCK_KV, not the full DIM) and + // BLOCK_KV substituted for BLOCK_Q as the row extent being tiled across + // warps (numerically identical here since BLOCK_Q == BLOCK_KV == 64). + uint32_t PT_smem_thread; + { + const int row_off = warp_id * WARP_Q + (lane_id % 16); + const int col_off = (lane_id / 16) * 8; + PT_smem_thread = sPT_base + (row_off * BLOCK_Q + col_off) * sizeof(nv_bfloat16); + } + + uint32_t K_rmem_nat[BLOCK_KV / MMA_N][DIM / MMA_K][2]; +#pragma unroll + for (int nkv = 0; nkv < BLOCK_KV / MMA_N; nkv++) +#pragma unroll + for (int kd = 0; kd < DIM / MMA_K; kd += 2) { + uint32_t addr = K_smem_thread_nat + nkv * MMA_N * DIM * sizeof(nv_bfloat16) + + kd * MMA_K * sizeof(nv_bfloat16); + ldmatrix_x4(K_rmem_nat[nkv][kd], addr); + } + uint32_t V_rmem_nat[BLOCK_KV / MMA_N][DIM / MMA_K][2]; +#pragma unroll + for (int nkv = 0; nkv < BLOCK_KV / MMA_N; nkv++) +#pragma unroll + for (int kd = 0; kd < DIM / MMA_K; kd += 2) { + uint32_t addr = V_smem_thread_nat + nkv * MMA_N * DIM * sizeof(nv_bfloat16) + + kd * MMA_K * sizeof(nv_bfloat16); + ldmatrix_x4(V_rmem_nat[nkv][kd], addr); + } + uint32_t K_rmem_trans[BLOCK_KV / MMA_K][DIM / MMA_N][2]; +#pragma unroll + for (int nkv = 0; nkv < BLOCK_KV / MMA_K; nkv++) +#pragma unroll + for (int d = 0; d < DIM / MMA_N; d += 2) { + uint32_t addr = K_smem_thread_trans + nkv * MMA_K * DIM * sizeof(nv_bfloat16) + + d * MMA_N * sizeof(nv_bfloat16); + ldmatrix_x4_trans(K_rmem_trans[nkv][d], addr); + } + + const uint32_t qdo_tile_bytes = 2 * BLOCK_Q * DIM * sizeof(nv_bfloat16); + auto issue_qdo = [&](int q_id) { + const int start_m = lo + q_id * BLOCK_Q; + if (start_m < 0 || start_m >= seqlen_q) + return; + const int buf = q_id % STAGES; + const uint32_t off = buf * BLOCK_Q * DIM * sizeof(nv_bfloat16); + mbarrier_arrive_expect_tx(mbar_qdo[buf], qdo_tile_bytes); + tma_2d_g2s(sQ_tma + off, &q_tmap, h * DIM, q_start + start_m, mbar_qdo[buf]); + tma_2d_g2s(sDO_tma + off, &do_tmap, h * DIM, q_start + start_m, mbar_qdo[buf]); + }; + + int phase[STAGES]; +#pragma unroll + for (int s = 0; s < STAGES; ++s) + phase[s] = 0; + if (tid == 0) { +#pragma unroll + for (int s = 0; s < STAGES - 1; ++s) + issue_qdo(s); + } + + float dk_out[WARP_Q / MMA_M][DIM / MMA_N][4] = {}; + float dv_out[WARP_Q / MMA_M][DIM / MMA_N][4] = {}; + + int q_id = 0; + for (int start_m = lo; start_m < seqlen_q; start_m += BLOCK_Q, q_id++) { + const int buf = q_id % STAGES; + const uint32_t buf_off = buf * BLOCK_Q * DIM * sizeof(nv_bfloat16); + + if (tid == 0) + issue_qdo(q_id + STAGES - 1); + + mbarrier_wait(mbar_qdo[buf], phase[buf]); + phase[buf] ^= 1; + __syncthreads(); + + uint32_t Q_rmem_nat[WARP_Q / MMA_M][DIM / MMA_K][4]; +#pragma unroll + for (int mi = 0; mi < WARP_Q / MMA_M; mi++) +#pragma unroll + for (int kd = 0; kd < DIM / MMA_K; kd++) { + uint32_t addr = Q_smem_thread_nat + buf_off + mi * MMA_M * DIM * sizeof(nv_bfloat16) + + kd * MMA_K * sizeof(nv_bfloat16); + ldmatrix_x4(Q_rmem_nat[mi][kd], addr); + } + + // S = Q @ K^T (identical pattern to forward's QK^T). + float S_rmem[WARP_Q / MMA_M][BLOCK_KV / MMA_N][4] = {}; +#pragma unroll + for (int mi = 0; mi < WARP_Q / MMA_M; mi++) +#pragma unroll + for (int nkv = 0; nkv < BLOCK_KV / MMA_N; nkv++) +#pragma unroll + for (int kd = 0; kd < DIM / MMA_K; kd++) + mma_m16n8k16(Q_rmem_nat[mi][kd], K_rmem_nat[nkv][kd], S_rmem[mi][nkv]); + + // dP = DO @ V^T (same operand shapes as S = Q @ K^T: DO as A, V-nat as B). + uint32_t DO_rmem_nat[WARP_Q / MMA_M][DIM / MMA_K][4]; +#pragma unroll + for (int mi = 0; mi < WARP_Q / MMA_M; mi++) +#pragma unroll + for (int kd = 0; kd < DIM / MMA_K; kd++) { + uint32_t addr = DO_smem_thread_nat + buf_off + + mi * MMA_M * DIM * sizeof(nv_bfloat16) + kd * MMA_K * sizeof(nv_bfloat16); + ldmatrix_x4(DO_rmem_nat[mi][kd], addr); + } + float dP_rmem[WARP_Q / MMA_M][BLOCK_KV / MMA_N][4] = {}; +#pragma unroll + for (int mi = 0; mi < WARP_Q / MMA_M; mi++) +#pragma unroll + for (int nkv = 0; nkv < BLOCK_KV / MMA_N; nkv++) +#pragma unroll + for (int kd = 0; kd < DIM / MMA_K; kd++) + mma_m16n8k16(DO_rmem_nat[mi][kd], V_rmem_nat[nkv][kd], dP_rmem[mi][nkv]); + + uint32_t P_pack[WARP_Q / MMA_M][BLOCK_KV / MMA_K][4]; + uint32_t ds_pack[WARP_Q / MMA_M][BLOCK_KV / MMA_K][4]; + +#pragma unroll + for (int mi = 0; mi < WARP_Q / MMA_M; mi++) { + const int row_a = start_m + warp_id * WARP_Q + mi * MMA_M + lane_id / 4; + const int row_b = row_a + 8; + const float lse_a = (row_a < seqlen_q) ? Lse[(q_start + row_a) * H + h] : 0.0f; + const float lse_b = (row_b < seqlen_q) ? Lse[(q_start + row_b) * H + h] : 0.0f; + const float delta_a = (row_a < seqlen_q) ? Delta[(q_start + row_a) * H + h] : 0.0f; + const float delta_b = (row_b < seqlen_q) ? Delta[(q_start + row_b) * H + h] : 0.0f; + +#pragma unroll + for (int nkv = 0; nkv < BLOCK_KV / MMA_N; nkv++) { + float *s = S_rmem[mi][nkv]; + float *dp = dP_rmem[mi][nkv]; + // Local (0..BLOCK_KV-1, for sPT scratch addressing) vs. + // global-within-sequence (for seqlen_k/causal comparisons) + // column index -- sPT is only BLOCK_KV rows, so it must be + // indexed local-to-this-CTA's-KV-tile, not by the global + // column (which can exceed BLOCK_KV once kv_tile_row_base > 0). + const int col_local_a = nkv * MMA_N + (lane_id % 4) * 2; + const int col_local_b = col_local_a + 1; + const int col_a = kv_tile_row_base + col_local_a; + const int col_b = col_a + 1; + + const bool keep00 = col_a < seqlen_k && row_a < seqlen_q && + (!causal || row_a >= col_a - causal_offset); + const bool keep01 = col_b < seqlen_k && row_a < seqlen_q && + (!causal || row_a >= col_b - causal_offset); + const bool keep10 = col_a < seqlen_k && row_b < seqlen_q && + (!causal || row_b >= col_a - causal_offset); + const bool keep11 = col_b < seqlen_k && row_b < seqlen_q && + (!causal || row_b >= col_b - causal_offset); + + const float p00 = keep00 ? __expf(s[0] * sm_scale - lse_a) : 0.0f; + const float p01 = keep01 ? __expf(s[1] * sm_scale - lse_a) : 0.0f; + const float p10 = keep10 ? __expf(s[2] * sm_scale - lse_b) : 0.0f; + const float p11 = keep11 ? __expf(s[3] * sm_scale - lse_b) : 0.0f; + + const float ds00 = p00 * (dp[0] - delta_a) * sm_scale; + const float ds01 = p01 * (dp[1] - delta_a) * sm_scale; + const float ds10 = p10 * (dp[2] - delta_b) * sm_scale; + const float ds11 = p11 * (dp[3] - delta_b) * sm_scale; + + // Natural-orientation ds pack (A operand for dQ = ds @ K), + // same 2-nkv-per-slot packing convention forward uses for P. + nv_bfloat162 *ds_frag = reinterpret_cast(ds_pack[mi][nkv / 2]); + ds_frag[(nkv % 2) * 2] = __float22bfloat162_rn({ds00, ds01}); + ds_frag[(nkv % 2) * 2 + 1] = __float22bfloat162_rn({ds10, ds11}); + + // Transposed-write scratch for P (this iteration's dV) and, + // after dV is consumed below, ds (this iteration's dK) -- + // key-major: sPT[key][query], 4 independent scalar stores. + const int local_row_a = warp_id * WARP_Q + mi * MMA_M + lane_id / 4; + const int local_row_b = local_row_a + 8; + sPT[col_local_a * BLOCK_Q + local_row_a] = __float2bfloat16_rn(p00); + sPT[col_local_b * BLOCK_Q + local_row_a] = __float2bfloat16_rn(p01); + sPT[col_local_a * BLOCK_Q + local_row_b] = __float2bfloat16_rn(p10); + sPT[col_local_b * BLOCK_Q + local_row_b] = __float2bfloat16_rn(p11); + } + } + __syncthreads(); + + // dV += P^T @ DO (P read back key-major/natural from sPT as A; + // DO's transposed read as B). + uint32_t PT_rmem[WARP_Q / MMA_M][BLOCK_Q / MMA_K][4]; +#pragma unroll + for (int mi = 0; mi < WARP_Q / MMA_M; mi++) +#pragma unroll + for (int kd = 0; kd < BLOCK_Q / MMA_K; kd++) { + uint32_t addr = + PT_smem_thread + mi * MMA_M * BLOCK_Q * sizeof(nv_bfloat16) + kd * MMA_K * sizeof(nv_bfloat16); + ldmatrix_x4(PT_rmem[mi][kd], addr); + } + uint32_t DO_rmem_trans[BLOCK_Q / MMA_K][DIM / MMA_N][2]; +#pragma unroll + for (int nkv = 0; nkv < BLOCK_Q / MMA_K; nkv++) +#pragma unroll + for (int d = 0; d < DIM / MMA_N; d += 2) { + uint32_t addr = DO_smem_thread_trans + buf_off + + nkv * MMA_K * DIM * sizeof(nv_bfloat16) + d * MMA_N * sizeof(nv_bfloat16); + ldmatrix_x4_trans(DO_rmem_trans[nkv][d], addr); + } +#pragma unroll + for (int mi = 0; mi < WARP_Q / MMA_M; mi++) +#pragma unroll + for (int d = 0; d < DIM / MMA_N; d++) +#pragma unroll + for (int kd = 0; kd < BLOCK_Q / MMA_K; kd++) + mma_m16n8k16(PT_rmem[mi][kd], DO_rmem_trans[kd][d], dv_out[mi][d]); + + __syncthreads(); // sPT fully consumed by every thread before ds overwrites it + + // Re-issue the ds transposed write into the same sPT scratch. +#pragma unroll + for (int mi = 0; mi < WARP_Q / MMA_M; mi++) { + const int local_row_a = warp_id * WARP_Q + mi * MMA_M + lane_id / 4; + const int local_row_b = local_row_a + 8; +#pragma unroll + for (int nkv = 0; nkv < BLOCK_KV / MMA_N; nkv++) { + // Local (0..BLOCK_KV-1) column -- sPT-only, no masking use + // here, so no global-within-sequence variant is needed. + const int col_a = nkv * MMA_N + (lane_id % 4) * 2; + const int col_b = col_a + 1; + const nv_bfloat162 *ds_frag = + reinterpret_cast(ds_pack[mi][nkv / 2]) + (nkv % 2) * 2; + const nv_bfloat162 v0 = ds_frag[0]; + const nv_bfloat162 v1 = ds_frag[1]; + sPT[col_a * BLOCK_Q + local_row_a] = v0.x; + sPT[col_b * BLOCK_Q + local_row_a] = v0.y; + sPT[col_a * BLOCK_Q + local_row_b] = v1.x; + sPT[col_b * BLOCK_Q + local_row_b] = v1.y; + } + } + __syncthreads(); + + // dK += ds^T @ Q (ds read back key-major/natural from sPT as A; + // Q's transposed read as B). + uint32_t dsT_rmem[WARP_Q / MMA_M][BLOCK_Q / MMA_K][4]; +#pragma unroll + for (int mi = 0; mi < WARP_Q / MMA_M; mi++) +#pragma unroll + for (int kd = 0; kd < BLOCK_Q / MMA_K; kd++) { + uint32_t addr = + PT_smem_thread + mi * MMA_M * BLOCK_Q * sizeof(nv_bfloat16) + kd * MMA_K * sizeof(nv_bfloat16); + ldmatrix_x4(dsT_rmem[mi][kd], addr); + } + uint32_t Q_rmem_trans[BLOCK_Q / MMA_K][DIM / MMA_N][2]; +#pragma unroll + for (int nkv = 0; nkv < BLOCK_Q / MMA_K; nkv++) +#pragma unroll + for (int d = 0; d < DIM / MMA_N; d += 2) { + uint32_t addr = Q_smem_thread_trans + buf_off + + nkv * MMA_K * DIM * sizeof(nv_bfloat16) + d * MMA_N * sizeof(nv_bfloat16); + ldmatrix_x4_trans(Q_rmem_trans[nkv][d], addr); + } +#pragma unroll + for (int mi = 0; mi < WARP_Q / MMA_M; mi++) +#pragma unroll + for (int d = 0; d < DIM / MMA_N; d++) +#pragma unroll + for (int kd = 0; kd < BLOCK_Q / MMA_K; kd++) + mma_m16n8k16(dsT_rmem[mi][kd], Q_rmem_trans[kd][d], dk_out[mi][d]); + + // dQ += ds @ K (ds natural/registers as A -- no transpose needed, + // output row = query matches ds's natural row already; K's + // transposed read as B). Atomic-add straight into bf16 DQ (native + // atom.add.noftz.bf16 on SM90). + float dQ_acc[WARP_Q / MMA_M][DIM / MMA_N][4] = {}; +#pragma unroll + for (int mi = 0; mi < WARP_Q / MMA_M; mi++) +#pragma unroll + for (int d = 0; d < DIM / MMA_N; d++) +#pragma unroll + for (int kd = 0; kd < BLOCK_KV / MMA_K; kd++) + mma_m16n8k16(ds_pack[mi][kd], K_rmem_trans[kd][d], dQ_acc[mi][d]); + +#pragma unroll + for (int mi = 0; mi < WARP_Q / MMA_M; mi++) { + const int row_a = start_m + warp_id * WARP_Q + mi * MMA_M + lane_id / 4; + const int row_b = row_a + 8; +#pragma unroll + for (int d = 0; d < DIM / MMA_N; d++) { + const int col = d * MMA_N + (lane_id % 4) * 2; + float *regs = dQ_acc[mi][d]; + if (row_a < seqlen_q) { + atomicAdd(&DQ[((q_start + row_a) * H + h) * DIM + col], __float2bfloat16_rn(regs[0])); + atomicAdd(&DQ[((q_start + row_a) * H + h) * DIM + col + 1], + __float2bfloat16_rn(regs[1])); + } + if (row_b < seqlen_q) { + atomicAdd(&DQ[((q_start + row_b) * H + h) * DIM + col], __float2bfloat16_rn(regs[2])); + atomicAdd(&DQ[((q_start + row_b) * H + h) * DIM + col + 1], + __float2bfloat16_rn(regs[3])); + } + } + } + + __syncthreads(); + } + +#pragma unroll + for (int mi = 0; mi < WARP_Q / MMA_M; mi++) { + const int row = kv_tile_row_base + warp_id * WARP_Q + mi * MMA_M + lane_id / 4; + const int row8 = row + 8; +#pragma unroll + for (int d = 0; d < DIM / MMA_N; d++) { + const int col = d * MMA_N + (lane_id % 4) * 2; + float *dk_regs = dk_out[mi][d]; + float *dv_regs = dv_out[mi][d]; + if (row < seqlen_k) { + DK[((k_start + row) * H + h) * DIM + col] = __float2bfloat16_rn(dk_regs[0]); + DK[((k_start + row) * H + h) * DIM + col + 1] = __float2bfloat16_rn(dk_regs[1]); + DV[((k_start + row) * H + h) * DIM + col] = __float2bfloat16_rn(dv_regs[0]); + DV[((k_start + row) * H + h) * DIM + col + 1] = __float2bfloat16_rn(dv_regs[1]); + } + if (row8 < seqlen_k) { + DK[((k_start + row8) * H + h) * DIM + col] = __float2bfloat16_rn(dk_regs[2]); + DK[((k_start + row8) * H + h) * DIM + col + 1] = __float2bfloat16_rn(dk_regs[3]); + DV[((k_start + row8) * H + h) * DIM + col] = __float2bfloat16_rn(dv_regs[2]); + DV[((k_start + row8) * H + h) * DIM + col + 1] = __float2bfloat16_rn(dv_regs[3]); + } + } + } +} + +} // namespace + +// Forward-only packed varlen FlashAttention (SM90, TMA + mma.sync). q/k/v are +// bf16 [total, H, DIM=128], contiguous, no GQA. cu_seqlens_{q,k} are int32 +// [batch+1] cumulative offsets (cu_seqlens[0]==0), same convention as +// rl_engine/kernels/ops/triton/triton_attn.py's varlen path. Returns +// (out [total_q, H, DIM] bf16, lse [total_q, H] f32) -- lse is the +// attention-domain log-sum-exp (M + log(L)), not the vocab-domain LSE used by +// the fused_logp/linear_logp kernels. +std::vector flash_attention_varlen_sm90_forward( + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor cu_seqlens_q, + torch::Tensor cu_seqlens_k, + int64_t max_seqlen_q, + int64_t max_seqlen_k, + bool causal, + double sm_scale) { + TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda(), "q/k/v must be CUDA tensors"); + TORCH_CHECK(q.device() == k.device() && q.device() == v.device(), + "q/k/v must be on the same device"); + TORCH_CHECK(q.scalar_type() == at::kBFloat16 && k.scalar_type() == at::kBFloat16 && + v.scalar_type() == at::kBFloat16, + "flash_attention_varlen_sm90 requires bfloat16 q/k/v"); + TORCH_CHECK(q.dim() == 3 && k.dim() == 3 && v.dim() == 3, "q/k/v must be [total, H, D]"); + TORCH_CHECK(q.is_contiguous() && k.is_contiguous() && v.is_contiguous(), + "flash_attention_varlen_sm90 requires contiguous q/k/v (no transposed/strided " + "views); the Triton varlen path tolerates non-contiguous inputs, this kernel " + "does not"); + + const int total_q = q.size(0); + const int H = q.size(1); + const int D = q.size(2); + TORCH_CHECK(D == DIM, "flash_attention_varlen_sm90 only supports head_dim=", DIM, + " in this milestone, got ", D); + TORCH_CHECK(k.size(1) == H && v.size(1) == H, + "GQA is not supported: k/v head count must equal q's"); + TORCH_CHECK(k.size(2) == D && v.size(2) == D, "k/v head_dim must match q's"); + const int total_k = k.size(0); + TORCH_CHECK(v.size(0) == total_k, "k and v must have the same total token count"); + + TORCH_CHECK(cu_seqlens_q.is_cuda() && cu_seqlens_k.is_cuda(), "cu_seqlens must be on CUDA"); + TORCH_CHECK(cu_seqlens_q.scalar_type() == at::kInt && cu_seqlens_k.scalar_type() == at::kInt, + "cu_seqlens_q/cu_seqlens_k must be int32"); + TORCH_CHECK(cu_seqlens_q.dim() == 1 && cu_seqlens_k.dim() == 1, "cu_seqlens must be 1-D"); + TORCH_CHECK(cu_seqlens_q.is_contiguous() && cu_seqlens_k.is_contiguous(), + "cu_seqlens must be contiguous"); + const int batch = cu_seqlens_q.numel() - 1; + TORCH_CHECK(batch >= 1, "cu_seqlens_q must have at least 2 elements"); + TORCH_CHECK(cu_seqlens_k.numel() - 1 == batch, "cu_seqlens_q/cu_seqlens_k batch mismatch"); + + at::cuda::CUDAGuard device_guard(q.device()); + + auto out = torch::empty_like(q); + auto lse = torch::empty({total_q, H}, q.options().dtype(torch::kFloat)); + + CUtensorMap q_tmap, k_tmap, v_tmap; + init_tensor_map_noswizzle(&q_tmap, reinterpret_cast(q.data_ptr()), + total_q, H * D, BLOCK_Q, D); + init_tensor_map_noswizzle(&k_tmap, reinterpret_cast(k.data_ptr()), + total_k, H * D, BLOCK_KV, D); + init_tensor_map_noswizzle(&v_tmap, reinterpret_cast(v.data_ptr()), + total_k, H * D, BLOCK_KV, D); + + const int smem = static_cast((BLOCK_Q * DIM + 2 * STAGES * BLOCK_KV * DIM) * + sizeof(nv_bfloat16) + + (1 + STAGES) * 8); + if (smem > 48 * 1024) { + cudaFuncSetAttribute(flash_attention_varlen_sm90_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem); + } + + dim3 grid(cdiv(static_cast(max_seqlen_q), BLOCK_Q), batch * H, 1); + flash_attention_varlen_sm90_kernel<<>>( + q_tmap, k_tmap, v_tmap, cu_seqlens_q.data_ptr(), cu_seqlens_k.data_ptr(), + reinterpret_cast(out.data_ptr()), lse.data_ptr(), H, + static_cast(sm_scale), causal); + + return {out, lse}; +} + +// Backward for flash_attention_varlen_sm90_forward: dQ/dK/dV via the +// recompute-based algorithm in triton_attn.py's `_bwd_kernel_varlen` +// (grid over KV-tiles, Q-tiles looped internally), using the forward's +// saved `lse` directly (p = exp(qk*scale - lse), algebraically identical to +// Triton's separate-M/L form -- see the comment on this in the plan/PR). +// `dq` is atomic-accumulated in bf16 (native SM90 `atom.add.noftz.bf16`), +// `dk`/`dv` are written once per KV-tile CTA (no atomics needed there). +std::vector flash_attention_varlen_sm90_backward( + torch::Tensor do_, + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor out, + torch::Tensor lse, + torch::Tensor cu_seqlens_q, + torch::Tensor cu_seqlens_k, + int64_t max_seqlen_k, + bool causal, + double sm_scale) { + TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda() && do_.is_cuda() && out.is_cuda() && + lse.is_cuda(), + "flash_attention_varlen_sm90_backward requires CUDA tensors"); + TORCH_CHECK(q.scalar_type() == at::kBFloat16 && k.scalar_type() == at::kBFloat16 && + v.scalar_type() == at::kBFloat16 && do_.scalar_type() == at::kBFloat16 && + out.scalar_type() == at::kBFloat16, + "flash_attention_varlen_sm90_backward requires bfloat16 q/k/v/do/out"); + TORCH_CHECK(lse.scalar_type() == at::kFloat, "lse must be float32"); + TORCH_CHECK(q.dim() == 3 && k.dim() == 3 && v.dim() == 3 && do_.dim() == 3 && out.dim() == 3, + "q/k/v/do/out must be [total, H, D]"); + TORCH_CHECK(q.is_contiguous() && k.is_contiguous() && v.is_contiguous() && do_.is_contiguous() && + out.is_contiguous() && lse.is_contiguous(), + "flash_attention_varlen_sm90_backward requires contiguous q/k/v/do/out/lse"); + + const int total_q = q.size(0); + const int H = q.size(1); + const int D = q.size(2); + TORCH_CHECK(D == DIM, "flash_attention_varlen_sm90_backward only supports head_dim=", DIM, + " in this milestone, got ", D); + TORCH_CHECK(k.size(1) == H && v.size(1) == H, + "GQA is not supported: k/v head count must equal q's"); + TORCH_CHECK(k.size(2) == D && v.size(2) == D, "k/v head_dim must match q's"); + TORCH_CHECK(do_.sizes() == q.sizes() && out.sizes() == q.sizes(), + "do/out must have the same shape as q"); + TORCH_CHECK(lse.size(0) == total_q && lse.size(1) == H, "lse must be [total_q, H]"); + const int total_k = k.size(0); + TORCH_CHECK(v.size(0) == total_k, "k and v must have the same total token count"); + + TORCH_CHECK(cu_seqlens_q.is_cuda() && cu_seqlens_k.is_cuda(), "cu_seqlens must be on CUDA"); + TORCH_CHECK(cu_seqlens_q.scalar_type() == at::kInt && cu_seqlens_k.scalar_type() == at::kInt, + "cu_seqlens_q/cu_seqlens_k must be int32"); + TORCH_CHECK(cu_seqlens_q.dim() == 1 && cu_seqlens_k.dim() == 1, "cu_seqlens must be 1-D"); + TORCH_CHECK(cu_seqlens_q.is_contiguous() && cu_seqlens_k.is_contiguous(), + "cu_seqlens must be contiguous"); + const int batch = cu_seqlens_q.numel() - 1; + TORCH_CHECK(batch >= 1, "cu_seqlens_q must have at least 2 elements"); + TORCH_CHECK(cu_seqlens_k.numel() - 1 == batch, "cu_seqlens_q/cu_seqlens_k batch mismatch"); + + at::cuda::CUDAGuard device_guard(q.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + + auto dq = torch::zeros_like(q); // atomic-accumulated, must start at zero + auto dk = torch::empty_like(k); + auto dv = torch::empty_like(v); + auto delta = torch::empty({total_q, H}, q.options().dtype(torch::kFloat)); + + flash_attention_bwd_preprocess_kernel<<>>( + reinterpret_cast(out.data_ptr()), + reinterpret_cast(do_.data_ptr()), delta.data_ptr(), + total_q, H); + + CUtensorMap q_tmap, k_tmap, v_tmap, do_tmap; + init_tensor_map_noswizzle(&q_tmap, reinterpret_cast(q.data_ptr()), + total_q, H * D, BLOCK_Q, D); + init_tensor_map_noswizzle(&k_tmap, reinterpret_cast(k.data_ptr()), + total_k, H * D, BLOCK_KV, D); + init_tensor_map_noswizzle(&v_tmap, reinterpret_cast(v.data_ptr()), + total_k, H * D, BLOCK_KV, D); + init_tensor_map_noswizzle(&do_tmap, + reinterpret_cast(do_.data_ptr()), total_q, + H * D, BLOCK_Q, D); + + const int smem = static_cast((2 * BLOCK_KV * DIM + 2 * STAGES * BLOCK_Q * DIM + + BLOCK_KV * BLOCK_Q) * + sizeof(nv_bfloat16) + + (1 + STAGES) * 8); + if (smem > 48 * 1024) { + cudaFuncSetAttribute(flash_attention_varlen_sm90_bwd_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem); + } + + dim3 grid(cdiv(static_cast(max_seqlen_k), BLOCK_KV), batch * H, 1); + flash_attention_varlen_sm90_bwd_kernel<<>>( + q_tmap, k_tmap, v_tmap, do_tmap, cu_seqlens_q.data_ptr(), cu_seqlens_k.data_ptr(), + lse.data_ptr(), delta.data_ptr(), + reinterpret_cast(dq.data_ptr()), + reinterpret_cast(dk.data_ptr()), + reinterpret_cast(dv.data_ptr()), H, static_cast(sm_scale), + causal); + + return {dq, dk, dv}; +} diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 61ba4a3b..3bcd4502 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -58,6 +58,26 @@ torch::Tensor linear_logp_logits_bf16_to_dlogits(torch::Tensor logits, torch::Tensor grad_logp, torch::Tensor lse, int64_t vocab_start_index); +std::vector flash_attention_varlen_sm90_forward(torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor cu_seqlens_q, + torch::Tensor cu_seqlens_k, + int64_t max_seqlen_q, + int64_t max_seqlen_k, + bool causal, + double sm_scale); +std::vector flash_attention_varlen_sm90_backward(torch::Tensor do_, + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor out, + torch::Tensor lse, + torch::Tensor cu_seqlens_q, + torch::Tensor cu_seqlens_k, + int64_t max_seqlen_k, + bool causal, + double sm_scale); #endif #if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA) @@ -150,6 +170,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "In-place local bf16 probs -> TP dlogits for selected log-prob backward"); m.def("linear_logp_logits_bf16_to_dlogits", &linear_logp_logits_bf16_to_dlogits, "Build bf16 dlogits from bf16 logits and fp32 lse"); + m.def("flash_attention_varlen_sm90", &flash_attention_varlen_sm90_forward, + "TMA+mma.sync causal packed varlen FlashAttention forward (SM90), returns (out, lse)"); + m.def("flash_attention_varlen_sm90_backward", &flash_attention_varlen_sm90_backward, + "TMA+mma.sync causal packed varlen FlashAttention backward (SM90), returns (dq, dk, dv)"); #endif #if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA) diff --git a/docs/.nav.yml b/docs/.nav.yml index 8794c731..3b77709b 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -16,6 +16,7 @@ nav: - operators/README.md - operators/activation.md - operators/attention.md + - operators/attention-varlen.md - operators/fused-logp.md - operators/linear-logp.md - operators/linear-logp-tp-test.md diff --git a/docs/operators/README.md b/docs/operators/README.md index bf38f368..e06cccf1 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -20,6 +20,7 @@ Every operator page should include: - [SiLU / SwiGLU Activation](activation.md) - [Standard Attention](attention.md) +- [FlashAttention: LSE Export + Variable-Length Packing (Triton)](attention-varlen.md) - [Fused LogP](fused-logp.md) - [Fused Linear LogP](linear-logp.md) - [Fused Linear LogP TP Test Runbook](linear-logp-tp-test.md) diff --git a/docs/operators/attention-varlen.md b/docs/operators/attention-varlen.md new file mode 100644 index 00000000..13d802e2 --- /dev/null +++ b/docs/operators/attention-varlen.md @@ -0,0 +1,196 @@ +# FlashAttention: LSE Export + Variable-Length Packing + +## Summary + +Extends the existing pure-Triton FlashAttention fallback +(`rl_engine/kernels/ops/triton/triton_attn.py`) with two capabilities needed for +long-context RL rollout/training workloads: + +- **Attention-domain LSE export** — the per-query-row log-sum-exp of the scaled, + masked `QKᵀ` logits, in the same fixed online-softmax reduction order the + kernel already used internally for the backward pass. Useful for backward + recomputation, diagnostics, and rollout/training attention alignment checks. + **This is not the vocab-domain LSE produced by the `fused_logp` / + `linear_logp` kernels** — same name, different tensor domain (per key-length + reduction vs. per-vocab reduction). Do not conflate the two. +- **Variable-length (packed) attention** — operates directly on `cu_seqlens`-packed + `[total_tokens, H, D]` tensors instead of a padded `[B, H, S, D]` batch, so RL + batches with wildly different response lengths (or empty/fully-masked + responses) don't pay for padding in either compute or memory. + +This module is the **cross-platform semantic baseline**: the SM90 (forward, +now implemented — see Backends), SM80 `mma.sync`, and ROCm MFMA +fused-attention kernels are checked against it (and, transitively, against +`NativeAttentionOp`) rather than against each other. See +`docs/operators/attention.md` for the pure-PyTorch WS1 ground-truth reference +this whole family is validated against. + +## Entry Point + +```python +from rl_engine.kernels.ops.triton.triton_attn import ( + triton_flash_attention, + triton_flash_attention_varlen, +) + +# Dense — unchanged default behavior, opt into LSE: +out = triton_flash_attention(q, k, v, causal=True) # [B, H, S, D] +out, lse = triton_flash_attention(q, k, v, causal=True, return_lse=True) # lse: [B, H, S] fp32 + +# Packed variable-length: +out = triton_flash_attention_varlen( + q, k, v, # [total_q, H, D], [total_k, H, D], [total_k, H, D] + cu_seqlens_q, cu_seqlens_k, # int32 [batch + 1], cu_seqlens[0] == 0 + max_seqlen_q, max_seqlen_k, # host ints, size the launch grid + causal=True, +) +out, lse = triton_flash_attention_varlen(..., return_lse=True) # lse: [total_q, H] fp32 +``` + +## Backends + +| Backend | Wrapper | Native symbol | Status | +| --- | --- | --- | --- | +| Triton (CUDA) | `triton_flash_attention` / `triton_flash_attention_varlen` | `_fwd_kernel[_varlen]`, `_bwd_kernel[_varlen]` | Cross-platform semantic baseline. | +| CUDA SM90 (TMA + `mma.sync`) | `flash_attention_sm90_varlen` (`rl_engine.kernels.ops.cuda.attention.flash_attn_sm90`) | `flash_attention_varlen_sm90[_backward]` | Forward + backward (causal, varlen, LSE export, `dQ`/`dK`/`dV`), fully differentiable via `torch.autograd.Function`. Validated against this Triton path. | +| CUDA SM80 (`mma.sync`) | — | — | Planned; validates against this Triton path. | +| ROCm (MFMA) | — | — | Planned; validates against this Triton path. | + +Not yet wired into `KernelRegistry` — this is the standalone kernel-development +stage (mirrors how `prefix_shared_attention` and the SM90 logp kernels were +built before registry integration). Registry dispatch is future work once a +production op_type contract for causal/varlen attention is settled. + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `q` (dense) | `[B, H, Sq, D]` | fp16/bf16/fp32 | `D ∈ {16,32,64,128,256}`. | +| `k`, `v` (dense) | `[B, H, Skv, D]` | same as `q` | **No GQA**: `k`/`v` head count must equal `q`'s (matches the pre-existing dense kernel's limitation, not new). | +| `q` (varlen) | `[total_q, H, D]` | fp16/bf16/fp32 | Packed, no padding. | +| `k`, `v` (varlen) | `[total_k, H, D]` | same as `q` | Same no-GQA constraint. | +| `cu_seqlens_q`, `cu_seqlens_k` | `[batch + 1]` | int32 | Cumulative offsets, `cu_seqlens[0] == 0`. Same convention as `flash_attn_varlen_func` and this repo's `pack` op (#182). | +| `max_seqlen_q`, `max_seqlen_k` | — | Python `int` | Host-side; sizes the launch grid. | +| `causal` | — | bool | Per-sequence anchor `Skv - Sq` (see Accuracy) — identical formula for prefill (`Sq == Skv`) and decode (`Sq < Skv`). | +| `return_lse` | — | bool | If `True`, also returns the attention-domain LSE, float32, **non-differentiable** (`ctx.mark_non_differentiable`) — diagnostics/backward-recompute only, matching `flash_attn`'s `softmax_lse` external contract. | +| output | dense: `[B, H, Sq, D]`; varlen: `[total_q, H, D]` | input dtype | — | + +The table above states the **Triton** path's contract. The **CUDA SM90** +backend (`flash_attention_sm90_varlen`) is narrower: **varlen only** (no dense +entry point), **bf16 only** (no fp16/fp32 — the reused `mma.sync` PTX +hardcodes bf16 operands), **`D == 128` only** (no other head dim this +milestone), and **contiguous `q`/`k`/`v` only** (raises `RuntimeError` on a +non-contiguous/transposed view — TMA descriptors assume a fixed row stride; +this is a deliberate, tested divergence from the Triton path, which tolerates +non-contiguous inputs). Inputs outside this support surface fall back to +`triton_flash_attention_varlen` automatically rather than erroring. + +## Dispatch Behavior + +Direct function calls only (see Backends). No registry op_type yet. + +## Accuracy + +Both paths are checked against an **independent** fp32 masked-softmax + LSE +closed form (not against each other, so a shared online-softmax bug can't +hide): + +```python +scores = einsum("hqd,hkd->hqk", q.float(), k.float()) * scale +if causal: + mask = torch.triu(torch.ones(Sq, Skv, bool), diagonal=Skv - Sq + 1) + scores = scores.masked_fill(mask, -inf) +probs = softmax(scores, dim=-1) +out = einsum("hqk,hkd->hqd", probs, v.float()) +lse = torch.logsumexp(scores, dim=-1) +``` + +Measured on an H100 SXM5 (fp16 inputs, fp32 accumulation in-kernel): + +- Dense: `out` max-abs-diff ≈ `1.4e-3`, `lse` max-abs-diff ≈ `1e-6`. +- Varlen: `out`/`dq`/`dk`/`dv` max-abs-diff in the `5e-4`–`1.1e-2` range across + uneven, non-block-aligned sequence lengths; `lse` ≈ `1e-6`. + +CUDA SM90 (bf16 inputs, bf16×bf16→fp32 `mma.sync` accumulation, checked +against the same fp32 reference): `out`/`lse` max-abs-diff up to ≈ `1.6e-2`; +`dq`/`dk`/`dv` max-abs-diff up to ≈ `1.6e-2` across the full shape matrix, +including the 24-KV-tile long-sequence and `dQ` atomic-accumulation stress +cases — bf16 (fewer mantissa bits than fp16) accounts for the looser bound +than the Triton fp16 numbers above. + +Varlen boundary handling: since packed sequences sit back-to-back in memory (no +padding), reading past a sequence's own `cu_seqlens` bound is a genuine +correctness bug (you'd read the *next* sequence's tokens), not just wasted +compute. All loads/stores in the varlen kernels are therefore explicitly masked +against `seqlen_q`/`seqlen_k` (not `boundary_check` on dynamically-shaped block +pointers) — this is deliberate and was chosen specifically to make the +padding-vs-next-sequence distinction unambiguous. + +## Performance Notes + +No standalone benchmark script yet (unlike `fused-logp.md` / `linear-logp.md`). +Follow-up: add `benchmarks/benchmark_attention_varlen.py` measuring +packed-vs-padded latency/VRAM at representative RL group sizes, alongside the +existing `benchmarks/benchmark_attention.py`. + +## Tests + +```bash +python -m pytest tests/test_triton_attention_varlen.py -v +python -m pytest tests/test_flash_attention_sm90.py -v +``` + +`test_triton_attention_varlen.py` covers: dense LSE vs. independent reference, +LSE non-differentiability with backward still populating `q`/`k`/`v` grads, +default-call backward compatibility (bare tensor when `return_lse=False`), +varlen forward+backward across uneven non-block-aligned seqlens, non-causal, +decode-style (`Sq < Skv` varying per sequence), `head_dim=128`, and a +zero-length sequence in the batch (a real occurrence for fully-masked/empty RL +responses). + +`test_flash_attention_sm90.py` (requires a Hopper GPU with the extension built +`KERNEL_ALIGN_FORCE_SM90=1`) covers the same varlen shape matrix for the SM90 +kernel — uneven non-block-aligned seqlens, block-aligned seqlens, non-causal, +decode-style, `head_dim=128`, zero-length sequence, long/many-KV-tile +sequences, single-head, and 8-entry mixed batches — with every case also +exercising `.backward()` and checking `dQ`/`dK`/`dV`. Forward output/LSE is +checked against both the independent fp32 reference and +`triton_flash_attention_varlen` on the same bf16 tensors; gradients are +checked against the fp32 reference only (`triton_flash_attention_varlen`'s +own backward can't run on bf16 tensors at all in this Triton version — +`tl.atomic_add` doesn't support bf16, a compiler-level restriction, not a +numerical one; this repo's own Triton backward tests avoid it too, always +using fp16). Additional dedicated cases: a single-KV-tile-CTA case (isolates +the `P`/`ds` shared-memory transpose-trick correctness from the causal +`lo`-bound loop and cross-CTA `dQ` atomics), a `dQ` atomic-accumulation stress +case (short query sequence attended to by ~16 KV-tile CTAs, non-causal, so +every CTA atomic-adds into the same few rows), an LSE-non-differentiability +check, a dense-vs-varlen consistency check, and validation-guard tests for +both the forward and backward native symbols (bad head_dim/dtype/GQA/ +cu_seqlens-mismatch/non-contiguous input). + +## Known Limitations + +- No GQA support on any path (`Hk` must equal `Hq`) — same constraint the + pre-existing dense kernel already had; not introduced by this change. +- Not wired into `KernelRegistry`; no automatic backend selection yet. +- CUDA SM90 forward + backward is implemented (TMA + `mma.sync`, causal, + varlen, LSE export, `dQ`/`dK`/`dV`) and fully differentiable. It is + **bf16-only**, **`D=128`-only**, **varlen-only** (no dense entry point), and + **rejects non-contiguous inputs** (see Tensor Contract) — inputs outside + that surface fall back to `triton_flash_attention_varlen` entirely + (forward *and* backward together; a build with only the forward native + symbol compiled also falls back entirely, rather than mixing an SM90 + forward with a Triton backward). Neither the forward nor backward kernel + uses split-KV/split-Q parallelism or a persistent-kernel schedule (a + sequential loop per CTA, same as `prefix_shared_attention.cu`) — a possible + future perf pass, not attempted here. +- `seqlen_q > seqlen_k` under `causal=True` is out-of-contract (violates the + `Skv - Sq` causal-offset convention, which assumes a query can always + attend at least to itself) and produces `NaN`/degenerate output — this is + not specific to the SM90 kernel; `triton_flash_attention_varlen` produces + `NaN` for the identical input too, confirmed directly. Neither + implementation is designed to handle it; callers must ensure + `seqlen_k >= seqlen_q` per sequence whenever `causal=True`. +- SM80 `mma.sync` and ROCm MFMA kernels are not implemented yet. +- No standalone performance benchmark yet (see Performance Notes). diff --git a/rl_engine/kernels/ops/cuda/attention/flash_attn_sm90.py b/rl_engine/kernels/ops/cuda/attention/flash_attn_sm90.py new file mode 100644 index 00000000..f55fc792 --- /dev/null +++ b/rl_engine/kernels/ops/cuda/attention/flash_attn_sm90.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import torch + +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.utils.logger import logger + +_SM90_HEAD_DIM = 128 + + +def _sm90_hardware_available() -> bool: + return ( + _EXT_AVAILABLE + and hasattr(_C, "flash_attention_varlen_sm90") + and hasattr(_C, "flash_attention_varlen_sm90_backward") + ) + + +def _sm90_supported(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> bool: + return ( + _sm90_hardware_available() + and q.is_cuda + and q.dtype == torch.bfloat16 + and k.dtype == torch.bfloat16 + and v.dtype == torch.bfloat16 + and q.dim() == 3 + and k.dim() == 3 + and v.dim() == 3 + and q.shape[-1] == _SM90_HEAD_DIM + and k.shape[-1] == _SM90_HEAD_DIM + and v.shape[-1] == _SM90_HEAD_DIM + and k.shape[1] == q.shape[1] + and v.shape[1] == q.shape[1] + and q.is_contiguous() + and k.is_contiguous() + and v.is_contiguous() + ) + + +class _FlashAttentionVarlenSM90Function(torch.autograd.Function): + """SM90 TMA+mma.sync forward and backward for packed varlen FlashAttention. + + forward() calls the compiled `_C.flash_attention_varlen_sm90` kernel and + saves what backward's recompute-based algorithm needs. backward() calls + `_C.flash_attention_varlen_sm90_backward`, which reuses the saved `lse` + directly -- `p = exp(qk*scale - lse)` is algebraically identical to the + separate-M/L form triton_attn.py's `_bwd_kernel_varlen` uses, so no + additional forward-side state is needed. + """ + + @staticmethod + def forward( + ctx, + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + causal, + sm_scale, + return_lse, + ): + cu_seqlens_q = cu_seqlens_q.to(device=q.device, dtype=torch.int32).contiguous() + cu_seqlens_k = cu_seqlens_k.to(device=q.device, dtype=torch.int32).contiguous() + + out, lse = _C.flash_attention_varlen_sm90( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + int(max_seqlen_q), + int(max_seqlen_k), + bool(causal), + float(sm_scale), + ) + + ctx.save_for_backward(q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k) + ctx.max_seqlen_k = int(max_seqlen_k) + ctx.causal = bool(causal) + ctx.sm_scale = float(sm_scale) + + if not return_lse: + return out, None + ctx.mark_non_differentiable(lse) + return out, lse + + @staticmethod + def backward(ctx, do, _dlse): + q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k = ctx.saved_tensors + do = do.contiguous() + + dq, dk, dv = _C.flash_attention_varlen_sm90_backward( + do, + q, + k, + v, + out, + lse, + cu_seqlens_q, + cu_seqlens_k, + ctx.max_seqlen_k, + ctx.causal, + ctx.sm_scale, + ) + # Inputs: q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, + # max_seqlen_k, causal, sm_scale, return_lse. + return dq, dk, dv, None, None, None, None, None, None, None + + +class FlashAttentionVarlenSM90Op: + """SM90 (Hopper) packed variable-length FlashAttention, forward + backward. + + TMA + `mma.sync` causal attention with attention-domain LSE export + (M + log(L), not the vocab-domain LSE used by the fused_logp / + linear_logp kernels), matching the semantics of + `triton_flash_attention_varlen` -- the cross-platform baseline this + kernel is validated against. + + Requires the extension built with `KERNEL_ALIGN_FORCE_SM90=1` on a Hopper + (SM90) device; bfloat16 q/k/v, head_dim=128, contiguous, no GQA only, and + *both* the forward and backward native symbols present -- a build with + only the forward kernel compiled falls back to Triton entirely rather + than mixing an SM90 forward with a Triton backward. Falls back to + `triton_flash_attention_varlen` (itself a proper autograd Function, so + gradients still work through the fallback) for anything outside that + support surface. + """ + + def __init__(self) -> None: + self.has_hardware_op = _sm90_hardware_available() + if self.has_hardware_op: + logger.info( + "Successfully linked to RL-Kernel _C.flash_attention_varlen_sm90" + " (+ backward)." + ) + else: + logger.warning( + "RL-Kernel _C.flash_attention_varlen_sm90[_backward] is unavailable. " + "FlashAttentionVarlenSM90Op will fall back to triton_flash_attention_varlen " + "(rebuild with KERNEL_ALIGN_FORCE_SM90=1 on a Hopper GPU for the fused kernel)." + ) + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + causal: bool = True, + sm_scale: float | None = None, + return_lse: bool = False, + ): + if sm_scale is None: + sm_scale = 1.0 / (q.shape[-1] ** 0.5) + + if self.has_hardware_op and _sm90_supported(q, k, v): + out, lse = _FlashAttentionVarlenSM90Function.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + causal, + sm_scale, + return_lse, + ) + return (out, lse) if return_lse else out + + from rl_engine.kernels.ops.triton.triton_attn import triton_flash_attention_varlen + + return triton_flash_attention_varlen( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + causal=causal, + sm_scale=sm_scale, + return_lse=return_lse, + ) + + +_OP: FlashAttentionVarlenSM90Op | None = None + + +def _get_op() -> FlashAttentionVarlenSM90Op: + global _OP + if _OP is None: + _OP = FlashAttentionVarlenSM90Op() + return _OP + + +def flash_attention_sm90_varlen( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + causal: bool = True, + sm_scale: float | None = None, + return_lse: bool = False, +): + """SM90 packed variable-length FlashAttention (forward + backward). + + Signature mirrors `triton_flash_attention_varlen` + (`rl_engine.kernels.ops.triton.triton_attn`) exactly, so call sites and + tests can swap between the two backends with no changes. Differentiable: + `q`/`k`/`v` gradients flow correctly whether the SM90 hardware kernel or + the Triton fallback services the call. + + Args: + q: [total_q, H, D] packed queries. D must be 128 for the fused SM90 + kernel to engage; other shapes/dtypes fall back to Triton. + k, v: [total_k, H, D] packed keys/values. GQA (Hk != Hq) unsupported. + cu_seqlens_q, cu_seqlens_k: int32 [batch + 1] cumulative offsets, + cu_seqlens[0] == 0. + max_seqlen_q, max_seqlen_k: max per-sequence length in the batch. + causal: causal masking, anchored per-sequence via `Skv - Sq`. + sm_scale: defaults to `1/sqrt(D)`. + return_lse: if True, also return the packed `[total_q, H]` float32 + attention-domain LSE (non-differentiable). + Returns: + `out` of shape [total_q, H, D] if `return_lse` is False, else + `(out, lse)`. + """ + return _get_op()( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + causal=causal, + sm_scale=sm_scale, + return_lse=return_lse, + ) diff --git a/rl_engine/kernels/ops/triton/triton_attn.py b/rl_engine/kernels/ops/triton/triton_attn.py index c1bd4809..69430afa 100644 --- a/rl_engine/kernels/ops/triton/triton_attn.py +++ b/rl_engine/kernels/ops/triton/triton_attn.py @@ -333,7 +333,7 @@ def _bwd_kernel( class _TritonAttention(torch.autograd.Function): @staticmethod - def forward(ctx, q, k, v, causal, sm_scale): + def forward(ctx, q, k, v, causal, sm_scale, return_lse): # [batch, num_heads, seq_len, head_dim] # Triton tutorial standard layout requires specific contiguity Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1] @@ -390,10 +390,20 @@ def forward(ctx, q, k, v, causal, sm_scale): ) ctx.save_for_backward(q, k, v, out, L, M) - return out + + if not return_lse: + return out, None + + # Attention-domain LSE: log-sum-exp of the (scaled, masked) QK^T logits per + # query row, in the same fixed reduction order the fwd kernel already used to + # accumulate M (running max) and L (running sum-exp) online. This is distinct + # from the vocab-domain LSE produced by the logp/linear_logp kernels. + lse = M + torch.log(L) + ctx.mark_non_differentiable(lse) + return out, lse @staticmethod - def backward(ctx, do): + def backward(ctx, do, _dlse): q, k, v, out, L, M = ctx.saved_tensors do = do.contiguous() @@ -468,7 +478,7 @@ def backward(ctx, do): num_stages=1, ) - return dq, dk, dv, None, None + return dq, dk, dv, None, None, None def triton_flash_attention( @@ -477,6 +487,7 @@ def triton_flash_attention( v: torch.Tensor, causal: bool = True, sm_scale: float | None = None, + return_lse: bool = False, ): """ Universal backup Triton FlashAttention (support Forward / Backward) @@ -484,12 +495,521 @@ def triton_flash_attention( Args: q, k, v: Tensors of shape [batch, num_heads, seq_len, head_dim]. It is recommended to switch to contiguous memory (.contiguous()). - causal: Whether to turn on causal masking. - sm_scale: Softmax scaling factor, default to 1.0 / sqrt(head_dim). + causal: Whether to turn on causal masking. + sm_scale: Softmax scaling factor, default to 1.0 / sqrt(head_dim). + return_lse: If True, also return the per-query-row attention-domain LSE + (log-sum-exp of the scaled, masked QK^T logits), shape + [batch, num_heads, seq_len], float32. Not a vocab-domain logprob LSE. + The returned LSE is non-differentiable (diagnostics / backward-recompute + use only), matching the external contract of `flash_attn`'s `softmax_lse`. + Returns: + `out` of shape [batch, num_heads, seq_len, head_dim] if `return_lse` is False, + else `(out, lse)`. + """ + if sm_scale is None: + sm_scale = 1.0 / (q.shape[-1] ** 0.5) + + out, lse = _TritonAttention.apply(q, k, v, causal, sm_scale, return_lse) + return (out, lse) if return_lse else out + + +# --------------------------------------------------------------------------- +# Variable-length (packed) attention. +# +# Q/K/V are packed along the token dimension: [total_tokens, H, D], with no +# per-sequence padding. `cu_seqlens_{q,k}` are int32 [batch + 1] cumulative +# sequence-length offsets (cu_seqlens[0] == 0), the same convention as +# `flash_attn_varlen_func` and this repo's `pack` op (#182). Each program +# handles one (query-block, batch, head) triple; `causal_offset = seqlen_k - +# seqlen_q` reproduces the `Skv - Sq` causal anchor from the WS1 +# `NativeAttentionOp` reference (docs/operators/attention.md) so prefill +# (Sq == Skv) and decode (Sq < Skv) share one formula. Boundary handling is +# via explicit masks (not `boundary_check`) because a block that runs past a +# sequence's length would otherwise read into the *next* packed sequence. +# --------------------------------------------------------------------------- + + +@triton.jit +def _fwd_kernel_varlen( + Q, + K, + V, + sm_scale, + cu_seqlens_q, + cu_seqlens_k, + L, + M, + Out, + stride_qm, + stride_qh, + stride_qk, + stride_kn, + stride_kh, + stride_kk, + stride_vn, + stride_vh, + stride_vk, + stride_om, + stride_oh, + stride_ok, + H, + BLOCK_M: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_N: tl.constexpr, + IS_CAUSAL: tl.constexpr, +): + start_m = tl.program_id(0) + pid_bh = tl.program_id(1) + b = pid_bh // H + h = pid_bh % H + + q_start = tl.load(cu_seqlens_q + b) + q_end = tl.load(cu_seqlens_q + b + 1) + seqlen_q = q_end - q_start + if start_m * BLOCK_M >= seqlen_q: + return + + k_start = tl.load(cu_seqlens_k + b) + k_end = tl.load(cu_seqlens_k + b + 1) + seqlen_k = k_end - k_start + causal_offset = seqlen_k - seqlen_q + + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_d = tl.arange(0, BLOCK_DMODEL) + valid_m = offs_m < seqlen_q + + q_ptrs = ( + Q + (q_start + offs_m[:, None]) * stride_qm + h * stride_qh + offs_d[None, :] * stride_qk + ) + q = tl.load(q_ptrs, mask=valid_m[:, None], other=0.0) + + m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32) + + if IS_CAUSAL: + hi = tl.minimum(seqlen_k, (start_m + 1) * BLOCK_M + causal_offset) + else: + hi = seqlen_k + + for start_n in range(0, hi, BLOCK_N): + offs_n = start_n + tl.arange(0, BLOCK_N) + valid_n = offs_n < seqlen_k + + k_ptrs = ( + K + + (k_start + offs_n[:, None]) * stride_kn + + h * stride_kh + + offs_d[None, :] * stride_kk + ) + k = tl.load(k_ptrs, mask=valid_n[:, None], other=0.0) + + qk = tl.dot(q, tl.trans(k)) * sm_scale + + mask = valid_n[None, :] + if IS_CAUSAL: + mask = mask & (offs_m[:, None] >= (offs_n[None, :] - causal_offset)) + qk = tl.where(mask, qk, float("-inf")) + + m_ij = tl.max(qk, 1) + m_i_new = tl.maximum(m_i, m_ij) + alpha = tl.exp(m_i - m_i_new) + beta = tl.exp(qk - m_i_new[:, None]) + l_i_new = alpha * l_i + tl.sum(beta, 1) + + p_scale = beta / l_i_new[:, None] + acc_scale = l_i / l_i_new * alpha + acc = acc * acc_scale[:, None] + + v_ptrs = ( + V + + (k_start + offs_n[:, None]) * stride_vn + + h * stride_vh + + offs_d[None, :] * stride_vk + ) + v = tl.load(v_ptrs, mask=valid_n[:, None], other=0.0) + p = p_scale.to(v.dtype) + acc += tl.dot(p, v) + + l_i = l_i_new + m_i = m_i_new + + acc = acc.to(Out.dtype.element_ty) + o_ptrs = ( + Out + (q_start + offs_m[:, None]) * stride_om + h * stride_oh + offs_d[None, :] * stride_ok + ) + tl.store(o_ptrs, acc, mask=valid_m[:, None]) + + l_ptrs = L + (q_start + offs_m) * H + h + m_ptrs = M + (q_start + offs_m) * H + h + tl.store(l_ptrs, l_i, mask=valid_m) + tl.store(m_ptrs, m_i, mask=valid_m) + + +@triton.jit +def _bwd_preprocess_varlen( + Out, + DO, + Delta, + stride_om, + stride_oh, + stride_ok, + stride_dom, + stride_doh, + stride_dok, + total_q, + H, + BLOCK_M: tl.constexpr, + D_HEAD: tl.constexpr, +): + pid_m = tl.program_id(0) + h = tl.program_id(1) + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_d = tl.arange(0, D_HEAD) + valid_m = offs_m < total_q + + o_ptrs = Out + offs_m[:, None] * stride_om + h * stride_oh + offs_d[None, :] * stride_ok + do_ptrs = DO + offs_m[:, None] * stride_dom + h * stride_doh + offs_d[None, :] * stride_dok + + o = tl.load(o_ptrs, mask=valid_m[:, None], other=0.0) + do = tl.load(do_ptrs, mask=valid_m[:, None], other=0.0).to(o.dtype) + + delta = tl.sum(o * do, axis=1) + tl.store(Delta + offs_m * H + h, delta, mask=valid_m) + + +@triton.jit +def _bwd_kernel_varlen( + Q, + K, + V, + sm_scale, + DO, + DQ, + DK, + DV, + L, + M, + Delta, + cu_seqlens_q, + cu_seqlens_k, + stride_qm, + stride_qh, + stride_qk, + stride_kn, + stride_kh, + stride_kk, + stride_vn, + stride_vh, + stride_vk, + stride_dom, + stride_doh, + stride_dok, + H, + BLOCK_M: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_N: tl.constexpr, + IS_CAUSAL: tl.constexpr, +): + start_n = tl.program_id(0) + pid_bh = tl.program_id(1) + b = pid_bh // H + h = pid_bh % H + + k_start = tl.load(cu_seqlens_k + b) + k_end = tl.load(cu_seqlens_k + b + 1) + seqlen_k = k_end - k_start + if start_n * BLOCK_N >= seqlen_k: + return + + q_start = tl.load(cu_seqlens_q + b) + q_end = tl.load(cu_seqlens_q + b + 1) + seqlen_q = q_end - q_start + causal_offset = seqlen_k - seqlen_q + + offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, BLOCK_DMODEL) + valid_n = offs_n < seqlen_k + + k_ptrs = ( + K + (k_start + offs_n[:, None]) * stride_kn + h * stride_kh + offs_d[None, :] * stride_kk + ) + v_ptrs = ( + V + (k_start + offs_n[:, None]) * stride_vn + h * stride_vh + offs_d[None, :] * stride_vk + ) + k = tl.load(k_ptrs, mask=valid_n[:, None], other=0.0) + v = tl.load(v_ptrs, mask=valid_n[:, None], other=0.0) + + dk = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32) + dv = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32) + + # Causal: query rows before `offs_n - causal_offset` never attend to this K/V + # block, so the M-loop can start there instead of at row 0. + lo = tl.maximum(0, start_n * BLOCK_N - causal_offset) if IS_CAUSAL else 0 + lo = (lo // BLOCK_M) * BLOCK_M + + for start_m in range(lo, seqlen_q, BLOCK_M): + offs_m = start_m + tl.arange(0, BLOCK_M) + valid_m = offs_m < seqlen_q + + q_ptrs = ( + Q + + (q_start + offs_m[:, None]) * stride_qm + + h * stride_qh + + offs_d[None, :] * stride_qk + ) + do_ptrs = ( + DO + + (q_start + offs_m[:, None]) * stride_dom + + h * stride_doh + + offs_d[None, :] * stride_dok + ) + q = tl.load(q_ptrs, mask=valid_m[:, None], other=0.0) + do = tl.load(do_ptrs, mask=valid_m[:, None], other=0.0) + + m_i = tl.load(M + (q_start + offs_m) * H + h, mask=valid_m, other=0.0) + l_i = tl.load(L + (q_start + offs_m) * H + h, mask=valid_m, other=1.0) + delta = tl.load(Delta + (q_start + offs_m) * H + h, mask=valid_m, other=0.0) + + qk = tl.dot(q, tl.trans(k)) * sm_scale + + mask = valid_n[None, :] & valid_m[:, None] + if IS_CAUSAL: + mask = mask & (offs_m[:, None] >= (offs_n[None, :] - causal_offset)) + qk = tl.where(mask, qk, float("-inf")) + + p = tl.exp(qk - m_i[:, None]) / l_i[:, None] + p = tl.where(mask, p, 0.0) + + dv += tl.dot(tl.trans(p.to(do.dtype)), do) + + dp = tl.dot(do, tl.trans(v)) + ds = p * (dp - delta[:, None]) * sm_scale + ds = tl.where(mask, ds, 0.0) + + dq_val = tl.dot(ds.to(q.dtype), k) + dq_ptrs = ( + DQ + + (q_start + offs_m[:, None]) * stride_qm + + h * stride_qh + + offs_d[None, :] * stride_qk + ) + tl.atomic_add(dq_ptrs, dq_val.to(q.dtype), mask=valid_m[:, None]) + + dk += tl.dot(tl.trans(ds.to(q.dtype)), q) + + dk_ptrs = ( + DK + (k_start + offs_n[:, None]) * stride_kn + h * stride_kh + offs_d[None, :] * stride_kk + ) + dv_ptrs = ( + DV + (k_start + offs_n[:, None]) * stride_vn + h * stride_vh + offs_d[None, :] * stride_vk + ) + tl.store(dk_ptrs, dk.to(k.dtype), mask=valid_n[:, None]) + tl.store(dv_ptrs, dv.to(v.dtype), mask=valid_n[:, None]) + + +class _TritonAttentionVarlen(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + causal, + sm_scale, + return_lse, + ): + total_q, H, head_dim = q.shape + assert k.shape[1] == H and v.shape[1] == H, "GQA is not supported by the varlen path yet" + assert head_dim in {16, 32, 64, 128, 256} + + batch = cu_seqlens_q.numel() - 1 + assert cu_seqlens_k.numel() - 1 == batch + + cu_seqlens_q = cu_seqlens_q.to(device=q.device, dtype=torch.int32) + cu_seqlens_k = cu_seqlens_k.to(device=q.device, dtype=torch.int32) + + out = torch.empty_like(q) + # Packed [total_q, H] layout (not [batch, H, seq_len]): total_q varies per + # batch, so there is no fixed per-sequence stride to lay these out densely. + M = torch.empty((total_q, H), device=q.device, dtype=torch.float32) + L = torch.empty((total_q, H), device=q.device, dtype=torch.float32) + + BLOCK_M = 64 + BLOCK_N = 64 if head_dim > 64 else 128 + + grid = (triton.cdiv(max_seqlen_q, BLOCK_M), batch * H) + _fwd_kernel_varlen[grid]( + q, + k, + v, + sm_scale, + cu_seqlens_q, + cu_seqlens_k, + L, + M, + out, + q.stride(0), + q.stride(1), + q.stride(2), + k.stride(0), + k.stride(1), + k.stride(2), + v.stride(0), + v.stride(1), + v.stride(2), + out.stride(0), + out.stride(1), + out.stride(2), + H, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + BLOCK_DMODEL=head_dim, + IS_CAUSAL=causal, + num_warps=4, + num_stages=2, + ) + + ctx.sm_scale = sm_scale + ctx.causal = causal + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.total_q = total_q + ctx.save_for_backward(q, k, v, out, L, M, cu_seqlens_q, cu_seqlens_k) + + if not return_lse: + return out, None + + lse = M + torch.log(L) + ctx.mark_non_differentiable(lse) + return out, lse + + @staticmethod + def backward(ctx, do, _dlse): + q, k, v, out, L, M, cu_seqlens_q, cu_seqlens_k = ctx.saved_tensors + + do = do.contiguous() + dq = torch.zeros_like(q) + dk = torch.empty_like(k) + dv = torch.empty_like(v) + + total_q, H, head_dim = q.shape + delta = torch.empty_like(L) + + BLOCK_M = 64 + BLOCK_N = 64 + + grid_prep = (triton.cdiv(total_q, BLOCK_M), H) + _bwd_preprocess_varlen[grid_prep]( + out, + do, + delta, + out.stride(0), + out.stride(1), + out.stride(2), + do.stride(0), + do.stride(1), + do.stride(2), + total_q, + H, + BLOCK_M=BLOCK_M, + D_HEAD=head_dim, + ) + + batch = cu_seqlens_q.numel() - 1 + grid_bwd = (triton.cdiv(ctx.max_seqlen_k, BLOCK_N), batch * H) + _bwd_kernel_varlen[grid_bwd]( + q, + k, + v, + ctx.sm_scale, + do, + dq, + dk, + dv, + L, + M, + delta, + cu_seqlens_q, + cu_seqlens_k, + q.stride(0), + q.stride(1), + q.stride(2), + k.stride(0), + k.stride(1), + k.stride(2), + v.stride(0), + v.stride(1), + v.stride(2), + do.stride(0), + do.stride(1), + do.stride(2), + H, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + BLOCK_DMODEL=head_dim, + IS_CAUSAL=ctx.causal, + num_warps=4, + num_stages=1, + ) + + return dq, dk, dv, None, None, None, None, None, None, None + + +def triton_flash_attention_varlen( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + causal: bool = True, + sm_scale: float | None = None, + return_lse: bool = False, +): + """ + Packed variable-length FlashAttention (Triton), the cross-platform semantic + baseline for RL rollout/training batches where sequences are concatenated + rather than padded. + + Args: + q: [total_q, H, D] packed queries. + k, v: [total_k, H, D] packed keys/values. GQA (Hk != Hq) is not supported + by this path yet (matches the existing dense Triton kernel's limitation). + cu_seqlens_q, cu_seqlens_k: int32 [batch + 1] cumulative sequence-length + offsets, cu_seqlens[0] == 0 (the `flash_attn_varlen_func` convention; + also what `pack`, #182, produces). + max_seqlen_q, max_seqlen_k: max per-sequence length in the batch (host + ints), used to size the launch grid. + causal: causal masking, anchored per-sequence via `Skv - Sq` (same + convention as `NativeAttentionOp` in docs/operators/attention.md, so + Sq == Skv is prefill and Sq < Skv is decode-with-cache). + sm_scale: defaults to `1/sqrt(D)`. + return_lse: if True, also return the packed `[total_q, H]` float32 + attention-domain LSE (non-differentiable; see `triton_flash_attention`). Returns: - Attention Output tensor of shape [batch, num_heads, seq_len, head_dim] + `out` of shape [total_q, H, D] if `return_lse` is False, else `(out, lse)`. """ if sm_scale is None: sm_scale = 1.0 / (q.shape[-1] ** 0.5) - return _TritonAttention.apply(q, k, v, causal, sm_scale) + out, lse = _TritonAttentionVarlen.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + causal, + sm_scale, + return_lse, + ) + return (out, lse) if return_lse else out diff --git a/setup.py b/setup.py index a4bd3b61..92070177 100644 --- a/setup.py +++ b/setup.py @@ -139,6 +139,7 @@ def get_extensions(): sm90_srcs = [ "csrc/cuda/fused_logp_sm90.cu", "csrc/cuda/fused_linear_logp_sm90.cu", # TMA + WGMMA fused linear log-prob + "csrc/cuda/attention/flash_attention_varlen_sm90.cu", # TMA + mma.sync varlen attention fwd ] enable_sm90 = envs.env_flag(envs.KERNEL_ALIGN_FORCE_SM90) present_sm90 = [s for s in sm90_srcs if os.path.exists(s)] diff --git a/tests/test_flash_attention_sm90.py b/tests/test_flash_attention_sm90.py new file mode 100644 index 00000000..3b0847f4 --- /dev/null +++ b/tests/test_flash_attention_sm90.py @@ -0,0 +1,631 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Tests for the SM90 (Hopper) TMA + `mma.sync` packed varlen FlashAttention +kernel: causal masking, `cu_seqlens` packing, attention-domain LSE export, and +a full autograd-differentiable forward + backward (dQ/dK/dV). + +Validated against two independent references: (a) an fp32 masked-softmax + +logsumexp closed form, differentiated via autograd (the same reference +tests/test_triton_attention_varlen.py uses), and (b) +`triton_flash_attention_varlen`, the cross-platform semantic baseline this +kernel is checked against -- both for forward output/LSE and for gradients. + +The exported LSE is attention-domain (per query row, over the key dimension), +not the vocab-domain LSE produced by the logp/linear_logp kernels. + +Note: `seqlen_q > seqlen_k` under `causal=True` is out-of-contract (violates +the documented `Skv - Sq` causal-offset convention, which assumes a query can +always attend at least to itself) and is not exercised here -- confirmed by +direct comparison that `triton_flash_attention_varlen` itself produces NaN +for that input too, so it isn't a case either implementation is designed to +handle, not a gap specific to this kernel. +""" + +import math + +import pytest +import torch + +from rl_engine.kernels.ops.cuda.attention.flash_attn_sm90 import flash_attention_sm90_varlen +from rl_engine.kernels.ops.triton.triton_attn import triton_flash_attention_varlen + + +def _sm90_available(): + """SM90 forward needs a Hopper GPU and the kernel compiled into the extension.""" + if not torch.cuda.is_available(): + return False + try: + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + + if not (_EXT_AVAILABLE and hasattr(_C, "flash_attention_varlen_sm90")): + return False + except Exception: # pragma: no cover + return False + return torch.cuda.get_device_capability()[0] == 9 + + +requires_sm90 = pytest.mark.skipif( + not _sm90_available(), + reason="flash_attention_varlen_sm90 requires a Hopper (sm_90) GPU with the extension " + "built KERNEL_ALIGN_FORCE_SM90=1.", +) + +_HEAD_DIM = 128 # only head_dim supported by the SM90 kernel this milestone + +_ATOL_OUT_REF = 2e-2 # bf16-mma kernel vs. independent fp32 reference +_ATOL_LSE_REF = 2e-2 +_ATOL_OUT_TRITON = 3e-2 # two independent bf16-precision compute paths, different tile schedules +_ATOL_LSE_TRITON = 3e-2 +_ATOL_GRAD_REF = 3e-2 # dq/dk/dv vs. fp32-autograd reference + + +def _ref_attn(q, k, v, causal, sm_scale): + """[H, Sq, D] / [H, Skv, D] fp32 masked-softmax reference with LSE.""" + H, Sq, D = q.shape + Skv = k.shape[1] + scores = torch.einsum("hqd,hkd->hqk", q, k) * sm_scale + if causal: + mask = torch.triu( + torch.ones(Sq, Skv, dtype=torch.bool, device=q.device), diagonal=Skv - Sq + 1 + ) + scores = scores.masked_fill(mask, float("-inf")) + probs = torch.softmax(scores, dim=-1) + out = torch.einsum("hqk,hkd->hqd", probs, v) + lse = torch.logsumexp(scores, dim=-1) + return out, lse + + +def _cu_seqlens(seqlens, device): + return torch.tensor( + [0, *torch.tensor(seqlens).cumsum(0).tolist()], dtype=torch.int32, device=device + ) + + +def _run_case(seqlens_q, seqlens_k, heads, causal, check_backward=True): + device = "cuda" + batch = len(seqlens_q) + total_q = sum(seqlens_q) + total_k = sum(seqlens_k) + cu_q = _cu_seqlens(seqlens_q, device) + cu_k = _cu_seqlens(seqlens_k, device) + + gen = torch.Generator(device=device).manual_seed(0) + q = torch.randn( + total_q, heads, _HEAD_DIM, device=device, dtype=torch.bfloat16, generator=gen + ).requires_grad_(check_backward) + k = torch.randn( + total_k, heads, _HEAD_DIM, device=device, dtype=torch.bfloat16, generator=gen + ).requires_grad_(check_backward) + v = torch.randn( + total_k, heads, _HEAD_DIM, device=device, dtype=torch.bfloat16, generator=gen + ).requires_grad_(check_backward) + + sm_scale = 1.0 / math.sqrt(_HEAD_DIM) + out, lse = flash_attention_sm90_varlen( + q, + k, + v, + cu_q, + cu_k, + max(seqlens_q), + max(seqlens_k), + causal=causal, + sm_scale=sm_scale, + return_lse=True, + ) + assert out.shape == (total_q, heads, _HEAD_DIM) + assert out.dtype == torch.bfloat16 + assert lse.shape == (total_q, heads) + assert lse.dtype == torch.float32 + assert not lse.requires_grad + + out_triton, lse_triton = triton_flash_attention_varlen( + q, + k, + v, + cu_q, + cu_k, + max(seqlens_q), + max(seqlens_k), + causal=causal, + sm_scale=sm_scale, + return_lse=True, + ) + + if check_backward: + # Gradient check is against the fp32 autograd reference only, not + # Triton's own backward: Triton's `_bwd_kernel_varlen` uses + # `tl.atomic_add` for dQ, which this Triton version (3.2.0) does not + # support on bf16 pointers at all (compiler error, not a numerical + # issue) -- confirmed directly, and consistent with + # tests/test_triton_attention_varlen.py itself only ever exercising + # backward with float16 tensors, never bfloat16. The SM90 kernel is + # bf16-only by design, so this comparison isn't possible here; the + # fp32 reference below is the rigorous check regardless. + do = torch.randn_like(out) + out.backward(do) + dq, dk, dv = q.grad.clone(), k.grad.clone(), v.grad.clone() + q.grad = k.grad = v.grad = None + + out_ref = torch.empty(total_q, heads, _HEAD_DIM, device=device, dtype=torch.float32) + lse_ref = torch.empty(total_q, heads, device=device, dtype=torch.float32) + q_ref = q.detach().clone().float().requires_grad_(check_backward) + k_ref = k.detach().clone().float().requires_grad_(check_backward) + v_ref = v.detach().clone().float().requires_grad_(check_backward) + qs = ks = 0 + for b in range(batch): + sq, sk = seqlens_q[b], seqlens_k[b] + if sq > 0: + o, lval = _ref_attn( + q_ref[qs : qs + sq].transpose(0, 1), + k_ref[ks : ks + sk].transpose(0, 1), + v_ref[ks : ks + sk].transpose(0, 1), + causal, + sm_scale, + ) + out_ref[qs : qs + sq] = o.transpose(0, 1) + lse_ref[qs : qs + sq] = lval.transpose(0, 1) + qs += sq + ks += sk + + torch.testing.assert_close(out.float(), out_ref, atol=_ATOL_OUT_REF, rtol=0.0) + torch.testing.assert_close(lse, lse_ref, atol=_ATOL_LSE_REF, rtol=0.0) + torch.testing.assert_close(out.float(), out_triton.float(), atol=_ATOL_OUT_TRITON, rtol=0.0) + torch.testing.assert_close(lse, lse_triton, atol=_ATOL_LSE_TRITON, rtol=0.0) + + if check_backward: + out_ref.backward(do.float()) + torch.testing.assert_close(dq.float(), q_ref.grad, atol=_ATOL_GRAD_REF, rtol=0.0) + torch.testing.assert_close(dk.float(), k_ref.grad, atol=_ATOL_GRAD_REF, rtol=0.0) + torch.testing.assert_close(dv.float(), v_ref.grad, atol=_ATOL_GRAD_REF, rtol=0.0) + torch.testing.assert_close(out.float(), out_triton.float(), atol=_ATOL_OUT_TRITON, rtol=0.0) + torch.testing.assert_close(lse, lse_triton, atol=_ATOL_LSE_TRITON, rtol=0.0) + + +@requires_sm90 +class TestFlashAttentionSM90Varlen: + def test_prefill_uneven_seqlens_not_block_aligned(self): + # BLOCK_Q/BLOCK_KV are 64/64; deliberately not multiples of either. + _run_case([37, 128, 200, 5], [37, 128, 200, 5], heads=4, causal=True) + + def test_non_causal(self): + _run_case([37, 130, 61], [37, 130, 61], heads=2, causal=False) + + def test_decode_style_sq_less_than_skv(self): + # Small new-query chunk (e.g. rollout decode step) against a longer KV cache, + # varying independently per sequence in the batch. + _run_case([1, 3, 1], [50, 91, 17], heads=4, causal=True) + + def test_head_dim_128(self): + # head_dim=128 is the only dim this kernel supports -- kept as its own + # named case for documentation value even though every other case + # above already uses it. + _run_case([100, 260], [100, 260], heads=2, causal=True) + + def test_zero_length_sequence_in_batch(self): + # A fully-masked / empty response is a real occurrence in packed RL + # batches; exercises the per-CTA early-return-on-seqlen_q==0 path. + _run_case([0, 128, 0, 37], [50, 128, 0, 37], heads=2, causal=True) + + def test_dense_vs_varlen_consistency(self): + # batch>1 packed in one call must match calling the same kernel once + # per sequence (batch=1 each time) -- confirms the packed (b,h) grid + # decomposition doesn't leak state across CTAs. + device = "cuda" + heads = 2 + seqlens = [40, 91, 17] + sm_scale = 1.0 / math.sqrt(_HEAD_DIM) + total = sum(seqlens) + gen = torch.Generator(device=device).manual_seed(1) + q = torch.randn(total, heads, _HEAD_DIM, device=device, dtype=torch.bfloat16, generator=gen) + k = torch.randn(total, heads, _HEAD_DIM, device=device, dtype=torch.bfloat16, generator=gen) + v = torch.randn(total, heads, _HEAD_DIM, device=device, dtype=torch.bfloat16, generator=gen) + cu = _cu_seqlens(seqlens, device) + + out_packed, lse_packed = flash_attention_sm90_varlen( + q, k, v, cu, cu, max(seqlens), max(seqlens), causal=True, sm_scale=sm_scale, + return_lse=True, + ) + + qs = 0 + for s in seqlens: + q_i = q[qs : qs + s].contiguous() + k_i = k[qs : qs + s].contiguous() + v_i = v[qs : qs + s].contiguous() + cu_i = torch.tensor([0, s], dtype=torch.int32, device=device) + out_i, lse_i = flash_attention_sm90_varlen( + q_i, k_i, v_i, cu_i, cu_i, s, s, causal=True, sm_scale=sm_scale, return_lse=True + ) + torch.testing.assert_close( + out_packed[qs : qs + s].float(), out_i.float(), atol=1e-3, rtol=0.0 + ) + torch.testing.assert_close(lse_packed[qs : qs + s], lse_i, atol=1e-3, rtol=0.0) + qs += s + + def test_non_contiguous_rejected(self): + # Deliberate divergence from the Triton path (which tolerates + # non-contiguous q/k/v, see test_non_contiguous_q_k_v_backward in + # tests/test_triton_attention_varlen.py): the compiled SM90 kernel's + # TMA descriptors assume a fixed contiguous row stride, so the C++ + # wrapper TORCH_CHECKs contiguity and raises rather than silently + # calling .contiguous() to hide a caller bug. + device = "cuda" + heads = 2 + seqlens = [37, 61] + total = sum(seqlens) + cu = _cu_seqlens(seqlens, device) + base = torch.randn(heads, total, _HEAD_DIM, device=device, dtype=torch.bfloat16) + q = base.transpose(0, 1) + assert not q.is_contiguous() + + from rl_engine.kernels.ops.base import _C + + with pytest.raises(RuntimeError): + _C.flash_attention_varlen_sm90( + q, q, q, cu, cu, max(seqlens), max(seqlens), True, 1.0 / math.sqrt(_HEAD_DIM) + ) + + def test_block_aligned_seqlens(self): + # BLOCK_Q/BLOCK_KV are 64/64; every other case deliberately avoids + # exact multiples. An off-by-one in the `row < seqlen_q` boundary + # guard (e.g. `<=` instead of `<`) would only surface when a Q-tile + # is exactly full -- no partial tail row to mask -- so this needs its + # own case distinct from the non-aligned ones. + _run_case([64, 128, 192], [64, 128, 192], heads=2, causal=True) + _run_case([64, 128, 192], [64, 128, 192], heads=2, causal=False) + + def test_long_sequence_many_kv_iterations(self): + # STAGES=2 double-buffering is only exercised across a handful of + # wraparounds by the other cases (seqlen_k up to 260 -> ~5 KV tiles). + # Push into the dozens of iterations to stress the mbarrier + # prefetch/consume pipeline for a buffer-reuse bug that only shows up + # after many wraparounds. + _run_case([1536, 777], [1536, 777], heads=2, causal=True) + + def test_single_head(self): + # H=1 collapses the blockIdx.y = b*H+h decomposition to pure `b`; + # make sure that degenerate case (division/modulo by 1) isn't broken. + _run_case([37, 128, 61], [37, 128, 61], heads=1, causal=True) + + def test_many_batch_entries(self): + # Stress cu_seqlens indexing / grid.y decomposition across more + # batch entries than any other case (8, mixed short/long/zero). + _run_case( + [3, 0, 64, 129, 1, 200, 0, 17], + [3, 0, 64, 129, 1, 200, 0, 17], + heads=2, + causal=True, + ) + + def test_invalid_head_dim_rejected(self): + # D=128 is the only head_dim this milestone; the C++ wrapper should + # reject anything else rather than silently misinterpreting the + # tensor map layout. + device = "cuda" + seqlens = [16, 20] + cu = _cu_seqlens(seqlens, device) + total = sum(seqlens) + q = torch.randn(total, 2, 64, device=device, dtype=torch.bfloat16) + + from rl_engine.kernels.ops.base import _C + + with pytest.raises(RuntimeError): + _C.flash_attention_varlen_sm90( + q, q, q, cu, cu, max(seqlens), max(seqlens), True, 1.0 / math.sqrt(64) + ) + + def test_invalid_dtype_rejected(self): + device = "cuda" + seqlens = [16, 20] + cu = _cu_seqlens(seqlens, device) + total = sum(seqlens) + q = torch.randn(total, 2, _HEAD_DIM, device=device, dtype=torch.float16) + + from rl_engine.kernels.ops.base import _C + + with pytest.raises(RuntimeError): + _C.flash_attention_varlen_sm90( + q, q, q, cu, cu, max(seqlens), max(seqlens), True, 1.0 / math.sqrt(_HEAD_DIM) + ) + + def test_gqa_mismatch_rejected(self): + # k/v with fewer heads than q (GQA) is unsupported; the wrapper + # should reject rather than reading past k/v's head dimension. + device = "cuda" + seqlens = [16, 20] + cu = _cu_seqlens(seqlens, device) + total = sum(seqlens) + q = torch.randn(total, 4, _HEAD_DIM, device=device, dtype=torch.bfloat16) + kv = torch.randn(total, 2, _HEAD_DIM, device=device, dtype=torch.bfloat16) + + from rl_engine.kernels.ops.base import _C + + with pytest.raises(RuntimeError): + _C.flash_attention_varlen_sm90( + q, kv, kv, cu, cu, max(seqlens), max(seqlens), True, 1.0 / math.sqrt(_HEAD_DIM) + ) + + def test_cu_seqlens_batch_mismatch_rejected(self): + device = "cuda" + q = torch.randn(36, 2, _HEAD_DIM, device=device, dtype=torch.bfloat16) + cu_q = _cu_seqlens([16, 20], device) # batch=2 + cu_k = _cu_seqlens([12, 12, 12], device) # batch=3 + + from rl_engine.kernels.ops.base import _C + + with pytest.raises(RuntimeError): + _C.flash_attention_varlen_sm90( + q, q, q, cu_q, cu_k, 20, 12, True, 1.0 / math.sqrt(_HEAD_DIM) + ) + + def test_wrapper_falls_back_for_unsupported_dtype(self): + # End-to-end check of the Python-level dispatch gate (`_supported`), + # not just that the raw C++ op rejects bad input: an fp16 tensor + # should silently route through `flash_attention_sm90_varlen` to the + # Triton fallback and still produce a correct result, not raise. + device = "cuda" + heads = 2 + seqlens = [37, 61] + total = sum(seqlens) + cu = _cu_seqlens(seqlens, device) + sm_scale = 1.0 / math.sqrt(_HEAD_DIM) + gen = torch.Generator(device=device).manual_seed(2) + q = torch.randn(total, heads, _HEAD_DIM, device=device, dtype=torch.float16, generator=gen) + k = torch.randn(total, heads, _HEAD_DIM, device=device, dtype=torch.float16, generator=gen) + v = torch.randn(total, heads, _HEAD_DIM, device=device, dtype=torch.float16, generator=gen) + + out, lse = flash_attention_sm90_varlen( + q, k, v, cu, cu, max(seqlens), max(seqlens), causal=True, sm_scale=sm_scale, + return_lse=True, + ) + + out_triton, lse_triton = triton_flash_attention_varlen( + q, k, v, cu, cu, max(seqlens), max(seqlens), causal=True, sm_scale=sm_scale, + return_lse=True, + ) + # Should be the *same* Triton call under the hood (fp16 is unsupported + # by the fused kernel), so this should match tightly, not just within + # the cross-implementation tolerance used elsewhere in this file. + torch.testing.assert_close(out, out_triton) + torch.testing.assert_close(lse, lse_triton) + + def test_single_kv_tile_backward(self): + # Isolates the P/ds shared-memory transpose-trick correctness (dV, + # dK, dQ) from the causal `lo`-bound loop logic and cross-CTA dQ + # atomics -- both seqlen_q and seqlen_k fit in one 64-row tile. + _run_case([32], [32], heads=1, causal=True) + + def test_dq_atomic_accumulation_stress(self): + # A short query sequence attended to by many KV-tiles (long + # seqlen_k, short seqlen_q, non-causal so every one of the ~16 + # KV-tile CTAs contributes an atomic-add to the same few dQ rows) -- + # the one scenario no other case exercises, since it's purely about + # backward's cross-CTA accumulation correctness. + _run_case([4], [1000], heads=2, causal=False) + + def test_dq_atomic_accumulation_stress_multi_batch(self): + # Same stress, but across three interleaved batch entries -- checks + # atomics from different (b, h) CTAs never cross into the wrong + # batch's dQ rows (packed layout means adjacent batches' dQ ranges + # are directly contiguous in memory, so a k_start/q_start indexing + # slip here would corrupt a neighboring batch's gradient, not just + # this one's). + _run_case([4, 2, 8], [1000, 500, 700], heads=2, causal=False) + + def test_backward_without_return_lse(self): + # lse is computed unconditionally by the forward kernel and always + # saved for backward's recompute, regardless of whether the caller + # asked for it back -- `return_lse=False` must not silently break + # differentiability. + device = "cuda" + seqlens = [37, 61] + cu = _cu_seqlens(seqlens, device) + total = sum(seqlens) + sm_scale = 1.0 / math.sqrt(_HEAD_DIM) + gen = torch.Generator(device=device).manual_seed(1) + q = torch.randn( + total, 2, _HEAD_DIM, device=device, dtype=torch.bfloat16, generator=gen + ).requires_grad_() + k = torch.randn( + total, 2, _HEAD_DIM, device=device, dtype=torch.bfloat16, generator=gen + ).requires_grad_() + v = torch.randn( + total, 2, _HEAD_DIM, device=device, dtype=torch.bfloat16, generator=gen + ).requires_grad_() + + out = flash_attention_sm90_varlen( + q, k, v, cu, cu, max(seqlens), max(seqlens), causal=True, sm_scale=sm_scale, + return_lse=False, + ) + assert isinstance(out, torch.Tensor) # bare tensor, not a tuple + do = torch.randn_like(out) + out.backward(do) + assert q.grad is not None and k.grad is not None and v.grad is not None + + q_ref = q.detach().clone().float().requires_grad_() + k_ref = k.detach().clone().float().requires_grad_() + v_ref = v.detach().clone().float().requires_grad_() + out_ref = torch.empty(total, 2, _HEAD_DIM, device=device, dtype=torch.float32) + qs = 0 + for s in seqlens: + o, _ = _ref_attn( + q_ref[qs : qs + s].transpose(0, 1), + k_ref[qs : qs + s].transpose(0, 1), + v_ref[qs : qs + s].transpose(0, 1), + True, + sm_scale, + ) + out_ref[qs : qs + s] = o.transpose(0, 1) + qs += s + out_ref.backward(do.float()) + torch.testing.assert_close(q.grad.float(), q_ref.grad, atol=_ATOL_GRAD_REF, rtol=0.0) + torch.testing.assert_close(k.grad.float(), k_ref.grad, atol=_ATOL_GRAD_REF, rtol=0.0) + torch.testing.assert_close(v.grad.float(), v_ref.grad, atol=_ATOL_GRAD_REF, rtol=0.0) + + def test_partial_requires_grad(self): + # Frozen-KV-cache-style training (q requires grad, k/v don't) and the + # reverse; PyTorch's autograd.Function contract allows returning + # gradients for inputs that don't need them (just extra, discarded + # work), but this confirms the kernel doesn't crash or corrupt the + # requested gradients when only some of q/k/v participate. + device = "cuda" + seqlens = [37, 61] + cu = _cu_seqlens(seqlens, device) + total = sum(seqlens) + sm_scale = 1.0 / math.sqrt(_HEAD_DIM) + + def make(seed, rq, rk, rv): + gen = torch.Generator(device=device).manual_seed(seed) + q = torch.randn( + total, 2, _HEAD_DIM, device=device, dtype=torch.bfloat16, generator=gen + ).requires_grad_(rq) + k = torch.randn( + total, 2, _HEAD_DIM, device=device, dtype=torch.bfloat16, generator=gen + ).requires_grad_(rk) + v = torch.randn( + total, 2, _HEAD_DIM, device=device, dtype=torch.bfloat16, generator=gen + ).requires_grad_(rv) + return q, k, v + + for rq, rk, rv in [(True, False, False), (False, True, True), (True, True, False)]: + q, k, v = make(2, rq, rk, rv) + out, _ = flash_attention_sm90_varlen( + q, k, v, cu, cu, max(seqlens), max(seqlens), causal=True, sm_scale=sm_scale, + return_lse=True, + ) + out.backward(torch.randn_like(out)) + assert (q.grad is not None) == rq + assert (k.grad is not None) == rk + assert (v.grad is not None) == rv + + def test_backward_non_contiguous_do(self): + # The upstream gradient `do` (unlike q/k/v) is not under this + # kernel's contiguity contract -- backward calls `do.contiguous()` + # internally, matching the Triton path's handling of the same + # argument. Build `do` as a genuinely non-contiguous transposed view. + device = "cuda" + seqlens = [37, 61] + cu = _cu_seqlens(seqlens, device) + total = sum(seqlens) + sm_scale = 1.0 / math.sqrt(_HEAD_DIM) + gen = torch.Generator(device=device).manual_seed(4) + q = torch.randn( + total, 2, _HEAD_DIM, device=device, dtype=torch.bfloat16, generator=gen + ).requires_grad_() + k = torch.randn( + total, 2, _HEAD_DIM, device=device, dtype=torch.bfloat16, generator=gen + ).requires_grad_() + v = torch.randn( + total, 2, _HEAD_DIM, device=device, dtype=torch.bfloat16, generator=gen + ).requires_grad_() + + out, _ = flash_attention_sm90_varlen( + q, k, v, cu, cu, max(seqlens), max(seqlens), causal=True, sm_scale=sm_scale, + return_lse=True, + ) + do = torch.randn(2, total, _HEAD_DIM, device=device, dtype=torch.bfloat16).transpose(0, 1) + assert not do.is_contiguous() + out.backward(do) + dq, dk, dv = q.grad.clone(), k.grad.clone(), v.grad.clone() + + q_ref = q.detach().clone().float().requires_grad_() + k_ref = k.detach().clone().float().requires_grad_() + v_ref = v.detach().clone().float().requires_grad_() + out_ref = torch.empty(total, 2, _HEAD_DIM, device=device, dtype=torch.float32) + qs = 0 + for s in seqlens: + o, _ = _ref_attn( + q_ref[qs : qs + s].transpose(0, 1), + k_ref[qs : qs + s].transpose(0, 1), + v_ref[qs : qs + s].transpose(0, 1), + True, + sm_scale, + ) + out_ref[qs : qs + s] = o.transpose(0, 1) + qs += s + out_ref.backward(do.float()) + torch.testing.assert_close(dq.float(), q_ref.grad, atol=_ATOL_GRAD_REF, rtol=0.0) + torch.testing.assert_close(dk.float(), k_ref.grad, atol=_ATOL_GRAD_REF, rtol=0.0) + torch.testing.assert_close(dv.float(), v_ref.grad, atol=_ATOL_GRAD_REF, rtol=0.0) + + def test_lse_is_non_differentiable_and_backward_still_works(self): + device = "cuda" + heads = 2 + seqlens = [40, 91] + cu = _cu_seqlens(seqlens, device) + total = sum(seqlens) + sm_scale = 1.0 / math.sqrt(_HEAD_DIM) + q = torch.randn( + total, heads, _HEAD_DIM, device=device, dtype=torch.bfloat16, requires_grad=True + ) + k = torch.randn( + total, heads, _HEAD_DIM, device=device, dtype=torch.bfloat16, requires_grad=True + ) + v = torch.randn( + total, heads, _HEAD_DIM, device=device, dtype=torch.bfloat16, requires_grad=True + ) + + out, lse = flash_attention_sm90_varlen( + q, k, v, cu, cu, max(seqlens), max(seqlens), causal=True, sm_scale=sm_scale, + return_lse=True, + ) + assert not lse.requires_grad + + out.float().sum().backward() + assert q.grad is not None and k.grad is not None and v.grad is not None + + def test_backward_validation_guards(self): + # The C++ backward wrapper's TORCH_CHECKs, mirroring the forward + # guard tests -- previously unverified for the backward symbol + # specifically (bad head_dim/dtype/GQA/cu_seqlens-mismatch/ + # non-contiguous should all be rejected, not just forward's). + device = "cuda" + seqlens = [16, 20] + cu = _cu_seqlens(seqlens, device) + total = sum(seqlens) + sm_scale = 1.0 / math.sqrt(_HEAD_DIM) + + from rl_engine.kernels.ops.base import _C + + def call(q, k, v, do, out, lse): + return _C.flash_attention_varlen_sm90_backward( + do, q, k, v, out, lse, cu, cu, max(seqlens), True, sm_scale + ) + + q_ok = torch.randn(total, 2, _HEAD_DIM, device=device, dtype=torch.bfloat16) + do_ok = torch.randn_like(q_ok) + out_ok = torch.randn_like(q_ok) + lse_ok = torch.randn(total, 2, device=device, dtype=torch.float32) + + # Bad head_dim. + q_bad_dim = torch.randn(total, 2, 64, device=device, dtype=torch.bfloat16) + with pytest.raises(RuntimeError): + call(q_bad_dim, q_bad_dim, q_bad_dim, q_bad_dim, q_bad_dim, + torch.randn(total, 2, device=device, dtype=torch.float32)) + + # Bad dtype (fp16 instead of bf16). + q_fp16 = torch.randn(total, 2, _HEAD_DIM, device=device, dtype=torch.float16) + with pytest.raises(RuntimeError): + call(q_fp16, q_fp16, q_fp16, q_fp16, q_fp16, + torch.randn(total, 2, device=device, dtype=torch.float32)) + + # GQA mismatch. + q_4h = torch.randn(total, 4, _HEAD_DIM, device=device, dtype=torch.bfloat16) + do_4h = torch.randn_like(q_4h) + out_4h = torch.randn_like(q_4h) + lse_4h = torch.randn(total, 4, device=device, dtype=torch.float32) + with pytest.raises(RuntimeError): + call(q_4h, q_ok, q_ok, do_4h, out_4h, lse_4h) + + # Non-contiguous. + q_noncontig = torch.randn(2, total, _HEAD_DIM, device=device, dtype=torch.bfloat16).transpose(0, 1) + with pytest.raises(RuntimeError): + call(q_noncontig, q_ok, q_ok, do_ok, out_ok, lse_ok) + + # cu_seqlens batch mismatch. + cu_q_bad = _cu_seqlens([16, 20], device) + cu_k_bad = _cu_seqlens([12, 12, 12], device) + with pytest.raises(RuntimeError): + _C.flash_attention_varlen_sm90_backward( + do_ok, q_ok, q_ok, q_ok, out_ok, lse_ok, cu_q_bad, cu_k_bad, 12, True, sm_scale + ) diff --git a/tests/test_triton_attention_varlen.py b/tests/test_triton_attention_varlen.py new file mode 100644 index 00000000..45c7d165 --- /dev/null +++ b/tests/test_triton_attention_varlen.py @@ -0,0 +1,257 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Tests for the Triton FlashAttention LSE export and packed varlen path. + +Extends `rl_engine.kernels.ops.triton.triton_attn` (the "cross-platform semantic +baseline" the SM90/SM80/ROCm fused-attention kernels will validate against) with: + + * `return_lse=True` on the existing dense kernel: exposes the attention-domain + LSE (log-sum-exp of the scaled, masked QK^T logits) already accumulated + online as `M`/`L` for the backward pass, but previously discarded. + * `triton_flash_attention_varlen`: packed `[total_tokens, H, D]` + `cu_seqlens` + attention, for RL rollout/training batches that are concatenated rather than + padded (same convention as `flash_attn_varlen_func` and this repo's `pack` + op, #182). + +Both are checked against an independent per-sequence fp32 reference (`_ref_attn`, +a masked-softmax + logsumexp closed form matching the causal-offset convention +`Skv - Sq` documented for `NativeAttentionOp` in docs/operators/attention.md), +not against the dense Triton kernel itself, so a shared bug in the online-softmax +loop cannot hide from these tests. + +The exported LSE is attention-domain (per query row, over the key dimension), +not the vocab-domain LSE produced by the logp/linear_logp kernels. +""" + +import math + +import pytest +import torch + +from rl_engine.kernels.ops.triton.triton_attn import ( + triton_flash_attention, + triton_flash_attention_varlen, +) + +try: + import triton # noqa: F401 + + _HAS_TRITON = True +except ImportError: # pragma: no cover + _HAS_TRITON = False + +requires_triton_cuda = pytest.mark.skipif( + not (_HAS_TRITON and torch.cuda.is_available()), + reason="Triton attention requires a CUDA device and Triton.", +) + +_ATOL_OUT = 0.05 # fp16 accumulation tolerance, dense and varlen paths alike +_ATOL_LSE = 1e-2 + + +def _ref_attn(q, k, v, causal, sm_scale): + """[H, Sq, D] / [H, Skv, D] fp32 masked-softmax reference with LSE.""" + H, Sq, D = q.shape + Skv = k.shape[1] + scores = torch.einsum("hqd,hkd->hqk", q, k) * sm_scale + if causal: + mask = torch.triu( + torch.ones(Sq, Skv, dtype=torch.bool, device=q.device), diagonal=Skv - Sq + 1 + ) + scores = scores.masked_fill(mask, float("-inf")) + probs = torch.softmax(scores, dim=-1) + out = torch.einsum("hqk,hkd->hqd", probs, v) + lse = torch.logsumexp(scores, dim=-1) + return out, lse + + +def _run_varlen_case(seqlens_q, seqlens_k, heads, head_dim, causal, dtype=torch.float16): + device = "cuda" + batch = len(seqlens_q) + total_q = sum(seqlens_q) + total_k = sum(seqlens_k) + cu_q = torch.tensor( + [0, *torch.tensor(seqlens_q).cumsum(0).tolist()], dtype=torch.int32, device=device + ) + cu_k = torch.tensor( + [0, *torch.tensor(seqlens_k).cumsum(0).tolist()], dtype=torch.int32, device=device + ) + + q = torch.randn(total_q, heads, head_dim, device=device, dtype=dtype, requires_grad=True) + k = torch.randn(total_k, heads, head_dim, device=device, dtype=dtype, requires_grad=True) + v = torch.randn(total_k, heads, head_dim, device=device, dtype=dtype, requires_grad=True) + + sm_scale = 1.0 / math.sqrt(head_dim) + out, lse = triton_flash_attention_varlen( + q, + k, + v, + cu_q, + cu_k, + max(seqlens_q), + max(seqlens_k), + causal=causal, + sm_scale=sm_scale, + return_lse=True, + ) + do = torch.randn_like(out) + out.backward(do) + dq, dk, dv = q.grad.clone(), k.grad.clone(), v.grad.clone() + + q_ref = q.detach().clone().float().requires_grad_() + k_ref = k.detach().clone().float().requires_grad_() + v_ref = v.detach().clone().float().requires_grad_() + out_ref = torch.empty(total_q, heads, head_dim, device=device, dtype=torch.float32) + lse_ref = torch.empty(total_q, heads, device=device, dtype=torch.float32) + + qs = ks = 0 + for b in range(batch): + sq, sk = seqlens_q[b], seqlens_k[b] + o, lval = _ref_attn( + q_ref[qs : qs + sq].transpose(0, 1), + k_ref[ks : ks + sk].transpose(0, 1), + v_ref[ks : ks + sk].transpose(0, 1), + causal, + sm_scale, + ) + out_ref[qs : qs + sq] = o.transpose(0, 1) + lse_ref[qs : qs + sq] = lval.transpose(0, 1) + qs += sq + ks += sk + out_ref.backward(do.float()) + + torch.testing.assert_close(out.float(), out_ref, atol=_ATOL_OUT, rtol=0.0) + torch.testing.assert_close(lse, lse_ref, atol=_ATOL_LSE, rtol=0.0) + torch.testing.assert_close(dq.float(), q_ref.grad, atol=_ATOL_OUT, rtol=0.0) + torch.testing.assert_close(dk.float(), k_ref.grad, atol=_ATOL_OUT, rtol=0.0) + torch.testing.assert_close(dv.float(), v_ref.grad, atol=_ATOL_OUT, rtol=0.0) + + +@requires_triton_cuda +class TestDenseLSEExport: + def test_lse_matches_reference(self): + device = "cuda" + batch, heads, seq, head_dim = 2, 4, 256, 64 + sm_scale = 1.0 / math.sqrt(head_dim) + q = torch.randn(batch, heads, seq, head_dim, device=device, dtype=torch.float16) + k = torch.randn(batch, heads, seq, head_dim, device=device, dtype=torch.float16) + v = torch.randn(batch, heads, seq, head_dim, device=device, dtype=torch.float16) + + out, lse = triton_flash_attention(q, k, v, causal=True, sm_scale=sm_scale, return_lse=True) + assert lse.shape == (batch, heads, seq) + assert lse.dtype == torch.float32 + + ref_out = torch.empty_like(q, dtype=torch.float32) + ref_lse = torch.empty(batch, heads, seq, device=device) + for b in range(batch): + o, lval = _ref_attn(q[b].float(), k[b].float(), v[b].float(), True, sm_scale) + ref_out[b], ref_lse[b] = o, lval + + torch.testing.assert_close(out.float(), ref_out, atol=_ATOL_OUT, rtol=0.0) + torch.testing.assert_close(lse, ref_lse, atol=_ATOL_LSE, rtol=0.0) + + def test_lse_is_non_differentiable_and_backward_still_works(self): + device = "cuda" + q = torch.randn(1, 2, 64, 64, device=device, dtype=torch.float16, requires_grad=True) + k = torch.randn(1, 2, 64, 64, device=device, dtype=torch.float16, requires_grad=True) + v = torch.randn(1, 2, 64, 64, device=device, dtype=torch.float16, requires_grad=True) + + out, lse = triton_flash_attention(q, k, v, causal=True, return_lse=True) + assert not lse.requires_grad + + out.float().sum().backward() + assert q.grad is not None and k.grad is not None and v.grad is not None + + def test_default_call_is_backward_compatible(self): + device = "cuda" + q = torch.randn(1, 2, 64, 64, device=device, dtype=torch.float16) + k = torch.randn(1, 2, 64, 64, device=device, dtype=torch.float16) + v = torch.randn(1, 2, 64, 64, device=device, dtype=torch.float16) + + out = triton_flash_attention(q, k, v, causal=True) + assert isinstance(out, torch.Tensor) + + +@requires_triton_cuda +class TestVarlenAttention: + def test_prefill_uneven_seqlens_not_block_aligned(self): + # BLOCK_M/BLOCK_N default to 64/128; deliberately not multiples of either. + _run_varlen_case([37, 128, 200, 5], [37, 128, 200, 5], heads=4, head_dim=64, causal=True) + + def test_non_causal(self): + _run_varlen_case([37, 130, 61], [37, 130, 61], heads=2, head_dim=64, causal=False) + + def test_decode_style_sq_less_than_skv(self): + # Small new-query chunk (e.g. rollout decode step) against a longer KV cache, + # varying independently per sequence in the batch. + _run_varlen_case([1, 3, 1], [50, 91, 17], heads=4, head_dim=64, causal=True) + + def test_head_dim_128(self): + _run_varlen_case([100, 260], [100, 260], heads=2, head_dim=128, causal=True) + + def test_zero_length_sequence_in_batch(self): + # A fully-masked / empty response is a real occurrence in packed RL batches. + _run_varlen_case([0, 128, 0, 37], [50, 128, 0, 37], heads=2, head_dim=64, causal=True) + + def test_non_contiguous_q_k_v_backward(self): + # Regression test (review finding): `do` was forced `.contiguous()` in + # backward, but the kernels indexed DO using Q's/Out's strides instead of + # DO's own -- silently correct only when Q/Out/DO all happen to share the + # same contiguous layout, as every other case in this file does by + # construction (fresh `torch.randn(...)`). + # + # Building q/k/v as `[H, total, D].transpose(0, 1)` gives a non-contiguous + # but non-overlapping-and-dense view. `torch.empty_like` (used for `out` + # and `dq`) preserves that exact stride pattern rather than falling back to + # a contiguous layout, so `out`'s strides differ from `do`'s (`do` is a + # plain contiguous tensor here) -- exactly the mismatch the review flagged. + device = "cuda" + heads, head_dim = 2, 64 + seqlens = [37, 61] + total = sum(seqlens) + cu = torch.tensor( + [0, *torch.tensor(seqlens).cumsum(0).tolist()], dtype=torch.int32, device=device + ) + sm_scale = 1.0 / math.sqrt(head_dim) + + def make(seed): + gen = torch.Generator(device=device).manual_seed(seed) + base = torch.randn( + heads, total, head_dim, device=device, dtype=torch.float16, generator=gen + ) + return base.transpose(0, 1).detach().requires_grad_() + + q, k, v = make(1), make(2), make(3) + assert not (q.is_contiguous() or k.is_contiguous() or v.is_contiguous()) + + out = triton_flash_attention_varlen( + q, k, v, cu, cu, max(seqlens), max(seqlens), causal=True, sm_scale=sm_scale + ) + assert not out.is_contiguous() # confirms empty_like preserved q's layout + + do = torch.randn(total, heads, head_dim, device=device, dtype=out.dtype) # plain contiguous + assert do.stride() != out.stride() + out.backward(do) + dq, dk, dv = q.grad.clone(), k.grad.clone(), v.grad.clone() + + q_ref = q.detach().clone().float().requires_grad_() + k_ref = k.detach().clone().float().requires_grad_() + v_ref = v.detach().clone().float().requires_grad_() + out_ref = torch.empty(total, heads, head_dim, device=device, dtype=torch.float32) + qs = 0 + for sq in seqlens: + o, _ = _ref_attn( + q_ref[qs : qs + sq].transpose(0, 1), + k_ref[qs : qs + sq].transpose(0, 1), + v_ref[qs : qs + sq].transpose(0, 1), + True, + sm_scale, + ) + out_ref[qs : qs + sq] = o.transpose(0, 1) + qs += sq + out_ref.backward(do.float()) + + torch.testing.assert_close(out.float(), out_ref, atol=_ATOL_OUT, rtol=0.0) + torch.testing.assert_close(dq.float(), q_ref.grad, atol=_ATOL_OUT, rtol=0.0) + torch.testing.assert_close(dk.float(), k_ref.grad, atol=_ATOL_OUT, rtol=0.0) + torch.testing.assert_close(dv.float(), v_ref.grad, atol=_ATOL_OUT, rtol=0.0)