Module ml::gpu_cuda_attention

attention on NVIDIA GPUs (CUDA): flash attention and linear attention as this library's own CUDA kernels.

Batched multi-head versions of attention.h's flash_attention_forward and linear_attention_forward, on gpu_cuda.h's device buffers and stream:

• Flash attention: softmax(Q K^T * scale) V without materializing the score matrix. fp16/bf16 go to cuDNN's fused attention (the kernels PyTorch's scaled_dot_product_attention uses on Hopper) when cuDNN 9 is installed; otherwise, and for anything cuDNN rejects, this library's FlashAttention-2 kernel runs them on tensor cores (mma.sync, fp32 accumulation) on sm_80+ for head dims up to 128. fp32 runs the same kernel in 3xTF32 (close to full fp32 accuracy) on sm_80+ up to head dim 128; head dim 256 and older GPUs use a CUDA-core kernel with the same algorithm.
• Linear attention: Katharopoulos et al.'s kernel-trick attention with phi(x) = elu(x) + 1, in O((seqLen + kvLen) D^2) instead of O(seqLen kvLen D), with fp32 accumulation. Up to head dim 128 the keys are cut into 64-row chunks: a per-head pass walks them with the D x D state phi(K)^T V in registers (causal: storing the state before each chunk), and every output chunk then runs in parallel from its state plus, when causal, the masked phi(Q_c) phi(K_c)^T V_c inside the chunk. Head dim 256 walks the sequence in 32-row chunks with a running state, as a linear RNN would.

Two calling forms:

• cuda_flash_attention / cuda_linear_attention (and their host-array _f32 versions) take the same arguments as gpu_rocm.h's rocm_flash_attention / rocm_linear_attention, so code moves between AMD and NVIDIA GPUs unchanged: fp32, q and out [heads][seqLen][headDim], k and v [heads][kvLen][headDim] (fold a batch into heads).
• The ex and cuda_half versions take a CudaAttnShape: a batch, grouped- query K/V (kvHeads dividing heads; query head h reads kv head h / (heads / kvHeads)), and a sliding window. Q and out are then [batch][heads][seqLen][headDim], K and V [batch][kvHeads][kvLen][headDim]. All buffers are row-major and contiguous.

Masking: query i sits at position pos = i + (kvLen - seqLen) -- the last seqLen keys are the queries' own tokens, any before them a KV cache (as gpu_mps.h's mps_flash_attn_bf16 and the ROCm ops count it). causal hides keys j > pos; a sliding window w > 0 (flash only) also hides keys j <= pos - w, i.e. the query sees its last w positions.

Head dims: 32, 64, 96, 128 or 256. All math is fp32; the half versions read and write fp16 or bf16 (all four buffers of one dtype). Ops are queued on gpu_cuda's stream like every other op there and can be recorded into a CudaGraph. Each (kernel, dtype, head dim) variant is compiled with NVRTC on its first call (~0.3 s), so make one call of each shape before capturing a graph; non-causal linear attention also sizes its workspace on that call. Forward only, like attention.h. 1 = enqueued (the _f32 forms: done), 0 = bad shape/buffer or no GPU. Link with gpu_cuda.h's libraries (-lcuda -lcublas -lcublasLt -lnvrtc).