diff --git a/benchmarks/benchmark_deterministic_attention.py b/benchmarks/benchmark_deterministic_attention.py new file mode 100644 index 00000000..424ebf5f --- /dev/null +++ b/benchmarks/benchmark_deterministic_attention.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Benchmark for CUDA deterministic standard-softmax attention (issue #147). + +Reports latency and peak memory for Qwen3-8B representative shapes, +including the cost of full scores/P materialization. +""" + +import argparse +import time + +import torch + + +def benchmark_attention( + B: int, Hq: int, Hkv: int, Sq: int, Skv: int, D: int, dtype, warmup: int = 5, iters: int = 20 +): + from rl_engine.kernels.ops.cuda.attention.deterministic_attn import DeterministicAttentionOp + + op = DeterministicAttentionOp() + device = "cuda" + + q = torch.randn(B, Hq, Sq, D, device=device, dtype=dtype) + k = torch.randn(B, Hkv, Skv, D, device=device, dtype=dtype) + v = torch.randn(B, Hkv, Skv, D, device=device, dtype=dtype) + + torch.cuda.reset_peak_memory_stats() + + with torch.no_grad(): + for _ in range(warmup): + op.forward(q, k, v, causal=True) + torch.cuda.synchronize() + + start = time.perf_counter() + for _ in range(iters): + op.forward(q, k, v, causal=True) + torch.cuda.synchronize() + elapsed = (time.perf_counter() - start) / iters + + peak_mem_mb = torch.cuda.max_memory_allocated() / (1024 * 1024) + + scores_mem_mb = (B * Hq * Sq * Skv * 4) / (1024 * 1024) + + return { + "latency_ms": elapsed * 1000, + "peak_memory_mb": peak_mem_mb, + "scores_materialization_mb": scores_mem_mb, + } + + +QWEN3_8B_SHAPES = [ + {"B": 1, "Hq": 32, "Hkv": 8, "Sq": 1, "Skv": 128, "D": 128, "label": "decode-128"}, + {"B": 1, "Hq": 32, "Hkv": 8, "Sq": 1, "Skv": 1024, "D": 128, "label": "decode-1k"}, + {"B": 1, "Hq": 32, "Hkv": 8, "Sq": 128, "Skv": 128, "D": 128, "label": "prefill-128"}, + {"B": 1, "Hq": 32, "Hkv": 8, "Sq": 512, "Skv": 512, "D": 128, "label": "prefill-512"}, + {"B": 1, "Hq": 32, "Hkv": 8, "Sq": 1024, "Skv": 1024, "D": 128, "label": "prefill-1k"}, + {"B": 4, "Hq": 32, "Hkv": 8, "Sq": 128, "Skv": 128, "D": 128, "label": "batch4-prefill-128"}, + {"B": 8, "Hq": 32, "Hkv": 8, "Sq": 64, "Skv": 64, "D": 128, "label": "batch8-prefill-64"}, +] + + +def main(): + parser = argparse.ArgumentParser(description="Benchmark deterministic attention") + parser.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16") + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--iters", type=int, default=20) + args = parser.parse_args() + + dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16 + + print(f"{'Shape':<25} {'Latency(ms)':>12} {'PeakMem(MB)':>12} {'Scores(MB)':>11}") + print("-" * 65) + + for shape in QWEN3_8B_SHAPES: + kwargs = shape.copy() + label = kwargs.pop("label") + try: + result = benchmark_attention( + **kwargs, + dtype=dtype, + warmup=args.warmup, + iters=args.iters, + ) + print( + f"{label:<25} {result['latency_ms']:>12.3f} " + f"{result['peak_memory_mb']:>12.1f} " + f"{result['scores_materialization_mb']:>11.1f}" + ) + except RuntimeError as exc: + print(f"{label:<25} {'OOM' if 'out of memory' in str(exc) else 'ERROR':>12}") + if "out of memory" in str(exc): + torch.cuda.empty_cache() + + +if __name__ == "__main__": + main() diff --git a/csrc/cuda/attention/deterministic_attention.cu b/csrc/cuda/attention/deterministic_attention.cu new file mode 100644 index 00000000..973b07a8 --- /dev/null +++ b/csrc/cuda/attention/deterministic_attention.cu @@ -0,0 +1,663 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Deterministic standard-softmax attention (issue #147). +// Forward: QK kernel → masked softmax+LSE kernel → PV kernel. +// All reductions use fixed order; no split-KV or dynamic dispatch. + +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace { + +constexpr int64_t kDeterministicAttentionHeadDim = 128; +constexpr int kSoftmaxThreads = 256; + +// --------------------------------------------------------------------------- +// QK Kernel: scores[b, hq, q, k] = scale * sum_{d=0}^{D-1} Q[b,hq,q,d]*K[b,kv_head,k,d] +// Grid: (Skv_blocks, Sq_blocks, B * Hq) +// Block: (TILE_K, TILE_Q) threads, each thread computes one score element. +// --------------------------------------------------------------------------- +constexpr int kQKTileQ = 16; +constexpr int kQKTileK = 16; + +template +__global__ void qk_kernel( + const scalar_t* __restrict__ Q, // [B, Hq, Sq, D] + const scalar_t* __restrict__ K, // [B, Hkv, Skv, D] + float* __restrict__ scores, // [B, Hq, Sq, Skv] + int64_t B, int64_t Hq, int64_t Hkv, + int64_t Sq, int64_t Skv, int64_t D, + float scale) { + + const int k_idx = blockIdx.x * kQKTileK + threadIdx.x; + const int q_idx = blockIdx.y * kQKTileQ + threadIdx.y; + const int bh = blockIdx.z; // flattened (b * Hq + hq) + const int b = bh / Hq; + const int hq = bh % Hq; + + if (q_idx >= Sq || k_idx >= Skv) return; + + const int kv_head = hq / (Hq / Hkv); + + const scalar_t* q_ptr = Q + ((int64_t)b * Hq * Sq * D + (int64_t)hq * Sq * D + (int64_t)q_idx * D); + const scalar_t* k_ptr = K + ((int64_t)b * Hkv * Skv * D + (int64_t)kv_head * Skv * D + (int64_t)k_idx * D); + + float acc = 0.0f; + #pragma unroll 8 + for (int64_t d = 0; d < D; ++d) { + acc += (float)q_ptr[d] * (float)k_ptr[d]; + } + + const int64_t out_idx = (int64_t)b * Hq * Sq * Skv + (int64_t)hq * Sq * Skv + (int64_t)q_idx * Skv + k_idx; + scores[out_idx] = scale * acc; +} + +// --------------------------------------------------------------------------- +// Masked Softmax + LSE Kernel +// One CTA per (b, hq, q) row. Fixed 256 threads. +// Applies causal + padding mask, computes max, sum-exp, writes P and LSE. +// --------------------------------------------------------------------------- +__global__ void masked_softmax_lse_kernel( + float* __restrict__ scores, // [B, Hq, Sq, Skv] in-place -> P + float* __restrict__ lse, // [B, Hq, Sq] + const bool* __restrict__ pad_mask, // [B, Skv] or nullptr + int64_t B, int64_t Hq, int64_t Sq, int64_t Skv, + bool causal) { + + const int row_idx = blockIdx.x; // flattened (b * Hq * Sq + hq * Sq + q) + const int b = row_idx / (Hq * Sq); + const int hq = (row_idx / Sq) % Hq; + const int q = row_idx % Sq; + + float* row = scores + (int64_t)row_idx * Skv; + + // Causal boundary: key_index <= Skv - Sq + q + const int64_t causal_limit = causal ? (Skv - Sq + q + 1) : Skv; + + // Phase 1: Apply masks and find max + __shared__ float smax[kSoftmaxThreads]; + float thread_max = -INFINITY; + for (int k = threadIdx.x; k < Skv; k += kSoftmaxThreads) { + bool valid = (k < causal_limit); + if (valid && pad_mask != nullptr) { + valid = pad_mask[(int64_t)b * Skv + k]; + } + if (!valid) { + row[k] = -INFINITY; + } + if (valid) { + thread_max = fmaxf(thread_max, row[k]); + } + } + + // Warp reduction for max + smax[threadIdx.x] = thread_max; + __syncthreads(); + // Tree reduction + for (int stride = kSoftmaxThreads / 2; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + smax[threadIdx.x] = fmaxf(smax[threadIdx.x], smax[threadIdx.x + stride]); + } + __syncthreads(); + } + float row_max = smax[0]; + + // Phase 2: Compute sum of exp(s - max) + __shared__ float ssum[kSoftmaxThreads]; + float thread_sum = 0.0f; + for (int k = threadIdx.x; k < Skv; k += kSoftmaxThreads) { + bool valid = (k < causal_limit); + if (valid && pad_mask != nullptr) { + valid = pad_mask[(int64_t)b * Skv + k]; + } + if (valid) { + float val = expf(row[k] - row_max); + row[k] = val; + thread_sum += val; + } else { + row[k] = 0.0f; + } + } + + ssum[threadIdx.x] = thread_sum; + __syncthreads(); + for (int stride = kSoftmaxThreads / 2; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + ssum[threadIdx.x] += ssum[threadIdx.x + stride]; + } + __syncthreads(); + } + float row_sum = ssum[0]; + + // Phase 3: Normalize to get P, compute LSE + float lse_val; + if (row_sum == 0.0f) { + // Fully masked row + lse_val = -INFINITY; + for (int k = threadIdx.x; k < Skv; k += kSoftmaxThreads) { + row[k] = 0.0f; + } + } else { + lse_val = row_max + logf(row_sum); + float inv_sum = 1.0f / row_sum; + for (int k = threadIdx.x; k < Skv; k += kSoftmaxThreads) { + row[k] *= inv_sum; + } + } + + if (threadIdx.x == 0) { + lse[row_idx] = lse_val; + } +} + +// --------------------------------------------------------------------------- +// PV Kernel: out[b, hq, q, d] = sum_{k=0}^{Skv-1} P[b,hq,q,k] * V[b,kv_head,k,d] +// Grid: (D_blocks, Sq_blocks, B * Hq) +// Each thread computes one output element with sequential k accumulation. +// --------------------------------------------------------------------------- +constexpr int kPVTileQ = 16; +constexpr int kPVTileD = 16; + +template +__global__ void pv_kernel( + const float* __restrict__ P, // [B, Hq, Sq, Skv] + const scalar_t* __restrict__ V, // [B, Hkv, Skv, D] + scalar_t* __restrict__ out, // [B, Hq, Sq, D] + int64_t B, int64_t Hq, int64_t Hkv, + int64_t Sq, int64_t Skv, int64_t D) { + + const int d_idx = blockIdx.x * kPVTileD + threadIdx.x; + const int q_idx = blockIdx.y * kPVTileQ + threadIdx.y; + const int bh = blockIdx.z; + const int b = bh / Hq; + const int hq = bh % Hq; + + if (q_idx >= Sq || d_idx >= D) return; + + const int kv_head = hq / (Hq / Hkv); + + const float* p_row = P + ((int64_t)b * Hq * Sq * Skv + (int64_t)hq * Sq * Skv + (int64_t)q_idx * Skv); + const scalar_t* v_base = V + ((int64_t)b * Hkv * Skv * D + (int64_t)kv_head * Skv * D); + + float acc = 0.0f; + for (int64_t k = 0; k < Skv; ++k) { + acc += p_row[k] * (float)v_base[k * D + d_idx]; + } + + const int64_t out_idx = (int64_t)b * Hq * Sq * D + (int64_t)hq * Sq * D + (int64_t)q_idx * D + d_idx; + out[out_idx] = (scalar_t)acc; +} + +void check_deterministic_attention_inputs( + const torch::Tensor& q, + const torch::Tensor& k, + const torch::Tensor& v, + const torch::optional& key_padding_mask) { + TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda(), + "deterministic_attention: q, k, v must be CUDA tensors"); + TORCH_CHECK(q.device() == k.device() && q.device() == v.device(), + "deterministic_attention: q, k, v must be on the same device"); + TORCH_CHECK(q.dim() == 4 && k.dim() == 4 && v.dim() == 4, + "deterministic_attention: q/k/v must be 4-D [B, H, S, D]"); + TORCH_CHECK( + q.scalar_type() == at::kHalf || q.scalar_type() == at::kBFloat16, + "deterministic_attention: only FP16 and BF16 are supported, got ", + q.scalar_type()); + TORCH_CHECK(k.scalar_type() == q.scalar_type() && v.scalar_type() == q.scalar_type(), + "deterministic_attention: q, k, v must share the same dtype"); + + const int64_t B = q.size(0); + const int64_t Hq = q.size(1); + const int64_t Sq = q.size(2); + const int64_t D = q.size(3); + const int64_t Hkv = k.size(1); + const int64_t Skv = k.size(2); + + TORCH_CHECK(D == kDeterministicAttentionHeadDim, + "deterministic_attention: head dim D must be ", + kDeterministicAttentionHeadDim, + ", got ", + D); + TORCH_CHECK(k.size(0) == B && v.size(0) == B, + "deterministic_attention: batch size mismatch between q/k/v"); + TORCH_CHECK(v.size(1) == Hkv && v.size(2) == Skv && k.size(3) == D && v.size(3) == D, + "deterministic_attention: k/v shape mismatch"); + TORCH_CHECK(Hq % Hkv == 0, + "deterministic_attention: Hq (", + Hq, + ") must be divisible by Hkv (", + Hkv, + ") for GQA"); + TORCH_CHECK(Sq >= 1 && Skv >= 1, + "deterministic_attention: Sq and Skv must be positive"); + + if (key_padding_mask.has_value() && key_padding_mask->defined()) { + const auto& mask = *key_padding_mask; + TORCH_CHECK(mask.is_cuda(), "deterministic_attention: key_padding_mask must be CUDA"); + TORCH_CHECK(mask.device() == q.device(), + "deterministic_attention: key_padding_mask must match q device"); + TORCH_CHECK(mask.scalar_type() == at::kBool, + "deterministic_attention: key_padding_mask must be bool"); + TORCH_CHECK(mask.dim() == 2 && mask.size(0) == B && mask.size(1) == Skv, + "deterministic_attention: key_padding_mask must be [B, Skv]"); + } +} + +} // namespace + +// Returns {out, lse, P}: +// out: [B, Hq, Sq, D] same dtype as q +// lse: [B, Hq, Sq] FP32 +// P: [B, Hq, Sq, Skv] FP32 (softmax probabilities, saved for backward) +std::vector deterministic_attention_forward( + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + bool causal, + double scale, + torch::optional key_padding_mask) { + check_deterministic_attention_inputs(q, k, v, key_padding_mask); + + const at::cuda::OptionalCUDAGuard device_guard(at::device_of(q)); + + auto q_contig = q.contiguous(); + auto k_contig = k.contiguous(); + auto v_contig = v.contiguous(); + torch::optional mask_contig; + if (key_padding_mask.has_value() && key_padding_mask->defined()) { + mask_contig = key_padding_mask->contiguous(); + } + + const int64_t B = q_contig.size(0); + const int64_t Hq = q_contig.size(1); + const int64_t Sq = q_contig.size(2); + const int64_t D = q_contig.size(3); + const int64_t Hkv = k_contig.size(1); + const int64_t Skv = k_contig.size(2); + + auto stream = at::cuda::getCurrentCUDAStream(); + + // Allocate scores [B, Hq, Sq, Skv] FP32 + auto scores = torch::empty({B, Hq, Sq, Skv}, q_contig.options().dtype(at::kFloat)); + + // --- Launch QK kernel --- + { + dim3 block(kQKTileK, kQKTileQ); + dim3 grid( + (Skv + kQKTileK - 1) / kQKTileK, + (Sq + kQKTileQ - 1) / kQKTileQ, + B * Hq); + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, + q_contig.scalar_type(), "qk_kernel", [&] { + qk_kernel<<>>( + q_contig.data_ptr(), + k_contig.data_ptr(), + scores.data_ptr(), + B, Hq, Hkv, Sq, Skv, D, + (float)scale); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + }); + } + + // --- Launch Masked Softmax + LSE kernel --- + auto lse = torch::empty({B, Hq, Sq}, q_contig.options().dtype(at::kFloat)); + { + const int64_t num_rows = B * Hq * Sq; + dim3 block(kSoftmaxThreads); + dim3 grid(num_rows); + const bool* pad_mask_ptr = nullptr; + if (mask_contig.has_value()) { + pad_mask_ptr = mask_contig->data_ptr(); + } + masked_softmax_lse_kernel<<>>( + scores.data_ptr(), + lse.data_ptr(), + pad_mask_ptr, + B, Hq, Sq, Skv, causal); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + + // --- Launch PV kernel --- + auto out = torch::empty_like(q_contig); + { + dim3 block(kPVTileD, kPVTileQ); + dim3 grid( + (D + kPVTileD - 1) / kPVTileD, + (Sq + kPVTileQ - 1) / kPVTileQ, + B * Hq); + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, + q_contig.scalar_type(), "pv_kernel", [&] { + pv_kernel<<>>( + scores.data_ptr(), + v_contig.data_ptr(), + out.data_ptr(), + B, Hq, Hkv, Sq, Skv, D); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + }); + } + + return {out, lse, scores}; +} + +// =========================================================================== +// BACKWARD +// =========================================================================== +namespace { + +// --------------------------------------------------------------------------- +// dP kernel: dP[b,hq,q,k] = sum_{d=0}^{D-1} dO[b,hq,q,d] * V[b,kv_head,k,d] +// Grid: (Skv_blocks, Sq_blocks, B*Hq) +// --------------------------------------------------------------------------- +template +__global__ void dp_kernel( + const scalar_t* __restrict__ dO, // [B, Hq, Sq, D] + const scalar_t* __restrict__ V, // [B, Hkv, Skv, D] + float* __restrict__ dP, // [B, Hq, Sq, Skv] + int64_t B, int64_t Hq, int64_t Hkv, + int64_t Sq, int64_t Skv, int64_t D) { + + const int k_idx = blockIdx.x * kQKTileK + threadIdx.x; + const int q_idx = blockIdx.y * kQKTileQ + threadIdx.y; + const int bh = blockIdx.z; + const int b = bh / Hq; + const int hq = bh % Hq; + + if (q_idx >= Sq || k_idx >= Skv) return; + + const int kv_head = hq / (Hq / Hkv); + + const scalar_t* do_ptr = dO + ((int64_t)b * Hq * Sq * D + (int64_t)hq * Sq * D + (int64_t)q_idx * D); + const scalar_t* v_ptr = V + ((int64_t)b * Hkv * Skv * D + (int64_t)kv_head * Skv * D + (int64_t)k_idx * D); + + float acc = 0.0f; + #pragma unroll 8 + for (int64_t d = 0; d < D; ++d) { + acc += (float)do_ptr[d] * (float)v_ptr[d]; + } + + const int64_t out_idx = (int64_t)b * Hq * Sq * Skv + (int64_t)hq * Sq * Skv + (int64_t)q_idx * Skv + k_idx; + dP[out_idx] = acc; +} + +// --------------------------------------------------------------------------- +// Softmax backward kernel: one CTA per (b, hq, q) row +// delta[row] = sum_k(dP[row,k] * P[row,k]) +// dS[row,k] = P[row,k] * (dP[row,k] - delta) +// Writes dS in-place over the dP buffer. +// --------------------------------------------------------------------------- +__global__ void softmax_backward_kernel( + float* __restrict__ dP_dS, // [B, Hq, Sq, Skv] - input dP, output dS + const float* __restrict__ P, // [B, Hq, Sq, Skv] + int64_t Skv) { + + const int row_idx = blockIdx.x; + float* ds_row = dP_dS + (int64_t)row_idx * Skv; + const float* p_row = P + (int64_t)row_idx * Skv; + + // Compute delta = sum_k(dP * P) + __shared__ float sdelta[kSoftmaxThreads]; + float thread_delta = 0.0f; + for (int k = threadIdx.x; k < Skv; k += kSoftmaxThreads) { + thread_delta += ds_row[k] * p_row[k]; + } + sdelta[threadIdx.x] = thread_delta; + __syncthreads(); + for (int stride = kSoftmaxThreads / 2; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + sdelta[threadIdx.x] += sdelta[threadIdx.x + stride]; + } + __syncthreads(); + } + float delta = sdelta[0]; + + // dS = P * (dP - delta) + for (int k = threadIdx.x; k < Skv; k += kSoftmaxThreads) { + ds_row[k] = p_row[k] * (ds_row[k] - delta); + } +} + +// --------------------------------------------------------------------------- +// dQ kernel: dQ[b,hq,q,d] = scale * sum_{k=0}^{Skv-1} dS[b,hq,q,k] * K[b,kv_head,k,d] +// Grid: (D_blocks, Sq_blocks, B*Hq) +// --------------------------------------------------------------------------- +template +__global__ void dq_kernel( + const float* __restrict__ dS, // [B, Hq, Sq, Skv] + const scalar_t* __restrict__ K, // [B, Hkv, Skv, D] + scalar_t* __restrict__ dQ, // [B, Hq, Sq, D] + int64_t B, int64_t Hq, int64_t Hkv, + int64_t Sq, int64_t Skv, int64_t D, + float scale) { + + const int d_idx = blockIdx.x * kPVTileD + threadIdx.x; + const int q_idx = blockIdx.y * kPVTileQ + threadIdx.y; + const int bh = blockIdx.z; + const int b = bh / Hq; + const int hq = bh % Hq; + + if (q_idx >= Sq || d_idx >= D) return; + + const int kv_head = hq / (Hq / Hkv); + + const float* ds_row = dS + ((int64_t)b * Hq * Sq * Skv + (int64_t)hq * Sq * Skv + (int64_t)q_idx * Skv); + const scalar_t* k_base = K + ((int64_t)b * Hkv * Skv * D + (int64_t)kv_head * Skv * D); + + float acc = 0.0f; + for (int64_t k = 0; k < Skv; ++k) { + acc += ds_row[k] * (float)k_base[k * D + d_idx]; + } + + const int64_t out_idx = (int64_t)b * Hq * Sq * D + (int64_t)hq * Sq * D + (int64_t)q_idx * D + d_idx; + dQ[out_idx] = (scalar_t)(scale * acc); +} + +// --------------------------------------------------------------------------- +// dK kernel: dK[b,hkv,k,d] = scale * sum_{local=0..g-1} sum_{q=0..Sq-1} dS[b,hq,q,k]*Q[b,hq,q,d] +// Grid: (D_blocks, Skv_blocks, B*Hkv) +// Each thread: single writer for one dK element (§4.1 fixed order). +// --------------------------------------------------------------------------- +template +__global__ void dk_kernel( + const float* __restrict__ dS, // [B, Hq, Sq, Skv] + const scalar_t* __restrict__ Q, // [B, Hq, Sq, D] + scalar_t* __restrict__ dK, // [B, Hkv, Skv, D] + int64_t B, int64_t Hq, int64_t Hkv, + int64_t Sq, int64_t Skv, int64_t D, + float scale) { + + const int d_idx = blockIdx.x * kPVTileD + threadIdx.x; + const int k_idx = blockIdx.y * kPVTileQ + threadIdx.y; + const int b_hkv = blockIdx.z; + const int b = b_hkv / Hkv; + const int hkv = b_hkv % Hkv; + + if (k_idx >= Skv || d_idx >= D) return; + + const int64_t g = Hq / Hkv; + + float acc = 0.0f; + for (int64_t local = 0; local < g; ++local) { + int64_t hq = hkv * g + local; + for (int64_t qi = 0; qi < Sq; ++qi) { + float ds_val = dS[(int64_t)b * Hq * Sq * Skv + hq * Sq * Skv + qi * Skv + k_idx]; + float q_val = (float)Q[(int64_t)b * Hq * Sq * D + hq * Sq * D + qi * D + d_idx]; + acc += ds_val * q_val; + } + } + + const int64_t out_idx = (int64_t)b * Hkv * Skv * D + (int64_t)hkv * Skv * D + (int64_t)k_idx * D + d_idx; + dK[out_idx] = (scalar_t)(scale * acc); +} + +// --------------------------------------------------------------------------- +// dV kernel: dV[b,hkv,k,d] = sum_{local=0..g-1} sum_{q=0..Sq-1} P[b,hq,q,k]*dO[b,hq,q,d] +// Grid: (D_blocks, Skv_blocks, B*Hkv) +// Each thread: single writer for one dV element (§4.1 fixed order). +// --------------------------------------------------------------------------- +template +__global__ void dv_kernel( + const float* __restrict__ P, // [B, Hq, Sq, Skv] + const scalar_t* __restrict__ dO, // [B, Hq, Sq, D] + scalar_t* __restrict__ dV, // [B, Hkv, Skv, D] + int64_t B, int64_t Hq, int64_t Hkv, + int64_t Sq, int64_t Skv, int64_t D) { + + const int d_idx = blockIdx.x * kPVTileD + threadIdx.x; + const int k_idx = blockIdx.y * kPVTileQ + threadIdx.y; + const int b_hkv = blockIdx.z; + const int b = b_hkv / Hkv; + const int hkv = b_hkv % Hkv; + + if (k_idx >= Skv || d_idx >= D) return; + + const int64_t g = Hq / Hkv; + + float acc = 0.0f; + for (int64_t local = 0; local < g; ++local) { + int64_t hq = hkv * g + local; + for (int64_t qi = 0; qi < Sq; ++qi) { + float p_val = P[(int64_t)b * Hq * Sq * Skv + hq * Sq * Skv + qi * Skv + k_idx]; + float do_val = (float)dO[(int64_t)b * Hq * Sq * D + hq * Sq * D + qi * D + d_idx]; + acc += p_val * do_val; + } + } + + const int64_t out_idx = (int64_t)b * Hkv * Skv * D + (int64_t)hkv * Skv * D + (int64_t)k_idx * D + d_idx; + dV[out_idx] = (scalar_t)acc; +} + +} // namespace (backward kernels) + +// Returns {dQ, dK, dV} +std::vector deterministic_attention_backward( + torch::Tensor grad_output, + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor P, // saved from forward [B, Hq, Sq, Skv] FP32 + bool causal, + double scale, + torch::optional key_padding_mask) { + + const at::cuda::OptionalCUDAGuard device_guard(at::device_of(q)); + + auto dO = grad_output.contiguous(); + auto q_c = q.contiguous(); + auto k_c = k.contiguous(); + auto v_c = v.contiguous(); + auto P_c = P.contiguous(); + + const int64_t B = q_c.size(0); + const int64_t Hq = q_c.size(1); + const int64_t Sq = q_c.size(2); + const int64_t D = q_c.size(3); + const int64_t Hkv = k_c.size(1); + const int64_t Skv = k_c.size(2); + + auto stream = at::cuda::getCurrentCUDAStream(); + + // dP = dO @ V^T [B, Hq, Sq, Skv] + auto dP = torch::empty({B, Hq, Sq, Skv}, q_c.options().dtype(at::kFloat)); + { + dim3 block(kQKTileK, kQKTileQ); + dim3 grid( + (Skv + kQKTileK - 1) / kQKTileK, + (Sq + kQKTileQ - 1) / kQKTileQ, + B * Hq); + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, + q_c.scalar_type(), "dp_kernel", [&] { + dp_kernel<<>>( + dO.data_ptr(), + v_c.data_ptr(), + dP.data_ptr(), + B, Hq, Hkv, Sq, Skv, D); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + }); + } + + // Softmax backward: dS = P * (dP - delta), writes in-place over dP + { + const int64_t num_rows = B * Hq * Sq; + softmax_backward_kernel<<>>( + dP.data_ptr(), + P_c.data_ptr(), + Skv); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + // dP buffer now contains dS + + // dQ = scale * dS @ K + auto dQ = torch::empty_like(q_c); + { + dim3 block(kPVTileD, kPVTileQ); + dim3 grid( + (D + kPVTileD - 1) / kPVTileD, + (Sq + kPVTileQ - 1) / kPVTileQ, + B * Hq); + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, + q_c.scalar_type(), "dq_kernel", [&] { + dq_kernel<<>>( + dP.data_ptr(), + k_c.data_ptr(), + dQ.data_ptr(), + B, Hq, Hkv, Sq, Skv, D, + (float)scale); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + }); + } + + // dK = scale * dS^T @ Q (per kv_head, accumulate over query heads in group) + auto dK = torch::empty_like(k_c); + { + dim3 block(kPVTileD, kPVTileQ); + dim3 grid( + (D + kPVTileD - 1) / kPVTileD, + (Skv + kPVTileQ - 1) / kPVTileQ, + B * Hkv); + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, + q_c.scalar_type(), "dk_kernel", [&] { + dk_kernel<<>>( + dP.data_ptr(), + q_c.data_ptr(), + dK.data_ptr(), + B, Hq, Hkv, Sq, Skv, D, + (float)scale); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + }); + } + + // dV = P^T @ dO (per kv_head, accumulate over query heads in group) + auto dV = torch::empty_like(v_c); + { + dim3 block(kPVTileD, kPVTileQ); + dim3 grid( + (D + kPVTileD - 1) / kPVTileD, + (Skv + kPVTileQ - 1) / kPVTileQ, + B * Hkv); + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, + q_c.scalar_type(), "dv_kernel", [&] { + dv_kernel<<>>( + P_c.data_ptr(), + dO.data_ptr(), + dV.data_ptr(), + B, Hq, Hkv, Sq, Skv, D); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + }); + } + + return {dQ, dK, dV}; +} diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 61ba4a3b..715d200e 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -75,6 +75,25 @@ torch::Tensor deterministic_logp_forward_fp32(torch::Tensor logits, torch::Tenso torch::Tensor deterministic_logp_forward_indexed_out(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor row_indices, torch::Tensor output); torch::Tensor deterministic_logp_forward_indexed_fp32(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor row_indices); +// Deterministic standard-softmax attention (issue #147) +std::vector deterministic_attention_forward( + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + bool causal, + double scale, + torch::optional key_padding_mask); + +std::vector deterministic_attention_backward( + torch::Tensor grad_output, + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor P, + bool causal, + double scale, + torch::optional key_padding_mask); + // Prefix-Shared Attention Declarations & Wrappers void prefix_shared_attention_forward( @@ -169,5 +188,15 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { // registry Prefix-Shared Attention m.def("prefix_shared_attention", &prefix_shared_attention, "Prefix-Shared Fused Attention for GRPO"); + + // Deterministic standard-softmax attention (issue #147) + m.def( + "deterministic_attention_forward", + &deterministic_attention_forward, + "Deterministic standard softmax attention forward (out, lse)"); + m.def( + "deterministic_attention_backward", + &deterministic_attention_backward, + "Deterministic standard softmax attention backward (dQ, dK, dV)"); #endif } diff --git a/docs/operators/attention.md b/docs/operators/attention.md index e3cb5f9b..ebff9a58 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -48,7 +48,7 @@ The op exposes the WS1 dual-path contract: | Backend | Wrapper | Native symbol | Status | | --- | --- | --- | --- | | PyTorch fallback | `NativeAttentionOp` | None | fp32 ground-truth reference; CPU and any GPU. | -| CUDA / ROCm / Triton | — | — | Planned: downstream fused attention kernels validate against this reference. | +| CUDA deterministic | `DeterministicAttentionOp` | `_C.deterministic_attention_forward/backward` | Batch-invariant CUDA implementation (issue #147). | ## Tensor Contract @@ -76,12 +76,14 @@ the inputs' device. ## Dispatch Behavior `kernel_registry.get_op("attention")` resolves through the `OpBackend` priority map. On -`cuda` / `rocm` / `cpu` the only registered backend today is the PyTorch native op -(`PYTORCH_NATIVE_ATTENTION`), so every device dispatches to this op. Calling it (`__call__` -> -`forward(...)`) computes in the input dtype; `forward_fp32(...)` is the explicit fp32 golden -path. When fused attention kernels land, they are prepended to the priority list and the native -op becomes the fallback. The production `"attn"` op_type (SDPA-based `PYTORCH_ATTN`, FlashAttention, etc.) is -a separate dispatch chain and is unaffected. +`cuda` the priority is: + +1. `CUDA_DETERMINISTIC_ATTENTION` — `DeterministicAttentionOp` (batch-invariant, fixed-order). +2. `PYTORCH_NATIVE_ATTENTION` — `NativeAttentionOp` (fallback). + +Calling it (`__call__` -> `forward(...)`) computes in the input dtype; `forward_fp32(...)` is +the explicit fp32 golden path (NativeAttentionOp only). The production `"attn"` op_type +(SDPA-based `PYTORCH_ATTN`, FlashAttention, etc.) is a separate dispatch chain and is unaffected. ## Accuracy @@ -146,14 +148,76 @@ GPU-only LARGE Qwen3-8B real-shape smoke test. ## Implementation Files -- `rl_engine/kernels/ops/pytorch/attention/standard_attn.py` +- `rl_engine/kernels/ops/pytorch/attention/standard_attn.py` — ground-truth reference +- `rl_engine/kernels/ops/cuda/attention/deterministic_attn.py` — CUDA deterministic op +- `csrc/cuda/attention/deterministic_attention.cu` — CUDA kernels - `rl_engine/kernels/registry.py` - `tests/test_attention.py` +- `tests/test_deterministic_attention_cuda.py` + +## Fixed Reduction Order (CUDA Deterministic Backend) + +The CUDA `DeterministicAttentionOp` pins reduction order for batch-invariance: + +**QK kernel**: D-dimension FP32 accumulation, `d = 0 .. D-1`. Each `scores[b,hq,q,k]` has +exactly one writer thread. + +**Softmax + LSE kernel**: Each `(b, hq, q)` row processed by one CTA with fixed 256 threads. +Max and sum-exp use a fixed shared-memory tree reduction (power-of-two stride). No split by +batch size, sequence length, or SM count. + +**PV kernel**: K-dimension FP32 accumulation, `k = 0 .. Skv-1`. Each `out[b,hq,q,d]` has +exactly one writer thread. + +**Backward dK/dV**: Per `(b, hkv, k, d)` output element, a single thread accumulates over +query heads in group order then query positions: +```text +for local = 0 .. g-1: # g = Hq / Hkv + hq = hkv * g + local + for q = 0 .. Sq-1: + acc += ... +``` +No cross-CTA atomics. No launch-order dependent accumulation. + +## Prefill / Decode / KV-cache Shared Contract + +All inference modes use **the same standard attention kernels** (only Sq/Skv differ): + +- **Prefill**: `Sq == Skv`. Causal mask `key_index <= Skv - Sq + query_index`. +- **Chunked-prefill**: each chunk uses `Sq = chunk_size`, `Skv = past + chunk_size`. + Same causal offset formula produces identical results to full prefill at matching positions. +- **Decode**: `Sq = 1` (or few), `Skv = full_context`. Same kernel, same offset. +- **KV-cache**: caller does `k_full = cat([k_cache, k_new], dim=2)` then calls this op. + No separate KV-cache softmax implementation allowed. + +Hooks: +- `forward(q, k, v, ...)` — main path (registry, #108 harness). Differentiable. +- `forward_with_lse(q, k, v, ...)` — returns `(out, lse)` for LSE verification, debugging, + and future KV-cache / training integration. + +## Tolerance + +| Scenario | Comparison | Tolerance | +| --- | --- | --- | +| Same physical shape, varying batch position/size/chunk | bitwise | `batch_invariance` (atol=0, rtol=0) | +| Chunked-prefill on/off at same position | bitwise | `batch_invariance` | +| Prefill tail vs decode slice | bitwise | `batch_invariance` | +| CUDA vs `forward_fp32` output/grad | tolerance | `accuracy.default.attention` | +| Valid-only vs padded (reduction width differs) | near-equal | accuracy tolerance (NOT bitwise) | + +## Memory Tradeoff (First Version) + +The first version materializes full FP32 `scores [B, Hq, Sq, Skv]` and `P [B, Hq, Sq, Skv]`. +Memory cost: `4 * B * Hq * Sq * Skv` bytes per tensor. For Qwen3-8B at B=8, Sq=Skv=4096, +Hq=32: each tensor is ~17 GB. This is acceptable for correctness verification and moderate +sequence lengths but OOM-prone for long sequences. See `benchmarks/benchmark_deterministic_attention.py` +for measured peak memory at representative shapes. ## Known Limitations -- PyTorch fallback only; no fused CUDA/Triton backend yet (downstream work). +- First version: `D=128` only (Qwen3-8B alignment). +- Supported dtypes: BF16, FP16. +- Full materialization of scores/P limits practical sequence length. - `Hq` must be divisible by `Hkv` (raises `ValueError` otherwise). -- The naive path materializes the full `[B, Hq, Sq, Skv]` scores tensor — no query-chunking, - so the LARGE load point is memory-heavy and GPU-only. -- Covers softmax attention only; QK-Norm and RoPE are applied before the call. +- CUDA KV-cache op wrapper is not in scope (caller does cat + calls this op). +- No FP8, no multi-GPU / sequence-parallel. diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 6ab592c6..801d9d36 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -130,3 +130,24 @@ def deterministic_logp_forward_indexed_fp32( token_ids: torch.Tensor, row_indices: torch.Tensor, ) -> torch.Tensor: ... +def deterministic_attention_forward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + causal: bool, + scale: float, + key_padding_mask: torch.Tensor | None, +) -> list[torch.Tensor]: + """Returns [out, lse, P].""" + ... + +def deterministic_attention_backward( + grad_output: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + P: torch.Tensor, + causal: bool, + scale: float, + key_padding_mask: torch.Tensor | None, +) -> list[torch.Tensor]: ... diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index c2ba34df..7b87ea83 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -95,19 +95,34 @@ def _make_attention_inputs( args: argparse.Namespace, dtype: torch.dtype, device: torch.device ) -> dict[str, Any]: batch, seq = _batch_seq(args) - return { - "q": _floating_tensor( - (batch, DEFAULT_N_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 0 - ), - "k": _floating_tensor( - (batch, DEFAULT_N_KV_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 1 - ), - "v": _floating_tensor( - (batch, DEFAULT_N_KV_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 2 - ), - "causal": True, + skv = _arg_int(args, "skv", seq) + n_heads = _arg_int(args, "n_heads", DEFAULT_N_HEADS) + n_kv_heads = _arg_int(args, "n_kv_heads", DEFAULT_N_KV_HEADS) + causal = bool(_arg_int(args, "causal", 1)) + use_padding = bool(_arg_int(args, "use_padding", 0)) + scale_mode = _arg_str(args, "scale_mode", "default") + + inputs: dict[str, Any] = { + "q": _floating_tensor((batch, n_heads, seq, DEFAULT_HEAD_DIM), args, dtype, device, 0), + "k": _floating_tensor((batch, n_kv_heads, skv, DEFAULT_HEAD_DIM), args, dtype, device, 1), + "v": _floating_tensor((batch, n_kv_heads, skv, DEFAULT_HEAD_DIM), args, dtype, device, 2), + "causal": causal, } + if scale_mode == "zero": + inputs["scale"] = 0.0 + elif scale_mode == "custom": + inputs["scale"] = 0.05 + # else: scale_mode == "default" -> no scale kwarg (uses 1/sqrt(D)) + + if use_padding: + generator = _generator(args, device, offset=42) + key_padding_mask = torch.rand((batch, skv), generator=generator, device=device) > 0.3 + key_padding_mask[:, 0] = True + inputs["key_padding_mask"] = key_padding_mask + + return inputs + def _make_logp_inputs( args: argparse.Namespace, dtype: torch.dtype, device: torch.device diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index 4051c276..9ad7bea9 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -57,6 +57,20 @@ def _load_object(path: str) -> Any: }, grad_input_names=("hidden", "lm_head_weight"), ), + "attention": OperatorSpec( + name="attention", + op_class="attention", + gold_path="rl_engine.kernels.ops.pytorch.attention.standard_attn.NativeAttentionOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.kernels.ops.pytorch.attention.standard_attn.NativeAttentionOp", + "cuda": ( + "rl_engine.kernels.ops.cuda.attention.deterministic_attn." + "DeterministicAttentionOp" + ), + }, + grad_input_names=("q", "k", "v"), + ), "batch_invariant_logp": OperatorSpec( name="batch_invariant_logp", op_class="logprob", diff --git a/rl_engine/kernels/gtest/tolerance_contract.json b/rl_engine/kernels/gtest/tolerance_contract.json index 7a2ce2b4..975ae450 100644 --- a/rl_engine/kernels/gtest/tolerance_contract.json +++ b/rl_engine/kernels/gtest/tolerance_contract.json @@ -16,6 +16,11 @@ "float32": {"atol": 1.0e-5, "rtol": 0.0}, "bfloat16": {"atol": 5.0e-2, "rtol": 0.0}, "float16": {"atol": 5.0e-3, "rtol": 0.0} + }, + "attention": { + "float32": {"atol": 1.0e-4, "rtol": 1.0e-4}, + "bfloat16": {"atol": 5.0e-2, "rtol": 2.0e-2}, + "float16": {"atol": 1.0e-3, "rtol": 1.0e-3} } }, "arch_overrides": { diff --git a/rl_engine/kernels/ops/cuda/attention/__init__.py b/rl_engine/kernels/ops/cuda/attention/__init__.py index 8d6addd9..09775c8e 100644 --- a/rl_engine/kernels/ops/cuda/attention/__init__.py +++ b/rl_engine/kernels/ops/cuda/attention/__init__.py @@ -1,9 +1,11 @@ # File: rl_engine/kernels/ops/cuda/attention/__init__.py +from .deterministic_attn import DeterministicAttentionOp from .flash_attn import FlashAttentionOp from .prefix_shared_attn import PrefixSharedAttentionOp __all__ = [ + "DeterministicAttentionOp", "FlashAttentionOp", "PrefixSharedAttentionOp", ] diff --git a/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py b/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py new file mode 100644 index 00000000..81f80a7f --- /dev/null +++ b/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA deterministic standard-softmax attention (issue #147). + +Forward: QK → masked softmax+LSE → PV (all FP32 intermediate). +Backward: dP → softmax_bwd → dQ/dK/dV with §4.1 fixed GQA order. +Wrapped in autograd.Function so #108 harness can .backward() through it. +""" + +from __future__ import annotations + +import math +from typing import Optional + +import torch +from torch.autograd import Function +from torch.autograd.function import once_differentiable + +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.utils.logger import logger + +_HEAD_DIM = 128 + + +class _DeterministicAttentionFn(Function): + @staticmethod + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + causal: bool, + scale: float, + key_padding_mask: Optional[torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor]: + q_c = q.contiguous() + k_c = k.contiguous() + v_c = v.contiguous() + mask_c = key_padding_mask.contiguous() if key_padding_mask is not None else None + + results = _C.deterministic_attention_forward(q_c, k_c, v_c, causal, float(scale), mask_c) + out, lse, P = results[0], results[1], results[2] + + ctx.save_for_backward(q_c, k_c, v_c, P, mask_c) + ctx.causal = causal + ctx.scale = scale + ctx.mark_non_differentiable(lse) + + return out, lse + + @staticmethod + @once_differentiable + def backward(ctx, grad_out: torch.Tensor, grad_lse: torch.Tensor): + q_c, k_c, v_c, P, mask_c = ctx.saved_tensors + + dQ, dK, dV = _C.deterministic_attention_backward( + grad_out.contiguous(), + q_c, + k_c, + v_c, + P, + ctx.causal, + float(ctx.scale), + mask_c, + ) + + return dQ, dK, dV, None, None, None + + +class DeterministicAttentionOp: + """Batch-invariant standard softmax attention on CUDA. + + Materializes full FP32 scores/P. Public surface matches NativeAttentionOp + so #108 harness can call forward(**inputs) with key_padding_mask. + """ + + def __init__(self) -> None: + if not _EXT_AVAILABLE or not hasattr(_C, "deterministic_attention_forward"): + raise RuntimeError( + "Deterministic CUDA attention kernel is unavailable. " + "Rebuild the extension with `pip install -e .` on a CUDA build." + ) + if not hasattr(_C, "deterministic_attention_backward"): + raise RuntimeError( + "Deterministic CUDA attention backward kernel is unavailable. " + "Rebuild the extension with `pip install -e .` on a CUDA build." + ) + logger.info("Successfully linked to _C.deterministic_attention_forward/backward.") + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return self.forward(q, k, v, causal=causal, scale=scale, key_padding_mask=key_padding_mask) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Harness / registry main path: return out only. Differentiable.""" + out, _lse = self.forward_with_lse( + q, k, v, causal=causal, scale=scale, key_padding_mask=key_padding_mask + ) + return out + + def forward_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return (out, lse) with FP32 LSE for debug / handoff hooks.""" + self._validate_inputs(q, k, v, key_padding_mask) + resolved_scale = scale if scale is not None else (1.0 / math.sqrt(q.shape[-1])) + out, lse = _DeterministicAttentionFn.apply( + q, k, v, causal, resolved_scale, key_padding_mask + ) + return out, lse + + @staticmethod + def _validate_inputs( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + key_padding_mask: Optional[torch.Tensor], + ) -> None: + if q.dim() != 4 or k.dim() != 4 or v.dim() != 4: + raise ValueError( + f"q/k/v must be 4-D [B, H, S, D], got q={tuple(q.shape)}, " + f"k={tuple(k.shape)}, v={tuple(v.shape)}" + ) + b, hq, sq, d = q.shape + hkv, skv = k.shape[1], k.shape[2] + if k.shape[0] != b or v.shape[0] != b: + raise ValueError("batch size mismatch between q/k/v") + if v.shape[1] != hkv or v.shape[2] != skv or k.shape[3] != d or v.shape[3] != d: + raise ValueError( + f"k/v shape mismatch: k={tuple(k.shape)}, v={tuple(v.shape)}, " + f"expected k/v [B={b}, Hkv, Skv, D={d}]" + ) + if d != _HEAD_DIM: + raise ValueError(f"head dim D must be {_HEAD_DIM}, got {d}") + if hq % hkv != 0: + raise ValueError(f"Hq={hq} not divisible by Hkv={hkv} (GQA group)") + if q.dtype not in (torch.float16, torch.bfloat16): + raise ValueError(f"only FP16/BF16 supported, got {q.dtype}") + if k.dtype != q.dtype or v.dtype != q.dtype: + raise ValueError("q, k, v must share the same dtype") + if not (q.is_cuda and k.is_cuda and v.is_cuda): + raise ValueError("q, k, v must be CUDA tensors") + if key_padding_mask is not None: + if key_padding_mask.dtype != torch.bool: + raise ValueError("key_padding_mask must be bool") + if key_padding_mask.shape != (b, skv): + raise ValueError( + f"key_padding_mask must be [B, Skv]=[{b}, {skv}], " + f"got {tuple(key_padding_mask.shape)}" + ) + if sq < 1 or skv < 1: + raise ValueError(f"Sq and Skv must be positive, got Sq={sq}, Skv={skv}") diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 2b9b0c30..b5115b94 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -30,6 +30,10 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): CUDA_FUSED_LOGP_SM90 = "rl_engine.kernels.ops.cuda.loss.logp.FusedLogpSM90Op" CUDA_FUSED_LOGP_GENERIC = "rl_engine.kernels.ops.cuda.loss.logp.FusedLogpGenericOp" CUDA_DETERMINISTIC_LOGP = "rl_engine.kernels.ops.cuda.loss.logp.DeterministicLogpCUDAOp" + # Deterministic standard-softmax attention (issue #147); not FlashAttention. + CUDA_DETERMINISTIC_ATTENTION = ( + "rl_engine.kernels.ops.cuda.attention.deterministic_attn.DeterministicAttentionOp" + ) # AMD ROCm optimized stack ROCM_AITER = "rl_engine.kernels.ops.rocm.aiter.AiterOp" @@ -171,7 +175,10 @@ def __init__(self): OpBackend.PYTORCH_NATIVE, ], "attn": [OpBackend.FLASH_ATTN, OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_ATTN], - "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION], + "attention": [ + OpBackend.CUDA_DETERMINISTIC_ATTENTION, + OpBackend.PYTORCH_NATIVE_ATTENTION, + ], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], "linear_logp": [OpBackend.TRITON_LINEAR_LOGP, OpBackend.PYTORCH_LINEAR_LOGP], diff --git a/setup.py b/setup.py index 6f94e040..725d20ab 100644 --- a/setup.py +++ b/setup.py @@ -79,6 +79,7 @@ def get_extensions(): "csrc/fused_logp_kernel.cu", "csrc/deterministic_logp_kernel.cu", "csrc/cuda/attention/prefix_shared_attention.cu", + "csrc/cuda/attention/deterministic_attention.cu", ] cc_major, cc_minor = torch.cuda.get_device_capability() diff --git a/tests/test_attention.py b/tests/test_attention.py index ffb5558d..469c6d30 100644 --- a/tests/test_attention.py +++ b/tests/test_attention.py @@ -434,8 +434,13 @@ def test_gradient_matches_reference(): def test_registry_dispatches_native_attention_op(): - """The registry resolves "attention" to the ground-truth NativeAttentionOp.""" - assert isinstance(kernel_registry.get_op("attention"), NativeAttentionOp) + """Resolve attention to the deterministic CUDA op or native fallback.""" + op = kernel_registry.get_op("attention") + # On CUDA with the extension built, the registry prefers DeterministicAttentionOp. + # On CPU or without the CUDA extension, it falls back to NativeAttentionOp. + from rl_engine.kernels.ops.cuda.attention.deterministic_attn import DeterministicAttentionOp + + assert isinstance(op, (NativeAttentionOp, DeterministicAttentionOp)) # --------------------------------------------------------------------------- # diff --git a/tests/test_deterministic_attention_cuda.py b/tests/test_deterministic_attention_cuda.py new file mode 100644 index 00000000..f08315a8 --- /dev/null +++ b/tests/test_deterministic_attention_cuda.py @@ -0,0 +1,657 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Deterministic standard-softmax attention CUDA tests (issue #147). + +Covers (per §7 and §8 of the implementation plan): +- Forward correctness via #108 harness (run_operator_suite) +- Backward correctness via #108 harness (check_grad=True, grad_mode="random") +- LSE correctness +- Batch invariance — Axis-A bitwise (slice, position permutation, batch-chunk) +- Sequence-dim chunked-prefill invariance (§7.4) +- Prefill/decode handoff (§7.5) +- KV-cache cat handoff (§7.5.2) +- Scale: None / 0.0 / custom +- GQA dK/dV order validation +- FP64 high-precision gradient comparison +- Gradient batch invariance +- Valid-only vs padded accuracy (not bitwise) +""" + +import math + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + +try: + from rl_engine.kernels.gtest.op_checks import CandidateSpec, OperatorCase, run_operator_suite + from rl_engine.kernels.ops.cuda.attention.deterministic_attn import DeterministicAttentionOp + from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp + + _OP_AVAILABLE = True +except (ImportError, RuntimeError): + _OP_AVAILABLE = False + +pytestmark = [ + pytestmark, + pytest.mark.skipif(not _OP_AVAILABLE, reason="CUDA attention op not built"), +] + +DEVICE = "cuda" +D = 128 + + +@pytest.fixture +def cuda_op(): + return DeterministicAttentionOp() + + +@pytest.fixture +def gold_op(): + return NativeAttentionOp() + + +def _tol(dtype): + if dtype == torch.bfloat16: + return 5e-2, 2e-2 + return 1e-3, 1e-3 + + +# ============================================================================= +# §8.4 — #108 Harness integration: forward + backward via run_operator_suite +# ============================================================================= + + +def _make_case(name, dtype, hq, hkv, sq, skv, causal, scale=None, padding=False): + torch.manual_seed(42) + B = 2 + inputs = { + "q": torch.randn(B, hq, sq, D, device=DEVICE, dtype=dtype), + "k": torch.randn(B, hkv, skv, D, device=DEVICE, dtype=dtype), + "v": torch.randn(B, hkv, skv, D, device=DEVICE, dtype=dtype), + "causal": causal, + } + if scale is not None: + inputs["scale"] = scale + if padding: + mask = torch.ones(B, skv, device=DEVICE, dtype=torch.bool) + mask[0, skv // 2 :] = False + mask[1, skv * 3 // 4 :] = False + inputs["key_padding_mask"] = mask + gold = NativeAttentionOp() + return OperatorCase( + name=name, + op_class="attention", + dtype=dtype, + inputs=inputs, + gold_fn=gold.forward_fp32, + grad_input_names=("q", "k", "v"), + ) + + +def _build_harness_cases(): + return [ + _make_case("bf16-gqa4x1-16x32-causal", torch.bfloat16, 4, 1, 16, 32, True), + _make_case("bf16-gqa32x8-16x32-causal", torch.bfloat16, 32, 8, 16, 32, True), + _make_case("bf16-gqa1x1-3x31-nocausal", torch.bfloat16, 1, 1, 3, 31, False), + _make_case("fp16-gqa4x2-17x33-causal", torch.float16, 4, 2, 17, 33, True), + _make_case("fp16-gqa32x8-1x64-decode", torch.float16, 32, 8, 1, 64, True), + _make_case("bf16-gqa4x1-16x16-nocausal", torch.bfloat16, 4, 1, 16, 16, False), + _make_case("bf16-gqa32x8-64x65-causal", torch.bfloat16, 32, 8, 64, 65, True), + _make_case("fp16-gqa4x1-65x127-causal", torch.float16, 4, 1, 65, 127, True), + _make_case("bf16-pad-gqa4x1-8x16", torch.bfloat16, 4, 1, 8, 16, True, padding=True), + _make_case("bf16-scale0-4x1-8x16", torch.bfloat16, 4, 1, 8, 16, True, scale=0.0), + _make_case("fp16-scale-custom-4x1-8x16", torch.float16, 4, 1, 8, 16, True, scale=0.05), + ] + + +def test_harness_forward(): + """§8.4: run_operator_suite forward — candidate vs gold (accuracy tolerance).""" + cuda = DeterministicAttentionOp() + candidate = CandidateSpec(name="cuda-attention", fn=cuda, backend="cuda") + report = run_operator_suite("attention", candidates=[candidate], cases=_build_harness_cases()) + for cr in report.candidates: + for case in cr.cases: + if not case.passed: + msgs = [ + o.message or f"idx={o.output_index} max_abs={o.max_abs_error}" + for o in case.outputs + if not o.passed + ] + pytest.fail(f"Forward failed {case.case_name}: {msgs}") + assert report.passed + + +def test_harness_backward(): + """§8.4: run_operator_suite backward — grad comparison vs gold fp32 autograd.""" + cuda = DeterministicAttentionOp() + candidate = CandidateSpec(name="cuda-attention", fn=cuda, backend="cuda") + report = run_operator_suite( + "attention", + candidates=[candidate], + cases=_build_harness_cases(), + check_grad=True, + grad_mode="random", + ) + for cr in report.candidates: + for case in cr.cases: + if not case.passed: + msgs = [ + o.message or f"idx={o.output_index} max_abs={o.max_abs_error}" + for o in case.outputs + if not o.passed + ] + pytest.fail(f"Backward failed {case.case_name}: {msgs}") + assert report.passed + + +# ============================================================================= +# §7.1 — Full correctness sweep (covers Sq=65, Skv=65/127, non-causal sq>skv) +# ============================================================================= + +SWEEP_CONFIGS = [] +for dtype in [torch.bfloat16, torch.float16]: + for hq, hkv in [(1, 1), (4, 1), (4, 2), (32, 8)]: + for sq in [1, 3, 16, 17, 64, 65]: + for skv in [1, 31, 32, 33, 64, 65, 127]: + for causal in [True, False]: + if causal and sq > skv: + continue + SWEEP_CONFIGS.append((dtype, hq, hkv, sq, skv, causal)) + + +@pytest.mark.parametrize("dtype,hq,hkv,sq,skv,causal", SWEEP_CONFIGS) +def test_forward_sweep(cuda_op, gold_op, dtype, hq, hkv, sq, skv, causal): + B = 2 + torch.manual_seed(42) + q = torch.randn(B, hq, sq, D, device=DEVICE, dtype=dtype) + k = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=dtype) + v = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=dtype) + + out_cuda = cuda_op.forward(q, k, v, causal=causal) + out_gold = gold_op.forward_fp32(q, k, v, causal=causal) + + atol, rtol = _tol(dtype) + torch.testing.assert_close(out_cuda.float(), out_gold.float(), atol=atol, rtol=rtol) + + +# ============================================================================= +# §7.1 — Scale tests (None, 0.0, custom) +# ============================================================================= + + +@pytest.mark.parametrize("scale", [None, 0.0, 0.05]) +def test_scale(cuda_op, gold_op, scale): + B, hq, hkv, sq, skv = 2, 4, 1, 8, 16 + torch.manual_seed(42) + q = torch.randn(B, hq, sq, D, device=DEVICE, dtype=torch.bfloat16) + k = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) + v = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) + + out_cuda = cuda_op.forward(q, k, v, causal=True, scale=scale) + out_gold = gold_op.forward_fp32(q, k, v, causal=True, scale=scale) + + atol, rtol = _tol(torch.bfloat16) + torch.testing.assert_close(out_cuda.float(), out_gold.float(), atol=atol, rtol=rtol) + + +# ============================================================================= +# Padding and fully-masked row +# ============================================================================= + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_forward_with_padding(cuda_op, gold_op, dtype): + B, hq, hkv, sq, skv = 2, 4, 1, 8, 16 + torch.manual_seed(7) + q = torch.randn(B, hq, sq, D, device=DEVICE, dtype=dtype) + k = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=dtype) + v = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=dtype) + mask = torch.ones(B, skv, device=DEVICE, dtype=torch.bool) + mask[0, 10:] = False + mask[1, 12:] = False + + out_cuda = cuda_op.forward(q, k, v, causal=True, key_padding_mask=mask) + out_gold = gold_op.forward_fp32(q, k, v, causal=True, key_padding_mask=mask) + + atol, rtol = _tol(dtype) + torch.testing.assert_close(out_cuda.float(), out_gold.float(), atol=atol, rtol=rtol) + + +def test_fully_masked_row(cuda_op): + B, hq, hkv, sq, skv = 1, 1, 1, 2, 4 + q = torch.randn(B, hq, sq, D, device=DEVICE, dtype=torch.bfloat16) + k = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) + v = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) + mask = torch.zeros(B, skv, device=DEVICE, dtype=torch.bool) + + out, lse = cuda_op.forward_with_lse(q, k, v, causal=False, key_padding_mask=mask) + assert (out == 0).all() + assert (lse == float("-inf")).all() + + +# ============================================================================= +# §7.2/8.5 — LSE correctness +# ============================================================================= + + +def test_lse_correctness(cuda_op): + B, hq, hkv, sq, skv = 2, 4, 1, 8, 16 + torch.manual_seed(99) + q = torch.randn(B, hq, sq, D, device=DEVICE, dtype=torch.bfloat16) + k = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) + v = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) + + _, lse_cuda = cuda_op.forward_with_lse(q, k, v, causal=True) + + scale = 1.0 / math.sqrt(D) + g = hq // hkv + k_exp = k.float().repeat_interleave(g, dim=1) + scores = scale * torch.einsum("bhqd,bhkd->bhqk", q.float(), k_exp) + sq_idx = torch.arange(sq, device=DEVICE).unsqueeze(1) + kv_idx = torch.arange(skv, device=DEVICE).unsqueeze(0) + causal_mask = kv_idx <= (skv - sq + sq_idx) + scores = scores.masked_fill(~causal_mask.unsqueeze(0).unsqueeze(0), float("-inf")) + lse_gold = torch.logsumexp(scores, dim=-1) + + torch.testing.assert_close(lse_cuda, lse_gold, atol=5e-2, rtol=2e-2) + + +# ============================================================================= +# §7.2 — Valid-only vs padded: near-equal, NOT bitwise +# ============================================================================= + + +def test_valid_only_vs_padded_accuracy(cuda_op): + """Padding changes reduction width; result is near-equal, not bitwise. + + We use causal=False here so that all valid keys are equally visible + regardless of Skv. The padded path has extra columns masked to -inf, + which changes the sum-exp denominator and can shift results slightly. + """ + B, hq, hkv, sq, skv_valid = 1, 4, 1, 4, 8 + skv_padded = 12 + torch.manual_seed(77) + q = torch.randn(B, hq, sq, D, device=DEVICE, dtype=torch.bfloat16) + k_valid = torch.randn(B, hkv, skv_valid, D, device=DEVICE, dtype=torch.bfloat16) + v_valid = torch.randn(B, hkv, skv_valid, D, device=DEVICE, dtype=torch.bfloat16) + + k_padded = torch.cat( + [ + k_valid, + torch.zeros(B, hkv, skv_padded - skv_valid, D, device=DEVICE, dtype=torch.bfloat16), + ], + dim=2, + ) + v_padded = torch.cat( + [ + v_valid, + torch.zeros(B, hkv, skv_padded - skv_valid, D, device=DEVICE, dtype=torch.bfloat16), + ], + dim=2, + ) + mask = torch.ones(B, skv_padded, device=DEVICE, dtype=torch.bool) + mask[:, skv_valid:] = False + + out_valid = cuda_op.forward(q, k_valid, v_valid, causal=False) + out_padded = cuda_op.forward(q, k_padded, v_padded, causal=False, key_padding_mask=mask) + + atol, rtol = _tol(torch.bfloat16) + torch.testing.assert_close(out_valid.float(), out_padded.float(), atol=atol, rtol=rtol) + + +# ============================================================================= +# §7.3 — Batch invariance (Axis-A bitwise) +# ============================================================================= + + +def test_batch_invariance_single(cuda_op): + """Same sample in full batch vs extracted single — bitwise.""" + B, hq, hkv, sq, skv = 4, 4, 1, 8, 16 + torch.manual_seed(11) + q = torch.randn(B, hq, sq, D, device=DEVICE, dtype=torch.bfloat16) + k = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) + v = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) + + out_full, lse_full = cuda_op.forward_with_lse(q, k, v, causal=True) + + for i in range(B): + out_single, lse_single = cuda_op.forward_with_lse( + q[i : i + 1], k[i : i + 1], v[i : i + 1], causal=True + ) + assert torch.equal(out_full[i : i + 1], out_single), f"Output batch invariance failed i={i}" + assert torch.equal(lse_full[i : i + 1], lse_single), f"LSE batch invariance failed i={i}" + + +def test_batch_invariance_position_permutation(cuda_op): + """Same sample at different batch positions — bitwise.""" + B, hq, hkv, sq, skv = 4, 32, 8, 8, 16 + torch.manual_seed(12) + q = torch.randn(B, hq, sq, D, device=DEVICE, dtype=torch.bfloat16) + k = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) + v = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) + + out_full = cuda_op.forward(q, k, v, causal=True) + + perm = [2, 0, 3, 1] + q_perm = q[perm] + k_perm = k[perm] + v_perm = v[perm] + out_perm = cuda_op.forward(q_perm, k_perm, v_perm, causal=True) + + for new_pos, orig_pos in enumerate(perm): + assert torch.equal( + out_full[orig_pos], out_perm[new_pos] + ), f"Position permutation invariance failed: orig={orig_pos} new={new_pos}" + + +def test_batch_invariance_chunk(cuda_op): + """Batch-dim chunking — bitwise.""" + B, hq, hkv, sq, skv = 4, 4, 1, 8, 16 + torch.manual_seed(13) + q = torch.randn(B, hq, sq, D, device=DEVICE, dtype=torch.bfloat16) + k = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) + v = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) + + out_full = cuda_op.forward(q, k, v, causal=True) + out_chunk = torch.cat( + [ + cuda_op.forward(q[:2], k[:2], v[:2], causal=True), + cuda_op.forward(q[2:], k[2:], v[2:], causal=True), + ], + dim=0, + ) + assert torch.equal(out_full, out_chunk), "Batch-chunk invariance failed" + + +# ============================================================================= +# §7.4 — Sequence-dim chunked-prefill invariance +# ============================================================================= + + +@pytest.mark.parametrize( + "chunk_size,hq,hkv", + [ + (1, 4, 1), + (3, 4, 1), + (8, 4, 1), + (1, 32, 8), + (3, 32, 8), + ], +) +def test_chunked_prefill(cuda_op, chunk_size, hq, hkv): + B, T = 1, 16 + torch.manual_seed(22) + q = torch.randn(B, hq, T, D, device=DEVICE, dtype=torch.bfloat16) + k = torch.randn(B, hkv, T, D, device=DEVICE, dtype=torch.bfloat16) + v = torch.randn(B, hkv, T, D, device=DEVICE, dtype=torch.bfloat16) + + out_full = cuda_op.forward(q, k, v, causal=True) + + outs = [] + for t in range(0, T, chunk_size): + c = min(chunk_size, T - t) + outs.append( + cuda_op.forward(q[:, :, t : t + c], k[:, :, : t + c], v[:, :, : t + c], causal=True) + ) + out_chunked = torch.cat(outs, dim=2) + assert torch.equal( + out_full, out_chunked + ), f"Chunked-prefill invariance failed chunk={chunk_size} hq={hq} hkv={hkv}" + + +@pytest.mark.parametrize("chunk_size", [1, 3, 8]) +def test_chunked_prefill_with_padding(cuda_op, chunk_size): + """§7.4: chunked-prefill with key_padding_mask (mask sliced with Skv).""" + B, hq, hkv, T = 1, 4, 1, 16 + torch.manual_seed(23) + q = torch.randn(B, hq, T, D, device=DEVICE, dtype=torch.bfloat16) + k = torch.randn(B, hkv, T, D, device=DEVICE, dtype=torch.bfloat16) + v = torch.randn(B, hkv, T, D, device=DEVICE, dtype=torch.bfloat16) + mask = torch.ones(B, T, device=DEVICE, dtype=torch.bool) + mask[0, 12:] = False + + out_full = cuda_op.forward(q, k, v, causal=True, key_padding_mask=mask) + + outs = [] + for t in range(0, T, chunk_size): + c = min(chunk_size, T - t) + outs.append( + cuda_op.forward( + q[:, :, t : t + c], + k[:, :, : t + c], + v[:, :, : t + c], + causal=True, + key_padding_mask=mask[:, : t + c], + ) + ) + out_chunked = torch.cat(outs, dim=2) + assert torch.equal( + out_full, out_chunked + ), f"Chunked-prefill with padding failed chunk={chunk_size}" + + +@pytest.mark.parametrize("chunk_size", [1, 3]) +def test_chunked_prefill_lse(cuda_op, chunk_size): + """§7.4: LSE chunked-prefill invariance.""" + B, hq, hkv, T = 1, 4, 1, 12 + torch.manual_seed(24) + q = torch.randn(B, hq, T, D, device=DEVICE, dtype=torch.bfloat16) + k = torch.randn(B, hkv, T, D, device=DEVICE, dtype=torch.bfloat16) + v = torch.randn(B, hkv, T, D, device=DEVICE, dtype=torch.bfloat16) + + _, lse_full = cuda_op.forward_with_lse(q, k, v, causal=True) + + lses = [] + for t in range(0, T, chunk_size): + c = min(chunk_size, T - t) + _, lse_chunk = cuda_op.forward_with_lse( + q[:, :, t : t + c], k[:, :, : t + c], v[:, :, : t + c], causal=True + ) + lses.append(lse_chunk) + lse_chunked = torch.cat(lses, dim=2) + assert torch.equal(lse_full, lse_chunked), "LSE chunked-prefill invariance failed" + + +# ============================================================================= +# §7.5 — Prefill/decode handoff +# ============================================================================= + + +def test_prefill_decode_slice(cuda_op): + """§7.5.1: prefill[:, :, -1:] == decode(q[-1:], k_full, v_full).""" + B, hq, hkv, sq, skv = 1, 4, 1, 8, 16 + torch.manual_seed(33) + q = torch.randn(B, hq, sq, D, device=DEVICE, dtype=torch.bfloat16) + k = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) + v = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) + + prefill = cuda_op.forward(q, k, v, causal=True) + decode = cuda_op.forward(q[:, :, -1:], k, v, causal=True) + assert torch.equal(prefill[:, :, -1:], decode) + + +@pytest.mark.parametrize( + "S_new,hq,hkv", + [ + (1, 4, 1), + (1, 32, 8), + (3, 4, 1), + (3, 32, 8), + ], +) +def test_kv_cache_handoff(cuda_op, S_new, hq, hkv): + """§7.5.2: cat(k_cache, k_new) handoff == prefill tail.""" + B, S_past = 1, 12 + torch.manual_seed(44) + q_full = torch.randn(B, hq, S_past + S_new, D, device=DEVICE, dtype=torch.bfloat16) + k_full = torch.randn(B, hkv, S_past + S_new, D, device=DEVICE, dtype=torch.bfloat16) + v_full = torch.randn(B, hkv, S_past + S_new, D, device=DEVICE, dtype=torch.bfloat16) + + q_new = q_full[:, :, -S_new:] + prefill_tail = cuda_op.forward(q_full, k_full, v_full, causal=True)[:, :, -S_new:] + decode_path = cuda_op.forward(q_new, k_full, v_full, causal=True) + assert torch.equal(decode_path, prefill_tail) + + +def test_kv_cache_handoff_with_padding(cuda_op): + """§7.5.2: cat handoff with padding mask.""" + B, hq, hkv, S_past, S_new = 1, 4, 1, 12, 1 + Skv = S_past + S_new + torch.manual_seed(45) + q_full = torch.randn(B, hq, Skv, D, device=DEVICE, dtype=torch.bfloat16) + k_full = torch.randn(B, hkv, Skv, D, device=DEVICE, dtype=torch.bfloat16) + v_full = torch.randn(B, hkv, Skv, D, device=DEVICE, dtype=torch.bfloat16) + mask = torch.ones(B, Skv, device=DEVICE, dtype=torch.bool) + mask[0, 8:10] = False + + q_new = q_full[:, :, -S_new:] + prefill_tail = cuda_op.forward(q_full, k_full, v_full, causal=True, key_padding_mask=mask)[ + :, :, -S_new: + ] + decode_path = cuda_op.forward(q_new, k_full, v_full, causal=True, key_padding_mask=mask) + assert torch.equal(decode_path, prefill_tail) + + +# ============================================================================= +# §7.6 — Backward +# ============================================================================= + + +def test_backward_smoke(cuda_op): + """Backward runs and produces gradients with correct shapes.""" + B, hq, hkv, sq, skv = 1, 4, 1, 4, 8 + torch.manual_seed(55) + q = torch.randn(B, hq, sq, D, device=DEVICE, dtype=torch.bfloat16, requires_grad=True) + k = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16, requires_grad=True) + v = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16, requires_grad=True) + + out = cuda_op.forward(q, k, v, causal=True) + loss = out.sum() + loss.backward() + + assert q.grad is not None and q.grad.shape == q.shape + assert k.grad is not None and k.grad.shape == k.shape + assert v.grad is not None and v.grad.shape == v.shape + + +def test_backward_fp64_reference(cuda_op): + """§7.6: FP64 high-precision gradient comparison.""" + B, hq, hkv, sq, skv = 1, 4, 1, 4, 8 + torch.manual_seed(56) + q_bf16 = torch.randn(B, hq, sq, D, device=DEVICE, dtype=torch.bfloat16, requires_grad=True) + k_bf16 = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16, requires_grad=True) + v_bf16 = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16, requires_grad=True) + + grad_out = torch.randn(B, hq, sq, D, device=DEVICE, dtype=torch.float32) + + out_cuda = cuda_op.forward(q_bf16, k_bf16, v_bf16, causal=True) + out_cuda.backward(grad_out.to(out_cuda.dtype)) + dq_cuda = q_bf16.grad.float() + dk_cuda = k_bf16.grad.float() + dv_cuda = v_bf16.grad.float() + + scale = 1.0 / math.sqrt(D) + g = hq // hkv + q64 = q_bf16.detach().double().requires_grad_(True) + k64 = k_bf16.detach().double().requires_grad_(True) + v64 = v_bf16.detach().double().requires_grad_(True) + k64_exp = k64.repeat_interleave(g, dim=1) + v64_exp = v64.repeat_interleave(g, dim=1) + scores64 = scale * torch.einsum("bhqd,bhkd->bhqk", q64, k64_exp) + sq_idx = torch.arange(sq, device=DEVICE).unsqueeze(1) + kv_idx = torch.arange(skv, device=DEVICE).unsqueeze(0) + causal_mask = kv_idx <= (skv - sq + sq_idx) + scores64 = scores64.masked_fill(~causal_mask.unsqueeze(0).unsqueeze(0), float("-inf")) + P64 = torch.softmax(scores64, dim=-1) + out64 = torch.einsum("bhqk,bhkd->bhqd", P64, v64_exp) + out64.backward(grad_out.double()) + dq_gold = q64.grad.float() + dk_gold = k64.grad.float() + dv_gold = v64.grad.float() + + torch.testing.assert_close(dq_cuda, dq_gold, atol=5e-2, rtol=2e-2) + torch.testing.assert_close(dk_cuda, dk_gold, atol=5e-2, rtol=2e-2) + torch.testing.assert_close(dv_cuda, dv_gold, atol=5e-2, rtol=2e-2) + + +def test_gradient_batch_invariance(cuda_op): + """§7.6: dQ/dK/dV bitwise identical for same sample at different batch positions.""" + B, hq, hkv, sq, skv = 3, 4, 1, 4, 8 + torch.manual_seed(57) + q = torch.randn(B, hq, sq, D, device=DEVICE, dtype=torch.bfloat16) + k = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) + v = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) + grad_out = torch.randn(B, hq, sq, D, device=DEVICE, dtype=torch.bfloat16) + + q_full = q.clone().requires_grad_(True) + k_full = k.clone().requires_grad_(True) + v_full = v.clone().requires_grad_(True) + out_full = cuda_op.forward(q_full, k_full, v_full, causal=True) + out_full.backward(grad_out) + + for i in range(B): + qi = q[i : i + 1].clone().requires_grad_(True) + ki = k[i : i + 1].clone().requires_grad_(True) + vi = v[i : i + 1].clone().requires_grad_(True) + out_i = cuda_op.forward(qi, ki, vi, causal=True) + out_i.backward(grad_out[i : i + 1]) + assert torch.equal(q_full.grad[i : i + 1], qi.grad), f"dQ batch invariance failed i={i}" + assert torch.equal(k_full.grad[i : i + 1], ki.grad), f"dK batch invariance failed i={i}" + assert torch.equal(v_full.grad[i : i + 1], vi.grad), f"dV batch invariance failed i={i}" + + +def test_gqa_dk_dv_order(cuda_op): + """§7.6/§4.1: GQA dK/dV must follow fixed (hq_local, query_index) order. + + Two checks: + 1. Batch-size invariance: same sample at B=1 vs B=4 gives bitwise-identical dK/dV + (catches unordered atomics or grid-shape-dependent accumulation). + 2. Correctness vs FP64 gold: with asymmetric per-head Q, a reversed local order + would produce different numerical results. This catches "deterministic but wrong". + """ + hq, hkv, sq, skv = 32, 8, 4, 8 + torch.manual_seed(58) + q_data = torch.randn(1, hq, sq, D, device=DEVICE, dtype=torch.bfloat16) + k_data = torch.randn(1, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) + v_data = torch.randn(1, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) + grad_data = torch.randn(1, hq, sq, D, device=DEVICE, dtype=torch.bfloat16) + + # Check 1: batch-size invariance + q1 = q_data.clone().requires_grad_(True) + k1 = k_data.clone().requires_grad_(True) + v1 = v_data.clone().requires_grad_(True) + cuda_op.forward(q1, k1, v1, causal=True).backward(grad_data) + + q_batch = q_data.expand(4, -1, -1, -1).contiguous().clone().requires_grad_(True) + k_batch = k_data.expand(4, -1, -1, -1).contiguous().clone().requires_grad_(True) + v_batch = v_data.expand(4, -1, -1, -1).contiguous().clone().requires_grad_(True) + grad_batch = grad_data.expand(4, -1, -1, -1).contiguous() + cuda_op.forward(q_batch, k_batch, v_batch, causal=True).backward(grad_batch) + + assert torch.equal(k1.grad, k_batch.grad[0:1]), "dK GQA order depends on batch size" + assert torch.equal(v1.grad, v_batch.grad[0:1]), "dV GQA order depends on batch size" + + # Check 2: correctness vs FP64 reference (catches reversed hq_local order) + scale = 1.0 / math.sqrt(D) + g = hq // hkv + q64 = q_data.detach().double().requires_grad_(True) + k64 = k_data.detach().double().requires_grad_(True) + v64 = v_data.detach().double().requires_grad_(True) + k64_exp = k64.repeat_interleave(g, dim=1) + v64_exp = v64.repeat_interleave(g, dim=1) + scores64 = scale * torch.einsum("bhqd,bhkd->bhqk", q64, k64_exp) + sq_idx = torch.arange(sq, device=DEVICE).unsqueeze(1) + kv_idx = torch.arange(skv, device=DEVICE).unsqueeze(0) + causal_mask = kv_idx <= (skv - sq + sq_idx) + scores64 = scores64.masked_fill(~causal_mask.unsqueeze(0).unsqueeze(0), float("-inf")) + P64 = torch.softmax(scores64, dim=-1) + out64 = torch.einsum("bhqk,bhkd->bhqd", P64, v64_exp) + out64.backward(grad_data.double()) + dk_gold = k64.grad.float() + dv_gold = v64.grad.float() + + torch.testing.assert_close(k1.grad.float(), dk_gold, atol=5e-2, rtol=2e-2) + torch.testing.assert_close(v1.grad.float(), dv_gold, atol=5e-2, rtol=2e-2) diff --git a/tests/test_tolerance_contract.py b/tests/test_tolerance_contract.py index fb429d81..f9fe9a1b 100644 --- a/tests/test_tolerance_contract.py +++ b/tests/test_tolerance_contract.py @@ -9,12 +9,12 @@ def test_load_contract_contains_expected_operator_classes(): contract = load_contract() accuracy = contract["accuracy"]["default"] - assert set(accuracy) == {"elementwise", "reduction", "logprob"} + assert set(accuracy) == {"elementwise", "reduction", "logprob", "attention"} def test_load_contract_contains_expected_dtypes(): contract = load_contract() - for op_class in ("elementwise", "reduction", "logprob"): + for op_class in ("elementwise", "reduction", "logprob", "attention"): assert set(contract["accuracy"]["default"][op_class]) == { "float32", "bfloat16",