We need to write a custom CUDA kernel for scaled_dot_product_attention (SDPA) forward, matching the reference output for fp32. Inputs are [batch_size=32, num_heads=32, seq_len=512, embed_dim=1024], i.e., Q, K, V are each [B, H, N, D] with D=1024. Wait, typical scaled dot-product attention: compute scores = Q*K^T / sqrt(D), then softmax over last dim, then *V. But the standard PyTorch `scaled_dot_product_attention` uses scale factor = 1/sqrt(E) where E = query.size(-1) = embedding_dimension. Wait, for PyTorch scaled_dot_product_attention, the shape is (N, ..., L, E) for query and key, value is (N, ..., S, E), with E being embed dim. So scale = head_dim? Actually for multi-head attention with shape (N, num_heads, L, E_per_head). Here E=1024? That means each head dimension is 1024, which is huge. Usually head_dim = embed_dim / num_heads = 32. But here they say num_heads=32, embedding_dimension=1024. Then each head size is 1024, but typical attention would have head_dim = embed_dim. Wait, shape (batch_size, num_heads, seq_len, embedding_dimension) where embedding_dimension=1024. That means Q, K, V are each [32, 32, 512, 1024]. That is enormous: total elements per tensor = 32*32*512*1024 = 536,870,912 ~ 2 GiB per tensor in fp32 (4 bytes). That's huge, total maybe >6 GB. But likely they just used these shapes; each head dimension = 1024, which is unusual. But we need to write kernel that can handle this? Actually maybe they treat embedding_dimension = head dimension. But then scaled_dot_product_attention will compute Q*K^T where Q, K are [32, 32, 512, 1024]. So matmuls are huge: 512x1024 * 1024x512 = O(512*1024*512) = 268M per head, times 1024 heads = 274B ops. That's huge, maybe memory bandwidth bound but still. However, perhaps the test runs with different sizes? Usually scaled_dot_product_attention expects the last dimension to be head_dim, not total embedding. But the reference code as given uses those shapes. Maybe we should double-check standard PyTorch shape expectations. In `torch.nn.functional.scaled_dot_product_attention`, query shape is `(N, ..., L, E)`; key shape is `(N, ..., S, E)`. It does not require `E = embed_dim // num_heads` explicitly; it just treats E as the head dimension. So with Q [32,32,512,1024], E=1024. So it's huge. But memory is 32*32*512*1024*4 = 2,147,483,648 bytes = 2 GB per tensor. That's okay on an A100 80GB, but on a smaller GPU maybe not. However, the evaluation likely uses the provided shapes. They didn't say it will fit in memory; but presumably yes. Our job: write a custom CUDA kernel that implements SDPA forward. Since the shapes are huge, we need to be very careful about performance. However, we are constrained by the time: we need to write a CUDA kernel that is faster than PyTorch's reference `scaled_dot_product_attention`. PyTorch likely uses cuDNN or FlashAttention (for supported shapes). FlashAttention typically requires head_dim <= 128 or 256 (newer). Here head_dim = 1024, so FlashAttention won't be used by PyTorch? Actually PyTorch's SDPA will use the `math` or `mem_efficient` backend if head_dim > 128. For head_dim 1024, it likely uses a straightforward cuBLAS-based approach: compute S = Q*K^T via batched GEMM, scale, softmax, then compute O = attn*V via batched GEMM. That's O(N^2 D) memory and compute. We need to write a custom CUDA kernel that outperforms standard PyTorch SDPA for these shapes. We have to be really efficient. Since D=1024 is large, maybe a fused kernel isn't straightforward. But we can try to use a custom kernel that splits along the sequence length dimension for the attention score matrix. Or we can do the same as PyTorch but with better memory layout? However, PyTorch uses highly optimized cuBLAS GEMMs. It might be hard to beat with a custom kernel unless we fuse operations to reduce memory bandwidth: doing scale + softmax fused with the QK^T GEMM output, and then multiply with V in a fused manner. For large D, perhaps we can use a kernel that loads Q and K into shared memory and computes tiles. But a full custom implementation of a large batched GEMM is hard to beat cuBLAS. Alternative approach: Since D is large, we can think of using tensor cores with WMMA? That's too complex. Wait, maybe we can treat the attention as two group GEMMs but fuse the intermediate softmax to avoid writing and reading the large attention matrix (32*32*512*512 = 268M elements = 1GB). Fusing can save 2-4 GB/s of memory bandwidth. Since PyTorch likely writes S to a temporary buffer, then reads it for softmax, then writes P, then reads P for GEMM with V. That's several passes over 1GB. A custom kernel that computes tiles of P on the fly and accumulates O can save significant bandwidth. Since D=1024 is large, the tile size for K and V in shared memory might be limited. But we can process the sequence length N=512 in tiles, and for each tile compute scores and update output. Consider the "flash attention" algorithm: We partition O calculation along N (query length) and loop over K,V sequence blocks. For each query block, we maintain running max and sum for softmax. As we loop over K,V blocks, we compute S = Q_block * K_block^T, update the softmax statistics, and accumulate weighted V_block contributions. This algorithm is IO efficient: it loads Q_block once and streams K,V blocks. However, FlashAttention typically assumes head_dim is small enough to fit in SRAM. With head_dim=1024, each Q_block row is 1024 floats = 4KB. If we have a block of 64 queries, Q_block is 256KB. K_block and V_block similarly. For D=1024, one tile of K_block for 64 keys is 256KB. Shared memory per SM is 164KB or 228KB (A100). That's too large. We could use smaller block sizes: e.g., query block size Bq = 16 (64KB), key block size Bkv = 16 (64KB). That's manageable. The problem is that the loop over 512 keys with block size 16 gives 32 iterations, each computing a 16x16 score tile. That's many iterations, but each iteration is tiny (16x16 * 1024 multiplications), leading to low arithmetic intensity but acceptable. The total compute is still the same, but memory reads of K and V are reduced because we keep them in registers or shared memory. Actually we still need to load K and V from global memory for each key block. But maybe the simplest approach is to use `torch.utils.cpp_extension.load_inline` to call `cublasGemmStridedBatchedEx` or `cublasGemmBatchedEx` for the two big GEMMs, but write a custom fused kernel for the scale + softmax + dropout (no dropout here) + multiply by V. However, the standard fused path could be: 1. Compute S = Q * K^T via `torch.matmul` or cublas. [B*H, N, N] 2. Scale S in-place. 3. Compute softmax over last dim in fused kernel that writes P = softmax(S) directly. 4. Compute O = P * V via `torch.matmul` or cublas. This doesn't beat PyTorch's memory-efficient attention because PyTorch likely already does something similar: depending on the chosen algorithm (math, memory-efficient, flash, cudnn). For head_dim=1024, PyTorch's `scaled_dot_product_attention` may fall back to `math` or `memory_efficient` (the xformers style) if available. But the memory_efficient implementation from xformers splits along N and uses custom kernels to reduce memory. It might not support head_dim=1024 either? Actually xformers memory_efficient_attention supports arbitrary head_dims, but it splits into blocks. It might still be able to handle 1024 but performance may degrade. Maybe we can use a kernel that simply calls cuBLAS for QK^T, then uses a custom blockwise softmax kernel that computes the max and sum in shared memory to be very fast, then cuBLAS for PV. The softmax kernel can be more efficient than PyTorch's generic softmax, especially for 512-length rows. And maybe we can avoid reserving a large intermediate S by overwriting it. That saves a copy. Actually, PyTorch's math backend likely does exactly: S = matmul(Q,K^T); S = S / sqrt(D); attn = softmax(S, dim=-1); out = matmul(attn, V). It might use ATen ops that are already quite optimized. But maybe their generic softmax isn't as fast as a custom kernel tuned for this shape. The bottleneck could be the N x N matrix (512x512 = 262K per head, 1K heads => 268M elements). That's not too huge (1 GB). For A100 with 1.5-2 TB/s bandwidth, reading and writing that matrix a few times takes ~1-2 ms. The GEMMs are compute-heavy: QK^T is [1024 heads*batch=1024? Actually B*H = 1024] of (512x1024)*(1024x512). The total FLOPs = 1024 * 2*512*1024*512 = 1024 * 536,870,912 = 549,755,813,888 ~ 5.5e11 FLOPs. In FP32 on A100 19.5 TFLOPS => ~28 ms. But memory bandwidth bound might be less. Actually for matrix multiply with those dimensions (M=N=512, K=1024, batch=1024), the effective throughput could be high. Let's check: M=N=512, K=1024. The square of M is 262K; K is 1024. The operation is M*K*N per batch = 512*1024*512 = 268M FLOPs. Batch=1024 => total ~275B FLOPs. A100 FP32 peak 19.5 TFLOPS => theoretical min 14 ms. In practice maybe 20-30 ms for one matmul? The second matmul is same size. So total maybe 40-60 ms. To beat PyTorch, we need to be close to cuBLAS performance on the GEMMs, and more efficient on the softmax. Since writing custom GEMM kernel to match cuBLAS is hard, perhaps we can still use cuBLAS for the GEMMs via `torch.matmul` or `torch::matmul` in C++ but fuse the scale and softmax in between. However, if we call torch ops in C++, it's essentially same as Python, and the overhead is negligible. But we can't easily fuse cuBLAS matmul with softmax unless we write our own. But maybe we can use `cublasSgemmStridedBatched` for QK^T, then a custom kernel for scale+softmax that writes P in-place or to a new tensor, then another `cublasSgemmStridedBatched` for P*V. The question is: does this beat PyTorch? The reference `scaled_dot_product_attention` might already use `torch._C_ops` that are fused and use cuDNN or FlashAttention? For this shape (large head_dim), maybe not. Actually, `scaled_dot_product_attention` on A100 will try to use `flash_attention` for head_dim <=64 or 128 etc. If head_dim=1024, it will fall back to `math` or `memory_efficient`. The memory_efficient backend uses custom cutlass kernels for arbitrary head_dims but is best for head_dims <= 128 as well. Actually the PyTorch memory_efficient_attention uses Phil Tillet's kernel which supports arbitrary head_dims but is efficient for moderate sizes. For head_dim=1024, it might still be okay but not as optimized as cuBLAS GEMM. Let's check the GPU specs to know the hardware.

The reference is `torch.nn.functional.scaled_dot_product_attention` (SDPA) for fp32 on H100. We need to write a custom CUDA kernel that is faster. The dimensions are: batch=32, heads=32, seq_len=512, head_dim=1024. Wait: in standard attention, embedding_dimension is head_dim. But typically head_dim = embedding_dim / num_heads. Here embedding_dimension = 1024 directly, so each head has dimension 1024. That's actually (32, 32, 512, 1024) Q, K, V. This is a large head dimension (D=1024). So attention is memory-bound: Q,K,V are each 32*32*512*1024 = 536,870,912 elements ~ 2GB each? Actually fp32: 4 bytes per element => 2.147 GB per tensor. That's large but within memory. The O(N^2) attention score matrix S = Q*K^T has shape (32,32,512,512) = 268M elements, ~1 GB. The O(N) reduction over D dimension is large. Given D=1024, math intensity of attention is: - Q*K^T: B*H*S*S*D ops = 32*32*512*512*1024 = 274,877,906,944 FMAs ~ 275 GFLOPs. - Softmax: B*H*S*S ~ 268M ops. - P*V: same as Q*K^T: 275 GFLOPs. Total ~550 GFLOPs. Memory traffic: reading Q, K, V, writing O, and intermediate S, P etc. Best fused kernel reads Q,K,V once and writes O once: 4 * B*H*S*D * 4 bytes = ~8.6 GB reads? Actually reads: Q+K+V is 3 tensors, write O is 1: total ~11.4 GB if just those. Memory bandwidth is 3.35 TB/s, so theoretical minimum time for memory-only fused kernel is ~3.4 ms. Compute is 550 GFLOPs, at 67 TFLOPS (FP32) that's ~8.2 ms. So compute-bound? Wait 550 GFLOPs / 67 TFLOPS = 8.2 ms. So it's actually compute-bound if peak FP32 is 67 TFLOPS. But with tensor cores (TF32 is 989 TFLOPS), if we can use tensor cores we can be much faster. PyTorch SDPA likely uses cuDNN or FlashAttention, which uses tensor cores. Wait, H100 has 1979 TFLOPS BF16 Tensor Core, but FP32 non-tc is 67. The reference uses fp32 precision. The prompt says "at fp32 precision". The inputs are fp32. Our kernel must produce numerically equivalent outputs at fp32 precision. Can we use FP32 tensor cores? H100 FP32 tensor core TFLOPS is 67 (same as non-TC). Wait H100 specs above: FP64 Tensor Core 67, FP32 67 (non-TC), TF32 Tensor Core 989. So normal FP32 doesn't get Tensor Core speedup unless using TF32, which changes precision. Since we need fp32 precision, we cannot use TF32-ish approximations. So compute is actually limited to 67 TFLOPS. But wait, is the reference using cuda kernels that leverage tensor cores via accumulation in higher precision? Actually torch SDPA likely calls into cuDNN or memory-efficient attention. For FP32, it might use standard GEMM (not tensor core) or perhaps it does use tensor cores but keeps accumulation in FP32? TF32 has fp32 accumulation but mantissa is only 10 bits instead of 23, so it's lower precision. But TF32 with fp32 accumulation might still pass the tolerance? Tolerance is 1e-4. TF32 absolute error can be higher. To be safe and definitely correct, we should compute in FP32 exact. But can we beat PyTorch SDPA for FP32 on H100? Observations: - seq_len = 512 is relatively small. - head_dim = 1024 is large. - This means Q*K^T is (512 x 1024) @ (1024 x 512) -> 512x512. And P*V is (512x512) @ (512x1024). These are small square-ish GEMMs. cuBLAS for these sizes might not be super efficient due to small M/N dimensions. Actually M=N=512, K=1024. FlashAttention-style kernel fuses everything and does tile-wise computation to reduce HBM traffic. Given small seq_len (512) and large head_dim (1024), the classic FlashAttention with small tiles might have overhead but could still be faster due to fewer HBM roundtrips. However, writing a fully correct and fast FlashAttention kernel manually is complex. Let's think about simpler approaches. Option 1: Use a custom fused kernel that computes Q*K^T, softmax, and P*V in shared memory, but given seq_len=512 and D=1024, storing a full 512x512 tile in smem for P is fine (1M floats = 4MB). But each SM only has 228KB shared memory. So we cannot store the full 512x512 tile in smem. However we can tile across seq_len dimension. Option 2: Since seq_len=512 is small, maybe we can transpose the order of operations. Specifically, attention with large D can be reformulated: compute K^T * V first? No, standard. Option 3: Simply write a kernel that computes row-wise attention. For each query position i, we need: - scores_j = dot(Q_i, K_j) for j=0..511 - softmax of scores - output_i = sum_j P_j * V_j We can parallelize over batch*heads and over i (queries). Each thread block handles one (batch, head, i) triple. There are 32*32*512 = 524,288 such rows. That's huge. Each row length is 512. D=1024. Memory per thread block: loads Q_i (1024 floats, 4KB) once. Loads K_j and V_j for all j (512 * 1024 * 2 = 2M floats = 8MB). Too much for smem, but we can stream K_j and V_j through registers or smem in chunks. We can process K and V in chunks of e.g., 64. For each chunk of K:V (64x1024): - compute scores for j in chunk using shared memory K tile. - compute max and sum for softmax online. - accumulate output_i weighted by softmax probabilities. This is basically FlashAttention-2 style but extremely simple because we have one query per thread block. Actually each thread block handling a single row might be under-utilized due to only 512 QK dots per block (512*1024 FMAs = 524k FMAs). If we use many threads per block, that's okay. Wait, with seq_len=512 and D=1024, total QK operations per head per batch is 512*512*1024 = 268M FMAs. Dividing by 524,288 blocks gives 512 FMAs per block. That's nothing; each block would finish instantly, but we have huge grid (524k blocks). H100 has ~132 SMs? Actually H100 SXM has 132 SMs (or 144?). Each SM can handle many blocks. But the kernel launch overhead and latency might be significant if each block is tiny. Better: each thread block handles multiple queries (e.g., a tile of 8 or 16 queries) to amortize loading K and V. Then there are 524,288 / 16 = 32,768 blocks. Still plenty. Alternatively, we can parallelize across the D dimension as well. Let's analyze what PyTorch likely does. With D=1024, M=N=512, and fp32, it might call into `memory_efficient_attention` which uses a specific algorithm. On H100, `scaled_dot_product_attention` likely uses cuDNN's fused attention or xFormers. Actually Pytorch's SDPA in Hopper for FP32 might dispatch to FlashAttention if available? FlashAttention typically requires head_dim up to 256 for FA2. FA3 might support 256 or 512. With D=1024, FlashAttention won't be used because it tiles over D dimension? Actually standard FlashAttention tiles over D for the Q and K/V. For large D like 1024, you can still do it but the tile sizes are large; FA2 usually supports up to 256, so PyTorch probably falls back to a kernel or cuDNN that handles larger D via other algorithms, maybe raw GEMM + softmax + GEMM. If it does GEMM + softmax + GEMM, then memory traffic is high and HBM bound. Given D=1024, K=1024. GEMM sizes: 512x1024 @ 1024x512 -> 512x512. This is a very small M,N=512, K=1024 GEMM. cuBLAS performance for such a batched GEMM might not be great. Then second GEMM: 512x512 @ 512x1024 -> 512x1024. Also small. So our custom kernel could fuse everything and be significantly faster than the unfused reference due to better memory access patterns, even if our compute throughput per thread isn't perfect. Let's try to write a simple fused kernel. Kernel Design: - Grid: [batch * heads, ceil(NQ / BQ)] where BQ is queries per block. - Threads per block: e.g., 256. - Each block processes BQ queries and all K/V keys. - We'll tile over the key dimension (N_K = 512) and optionally over the head dimension D=1024. Memory constraints per block: - Q tile: BQ x D. With BQ=8, D=1024, that's 8*1024 = 8192 floats = 32KB. - K/V tiles: maybe tile D into chunks, say D_TILE=64. Then K tile for a key chunk is e.g., BK=64 keys x D_TILE=64 -> 4K floats = 16KB. But D=1024, so we'd need to loop over D_TILE. Wait, we need to compute dot(Q_i, K_j) which sums over D. If we tile D, we compute partial sums and keep a running max/sum? For softmax, that is tricky because max depends on the full dot product. However we can compute the full dot product if we keep accumulators for each query-key pair in registers. With BQ=8, N_K=512, we need 8*512 = 4096 accumulators. That's too many for registers (maybe 4096 * 4 bytes = 16KB, which is possible if distributed among 256 threads -> 16 floats per thread, manageable). Wait, 256 threads * 16 registers = 4096 registers. That's okay if we have 255 max per thread. Actually 16 32-bit registers per thread is trivial. But we also need registers for other things. Alternatively, we can tile over keys: for each key chunk of size BK (say 64), compute scores for BQ queries x BK keys, do online softmax over chunks, and accumulate weighted V. This is the standard FlashAttention algorithm. Given D=1024, Q tile is BQ x D. Let's set BQ = 4 (or 8). BK = 64 (keys per chunk). Algorithm per block: 1. Load Q tile (BQ x D) into shared memory. (4*1024 floats = 16KB) 2. Iterate over keys in chunks of BK (512 / 64 = 8 chunks): a. Load K chunk (BK x D) into shared memory. (64*1024 = 256KB -> too big! 256KB > 228KB shared memory max). So we cannot load full D for BK=64 in smem at once if we want to fit in smem. So we must also tile over D dimension. Let's tile D into D_TILE (say 64). Then: - Q tile: BQ x D_TILE floats per D tile. For BQ=8, D_TILE=64, Q tile = 512 floats = 2KB. - K chunk: BK x D_TILE = 64*64 = 4096 floats = 16KB. - V chunk: BK x D_TILE = same size, but we can load after K or reuse? Actually V is needed for output accumulation weighted by P. V is D-dimensional. So once we have the full P for a key chunk, we need to weight V_j and sum over D dimension for output O_i. This also suggests tiling D for V. FlashAttention typically loads K and V in tiles of D_TILE x BK, and Q in BQ x D_TILE. For each D tile, they compute partial QK^T dot products (accumulating across D tiles), keep the full QK^T accumulators in registers? Actually in FlashAttention-2, they compute QK^T for a tile of Q and a tile of K (after loading full D into SRAM in pieces), store the S matrix in shared memory or registers? Let's recall. Standard FA2: - Tiles Q of size Br x d, K/V of size Bc x d. - It loads Q tile and K tile into SRAM. To handle large d, it actually loads the full d dimension by looping over d and accumulating dot products into an SRAM buffer of size Br x Bc. But if d is large, the SRAM needed for K/V would be Bc x d, which is too large if d=1024 and Bc=128 -> 512KB. So FlashAttention only supports d up to 256 because then Bc=128 => 128*256=32KB fits in SRAM. Thus for d=1024, FlashAttention-2 is not directly applicable in a single block without streaming over d for QK and separately streaming over d for PV. Wait, we can restructure: - For each block (batch, head, q_tile): - Initialize registers for output O_i (size BQ x D) and softmax stats (m_i, l_i) for each query. - Iterate over key tiles (BK keys). For each key tile: - Compute scores S_{i,j} = dot(Q_i, K_j) for i in q_tile, j in key_tile. - To compute dot product, we need to sum over D=1024. We can do this by iterating over D_TILE chunks. - For each D_TILE chunk (e.g., 64 floats): - Load Q_i segment and K_j segment into registers or smem. - Compute partial dot products. Since we want the full S before softmax, we need the entire sum. We can accumulate in a register array of size BQ x BK. Let's say BQ=16, BK=16. Then BQ*BK=256 partial accumulators. We can loop over D_TILE chunks (1024/64 = 16) and accumulate. - After finishing D loop, we have S_{i,j}. - Now we have S matrix. We can apply softmax scaling (1/sqrt(D)), and then do online update with previous chunks. - Then we need to accumulate weighted V into O_i for all D dimensions. For each D_TILE chunk: - Load V_j segment (BK x D_TILE). - Load the softmax weights P_{i,j} (computed from S). - Update O_i segment += sum_j P_{i,j} * V_j segment. This is possible. But memory for S registers: BQ*BK floats. If BQ=16, BK=32 => 512 floats = 2KB => 512 registers. For 256 threads, that's 2 registers per thread. Very manageable. Let's choose parameters: - BQ = 16 (queries per block) - BK = 32 (keys per iteration over K) - D_TILE = 64 (dimension tiles) Then: - Q tile: 16 x 64 = 1024 floats = 4KB smem. - K tile: 32 x 64 = 2048 floats = 8KB smem. - V tile: 32 x 64 = 2048 floats = 8KB smem. Total smem = 20KB, well under limit. Number of queries per block = 16. Number of blocks = B*H*ceil(N/BQ) = 32*32*ceil(512/16) = 32*32*32 = 32768 blocks. That's a lot but okay. For each block: - Outer loop over key chunks: NK/BK = 512/32 = 16 chunks. - Inside each key chunk: - Inner loop over D tiles: D/D_TILE = 1024/64 = 16 tiles to compute S. - For each D tile: load Q[...], K[...] into smem, compute partial S. - After S computed, do softmax online update (requiring S max over j, then sum, etc.) - Then inner loop over D tiles again: for each D tile, load V[...] into smem, update O using P matrix. So total number of memory loads per block: - Q: loaded once (16*1024 floats = 16k floats) -> but tiled into 16 tiles: 16 loads. - K: loaded per key chunk, per D tile: 16 chunks * 16 tiles = 256 loads of 32*64 floats = 2048 floats each -> total 524k floats per block. - V: same as K: 524k floats per block. - O: written once (16*1024 floats = 16k floats). Total memory traffic per block: ~1M floats = 4MB. For 32768 blocks, total ~131GB. That's a lot but not crazy compared to 3.35 TB/s. Wait, total memory traffic is traffic to HBM. In this algorithm, each block reads its own Q, K, V tiles. There is no sharing across blocks. So total DRAM reads: Q: B*H*N*D = 537M floats (total Q once) K: B*H*N*D = 537M floats (but each block reads the entire K for its head? Wait no, each block reads only K tiles for its head. But all blocks for a given head collectively cover all keys. Since blocks are independent, total K reads = B*H * (N/BQ) * N * D? Actually each block processes BQ queries and reads all N keys. So total K reads = (number of blocks) * (N * D) = (B*H*N/BQ) * N*D = B*H*N^2*D / BQ. With BQ=16, that's 32*32*512*512*1024 / 16 = 32*32*512*512*64 = 32*32*16,777,216 = 17.2B floats? That's huge. Wait, that's because each query block reads the entire K and V. But since each query block is for different queries, it must read K and V. However K and V are the same for all queries within a head. So total K/V reads are read once per block. But there are many blocks. The total HBM reads for K/V would be B*H * (num_q_blocks) * N*D. Since num_q_blocks = N/BQ, total K reads = B*H * (N/BQ) * N*D. That's B*H*N^2*D / BQ. For our numbers, that's 32*32*512^2*1024 / 16 = 32*32*512*512*64 = 17179869184 floats = 68.7 GB of K reads. Same for V: 68.7 GB. Q reads: each block reads BQ*D, total B*H*N*D = 2GB. O writes: 2GB. Total memory traffic: ~141GB. That's way more than the non-fused trio which is ~8GB. But HBM bandwidth is 3.35TB/s, so 141GB takes ~42ms. That's terrible. Wait, my math above must be wrong because each block reads K and V for its own head. Actually each block corresponds to a specific head; all blocks for the same head read the same K and V. But they are independent, so there is no caching across blocks in different SMs unless the L2 cache can help. With 32*32=1024 heads, the total K matrix per head is 512*1024=524k floats = 2MB. If L2 cache is ~50MB, only about 25 heads' K data can fit. So K is mostly read from HBM for each block. This means total K reads = num_blocks * (N*D) = 32768 * 524288 = 17.2B floats. That's indeed huge. Actually this algorithm is terrible for HBM bandwidth because it reads K and V (N*D) per query block, which is O(N^2*D/BQ) rather than O(N*D). For N=512 and D=1024, the standard unfused approach reads K and V once overall: total HBM reads: Q: 2GB, K: 2GB, V: 2GB. S=(QK^T) written and read: 1GB. P written and read: 1GB. O written: 2GB. Total ~10GB. The fused kernel's goal is to reduce HBM traffic by keeping intermediate results on-chip. But if we re-read K and V many times (once per query block), we might increase traffic compared to unfused! However for FlashAttention, they read K and V once per block because there is no other way. But they make the number of query blocks small enough that K/V re-read is okay compared to the O(N) scheme? Let's check standard FlashAttention analysis: total HBM accesses = O(N*D) for Q,K,V,O plus O(N^2) for softmax but the N^2 factor is reduced by block size. Actually the key is that FlashAttention reads K and V once per query block, but the number of query blocks is N/Br. So total K reads = (N/Br) * Nc * (Bc * D)? Wait, no. Let's re-derive: - Grid: N/Br query blocks. - Each query block reads ALL K blocks (Nc = N/Bc) and V blocks. - So total K reads = (N/Br) * N * D. - Standard: K reads = N*D. - Ratio = N/Br. So if Br is small, we read K many times. For N=512, if Br=16, ratio=32. So we read K 32 times! That's bad. But FlashAttention is still faster because it avoids materializing the NxN S and P matrices in HBM. For large N (e.g., 4096 or 8192), the S and P matrices are huge relative to K and V. Here N=512, so S and P are only 512*512=262k floats = 1MB per head, much smaller than K and V (2MB each). In fact, HBM traffic for unfused: Q (2MB), K (2MB), V (2MB), S (1MB), P (1MB), O (2MB) = 10MB per head. Total 10GB. For fused with Br=16: Q reads = 2MB, K reads = 32*2MB=64MB, V reads=64MB, O writes=2MB. Total 132MB per head, 135GB overall. That's 13.5x more HBM traffic than unfused! So a naive fused kernel would be much slower on H100 if it's HBM bound. But compute-bound? Unfused does two GEMMs and a softmax. With small M,N=512, K=1024 GEMMs, maybe cublas is inefficient and the kernel is latency-bound. A custom fused kernel could use more parallelism and avoid kernel launch overhead. But if we have 13x more HBM traffic, it's doomed. We need to design a kernel that reads K and V roughly once, not N/Br times. That means we should parallelize over keys or over the sequence dimension differently. Wait, what if each block handles a tile of keys and a tile of queries, and we compute the full attention output for that query tile? That inherently requires reading K and V for all keys per query tile. But what if we parallelize over D and over batch*heads, and within each block we process the entire sequence? E.g., each block handles one (batch, head) and one query i. Then K and V are read once per block but there are N blocks per head, leading to N times reads of K and V. That's the same problem. Actually, in the unfused approach, PyTorch does batched GEMMs across all batch*heads. These GEMMs are large in batch dimension (1024 batches of 512x512x1024). cuBLAS can process this efficiently. The total memory reads are indeed just Q,K,V once plus intermediates. The GEMM is compute-bound (since it's batched and total FLOPs are high). Wait, compute intensity for first GEMM: total FLOPs = B*H*N*N*D = 275 GFLOPs. Memory reads for Q and K = 4MB per head * 1024 heads = 4GB. Wait, Q is 2MB per head, K is 2MB per head. Total reads = 4GB. Writes S = 1MB per head = 1GB. Total HBM = 5GB. FLOPs/HBM byte = 275G / 5GB = 55 FLOPs/byte. FP32 throughput is 67 TFLOPS, bandwidth 3.35 TB/s => 20 FLOPs/byte for compute-bound. So 55 FLOPs/byte is compute-bound! Indeed, since the batch dimension is huge (1024), the GEMM is effectively a very large operation. cuBLAS will handle this well. Similarly the second GEMM is also compute-bound. So why would SDPA be slow? Let's think: the operation is batched matrix multiplication with: - BatchCount = B*H = 1024. - M=N=512, K=1024. This is a BMM. cuBLAS performance for BMM with these dimensions in FP32 on H100 might actually be decent but not optimal due to the M and N being relatively small (512) compared to K (1024). The arithmetic intensity per batch is M*N*K / (M*K + N*K) = 512*512*1024 / (512*1024 + 512*1024) = 256 FLOPs per byte of input. That's very compute-bound. And with 1024 batches, occupancy is high. Wait, why would a custom kernel beat cuBLAS + elementwise fused softmax + cuBLAS? Perhaps PyTorch's SDPA for this specific configuration (D=1024) does not use an optimal backend. Or maybe it does use cuDNN which might have some overhead. And the prompt says "faster than the reference", so a moderate speedup is likely possible. What if we simply fuse the entire attention into a single kernel but arrange blocks to process one head at a time and have all threads in a block cooperatively compute the entire output O for that head? That would read Q,K,V once per head. For each head, data size is small (Q=2MB, K=2MB, V=2MB). We can have one thread block per head that reads Q,K,V into shared memory in tiles? 2MB is too big for shared memory (228KB). So we can't read an entire head's data into smem. But we can have multiple thread blocks per head, each computing a tile of O. However K and V would be read multiple times. The key is to ensure that K and V reads from HBM are minimized by L2 cache reuse. If we have multiple thread blocks in the same head that share K and V through L2, the total HBM traffic might be close to reading them once, depending on L2 capacity and number of blocks. L2 size on H100 is 50MB. Q,K,V for one head = 6MB. S and P for one head = 2MB. Total per head active dataset is much less than 50MB. If we schedule all blocks for a head to run concurrently or near-concurrently, K and V will be resident in L2 and HBM reads will be low. But if we have 1024 heads and many SMs, L2 might get thrashed. However each head's K/V is 4MB. L2 is 50MB, can hold ~12 heads' K/V. With 1024 heads, blocks from different heads might evict each other. But if we use multiple thread blocks per head (e.g., 32 blocks per head computing tiles of O), the total memory traffic might still be high because L2 cannot keep all blocks' K/V. However if each block accesses K/V sequentially and L2 is large enough to hold a head's K/V, then later blocks reusing K/V will hit L2. But if blocks are from different heads interleaved, L2 thrashing occurs. We can try to help this by using grid swizzling? Pytorch's kernel has the same issue. Maybe the issue is simpler: PyTorch SDPA on fp32 with D=1024 doesn't use FlashAttention and instead does plain GEMM + softmax + GEMM. The first GEMM is Q @ K^T, producing S. Then softmax is applied. Then second GEMM is P @ V. There might be three separate kernel launches, and small M/N dimensions might lead to suboptimal utilization, especially if cuBLAS for small M/N is poor. But the batch count is high. Actually the cuBLAS batched GEMM API (cublasGemmBatchedEx or strided BMM) is efficient for these sizes. Let's check the theoretical performance: - First BMM: 1024 batches of 512x512x1024 GEMM. - Total FLOPs = 1024 * 2 * 512 * 512 * 1024 = 549.8 GFLOPs? Wait 2*M*N*K per batch (because C+=A*B). Actually for C = A*B, it's approx 2*M*N*K FLOPs. Wait FMA counts as 2 FLOPs? Usually 1 FMA = 2 FLOPs. But in GEMM throughput, FMAs per second is used. Let's use 1 FLOP = 1 FMA for simplicity? Actually no, 1 FMA = 2 FLOPs. cuBLAS performance is usually reported in TFLOPS (where 1 FMA = 2 FLOPs). For attention, QK^T = 275 GFLOPs (if 1 FMA = 1 FLOP) or 550 GFLOPS (if 1 FMA = 2 FLOPs). Let's use 1 FMA = 2 FLOPs. So total FLOPs ≈ 550 GFLOPs for both GEMMs + negligible softmax. - Peak FP32 non-TC = 67 TFLOPS. So lower bound on compute is 550 / 67 = 8.2 ms. If we actually achieve 50% of peak, it's ~16ms. Wait, 67 TFLOPS is the whole GPU. Can we achieve close to that for FP32 non-TC on batched GEMM? Usually cuBLAS FP32 can get ~80-90% of theoretical for large GEMMs. For BMM with M=N=512, K=1024, each GEMM is decent but not huge. The total operations are 1024 * (512*512*1024) = 274,877,906,944 FMAs (274B FMAs). That's a lot. Let's compute time at 80% of 33.5 TFLOPS? Wait, 67 TFLOPS is for the whole GPU. For FP32 on H100, throughput is 67 TFLOPS. For tensor core TF32, it's 989. For FP32 non-TC, it's 67. This is for the whole GPU (132 SMs). So per SM, it's 67/132 ≈ 0.5 TFLOPS = 512 GFLOPS. For a GEMM with M=512,N=512,K=1024 on one SM, it might take a while, but with many SMs working in parallel across batches, it's fine. If we can get 70% of peak, time = 550 GFLOPs / (0.7 * 67 TFLOPS) = 11.7 ms. Memory time at 3.35 TB/s for 10GB is 3ms. So compute bound, around 12ms. Can a custom kernel beat cuBLAS? It seems hard if cuBLAS is well-optimized. However, maybe PyTorch's SDPA does not use cuBLAS GEMM but some less efficient path (like falling back to a generic kernel because of head_dim 1024). Or maybe cuDNN's fused attention is used. In either case, we need to write something faster. Alternative idea: use Tensor Cores with FP32 accumulation but in a way that is still exact? No, tensor cores on H100 for FP32 mode (called FP32 TC) do exact FP32? Wait, the spec says FP64 TC = 67, FP32 = 67. Is there a separate FP32 tensor core path? On Ampere/Hopper, there is no FP32 tensor core mode that does exact FP32 multiplication with FP32 accumulation? Actually there is: FP64 TC exists, but FP32 TC does not. The "FP32" in the spec is via standard CUDA cores. TF32 uses tensor cores but with truncated mantissa. So for exact FP32, we must use CUDA cores. Thus peak is 67 TFLOPS. Wait, I just realized: maybe the reference is using `torch.nn.functional.scaled_dot_product_attention` with fp32, and on H100 it might dispatch to a kernel that is actually slower than expected because it chooses a suboptimal backend (e.g., trying to use FlashAttention but failing due to D=1024 and falling back to a generic eager implementation, or even the math implementation). In PyTorch 2.x, SDPA selects a backend. On CUDA, for sm80+ (Ampere/Hopper), available backends include flash_attention, memory_efficient_attention, and cudnn. For FP32 and D=1024, I think no flash_attention supports D=1024. So it will likely fall back to `memory_efficient_attention` from xFormers or plain BMM+softmax+BMM. If it falls back to plain BMM+softmax+BMM, that's exactly three ops. Could we fuse the softmax and the second GEMM? Or fuse all three? Fusing across the sequence dimension with shared memory can reduce HBM traffic and kernel launch overhead. But as discussed, reading K and V many times can increase traffic. However if we choose the tiling carefully, perhaps the L2 cache mitigates this, and the elimination of writing S and P to HBM yields a net win. For N=512, S and P are small (1MB per head). Savings from not writing S and P are 2MB per head. K and V re-read overhead depends on tile size. If we tile with Br = 128 (query tile), then number of query blocks per head = 512/128 = 4. So K and V are read 4 times per head. K+V size = 4MB per head. Re-reading 4 times = 16MB per head. Q reads = 2MB. O writes = 2MB. Total = 20MB per head. Unfused total = Q+K+V+S+P+O = 2+2+2+1+1+2 = 10MB per head? Wait S and P are intermediate. Actually unfused writes S to HBM (1MB), reads it for softmax (1MB), writes P to HBM (1MB), reads it for second GEMM (1MB). Total HBM traffic per head = read Q(2) + read K(2) + read V(2) + write S(1) + read S(1) + write P(1) + read P(1) + write O(2) = 12MB. Plus read Q again? In first GEMM, Q and K are read. Softmax reads S and writes P. Second GEMM reads P and V, writes O. So total = 2+2+2+1+1+1+1+2 = 12MB. (Actually some reads/writes might be cached and not go to HBM, but worst-case distinct is 12MB). Our fused kernel with Br=128 reads Q once (2MB), K 4 times (8MB), V 4 times (8MB), writes O once (2MB). Total = 20MB. That's actually more traffic than unfused! So fusing without also fusing the softmax and second GEMM might not reduce HBM traffic. Actually our fused kernel does fuse all three, eliminating S and P entirely from HBM. But K and V are re-read. To reduce K/V re-reads, we can increase Br to 512. Then each block handles all queries, and K/V are read once per head. But can a single block handle all 512 queries? Threads per block limited to 1024. We could have the block loop over queries in groups, using shared memory for K and V tiles. With one block per head, we have 1024 blocks total. That's not enough to saturate H100 (which needs many thread blocks). So we'd need multiple blocks per head. What if we parallelize over D dimension as well? E.g., each block handles a subset of queries and a subset of D features. Suppose we have blocks for (head, q_tile, d_tile). Then each block reads K and V for all keys but only for its d_tile. If d_tile is e.g., 64, then K/V read per block = N * d_tile * 4 bytes = 512*64*4 = 128KB. Number of d_tiles = 1024/64 = 16. Number of q_tiles = 512/128 = 4. Total blocks per head = 4 * 16 = 64. Total blocks = 1024*64 = 65536. Good parallelism. K reads per head = q_tiles * d_tiles * N * d_tile = 4 * 16 * 512 * 64 = 2,097,152 floats = 8MB. But K size per head is N*D = 512*1024 = 524,288 floats = 2MB. So K is read 4 times per head (because q_tiles=4). Since d_tiles loop over D dimension within each block, but all blocks across d_tiles for the same q_tile read the same K rows but different columns. So K is effectively read q_tiles times. Same for V. So K/V re-read factor = q_tiles = 4. Total K+V reads = 4*2MB + 4*2MB = 16MB. Q reads: each block reads its q_tile and d_tile of Q. Total Q reads = q_tiles * d_tiles * (q_tile_size * d_tile) per head = 4*16*128*64 = 524,288 floats = 2MB. That's just once per head! Because Q is partitioned disjointly. O writes similarly once per head (2MB). Total HBM traffic = 2 + 16 + 2 = 20MB again. Still more than 12MB. Wait, but if q_tiles=4, then K/V are read 4 times. Unfused reads K and V once each. S and P are read/written twice each. So fused saves 4MB (S+P I/O) but costs 6MB extra (K+V re-read factor 3 extra times). Net worse by 2MB. Actually 20MB vs 12MB. Not good. What if we make q_tiles=1? Then one block per head handles all queries but only a d_tile of D. But 512 queries with d_tile=64 means each block computes partial outputs for all queries but only 64 features. How to compute attention scores? Scores require the full dot product over D. So if we split across D, we need to reduce across blocks. That's awkward because softmax depends on the full dot product. So splitting D for output is easy, but splitting D for score computation requires a reduction. Alternatively, we can split across the key dimension (N). Each block handles a chunk of keys for all queries, computes local softmax, and then we do a cross-block reduction. That's the tree reduction approach used in some kernels. But that adds complexity and synchronization. Another approach: Transpose the attention. Since D=1024 >> N=512, we can compute K^T @ V first? No, that's not attention. What about using the identity: attention(Q,K,V) = (Q @ K^T) @ V / sqrt(D). This is just two GEMMs. Since cuBLAS is already fast, maybe we can just call it directly but in a more efficient way? Or maybe the PyTorch eager fallback is not using batched strided GEMM but some slower path. Actually, let's consider if the SDPA reference uses `math` backend because the dimensions don't fit `flash_attention` or `memory_efficient_attention`. With `math` backend, it computes: S = matmul(Q, K.transpose(-2,-1)) / sqrt(D) P = softmax(S, dim=-1) out = matmul(P, V) This involves two `bmm` calls + elementwise ops. Could we write a kernel that uses the CUTLASS-like approach but is a single fused kernel? Even if memory traffic is slightly higher, the benefit of fusion might be in avoiding kernel launch overhead and launching more thread blocks to keep SMs busy. But for such a large batch count (1024), the GEMMs themselves launch enough thread blocks. Let's look at this from a different angle. The problem says "Produce numerically equivalent outputs at fp32 precision. Run faster than the PyTorch reference." So there is a reference implementation running on H100. We need to beat it. What is the actual reference runtime? We don't know but we can guess it's maybe around 10-20ms or more if it's using the math fallback. If PyTorch dispatches to a bad backend, even a simple custom kernel could be faster. What if we simply implement a custom CUDA kernel that for each (batch, head, query) uses a loop over keys and computes the dot product over D dimension in a shared-memory-friendly way, but with each block handling multiple queries to utilize the GPU? Even if HBM traffic is higher, the kernel might be faster if the reference is doing a lot of memory-bound intermediate steps or is heavily kernel-launch-limited. Wait, I should test this. But I cannot test without writing code. Let's think: with B*H=1024 batches, a naive custom kernel that processes each (batch, head, query) independently might launch 512*1024 = 524,288 thread blocks. H100 can handle that. Each thread block has e.g., 128 threads. The work per thread block: 512 dot products of length 1024, plus weighted sum over 512 vectors of length 1024. A thread block can compute this by having each thread compute one query-key dot product? That would be 512 threads per query. But we can do better. Let's design a kernel where each thread block handles one (batch, head, query). Total blocks = 524,288. Good occupancy? 132 SMs, each can run up to 32 blocks -> 4224 blocks concurrently. So many waves. Inside block: - Load Q_i (1024 floats) into shared memory. (4KB) - Iterate over K_j in groups of e.g., 32 or 64. For each group: - Load K_j (group_size * 1024 floats) chunk? That's 32*1024 = 128KB, too large for smem. So we need to chunk over D. - But we can compute dot products without smem: each thread reads elements of Q_i and K_j from global memory. But global memory loads would be very inefficient. Alternatively, tile over D: smem for Q_i and K_j segments. e.g., D_TILE=64. - Load Q_i segment into smem (64 floats). Load K_j segment for all j in group into smem (group_size * 64 floats). With group_size=64, that's 4096 floats = 16KB. Total smem = 16KB + 256B = fine. - Compute partial dot products for all j in group. Each thread computes one or more dot products. - Accumulate in registers. After full D loop: - We have scores for all j (512). Do softmax: need max, then sum. - Then compute output_i: again iterate over D tiles. For each D tile, load V_j segment for all j (or chunk). Use softmax weights to accumulate. This is exactly the per-query block approach. Total memory traffic: Q_i read once per block (4KB). K_j read N*(D) per block? Actually K_j is read group by group. All keys are read: 512*1024*4 = 2MB per block. V_j same: 2MB per block. O_i write: 4KB. Total per block: 4MB + 4KB. For 524,288 blocks: 2,097,152 MB + 2GB = ~2,099GB = 2TB of K/V reads. That's insane. HBM bandwidth 3.35 TB/s yields ~0.6s. This would be way slower than reference. So we absolutely cannot have each query block read all K and V independently without L2 reuse. L2 reuse would help only if multiple query blocks for the same head run serially on the same SM and L2 retains K/V. But 524,288 blocks is huge; K/V per head is 2MB. L2 is 50MB, holds 25 heads. So L2 can serve 25 heads worth of K/V without HBM reads for subsequent blocks of those heads. But there are 1024 heads. So only 2.4% of blocks get L2 hit? Actually if blocks from the same head are scheduled to the same SM consecutively, L2 can keep the head's K/V. But with 524k blocks, scheduling is interleaved across heads. We could use grid stride to ensure blocks from the same head are contiguous and run on same SMs? CUDA scheduling is not guaranteed, but generally blocks with consecutive IDs run on the same SM if they are from the same kernel launch, until the SM is full. However with 524k blocks, they won't all fit. But an SM processes a block, then finishes, then takes another block. If the new block is from the same head, L2 may still have K/V. But grid order is linear: blockIdx.x corresponds to query. So within a head, query blocks are consecutive. An SM might process query 0, then query 1, etc. It can hold K/V in L2 across these. Then K/V reads hit L2! That's crucial. If K and V are read from L2 instead of HBM, memory bandwidth is much higher (maybe 10-20 TB/s). If L2 bandwidth acts as a cache, the latency is lower. But still, reading 2MB per query block from L2 for 512 queries means 1GB of L2 traffic per head. L2 bandwidth is maybe 10 TB/s, so 1GB / 10TB/s = 0.1ms per head? Not sure. Wait, but if L2 holds K/V for a head (4MB), then each query block reads K/V from L2. 512 query blocks per head * 4MB = 2GB of L2 reads per head. L2 bandwidth on H100 is around 20 TB/s? Let's assume 10 TB/s. Then time = 2GB / 10TB/s = 0.2ms per head. For 1024 heads, total 204ms if serialized. But heads run in parallel across 132 SMs. So 1024/132 ≈ 8 waves per SM. Each wave handles ~4 heads (since each SM can run multiple blocks concurrently? Each block is per query, so many blocks per head). Actually occupancy: each block maybe 128 threads. H100 can run many. Let's say 32 blocks per SM. If each block is one query, then an SM can run 32 queries concurrently. That's from 1 head. So per SM, we process 32 queries from a head. They all read K and V from L2. L2 bandwidth is shared. 132 SMs * 32 blocks = 4224 blocks concurrent. All read K/V. So total concurrent L2 read bandwidth demand = 4224 blocks * (K+V read rate). But K and V are loaded cooperatively and cached. Since K/V per head is 4MB and all concurrent blocks for the same head use the same K/V, the L2 bandwidth per head is shared. If 32 blocks on one SM read K/V, they access the same lines, so L2 bandwidth is not multiplied by 32; it's roughly the bandwidth to read K/V once per block (but with coalescing, the SM reads each line and broadcasts). Actually within an SM, threads from different blocks are resident. Loads from global memory go through L1 per SM. So 32 blocks on same SM each load K/V; the L1 cache can only hold so much. But if they all access the same K/V data, it might be in L1 after the first block loads it, and subsequent blocks hit L1? L1 is 128KB per SM. K+V is 4MB. So no, K+V doesn't fit in L1. So each block's loads will go to L2 for K/V. If the data is in L2, it's an L2 hit. But 32 blocks issuing loads concurrently will each read the same data from L2, causing L2 bandwidth contention per SM. However H100 L2 bandwidth is high. This is getting complicated. Let's step back. Given the problem setup, the target speedup might be modest, and a well-written TILE-based kernel could be sufficient. We need to produce something reasonable. The most common "fused attention" kernel that is relatively simple is a block-sparse or tile-based approach. Since N=512 and D=1024, a simple approach is: - Each thread block computes O for a tile of Br queries and all keys, using shared memory for K and V tiles. Since N=512 is small, we can load K and V in tiles of Bc keys and accumulate. - To minimize K/V re-reads from HBM, we choose Br as large as possible to reduce the number of query blocks per head: Br = 64 or 128. Then K/V are read N/Br = 8 or 4 times per head. This might be okay if the reference is doing more HBM traffic or has worse compute utilization. Wait, but we saw fused with Br=128 has more HBM traffic (20MB per head vs 12MB unfused). However the unfused reference does two GEMMs and elementwise. If the reference achieves high compute utilization (say 60% of 67 TFLOPS = 40 TFLOPS), time = 550 GFLOPs / 40 TFLOPS = 13.75ms. If our fused kernel is memory bound with 20MB per head and 1024 heads: total 20GB. At 3.35 TB/s = 6ms. Wait, 20MB * 1024 = 20,480 MB = 20GB. At 3.35 TB/s = 6ms. Actually 20GB / 3.35 TB/s = 0.006s = 6ms. That's faster than 13.75ms! But wait, we are not taking into account K/V re-read from HBM, not L2. If K/V are actually read from HBM 4 times per head, total per head is 4*4MB + 2MB + 2MB = 20MB. 20MB * 1024 = 20GB. At 3.35 TB/s, 6ms. But if we also account for Q and O traffic, yes. But the compute in our fused kernel is the same: we still do 550 GFLOPs. If we are compute-bound? Let's see arithmetic intensity for the fused kernel: FLOPs = 550 GFLOPs. HBM traffic = 20GB. FLOPs/byte = 27.5. This is below the compute-bound threshold (peak FP32 / bandwidth = 67 / 3.35 = 20). Actually 27.5 > 20, so it is slightly compute-bound! Wait, compute intensity needed to be compute-bound is > 20. So we are compute-bound. Thus if we achieve e.g., 50% compute efficiency, time = 550 / (0.5*67) = 16ms. But if we are slightly compute-bound, memory time is 6ms, so the bottleneck is compute. So our fused kernel would take ~16ms if compute efficiency is 50%. For the reference, if it uses cuBLAS, it likely has much better compute efficiency (maybe 70%+) and thus takes ~11ms. So the reference might be faster. But if the reference doesn't use cuBLAS (e.g., uses the math fallback with `baddbmm` etc.), its compute efficiency might be lower, maybe 30%? Then it takes ~27ms. Our optimized kernel could take 15ms and win. We need to know what PyTorch actually does. SDPA on CUDA for FP32 might dispatch to `CUDNN_SDPA`, `EFFICIENT_ATTENTION`, or `FLASH_ATTENTION`. For D=1024, FlashAttention is out. Efficient attention (xFormers) typically also has dimension limits (D<=128 or 256). So likely it falls back to either `CUDNN_SDPA` or `MATH`. On Hopper, cuDNN likely supports fused attention for these dimensions. cuDNN's fused attention might be very fast, maybe close to optimal. If cuDNN is used, beating it will be extremely hard. If math fallback is used, beating it should be easy. How to know? We can look at the dimensions. In PyTorch 2.4 (current), SDPA on CUDA with sm90 (Hopper) will try cuDNN if available and dimensions fit. I believe cuDNN supports head_dim up to 128, maybe 256 for some configs, but 1024 is large. Actually cuDNN's fused attention supports head_dim up to 64 for some sequences? I'm not sure. Alternatively, maybe it uses the `math` fallback. We could check by writing a dummy kernel that does something simple and see if it passes. But first we need to write something. Let's reconsider the possibility of using Tensor Cores. The tolerance is `atol=rtol=1e-4` for fp32. Wait, the prompt says: "Outputs must match the reference within tolerance: fp32 uses atol=rtol=1e-4; fp16/bf16 use atol=rtol=1e-2." If we use standard CUDA cores in FP32, it's exact (up to associativity). If we use Tensor Cores in TF32, the error is larger than FP32 but maybe still within 1e-4? TF32 has 10 bits of mantissa, giving relative error ~2^{-10} = 0.001. For typical values around 1, absolute error could be ~0.001. That is within 1e-4? No, 1e-4 is 0.0001. 0.001 > 0.0001, so TF32 would likely fail the tolerance for many outputs, especially since there are many accumulations. So we cannot use TF32 tensor cores. We must use FP32 CUDA cores. But what about Hopper's FP32 Tensor Cores? Actually Hopper introduced FP8, but does it have FP32 TC? The spec says FP64 TC = 67, FP32 = 67 (non-TC). So likely no FP32 Tensor Cores. Thus exact FP32 means CUDA cores. So compute peak is 67 TFLOPS. Reference likely also uses CUDA cores for FP32 (since cuBLAS FP32 uses CUDA cores). Therefore the difference in speed will come from: 1. Better utilization of SMs for small M,N dimensions. 2. Fusion reducing kernel overhead and memory traffic. 3. Better occupancy/warp scheduling. Wait, but cuBLAS FP32 for batched GEMM with M=512,N=512,K=1024 and batch=1024 is probably already very well optimized. Let's estimate its performance. Each GEMM in the batch is 512x512x1024. This is M=N=512, K=1024. A single GEMM of this size can achieve quite good cuBLAS performance. In fact, 512 is a nice number. cuBLAS will use a NK dispatch or similar. On H100, a single 512x512x512 GEMM might take ~10-20 microseconds? Actually let's use roofline: A single 512x512x1024 GEMM: FLOPs = 2*512*512*1024 = 536M FLOPs. At 67 TFLOPS, min time = 8 microseconds. But that's full GPU. For batched, if you run many in parallel, you might approach that. However, the kernel launch overhead and small per-GEMM size might limit occupancy. With batch=1024, total FLOPs = 549G. If you can pipeline them perfectly across all SMs, total time ~8ms. In practice, cuBLAS batched strided GEMM can be very efficient. But then you have two GEMMs (17-20ms total) + softmax overhead. Our custom kernel is 550 GFLOPs fused. If we match cuBLAS compute efficiency, we would take ~8-10ms. But achieving cuBLAS-level compute efficiency in a handwritten kernel is nearly impossible for these sizes. Alternative idea: can we use a single `torch.bmm` call? Not for the whole thing. Wait, is there a mathematical optimization for D > N? Since D=1024, N=512. The first GEMM is Q @ K^T: (N x D) @ (D x N) = (N x N). The second is (N x N) @ (N x D) = (N x D). Note that D is large. What if we compute attention as: O = V @ softmax(K^T @ Q) ? No. What if we explicitly compute K^T @ V first? (D x N) @ (N x D) = D x D. Then what? No. We could consider block-wise algorithms that exploit the fact that D is large by splitting D across thread blocks such that each thread block computes a partial dot product and a partial weighted sum, followed by a global reduction using atomics or a second kernel. This might allow K and V to be read only once per D tile, and Q to be read Q_tiles times. Let's formalize: - Partition D into D_TILE chunks (e.g., 64). - For each (batch, head, query), we need dot(Q_i, K_j) over full D. - We can have each block compute the dot product over one D_TILE for all queries and keys? - Actually, we can have blocks compute partial scores S^{(d)} for Q_tile x K_tile over a D_TILE. - Then we need to sum these partial scores across D_TILEs to get full S. - After getting full S, we need softmax, then weight V. - This is basically the classic "split-k" for attention. We could use two kernels: 1. Compute partial attention outputs for each D tile: for each (b,h, q_tile, k_tile, d_tile), compute: S_partial = Q[q_tile, d_tile] @ K[k_tile, d_tile]^T Then maybe compute the contribution to O: O_partial = softmax? But softmax requires full S. We could compute local max and sum per D tile and then do a multi-pass softmax reduction? That's very complex. Let's go back to the fused kernel idea, but try to make it compute-bound and efficient. We can use a large tile size for queries (Br) and keys (Bc) while tiling over D. The key insight is that cuBLAS's performance for FP32 might not be as amazing as I thought for these specific shapes because the leading dimensions are not large. But the huge batch count might help. Actually for BMM, cuBLAS dispatches to a single kernel that processes all batches. The efficiency is good. However, there is another path: PyTorch's `scaled_dot_product_attention` is called directly and likely maps to a single `cudnn` or `flash_attn` kernel. If it's the SDPA generic dispatch, it might be slower than raw cuBLAS BMM because of the dispatch overhead or because it chooses a suboptimal algorithm. If `cudnn` is available and chosen, it's probably very fast and hard to beat. What if we write a kernel that explicitly calls `cublasSgemmStridedBatched` via C++ extension? That wouldn't be faster than PyTorch's default which already does this. What if we write a very simple kernel that just implements the math ops but tries to fuse the softmax with the second GEMM by doing a custom loop? It likely won't beat cuBLAS. Wait, I should examine the possibility that the `embedding_dimension = 1024` is the total embedding dimension, not the per-head dimension. The reference code says: ```python Q = torch.rand(batch_size, num_heads, sequence_length, embedding_dimension) ``` In standard multi-head attention, the last dimension should be `embedding_dimension // num_heads = 1024 // 32 = 32`. But here they pass 1024 directly. That means each head is 1024-dimensional. This is unusual. In PyTorch's `scaled_dot_product_attention`, it treats the last dimension as the head dimension. So indeed head_dim = 1024. This unusual choice might be intentional to break optimized backends. Many optimized SDPA kernels (FlashAttention, xFormers memory-efficient attention) have head_dim limits (e.g., <=128). With head_dim=1024, they will likely fall back to the MATH backend. The MATH backend in PyTorch does: ```python attn_mask = None L, S = Q.size(-2), K.size(-2) scale_factor = 1 / math.sqrt(Q.size(-1)) attn_weight = Q @ K.transpose(-2, -1) * scale_factor attn_weight = torch.softmax(attn_weight, dim=-1) attn_weight = torch.dropout(attn_weight, dropout_p, train) output = attn_weight @ V ``` This calls `bmm` twice and `softmax`. For FP32, `bmm` calls cuBLAS strided batched GEMM. The intermediate `attn_weight` is materialized in HBM. Is it possible to write a better kernel? Maybe yes, if we fuse the entire thing. But as discussed, fusion might increase K/V reads. However, for MATH fallback, the two BMMs are `torch.matmul` for 4D tensors. PyTorch might call `at::native::bmm` which calls cuBLAS. The first BMM is Q @ K^T: `batch1` shape (1024, 512, 1024), `batch2` shape (1024, 1024, 512). This is `cublasGemmStridedBatchedEx` with transB. The second BMM is (1024, 512, 512) @ (1024, 512, 1024). cuBLAS is very efficient for strided batched GEMM when the batch size is large. But the first GEMM's K dimension is 1024, M=N=512. The second GEMM's M=512,N=1024,K=512. These are somewhat small inner dimensions. Another factor: the MATH fallback computes `attn_weight = Q @ K^T * scale_factor` which is elementwise multiply on a 1024x512x512 tensor (1.07GB). Then `softmax` over the last dim (512). Softmax on such a large tensor (1024*512*512 = 268M elements) involves reading/writing HBM. Then dropout is skipped (train=False). Then second BMM. Total HBM traffic: read Q(2GB), K(2GB) -> compute S -> write S(1GB). Read S -> compute softmax -> write P(1GB). Read P, V(2GB) -> compute O -> write O(2GB). Total = ~10GB if we count unique. But if we write a custom kernel that loads Q,K,V into registers/shared memory and fuses all ops, can we reduce this traffic? Yes: read Q,K,V once and write O once. That's 6GB. So we cut HBM traffic in half! The catch is that because D=1024 and N=512, the intermediate S and P are relatively small compared to Q,K,V. Wait, compute: Q,K,V each = B*H*N*D = 32*32*512*1024 = 536,870,912 elements * 4 bytes = 2,147,483,648 bytes each = 2GB. So three of them = 6GB. S,P each = B*H*N*N = 32*32*512*512 = 268,435,456 elements * 4 bytes = 1,073,741,824 bytes each = 1GB. So S+P = 2GB. Total unfused: 6GB + 2GB = 8GB reads/writes (assuming reads+writes counted). Actually reads: Q+K+V+S+P = 6GB? Wait, in unfused: 1. Q,K read (4GB), S written (1GB). Total 5GB. 2. S read (1GB), P written (1GB). Total 2GB. 3. P,V read (3GB), O written (2GB). Total 5GB. Sum = 12GB. Fused: read Q,K,V (6GB), write O (2GB). Total = 8GB. So fused saves 4GB (33% reduction). That's not huge but could help. But as noted earlier, if we tile along queries and each tile reads K and V again, we might not save. However if we parallelize over queries and D dimension such that K and V are not re-read... Wait, if a kernel block computes a subset of queries but needs all keys for softmax, it must read all K and V for those keys. So any block that computes a query subset reads K and V. The only way to read K and V exactly once overall is if there is one block per head or less. But then parallelism is limited. However, we can have one block per head but multiple warps within the block computing different queries and features. A single block has up to 1024 threads. To process 512 queries and 1024 features, we need 512*1024 = 524,288 dot product accumulators? No. Let's think of mapping: - BlockDim = (threads). Each thread can compute one complete output element O[b,h,q,d]. To compute O[b,h,q,d], a thread needs: - Q[b,h,q,d] (float) - K[b,h,:,d] (512 floats) - V[b,h,:,d] (512 floats) (after softmax weights) But to compute softmax weight for key j, the thread needs dot(Q_q, K_j) which depends on all d. So a thread computing O[b,h,q,d] cannot compute softmax locally unless it knows all scores. So a thread block must compute all scores for its queries, then softmax, then compute all O for its queries. If one block per head handles all queries, it needs to compute the N x N score matrix (512x512 = 262k floats). This fits in registers/shared memory across threads (e.g., 256 threads * 1024 regs = 256k regs, but 262k floats * 4 bytes = 1MB. With 255 regs/thread and 1024 threads, we have 260k 32-bit regs, just enough for the score matrix but not much else). Actually 1024 threads * 255 registers = 261,120 registers. 512*512 = 262,144 floats. That's exactly one register per float. That's too tight and would leave no registers for other variables. Compiler would spill, causing terrible performance. We can use smaller number of threads and store S in shared memory. 512x512 = 262k floats = 1MB. A block has up to 48KB or with dynamic shared maybe 227KB max. So 1MB doesn't fit in shared memory. Therefore a single block cannot hold the entire S matrix in shared memory. It could compute S tile by tile and accumulate O. So the typical FlashAttention approach is necessary. But maybe we can choose a different parallelization: instead of one block per head, we can have blocks that handle a range of keys and do a reduction. This is similar to the "split-K" but with softmax. However, implementing the softmax reduction across blocks correctly is possible and not too complex if we use a two-pass or tree reduction. Specifically: - Partition keys into K_tiles (e.g., NK_blocks = 8). - Each block computes for all queries (or a query tile) and its key tile: - S_tile = Q[tile_q, :] @ K[key_tile, :]^T - local_max = max_j in key_tile S_tile - local_sum = sum_j exp(S_tile - local_max) - local_O = sum_j exp(S_tile - local_max) * V[key_tile, :] - Then we need to combine across key tiles for each query. We can use an atomic or a second reduction kernel. Paper: "Online normalizer calculation for softmax" (Milakov et al.) and "FlashAttention-2". The online softmax trick allows combining two chunks: Given chunk1 with m1, l1, o1; chunk2 with m2, l2, o2. m_new = max(m1, m2) l_new = exp(m1 - m_new) * l1 + exp(m2 - m_new) * l2 o_new = exp(m1 - m_new) * o1 + exp(m2 - m_new) * o2 So blocks can stream over key tiles and maintain running m, l, o. Within a single block, we can have it loop over key tiles. But the block still reads K and V for each key tile. If a block handles all queries, and loops over all key tiles, then K and V are read once. But the block would be very large (needs to handle 512 queries and 1024 features). Wait, a block can handle all queries but loop over D and key tiles. Let's design this: - Block size: 256 threads. - Each block handles one head (b,h). - The 256 threads cooperate to compute all 512 queries. - They process queries in groups of, say, 8. For each query group (8 queries): - Load Q_group into smem (8 * D = 8192 floats = 32KB). D=1024 is too large to load all at once into 32KB? 8*1024*4 = 32KB. That fits in shared memory (max 227KB). Good. Actually we can load the full 8 queries into smem. - We then iterate over key tiles of size NK_tile (say 64 keys). - For each key tile, we need K tile (64 * D floats = 256KB) -> too big. - So we also tile over D: D_TILE = 64. - Process D tiles: for each D tile: - Load Q_group segment (8 x 64 = 512 floats = 2KB) - Load K_tile segment (64 x 64 = 4096 floats = 16KB) - Compute partial S (8 x 64) accumulating in registers. - After all D tiles, we have S (8x64). Do online softmax update for these 8 queries with this key tile. - Then for D tiles of V: - Load V_tile segment (64 x 64 = 16KB) - For each of the 8 queries, compute weighted sum using P weights (8x64). Update O_group segment. - Write O_group for all D. This is classic FlashAttention but with one block per head, handling all queries in subgroups. Number of blocks = 1024. With 132 SMs, that's ~8 blocks per SM. Each block does a lot of work: loops over 8 query groups, 8 key tiles, 16 D tiles. The work is highly compute-bound? Let's compute. Per block FLOPs: N^2 * D = 512*512*1024 = 268M FMAs (for one head). That's the same as the two GEMMs combined (since QK and PV are combined). Actually total FLOPs per head for standard attention = 2*N^2*D = 536M FMAs (1 FMA per multiply-add). Wait, QK^T: N*N*D FMAs. PV: N*N*D FMAs. So 2*N^2*D = 536M FMAs per head. Our fused block does the same 536M FMAs. At 67 TFLOPS for the whole GPU, one block at full speed would take 536M / 67T = 0.008ms. But 1024 blocks would take ~8ms sequentially if they share SMs perfectly (which they don't). If 132 SMs each do ~8 blocks, total time ~8 * 0.008 = 0.064ms? No, blocks are parallel. 132 SMs * throughput. Let's just say total FLOPs across all heads = 1024 * 536M = 549B FMAs. At 67 TFLOPS = 8.2ms. Now, does each block keep K and V in L2? With 1024 blocks, each head's K+V is 4MB. L2 is 50MB. So at most ~12 heads' data stays in L2. But if each block loops over key tiles and loads K/V sequentially, the data might be cached in L2 for the duration of the block. Since a block takes some time, other blocks may evict it. But the active set size per block is Q (32KB) + K/V tiles (16KB each) = 64KB. L2 can hold many of these tiles. The key issue is HBM traffic for K and V. If 1024 blocks each read K and V once, total K+V reads = 1024 * 4MB = 4GB. Plus Q reads = 2GB, O writes = 2GB. Total 8GB. That's less than the 12GB of the math fallback! Wait, why is it 4GB and not 8GB? K size per head = 2MB. 1024 heads -> 2GB. V same -> 2GB. Total K+V = 4GB. Q = 2GB, O = 2GB. Sum = 8GB. This matches the fused "read once" ideal. The trick is that each block reads K and V for exactly one head. Since there's one block per head, K and V are read exactly once per head. This is the minimal HBM traffic! The previous traffic analysis for the fused kernel assumed multiple query blocks per head, which caused K/V re-reads. With one block per head, we avoid re-reads. So the key design choice: **One thread block per head**. Is one block per head feasible given the work per head (536M FMAs)? Yes. With 256 threads, each thread does ~2M FMAs. That's a lot of work per thread but loops handle it. Register pressure and shared memory usage need careful management. Shared memory: - Q_group: we can process e.g., 8 queries at a time. Q_group segment per D_TILE: 8*64 floats = 512 floats = 2KB. - K_tile segment: BK * D_TILE. Let's choose BK=64, D_TILE=64: 64*64 = 4096 floats = 16KB. - V_tile segment: same = 16KB. - S matrix for 8 queries x 64 keys: 512 floats. Can be in registers or smem. We might keep it in registers. Total smem: Q(2) + K(16) + V(16) = 34KB + some for reduction. That's fine under 48KB default, but we can request more if needed (up to 227KB). Should we use a `__shared__` buffer that is reused for K and V to reduce smem usage? Yes, we can overlap. Smem max needed at one time = max(Q+K, Q+V, ...). Q is 2KB, K is 16KB, V is 16KB. If we load K, compute S, then load V, compute O, we can reuse the same buffer for K and V. So smem = Q(2) + KV(16) = 18KB. Plus maybe a workspace for reduction (e.g., smem for per-thread max/sum reduction). Let's keep it simple. But soft max requires max over the 64 keys for each query. We can compute max locally per thread and then warp reduce. Let's map threads: - Thread block = 256 threads. - Warp size = 32. - We have 8 warps. Tasks: 1. Compute S_t = Q_group (BQ=8) x K_tile (BK=64) for the current D_TILE. - Each dot product is over D_TILE=64. We can have each thread compute one or more (q,k) pairs. - Total pairs = 8*64 = 512. With 256 threads, each thread computes 2 pairs. That's great. - D loop over 16 tiles. For each tile, we load Q_group[8][64] and K_tile[64][64] into smem. - Then each thread computes its two partial dot products by looping over d=0..63 and accumulating in registers. - After 16 D tiles, we have full S_t (8x64). Each thread has 2 accumulators. 2. Online softmax update: - For each query q in 0..7: - Find max over k in 0..63. Since there are 64 keys, and each thread holds values for 2 keys, we need warp reduce for max. - Actually 8 queries x 64 keys. We can have the 256 threads arranged as (q, k) pairs. Wait, 256 threads = 8 warps * 32 threads. - Let's assign threadIdx.x % 64 to k index, and threadIdx.x / 64 to q index? But 64*8=512, not 256. - Let's assign threadIdx.x % 32 to some lanes, and warps to queries and key groups. - Simplest: Use 8 warps. Each warp handles one query (q). Within a warp (32 threads), each thread handles 2 keys (since BK=64). So warp q handles keys [0..63], with thread t handling keys t and t+32. - We can then do warp-level reduction for max and sum across the 64 keys. But warp size is 32, so to reduce 64 elements, we can have each thread know two values. - Actually, for query q, we have 64 score values distributed across the warp (2 per thread). To get max across 64, each thread can compute local max of its two values. Then warp reduce (`__shfl_xor_sync`) over 32 lanes. But this gives max across 32 threads' local max, not across the 64 values if the two values per thread are treated as a local max. Wait, each thread has two scores: s0 and s1. Thread computes m_local = max(s0, s1). Then warp reduces m_local to get global max. This correctly finds max across all 64 values because each thread contributes its max, and warp reduce takes max of all contributions. - Similarly for sum of exp(s - m_global): each thread computes sum_local = exp(s0 - m_global) + exp(s1 - m_global). Then warp reduce sum. - So we can partition threads as: 8 queries * 32 threads = 256 threads. Perfect! - Thus: - warp_id = tid / 32 (0..7) -> query index within group (0..BQ-1). - lane_id = tid % 32. - This warp computes scores for query `q = warp_id` and keys `k = lane_id` and `k = lane_id + 32`. Wait, lane_id 0..31. Keys handled: lane_id and lane_id + 32 => 64 keys total. - This mapping is beautiful. 3. Accumulating output O: - After computing S and doing online softmax, we have weights P = exp(S - m) / l (but we maintain O incrementally). - We need to update the output for query q. - The running O_q is a vector of length D=1024. We can accumulate O_q in registers across key tiles? D=1024. If we keep the full O_q in registers for 8 queries, that's 8*1024 = 8192 floats. Too many registers. So we must tile over D for O accumulation as well. - We can process key tiles in outer loop. For each key tile, after computing P weights (64 per query), we need to update O_q for all D. We can process D in tiles of D_TILE=64. - For each D tile d: - Load V_tile of keys in current key tile and D tile into smem: shape BK x D_TILE = 64 x 64 = 4096 floats = 16KB. - For query q, O_q[d] += sum_{k in tile} P_{q,k} * V_{k,d}. - Distribution among threads: 256 threads. We need to compute 8 queries * 64 features = 512 updates. Each update is a sum over 64 keys. - Wait, O_q[d] for one D tile is 64 values. For 8 queries, 512 dot products. Each dot product is sum over BK=64 keys. - We can map each thread to one (q,d) pair. 256 threads => each thread handles 2 (q,d) pairs. For each pair, sum over 64 keys: P_{q,k} * V_{k,d}. P is stored per thread? In the previous step, warp q had thread lane handling 2 keys. So thread has P_{q,k0} and P_{q,k1}. To compute sum over 64 keys, each thread can compute a partial sum of its two keys, then warp reduce. This is perfect! - Wait: if each thread handles two keys (k0, k1), then for a given (q,d), thread computes p0*v_{k0,d} + p1*v_{k1,d}. Then warp reduce across 32 lanes to sum over all 64 keys. This gives O_q[d]. - So for O accumulation: - We iterate over d in D tile. - We can have threads compute for their assigned (q,d) pairs. But 8 queries * 64 d = 512 pairs, 256 threads => 2 pairs per thread. Thread t (tid) handles q = (tid / 64) ??? Let's define mapping: - We have 8 queries. Each query needs 64 outputs. Total 512 outputs. - Let's assign tid % 64 = d_index in [0,63]. tid / 64 = q_index in [0,7]? But 8*64=512 > 256. So we need to split queries and features differently. - Alternative: use a 2D thread mapping. 256 threads = 8 warps * 32 threads. We already used warps for queries. For O accumulation, we can keep warp for query q. Within warp, each lane handles 2 d values: d = lane_id and d = lane_id + 32. This is symmetric to the key mapping. - So lane_id handles d0 and d1. Then for each d, we need to sum over 64 keys using the P values we have. - Thread computes local_sum_d0 = P_q,k0 * V_k0,d0 + P_q,k1 * V_k1,d0. Then warp reduce (sum) over 32 lanes to get O_q,d0. Same for d1. - But we also need to iterate over D tiles. The D tile is 64 values. With lanes handling 2 values each, we cover all 64. Great. - This means we don't need to store the P matrix in smem or registers beyond the two P values per thread! We just compute them, do max/sum reduction (updating online softmax stats), keep the two P values in registers, and then immediately use them for the V accumulation across D tiles. - To do online softmax correctly across key tiles, we need to maintain running max m and running sum l (or running O). Pseudo: - Initialize m = -inf, l = 0, O = 0 for all d. - For each key tile (with 64 keys): - Compute scores S_qk for this tile (full dot products over D=1024). - Let m_local = max_k S_qk. - m_new = max(m, m_local) - For each key k in tile: - P_qk = exp(S_qk - m_new) [but we need this for O] - l_new = exp(m - m_new) * l + sum_k exp(S_qk - m_new) - For each d: - O_new[d] = exp(m - m_new) * O[d] + sum_k P_qk * V_k,d - Set m = m_new, l = l_new, O = O_new. - After all key tiles, O_final = O / l? Actually if we keep O as the weighted sum of V with unnormalized weights, then at the end O = O / l. Or we can divide at the end. - Wait, the standard online softmax formula: - We want O = sum_k exp(S_k - m_global) * V_k / sum_k exp(S_k - m_global) - Maintain m (max so far), l (sum of exp with respect to current m), and o (weighted sum with respect to current m). - When new tile arrives with local m_tile, l_tile, o_tile (computed with m_tile as local max): - m_new = max(m, m_tile) - o_new = exp(m - m_new) * o + exp(m_tile - m_new) * o_tile - l_new = exp(m - m_new) * l + exp(m_tile - m_new) * l_tile - Here o_tile = sum_{k in tile} exp(S_k - m_tile) * V_k. - l_tile = sum_{k in tile} exp(S_k - m_tile). - We can compute S_k, then m_tile, then compute exp(S_k - m_tile) directly. - At the end, O_final = o / l. - This algorithm requires storing o for all D dimensions. Since D=1024, 8 queries * 1024 = 8192 floats. This is 8192 * 4 = 32KB. Can we store this in registers? 256 threads * 32 registers = 8192 registers. That's exactly enough if we use one register per float. But we also need other registers. However we can tile over D for the O accumulation too! - Wait, we already tile over D for V. The outermost loop is over key tiles. We need to update O for all D. But we process D tiles sequentially. - Let's restructure: - Outer loop over D tiles (DT=64). There are 16 D tiles. - For each D tile, we process all key tiles to compute the contribution of this D tile to the scores and to O? No, the score computation requires the full D. So we must compute scores over all D before softmax. - So we can't tile D in the outermost loop. The loop order must be: 1. Loop over key tiles to compute S (requires full D). During this, maintain online softmax stats m and l (per query). We also need to keep P weights for each key in the tile to later weight V. But we can process one key tile at a time and for each key tile, after computing the scores and softmax update, we immediately accumulate V for that key tile across all D tiles. This way, we don't need to store P for all keys, just for the current tile. - So algorithm: - For query group of 8 queries: - Initialize m[8] = -inf, l[8] = 0. - Initialize O[8][D] = 0. But storing all D is too much. So we must process D tiles in the innermost loop after computing P for the current key tile. - For each key tile (BK=64): - Compute S_tile (8x64) using full D=1024 (loop over D_TILE=64). - Compute m_tile, l_tile, and P_tile weights (8x64). Each warp handles one query; 32 lanes handle 2 keys each. Thread computes P_qk0, P_qk1 = exp(S_qk - m_new). Actually: - First, compute local max per thread (max of its two S). Warp reduce to get m_tile. - Then compute exp(S - m_tile) for each key. - Compute local sum per thread (sum of its two exp). Warp reduce to get l_tile. - m_new = max(m_prev, m_tile). - correction_prev = exp(m_prev - m_new) - correction_tile = exp(m_tile - m_new) - l_new = correction_prev * l_prev + correction_tile * l_tile - Now for each key, we need the actual weight to multiply V. The weight for key k is exp(S_qk - m_new) / l_new? Wait, for the next step we need to update o. The update is: o_new = correction_prev * o_prev + correction_tile * o_tile where o_tile = sum_k (exp(S_qk - m_tile) / ???). Let's derive carefully. We want o_new to be sum over all processed keys of exp(S_k - m_new) * V_k. So o_new = sum_{keys in prev tiles} exp(S_k - m_new) * V_k + sum_{keys in tile} exp(S_k - m_new) * V_k = exp(m_prev - m_new) * sum_{prev} exp(S_k - m_prev) * V_k + sum_{tile} exp(S_k - m_new) * V_k = exp(m_prev - m_new) * o_prev + sum_{tile} exp(S_k - m_new) * V_k. Here o_prev = sum_{prev} exp(S_k - m_prev) * V_k. Similarly, l_new = exp(m_prev - m_new) * l_prev + sum_{tile} exp(S_k - m_new). So we need to compute `weight = exp(S_qk - m_new)` for each key in the tile. Note that `exp(S_qk - m_new) = exp(S_qk - m_tile) * exp(m_tile - m_new)`. We have exp(S_qk - m_tile) already (from l_tile computation). And exp(m_tile - m_new) is a scalar per query. So `weight_qk = exp_val_qk * exp(m_tile - m_new)`. But since correction_tile = exp(m_tile - m_new) is a scalar per query (same for all k in tile), we can just scale the local sums by it. To update o_prev for all D, we do: o_new[d] = correction_prev * o_prev[d] + correction_tile * sum_{k in tile} exp(S_qk - m_tile) * V_k[d] The inner sum is over keys in the tile. For each D tile, we can compute this sum using the P local values. - So after computing m_new, l_new, we know correction_prev and correction_tile. - Now we need to update o_prev[d] for all d in D. - We can process D in tiles of DT=64. - For each D tile: - Load V_tile[64][64] into smem. - For query q (warp), and d handled by lane (two d's): - Compute local_sum_d0 = sum_{k in my_keys} exp(S_qk - m_tile) * V_k,d0 - Warp reduce (sum) local_sum_d0 across 32 lanes -> tile_sum_d0. - Then o_prev_d0 = correction_prev * o_prev_d0 + correction_tile * tile_sum_d0. - Store updated o_prev_d0 in some accumulator. What memory to use for o_prev across key tiles? Since there are 16 D tiles and 8 queries, total 8192 floats. We can store them in shared memory! Shared memory size: 8192 floats = 32KB. Plus K/V tile (16KB) and Q tile (2KB). Total 50KB. That's within 48KB? Actually slightly above. But we can use the max shared memory per block of 48KB by default. If we need a bit more, we can request up to 227KB via `cudaFuncSetAttribute`. But let's try to fit within 48KB. - 32KB for O accumulator + 18KB for Q,K/V = 50KB. Too much for default 48KB. We can reduce DT or BQ or BK. - Let's adjust parameters to fit in 48KB: - BQ=8 queries per group. - BK=32 keys per tile. - DT=64. - O accumulator: BQ * D_DT floats = 8 * 64 = 512 floats = 2KB per D tile? Wait, we process D tile by tile. We only need to hold the accumulator for the current D tile across key tiles. So O accumulator per block = BQ * DT = 8*64 = 512 floats = 2KB! That's tiny. We don't need to hold all D at once. - Yes! Because we loop over D tiles in the inner loop after computing P. Wait, can we loop over D tiles before moving to the next key tile? No: the key tile loop must happen first to compute S. But after S is computed for the key tile, we can loop over all D tiles to update O for that key tile. Then we move to the next key tile. So we only need O accumulator for one D tile at a time. - But we need the running O accumulator to persist across key tiles for each D tile. If we process D tile 0 for key tile 0, then key tile 1, then D tile 0 again, the accumulator for D tile 0 must persist. So we can either: a. Loop over key tiles outermost, and within each key tile, loop over all D tiles. This requires O accumulator for all D simultaneously while processing key tiles, OR b. Loop over D tiles outermost, and within each D tile, loop over key tiles. But to compute S for a key tile, we need the full D. However S is independent of D tile? No, S = sum over full D. If we process D tiles outermost, we compute partial S contributions per D tile, accumulate them into a temporary S array, and after finishing all D tiles for a key tile, do softmax. Then to accumulate O, we process D tiles again: for each D tile, accumulate V weighted by P. The O accumulator must persist across key tiles. So if we process all key tiles for D tile 0 first, then move to D tile 1, we need to store the partial O for D tile 0. That partial O must persist while we process all other D tiles for all key tiles. That's effectively storing O for all D. c. Alternatively, process key tile outermost, and within each key tile, process D tiles for O update. To persist O across key tiles, we can store O in shared memory or global memory. Let's size it: - BQ=8, D=1024. O = 8*1024 floats = 32KB. This fits in shared memory! Yes. 32KB is within 48KB. Plus K/V and Q tiles (18KB) gives 50KB. Slightly over 48KB but we can reduce tile sizes. d. We can write O to global memory after each key tile and read it back for the next key tile? That defeats the purpose. e. Better: store O in global memory but keep it in L2? Not helpful. f. Reduce Q tile and K/V tile sizes. Let's try to fit in 48KB. - Let's compute smem needed if we store full O for BQ=8, D=1024: - O_acc: 8*1024*4 = 32,768 bytes = 32KB. - Q_tile: BQ*DT*4 = 8*64*4 = 2KB. - K_tile: BK*DT*4 = let's say BK=32, DT=64 -> 32*64*4 = 8KB. - V_tile: same = 8KB. - Total = 32 + 2 + 8 + 8 = 50KB. Ugh, 2KB over. - If we set BK=16: - K_tile = 16*64*4 = 4KB. V_tile = 4KB. - Total = 32 + 2 + 4 + 4 = 42KB. Good! Fits in 48KB. - But BK=16 means more key tiles (512/16 = 32). That's okay. - If we set DT=32: - Q_tile = 8*32*4 = 1KB. - K_tile = 32*32*4 = 4KB (BK=32). - V_tile = 4KB. - O_acc = 8*1024*4 = 32KB. - Total = 32 + 1 + 4 + 4 = 41KB. Good. - But DT=32 means more D tiles (32). That's okay. - Let's choose DT=64, BK=16, BQ=8. - O_acc = 32KB. - Q_seg = 8*64*4 = 2KB. - K_seg = 16*64*4 = 4KB. - V_seg = 4KB. - Total = 42KB + some small arrays (max/sum per query) = maybe 43KB. Under 48KB. Perfect. - Or we can keep BK=32 and DT=32: - O_acc = 32KB. - Q_seg = 8*32*4 = 1KB. - K_seg = 32*32*4 = 4KB. - V_seg = 4KB. - Total = 41KB. Also fine. - Which is better? Larger DT means fewer loops. BK=32 means fewer key tiles. - Let's try DT=64, BK=16. - D tiles = 1024/64 = 16. - Key tiles = 512/16 = 32. - Number of key tile loops = 32. - Within each key tile: - Compute S: loop over 16 D tiles. Each iteration: load Q_seg (2KB) and K_seg (4KB), compute partial S for 8 queries x 16 keys. - After S done: softmax update. - Update O for this key tile: loop over 16 D tiles. Each iteration: load V_seg (4KB), compute weighted sum for all 8 queries x 64 d values. - Notice Q_seg is loaded 16 times during S, and 16 times during O? Actually for O update, Q is not loaded. O is in smem. - V_seg is loaded 16 times per key tile. - K_seg is loaded 16 times per key tile. - Q_seg is loaded 16 times per key tile. - Total smem loads per key tile = 16 * (Q+K+V) = 16*10KB = 160KB of smem traffic. Over 32 key tiles = 5.12MB of smem traffic per block. - HBM traffic per head: Q read once (2MB), K read once (2MB), V read once (2MB), O read/write? We accumulate O in smem and write once (2MB). Plus scaling writes. Total HBM traffic ~8MB per head. Great. - Register usage: each thread handles 2 keys. For S computation, accumulators for 2 floats. For P computation, 2 floats. For O accumulation, thread handles 2 d values. But wait, O accumulation requires sum over 16 keys. Each thread has P for 2 keys. It can compute local sum for its 2 keys, then warp reduce. Then update O accumulator for its 2 d values. - Wait, in our mapping, warp_id handles one query q. Lane_id handles keys: with BK=16, lane_id 0..15 handles two keys? No, 32 lanes and 16 keys: each lane handles 0.5 key? Not integer. Let's instead map each warp to a query. Within warp (32 threads), we need to cover 16 keys and 64 features. - Let's redesign thread mapping for BK=16, BQ=8, DT=64. - Total threads = 256. - We have 8 warps. Each warp handles one query. - For S computation: there are 16 keys. 32 threads in warp. We can assign 2 threads per key? Or each thread handles one key (only 16 threads active) and compute dot product with Q. But dot product over DT=64; each thread can loop over d and accumulate. With 16 threads active, each does 64 FMAs per D tile. That's fine. We can have 16 threads compute for 16 keys, the other 16 threads idle. But then warp reduction for max/sum only needs 16 lanes. - Alternatively, we can assign 16 keys to 32 threads by having each thread handle half of D? No, we need the full dot product per key. A thread can compute partial sum over a subset of D and then shuffle reduce. But D is 1024 and we have 16 keys. If we use 32 threads, each thread can compute partial dot product for one key over D_TILE=64, then after 16 D tiles, we have the full dot product. With 32 threads and 16 keys, we can assign thread 0 and 16 to key 0, thread 1 and 17 to key 1, etc. Each thread computes partial dot over its assigned half of D_TILE (32 elements). Then warp reduce sum to get the full dot product for the key. This utilizes all 32 threads. - For O accumulation: we need to compute for 64 features (d). With 32 threads, each thread can handle 2 features. For each feature d, we need weighted sum over 16 keys. We can assign 16 threads to handle the keys for feature d0, and another 16 for d1? But each thread only has P for one or two keys. Let's think. - We have 16 keys, and P weights per key. The weights are distributed across the warp. If we use the same mapping as S (two threads per key), then each pair of threads has the same P for that key. To compute sum over keys of P_k * V_k,d, we can have each thread multiply its P by V_k,d and then warp reduce sum. Since two threads have the same P and V (for the same key), they will compute the same product. That's wasteful but simple. - Alternatively, reassign threads for O accumulation: 32 threads, each handles 2 d's. For each d assigned to thread t (d0=t, d1=t+32), the thread computes local_sum over its keys. But how many keys does each thread own? If we assign 2 keys per thread, then 32 threads cover 64 keys, but we only have 16 keys. So maybe not. - Let's instead use a different block shape for O accumulation. We can use a separate loop where threads are mapped to (d) features: 64 features, 256 threads = 8 warps * 32 threads. Since there are 8 queries (one per warp), and 64 features. Each warp has 32 threads, but only 64 features per query. We can assign 2 features per thread within the warp: thread t in warp q handles d = t and d = t+32? No, d ranges 0..63 for a D tile. 32 threads * 2 = 64. So each thread handles two features. Good. - For each feature d assigned to thread t: - The thread needs to compute sum_{k=0..15} P_{q,k} * V_{k,d}. - Where does thread t get P_{q,k}? It needs the P weights for all 16 keys. We can store the P weights for the tile in shared memory: 8 queries * 16 keys = 128 floats = 512 bytes. After computing S and softmax, we write P to smem. Then all threads can read P from smem for the O accumulation. - This is very clean! After computing S and the online softmax update, we compute the normalized weights `P = exp(S - m_new)`. Wait, for updating O we need `exp(S - m_new)` divided by nothing yet? Let's re-express. - We maintain per-query: m (max so far), l (sum so far), O[d] (weighted sum so far). - When new key tile arrives, compute S_tile. - m_tile = max_k S_tile. - m_new = max(m, m_tile). - exp_shift = exp(m - m_new). - For each k: p_raw = exp(S_k - m_new). // This is directly the weight to use for V Actually: sum_{k} exp(S_k - m_new) = sum_{k} exp(S_k - m_tile) * exp(m_tile - m_new). If we compute p_raw = exp(S_k - m_new), then l_new = exp_shift * l + sum p_raw. And O_new[d] = exp_shift * O[d] + sum_k p_raw * V_k,d. - So we just need `p_raw` for each key. - We can compute p_raw per key and store in smem (8x16). - Then for O accumulation, threads load p_raw from smem and V from smem. - This requires storing P_tile in smem: 128 floats = 512 bytes. Very little. - Thread mapping for O: - warp q, thread t in warp handles d = t and d = t+32? Wait, D tile is 64. Warp has 32 threads. So thread t handles d = t and d = t+32? That gives d in [0,31] and [32,63]. Perfect. - For each of its two d values, thread loads P_k for k=0..15 from smem? That would be 16 loads per d, 32 loads per thread. Instead, thread can compute: sum0 = 0, sum1 = 0. for k=0..15: p = P_smem[q][k]; v0 = V_smem[k][d0]; v1 = V_smem[k][d1]; sum0 += p * v0; sum1 += p * v1; - This is a loop of 16 iterations, each loading p (broadcast from smem) and two v's. This is efficient if V is in smem and P is in smem. - Then O_new[d0] = exp_shift * O_old[d0] + sum0. O_old is in a register or smem. - We can keep O_old for the current D tile in a register array per thread: 2 floats. Or in smem. - Since there are 8 warps * 32 threads * 2 floats = 512 floats = 2KB. This can be in registers easily. - This mapping is super clean. - Let's define the kernel precisely. - Grid: dim3(batch_size, num_heads) = (32, 32). Total 1024 blocks. - Block: 256 threads. - Each block processes one head. - Inside block: - Load Q for the head into some temporary? Q is 512x1024 = 2MB. We cannot load all at once. We will process queries in groups of BQ=8. There are 512/8 = 64 query groups. - But wait, if we process queries in groups, then the O accumulator in smem is only for 8 queries. We process the first 8 queries through all key tiles, write O to global memory, then load next 8 queries. - This means K and V are read 64 times per head (once per query group). But K and V are in L2 after the first group? Not guaranteed. However we can hope for L2 reuse? Actually 2MB K+V, L2 is 50MB, so it can easily hold K+V for the head across query groups. If the block executes query groups sequentially, K/V loads will hit L2. But H100 has L2 bandwidth around 10-20 TB/s, which is high. However, we still have 64 groups. K/V read per group = 2MB. Total K/V traffic to L2 = 64*2MB = 128MB per head. HBM traffic per head for K/V might still be just 2MB each if L2 holds them. But does L2 retain them? Yes, 4MB fits easily. HBM reads: 2MB for K, 2MB for V. Then L2 serves the rest. So HBM traffic is minimal. - This is a key insight! With one block per head, processing query groups sequentially, K and V are loaded from HBM once and stay in L2 (or even L1 if we try) for the entire head. So HBM traffic is close to ideal: Q (2MB), K (2MB), V (2MB), O (2MB) = 8MB per head. Total 8GB. - L2 traffic will be higher but L2 is fast. - This might be a great approach! - But is reading K/V from L2 64 times fast enough? L2 bandwidth is ~20 TB/s. 128MB per head * 1024 heads = 131GB of L2 traffic. At 20 TB/s = 6.5ms. Plus compute. Plus HBM traffic (8GB = 2.4ms). Total maybe 10-15ms. Might be competitive. - However, 64 query groups means an inner loop of 64 iterations, each with 32 key tiles. That's 2048 iterations of the core loop. Each iteration loads Q_seg and K_seg/V_seg, computes partials. This is a lot of instructions but doable. - Wait, there are 256 threads. Do we have enough registers? Let's estimate. - Each thread in the S loop accumulates a float for its key (if mapping is one key per thread pair). If we have 32 threads per warp and 16 keys, two threads per key. Each thread accumulates partial S (half of D_TILE). After 16 D tiles, it reduces with its pair. So each thread needs 1 float accumulator. After reduction, we have 16 floats per warp (one per key). We can store them in smem P_smem (16 floats per warp). - Actually, if we compute dot products across D_TILE=64 with two threads per key, each thread does sum over 32 floats. Then `__shfl_xor_sync` to sum the pair. Then store in smem. - For O update: Each thread handles 2 d values (d0=t, d1=t+32). It needs two accumulators `sum0, sum1`. It loops over 16 keys. For each key, it loads P_smem[q][k] (shared), loads V_smem[k][d0], V_smem[k][d1] (shared), does 2 FMAs. So registers needed: sum0, sum1, plus loop counter, plus maybe some pointers. Very low. - The O persistent values (the final O accumulator across key tiles) for the two d's per thread: `o0, o1`. So total registers per thread: ~10-15 floats. Very manageable. - The main issue is instruction count and loop overhead. - Wait, I just realized: with 1024 blocks total and 132 SMs, each SM processes ~8 blocks. Each block has 256 threads. So total resident threads = 132 * (some number). At max 32 blocks per SM, but limited by registers. If we use low registers, maybe 2-4 blocks per SM concurrently. That gives 300-500 blocks concurrently across the GPU. 1024 blocks is roughly 2-3 waves. That's fine. - But each block does a lot of work (64 query groups * 32 key tiles * 16 D tiles). The execution time per block might be high. If the number of waves is low, we might not fully utilize the GPU's compute. However, each block has 256 threads doing work. Even with few waves, if compute is dense, utilization might be okay. - Wait, 256 threads * 2-3 waves * 132 SMs = plenty of warps for latency hiding. Let's refine the kernel design. Kernel parameters: - BQ = 8 (queries per group) - BK = 16 (keys per tile) - DT = 64 (feature tile size) - Threads per block = 256 - Grid = (batch_size, num_heads) Mapping: - blockIdx.x = batch, blockIdx.y = head. - threadIdx.x 0..255. - warp_id = threadIdx.x / 32 (0..7) -> maps to query within group: q = warp_id. - lane_id = threadIdx.x % 32. Memory layout: - Global memory tensors are in BHSD format: Q, K, V, O are [B, H, S, D] (here S=seq_len=512, D=1024). - For block (b,h): - base offset = b * H * N * D + h * N * D. - Q_ptr = Q + base_offset. - K_ptr = K + base_offset. - V_ptr = V + base_offset. - O_ptr = O + base_offset. Query group loop: - for qg = 0; qg < N / BQ; qg++: - Query indices: q = qg * BQ + warp_id. - Load Q for these queries: For each d_tile in 0..(D/DT-1): - Load Q[q][d_tile*DT + tid_local] into smem or registers? We need Q for S computation. But S computation iterates over d_tiles. We can load Q for the current d_tile on demand. - Actually we process key tiles outermost? Let's decide loop order. Loop order Option A: for qg in query_groups: for kt in key_tiles: // compute S for this key tile for this query group for dt in d_tiles: load Q[qg][dt] into smem load K[kt][dt] into smem for each key in tile: compute partial dot product // after dt loop, we have full S (8x16) // compute softmax update // update O for all d_tiles for dt in d_tiles: load V[kt][dt] into smem update O_acc for this dt This is exactly as described. Each iteration of the inner `dt` loop for V uses the P weights computed for the key tile. S computation details: - dt loop: 16 iterations. - Shared mem: - float q_smem[BQ][DT]; // 8*64 = 512 floats = 2KB - float k_smem[BK][DT]; // 16*64 = 1024 floats = 4KB - Total for S stage: 6KB. - Thread mapping for dot products: - warp_id = q (0..7) - lane_id = 0..31 - We have BK=16 keys. Assign two threads per key: key k is handled by lane_id = k and lane_id = k+16. - Each thread computes partial_sum = sum_{i=0}^{DT/2 - 1} Q[q][dt*DT + i_offset + i] * K[k][dt*DT + i_offset + i]. - Wait, if two threads split DT=64, each handles 32 elements. - thread A (lane k): sums over d = dt*DT + 0 .. 31 - thread B (lane k+16): sums over d = dt*DT + 32 .. 63 - After loop, thread A gets partial_sum_A, thread B gets partial_sum_B. - They need to add them: `partial_sum = partial_sum_A + __shfl_xor_sync(0xFFFFFFFF, partial_sum, 16)`? No, we need to exchange between lane k and lane k+16. - Actually we can use `__shfl_sync` or `__shfl_xor_sync` with mask 0xFFFFFFFF and lane 16. Let's say val = partial_sum. val = val + __shfl_down_sync(0xFFFFFFFF, val, 16); Then after this, lane k (0..15) will have the sum? Wait, `__shfl_down_sync` copies from lane `lane + delta`. So lane k (0..15) receives from lane k+16. So lane k does: `val = val + __shfl_down_sync(0xFFFFFFFF, val, 16)`. Then lane 0..15 have the full dot product. Lane 16..31 have the old val plus something else? Actually after this instruction, lane k (0..15) gets `lane_k_val + lane_{k+16}_val`. Lane k+16 gets `lane_{k+16}_val + lane_{k+32}_val` (which is out of range, returns lane_{k+16}_val). So lane 16..31 just have their own val. So only lanes 0..15 have the full sum. That's fine. We can have only lanes 0..15 participate in the rest, or we can store the result for all lanes. - Alternatively, use a shared memory array of size [8][16] to store the dot products. - Procedure: - Initialize local_sum = 0. - For i = lane_id; i < DT; i += 32: // if lane_id >=16, it handles some other elements? Let's use a simpler mapping. - Actually we can just have each thread sum over DT elements with stride 32. Wait, lane_id ranges 0..31. If each thread sums elements `d = lane_id, lane_id+32` within the DT=64 tile. Then after 1 iteration, each thread has the sum of those 2 elements. Then we need to sum across all 32 threads to get the dot product for a specific key. That would require a warp reduction sum for each key. Since there are 16 keys, we can do 16 warp reductions concurrently if threads are partitioned. But a warp reduction is over the whole warp. If we have 16 keys, and we want to reduce 32 values into 16 sums, we can use a segmented reduction? - Or we can just assign each key to one thread (16 threads active, 16 idle). Thread k (lane_id=k) loads Q[q][d] and K[k][d] for d=0..63, computes the dot product. Since DT=64 and we process 16 D tiles, the dot product is accumulated as we iterate d_tiles. Each thread does 64 element loads per D tile, 16 tiles => 1024 element pairs. That's 1024 FMAs per key. 16 keys * 1024 = 16384 FMAs per warp. With 8 warps = 131072 FMAs for S per key tile. With 32 key tiles = 4,194,304 FMAs. Plus V update, doubling to ~8.4M FMAs per block. Wait, the total should be ~536M per head / 256 threads = 2M FMAs per thread. Something is off. - Let's recalculate: per head FLOPs = 2 * N^2 * D = 2 * 512 * 512 * 1024 = 536,870,912 FMAs? Wait, one FMA = multiply-add. So M*D*N for QK^T uses M*D*N multiplies and M*D*N adds. That's 2*N*D*N ops? Usually GEMM counts 2*M*N*K flops (counting multiply and add separately). So FMA count = M*N*K. For QK^T: 512*512*1024 = 268,435,456 FMAs. For PV: same. Total = 536,870,912 FMAs per head. - My per-key-tile calculation: key_tile size BK=16. For one query group (BQ=8), D_tiles=16. QK part: BQ * BK * D = 8 * 16 * 1024 = 131,072 FMAs. With 32 key tiles = 4,194,304 FMAs. But there are 64 query groups! Wait, qg loop is 64. So total QK FMAs = 64 * 32 * 131,072 = 268,435,456. Yes, exactly. Good. PV part: same. So total per block = 536M FMAs. With 256 threads, each thread does ~2M FMAs. Good. - Back to thread mapping: If we use only 16 threads active per warp for S computation, then those 16 threads do twice the work: each thread handles one key and loops over all D. Wait, 16 threads, each handles one key. For a key tile, the dot product is over D=1024. Each thread does 1024 FMAs per key tile. There are 32 key tiles and 64 query groups. Total FMAs per thread = 32*64*1024 = 2,097,152. That's exactly the budget! The idle 16 threads can be used for other things or just idle. But we have 256 threads total; 128 will be active (half), doing 2M FMAs each. The other 128 idle. Is this okay? It means we are using only half the threads for the main compute. But each warp has 16 active threads. On NVIDIA GPUs, warps execute with all 32 threads; inactive threads are masked. So we effectively waste half the potential throughput. But we can perhaps have each active thread handle two keys. Let's see: - Each warp handles one query. If we have 16 keys, we can assign each thread two keys by splitting the D dimension. - We want each of the 32 threads to be active. - For each key, we need to compute dot(Q, K) over D=1024. We can split D across the two threads assigned to the key. - Thread k (0..15) handles key k and d in [0, 511]? No, both threads would handle half of D for the same key. - Thread k and thread k+16 both work on key k. They split DT=64: thread k handles d in [0,31] of the tile, thread k+16 handles [32,63]. Then after each D tile, they reduce. Then after all D tiles, they have the full dot product. - This uses all 32 threads, each doing half the work per key tile. The reduction is a shuffle. - This is better. Modified S computation mapping: - warp_id = query within group (0..7) - lane_id = 0..31 - key_k = lane_id % 16 (for lane_id 0..15 and 16..31). Lane 0 and 16 handle key 0; lane 1 and 16 handle key 1, etc. - part = lane_id / 16 (0 or 1). part 0 handles d_offset = 0..31 within DT; part 1 handles d_offset = 32..63. - Algorithm per D tile dt: - Load Q[q][dt*DT + d_local] for d_local in thread's range into registers or smem. - Wait, to compute dot product, each thread needs Q for its query and its half of D. But Q is the same for both threads working on the same key (same query). We can load Q into shared memory to avoid redundant loads. Q_smem[BQ][DT] is loaded cooperatively. - Similarly, each thread needs K[k][dt*DT + d_local]. K_smem[BK][DT] is loaded cooperatively. - Then each thread computes `acc += Q_smem[q][d_thread] * K_smem[k][d_thread]`. - After all 16 D tiles, each thread has acc (sum over 512 elements of its half). - Then `float other = __shfl_xor_sync(0xFFFFFFFF, acc, 16);`? Wait, lane k and lane k+16 need to communicate. `__shfl_xor_sync(mask, var, 16)` exchanges between lane i and i^16. So lane k (0..15) exchanges with lane k+16. After `acc = acc + other`, both lanes have the full sum. This requires the `other` to come from lane k^16. - `float other = __shfl_down_sync(0xFFFFFFFF, acc, 16);` for lane k gets from lane k+16. For lane k+16, gets from k+32 (out of range => returns acc). So only lanes 0..15 get the sum. Let's just have lanes 0..15 store to smem. - `if (part == 0) acc = acc + __shfl_down_sync(0xFFFFFFFF, acc, 16);` // lane k gets lane k+16 value. - Then lane 0..15 writes acc to `s_smem[warp_id][key_k]`. This stores the full dot product for key_k. - After the dt loop, we have s_smem[q][k] for q=0..7, k=0..15. - Now we need max over k for each q. - For warp q: lanes 0..15 read s_smem[q][lane_id] (lane_id in 0..15). Lane 16..31 can also read or just idle. We need max over 16 values. - `float val = s_smem[warp_id][lane_id % 16];` - Warp reduce max over 32 lanes: we can do `for (int offset = 16; offset > 0; offset /= 2) val = max(val, __shfl_down_sync(0xFFFFFFFF, val, offset));` - After this, lane 0 has m_tile (max over all 16 keys). - Next, compute l_tile = sum_k exp(s_k - m_tile). - `float val = expf(s_smem[warp_id][lane_id % 16] - m_tile);` - Warp reduce sum: `for offset... val += __shfl_down_sync(...)`. After, lane 0 has l_tile. - Next, compute p_k = exp(s_k - m_new) ... Actually we need to update running stats. - We have running m (per query) and running l (per query). These can be stored in registers for lane 0 of each warp. Or in shared memory (small array of 8 floats each). - m_new = max(m, m_tile). - correction_prev = expf(m - m_new); // if m was -inf, this is 0? Actually exp(-inf) = 0. - correction_tile = expf(m_tile - m_new); - l_new = correction_prev * l + correction_tile * l_tile; - // We also need p_k_raw = exp(s_k - m_new) for each key k. This is `exp(s_k - m_tile) * exp(m_tile - m_new) = exp(s_k - m_tile) * correction_tile`. - So for each key k, p_weight = exp(s_k - m_tile) * correction_tile. - We can store p_weight in `p_smem[q][k]` for k=0..15. Since p_smem is 8*16=128 floats. - Computation: lane_id in 0..15 reads s_smem[warp_id][lane_id], computes `p_weight = expf(s_k - m_tile) * correction_tile`. Writes to p_smem[warp_id][lane_id]. - Lane_id >= 16 can write to p_smem too if we want, but they don't have valid s. So only 0..15. - Then update O accumulator for this query group across all D tiles. - O accumulator: for each D tile, we have `o_acc[q][d]` in some memory. - We can store o_acc in shared memory: size = BQ * D = 8*1024 = 8192 floats = 32KB. This is exactly the O accumulator for the entire query group across all D dimensions. - Wait, we process query group once per key tile. Between key tiles, we need o_acc to persist. Smem is perfect for this. - Initialize o_acc to 0 before key tile loop. - For each key tile, after computing p_weights: - For each dt in 0..(D/DT - 1): - Load V[kt][dt] into v_smem[BK][DT] (16*64=1024 floats=4KB). - Threads update o_acc for their assigned d values. - Thread mapping for O update: - We need to compute for all q (8) and all d in dt (64) the update: `o_acc[q][dt*DT + d] = correction_prev * o_acc[q][...] + sum_k p_smem[q][k] * v_smem[k][d]`. - warp_id = q. - lane_id = 0..31. - For dt=64, each warp handles 64 d values. With 32 threads, each thread handles 2 d values: `d0 = lane_id`, `d1 = lane_id + 32`. - For each of its d values: - Compute local_sum = sum_{k=0}^{BK-1} p_smem[q][k] * v_smem[k][d]. - Since BK=16, we can unroll or loop. - `float sum0 = 0, sum1 = 0;` - `for (int k=0; k<BK; k++) { float p = p_smem[warp_id][k]; sum0 += p * v_smem[k][d0]; sum1 += p * v_smem[k][d1]; }` - But each thread does this loop independently. That's 16 iterations per thread, 2 FMAs per iteration = 32 FMAs per thread per dt. There are 16 dt and 32 kt and 64 qg. Total FMAs per thread = 32 * 16 * 32 * 64 = 1,048,576. Wait, that's half the required 2M. Why? - Let's recalculate: For PV: BQ * N * D = 8 * 512 * 1024 = 4,194,304 FMAs per query group. With 32 key tiles, each key tile contributes 8*16*1024 = 131,072 FMAs. Per thread? 256 threads. 131,072 / 256 = 512 FMAs per thread per key tile? That seems too low. Wait, total FMAs for one key tile's PV contribution is: for each query (8), each key (16), each D (1024): 8*16*1024 = 131,072 FMAs. With 256 threads, average 512 FMAs per thread per key tile. Over 32 key tiles, that's 16,384 FMAs per thread per query group. Over 64 query groups, that's 1,048,576 FMAs. This is only for the PV part. QK part is another 1,048,576 FMAs. Total 2,097,152 FMAs per thread. Yes, that matches. - In the O update loop above, each thread did 16 iterations * 2 FMAs = 32 FMAs per dt. Over 16 dt => 512 FMAs per thread per key tile. Over 32 key tiles => 16,384 FMAs per thread per query group. Over 64 query groups => 1,048,576 FMAs for PV. Perfect. So mapping gives correct FMA count. - After computing sum0, sum1, we update o_acc: - Read old o_acc: `old0 = o_smem[q][dt*DT + d0]` (or from register if we kept it). - `new0 = correction_prev * old0 + sum0;` - Write back. - Note: `correction_prev` is the same for all d for a given q. - `correction_prev` and `correction_tile` and `l_new` can be stored in shared memory arrays `m_arr[8]`, `l_arr[8]`, and maybe computed per key tile. - Since lane 0 computed m_new, l_new, it can store them. Other lanes need to read `correction_prev`. We can broadcast via smem or shuffle. Optimizing shared memory usage: - Smem layout: - float smem_q[8][64]; // 2KB (used during S stage) - float smem_k[16][64]; // 4KB (used during S stage) - float smem_v[16][64]; // 4KB (used during O stage, can reuse smem_k) - float smem_s[8][16]; // 512 bytes (scores for current key tile) - float smem_p[8][16]; // 512 bytes (weights, can reuse smem_s) - float smem_o[8][1024]; // 32KB (accumulator for O for current query group) - float smem_m[8]; // 32 bytes - float smem_l[8]; // 32 bytes - Total with reuse: - smem_q (2KB) + smem_k/v (4KB) + smem_s/p (512B) + smem_o (32KB) + smem_m/l (64B) = 38.5KB. - Wait, can we overlap smem_q with smem_v? During O stage, we don't need Q. But we might need smem_q for next key tile's S? No, we finish all key tiles for current query group, write O, then load next query group. But within a query group, we do key tile loop: S stage then O stage. After O stage, we go to next key tile. The next iteration S stage loads Q again. So we can reuse smem_v for smem_k (they are never needed simultaneously). Actually during S stage, we need smem_q and smem_k. During O stage, we need smem_p and smem_v. No overlap. So max footprint = max(2+4, 0.5+4) + 32 = 38KB. Plus small arrays. So ~38KB. That fits in 48KB easily. - Wait, do we need smem_q and smem_k simultaneously during S? Yes. So 6KB. - Do we need smem_p and smem_v simultaneously during O? Yes. But smem_p is 512B, smem_v is 4KB. So 4.5KB. - So max = 32KB (o) + 6KB (q+k) = 38KB. Good. Programming details: - We need to be careful about __syncthreads() between loading smem and using it. - Since block size is 256 and we have 8 warps, some operations like loading smem can be done by all threads. - Loading Q_smem: for each query group qg and dt, we need to load Q for 8 queries and 64 features. - Q layout: [N, D] (for this head). Q[q][d]. - Offset: base + (qg * BQ + q) * D + dt * DT + d. - We can have all 256 threads load one element each: we need 8*64=512 elements. But we only have 256 threads. So each thread loads 2 elements. - Thread tid loads: - q_idx = tid / 64? But 256/64 = 4, not 8. Let's do: - For loading Q_smem, we can use a simple 1D loop: - for (int i = tid; i < BQ * DT; i += blockDim.x) { int q = i / DT; int d = i % DT; smem_q[q][d] = Q[...]; } - This is efficient coalescing? Threads access consecutive elements if i is consecutive. Since Q is row-major [N][D], consecutive i means consecutive d, same q until wrap. This is coalesced within each query. - Similarly for K_smem: for (int i = tid; i < BK * DT; i += blockDim.x), load K[kt*BK + k][dt*DT + d]. - For V_smem. - For O output: after all key tiles, write smem_o to global memory O for the query group. Wait, there is an issue. The query group qg iterates 64 times. The O accumulator for the query group is stored in smem_o (32KB). After finishing all key tiles for the query group, we must write O to global memory. Then reset smem_m, smem_l, smem_o for the next query group. This is straightforward. Let's carefully write out the kernel steps: ``` extern "C" __global__ void fused_attention_kernel( const float* __restrict__ Q, const float* __restrict__ K, const float* __restrict__ V, float* __restrict__ O, int B, int H, int N, int D, float scale) { // Block indices int b = blockIdx.x; int h = blockIdx.y; int tid = threadIdx.x; int warp_id = tid / 32; int lane_id = tid % 32; // Base offset for this head int head_offset = ((b * H + h) * N) * D; // Smem declarations __shared__ float smem_q[8][64]; // BQ x DT __shared__ float smem_k[16][64]; // BK x DT // Reuse smem_k for V later __shared__ float smem_s[8][16]; // scores const int BQ = 8; const int BK = 16; const int DT = 64; __shared__ float smem_o[8][1024]; // BQ x D __shared__ float smem_m[8]; __shared__ float smem_l[8]; // Each block processes one head for (int qg = 0; qg < N / BQ; qg++) { // Initialize smem_o, smem_m, smem_l for (int i = tid; i < BQ * D; i += blockDim.x) { smem_o[i / D][i % D] = 0.0f; } if (tid < BQ) { smem_m[tid] = -INFINITY; smem_l[tid] = 0.0f; } __syncthreads(); // Load Q for this query group? We will load per dt in key tile loop. // Actually we can load all Q for this group once? 8*1024=8192 floats. // 8192 * 4 = 32KB. We have smem_o (32KB) being used, so we can't load all Q at once. // So we load Q per D tile as planned. for (int kt = 0; kt < N / BK; kt++) { // Compute scores S for this key tile // Initialize smem_s // Actually we compute into registers first, accumulate, then write to smem_s at end. float local_acc = 0.0f; int q_local = warp_id; int k_local = lane_id % BK; // 0..15 int part = lane_id / BK; // 0 or 1, each handles 32 elements of DT for (int dt = 0; dt < D / DT; dt++) { // Load Q for (int i = tid; i < BQ * DT; i += blockDim.x) { int q = i / DT; int d = i % DT; int q_idx = qg * BQ + q; smem_q[q][d] = Q[head_offset + q_idx * D + dt * DT + d]; } // Load K for (int i = tid; i < BK * DT; i += blockDim.x) { int k = i / DT; int d = i % DT; int k_idx = kt * BK + k; smem_k[k][d] = K[head_offset + k_idx * D + dt * DT + d]; } __syncthreads(); // Compute partial dot product // q_local = warp_id, k_local = lane_id % BK // We access smem_q[q_local][d] and smem_k[k_local][d] // Each thread handles d from part*(DT/2) to part*(DT/2) + DT/2 - 1 int d_start = part * (DT / 2); #pragma unroll for (int d = 0; d < DT / 2; d++) { int d_offset = d_start + d; local_acc += smem_q[q_local][d_offset] * smem_k[k_local][d_offset]; } __syncthreads(); // before next load } // Reduce local_acc between the two threads of the same key // lane_id and lane_id^16 correspond to same key, different parts float other = __shfl_xor_sync(0xFFFFFFFF, local_acc, 16); // For lane_id 0..15, we want sum; for 16..31, also sum if we want. // Actually __shfl_xor_sync exchanges with lane_id ^ 16. // So both lanes get the other part. float s_val = local_acc + other; // Only store from lanes 0..15 to avoid duplication if (lane_id < BK) { smem_s[warp_id][lane_id] = s_val * scale; } __syncthreads(); // Now do softmax for this key tile // Read s_val for this lane's key (only lanes 0..15 valid) float my_s = (lane_id < BK) ? smem_s[warp_id][lane_id] : -INFINITY; // Warp reduce max over 32 lanes float m_val = my_s; #pragma unroll for (int offset = 16; offset > 0; offset /= 2) { float other_m = __shfl_down_sync(0xFFFFFFFF, m_val, offset); m_val = fmaxf(m_val, other_m); } // m_val now has max over all lanes. But shuffle_down duplicates; need careful. // Wait, __shfl_down_sync copies from lane + offset. For offset=16, lane 0 gets from lane 16. // Lane 0 and lane 16 both have my_s (for key 0 and 16? Actually lane 16 has key 0). // So lane 0 gets key0_part1's s_val. lane 0 fmaxs with its own key0_part0 s_val. Good. // But lane 1 and lane 17 both have key1. etc. // After offset=16, lanes 0..15 have max of their key and key+16? No, lane 0..15 have max(key_k, key_{k+16})? But key_{k+16} corresponds to lane 16+k which also has key k. So it's just max of the two copies, which are identical after shfl_xor sum? Wait, __shfl_xor gave the same s_val to both lanes (since local_acc+other is same). So m_val after offset=16 is just that same value. // Then offset=8: lane 0 gets from lane 8 (key 8). Lane 0 has key0. We want max over all keys 0..15. This works! Because lanes 0..15 represent keys 0..15. // But wait, lane 16 gets from lane 24 (key 8's duplicate). The values in lanes 16..31 are the same as 0..15. The reduction across all 32 lanes will compute max of all keys because values in 16..31 are duplicates of 0..15. So it's fine; we just need lane 0's result. // Actually `__shfl_down_sync` for offset>=16 for lane 16..31 will read out of range (wrap around or same lane). According to CUDA docs, if src lane is out of range, it returns the value from the source lane modulo warp size? Actually for `__shfl_down_sync`, if `srcLane >= 32`, the operation returns the value from the same lane (i.e., `var`). So for lane 16, offset=16 => src=32 >=32 => returns `var`. So lane 16..31 will effectively ignore higher offsets. That's fine because we only care about lane 0's result. // Let's just do the standard warp reduce using `__shfl_xor_sync` which is simpler for max. m_val = my_s; #pragma unroll for (int offset = 16; offset > 0; offset /= 2) { float other_m = __shfl_xor_sync(0xFFFFFFFF, m_val, offset); m_val = fmaxf(m_val, other_m); } // After this, all lanes have the same m_val (max over warp). Good. // Now compute l_tile = sum exp(my_s - m_val) float exp_val = (lane_id < BK) ? expf(my_s - m_val) : 0.0f; float l_val = exp_val; #pragma unroll for (int offset = 16; offset > 0; offset /= 2) { l_val += __shfl_xor_sync(0xFFFFFFFF, l_val, offset); } // l_val is sum over warp (all lanes), same in all lanes. // Update global softmax stats for this query float m_prev, l_prev; if (lane_id == 0) { m_prev = smem_m[warp_id]; l_prev = smem_l[warp_id]; } m_prev = __shfl_sync(0xFFFFFFFF, m_prev, 0); l_prev = __shfl_sync(0xFFFFFFFF, l_prev, 0); float m_new = fmaxf(m_prev, m_val); float corr_prev = expf(m_prev - m_new); float corr_tile = expf(m_val - m_new); float l_new = corr_prev * l_prev + corr_tile * l_val; // Compute p weights for each key: p = exp(my_s - m_new) = exp_val * corr_tile // Wait: exp(my_s - m_new) = exp(my_s - m_val) * exp(m_val - m_new) = exp_val * corr_tile. float p_weight = exp_val * corr_tile; // Store p_weight to smem_s or a dedicated smem_p // We can reuse smem_s for p weights since we don't need scores anymore. if (lane_id < BK) { smem_s[warp_id][lane_id] = p_weight; } // Store updated m and l if (lane_id == 0) { smem_m[warp_id] = m_new; smem_l[warp_id] = l_new; } __syncthreads(); // Now update O accumulator using p_weights and V tile for (int dt = 0; dt < D / DT; dt++) { // Load V tile for (int i = tid; i < BK * DT; i += blockDim.x) { int k = i / DT; int d = i % DT; int k_idx = kt * BK + k; smem_k[k][d] = V[head_offset + k_idx * D + dt * DT + d]; // Reuse smem_k buffer for V (since we are in O stage) } __syncthreads(); // Each warp handles one query q_local = warp_id // Each thread handles two d's: d0 = lane_id, d1 = lane_id + 32 int d0 = lane_id; int d1 = lane_id + 32; float sum0 = 0.0f; float sum1 = 0.0f; #pragma unroll for (int k = 0; k < BK; k++) { float p = smem_s[q_local][k]; sum0 += p * smem_k[k][d0]; sum1 += p * smem_k[k][d1]; } // Apply correction_prev to old o and add sum // o_acc is in smem_o[q_local][dt*DT + d] int d_base = dt * DT; float old0, old1; // We need to read old o. Each thread reads its own two positions. // But there is no atomic needed; they are distinct addresses. old0 = smem_o[q_local][d_base + d0]; old1 = smem_o[q_local][d_base + d1]; float new0 = corr_prev * old0 + sum0; float new1 = corr_prev * old1 + sum1; smem_o[q_local][d_base + d0] = new0; smem_o[q_local][d_base + d1] = new1; __syncthreads(); } } // end key tiles // After all key tiles, normalize O by l for (int i = tid; i < BQ * D; i += blockDim.x) { int q = i / D; int d = i % D; float l_val = smem_l[q]; float o_val = smem_o[q][d]; int q_idx = qg * BQ + q; O[head_offset + q_idx * D + d] = o_val / l_val; } __syncthreads(); } // end query groups } ``` Wait, there is a subtle bug in the max reduction. `my_s` is only valid for lanes 0..15; lanes 16..31 are set to -INFINITY (or 0 for exp but for max, -INFINITY). But the warp reduction `__shfl_xor_sync` will compare. Since lanes 16..31 have -INFINITY, they won't affect the max if lanes 0..15 have real values. So `m_val` computed by all lanes will be correct. However, there is a problem: `__shfl_xor_sync` with offset 16 will exchange lane 0 and lane 16. Lane 0 has `my_s` for key 0. Lane 16 has `my_s` but we set it to -INFINITY. So after the first XOR, lane 0's m_val becomes max(key0, -INF) = key0. But lane 16's m_val becomes max(-INF, key0) = key0. Actually both get key0. That means lanes 0..15 and 16..31 will have the values of the corresponding lanes they exchange with. After all offsets, all lanes will have the max of all keys where at least one lane had that value. Since lanes 0..15 have the actual values, and lanes 16..31 have -INF, the max across all 32 will be the max of keys 0..15. So m_val in all lanes will be correct. Good. But wait, `exp_val` for lanes 16..31 is 0. So l_val includes 0 for those lanes. Warp reduction sum will only sum the 16 valid lanes. That's correct. Now, a performance concern: the `__syncthreads()` inside the `dt` loop for V loading and update. Since we have 16 dt iterations per key tile and 32 key tiles = 512 __syncthreads() calls inside the inner loop (plus the ones in S stage). That's about 1024 __syncthreads per query group. Per block, 64 query groups => ~65k __syncthreads. This is a lot of synchronization overhead. However, `__syncthreads()` within a block on modern GPUs is relatively fast, but it might still add up. We can try to reduce syncs by having threads load V into smem without full sync? No, we need all threads to see the loaded V before using it. But what if we load V for the next dt while processing current dt? That's double buffering. We can allocate two buffers for V in smem and use them alternately. But the sync is still needed before consumption. However, we can maybe remove the `__syncthreads()` after writing O, because threads within a warp access distinct d indices, but they all read V_smem which was loaded and synced. We need to ensure that no thread proceeds to load the next V tile before all threads have finished reading the current one. Since the loop iteration is the same for all threads (they all do `for dt...`), and `__syncthreads()` after loading V is needed. But do we need a `__syncthreads()` after writing O? The next iteration loads V into the same smem_k. Some threads might be writing O (to smem_o) while others have moved to load V. But O writes go to smem_o, and V loads go to smem_k. These are disjoint smem arrays, so no conflict. However, we need a `__syncthreads()` at the end of the `dt` loop to ensure all writes to smem_o are done before any thread moves to the next key tile? No, the next key tile reads V and writes smem_k, and also reads smem_o. But all threads are in lockstep within the `dt` loop. After the last `dt`, we exit the dt loop and go to next `kt`. We have a `__syncthreads()` at the start of the next dt loop's V load. But between the last dt and the next key tile, we hit the top of the outer `kt` loop. There is no `__syncthreads()` after the last dt of the previous kt before the first dt of the next kt, except if we put one before `for dt=0`. Actually after the inner dt loop, we go to next `kt`. The first thing in the next `kt` is computing scores S, which needs `__syncthreads()` after loading Q and K. But we also need to make sure all threads have finished reading smem_k from the previous iteration before some thread starts loading Q/K into smem_k in the next S stage? No, in the S stage we load into smem_k. During O stage we use smem_k for V. In the next S stage, we overwrite smem_k with K. This is fine because the previous iteration is fully complete for all threads. There is an implicit sync because all threads execute the `kt` loop in lockstep. So just placing `__syncthreads()` after loading Q/K in S stage and after loading V in O stage should be sufficient. But actually in the `dt` loop of O stage: - load V - __syncthreads() - compute and write O // loop back, load next V. No need for __syncthreads after write O because next iteration loads V into smem_k, which is different from smem_o. Some warps might be ahead loading V while others are still writing O. But O writes are to smem_o; V loads are to smem_k. These don't alias. So it's safe! Thus we can remove the __syncthreads after O computation inside the dt loop. - However, we must be careful: inside the dt loop, after `__syncthreads()` following V load, all threads read smem_k. Then they write smem_o. Since smem_o is not read by any other thread in this iteration, no sync needed. Next iteration: we load V into smem_k. Threads that finished writing O can start loading next V into smem_k. But threads that are still writing O might be slower. If they catch up, they might start reading smem_k before it's fully loaded by the faster threads. So we need a `__syncthreads()` after V load at the start of each iteration. That's enough. - Wait, but the load into smem_k is done cooperatively by all threads. Faster threads might finish loading and then slower threads might still be loading. `__syncthreads()` after the load loop ensures all loads are visible. Good. - Similarly, in the S stage dt loop, we load Q and K, then __syncthreads(), then compute, then __syncthreads() before the next iteration's load. Actually after compute, we might not need __syncthreads if the next dt iteration loads into different smem? No, it loads into smem_q and smem_k again. Some threads might still be reading smem_q/k from previous dt while others are writing to them for next dt. So we need __syncthreads() at the end of dt loop iteration (after compute, before next load). However, Q and K loads in the next iteration will overwrite. But if some threads from the previous iteration are still reading smem_q (from the compute stage), they could be reading new data. So we need a __syncthreads() after compute. Since the compute stage is fully within registers (reading smem_q and smem_k into registers), and all threads in the block participate, after the compute loop we need __syncthreads() before the next load. Unless some threads are faster and start loading while slower threads still read. So yes, __syncthreads() is needed. Number of __syncthreads per dt for S stage: 2 (one after load, one after compute). For O stage: 1 (after load V). Actually for O stage, after O write, next iteration loads V. We need __syncthreads after O write. Let's keep it simple and safe: place __syncthreads after compute in S dt loop, and after write in O dt loop. Total per dt: 3. Total per key tile: (16*2 + 16*2) = 64? No: S dt loop (16 iterations): load Q+K, sync, compute, sync. -> 32 syncs. O dt loop (16 iterations): load V, sync, compute+write O, sync. -> 32 syncs. Total per key tile: 64 syncs. Over 32 key tiles: 2048 syncs per qg. Over 64 qg: 131k syncs per block. That seems like a lot but maybe okay? Each __syncthreads is a barrier instruction. On H100, it's cheap but not free. 131k * 1024 blocks = 134M syncs. That might be a bottleneck. Wait, we can avoid per-dt syncs in O stage by loading V for all D tiles at once? No, D=1024, V tile for all D would be 16*1024*4 = 64KB. That exceeds 48KB. But we can request more shared memory. H100 allows up to 227KB. If we request e.g., 96KB, we could load V for two D tiles at once? Or we could load K for a larger BK and process more keys per tile? But BK=16 fits resource constraints. Alternative: load V for the entire key tile (BK=16, D=1024) into registers? 16*1024 = 16k floats = 64KB per thread block. Distributed over 256 threads, that's 64 floats per thread. But with coalesced shared loads, it's better. Another idea: process query groups differently. Instead of 64 query groups of 8, maybe process larger query groups to reduce the number of query groups and thus the overhead of resetting smem. But the smem_o footprint grows with BQ. If BQ=16, smem_o = 16*1024*4 = 64KB. Plus Q+K/V -> 70KB. That's fine within 96KB or 128KB shared memory. But register pressure might increase? Not necessarily. What if BQ=16, BK=16, DT=64, and block size 256 or 512? - With BQ=16, we need 16 warps (512 threads) to assign one warp per query. Or we could keep 256 threads and assign half a warp per query (16 threads). That might be less efficient. - Let's try BQ=16, block size = 256 threads. Then each query gets 16 threads (half warp). That's not a full warp; warp reductions work on 32 threads. So we'd need different reductions. This adds complexity. Let's revisit the FLOPs and see if this kernel is compute-bound or latency-bound. Per block total FLOPs = 536M. Block time = e.g., 10ms. Then per dt sync cost is negligible. If block time = 20ms, maybe sync cost matters. But 256 threads doing 536M FMAs is 2M FMAs/thread. At 1 FMA/cycle (FP32 on H100?), each thread might take ~2M cycles. At ~2GHz, that's 1ms per thread. But thread throughput is lower because of memory accesses. The actual block time might be several ms. 131k syncs * maybe 10 cycles each = 1.3M cycles ≈ 0.6ms. So sync overhead is maybe 10-20%. Acceptable, but can be reduced. We can reduce dt syncs in S stage by using double buffering in smem: load next Q/K while computing current. But that doubles smem usage for Q+K from 6KB to 12KB. We have plenty of smem budget (48KB total, using 38KB). So 12KB is fine. Double buffering for S: - smem_q[2][8][64] - smem_k[2][16][64] - We load buffer 0, sync, start compute on buffer 0 while loading buffer 1? Can't overlap compute and load within same block easily without async. But we can use pipeline concept. With H100, we could use TMA or cp.async, but simple sync is okay. Wait, maybe we can reduce the dt syncs for S stage to just 1 per dt? After compute, we don't need to sync if next load is independent. But we need to ensure that all threads finish reading smem_q/k before any thread starts writing to them for next dt. Since compute reads from smem and all threads in the block participate in compute, after the last thread finishes compute, it's safe to load. But the load loop is executed by all threads. If we remove the sync after compute, threads that finish compute late will start the load for next dt later. Threads that finish early might start loading, but since the load is cooperative, all threads must eventually do it. If a fast thread overwrites smem_q while a slow thread is still reading it, that's a race. So yes, sync needed. Let's try to optimize by increasing DT to 128. Then D tiles = 8. This halves the number of dt iterations and thus syncs. - DT = 128: - Q_smem: 8*128 = 1024 floats = 4KB. - K_smem: 16*128 = 2048 floats = 8KB. - Total S stage = 12KB. - O_smem: 8*1024*4 = 32KB. - V_smem: 16*128 = 2048 floats = 8KB. - Total max = 32 + 12 = 44KB. Plus small arrays. Still under 48KB! Great. - Let's check thread compute: each thread in S stage does sum over DT/2 = 64 elements per dt. With 8 dts, total 512 elements, 512 FMAs. Per key tile, that's per thread. With 32 key tiles, 16k FMAs. With 64 qg, 1M FMAs per thread for QK. Good. - In O stage, each thread loop over 16 keys and 2 d's. 16 * 2 = 32 FMAs per dt. With 8 dts, 256 FMAs per key tile. 32*64 = 2048 key tiles? No, 32 key tiles per qg * 64 qg = 2048 key tiles. 2048 * 256 = 524k FMAs? Wait, earlier we had 32 FMAs per dt * 16 dt = 512 per key tile. Now 32 * 8 = 256 per key tile. That's half! Something is wrong. - PV FMAs per key tile: BQ * BK * D = 8 * 16 * 1024 = 131,072. With 256 threads, each thread should do 131,072 / 256 = 512 FMAs per key tile. Wait, previously with DT=64, each thread did 16 dt * 16 keys * 2 d = 512 FMAs per key tile. With DT=128, each thread does 8 dt * 16 keys * 2 d = 256 FMAs per key tile. But total FMAs didn't change. Why? Because each thread handles 2 d's out of 128, so still 2 d's. The number of dts halved, but the number of features per thread per dt is... wait, each thread still handles d0 = lane_id and d1 = lane_id + 32. Since DT=128, lane_id 0..31 handles 2 out of 128 features. That's a smaller fraction! Previously DT=64, lane_id handled 2 out of 64 = 1/32 of the features. Now it's 2 out of 128 = 1/64. So each thread's work halved. That means we'd need twice as many threads to cover the work, but we have fixed block size. So with DT=128, we only use half the parallelism for O update. That's because our thread mapping assigned only 32 threads per query (one warp). To cover 128 features, we'd need 4 warps per query, but we only have 1 warp. So DT=128 requires rethinking thread mapping. - Alternative mapping for O stage with DT=128: each query uses 64 threads to handle 128 features (2 per thread). With 8 queries, we need 512 threads. That's >256. So we can't. - Or each thread handles 4 features: d0=tid%32, d1=d0+32, d2=d0+64, d3=d0+96. Then each thread does 4*16=64 FMAs per dt. 8 dts -> 512 FMAs per key tile. That works with 256 threads. Each thread handles 4 features, 32 threads per query. - For S stage with DT=128, we split DT/2=64 per thread. That remains fine. - Let's check register pressure: sum0,sum1,sum2,sum3 = 4 registers. old0..3 = 4 registers. Plus loop counters. Still okay. - So DT=128 is possible. But let's see if it complicates code. - With DT=128, number of dt loops halves. Syncs reduce. This is good. But maybe simplest to stick with DT=64, which gives a nice round mapping and enough parallelism. Another option: increase BQ to 16, block size to 512. Then each warp handles one query (16 warps). O_smem = 16*1024*4 = 64KB. Total with Q+K = 64+6 = 70KB. Still under 96KB (if we request 96KB). But Hopper supports up to 227KB. If we request 80KB, we can. But we want to keep occupancy high. With 80KB smem per block, max blocks per SM is 227/80 = 2. So 2 blocks per SM. With 132 SMs, 264 blocks concurrently. Total blocks = 1024. So ~4 waves. 256 threads per block. That's fine. But BQ=16 means 32 query groups (512/16). So qg loop is 32. Halving query groups might reduce overhead. Let's see if BQ=16, blockDim=512 (16 warps), DT=64, BK=16 works. - S stage: each warp (32 threads) handles one query and 16 keys. Same mapping as before. Works. - O stage: each warp handles one query. 64 features per dt, 32 threads, 2 features per thread. Same as before. - O_smem = 16*1024*4 = 64KB. Q+K = 6KB. Total = 70KB. V_smem = 4KB. But V_smem can reuse K. So max = 64+6 = 70KB. Request 72KB. Should be fine. - Register usage: each thread still low. - Number of syncs: qg loop 32, kt loop 32, dt loop 16 -> 32*32*16*syncs. Syncs are proportional to qg*kt*dt. BQ=16 halves qg, so total syncs halved. Good. - HBM traffic still 8MB per head (Q=2, K=2, V=2, O=2). Because K/V read once from HBM, Q once, O written once. Perfect. - This seems better! Let's target this. - Wait, does it beat reference? Reference math fallback does two BMMs through cuBLAS. cuBLAS is likely highly optimized. Our kernel uses CUDA cores and has many __syncthreads and loops. It might be compute latency limited. But maybe the memory savings and fusion give advantage. Let's try both BQ=8,BK=16,DT=64 and BQ=16,BK=16,DT=64. The latter requires more smem but fewer query groups. Let's start with BQ=8 because it fits in 48KB without requesting more. If performance is poor, we can try larger. Actually, I realize there is a potential issue with the one-block-per-head approach: occupancy. 1024 blocks on 132 SMs. If each SM can run only 1 block at a time due to register or smem constraints, that's 132 blocks concurrently. 1024/132 ≈ 8 waves. If each wave takes ~1ms, total ~8ms. That's okay. But if each wave takes ~5ms, total ~40ms. The reference might be ~15ms. We need to maximize throughput. To maximize throughput, we want more blocks running concurrently. If we limit smem to 48KB, each SM can run at most floor(228/48)=4 blocks (actually max 32, but limited by smem). So 4 blocks per SM. 132*4 = 528 blocks concurrently. 1024/528 ≈ 2 waves. That's great! So occupancy is not a problem with 48KB. If we use 70KB, floor(228/70)=3 blocks per SM. 132*3=396 concurrently. 1024/396≈2.6 waves. Still okay. But let's stay with 48KB for better occupancy and lower risk. So BQ=8, BK=16, DT=64, blockDim=256. Let's think about the memory access pattern for K and V loads. When loading K for a key tile and dt: `K[head_offset + k_idx * D + dt * DT + d]` When iterating over dt, the address strides by DT. For fixed k_idx, consecutive dt accesses are separated by D=1024 floats = 4KB. This is not contiguous across dt. But within a dt loop iteration, consecutive threads load consecutive d, which is coalesced. The dt iterations access `d` in chunks. The access pattern is sequential within the chunk. When iterating over key tiles (kt), k_idx changes. For the same dt, consecutive key tiles access different keys. Since we have 32 key tiles, K is effectively read sequentially fully. Because one block reads the entire K matrix (512x1024), the pattern is good for caching in L2. But does it utilize the L1 cache? The Q and K are loaded into smem in tiles, which is good. Similarly for V. Potential issue: when loading Q for each dt and each qg, we access Q. Q is read sequentially in qg loop. Good. Let's code this kernel. One more detail: for the first key tile, `smem_m` is -inf and `smem_l` is 0. `m_new = max(-inf, m_val) = m_val`. `corr_prev = exp(-inf - m_val) = 0`. `corr_tile = exp(m_val - m_val) = 1`. `l_new = 0 * 0 + 1 * l_val = l_val`. `p_weight = exp(my_s - m_val) * 1 = exp(my_s - m_val)`. `new_o = 0 * old_o + sum = sum`. Correct. For subsequent key tiles, `corr_prev = exp(m_prev - m_new)`. If `m_prev` > `m_val`, then `corr_tile = exp(m_val - m_new)` where `m_new=m_prev`. So `corr_tile = exp(m_val - m_prev) < 1`. This correctly downweights the new tile. Wait, in the formula `l_new = corr_prev * l_prev + corr_tile * l_val`. If m_new = m_prev, corr_prev = 1. If m_new = m_val, corr_prev = exp(m_prev - m_val) < 1. This is correct. At the end, `O = o / l`. This gives the normalized weighted sum. Wait, is `smem_s[8][16]` valid for `smem_s[q_local][k]` when `k` is read by threads? In the O update stage, we have: `for (int k = 0; k < BK; k++) { float p = smem_s[q_local][k]; ... }` All threads in warp q_local read `smem_s[q_local][0..15]`. This is fine because shared memory broadcast for the same address within a warp is efficient. But `q_local` is the same for all threads in the warp. However all threads in the block read `smem_s[q_local][k]` where `q_local` differs per warp. This is also fine because each warp accesses its own row. Wait, what about `smem_m` and `smem_l` reads? In the S stage, after sync, all threads read `smem_m[warp_id]` and `smem_l[warp_id]`. We used shuffle from lane 0. But if we use `__shfl_sync`, we need lane 0 to have loaded it. Since `smem_m` and `smem_l` are only written by lane 0 after the first iteration, and in the first iteration they were initialized by all threads. After initialization we had a `__syncthreads()`. So all threads see it. But in the O stage, we read `corr_prev` which is computed by lane 0. All threads need it to update O. So lane 0 broadcast it via shuffle. Good. Wait, `corr_prev` and `corr_tile` are computed per key tile, not per dt. So inside the dt loop for O, we need their values. They are in registers (broadcasted). Good. Let's write the kernel. Note: The prompt says "Do not write your own PYBIND11_MODULE block". We will use `load_inline` with `functions=["forward"]`. Our C++ function should be named `forward` and accept pointers plus dimensions. But since we need to pass tensors, we can write a wrapper in C++ that takes `at::Tensor` and calls the kernel. Simpler: define the CUDA kernel as described, then write a C++ wrapper: ```cpp torch::Tensor forward(torch::Tensor Q, torch::Tensor K, torch::Tensor V) { auto options = Q.options(); auto O = torch::empty_like(Q); int B = Q.size(0); int H = Q.size(1); int N = Q.size(2); int D = Q.size(3); float scale = 1.0f / sqrtf((float)D); dim3 grid(B, H); dim3 block(256); fused_attention_kernel<<<grid, block>>>( Q.data_ptr<float>(), K.data_ptr<float>(), V.data_ptr<float>(), O.data_ptr<float>(), B, H, N, D, scale); return O; } ``` We need to include `<torch/extension.h>` and `<cuda_runtime.h>`. We should also add `cudaFuncSetAttribute` to allow shared memory size > 48KB if needed. For 38KB, it's fine without. But let's just include it for safety. One more thing: the kernel does a lot of shared memory accesses. We should check bank conflicts. - `smem_q[8][64]`: access pattern is `smem_q[q_local][d_offset]`. For a warp, q_local is the same for all threads (since warp_id is same). d_offset differs per thread. If `DT=64` and threads access with stride 1, there are no bank conflicts (assuming 32 banks, width=4 bytes; consecutive d map to consecutive banks). Good. - `smem_k[16][64]`: similar, access by `k_local` (same for threads that share key, different for different keys) and `d_offset` (different for the two parts). Since threads with same `k_local` but different `part` access `d_offset` that differ by 32. Bank index = (addr/4) % 32. For `d=0` and `d=32`, banks are 0 and 0? Wait, shared memory banking: 32 banks, each 4 bytes. Address `a` maps to bank `a / 4 % 32`. For `smem_k[k][d]`, address = (k*64 + d)*4. Bank = (k*64 + d) % 32 = (k*0 + d) % 32 since 64 is multiple of 32. So bank = d % 32. - For part 0: d in [0..31]. Banks 0..31. Each thread accesses a different bank. No conflict. - For part 1: d in [32..63]. Banks 0..31 again. Each thread accesses a different bank. No conflict. - Across the two parts, same banks but different threads? Actually lane k and lane k+16 both access same key. Part 0 thread accesses d=k? No, `d_offset` is a loop variable. Wait, my code uses a loop: ```cpp int d_start = part * (DT / 2); for (int d = 0; d < DT / 2; d++) { int d_offset = d_start + d; local_acc += smem_q[q_local][d_offset] * smem_k[k_local][d_offset]; } ``` This means all threads in the warp loop together. Within an iteration `d`, each thread accesses a different `d_offset` because `d_start` differs by part and lane? Wait, `d_offset = d_start + d`. `d_start` is either 0 or 32. It is the same for all threads with the same `part`. For a given `d`, all threads with `part=0` access `d_offset=d`. All threads with `part=1` access `d_offset=d+32`. But threads have different `k_local`. So `smem_k[k_local][d_offset]`. Bank for `smem_k[k][d]` = (k*64 + d) % 32 = d % 32. For `part=0`, bank = d % 32. For `part=1`, bank = (d+32)%32 = d%32. So both groups access the same banks! But they are in different threads. Since there are 32 threads, and for a given `d`, 16 threads (part=0) access banks d%32, and 16 threads (part=1) also access banks d%32. This means two threads access the same bank per bank! Bank conflict: 2-way conflict per bank. Shared memory can broadcast if addresses are the same, but here addresses are different (different keys, same d_offset). So it's a 2-way bank conflict. The throughput halves for these accesses. That's not terrible. We can fix this by ensuring `d_offset` maps to different banks for the two parts. For example, we can pad the second half: `smem_k[16][64];` but if we change DT to e.g., 65, it's messy. Alternatively, we can change the thread mapping for S computation: use all 32 threads to compute a single query-key dot product? No. Or we can have threads with part=1 iterate over d in reverse order? Bank is `d % 32`. Reversing doesn't change bank. We could pad the shared memory array: `smem_k[16][64+1]`? No, address = (k*(DT+1) + d)*4. Bank = (k*(DT+1) + d) % 32. If DT+1=65, then bank = (k*1 + d) % 32 = (k+d)%32. For part=0, lane_id = k (0..15). d_offset = d. Bank = (k+d)%32. For part=1, lane_id = k+16. d_offset = d+32. Wait, k = lane_id % 16. So bank = (k + d_offset) % 32. For part=0: bank = (k+d)%32. For part=1: bank = (k + d + 32)%32 = (k+d)%32. Same conflict! To avoid conflict, we need `d_offset` for part=1 to map to different banks. If we use `d_start = part * (DT/2) + part`, i.e., part 0: [0..31], part 1: [33..64]? Then DT must be 65, not 64. Or part 0: [0..31], part 1: [32..63] but banks same. Actually if DT=64, the bank for `d` is `d%32`. The second half uses the same banks. A standard trick is to have `smem_k[16][64+4]` or something to shift banks for different rows. But for a given `d`, different `k` also map to same banks if k*64 is multiple of 32. k*64 % 32 = 0. So all keys use the same bank for the same d. That means for a fixed d, all threads access the same bank but different addresses (different k). This is inherently a 16-way bank conflict! Wait, this is critical. Address = base + (k*64 + d) * 4. Bank = ((k*64 + d)*4 / 4) % 32 = (k*64 + d) % 32. Since 64 % 32 = 0, bank = d % 32, independent of k. So for a given d, all threads (with different k) access the same bank! If all 32 threads are in the same warp and access different k but same d, this is a 32-way bank conflict (or 16-way if only 16 threads active). Shared memory can handle 2-way? Actually each bank can only serve one request per clock. So 32 threads accessing the same bank but different addresses results in 32 serialized accesses. That would be terrible! Wait, is that true? Let's check: For a 2D array `float smem[16][64]`, `smem[k][d]`. The address offset is (k*64 + d)*sizeof(float). Bank index = (offset / 4) % 32 = (k*64 + d) % 32. For fixed d and varying k, it is indeed all the same bank! This is a huge bank conflict. How do we fix it? We need to pad the leading dimension: `float smem_k[16][64+1]` or `float smem_k[16][64+P]` where P makes stride not multiple of 32. Address = (k*(64+P) + d)*4. Bank = (k*(64+P) + d) % 32 = (k*P + d) % 32 (if 64 is multiple of 32). So if we choose P such that it's odd (e.g., P=1), bank = (k + d) % 32. Then for fixed d and different k, banks are different! No bank conflict. Great. But this wastes a little shared memory: `16*65*4 = 4160 bytes` instead of 4096 bytes. Negligible. For `smem_q[8][64]`: address = (q*65 + d)*4? Actually if we also pad Q: `smem_q[8][65]`. Then bank = (q*1 + d) % 32 = (q+d)%32. For a warp, q is the same, so bank = d % 32. Threads in the warp access different d, so different banks. No conflict. Good. So we should pad the second dimension of smem arrays to avoid bank conflicts. This is a standard trick. Let's define: - `float smem_q[8][64+1];` or `float smem_q[8][68];` to align? Actually 65 is enough. But maybe using 68 for alignment? Access pattern is `smem_q[q][d]`. If `d` goes up to 63, and row stride is 65, the extra elements are unused. That's fine. - Similarly `smem_k[16][65];` and reuse for V. - `smem_o[8][1024+4]`? Actually O is accessed by `(q_idx*D + d)`. If we store `smem_o[8][1025]`, then bank = (q*1025 + d) % 32 = (q*1 + d) % 32 (since 1024 % 32 = 0). For fixed q (same warp), bank = d % 32. Each thread accesses a different d (d0 and d1 differ by 32). Bank(d0) = d0%32, Bank(d1) = (d0+32)%32 = d0%32. Both accesses by the same thread have the same bank but they are sequential (same thread). Different threads access different banks (if d0 differs). So no warp-level bank conflict. Good. But wait, in O stage, threads within a warp access `d0 = lane_id` and `d1 = lane_id + 32`. For a given warp, `d0` values are 0..31 across lanes, mapping to banks 0..31. `d1` values are 32..63 across lanes, mapping to banks 0..31. At any given clock, is there a conflict between a thread accessing d0 and another thread accessing d1? Yes, lane 0 accesses bank 0 (d0=0) and lane 0 also accesses bank 0 later. But concurrently within the warp, lane 0 accesses d0, lane 1 accesses d0 (bank 1), etc. At the same instruction, all threads access d0 (banks 0..31). No conflict. Next instruction (or next loop iteration), they access d1 (banks 0..31). Still no conflict because the order is serialized per instruction for the same thread? Actually a single thread accesses d0 then d1; these are sequential instructions, so not concurrent. So no conflict. Good. However, when writing back to smem_o, threads in a warp write to d0 (banks 0..31) simultaneously. No conflict. Good. For `smem_s[8][16]` (scores and p weights), address = (q*16 + k)*4. Bank = (q*16 + k) % 32. For a warp, q is fixed. k values differ (0..15 for active threads). Banks = (q*16 + k) % 32. These are 16 distinct banks. No conflict. So the only arrays that needed padding were Q and K/V because their column stride was multiple of 32. Let's pad them. Let's define `SmemCols = 65` for Q and K/V. For Q we only need 64 columns, so `smem_q[8][65]`. Actually we can define constants: `constexpr int S = 64; // actual tile size` `constexpr int PAD = 1;` `__shared__ float smem_q[BQ][S+PAD];` `__shared__ float smem_kv[BK][S+PAD];` `__shared__ float smem_s[BQ][BK];` Let's write the kernel with these. Wait, in the Q loading loop: `smem_q[q][d] = ...` If d goes from 0 to S-1, index is q*(S+PAD)+d. In the compute loop: `smem_q[q_local][d_offset]` This is fine. One issue: `__shfl_xor_sync` for `local_acc`. I need to make sure the mapping is correct. `int k_local = lane_id % BK;` `int part = lane_id / BK;` Assume BK=16. For lane_id in 0..15: k_local = lane_id, part = 0. For lane_id in 16..31: k_local = lane_id-16, part = 1. Wait, `lane_id / 16` gives 0 or 1. `lane_id % 16` gives 0..15. But `lane_id % BK` is fine if BK=16. For the shuffle, I want the two halves for the same key to exchange. Lane k and lane k+16 both have same k_local. `__shfl_xor_sync(0xFFFFFFFF, local_acc, 16)` exchanges lane i and i^16. So lane k gets from lane k+16. Correct. But what if BK is not 16? Let's keep BK=16. Wait, I used `lane_id % BK` for k_local. If BK=16, lane_id 0..31 gives k_local 0..15, 0..15. And `lane_id / BK` gives part 0 or 1. However, `lane_id / BK` with integer division for lane_id 16..31 gives 1. Good. What about `__shfl_down_sync` for max and sum? I used `__shfl_xor_sync` for max and sum over 32 lanes. Since only lanes 0..15 have valid data, and lanes 16..31 have -inf or 0, the result is correct. But `__shfl_xor_sync` will exchange data between lanes. For max, lanes 16..31 get max(valid, -inf) = valid after some iterations. All lanes end up with the same max. For sum, same. That's fine. But for storing p_weights to `smem_s`, I do `if (lane_id < BK)`. Since only lanes 0..15 write, it's fine. When reading `smem_s[q_local][k]` in O stage, all threads in warp read the same k? No, each thread loops over k from 0 to BK-1. So within a warp, all threads read the same `smem_s[q_local][k]` at the same time. This is a broadcast! Shared memory broadcast means all threads in the warp reading the same address are served simultaneously. This is very efficient. Great. So the loop `for (int k=0; k<BK; k++) { float p = smem_s[q_local][k]; ... }` is optimal. Similarly, accessing `smem_k[k][d0]` and `smem_k[k][d1]` where `d0`/`d1` differ per thread but `k` is the same across warp. Each thread accesses a different bank (since `d0` differs). No conflict. This kernel design actually looks quite solid. Let's write the kernel. We should also consider `constexpr int BQ=8, BK=16, DT=64`. Wait, I need to be careful about `smem_o`. Writing `smem_o` with float indexing. Since D=1024 may not fit in a 2D static array with large compile-time constants? `__shared__ float smem_o[8][1024+1];` size is 8*1025*4 = 32,800 bytes. That's fine. But can we allocate `__shared__ float smem_o[8][1025];`? 1025 is a compile-time constant. Yes. Actually, let me re-evaluate the number of query groups. With BQ=8 and N=512, qg goes from 0 to 63. That's 64 iterations. In each iteration, we initialize 32KB of smem_o. This takes some cycles but is minor. One more performance note: the loop over `qg` processes one contiguous block of 8 queries. Since Q is stored as [N][D], loading Q for a query group means reading 8*D = 8192 floats (32KB). This is sequential and coalesced across threads. Let's code it up. Wait, I should also consider using `float4` loads to increase memory bandwidth and reduce instructions. Since all dimensions are multiples of 4 (D=1024, DT=64), we can use `float4` for global loads and shared stores. For loading Q: `float4* Q_ptr4 = (float4*)(Q + head_offset + q_idx * D + dt * DT);` Each thread loads one float4 (4 floats). We need 8*64/4 = 128 float4s. With 256 threads, 128 threads load one float4 each, others idle. That's fine. Similarly for K and V: BK*DT/4 = 16*64/4 = 256 float4s. Exactly one per thread. For O write: 8*1024/4 = 2048 float4s. 256 threads * 8 iterations = 2048. Perfect. Using float4 loads/stores is a big win. Let's incorporate float4. For shared memory, we need to ensure alignment for float4. If we pad Q_smem to 65 floats, 65 is not a multiple of 4, so float4 addressing might cross row boundaries. Better to pad to a multiple of 4, e.g., 68. Then `smem_q[q][d]` address offset = q*68 + d. `d` up to 63. Float4 load at d=0,4,8...60. Since 68 is multiple of 4, row alignment is preserved. Bank analysis: bank = (q*68 + d) % 32. For a fixed q and varying d (0..63), bank sequence is q*4 + d (mod 32) since 68%32=4. For consecutive threads loading consecutive d (e.g., thread 0 loads d=0..3, thread 1 loads d=4..7 via float4), wait with float4, a single thread loads 4 consecutive floats. Banks for those 4 floats: d, d+1, d+2, d+3. Since the thread loads them sequentially, no bank conflict. For multiple threads, thread t loads float4 at offset `d = t*4` (if mapping that way). Banks for thread t: (q*68 + t*4) % 32, (q*68 + t*4 + 1) % 32, etc. The first element determines bank for the transaction? For 4-byte banks, a float4 load is a 16-byte transaction. It uses 4 consecutive banks. If consecutive threads load consecutive float4s (t*4 apart), the starting banks are q*4 + t*4 (mod 32). For threads t=0..7: banks 0,4,8,12,16,20,24,28. For threads t=8..15: banks 0.. again. Wait, t*4 mod 32 cycles every 8 threads. So thread 0 and thread 8 both start at bank 0 (if q=0). This is a bank conflict! Threads 0..7 and 8..15 would conflict on banks. Specifically, a float4 load from smem by a warp is split into two 128-bit transactions if needed? Actually shared memory has 32 banks, each 4 bytes. A float4 (16 bytes) always accesses 4 consecutive banks. Two threads accessing the same bank conflict. For `smem_q[q][d]` with row stride 68: address = (q*68 + d)*4. Float4 address = (q*68 + t*4)*4 for thread t if each loads a float4. Bank of first element = (q*68 + t*4) % 32 = (q*4 + t*4) % 32. For q=0: (t*4)%32. Threads t and t+8 have same starting bank. So threads 0..7 conflict with 8..15. Since each float4 accesses 4 banks, the conflict is 2-way for each bank? Let's see: thread 0 accesses banks 0,1,2,3. Thread 8 accesses banks 0,1,2,3. So 2-way conflict. Threads 1 and 9 access banks 4,5,6,7. 2-way conflict. So overall 2-way conflict for all banks. H100 shared memory can handle 2-way conflict with some throughput hit, but maybe not serialized fully. To avoid this, we want row stride to be an odd multiple of something? Ideally row_stride is not divisible by #banks * element_size? Wait, to avoid bank conflicts for float4 access across rows, we need consecutive threads accessing consecutive float4s to map to different banks for the first element. If each thread accesses `float4` at `d = t * 4`, bank = (row_stride*q + t*4) % 32. For fixed q, we want t*4 % 32 to be unique for each t in [0..31]? That's impossible because 32 threads * 4 bytes = 128 bytes per warp, but a warp can have 32 banks. If t varies 0..31, t*4 mod 32 has period 8. So there are always 4 threads mapping to the same first bank. Actually standard float4 access in shared memory has bank conflicts unless row_stride is odd relative to something. Standard trick: use row_stride that is a multiple of 32 + 1? Let's test: row_stride = 64 + 1 = 65. Then address offset = (q*65 + t*4)*4. Bank = (q*1 + t*4) % 32 = (q + t*4) % 32. Still same period 8. Actually for 1D array access where consecutive threads access consecutive float4 elements, you always get conflicts if the number of elements per row is even. Wait, in 1D global memory coalescing, it's fine because memory transactions are per warp. In shared memory, each bank is 4 bytes. A float4 (16 bytes) occupies 4 banks. When a warp accesses consecutive float4s, e.g., `smem[tid]`, where `smem` is `float4 smem[]`, then thread t accesses banks 4t, 4t+1, 4t+2, 4t+3 (mod 32). Thread t+8 accesses banks 4(t+8) = 4t+32 = 4t. So thread t and t+8 conflict on the same 4 banks! This is a 4-way bank conflict. Standard shared memory float4 arrays always have 4-way bank conflicts when accessed consecutively by thread index. However, if we access `float4` as `smem[ tid / 4 ][ tid % 4 ]` or something? Let's not overcomplicate. The smem loads and stores happen in the load loops. Since the load loop is `for (int i=tid; i < count; i+=blockDim.x)`, consecutive threads handle consecutive float4 elements, which in 1D shared memory is fine in terms of coalescing? No, in shared memory it's about bank conflicts. Using float4 for shared memory often causes 4-way bank conflicts but is still better than 4 separate float loads because each float4 load is a single instruction and the hardware might handle it more efficiently than 4 separate loads with 4-way conflicts? Actually 4 separate float loads by consecutive threads have no bank conflict (each thread accesses a different bank). A single float4 load by a warp results in 4-way bank conflicts because each thread accesses 4 consecutive banks, and threads t and t+8 access the same set of 4 banks. However, some GPUs can broadcast or handle this with some throughput reduction. Alternatively, we can simply use `float` loads and rely on the compiler. For simplicity, I will use float loads. Avoid float4 complexity. The compiler might vectorize automatically. Let's just write scalar loads. It's easier to get correct. If performance is bad, we can revisit. But scalar loads with 256 threads and good memory patterns might be okay. Let's finalize the kernel with scalar loads. One more detail: `expf` can be slow. We compute `expf` for each key (16 per query) per key tile (32) per query group (64) = 16*32*64 = 32768 expf per block. With 1024 blocks, total 33M expf. That's fine. We could use `__expf` for faster approximate exponential if tolerance allows. 1e-4 tolerance might allow `__expf`? It has slightly lower precision but might be okay. However, `expf` is likely fast enough on H100. Let's write the full Python code. Wait, there is a potential correctness issue: the softmax normalization factor `scale` should be multiplied on scores. I placed `s_val * scale` when storing to smem_s. That's correct. Another issue: when I compute `sum0` and `sum1` in O update, I use `smem_k` which at this point is loaded with V. I reuse `smem_k` array but the variable name is `smem_k`. In code I can `#define smem_v smem_k` or just use the same name. Code structure in C++: ```cpp #include <torch/extension.h> #include <cuda_runtime.h> #include <math.h> #define BQ 8 #define BK 16 #define DT 64 #define PADDED_DT 68 __global__ void fused_attention_kernel(...) { // ... } torch::Tensor forward(torch::Tensor Q, torch::Tensor K, torch::Tensor V) { // check sizes, contiguous auto O = torch::empty_like(Q); int B = Q.size(0); int H = Q.size(1); int N = Q.size(2); int D = Q.size(3); float scale = 1.0f / sqrtf((float)D); dim3 grid(B, H); dim3 block(256); fused_attention_kernel<<<grid, block>>>( Q.data_ptr<float>(), K.data_ptr<float>(), V.data_ptr<float>(), O.data_ptr<float>(), B, H, N, D, scale); return O; } ``` We also need to ensure tensors are on CUDA and contiguous. We can use `.contiguous()` in Python wrapper, or in C++ check strides. Since the reference implementation likely passes contiguous tensors from `torch.rand`, they are contiguous. But to be safe, we can call `.contiguous()` in Python. Actually, `load_inline` auto-generates a pybind wrapper for `forward`. In our Python `ModelNew`, we will call the extension's `forward` function. One concern: the kernel might take a while to compile due to `__syncthreads()` in loops and unrolls. But 256 threads and 5 nested loops is fine for nvcc. Let's consider warp divergence. The only divergence is `if (lane_id < BK)` for writing to smem_s. This means 16 threads active and 16 threads idle. That's 50% utilization for that instruction. On H100, this is not a big deal because the warp is fully active, just some threads masked. Another thing: `for (int d = 0; d < DT / 2; d++)` in S stage is unrolled? We can add `#pragma unroll`. For DT/2=32, unrolling 32 iterations is fine. `#pragma unroll` for the inner D loop and the K loop (BK=16) will help. Let's write the kernel now. Wait, I realize `__shfl_xor_sync` for `local_acc` might not yield the correct pair if BK is not 16. I fixed BK=16. If lane_id 0..31, `lane_id % 16` is 0..15, 0..15. `lane_id / 16` is 0 or 1. `__shfl_xor_sync(0xFFFFFFFF, local_acc, 16)` exchanges lane i and i+16 (XOR 16 flips bit 4). So lane k (0..15) gets from lane k+16. Both have same `k_local`. Perfect. Now, after `s_val = local_acc + __shfl_xor_sync(...)`, both lanes 0..15 and 16..31 have the same `s_val`. We only need to store from one of them. `if (lane_id < BK) smem_s[warp_id][lane_id] = s_val * scale;`. That's fine. One more thing: what if the number of queries N is not divisible by BQ? Here N=512, BQ=8, divisible. What if D is not divisible by DT? D=1024, DT=64, divisible. So hardcoding BQ, BK, DT for these specific dimensions is acceptable, but to be robust we can use templating or dynamic loops based on arguments. The reference dimensions are fixed: N=512, D=1024. But we should write the kernel generally using `N` and `D`. However, if we use dynamic loops with `#pragma unroll`, the compiler might not know the bounds. Since the dimensions are fixed in this problem, we can make them template parameters or `#define` them. But for safety, let's derive loop bounds from arguments. Within the kernel, `N / BQ` should be integer? We can assert or use loops with `for (int qg = 0; qg < N; qg += BQ)`. Similarly for `kt` and `dt`. This is more robust. Let's write it with `for (int qg_start = 0; qg_start < N; qg_start += BQ)` etc. But the smem_o is sized `BQ x (D + PAD)`. D is up to 1024. We can size smem_o to max, e.g., `smem_o[8][1024 + PAD]`. That works as long as D <= 1024. Or we can use dynamic shared memory. But static is fine. Wait, if D is a compile-time unknown, `smem_o[BQ][D + PAD]` is not allowed. We can use a 1D array: `__shared__ float smem_o[BQ * 1024];` or dynamic extern. But since D is known at runtime but constant for this problem, we can just use a large enough static size. However, to be safe, we can allocate `smem_o` as dynamic shared memory using `extern __shared__ float smem_o[];` and manage offsets manually. That would be cleaner for generic D. But dynamic shared memory affects kernel launch (we need third arg). We can compute size and pass it. It's fine. But then we need to manage smem_q, smem_k, smem_s, smem_o offsets. Actually, static smem is fine. Let's just set `MAX_D = 1024` and `MAX_BQ = 8`. Then `smem_o[8][1024 + PAD]`. The compiler allocates 8*1028*4 = ~32KB. Good. For `smem_q`, size is `MAX_BQ * (MAX_DT + PAD)`. We can set `MAX_DT=64`. `smem_q[8][68]`. Good. For `smem_k`, `MAX_BK * (MAX_DT + PAD)`. `MAX_BK=16`. `smem_k[16][68]`. Good. Since the actual dimensions match, we are safe. But the loops will use `dt < D / DT` which is 16. That's compile-time not known, but it's fine. Let's just write it with static arrays and constant sizes, but use `D` for indexing global memory. This is fine. Wait, there is a subtlety: `smem_o` is declared as `__shared__ float smem_o[BQ][1028];`. When indexing `smem_o[q][d]`, if d goes up to D-1 (1023), it's fine. But if D were larger, we'd be out of bounds. But problem dims are fixed. Now, I should think about whether the O_smem is actually needed. After computing all key tiles, we write O to global memory. Between key tiles, O is updated via `smem_o`. If we don't use smem_o but instead accumulate in a register array per thread for a D tile, and then write to global memory after each key tile update, we'd need global memory atomics or read-modify-write. Because different key tiles update the same O elements. Read-modify-write per key tile is bad (global HBM traffic for O becomes 2 * N_K_tiles * size(O) = 64 * 2MB = 128MB per head, all to HBM). That's much worse. So smem_o is essential. Validation of the algorithm on a simple case: N=4, D=4, BQ=2, BK=2, DT=2. Q = [[1,0,0,0],[0,1,0,0]] (queries) K = [[1,0,0,0],[0,1,0,0]] (keys) V = [[1,2,3,4],[5,6,7,8]] Expected: S = [[1,0],[0,1]]. Softmax(S) ~ [[0.731,0.269],[0.269,0.731]]. O0 = 0.731*[1,2,3,4] + 0.269*[5,6,7,8] = [0.731+1.345, 1.462+1.614, ...] Let's trace online softmax for query 0: Key tile 0 (keys 0,1): S = [1,0]. m_tile = 1. l_tile = exp(0) + exp(-1) = 1 + 0.3679 = 1.3679. m_prev = -inf. m_new = 1. corr_prev = 0. corr_tile = 1. l_new = 1.3679. p_weights = [1, 0.3679]. O update: O[0] = 0 + 1*1 + 0.3679*5 = 2.8395 O[1] = 0 + 1*2 + 0.3679*6 = 4.2074 etc. No more key tiles. Final O = O / l = [2.8395/1.3679, 4.2074/1.3679, ...] = [2.076, 3.076, ...] Now let's compute expected: 0.731*1 + 0.269*5 = 0.731 + 1.345 = 2.076. Correct! 0.731*2 + 0.269*6 = 1.462 + 1.614 = 3.076. Correct. Great, algorithm is correct. Now, what about the scale factor? I multiply `s_val * scale` before storing. This means all further calculations use scaled scores. That's correct. One more thing: `expf` arguments can be large negative or positive? Scores are dot products of random [0,1] vectors of length 1024. Mean dot product ≈ 1024 * 0.25 = 256. Scale = 1/sqrt(1024) = 1/32. So scaled scores ≈ 8. The range is maybe [0, 32]. `exp(32)` is huge. But with random uniform [0,1], typical max is near 32. `exp(32)` is about 7.9e13, which fits in float32 (~3.4e38). But subtracting max before exp prevents overflow. We do `exp(my_s - m_val)` where `m_val` is max over the tile. For the first tile, `m_val` is the local max. Later tiles, `m_new` is global max. So `exp(...)` is always <=1. No overflow. Good. However, `exp(m_prev - m_new)` can be very small. If `m_new` is much larger than `m_prev`, this approaches 0. That's fine. Underflow to 0 might happen but is harmless. Let's implement in C++. We should also include `constexpr float NEG_INF = -1e20f;` or use `-INFINITY`. Wait, we should also consider using `__ldg` for read-only global memory loads. `Q`, `K`, `V` are read-only. We can add `__restrict__ const` and optionally use `__ldg`. Modern compilers often emit ld.global.nc automatically, but explicit `__ldg` can help. However, `__ldg` returns a float, and it's fine. Also, should we unroll loops? The inner `d` loop (32 iterations) is good to unroll: `#pragma unroll 8` or full unroll. NVCC full unroll of 32 is fine. The `dt` loop (16 iterations) could be unrolled: `#pragma unroll 4`. The `kt` loop (32) maybe not unrolled. The `qg` loop (64) not unrolled. The `k` loop in O update (16 iterations) should be unrolled: `#pragma unroll`. Let's also think about using `warp_reduce` inline functions. Implementation details for C++ wrapper: ```cpp #include <torch/extension.h> #include <cuda_runtime.h> #include <cuda_fp16.h> #include <math.h> ``` Wait, we are writing a Python file that wraps this. We need to be careful with string formatting. Let's draft the code now. I'll write the Python module. Actually, I'll just write the kernel and test it via compile and correctness.

Good, it compiled. Now let's check correctness.

The first version passed correctness. Now I need to think about performance. The problem size is: - B=32 - H=32 - N=512 - D=1024 With my custom kernel: - BQ=8, BK=16 - 512/8 = 64 query tiles per head - 512/16 = 32 key tiles per head - Each block processes BQ=8 queries - BlockDim=256 (8 warps) - Grid = (32, 32) = 1024 blocks total - Each block: D=1024, dt=64, so 16 iterations for dot-product, 16 for V update - Shared mem: approx 8*(1024+4)*4 = ~32KB for smem_o alone, plus others. Total around 40KB. This might be OK but there are some concerns: 1. Only 8 warps, meaning occupancy might be OK but threads aren't super utilized 2. The smem_o is 8x1024 = 8192 floats, that's a lot of shared memory. At 1024 blocks total, with shared mem per block ~40KB, that's about 40MB total dynamic, but occupancy is limited to maybe 2 blocks per SM = 2*40KB = 80KB per SM, which is fine (limit is 227KB). 3. But the main issue: the memory access pattern. Each block loads Q once and reuses across all key tiles. K and V are loaded N/BK = 32 times each. Total memory traffic is similar to standard attention. However, PyTorch's scaled_dot_product_attention is likely using cuDNN/FlashAttention on H100, which is highly optimized with Tensor Cores (even in FP32 it might use TF32 tensor cores or other optimizations). My custom kernel is using only FP32 scalar operations, so I'm unlikely to beat cuDNN unless I also use Tensor Cores. Actually, wait. The reference uses `torch.nn.functional.scaled_dot_product_attention`. On H100, this will dispatch to FlashAttention if available, cuDNN SDPA, or efficient attention. That will be VERY hard to beat with a simple CUDA kernel, especially without Tensor Cores. But I can try to use shared memory more effectively and perhaps optimize the kernel for this specific shape. Let me think... Actually, for fp32 on H100, FlashAttention might not be called (it often uses bf16/fp16). For fp32, it might use the math backend or cuDNN attention. But cuDNN attention for fp32 should still be quite fast. Wait, the task says "Run faster than the PyTorch reference." It might be possible that the reference is using a non-optimal backend for fp32, or that I can beat it by doing something clever. Actually, looking at the problem size again: - Q/K/V size: 32 * 32 * 512 * 1024 * 4 bytes = 2GB - That's a lot of data - The attention computation is memory bound For memory bound operations, I need to minimize memory traffic. Floppy attention minimizes memory traffic by focusing on SRAM. But the problem size (D=1024, N=512) means the standard algorithm might have high memory traffic. Wait, N=512 is relatively small. D=1024 is large. - Each Q row has 1024 floats = 4KB - For all 512 queries in a head: 2MB - Each head processes this... Actually, with N=512 and D=1024, for standard attention: - Q·K^T: each query dot products with all keys. - Memory reads per head: Q (2MB) + K (2MB) + V (2MB) + output (2MB) = 8MB if we compute exactly once. - But for the matrix multiply, we read K N/BK times or similar. For FlashAttention style, the main reduction is in writes to O: only one write per element instead of multiple writes through softmax. But for D=1024, the intermediate activations might be harder to fit. With my kernel, each block loads: - Q: 8 rows * 1024 cols = 8K floats = 32KB (loaded once per block) - K: 16 rows * 1024 cols = 16K floats = 64KB (loaded 32 times) = 2048KB - V: same as K, 2048KB - O: written once, 32KB Total traffic per block = ~4.1MB 1024 blocks * 4.1MB = 4.2GB total memory traffic. Wait, but K/V are shared across all query groups within a head? Actually no, in my kernel each block only processes 8 queries. So for all 64 query groups in a head, each reads the full K and V. So total K reads per head = 64 * 2MB = 128MB. Similarly V reads = 128MB. Q reads = 2MB once. Output = 2MB once. Total = 260MB per head * 1024 heads = 260GB? That can't be right. Let me recalculate: B=32, H=32, so total heads = 1024. Per head: Q = 512*1024*4 = 2MB, K = 2MB, V = 2MB, O = 2MB. In my kernel, per block (8 queries): - Q: 8*1024*4 = 32KB (loaded once) - K: 2MB * (512/16) iterations? Actually K is loaded tile by tile: 16*1024*4 = 64KB per tile, and there are 32 tiles, so 2MB total K reads per block. - V: same 2MB total reads per block. - O: 32KB written once. Total per block = 4.1MB. Total for all blocks (1024) = 4.2GB. Wait, but K and V are the same across all query blocks! So my kernel is reading K and V way too many times. In a well-optimized attention kernel, you'd want to load K and V once and compute dot products with multiple query tiles, but with D=1024 and only ~227KB shared memory, and 1024 features, you can't fit a whole K/V head in shared memory. Actually, wait. The standard flash attention loads k tiles and keeps them in shared memory while iterating over query tiles. Wait, no - the standard FA loads Q tiles and keeps them in register, then iterates over K/V tiles. Or it loads Q tiles in SRAM and iterates over K/V tiles. But optimized implementations use fused MMAs and load K/V through L1 cache, benefiting from cache hits when multiple blocks process different query tiles but load the same K/V tiles. On H100, with: - 1024 blocks total - 132 SMs - 1024/132 ≈ 7.7 blocks per SM But each block reads the full K and V for its head. Each SM might process blocks from different heads (with different K/V), so L1 cache effectiveness depends on whether blocks from the same head are co-located on the same SM. Actually, with a grid of (B, H) = (32, 32), blocks are launched in x-major order. All 32 blocks for b=0, h=0 to h=31, then b=1, etc. Or maybe y-major? blockIdx.x is fastest, so blocks (0,0), (1,0), ..., (31,0), (0,1), ..., (31,31). So within a wave of 132 SMs, we process all 32 batches for h=0, then some for h=1. So blocks from the same head may be quite spread out. Hmm, this is hard to optimize. Alternative approach: for N=512, D=1024, maybe I can transpose the loops? Let me think about using Tensor Cores... Actually, H100 has very fast Tensor Cores. Even for FP32, using Tensor Cores properly could be beneficial. But writing a Tensor Core kernel from scratch for attention in C++ CUDA is very complex. Wait, let me think about the reference. The PyTorch SDPA reference on H100 for FP32 might not even use Tensor Cores. Let me check what backends are available: - For FP32, SDPA often falls back to the math backend (plain cuBLAS) if `is_causal=False`, or it might use cuDNN attention. Actually, on H100, anything built into PyTorch would be well optimized. Beating it with a simple hand-written kernel is hard, but not impossible if the reference uses a less optimized backend for some reason. Let me reconsider the kernel. Maybe I should use a simpler approach: since N=512 and D=1024, I could compute S = Q @ K^T (32*32 blocks of 512x512), then softmax, then S @ V. For the matmuls I could call cuBLAS from the kernel? No, from cuda source I can call cublas if I link it, but load_inline might not link cublasSgemm by default. Actually, I can specify extra libraries. But the kernel overhead of separate kernel calls for matmul + softmax + matmul might be high. Still, cuBLAS matmuls are extremely optimized. Alternative: use a single kernel but write it better. Let me think about using more threads and better memory patterns. What if I use a warp-level approach where each warp computes one output element? For D=1024, with 32 threads per warp, each thread does 1024/32 = 32 MACs. That's feasible. Actually, a common approach for efficient attention on small sequence lengths is to have each thread compute a single attention output. But with D=1024, we might need multiple queries per block. Let me try a different approach: use `mma` via Tensor Cores. But that's very complex for inline CUDA. Alternatively, I can use `wmma` or the newer `mma.sync.aligned` PTX. With Hopper, there are specialized warpgroup-level matrix multiply instructions. But again, this is extremely complex for a single file submission. Wait, can I just use cuBLAS within my extension? Yes! I can create a CUDA extension that calls cuBLAS for the matrix multiplications, and write custom CUDA for softmax. This would likely be very fast. Let me think: 1. Compute S = Q @ K^T for all heads: this is a batched matrix multiply of [BH, M, K] @ [BH, N, K]^T -> [BH, M, N] 2. Apply softmax row-wise 3. Compute O = S @ V: [BH, M, N] @ [BH, N, K] -> [BH, M, K] For Q, K, V shapes: [32, 32, 512, 1024] We can reshape to [32*32, 512, 1024] = [1024, 512, 1024] Then S = [1024, 512, 512] O = [1024, 512, 1024] Memory usage: - Q: 32*32*512*1024*4 = 2GB - K: 2GB - V: 2GB - S: 32*32*512*512*4 = 1GB Total intermediate = 1GB. That's a lot but H100 has 80GB so it's fine. The overhead is one extra 1GB read/write. But cuBLAS GEMM is extremely fast on H100, likely saturating memory bandwidth or compute. However, PyTorch's SDPA probably does the exact same thing: call cuBLAS for the matmuls, but fused with the softmax in between. Actually, SDPA likely uses the "flash attention" algorithm where intermediate S isn't materialized to global memory. If I materialize S to global memory, I'll be doing 1GB extra read/writes. That might hurt performance. But wait - on H100, with 3.35 TB/s bandwidth: - Reading Q+K+V once: 6GB - Writing O once: 2GB - Extra S read/write: 2GB Total: 10GB vs 8GB for fused. 4GB * 1024/3.35TB/s ≈ 1.2 ms difference in bandwidth-bound time. But the compute savings might help. Actually, PyTorch SDPA for FP32... let me check if it might be using efficient_attention kernel from xFormers or something. On recent PyTorch with H100, `scaled_dot_product_attention` with FP32 and no mask might call cuDNN attention or flash attention. If it uses cuDNN attention, it's extremely hard to beat. However, maybe I can write a better fused kernel using proper attention algorithm but with better parallelism. Let me reconsider my first kernel. The issue is: 1. Only 8 warps per block 2. warp_id = tid/32, so warp_id can be 0..7 3. q_local = warp_id (0..7) which maps to one of the BQ=8 queries 4. lane_id < BK=16 means only 16 threads per warp are used 5. The remaining 16 threads are wasted for all score computations and softmax 6. That's only 50% thread utilization for the critical path This is a major issue. Let me redesign to use all 32 threads. New design: - Process 32 queries per block? Or process 16 queries but use all lanes? - With D=1024, let's keep BQ and BK smaller. Alternative approach inspired by FlashAttention: - Load Q tile into smem (or registers) - Iterate over KV tiles - Use all warps to parallelize over queries and keys What if: - Each block processes BQ queries (e.g., 16) - Each warp processes one query? No, 16 warps = 512 threads, which might be too many. - Or each warp processes part of the score matrix. Actually, for online softmax + matrix multiply, a common efficient approach: - Load Q tile: 16 rows x 1024 cols - Each thread computes partial dot products for one (q, k) pair - With 256 threads, 16 queries, 512 keys -> 16*512 = 8192 elements in score matrix. 8192/256 = 32 elements per thread. Each element needs 1024 MACs. 32*1024 = 32768 MACs per thread. At 50 TMAC/s on FP32... that's too much for one thread! Wait, I mean each key tile is 16 keys (BK=16). Each iteration of the outer loop processes one KV tile. Let me try: - BQ = 16, BK = 16 - Block = 256 threads = 8 warps - Each warp handles 2 queries? Actually let's reconsider. For score computation: - We need S = Q_tile @ K_tile^T where Q_tile is BQ x D and K_tile is BK x D - Result is BQ x BK = 256 elements - With 256 threads, each thread computes exactly one dot product of D=1024... that's 1024 MADs per thread per tile, very slow! Alternative: each thread computes partial dot product, then reduce. - Each thread handles 1 element in D dimension - For a (q, k) pair: 1024 threads would each multiply one pair, then reduce with warp shuffle / shared memory. With 256 threads, we'd need 1024/256 = 4 accumulations per thread for each (q, k). But there are 256 (q, k) pairs. Total parallel work = 256 pairs * 1024 elements = 262k multiplications per tile. With 256 threads, each thread does 1024 multiplications, then we need reductions. Wait, here's a better way: - Load Q_tile and K_tile into shared memory - For each query in Q_tile: - Have all threads compute its dot product with all keys in K_tile in parallel - Specifically, each warp computes scores for 1 query across 32 keys? But BK=16, so only 16 keys. Actually, let me think differently. Let's transpose: have BQ=64 but each query processed by fewer threads? No... What about: - BQ=32 queries - BK=32 keys - Block=256 threads - D=1024 - For dot product, use all 256 threads to compute partial dot products for 32 queries: - Each thread handles D/256 = 4 features for all queries - Wait, this doesn't parallelize across K. Actually, for a block of 32 queries x 32 keys x 1024 dims: - Total work: 32*32*1024 = 1M multiplications - 1M / 256 threads = 4096 mults per thread per tile - Number of key tiles: 512/32 = 16 - Total mults per thread: 65536 - At say 20 Tops FP32 on normal CUDA cores (not tensor cores), that's about 3.3 microseconds of compute per head... negligible? Actually 65536 mult-adds / (67e12 FLOPS / 132 SMs) ≈ very small. The bottleneck is memory. Hmm wait, 67 TFLOPS is the full chip FP32 throughput. 67e12 ops/sec. 65536 ops is 1 nanosecond. But memory is much slower. So the kernel is definitely memory bound. For memory-bound attention with D=1024, N=512: - Total data movement for fused kernel = roughly: load Q once, load K N/BK times... no, load K once per key tile per query tile. - In Flash Attention, each query tile loads all K tiles. To reduce this, you'd want larger query tiles to amortize K/V reads. With my kernel: - BQ=8, so K is read 64 times per head - Total K reads = 64 * 2MB = 128MB per head - Total for all heads = 128MB * 1024 = 131GB - Similarly for V - Q reads: 2MB * 1024 = 2GB - O writes: 2GB - Total traffic ≈ 264GB - At 3.35TB/s = 3350GB/s, time = 79ms. That's way too high! Wait, but that's not right because: 1. K and V are cached between query tiles if multiple blocks from the same head run on the same SM 2. Also, with HBM3 on H100, maybe actual sustained bandwidth is different 3. But 264GB for one forward pass on 2GB inputs seems plausible for naive attention Let me recalculate more carefully. The standard attention complexity is O(N^2*D). The memory operations if we materialize nothing: - Q: read once, K: read N/BK times per query block... but in theory, if we have enough cache, each K is read once and stays in L2. Actually, H100 has 50MB L2 cache. Per head, K = 2MB. So we can fit about 25 heads in L2. With 1024 heads total and ~132 concurrent blocks, if each block is from a different head, L2 caching doesn't help. But if multiple blocks from the same head are running concurrently on different SMs, they all compete for L2 cache but each reads the same K, so L2 can serve it without going to HBM after the first read. Actually, in a wave of 132 blocks, if they're all from different heads, then K isn't shared. But the next wave will also be from different heads. So every K read goes to HBM. This is terrible. PyTorch's FlashAttention would be much better because it: 1. Uses much larger tile sizes, amortizing K/V loads 2. Or uses the online softmax algorithm without materializing S to global memory (which my kernel also does, but PyTorch's does it better) 3. And it uses Tensor Cores, which are much more efficient for the GEMM-like operations Wait, but if PyTorch SDPA uses cuDNN attention, cuDNN may also materialize some intermediates internally in FP32. Or it might also be highly optimized. Actually, wait. If PyTorch uses SDPA for FP32, it might actually dispatch to the `math` backend or a less optimized kernel because FP32 flash attention often isn't available. Let me check: In PyTorch, `scaled_dot_product_attention` has the following backends: - `flash_attention` (typically FP16/BF16 only, not FP32) - `mem_eff_attention` (FP16/BF16/FP32, but for FP32 it might be less optimized) - `cudnn_attention` (available for FP16/BF16 on Ampere+, maybe FP32 too?) - `math` (always available) For FP32 input on H100, it will likely try cuDNN first, then memory efficient attention, then math backend. cuDNN attention for FP32 might still be available. If cuDNN attention is used, it's almost impossible to beat. But wait, the correctness test passed. Let me think about whether I can do something simpler that actually performs better. Actually, let me think about this more carefully. The trick in these kinds of challenges is often that the reference implementation doesn't use the best available backend because of some configuration issue, or because FP32 SDPA doesn't dispatch to the fastest path. For instance, if cuDNN attention requires FP16/BF16 for its optimized path, then the FP32 reference might fall back to `memory_efficient_attention` or even `math`. The `math` backend computes `softmax(Q @ K^T / sqrt(d)) @ V` using explicit PyTorch operations, which involves materializing the N x N attention matrix to memory. If the reference materializes S = Q @ K^T (1GB), then does softmax, then O = S @ V, that would be: 1. S = matmul(Q, K^T): read Q (2GB) + K (2GB) + write S (1GB) = 5GB 2. softmax(S): read S (1GB) + write S (1GB) = 2GB 3. O = matmul(S, V): read S (1GB) + V (2GB) + write O (2GB) = 5GB Total: 12GB My fused kernel avoids materializing S, so even with poor reuse, it might be faster if the reference does materialize S. But if the reference uses flash attention (even in FP32 with some support), it would be very fast. Hmm, let me try my kernel and see if it's actually faster. But I can't run timing directly - I can only submit. However, I can reason about it. Actually, I realize my kernel has a bug in the soft max update with corr_prev. In the kernel I wrote: ```cpp float p_weight = exp_val * corr_tile; ``` Wait, in online softmax, when updating O accumulator, we should use: ```cpp float corr_prev = expf(m_prev - m_new); float p_weight = expf(my_s - m_new); ``` Currently I compute: ```cpp float p_weight = exp_val * corr_tile; ``` where `exp_val = expf(my_s - m_val)` and `corr_tile = expf(m_val - m_new)`. So `p_weight = expf(my_s - m_new)`. That's correct! And `corr_prev = expf(m_prev - m_new)`. And O is updated with `corr_prev * old + sum(p_weight * V)`. That's correct. OK, the kernel is numerically correct. Now about performance... Actually, I just realized the kernel submission passed correctness, but it's probably quite slow. Let me try to optimize it. Main optimization ideas: 1. Use all 32 threads in each warp effectively 2. Increase tile sizes to reduce key tile count (fewer K/V loads) 3. Consider coalescing and vectorized loads For N=512, D=1024: - If we can load 32 keys at a time (BK=32), that's only 16 key tiles instead of 32. This halves K/V reads. - But shared memory for K: 32*1024*4 = 128KB, which fits in shared memory (227KB limit). - Q tile: 8*1024*4 = 32KB. - O tile: 8*1024*4 = 32KB. - p tile: 8*32*4 = 1KB. - Total: ~200KB, under the 227KB limit. Let's try BK=32. With BK=32 and block=256 threads (8 warps): - Each warp should process some queries. With BQ=8, we have exactly 8 warps, so each warp handles 1 query. - For score computation with BK=32, only lane_id < 32... wait, all 32 lanes are used since BK=32. Great! - For the dot product, each thread handles part of D=1024. With DT=64 and 2 parts (lane_id/BK = lane_id/32, that's 0 or not... wait, BK=32 means lane_id/32 = 0 for all lanes. So we need a different way to split D. Actually, with BQ=8, BK=32, block=256: - There are 8 warps. - Each warp handles 1 query, and all 32 lanes handle the 32 keys. - For dot product over D=1024: - We need to split the 1024 features among threads. - One approach: each thread computes 1024/32 = 32 features for its 1 query x 1 key pair? No, there are 32 keys, but each lane only handles 1 key. Wait, each warp handles 1 query (q_local = warp_id) and 32 keys (k_local = lane_id). So each warp computes 32 dot products, one per lane. Each dot product is D=1024. To parallelize the dot product, we can have each lane process only part of D, then shuffle reduce. For example: - load Q for query q_local into shared memory - load K for all 32 keys into shared memory - each lane computes its own dot product for key k_local - To parallelize within lane: each lane does partial dot product over a chunk of D - But a lane is a single thread. We can load D in chunks of DT=64. - Lane k_local accesses K[k_local][d] and multiplies by Q_shared[q_local][d]. - This is just a scalar loop over D. Wait, but each lane is computing its own dot product serially over D=1024. That's 1024 iterations, each doing one multiply-add. With ~32 FMAs per clock per thread on H100, and ~2GHz clock, that's about 16 cycles per FMA (throughput limited). Actually FP32 FMA throughput is 1 per thread per clock on Ampere/Hopper for simple FMAs (when not using tensor cores). At 2GHz, 1024 FMAs take ~500 ns. That's quite fast. Then with 16 key tiles (512/32), each warp/ thread does 16*1024 = 16384 FMAs. At 2GHz, ~8 microseconds. Not too bad. For the V update: - Each query needs O = P @ V where P is 512 keys, V is 512x1024. - With BK=32, we process 32 keys at a time. - For each V tile, we need to multiply P_tile (1x32) by V_tile (32x1024) and add to accumulator. - Our warp is computing for 1 query. - Lane k handles key k. P[q][k] is loaded. - Then for each output dimension d, we accumulate P[q][k] * V[k][d]. - With 32 lanes and D=1024, we need all lanes to help. But a single warp only has 32 threads. Hmm, we could have lane d (for d < 32) accumulate over k=0..31 for V[k][d] and then process multiple chunks of D. For D=1024, each lane handles 1024/32 = 32 dims. So lane i accumulates for dim i + 32*j where j=0..31. Wait, that's a lot of unrolling. Let me think: - Load P[q][k] for k=0..31 (in shared memory from earlier step) - Load V[0..31][0..1023] into shared memory (or directly access) - Each of the 32 lanes computes a partial sum for a block of 32 dims. - Specifically, lane d computes sum over k of P[q][k] * V[k][d_base + d] for a d_base. With D=1024 and 32 lanes, we can process 32 dims at a time. - iteration 0: dims 0..31, lane d handles dim d - iteration 1: dims 32..63, lane d handles dim d+32 - ... - Total: 1024/32 = 32 iterations. In each iteration, each lane does 32 multiply-adds (over k=0..31). So 32*32 = 1024 FMAs in 32 iterations. Again, 16 key tiles * 1024 FMAs = 16384 FMAs. Plus score computation: 16384 FMAs. Total ~32768 FMAs per warp per head. 8 warps * 32768 = 262k FMAs per block. Total blocks = 32*32*64 = 65536 blocks? Wait: B=32, H=32, N/BQ = 512/8 = 64. Total blocks = 32*32*64 = 65536. 65536 blocks * 262k FMAs = 17.2e9 FMAs total. At 67 TFLOPS, that's 257ms. But memory bandwidth might dominate. Actually, 65536 blocks is a lot. At 256 threads each, that's 16.7M threads. With 132 SMs at ~2000 concurrent threads each, occupancy is fine. But 65536 blocks means each K tile is loaded 65536 times. Hmm, no: for each head, there are 64 query blocks, and each loads all 16 key tiles. So per head: 64*16 = 1024 K tile loads. Total heads = 1024. Total K tile loads = 1,048,576. Each tile is 32*1024*4 = 131KB. Total K traffic alone = 1,048,576 * 131KB = 137GB. That's massive. 137GB / 3.35TB/s = 41ms just for K reads. That's actually not terrible, but there's also V reads and Q reads. So total time could be 80-100ms. Wait, but the reference might be doing something similar or worse. Actually, let me reconsider whether PyTorch's scaled_dot_product_attention uses FlashAttention for FP32. Looking at PyTorch code, FlashAttention in PyTorch typically only supports FP16 and BF16. For FP32, it likely falls back to either memory-efficient attention or the math backend. Memory-efficient attention (the xFormers-style one) does fuse the operations to avoid materializing S, but might still use a similar algorithm to mine, possibly worse. Let me actually check: can I write a kernel that's competitive? Given the constraints, maybe the simplest winning strategy is: 1. Keep the current fused kernel approach (no materialized S) 2. Optimize thread utilization and tile sizes 3. Try to make it faster But there's another trick: since D=1024 and N=512, I can fit the entire S matrix for a query tile in shared memory or registers? S for BQ queries and all N keys would be BQ x N = 8 x 512 = 4096 floats = 16KB. That's small. Actually, what if each thread block processes ALL queries for a head at once? - BQ=512, but we process them in chunks for shared memory. - Q for all 512 queries: 2MB. Too big for smem. - But we can load Q to registers one query at a time. Alternative: Each thread computes one complete attention output vector. - There are B*H*N = 32*32*512 = 524,288 output vectors. - Each output vector is D=1024 floats. - Total threads can work in parallel. - Each thread computes softmax attention weights for one query. - Steps for one thread (one query): 1. Load Q vector (1024 floats) into registers/local memory. 2. Iterate over all keys (512 keys): - Load K vector (1024 floats). - Compute dot product Q·K (1024 mult-adds). - Track max and exp-sum for softmax online. 3. Iterate over all keys again: - Load K and V vectors. - Compute weight = exp(Q·K - max) / sum - Accumulate weight * V to local output buffer. 4. Write output. This requires loading K twice and V once, but avoids any inter-thread synchronization. - It also avoids shared memory entirely. Each thread does: - Read Q: 1024 floats - Read K: 512 * 1024 floats - Read V: 512 * 1024 floats - Write O: 1024 floats Total per thread = 1025K floats = 4.1MB Total for all threads = 524288 * 4.1MB = 2.1TB. That's way too much global memory traffic. But wait, when multiple threads access the same K/V from global memory, the GPU coalesces and caches. If all threads in a warp access the same key (for different queries or different features), that would be broadcast. But if they access different keys, it's coalesced if they're contiguous in memory. Actually, let's organize threads differently: - A warp of 32 threads processes 32 different queries, and they all load the SAME key to compute dot products with their respective Q vectors. - For key k, warp loads K[k][:1024] (broadcast to all threads). - Thread t has Q[t][:1024] in registers or local mem. - Each thread computes Q[t]·K[k] (1024 FMAs). - After 32 queries x 512 keys, each thread has its attention scores. - Then for V accumulation: load V[k][:] for key k, each thread multiplies by its softmax weight and accumulates. This way: - K is loaded once per key per warp (32 queries), and broadcast. - V is loaded once per key per warp. - Q is loaded once per query per warp. - Memory traffic per 32 queries: Q(32*1024) + K(512*1024) + V(512*1024) + O(32*1024) floats = ~1056K floats = 4.2MB - Total warps = 524288/32 = 16384 - Total traffic = 16384 * 4.2MB = 68.8GB That's much better than my shared-memory kernel! 68.8GB is very reasonable. At 3.35TB/s = 3350GB/s, that's about 20.5ms just for memory bandwidth. Add compute, maybe 25-30ms. And every K/V access is perfectly broadcast across a warp, maximizing cache utilization. Each key is loaded once from HBM and then broadcast to 32 threads via L1/L2. But 32 threads computing 32 dot products simultaneously, each against the SAME K vector: - Each thread processes one query. - K is broadcast to all threads. - But Q must be in registers for each thread. - D=1024, so each thread needs 1024 registers for Q... that's 1024*32 = 32k floats, which is 128KB for the warp. Each SM has 64K registers total, so that's already 2x too much per warp! We can't keep all of Q in registers across all 1024 dims if we need registers for other things too. We can process Q in chunks. Load Q[d:d+step] and K[d:d+step], compute partial dot product. Then reload for next chunk. - Each thread keeps an accumulator for the score. - Register usage per thread: Q_partial (step floats) + score (1 float) + max,sum (2 floats) + V_partial (step floats) + O_partial (step floats)... - If step=64: Q_partial = 64, V_acc = 64, plus a few scalars ≈ 130 floats ≈ 130 registers. That's reasonable. So the algorithm per thread block: - BlockDim = warp_size * num_warps - In my earlier calculation, one warp handles 32 queries. We can use multiple warps per block. - Say 8 warps per block = 256 threads, handling 32*8 = 256 queries. - Batch size = 32, Heads = 32, N = 512. Total queries = 524288. - 524288/256 = 2048 blocks. Grid = (2048,) or split as (32, 32, 16) for clarity. For 256 threads (8 warps) in a block: - Each warp handles 32 queries. - Warp 0: queries [0..31] - Warp 1: queries [32..63] - etc. Synchronization: - No need for __syncthreads() within a block if we just want atomic writes or separate outputs. Each warp computes independently. - But we need to make sure warps can load K and V cooperatively if we want to use shared memory for K/V to reduce global memory traffic. Wait, with 8 warps loading the same K/V but for different queries, if each warp loads K/V independently, the cache will serve it. On H100, L1 cache is quite good for broadcast. But we could also put K/V in shared memory to guarantee one load per key per block. Shared memory approach for K/V: - Key tile size = say 1 key (since we process one key at a time for all warps in the block). - Load K[k][0:1024] and V[k][0:1024] to shared memory. - 1024 floats * 2 = 8KB per key. - With shared memory 227KB, we can load multiple keys. - But 8 warps * 32 queries each = 256 queries. All processing same key tile. - If we load 8 keys at a time: 8 * 8KB = 64KB of shared memory. - 512 keys / 8 = 64 key tiles. Let's design it: - Process BK=8 keys at a time. - Load K[0:8][0:1023] to smem_k[8][1024] - Load V[0:8][0:1023] to smem_v[8][1024] - Each warp handles 32 queries. - For key j in 0..7: - Load Q for each query into registers (or process in chunks). Actually, let's process D in chunks of DT=64 to keep register pressure low. For each query, we do: 1. Score accumulation: For each key tile (8 keys): For each chunk d in 0..1023 step 64: Load Q[q][d:d+64] into registers Load K[k][d:d+64] from shared memory Compute partial dot product Sum partials via shuffle to get score 2. softmax 3. V accumulation: For each key tile (8 keys): For each chunk d in 0..1023 step 64: Load V[k][d:d+64] from smem Multiply by exp(score-max)/sum Accumulate into O[q][d:d+64] But to do this, we need two passes (one for scores, one for V accumulation), which doubles Q reads. Flash style online softmax avoids this by doing it in one pass. But with warps processing different queries and sharing K/V, online softmax is harder because each thread needs its own running max and sum. With two passes though: - Pass 1: compute all scores for all 512 keys, per query. Store 512 scores in shared memory or local memory. 512*4 = 2KB per query. For 256 queries per block = 512KB, too much. - We can do it in tiles: compute scores for BK keys, do softmax update, compute V contribution for those keys, discard scores. This is the online softmax approach. Online softmax per query in one pass: - For each key tile: - Compute scores and find tile max - Update running global max - Update running L and O accumulators This works perfectly! Each thread maintains its own max, sum, and O accumulator in registers. Let's formalize the per-thread algorithm: Each thread handles exactly 1 query (1 out of 256 in a block). Wait, 256 threads = 8 warps * 32 threads. We have 256 queries in a block, so 1 thread per query. That's perfect! Algorithm per thread (handling query q): max_val = -inf sum_exp = 0 o_acc[DT] = {0} but we'll process in chunks For key_tile = 0 to N-1 step BK: // compute scores for this key tile score = 0 For d_chunk = 0 to D-1 step DT: load Q[q][d_chunk:d_chunk+DT] from global to registers For k_local in 0..BK-1: load K[key_tile+k_local][d_chunk:d_chunk+DT] from shared Compute partial score... Wait, each thread needs BK scores (one per key in tile). If we have 1 thread per query, then for BK=8 keys and D=1024: - Each thread needs to compute 8 dot products. - Each dot product = 1024 FMAs. - 8 * 1024 = 8192 FMAs per thread per tile. - 512/8 = 64 tiles. - Total FMAs for scores = 8192 * 64 = 524k FMAs per thread. - Plus V accumulation: BK * D = 8 * 1024 = 8192 FMAs per tile * 64 = 524k. - Total: ~1M FMAs per thread. At 67 TFLOPS for 524288 threads? Actually, the GPU parallelizes across threads. Total work = 524288 threads * 1M FMAs = 5.24e11 FMAs = 524B operations. At 67 TFLOPS = 67e12 ops/sec, time = 7.8ms just for compute. But we're memory bound so it'll be slower. Memory traffic: - Q: Each query loaded once = 524288 * 1024 * 4 = 2GB. Actually Q is 32*32*512*1024*4 = 2GB. - K: Each key loaded once per block? Wait: - 256 queries per block, all processing all 512 keys. - With BK=8 keys per tile, 64 tiles. - K is loaded 64 * 8 = 512 keys per block, which is the full K. - 2048 blocks total, but each block loads its head's K. - Total K traffic = 32*32*512*1024*4 = 2GB. - V: same 2GB. - O: 2GB. Total: 8GB. That's the absolute minimum! At 3.35TB/s, 8GB = 2.4ms. Plus compute overhead. This could be quite fast! But wait, with 1 thread per query, and 256 threads per block, how do we load K and V into shared memory cooperatively? Block organization: - blockDim.x = 256 threads - Each block processes 256 queries from one head - There are 512 queries per head, so 2 blocks per head. - Total blocks for one head = 2 - For B=32, H=32, total blocks = 32*32*2 = 2048 blocks. Wait, each head has 512 queries. With 256 queries per block, we need exactly 2 blocks per head. But each block loads all 512 keys of its head's K and V. So K is loaded twice per head (once per block). That doubles K traffic. K traffic = 2GB * 2 = 4GB. Similarly V = 4GB. Q = 2GB. O = 2GB. Total = 12GB. Still quite good! 12GB / 3.35TB/s = 3.6ms. For 2048 blocks and 132 SMs, that's 15.5 blocks per SM. With 256 threads per block, that's about 4000 threads per SM, which is fine. But wait, can we fit 256 threads per block? Yes, up to 1024. Can we schedule 15 blocks per SM concurrently? Max is 32 blocks per SM. So we're fine. However, each block processes 256 queries but only 1 at a time per thread. Let me think about how to structure the kernel: ```cpp __global__ void attention_kernel( const float* Q, const float* K, const float* V, float* O, int B, int H, int N, int D ) { int tid = threadIdx.x; int q_per_block = 256; // blockDim.x int q_block_id = blockIdx.z; // 0 or 1 for N=512 int q = q_block_id * q_per_block + tid; int b = blockIdx.x; int h = blockIdx.y; if (q >= N) return; // read Q[q][0:D] into registers in chunks // iterate over all keys // use shared memory for K and V tiles int head_offset = ((b * H + h) * N + q) * D; int head_offset_kv = ((b * H + h) * N) * D; float scale = 1.0f / sqrtf((float)D); float m = -1e20f; float l = 0.0f; // O accumulator: process in chunks // We can't keep 1024 floats in registers per thread. // We'll process D in chunks of, say, 64. for (int d_base = 0; d_base < D; d_base += DT) { float o_acc = 0.0f; // Wait, we need to interleave keys and D chunks simultaneously. // Standard approach: inner loop over keys, but that means we need Q for each key. // Better: outer loop over keys, inner loop over D chunks. // Load Q[q][d_base:d_base+DT] once, then iterate over key tiles. // But for online softmax, we need scores for ALL keys before updating O? // No, we can do online softmax at the key level. // Actually, if we process D in chunks, we need partial scores per key tile. // This gets complicated. // Alternative: keep all O partials in a small shared memory per thread, // but 256 threads * 1024 floats = 1MB, which is too much. } } ``` The problem is that for online softmax, the correction factor for O accumulators depends on the global max, which isn't known until all keys of a tile are processed. But we can process all keys for one D chunk at a time? No, that doesn't work because the dot product requires ALL D for each key. We need to process per key, over all D. But D=1024 is too big for per-thread registers. We could break D into chunks but then we need to solve the online softmax problem. Alternative 1: Two-pass algorithm for each query. Pass 1: compute all 512 scores, store in shared memory. Pass 2: compute softmax, then accumulate V. But 512 scores per query, and 256 queries per block = 131k floats = 512KB. That fits in shared memory! Wait, shared memory is 227KB per SM, but per block it's also limited to 227KB. 512KB > 227KB. So we can't store all scores for all 256 queries in shared memory. But we can do it in chunks: process 64 queries at a time within a block. - 64 queries * 512 scores = 32768 floats = 128KB. That fits! - Process 64 queries, store their 512 scores. - Then compute softmax and accumulate V for those 64 queries. - Repeat 4 times (for 256 total). But then we have 64 queries computing scores, while the block loads K and V. The other 192 threads are idle or could help with loads. Actually, a better approach: have ALL 256 threads work together to load K/V into shared memory, but only 64 threads compute scores at any given time. That's wasteful. Wait, each thread handles one query. We could have each thread compute its score and immediately keep track of max, then in a second pass compute exp and accumulate. But that requires storing 512 scores per thread in local memory. Local memory is backed by global memory, which is slow. 512 floats = 2KB per thread. With 256 threads, total local memory = 512KB per block. This spills to global memory and causes 512KB of global memory traffic per block, which is acceptable (2048 blocks * 512KB = 1GB extra). But local memory accesses are slow. Hmm, local memory: each thread accesses its own private area, so they should be coalesced if laid out nicely. But the latency is higher than registers. Actually, for CUDA, local memory is just global memory with known addressing, so each access is a regular global load/store. 512 floats accessed sequentially should be coalesced per warp. But if each thread needs to write 512 floats and then read them back, that's 2 * 512 * 4 = 4KB per thread, and with 256 threads running in parallel... Actually no, local memory is allocated per thread and the accesses go through L1/L2. Alternative: Use a smaller number of threads per block but more blocks. For instance: - 1 block = 1 warp = 32 threads, processing 32 queries. - 512 queries / 32 = 16 blocks per head. - Total blocks = 32 * 32 * 16 = 16384 blocks. - Score storage: 32 threads * 512 scores = 16384 floats = 64KB per block. This easily fits in shared memory! - Wait, 64KB < 227KB, yes! 16384 blocks is actually fine. It gives very good parallelism. Each block processes 32 queries and loads the full K/V for its head. - K is loaded 16 times per head. - V is loaded 16 times per head. - Total K traffic = 32GB, V = 32GB. That's too much! Wait: 16 blocks per head * 2MB K per head = 32MB per head. 1024 heads = 32GB. Yes, that repeats K and V 16x. But with L2 cache of 50MB, if each head's K (2MB) can stay in L2 while its 16 blocks run, then only the first block reads K from HBM and the rest use L2. But the blocks might be scattered across SMs. Actually, with 16384 blocks, and 132 SMs, there are ~124 blocks per SM. These blocks are from different heads. L2 cache can hold K for 25 heads. The remaining 1000 heads' K won't be cached. So most K/V blocks will be HBM reads. This is bad. We need to reduce per-head K/V reads. To reduce K/V reads per head, we need FEWER blocks per head that all load the same K. That means MORE queries per block, which brings us back to the shared memory score storage issue. What if we have 1 block per head that processes ALL 512 queries? - Block = 512 threads. Each thread handles 1 query. - 32*32 = 1024 blocks total. - Each block loads K and V once. Great! - 512 scores per thread. We need to store them somewhere. - If we keep them in local memory, each thread writes and reads 512 floats. - Total local memory traffic = 1024 blocks * 512 threads * 512 floats * 4 bytes * 2 (read+write) = 2GB. Not too bad! But can we have 512 threads per block? Yes, up to 1024. 512 threads * 512 scores * 4 bytes = 1MB for score storage if using shared memory. But 1MB exceeds 227KB per block. If we use local memory (spill to global memory), the compiler handles it. However, 512 threads per block means we can only have up to 32 blocks per SM (2048 threads / 512 = 4 concurrent blocks). But with 1024 blocks and 132 SMs, we need 7.8 blocks per SM. So we'd process in 2 waves. That's OK. Wait, with 512 threads per block, total threads = 512 * 1024 = 524288, which is exactly the number of queries. Each thread 1 query, perfect. Then K and V are loaded once per block. Total traffic: - Q: 2GB - K: 2GB - V: 2GB - O: 2GB - Score spill/read-back: 2GB (write + read = 1GB each way? Actually local memory read+write) Total: 10GB At 3.35TB/s = 3ms for memory. Plus compute. This could be very fast! But wait - will local memory be coalesced? With 512 threads, each accessing 512 floats, if they're stored contiguously per thread, the accesses are NOT coalesced within a warp because thread 0 accesses elements 0,1,2,... and thread 1 accesses 512,513,... These are 512*4=2048 bytes apart, which might be coalesced across warps but not within a warp. Wait, actually a warp of 32 threads accessing 32 different 2048-byte-offset elements: each thread accesses a different cache line (assuming 128B lines). So 32 cache lines per warp transaction. That's terrible. To make score storage coalesced, we need thread t to access element t, t+blockDim, t+2*blockDim, etc. So we need to store scores as transposed[score_idx][thread] instead of [thread][score_idx]. But these scores are in local memory, and the compiler decides the layout. The layout for local arrays is typically row-major with the thread index as the slowest varying dimension, which means consecutive elements in the array are consecutive in memory. Wait, for `float scores[512];` in a thread, the compiler allocates `512 * 4 = 2048` bytes per thread. Memory layout is thread 0: bytes 0-2047, thread 1: bytes 2048-4095, etc. (Actually it's often interleaved in groups of 128B for coalescing, but for spilled arrays I'm not sure.) If the layout is thread-interleaved, accessing `scores[i]` by all threads in a warp would be perfectly coalesced! Wait, how does the compiler store local arrays? I recall NVIDIA docs say that access to local memory is only ever fully coalesced when ALL threads in a warp access the same address. Otherwise, it's strided. Actually, for 1D arrays, the compiler tries to ensure that `array[i]` by thread `t` is coalesced with `array[i]` by thread `t+1`. This means the elements `scores[i]` for all threads in a warp are stored together. Actually, the standard compiler behavior for `float local[512]` is to allocate them such that `local[i]` for all threads in a warp are contiguous. More precisely, the layout is: for each `i`, `thread0_local[i], thread1_local[i], ..., thread31_local[i], padding to 128B boundary`. This makes per-warp access to `local[i]` coalesced! But a warp of threads all accessing `local[i]` at the same time would be coalesced. In our algorithm: Pass 1: - Thread t computes score for key k over all D, writes to local[k]. - At key k, all threads in a warp write to their local[k]. This IS coalesced! Pass 2: - Thread t reads local[k] for key k. - All threads in a warp read from their own local[k] at the same time, which is also coalesced. WAIT, but pass 2 also needs ALL scores for a thread to compute softmax. The softmax reads scores sequentially: ```cpp for (int k=0; k<N; k++) s[k] = scores[k]; ``` With `float scores[512]`, and all threads reading `scores[0]`, then `scores[1]`, etc., at each k all threads read their local[k]. This is coalesced. Then we need to find max over 512 scores per thread. This is done in registers, no memory issue. Then compute exp(scores[k] - max) and sum. Again, read scores[k] sequentially, coalesced. Then accumulate V: read scores[k] again for weights. Wait, do we need to store scores at all? If we do two passes: - Pass 1: compute all scores, track max. - We only need max from pass 1. We don't need to store scores! - Pass 2: recompute scores and accumulate V with softmax weights. This avoids storing scores entirely! But it recomputes the dot products: - Pass 1: Q·K for all 512 keys = 512*1024 = 524k FMAs. - Pass 2: Q·K again for all 512 keys = 524k FMAs. - Plus V accumulation: 512*1024 = 524k FMAs. Total: ~1.57M FMAs per query. Compute total: 524288 queries * 1.57M = 823B FMAs. At 67 TFLOPS = 12.3ms. Memory: - Q: 2GB (loaded twice, once per pass) = 4GB - K: 2GB (loaded twice) = 4GB - V: 2GB (loaded once) = 2GB - O: 2GB Total: 12GB. At 3.35TB/s = 3.6ms. The total could be around 12-16ms. But if we're compute bound at 12ms, that's the time. But PyTorch reference might do something slower, especially if it's not using flash attention. However, this approach reads Q and K twice each. Can we do better? Alternative: keep scores in registers or shared memory to avoid recomputation. With 512 threads per block processing 512 queries in one head: - We could have shared memory store a tile of K and a tile of scores. - Actually, instead of each thread storing its own 512 scores, we could process keys in tiles. - For each key tile of 16 keys: - All 512 threads compute their scores for these 16 keys. - For 512 threads, 16 keys each, that's 8192 scores. - Store in shared memory? 8192 * 4 = 32KB. That's very reasonable! - Run softmax on the 16 scores + running stats - Accumulate V for these 16 keys - Discard scores Wait, this is EXACTLY the online softmax tile approach, but with the twist that each thread handles one query and all keys in the tile, and scores for a tile are stored in shared memory or registers. Actually, with 16 keys per tile: - Each thread computes 16 scores (one per key in tile). - It can keep them in 16 registers! - Then do warp-level or thread-level softmax over 16 elements. - With online softmax, we update running stats and O accumulator. - No need to store scores in shared memory at all! This is the key insight. Let me formalize: Per thread (handles 1 query): ```cpp float m = -1e20f; // running max float l = 0.0f; // running sum exp float o[DT]; // partial output accumulator, for D processed in chunks of DT for (int kt = 0; kt < N; kt += BK) { // BK = 16 // compute BK scores for this query float s[BK]; // 16 registers #pragma unroll for (int k = 0; k < BK; ++k) { s[k] = 0; } for (int d = 0; d < D; ++d) { float q_val = Q[q*D + d]; // each thread loads its own Q once per D #pragma unroll for (int k = 0; k < BK; ++k) { float k_val = K[(kt+k)*D + d]; // broadcast or coalesced load from shared s[k] += q_val * k_val; } } // apply scale #pragma unroll for (int k = 0; k < BK; ++k) s[k] *= scale; // online softmax for this tile float tile_max = s[0]; #pragma unroll for (int k = 1; k < BK; ++k) tile_max = fmaxf(tile_max, s[k]); float new_m = fmaxf(m, tile_max); float corr_m = expf(m - new_m); float corr_tile = expf(tile_max - new_m); float tile_sum = 0; #pragma unroll for (int k = 0; k < BK; ++k) tile_sum += expf(s[k] - tile_max); // update O accumulator for each d_chunk: load O accumulators for k in 0..BK-1: p = expf(s[k] - new_m); V_val = V[(kt+k)*D + d]; o[d_chunk] = o[d_chunk] * corr_m + p * V_val; // wait, this isn't right. ``` Actually, the standard online softmax update for O is: ```cpp O_new = O_old * exp(m_old - m_new) + sum(exp(s_i - m_new) * V_i) ``` So for each key: ```cpp float p = expf(s[k] - new_m); #pragma unroll for (int d = 0; d < DT; ++d) { O[d] = O[d] * corr_m + p * V[(kt+k)*D + d]; } l = l * corr_m + tile_sum * corr_tile; ``` But wait, `corr_m` is applied to all elements of O. This is correct. And `p = expf(s[k] - new_m)` is the weight for each key. The issue is: we need to iterate over keys first, then D, or D first then keys? Standard approach: iterate over D chunks, and for each D chunk, iterate over keys to accumulate. But with 16 keys and DT=64, the inner loop would be: ```cpp for each D chunk: for each key k: p = expf(s[k] - new_m); for d in chunk: o[d] = o[d] * corr_m + p * V_val; ``` Wait, `corr_m` should be applied ONCE to all of O, not per key. So: ```cpp for each D chunk: #pragma unroll for (int d = 0; d < DT; ++d) o[d] *= corr_m; for each key k: p = expf(s[k] - new_m); #pragma unroll for (int d = 0; d < DT; ++d) o[d] += p * V[(kt+k)*D + base + d]; ``` This reduces the number of multiplications by factor BK. But we need `corr_m` per tile. Wait, is this correct? Online softmax update: At tile t, we have old accumulator `O_old` with old max `m_old`. We compute tile max `m_t` and new max `m_new = max(m_old, m_t)`. The full output should be: `O_new = O_old * exp(m_old - m_new) / Z_new * Z_old + sum_i exp(s_i - m_new) * V_i / Z_new` But in the standard flash attention formulation, the running O accumulator stores `O = sum(exp(s_i - m_running) * V_i)`, and we don't divide by L until the end. So when we update: ```cpp O_new = O_old * exp(m_old - m_new) + sum_i exp(s_i - m_new) * V_i L_new = L_old * exp(m_old - m_new) + sum_i exp(s_i - m_new) ``` Then at end, `O_final = O_new / L_new`. So yes, `corr_m = exp(m_old - m_new)` should be applied to the old accumulator once. Then add contributions from the new tile. But there's a subtlety: if the old accumulator O_old corresponds to previous keys, and then we add new key contributions, each with their own `exp(s_i - m_new) * V_i`, that's correct. So the algorithm per D chunk: ```cpp for (int d = 0; d < DT; ++d) o_local[d] *= corr_m; for (int k = 0; k < BK; ++k) { float p = expf(s[k] - new_m); for (int d = 0; d < DT; ++d) { o_local[d] += p * smem_v[k][d]; } } ``` With BK=16, we do 16 loads of V and 16 * DT FMAs. If DT=32, that's 512 FMAs. Plus the initial multiply by corr_m (32 FMAs). And the 16 exp() computations. For K loading, we can put K in shared memory to avoid each thread loading K independently. - For BK=16 and D=1024, K tile = 16*1024*4 = 64KB. This fits in shared memory! - Similarly V tile = 64KB. - Total shared memory = 128KB + a bit for smem_q etc. Under 227KB. Wait, each thread needs to load Q values. Each thread processes one query. Q for that query can be loaded directly from global memory. Actually, Q is only read once per query, and each thread handles one query, so Q can be in global memory. But to maximize throughput, since each query's D=1024 is sequential, it's fine. Actually, for the dot product, each thread loads Q[d] and multiplies by K[k][d]. K[k][d] should be in shared memory for good reuse. There are 512 threads per block, all processing queries from the same head but different queries. They all access the SAME K[k][d] for their respective dot products. So K should definitely be in shared memory to avoid 512 redundant global loads. Let's structure a block to process one head: - blockDim.x = 512 threads (or 256 if 512 is too many and limits occupancy) - Each thread handles 1 query. - 512 queries per head means exactly 1 block per head. - Grid = (B, H) = (32, 32) Wait, with 512 threads per block: - On H100, max threads per SM = 2048 (if using 32 registers, but we need maybe 64-128). - With 512 threads per block, we get at most 4 blocks per SM. - Total blocks = 1024. 1024 / 132 = 7.76 blocks per SM. - So occupancy = 4 blocks concurrently? Actually 7.76 total, so about 2 waves. Fine. Register usage: - Each thread needs: q_val (1), score array BK=16 (16), O accumulator DT (say 32), plus some temporaries. - 16 + 32 + ~10 = ~58 registers. Well within 255 limit. - With 512 threads * 58 regs = ~30K registers. SM has 64K, so we can fit 2 blocks = 1024 threads comfortably. Actually 64K/58 = 1100 threads per SM. So with 512 threads, we can fit 1 block = 512 threads. With 2 blocks = 1024 threads, regs = 1024*58=59K, which is OK. Wait, I also need to consider shared memory: - K tile: 64KB - V tile: 64KB - Total: 128KB - Max shared memory per SM = 227KB - With 1 block per SM: 128KB < 227KB, OK. - With 2 blocks per SM: 256KB > 227KB. Oops. So we can only have 1 block per SM concurrently because of shared memory. Total concurrent blocks = 132. But we have 1024 blocks. So it takes 1024/132 = 7.76 waves. That's about 8 waves. If each wave takes ~5ms, total = 40ms. Not great. But wait, can't we use less shared memory? If we have 2 blocks per SM, we need shared memory < 113KB per block. - If BK=8 instead of 16: K=32KB, V=32KB, total=64KB. Then 2 blocks = 128KB < 227KB. OK. - Or we could use a different approach. With BK=8: - Number of key tiles = 512/8 = 64. - Shared memory per block = 64KB. - 2 blocks per SM = 128KB. OK. - Total concurrent blocks = 264. - Waves = 1024/264 = 3.88. Better. But BK smaller means more tiles, more softmax iterations, more overhead. What if we use 256 threads per block, processing 256 queries per head, with 2 blocks per head? - Block = 256 threads. - BK=8. K=32KB, V=32KB. Total smem=64KB. - 2 blocks per head = 2 blocks loading K and V independently. - 2 blocks per SM = 128KB smem used. OK. - 2 blocks * 256 threads = 512 threads per SM concurrent. - Total blocks = 32*32*2 = 2048. - Waves = 2048 / (132*2) = 7.76. Same issue as 512 threads. Hmm, the issue is the total number of head-blocks and how many fit concurrently. What if we process multiple heads per block? No, because queries are independent. What if we reduce the number of heads per grid launch by merging the head dimension into the block? Actually, let me reconsider: why do we need K and V in shared memory at all? If each of the 512 threads loads K[k][d] independently from global memory, the cache will handle it! Specifically: - All 512 threads want K[k][d] at the same time (for different queries). - H100 L1 cache can broadcast a single value across a warp, but 512 threads = 16 warps. - L2 cache is shared across all SMs. If a block has 512 threads on one SM, they share the same L1. A single L1 cache line can serve at most... hmm, I don't think L1 broadcasts to multiple warps within the same SM simultaneously. Actually, L1 cache on Ampere/Hopper can serve multiple requests per cycle but not necessarily broadcast to 512 threads. If all 512 threads request the same address, it might take multiple cycles to broadcast. Shared memory, on the other hand, CAN broadcast to all threads in a block in a single cycle! Wait, can shared memory broadcast to all threads? Shared memory has multiple banks. For a broadcast, all threads access the same bank address, which causes bank conflicts but some hardware might handle it specially. Actually, reading the same shared memory address from all threads in a warp might be broadcast (no conflicts), but from 16 warps, you need 16 access ports or time-multiplexing. Wait, shared memory banks: 32 banks, each 32 bits wide. To broadcast, multiple threads access the same 32-bit address. A broadcast occurs when all threads in a warp read the same address - this is handled with no bank conflicts. But if 16 warps all try to read the same shared memory address simultaneously, they can't all be served in one cycle because shared memory only has 32 banks/ports. It would take multiple cycles, but still much faster than multiple L1/L2 accesses. So yes, shared memory is good for data that's read by many threads. But shared memory limits occupancy. Can we use L1 caching effectively without explicit shared memory? For 512 threads per block processing one head: - When processing key k, all threads load K[k][d] and V[k][d]. - If K and V are accessed with a stride pattern, are they cached? - The first warp loads K[k][d] for d=0. The next warp loads K[k][d] for d=0 as well (since all queries use the same key). All warps load the SAME address. This is not good for cache but might be handled. Actually, if we process D as the inner loop: ```cpp for (int d = 0; d < D; ++d) { float k_val = K[(kt + k) * D + d]; for (int k_local = 0; k_local < BK; ++k_local) { ... } } ``` Wait, no. Let me think about the memory pattern. If we don't use shared memory for K/V, and each thread does: ```cpp for each key k: float score = 0; for each d: score += Q[q*D + d] * K[k*D + d]; ``` Then all 512 threads (for 512 queries) are executing: ``` tid 0: load Q[0*D+d], K[k*D+d] tid 1: load Q[1*D+d], K[k*D+d] ... tid 511: load Q[511*D+d], K[k*D+d] ``` For a given d and k, all 512 threads load the SAME K[k*D+d] address, but different Q addresses. The Q loads are perfectly coalesced (thread t loads Q[t*D+d], consecutive threads load consecutive addresses if D is large... wait, no. Q is laid out as [N][D], so Q[0][d], Q[1][d] are D floats apart. That's a stride of D*4 = 4096 bytes. That's terrible for coalescing! Wait, the layout is [B, H, N, D]. So for a fixed b and h, the data for query q starts at offset q*D. Query q+1 is at (q+1)*D. So Q[q][d] and Q[q+1][d] are separated by D elements = D*4 bytes = 4096 bytes. A warp of 32 threads accessing Q[0][d], Q[1][d], ..., Q[31][d] has a stride of 4096 bytes between consecutive threads. That's NOT coalesced. Each thread in the warp accesses a different 128B cache line (4096/128 = 32 lines apart). Each warp transaction needs to fetch 32 cache lines = 4096 bytes. That's the worst possible pattern. This is a CRITICAL realization. For attention where each thread handles one query, we need the Q, K, V access pattern to be coalesced. To access Q coalesced, threads in a warp should access the same query's consecutive elements, or different elements of the SAME query. If a warp of 32 threads processes ONE query together, then they access Q[q][d+tid], which is perfectly coalesced. That's the flash attention style: one warp or block of warps processes a tile of queries, and threads within the tile access different D elements of the SAME query or queries. In my original kernel, I had a block processing BQ=8 queries, and threads collectively loaded Q[0..7][d] with coalesced accesses. That part was actually better. So the single-thread-per-query approach fails on Q coalescing. We need thread cooperation within queries. Back to the original kernel style, but optimized: - A block of threads cooperates to process BQ queries. - Threads in a warp access the same query or different D elements. Actually, the original kernel had: ```cpp for (int i = tid; i < BQ * DT; i += blockDim.x) { int q = i / DT; int d = i % DT; smem_q[q][d] = Q[head_offset + q_idx * D + dt + d]; } ``` This is coalesced because consecutive threads load consecutive memory addresses (outer loop over BQ*DT in thread-index-linear order, and memory is laid out as [N][D], so q_idx*D + dt + d are consecutive when d increments). So the original kernel's memory access for Q loading was actually perfect. The main performance issues with the original kernel were: 1. Wasted threads (only 16 of 32 lanes used for BK=16) 2. Small BQ (8 queries), meaning K and V are loaded many times. Can we increase BQ to amortize K/V loads? With shared memory = 227KB and D=1024: - Q tile: BQ * DT (but DT doesn't have to be all of D) - Actually, we load Q in chunks of DT into shared memory. So smem_q size = BQ * DT * 4 bytes. - K tile: BK * DT * 4 bytes. - O accumulator: BQ * D * 4 bytes. With D=1024, O accumulator is BQ * 4096 bytes. If BQ=16, O = 64KB. If BQ=32, O = 128KB. Q tile (for DT=64): 16*64*4 = 4KB. K tile: BK*64*4. If BK=32, K=8KB. P tile: BQ*BK*4 = 16*32*4 = 2KB. M/L: BQ*8 = 128 bytes. Total: 64 + 4 + 8 + 2 + ... = ~80KB for BQ=16, BK=32. That's fine! So we can use BQ=16 and BK=32. - Blocks per head = 512/16 = 32. - Total blocks = 32*32*32 = 32768 blocks. - K read per head = 32 blocks * (512/32) tiles * 32*1024*4 bytes? Wait. - Each block loads K 512/BK = 16 times. - Total K reads per head = 32 blocks * 16 tiles * 32*1024*4 bytes = 32 * 16 * 128KB = 64MB per head. - Total for 1024 heads = 64GB. - Plus V = 64GB. Plus Q = 2GB. O = 2GB. - Total = 132GB. That's... a lot, but the original BQ=8 would be 128GB of K/V + 2GB + 2GB = 264GB. Ah, I see! Doubling BQ halves the number of blocks per head, but each block still loads the full K and V. Wait, no: - BQ=8: 64 blocks per head, each loads all K and V. Total K reads = 64 * 2MB = 128MB per head. - BQ=16: 32 blocks per head, each loads all K and V. Total K reads = 32 * 2MB = 64MB per head. - BQ=32: 16 blocks per head, K reads = 32MB per head. - BQ=64: 8 blocks per head, K reads = 16MB per head. - BQ=128: 4 blocks per head, K reads = 8MB per head. - BQ=256: 2 blocks per head, K reads = 4MB per head. - BQ=512: 1 block per head, K reads = 2MB per head. So as BQ increases, K/V reads decrease. But so does parallelism (fewer blocks). With BQ=512, we have 1024 blocks. That's enough for 132 SMs (7.8 blocks per SM). Not bad. But with BQ=512, the block has 512 threads processing 512 queries. We need a good strategy for this. Let me design the BQ=512, BK=64 kernel. But BK=64 means K tile shared mem for one tile of 64*1024 = 256KB, which exceeds shared memory. We need to process K in smaller tiles but load them in shared memory chunks. What if we process D in chunks? For each D chunk of DT=64: - Load Q chunk: 512*64 = 32K floats = 128KB into shared memory. That's too much! 128KB is on the edge but actually fits. Wait, BQ=512, DT=64, Q in smem = 512*64*4 = 128KB. K in smem = BK*64*4. If BK=32, K=8KB. O accumulator = we don't need all of O in smem if we process D chunks and update output directly. Actually, we can do an interesting thing: Since N=512 and D=1024, and we want to process all queries in a head in one block. For O dimension, we can write directly to global memory without a separate accumulator tile. Or keep a partial O in registers per thread. But each thread handling one query would have the strided Q access problem. So we need warps to handle multiple queries with cooperative loading. Let me try a different approach altogether. What if we keep it simple AND fast: - Use `cublasGemmStridedBatched` for both matmuls via `torch`'s operations, but my own custom fused softmax in between. - Actually, since I can write C++ CUDA code with torch extensions, I can use cublas from within the CUDA extension. Let me check if I can link cublas with `load_inline`. The `extra_ldflags` parameter allows passing linker flags. And I can include `cublas_v2.h`. However, two separate cublas calls + a custom softmax kernel would involve multiple kernel launches and global memory round-trips: 1. S = Q @ K^T (writes 1GB) 2. softmax(S) (reads 1GB, writes 1GB) 3. O = S @ V (reads 1GB + 2GB, writes 2GB) Total extra traffic: 7GB. Not bad, plus the original 4GB for Q/K/V/O = 11GB total. But if PyTorch's SDPA uses flash attention, it avoids materializing S and would be ~8GB total. The 3GB extra from materializing S might make my approach slower. However, if PyTorch falls back to the math backend, it does explicitly materialize everything. In that case, using cublas might be just as fast or faster (because of highly optimized GEMM). Let me investigate: what does PyTorch `scaled_dot_product_attention` do for FP32? According to PyTorch docs and code: - It runs a backend selection algorithm: - If CUDA and cuDNN available and meets conditions, use cuDNN. - Else if flash attention available and dtype is fp16/bf16, use flash. - Else if mem_eff attention available and appropriate, use it. - Else use the math backend. For FP32, cuDNN SDPA might still be available. But cuDNN SDPA in FP32 might not be as optimized as flash attention. Still, it's likely very fast. Actually, I just realized: I don't actually know what backend the reference uses. But to beat the reference, I need to be faster than whatever it uses. If it uses flash attention, I probably can't beat it with a simple custom kernel. If it uses math backend, I might be able to beat it. But given that this is on H100 with the latest PyTorch, I bet it uses something quite optimized (cuDNN or flash attention). However, maybe there's a winning move: use Tensor Cores. Even for FP32, H100's Tensor Cores can process matrices very fast. But using them in custom CUDA is extremely hard. Wait, what if I just use a much simpler approach: transpose the dimensions so the GEMMs are more efficient, and use `wmma`? No, too complex. Let me reconsider my first kernel and think about how to optimize it for this specific hardware. In the first kernel: - BQ=8, BK=16 - Lane_id < BK means only 16 lanes used for score storage - This means 50% thread waste in the critical loops If I change BK=32, all 32 lanes are used. - q_local = warp_id (but only warp_id 0..7 with 8 warps) - For block=256 threads, 8 warps: yes, warp_id 0..7. With BQ=8 and BK=32: - Score matrix = 8 x 32 = 256 elements. - With 8 warps, each warp handles 1 query and all 32 keys. That's perfect! Each lane handles exactly 1 key. - Dot product parallelization: each lane computes its own dot product for its key. - But D=1024 is still done as scalar loop. Can we parallelize the dot product? For each key in the tile, each lane computes Q[q]·K[k]. With D=1024, a lane would do 1024 FMAs serially. For BK=32 keys, all lanes work in parallel. So each warp does 32*1024 = 32K FMAs. But actually, the loops in my first kernel load Q and K in chunks. That's fine. The issue is: can we load K more efficiently using vector loads? Vector loads: if we load 4 floats at a time (float4), we reduce memory instructions. For Q loading: BQ*D = 8*1024 = 8K floats. With float4, 2K loads. For K loading per tile: BK*D = 32*1024 = 32K floats = 8K float4 loads. For a block of 256 threads, 8K/256 = 32 loads per thread per K tile. Not bad. For O accumulator: BQ*D = 8*1024 = 8K floats. Writing these out is 2K float4 stores. Let me check: what if I increase BQ to 16 and BK to 32? Shared memory: - smem_q: 16*64*4 = 4096 bytes (DT=64) - smem_kv: 32*64*4 = 8192 bytes - smem_p: 16*32*4 = 2048 bytes - smem_m/l: 16*8 = 128 bytes - smem_o: 16*1024*4 = 65536 bytes Total = ~80KB. Under 227KB limit. Great! Can we have 256 threads (8 warps) handle BQ=16 queries? - Each warp would handle 2 queries. - warp 0: queries 0,1 - warp 1: queries 2,3 - etc. - But the score matrix is 16x32 = 512 elements. 512 elements / 256 threads = 2 elements per thread. - Each thread would compute 2 scores. Each score is a dot product of D=1024. For dot product chunk: load Q[0..15][0..63] and K[0..31][0..63]. - Q: 16*64 = 1024 floats. 256 threads load 4 floats each (1 float4). Coalesced! - K: 32*64 = 2048 floats. 256 threads load 8 floats each (2 float4). Coalesced! Then for the dot products: - Each thread handles... hmm, with 256 threads and 16*32=512 scores, each thread does 2 dot products. - Thread i handles score for query (i/32) and key (i%32)? But i/32 gives warp_id (0..7), which only gets 8 queries. We'd need to map carefully. Actually, 256 threads = 8 warps * 32 lanes. If BQ=16 and BK=32: - We could have each warp handle 2 queries * 32 keys = 64 scores. - 8 warps * 64 = 512 scores. Perfect. - Warp w handles queries 2w and 2w+1. - In warp w, lane k handles key k for both queries (or just one query?). Wait, there are 32 lanes and 32 keys. So each lane can handle exactly 1 key. But we need to handle 2 queries per warp. So each lane needs to handle 2 dot products: one for query_a, one for query_b. Dot product: ```cpp float acc_a = 0, acc_b = 0; for (int d = 0; d < DT; ++d) { float q_a = smem_q[2*w][d]; float q_b = smem_q[2*w+1][d]; float k = smem_kv[lane_id][d]; acc_a += q_a * k; acc_b += q_b * k; } ``` Each lane does 2 FMAs per iteration. That's actually good for instruction throughput! After all DT chunks, each lane has acc_a and acc_b for key k. Then lane k writes acc_a and acc_b to smem_p. But smem_p needs to store for 16 queries x 32 keys. Wait, instead of storing in shared memory, we could keep scores in registers and do the softmax directly. But then we need all scores for one query to compute its max. Within a warp, for query_a, lanes 0..31 have acc_a for keys 0..31. We can do warp shuffles to find the max and sum for query_a and query_b. Yes! This is very elegant: - Each warp handles 2 queries. - All 32 lanes participate. - After computing dot products, lane k has score for key k. - We do warp-level reduction for max and sum for both queries simultaneously. - Each lane computes its own exp and weight. - Then for V accumulation, each lane can help compute partial contributions. For V accumulation with BK=32 keys and D=1024: - We need O for 2 queries, each O is 1024 floats. - With 32 lanes, we can have lane d handle dimension d for both queries. - Wait, D=1024 and 32 lanes means each lane handles 32 dimensions. - For each key k, we need V[k][:] for both queries. - Lane d handles dim d, d+32, d+64, ..., d+992 (32 dimensions total). - For query_a and query_b, each lane computes: `O_a[dim_j] += p_a[k] * V[k][dim_j]` for j=0..31 This requires loading V[k][d] for the lane's 32 dims. But V is stored as [N][D], so V[k][d] and V[k][d+1] are contiguous. Good. We can either load V directly from global memory or use shared memory. If we use shared memory for V tile: - For a BK=32 key tile and DT=64 D chunk: smem_v[32][64] = 8KB. - But if each lane handles 32 scattered dims, we can't easily map to a fixed DT tile. Maybe we process V in chunks of 32 dims: - For d_base = 0, 32, 64, ..., 992: - Load V[k][d_base + d] for k=0..31, d=0..31 into registers or shared memory. - Each lane owns dim d_base + lane_id. - For both queries, accumulate p_a * V and p_b * V. - This assumes p is the same for both queries across keys. Wait, we don't need shared memory for V if we process D=1024 in chunks of 32 and let each lane load its own V values. Let's formalize this approach: Block configuration: - BQ=16, BK=32 - blockDim = 256 threads (8 warps) - Each warp handles 2 queries. Algorithm per warp (for query pair q0, q1): For each key tile kt = 0..N-1 step BK: // Load K tile for BK=32 keys, D=1024 dims // We process D in chunks of DT=64 or 128 float s0 = 0, s1 = 0; // scores for q0, q1 for (int dt = 0; dt < D; dt += DT) { // Load Q for q0, q1: 2*DT floats // This could be done cooperatively by all warps once per block. // Actually, all 8 warps need Q for all 16 queries. // Load K for this warp's 32 keys: 32*DT floats. // Warp loads its own K tile into registers or shared memory. // For each d in dt..dt+DT-1: // k_val = K[kt + lane_id][d] (if K in registers, each lane loads its key) // q0_val = Q[q0][d] // q1_val = Q[q1][d] // s0 += q0_val * k_val; // s1 += q1_val * k_val; } // s0 and s1 are now dot products for key = lane_id. // scale s0 *= scale; s1 *= scale; // warp-level softmax for q0: // max0 = warp_max(s0) // sum0 = warp_sum(exp(s0 - max0)) // warp-level softmax for q1: // max1 = warp_max(s1) // sum1 = warp_sum(exp(s1 - max1)) // update running m0, l0, m1, l1 // compute p0 = exp(s0 - new_m0), p1 = exp(s1 - new_m1) // For V accumulation: // Each lane also handles some D dimensions. // wt = lane_id % something? // Actually for V accumulation, all 32 lanes can help for BOTH queries. // We have 32 lanes and D=1024 dims. Each lane can handle 32 dims. // Lane d handles dim d, d+32, d+64, ..., d+992. for (int d_chunk = 0; d_chunk < D; d_chunk += 32) { float v_val = V[kt + lane_id][d_chunk + something]? No, we need all keys. // For key k, lane j handles dim d_chunk + j. // But each key has its own V vector. // Actually, for a fixed D chunk, we need V[0..31][dim] for all keys. // If we process one key at a time: for (int k = 0; k < BK; ++k) { float p0_k = expf(s0_for_k - new_m0); // But only the lane that computed key k has s0! // We need to broadcast p0_k to all lanes. float p0 = __shfl_sync(0xFFFFFFFF, p0_for_some_lane, k); // and similarly p1. float v_val = V[kt + k][d_chunk + lane_id]; // coalesced if V is in shared or global with proper layout O0_acc[d_chunk/32] += p0 * v_val; O1_acc[d_chunk/32] += p1 * v_val; } } Wait, there's an issue: `s0` is in each lane, but lane k has s0 for key k. To get s0 for key k, we broadcast from lane k: `float s0_k = __shfl_sync(0xFFFFFFFF, s0, k);`. Then we can compute p0_k. This works! And V[kt+k][d_chunk + lane_id] should be loaded coalesced if we load from global memory with consecutive lanes accessing consecutive dims. But we also need O0_acc to have 32 registers (for 32 dims). With O1_acc also 32 registers, that's 64 registers just for accumulators. Plus scores (2), plus various temporaries. That's maybe 80 registers per thread. With 256 threads per block = 20K registers. Under 64K per SM, so we can fit 3 blocks. Not bad. For the V load: `V[kt + k][d_chunk + lane_id]`, with lane_id 0..31, accesses contiguous memory. Perfectly coalesced! But it's a global memory load. The L1 cache might handle this, but it's still going to HBM for V. Total V traffic: each block reads all V once, but for D=1024 and 16 queries, that's the same as any fused approach. Let me check: with BQ=16, blocks per head = 32. Each block reads V once = 2MB per head. Total V traffic = 64MB per head * 1024 heads = 64GB. Hmm, that's still high. If we increase BQ to 64: - 8 blocks per head. - V traffic = 16MB per head. - Total V = 16GB. Can we do BQ=64 with a 256-thread block? - 256 threads / 64 queries = 4 threads per query. That's not much. - We want each warp to handle some queries. 256 threads = 8 warps. - 64 queries / 8 warps = 8 queries per warp. But a warp has 32 threads, so 4 threads per query. - 4 threads can compute 4 keys each for one query... or all 4 can cooperate on the same key. Actually, with 4 threads per query and BK=32, we can have each query handled by 4 threads: - Thread 0: keys 0..7 - Thread 1: keys 8..15 - Thread 2: keys 16..23 - Thread 3: keys 24..31 Then for dot product, all 4 threads can load Q cooperatively. But warp-level shuffle reductions become more complex since the 32 lanes are divided into 8 queries × 4 threads. Hmm, this is getting complicated. Let me think about whether we can do BQ=64 with 2 warps per query: - 8 warps / 64 queries = 0.125 warps per query. No, that's not even integer. OK, the issue is that D=1024 is fixed, and we need to amortize K/V loads with larger BQ, but that reduces the threads per query. What if we use BQ=32 with 8 warps? That's 0.25 warp per query... wait, 32 queries / 8 warps = 4 queries per warp. That's what we had for BQ=16: 2 queries per warp. With BQ=32 and 8 warps: 4 queries per warp. Each warp handles 4 queries, and 32 lanes / 4 = 8 keys per query per warp? No, that's not right. Actually, 4 queries in a warp, each query needs scores for BK=32 keys. That's 128 scores. But a warp only has 32 lanes. So each lane would need to handle 4 scores. Lane assignments for 4 queries per warp, BK=32 keys: - Query 0: lanes 0..7 handle keys 0..7, 8..15, 16..23, 24..31? No, if each lane handles 1 key, that's only 8 keys per query. - Query 0: lanes 0..7 handle keys 0..31, 4 keys each (0,4,8,...,28 by lane 0)? Maybe we should assign lanes like: - Lane j handles key j for all 4 queries. - Query i (in this warp) needs score for key j, computed by lane j. - So lane j has scores for query 0,1,2,3 all for key j. - `s[i] = Q[qi]·K[kj]` where i=0..3. This means each lane computes 4 dot products, one per query. That's 4 FMAs per d iteration: ```cpp acc0 += q0 * k; acc1 += q1 * k; acc2 += q2 * k; acc3 += q3 * k; ``` With 32 lanes processing 32 keys, and 4 queries per warp, we cover 4 * 32 = 128 scores. For BQ=32, we need 8 warps. 8 * 128 = 1024 scores total across all warps. But BQ=32 and BK=32 means only 32*32 = 1024 scores total. Perfect! So the score matrix exactly fits: 1024 scores = 8 warps * 128 scores/warp. Then for warp-level softmax: - For query 0, we need all 32 scores. These are in lanes 0..31, but each lane has score for query 0 with DIFFERENT keys. - We need `warp_max` across lanes for query 0: extract all query-0 scores via shuffle. - Similarly for queries 1,2,3. - But a warp only has 32 values total for query 0 (1 per lane). That's a 32-element max reduction. Easy with shuffles. - Same for sum. For V accumulation: - Query 0 needs O of size D=1024. - With 32 lanes and 4 queries in the warp, we need to divide D among the lanes. - Total dimensions to compute = 4 * 1024 = 4096. - 32 lanes can handle 4096 dims by each lane handling 128 dims. - Lane j handles dim j, j+32, j+64, ..., j+992 for one query? Actually, we need each lane to load V[k][dim] for all keys k to contribute to the output. The weight for key k and query i is only known to one lane (lane k). We need to broadcast weights. For query 0, lane k has weight p0k. For dim d_j handled by lane j, we need: `O0[d_j] += p0k * V[k][d_j]` for all k. We can have each lane broadcast its p weights one by one: ```cpp for (int k = 0; k < BK; ++k) { float p0 = __shfl_sync(0xFFFFFFFF, p0_lane, k); float p1 = __shfl_sync(0xFFFFFFFF, p1_lane, k); float p2 = __shfl_sync(0xFFFFFFFF, p2_lane, k); float p3 = __shfl_sync(0xFFFFFFFF, p3_lane, k); // lane j loads V[k][d_start + j] for its d_start } ``` But if each lane handles 128 dims, for a single key, each lane loads 128 elements of V. That's a lot of loads. Wait, can we process V accumulation in chunks too? Like DT=32 per iteration? If each lane handles dim `d = lane_id` in chunk `d_base`, then for all 32 lanes, they load V[k][d_base + 0..31] which is coalesced. Each lane gets one float. But with 4 queries, after loading the V chunk, each lane needs to update 4 output accumulators. And it needs to receive 4 p weights (p0k, p1k, p2k, p3k) for each key k. For BK=32 keys and DT=32 dims: - 32 iterations of k. - In each k: 4 broadcasts + load V[k][d_base+lane_id] (coalesced). - 4 FMAs per lane (one per query). - Then move to next k. This is quite efficient on memory (32 floats loaded per k, all coalesced) and compute. For one V tile (BK=32 keys), each lane processes 32 dims (DT=32). There are 1024/32 = 32 D tiles. So per key tile: 32 * 32 = 1024 iterations of k+d. Actually 32 keys * 32 D tiles = 1024 operations per lane. With 32 lanes * 4 queries = 128 accumulators total. Each lane keeps 4 * 32 (DT registers)? No, we process one D tile at a time and reuse accumulators across D tiles. But we need to keep the full O across all D tiles. Since D=1024 and DT=32, we process 32 D tiles and for each tile we have 4 query accumulators. Total registers = 4 * 32 = 128 floats. Wait, NO. For online softmax, the O accumulator needs to be corrected (`* corr_m`) for each new key tile. The correction is per query, not per dimension. So: ```cpp for each key tile: compute new_m, corr_m for each D tile: for each query i in warp: O_i[d] *= corr_m; for k in 0..BK-1: p_i_k = ... O_i[d] += p_i_k * V[k][d]; ``` For 4 queries and DT=32, O accumulators = 4 * 32 = 128 registers. That's a lot. Actually, we'd process one D tile at a time and move to the next, accumulating O for that tile. But we need O from previous key tiles to persist across D tile processing? No, we can process a D tile for ALL key tiles before moving to next D tile: ```cpp for d_base in 0..D-1 step DT: initialize O0[d_base..d_base+DT-1] = 0 m = -inf, l = 0 for each key tile: compute scores (using full D, not just d_base) update m, l for k in 0..BK-1: for d in 0..DT-1: O0[d] = O0[d] * corr_m + p0k * V[k][d_base+d]; ``` Wait, this works! Because the score computation always uses the FULL D dimension. We just process output dimensions in chunks. So the algorithm is: 1. For each key tile, compute scores (dot products using all D=1024). 2. Do online softmax update. 3. For each D chunk (32 dims), update O accumulators using weights from this key tile. 4. Move to next key tile. Register usage: - Scores: 4 floats (one per query in warp). - O accumulation: 4 queries * 32 dims = 128 floats. - Q loads: 4 * 32 (DT for Q chunk, but Q for score uses full D, so Q is loaded in chunks and accumulated). Actually, for full D dot products, we don't need all Q in registers simultaneously. We load Q in chunks: ```cpp float s0 = 0, s1 = 0, s2 = 0, s3 = 0; for (int dt = 0; dt < D; dt += DT) { load Q[q0][dt]..Q[q0][dt+DT-1] into q0_vals[DT]? ... load K[k][dt]..K[k][dt+DT-1] into k_val[DT]? } ``` Wait, if each lane handles 4 queries, and each query's score is sum over D of Q*K, then: - We load Q for the 4 queries into registers. - We load K for one key into registers. - Multiply and accumulate. If DT=32: - q0[0..31], q1[0..31], q2[0..31], q3[0..31] = 128 floats. - k[0..31] = 32 floats. - Total: 160 floats. That's a lot of registers per thread. If DT=64: 256 + 64 = 320 floats. Too much. So DT needs to be smaller. What if DT=16? - 4*16 = 64 + 16 = 80 floats for Q+K. - Plus scores (4), plus O accumulators (4 * 32 = 128 if we process D output chunks of 32). - Total: 212 floats. That's a lot but might still be OK (255 limit). Wait, O accumulators (128 floats) are only needed during V accumulation, not during score computation. So the compiler might reuse registers. Let me try to write this kernel more concretely: ```cpp #define BQ 32 #define BK 32 #define WQ 4 // queries per warp #define DT_Q 64 #define DT_O 32 __global__ void attention_kernel(...) { int tid = threadIdx.x; int warp_id = tid / 32; int lane_id = tid % 32; int b = blockIdx.x; int h = blockIdx.y; // this block processes BQ=32 queries int qg = 0; // Actually with 1 block per head and BQ=32, we need 16 blocks per head? No, grid=(B,H,16) for 512/32=16 query groups per head. int qg = blockIdx.z; // 0..15 for N=512/BQ=32 int qid = warp_id * WQ; // 0, 4, 8, ..., 28 for warp 0..7 // queries handled by this warp: qg*BQ + qid + i, for i=0..WQ-1 int q0 = ...; int q1 = q0 + 1; int q2 = q0 + 2; int q3 = q0 + 3; float m0=-inf, m1=-inf, m2=-inf, m3=-inf; float l0=0, l1=0, l2=0, l3=0; // O accumulators for DT_O dims float o0[DT_O], o1[DT_O], o2[DT_O], o3[DT_O]; // initialize to 0 for (int kt = 0; kt < N; kt += BK) { // compute scores for this key tile float s0 = 0, s1 = 0, s2 = 0, s3 = 0; for (int dt = 0; dt < D; dt += DT_Q) { // load Q for 4 queries and K for 1 key (lane_id's key) float q0v[DT_Q/WQ]; // wait, we need all DT_Q dims. } } } ``` This is getting complex, and I need to be careful about shared memory, coalescing, and register spilling. Let me reconsider and go for a simpler but well-optimized kernel. Given the complexity and time constraints, maybe I should try a kernel with: - BQ=16, BK=32, DT=64 - 8 warps in a block - Each warp handles 2 queries (16/8) - All 32 lanes are active - Process D=1024 in DT=64 chunks - Shared memory: Q[16][64]=4096B, K/V[32][64]=8192B, P[16][32]=2048B - O accumulator stays in shared memory: 16*1024*4=65536B - Total smem = ~80KB Wait, can each warp handle 2 queries with 32 lanes? For scores: - Lane j computes score for key j for BOTH queries. - So dot products: `acc0 += q0 * k_j`, `acc1 += q1 * k_j`. - Two FMAs per iteration. For V accumulation: - Lane j loads V[k][d_chunk + j] for some d_chunk? - But we need to update O for both queries. - Each lane can handle one dimension. For D=1024, lanes 0..31 handle dim 0..31 for the first chunk, then 32..63, etc. - For each key k, lane j gets p0k and p1k via shuffle. - Then lane j updates o0[j] += p0k * v[k][j], and o1[j] += p1k * v[k][j]. - Wait, o0 needs 1024 dims. We can keep just the current chunk in registers, but we need to persist across key tiles. With online softmax, O must persist across key tiles. So we need the full O[16][1024] somewhere. Options: 1. Shared memory (80KB for O alone, as noted above). 2. Global memory (slow). 3. Split D processing: process all key tiles for a D chunk, then write that chunk and move to next. But that means processing scores (which need all D) for every D chunk, which is terrible. Shared memory for O seems unavoidable if we want to avoid global round-trips. With 80KB total smem, we can only fit 1-2 blocks per SM (SM has 227KB). That's acceptable. 1024 blocks / (132 SMs * 1 block) = 7.76 waves. Each wave processes all SMs for one iteration. But wait, we can further reduce O size by processing fewer queries per block or using smaller accumulators. Actually, 80KB is not that large. Let's calculate total time with this kernel: - Memory traffic ≈ 132GB (from earlier calculation? Actually with BQ=16: K=64GB, V=64GB, Q=2GB, O=2GB = 132GB) - 132GB / 3.35TB/s ≈ 39ms - Plus compute. This seems too slow. But I think my traffic calculation was wrong because with 1024 blocks and only 132 SMs and 80KB smem, the actual number of concurrent blocks per SM is limited by shared memory to maybe 2 (160KB), but with warps, it's limited by registers too. Actually, at 132GB of traffic, the kernel is clearly bandwidth bound. Is the reference doing significantly less traffic? If the reference uses flash attention with the optimal algorithm, it also loads Q once, K once, V once per head (if it can fit everything in tiles effectively). But with D=1024, a flash attention tile for Q might be small. Wait, in the standard flash attention algorithm, each block loads Q tiles of size BQ x DT into SRAM, and iterates over KV tiles. The total K/V reads per head is (N/BQ) * D * N / BK? No, each query block loads all K and V. So total K reads = (N/BQ) * (N * D). Similarly for V. With BQ=8: K reads = 64 * 512 * 1024 = 33.5M floats = 134MB per head? Wait no: N/BQ = 512/8 = 64 query blocks per head. Each query block loads K in tiles of BK. Total K loaded per query block = N * D = 512 * 1024 floats. Total K reads per head = 64 * 512 * 1024 floats = 33.5M floats = 134MB. Total for all heads = 134MB * 1024 = 137GB. Similarly V = 137GB. Q = 2GB. O = 2GB. Total = 278GB. Ah! I was calculating 2MB per head earlier, but I forgot that each query block loads the FULL K of 2MB. Yes, 64 query blocks * 2MB = 128MB per head. That's correct. So with my first kernel (BQ=8, BK=16), total traffic is roughly: - K: 128MB/head * 1024 = 131GB - V: 131GB - Q: 2GB - O: 2GB Total: ~266GB At 3.35TB/s, that's 79ms. My kernel could never beat a flash attention kernel that does only 8GB. Wait, but does PyTorch SDPA use the optimal flash attention? For FP32 on H100, I think it might not even use flash attention (since flash attention typically only supports FP16/BF16). If it uses cuDNN attention, cuDNN might also use a similar tiling algorithm. But cuDNN is highly optimized and might use larger tiles or Tensor Cores. If the reference uses the `math` backend, it explicitly computes: `softmax(Q @ K^T / scale) @ V` This involves: - Q @ K^T: [1024, 512, 1024] batch GEMM. cuBLAS handles this very efficiently and might even use Tensor Cores. - Writes S = [1024, 512, 512] = 1GB - Softmax on S: 1GB read + 1GB write - S @ V: [1024, 512, 512] @ [1024, 512, 1024] = 1GB read + 2GB read + 2GB write Total: Q(2GB) + K(2GB) + S(2GB) + V(2GB) + O(2GB) = 10GB 10GB / 3.35TB/s = 3ms for memory. Plus compute for GEMMs. cuBLAS GEMM at these sizes... [1024, 512, 1024] is actually [1024 batches, M=512, K=1024, N=512]. This is a fairly small GEMM. cuBLAS might not achieve peak performance for such sizes. The compute for Q @ K^T is B*H*N*N*D = 32*32*512*512*1024 = 274B FMAs. 274B / 67 TFLOPS = 4ms. Plus the second GEMM: same 274B FMAs = 4ms. Plus softmax: negligible. Total reference = ~11ms if compute bound, or ~10ms if memory bound. If I write a custom fused kernel with traffic ~266GB, it would take ~80ms. That would be MUCH SLOWER than the reference. So the reference is very likely NOT using the math backend. It's probably using cuDNN attention or some other optimized path that avoids materializing S. Hmm, but if the reference uses the math backend, my custom kernel with 266GB traffic would still lose. Wait, maybe I should write a better custom kernel that materializes S and uses cuBLAS for GEMMs. Actually, hold on. Can I just use `torch.matmul` inside `ModelNew.forward`? The challenge says to use a custom CUDA kernel. But I can write a kernel that uses `cublas` for the GEMMs and a custom kernel for softmax. That might be fast. Or even simpler: maybe I can use the same approach as the math backend but optimize the softmax. Wait, let's reconsider. Maybe PyTorch FP32 SDPA actually dispatches to something unoptimized on this platform. I need to check what the reference timing would likely be. Actually, I can look at the problem from a different angle: the evaluator says "run faster than the PyTorch reference". If the reference takes X ms, I need < X ms. The main variable I control is the algorithm. Given that custom kernels for attention without Tensor Cores are usually slower than FlashAttention/cuDNN, but if the reference doesn't use those... Actually, wait. For FP32, SDPA might fall back to `math` because: - FlashAttention doesn't support FP32 - Memory-efficient attention might not support FP32 either (or it's disabled for some shapes) - cuDNN attention requires certain conditions If the reference falls back to `math`, it computes three separate operations: matmul, softmax, matmul. My custom fused kernel would likely be faster if I can make it reasonably efficient, even with 266GB traffic? No, 266GB is way too much. Unless... wait. My traffic calculation was for the worst case where K/V aren't cached. But with the custom kernel running on 132 SMs, and 1024 blocks, the K might stay in L2 if multiple blocks from the same head run on the same SM. Actually, let's recalculate more carefully. A head has 64 blocks (for BQ=8). Each block reads 2MB of K. 64 blocks need to run. H100 has 132 SMs. So it's entirely possible that, say, 10 blocks from the same head run on 10 different SMs, none sharing cache. Then each reads K from HBM. But maybe the scheduler puts blocks from the same head together? Actually, CUDA block scheduling for 2D grid (B, H) usually assigns blocks in order to SMs. If the grid is (32, 32) and each block processes multiple query groups... in my first kernel there was no z dimension, so grid=(32,32), each block processed 1 query group. The blocks would be scheduled as (0,0), (1,0), ..., (31,0), (0,1), etc. So consecutive blocks in x are from the same head. This means blocks from the same head will indeed run on the same SMs, potentially sharing L2 cache. But L2 cache is only 50MB. K for one head is 2MB. We can fit 25 heads in L2. With 32 blocks per head all wanting K, and they might not all fit in L2. But K is 2MB, and if we read it 64 times (for 64 query groups), L2 can serve most of those reads if the working set is 2MB. Actually, L2 has much higher bandwidth than HBM (maybe 10-20 TB/s). So if K stays in L2 after the first read, subsequent reads come from L2. Then total K traffic from HBM is only 2MB per head, not 128MB! Is this realistic? For a single head, all 64 blocks access the same K. The first block reads K from HBM into L2. The next 63 blocks read from L2. So K HBM traffic is indeed 2MB per head, not 128MB! Similarly for V. So my total traffic calculation was overly pessimistic. But what about Q? Each block reads a different Q tile. Those don't share across blocks. So total traffic with L2 caching of K/V: - Q: 2GB - K: 2GB (from HBM to L2, then served from L2) - V: 2GB - O: 2GB Total: 8GB At 3.35TB/s, that's 2.4ms. Plus some L2 cache overhead. This is excellent! But wait, is L2 cache big enough? Per head, K+V = 4MB. With 132 SMs and blocks racing through, if 132 blocks are from different heads, L2 needs 132*4MB = 528MB of L2 cache, but we only have 50MB. So most K/V will be evicted. However, let's think about the wave structure. 1024 blocks total. First wave of ~132 blocks. If these are from 132 different x-coordinates (batches) and the same y-coordinate (head), then L2 needs only 4MB (one head). But if they're from different heads, L2 needs 132*4MB. Given grid=(B,H)=(32,32), blocks are launched in x-major order. Block sequence: (0,0), (1,0), ..., (31,0), (0,1), (1,1), ..., (31,31) So the first 32 blocks are all head=0, batch=0..31. They all share the same head! They all need K/V for head 0. So L2 only needs 4MB for the first 32 blocks. Then blocks (0,1) to (31,1) need head 1. K for head 0 might be evicted. But within the first 32 blocks, all use head 0's K/V. If they all run concurrently across 32 SMs (out of 132), then L2 bandwidth is shared. The remaining 100 SMs work on other blocks from other heads. Actually, the block scheduler usually assigns blocks to SMs in a round-robin or similar. The exact order isn't strictly x-major across the chip, but within a cluster or GPC it might be. On H100, with 132 SMs, and 1024 blocks, a wave of 132 blocks would include roughly 4 blocks from head 0 (batches 0,1,2,3), 4 from head 1, etc. Actually no, if scheduler takes blocks in sequence: (0,0), (1,0), ..., (131,0) but batch only goes up to 31. So first 32 blocks: (0..31, 0). Next 32: (0..31, 1). Next 32: (0..31, 2). Etc. For a wave of 132 blocks, we'd get: - 32 blocks from head 0 (all batches) - 32 blocks from head 1 - 32 blocks from head 2 - 32 blocks from head 3 - 4 blocks from head 4 So 5 different heads are active. K+V per head = 4MB. 5 heads = 20MB. L2 is 50MB. So K/V for these heads might entirely fit in L2! As blocks from head 0 complete, new blocks from head 4, 5, etc. join the active set. If the scheduler maintains a pool of ~132 blocks and replaces completed ones, we might have ~10-20 heads active at a time. 20 heads * 4MB = 80MB > 50MB L2. So some eviction occurs. But still, K/V traffic from HBM is much less than the worst case. Maybe 10-20GB total for K+V instead of 256GB. This makes fused attention much more viable. The original kernel might actually be decent if L2 caching works well. Now, for performance optimization, I should: 1. Increase BQ to reduce the number of query blocks per head (reducing K/V reads, even from L2, and improving L2 hit rates). 2. Use all threads efficiently. 3. Use shared memory for K/V tiles to amortize the cost of loading from L2/L1. With BQ=64: - 8 blocks per head. - A wave of 132 blocks would include ~4 heads fully + some partial. - L2 pressure = ~4 heads * 4MB = 16MB. Excellent caching. - Per head, K loaded 8 times from L2, but only once from HBM (if cached). - Total K HBM traffic = 2GB. V = 2GB. So BQ=64 seems sweet. But can we fit the kernel with BQ=64? If each block processes 64 queries: - With 512 threads per block (16 warps), each warp handles 4 queries. - Warp w handles queries 4w, 4w+1, 4w+2, 4w+3. - Each warp has 32 lanes. - Lane j handles key j for all 4 queries? No, BK would be <=32. - If BK=32, lane j handles key j for all 4 queries. - Score per lane: 4 dot products of D=1024 each. This is similar to my earlier calculation. For V accumulation: - Lane j handles D dimensions: j, j+32, j+64, ... for all 4 queries. - With D=1024 and 32 lanes, each lane handles 32 dims. - For BK=32 keys and all D, each lane does 32*32 = 1024 FMAs per query? No, per 4 queries. Actually, let's think about the total compute intensity. Total FLOPs = 2 * B*H*N*N*D = 2 * 32*32*512*512*1024 = 549B FMAs. At 67 TFLOPS = 8.2ms. If memory traffic is ~8GB (Q+K+V+O), that's 2.4ms at 3.35TB/s. So the kernel is mostly compute bound. The 8ms compute suggests we need high FMA throughput. With a custom kernel using scalar FMAs on CUDA cores, can we achieve 67 TFLOPS? - 67 TFLOPS requires nearly 100% utilization of FP32 units. - A warp can do 1 FMA per thread per clock. At 2GHz, a warp does 64 GFLOPS. An SM with 4 warps (128 threads) does 256 GFLOPS. 132 SMs = 33.8 TFLOPS. Wait, that's half of 67 TFLOPS. Actually, modern GPUs have TWO FP32 units per SM per scheduler? Let me think. H100 has 128 FP32 CUDA cores per SM. At 2GHz, that's 256 GFLOPS per SM (128 * 2). 132 SMs = 33.8 TFLOPS. But the spec said 67 TFLOPS. How? Ah, H100 has FP32 Tensor Cores too? No, the 67 TFLOPS is peak FP32 throughput for the whole chip. That implies some SM-level parallelism I'm not accounting for. Actually, H100 has 128 FP32 cores per SM, each can do 1 FMA per cycle. 1 FMA = 2 FLOPs. So per SM: 128 * 2 * 2GHz = 512 GFLOPS. 132 * 512 GFLOPS = 67.5 TFLOPS. Yes! So an SM with 128 active threads (4 warps) can achieve peak if all threads issue FMAs every cycle. My kernel needs to issue enough FMAs to hide memory latency. With BQ=64 and BK=32, and each block having 512 threads (16 warps): Each warp processes 4 queries, 32 keys. Score computation: 4 queries * 32 keys * 1024 D = 131k FMAs per warp. V accumulation: 4 queries * 32 keys * 1024 D = 131k FMAs per warp. Total: 262k FMAs per warp. At 512 GFLOPS per SM, 16 warps on an SM = 16*32=512 threads. If all issue FMAs, that's peak utilization. 262k FMAs / (512 threads * 1 FMA/cycle/thread) = 262k cycles. At 2GHz = 0.13ms per block. 1024 blocks / 132 SMs = 7.76 waves. 7.76 * 0.13ms = 1.0ms. Hmm, that seems too optimistic. But memory will slow it down. Actually, I realize I'm overcomplicating this. Let me just try to write a solid BQ=64 kernel with BK=32, 16 warps (512 threads), and see if it compiles and works. But wait, I can't easily test performance without submitting. Let me try a more conservative approach: write a kernel that's known to be well-structured, use a moderate BQ like 32 or 64, and submit. Actually, I have another idea. Since the problem size N=512 is relatively small, and D=1024 is large, what if I use a different parallelism strategy: Each block = 1 warp = 32 threads. Grid = (num_heads * num_batches * num_queries, 1) = 32 * 32 * 512 = 524288 blocks. But CUDA has a limit on grid dimensions: 2^31-1 for x, so 524288 is fine. Each block/thread handles 1 query. Algorithm per thread: ```cpp float q_val = Q[qid * D + threadIdx.x]; // Load only 1 element of Q? No. ``` Wait, no. A single thread can't efficiently compute the full dot product of D=1024, but we can use warp-level reductions. Per warp (32 threads), for 1 query: 1. Load all of Q: 1024 floats. With 32 threads, each thread loads 32 floats. Coalesced (thread t loads Q[t*32:(t+1)*32]). 2. Iterate over all keys (512 keys): a. Load K[k][t*32:(t+1)*32] for thread t. Coalesced. b. Compute partial dot product per thread: thread t computes dot(Q_chunk, K_chunk). c. Warp shuffle reduce to get score. d. Track max for softmax. 3. Iterate over keys again for V accumulation: a. Load V[k][t*32:(t+1)*32]. b. Broadcast weights via shuffle. c. Accumulate to O. This is actually a very clean approach! Let me formalize: Each warp handles 1 query. 1 block = 32 warps = 32 queries. Grid = (32 heads * 32 batches * 16 query_groups) = 16384 blocks. Wait, 512 queries / 32 queries per block = 16 blocks per head. Total blocks = 32 * 32 * 16 = 16384. Algorithm per warp (query q): ```cpp // Load Q into shared or local memory float q_local[32]; // 32 floats per thread for (int i = 0; i < 32; ++i) { q_local[i] = Q[q * D + threadIdx.x * 32 + i]; // thread t loads elements [t*32, (t+1)*32) } // Wait, threadIdx.x is 0..31 for warp 0? But a block has 32 warps. // threadIdx.x % 32 = lane_id // We can load: q_local[i] = Q[q*D + lane_id*32 + i]; ``` For score computation: For each key k = 0..511: float dot = 0; for (int i = 0; i < 32; ++i) { float k_val = K[k*D + lane_id*32 + i]; // same pattern as Q dot += q_local[i] * k_val; } // warp sum #pragma unroll for (int offset = 16; offset > 0; offset /= 2) dot += __shfl_xor_sync(0xFFFFFFFF, dot, offset); dot *= scale; // update max and sum // store dot in local array? We need all 512 scores. ``` For 512 scores, storing in local memory: 512 floats = 2KB per thread. Local memory is global-backed and coalesced if laid out nicely. But 512 scores * 32 threads per warp = 64KB per warp. For 32 warps = 2MB per block. That's the same order as O output, so local memory spill might be OK. But actually, we can do online softmax WITHOUT storing all scores! The trick: process keys in groups, compute tile max, update running max, compute exp relative to tile max, accumulate to sum. For online softmax, we need to know the max to compute exp. If we process one key at a time: ```cpp float m = -inf; float l = 0; float o_local[32] = {0}; for (int k = 0; k < N; ++k) { float s = 0; for (int i = 0; i < 32; ++i) { s += q_local[i] * K[k*D + lane_id*32 + i]; } warp_reduce(s); s *= scale; float new_m = fmaxf(m, s); float corr = expf(m - new_m); float p = expf(s - new_m); l = l * corr + p; for (int i = 0; i < 32; ++i) { o_local[i] = o_local[i] * corr + p * V[k*D + lane_id*32 + i]; } m = new_m; } ``` Wait! This is BRILLIANT! With warp-level online softmax, each warp processes ONE query. All 32 threads in the warp cooperate: - Each thread handles a 32-dim chunk of Q, K, V. - The score is computed as a warp reduction. - The output is accumulated in each thread's 32-dim local accumulator. - No shared memory needed! - No need to store scores! Let me verify the algorithm carefully: 1. Thread t loads Q[q][t*32:(t+1)*32] into 32 registers. 2. For each key k: a. Thread t loads K[k][t*32:(t+1)*32] into 32 registers (or directly use). b. Thread t computes partial dot product: sum_i Q[i] * K[i] for its 32 dims. c. Warp reduce across 32 threads to get full score. d. Apply scale. e. `new_m = max(m, score)` f. `corr = exp(m - new_m)` g. `p = exp(score - new_m)` h. `l = l * corr + p` (but we don't actually need l if we keep unnormalized O) Wait, we DO need l for the final normalization. Or do we? Actually, with the standard online softmax formulation: ``` O_new = O_old * exp(m_old - m_new) + exp(s - m_new) * V l_new = l_old * exp(m_old - m_new) + exp(s - m_new) ``` At the end, O_final = O_new / l_new. This is correct. i. Thread t loads V[k][t*32:(t+1)*32]. j. Thread t updates each output dim: `o_local[i] = o_local[i] * corr + p * v_local[i]` 3. After all keys: `o_local[i] /= l;` // wait, l is the sum of exp(s_i - m_final) Actually, O_new already has the exp(s_i - m_new) weights. We just need to divide by l_new. 4. Thread t writes O[q][t*32:(t+1)*32] = o_local. But wait, there's a subtle issue: `corr = exp(m - new_m)` is the same for all threads in the warp. We should broadcast it. And `p = exp(s - new_m)` is the same for all threads (since s is the same for the whole warp after reduction). We can broadcast p too. Let's trace through carefully: ```cpp float q_local[32]; // load Q for (int i = 0; i < 32; ++i) q_local[i] = Q[q*D + lane*32 + i]; float m = -1e20f; float l = 0.0f; float o_local[32] = {0}; for (int k = 0; k < N; ++k) { float dot = 0.0f; // load K chunk into k_local? We can load directly. for (int i = 0; i < 32; ++i) { float k_val = K[k*D + lane*32 + i]; dot += q_local[i] * k_val; } // reduce dot across warp #pragma unroll for (int offset = 16; offset > 0; offset /= 2) dot += __shfl_xor_sync(0xFFFFFFFF, dot, offset); dot *= scale; // broadcast m and l? Actually all threads have the same m and l. // Wait, m and l are scalars, each thread keeps its own copy. // But they all update identically because dot is the same. float new_m = fmaxf(m, dot); float corr = expf(m - new_m); float p = expf(dot - new_m); l = l * corr + p; // update O accumulators for (int i = 0; i < 32; ++i) { float v_val = V[k*D + lane*32 + i]; o_local[i] = o_local[i] * corr + p * v_val; } m = new_m; } // normalize for (int i = 0; i < 32; ++i) { o_local[i] /= l; O[q*D + lane*32 + i] = o_local[i]; } ``` But wait, `m` and `l` are updated identically on all threads, but without synchronization. Is that OK? Let's see: - `dot` is the same on all threads after the warp shuffle reduction. - `new_m = fmaxf(m, dot)` where `m` started the same on all threads (and remains the same). - `corr = expf(m - new_m)` is the same on all threads. - `p = expf(dot - new_m)` is the same. - `l = l * corr + p` updates identically. - `m = new_m` updates identically. But what about divergence? All threads in the warp execute the same instructions because the data is identical after the shuffle. Actually, no divergence because `m`, `dot`, `new_m` are the same for all lanes. Even though each thread computes its own `dot` locally before the shuffle, after the shuffle all threads have the same `dot` value (because `__shfl_xor` exchanges and adds... wait, no. Let me re-examine. `dot` is the partial sum for each thread's chunk. After the warp reduction `dot += __shfl_xor(...)`, all threads in the warp hold the TOTAL sum. Yes, because `__shfl_xor` does a butterfly reduction. After all shuffles, each thread has the full sum. So yes, all threads have the same `dot` after reduction. Then the rest of the computation is identical for all threads. So no divergence! But there's a register issue: each thread needs: - q_local[32] - o_local[32] - k_val and v_val can be loaded directly without an array, or in a loop. - So registers ≈ 64 + ~10 temp = 74. That's fine. Actually, inside the loop, if we load one K element at a time: ```cpp for (int i = 0; i < 32; ++i) { float k_val = K[k*D + lane*32 + i]; dot += q_local[i] * k_val; } ``` `q_local` needs 32 registers. `k_val` is a temp. `dot` is one register. So total ~34 registers. Then for V accumulation: ```cpp for (int i = 0; i < 32; ++i) { float v_val = V[k*D + lane*32 + i]; o_local[i] = o_local[i] * corr + p * v_val; } ``` `o_local` needs 32 registers. This is fine. Total registers = 32 (q) + 32 (o) + m + l + corr + p + dot + k_val/v_val + loop counters. Max ≈ 80 registers. That's well within 255. Now, memory access pattern: - Q load: thread lane loads Q[q*D + lane*32 + i] for i=0..31. These are contiguous within the thread, but strided across threads by 32 floats = 128 bytes. - Thread 0 loads bytes 0..127 - Thread 1 loads bytes 128..255 - ... - Thread 31 loads bytes 3968..4095 - This is perfectly coalesced! A warp issues one transaction per cache line (128B), covering the full 4096B with 32 threads. Same for K and V loads. Same for O writes. This access pattern is EXCELLENT. Simplicity: very simple. No shared memory. No __syncthreads. Performance analysis: - 1 block = 32 warps = 32 queries. - Total blocks = 32 * 32 * 16 = 16384. - 16384 blocks on 132 SMs = 124 blocks per SM if perfectly distributed. - Each block processes 32 queries. Per query: 512 keys. - Per warp: for each key, 32 iterations of dot product + 32 iterations of V accumulate. - Each dot product iteration: 1 FMA + 1 load (K) + 1 load (q_local, already in register). - Each V iteration: 1 FMA + 2 multiplies + 1 load (V). - Plus warp shuffle for score. - Compute per warp: 512 keys * 32 dims * ~5 ops = 81920 ops. - Memory per warp: Q (32*4B) + K (512*32*4B) + V (512*32*4B) + O (32*4B) = 128B + 65536B + 65536B + 128B = 131KB per query. Wait, per warp (1 query): K loaded 512 times * 32 floats * 4 bytes = 65536 bytes. V same. Q and O negligible. - Total memory per block (32 queries): 32 * 131KB = 4.2MB. - K and V are the same across all 32 warps in a block! But each warp loads K and V independently. - However, with L1 cache, the first warp loads K[0]...K[511] and V[0]...V[511] into L1. The other 31 warps read from L1! - Wait, can L1 serve 31 other warps? The L1 cache on H100 is ~128KB or ~256KB per SM. K+V for 512 keys * 32 dims... no, that's for all keys but only 32 dims per thread. Actually, K for one key is 4096 bytes. For 512 keys, K = 2MB. That's too big for L1. But each warp only accesses K[k] sequentially. When warp 0 loads K[0], that goes to L1. When warp 1 also loads K[0] (for its query), it might still be in L1. But with 32 warps all loading K sequentially, it's likely that: - Warp 0 loads K[0] to L1. - Warp 1 loads K[0] from L1 (or L2). - Warp 2 loads K[0] from L1 (or L2). - ... - This is effectively broadcasting K[0] to all warps once, then they move on to K[1]. The L1 cache might not have enough bandwidth to serve all 32 warps simultaneously for the same address, but it can probably serve a few. They might serialize, but the cache hit means it's much faster than HBM. However, if 32 warps in a block all try to access the same K[k] at the same time, there will be contention. `__syncwarp()` doesn't help. Wait, the warps in a block execute independently and might not be perfectly synchronized. They'll naturally stagger. But here's a key optimization: let ALL warps in a block process the SAME key at roughly the same time. This happens naturally if they iterate over k in lockstep (though not guaranteed). But even without lockstep, the L2 cache can serve K and V. L2 bandwidth on H100 is much higher than HBM. If K and V are in L2, loading is fast. For this kernel, total HBM traffic: - Q: each query loaded once = 2GB - K: each key loaded once per block? No, each block loads K 32 times (once per warp). 16384 blocks * 2MB = 32GB - V: same 32GB - O: 2GB Total = 68GB At 3.35TB/s = 20ms. But with L2 caching, most K and V loads come from L2, so HBM traffic is just 2GB for K and 2GB for V (loaded once per key into L2). So effective HBM traffic might be closer to 8GB = 2.4ms. The 68GB total includes L2 traffic. L2 bandwidth is ~10 TB/s, so 68GB / 10TB/s = 6.8ms. Compute: each warp does 512 keys * 32 dims * ~8 ops = 131k ops per query. 524288 queries * 131k = 68.7B ops. At 67 TFLOPS = 1ms. But wait, that's the total FLOPs, not just this part. Actually full attention FLOPs is 549B. 549B / 67T = 8.2ms. Hmm, my per-query calculation might be wrong. Let me recalculate: For one query: - Scores: 512 keys * 1024 dims = 524k FMAs. - V accumulation: 512 keys * 1024 dims = 524k FMAs. - Total: ~1M FMAs = 2M FLOPs. - 524288 queries * 2M = 1.05 T FLOPs. - At 67 TFLOPS = 15.7ms. My earlier 549B came from 2*B*H*N^2*D = 2*32*32*512*512*1024 = 549B. Yes, that's correct. Each warp has 32 threads. The partial dot product for 32 dims per thread requires 32 FMAs. Then warp shuffle reduction (negligible). So 32 FMAs per thread per key for scores + 32 FMAs per thread per key for V = 64 FMAs. Total FMAs per thread = 512 keys * 64 FMAs = 32768 FMAs. Total threads = 16384 blocks * 32 warps * 32 threads = 16.7M threads. Total FMAs = 16.7M * 32768 = 549B. Matches. At 67 TFLOPS: 549B / 67T = 8.2ms. At L2-bound 68GB / 10TB = 6.8ms. So total time might be around 10-15ms. Is the reference faster or slower? If reference uses cuBLAS for two GEMMs, each GEMM [1024, 512, 1024, 512] is actually: - GEMM 1: [1024 batches, M=512, N=512, K=1024] - This is large enough for cuBLAS to achieve good efficiency. cuBLAS might get 80% of peak on H100. - But each GEMM only has 512x512 output, which is relatively small. - cuBLAS batch GEMM: 274B FMAs each = 4ms each at 67 TFLOPS. - Plus softmax: ~1ms. - Plus materializing S: extra memory traffic. Actually, for a batch size of 1024, a [512, 512, 1024] GEMM total FLOPs = 1024 * 512 * 512 * 1024 * 2 = 549B. cuBLAS on H100 might achieve say 50-100 TFLOPS for this size? 549B / 80TF = 6.8ms per GEMM. Two GEMMs = 13.6ms. Plus memory for S = maybe +2ms. Total ~16ms. If I write a fused kernel taking ~10-15ms, I might be competitive. But can I do better than 10ms? What if I use shared memory to explicitly prefetch K and V tiles, avoiding redundant loads entirely within a block? With 32 warps in a block all processing the same K and V: - K tile: 32 keys * 1024 dims = 128KB. Wait, BK=32, D=1024, that's 128KB. - With BK=32, we process 512/32 = 16 key tiles. - For each key tile, all 32 warps in the block need the same K[kt:kt+31][:] and V[kt:kt+31][:]. - Store K tile and V tile in shared memory: 2 * 128KB = 256KB. Exceeds 227KB limit. What if BK=16? K tile = 16*1024*4 = 64KB. V tile = 64KB. Total = 128KB. Plus smem_q for 32 queries? But in this warp design, each warp handles 1 query and loads its own Q into registers. So no need for Q in smem. Total smem = 128KB. Under 227KB. Good! With BK=16: - Number of key tiles = 512/16 = 32. - Each block loads each K tile once into smem (instead of 32 times via L2). - This reduces L2 traffic significantly, and ensures perfect reuse. Wait, how do we load K tile into smem? - Block has 1024 threads. K tile is 16*1024 = 16K floats. - Each thread loads 16 floats into smem. Perfectly coalesced. - __syncthreads() to ensure all K is loaded. - Then all 32 warps compute their dot products using smem_k. - Then all 32 warps do V accumulation using smem_v. - __syncthreads() before loading next tile. This is a much better use of shared memory! Let me design this: Block = 1024 threads? Actually 32 warps = 1024 threads. That's the max. But wait, with 32 warps and 1024 threads per block, and 2 blocks per SM (2048 threads), that's fine. Occupancy might be limited. Actually, for Ampere/Hopper, max threads per SM = 2048. So with 1024 threads per block, we can have 2 concurrent blocks per SM. That gives good occupancy. But 1024 threads means register pressure matters. 1024 threads * 80 registers = 81K registers per SM. But an SM only has 64K registers. So we'd be limited to 1 block per SM if each thread uses 64 registers. With 80 registers, 1024*80 = 81K > 64K, so the compiler might spill or we'd only fit 1 block = 1024 threads per SM. What if we use 512 threads per block (16 warps)? - 512 threads * 80 regs = 40K registers. We can fit 1 block (40K) comfortably, and maybe 1.5 but effectively 1 block per SM. - Or if we use 64 registers: 512*64 = 32K, so 2 blocks = 64K, perfect! - But can we keep register usage under 64? Let's count: - q_local: 32 floats - o_local: 32 floats - loop temp: 3 floats - m, l, corr, p, s, dot: 6 floats Total = 73 floats. That's a lot. The compiler might keep q_local and o_local in registers, but might spill some. Actually, q_local and o_local are arrays. The compiler stores them in registers if possible. 64 registers is tight. Alternative: process Q, K, V in chunks of 16 instead of 32: - q_local[16], o_local[16] - Total registers: 16 + 16 + 6 = 38. Very comfortable! - But we need 1024/16 = 64 chunks. More loop iterations. - D chunk = 16 means each thread handles 16 dims. - 512 keys * 64 chunks * 2 FMAs = 65536 FMAs. Same as before. With D chunk = 16, and 512 threads per block: - Q load: 512 threads * 16 floats = 8192 floats per block. Coalesced. - Block processes 16 queries? (16 warps, 1 query per warp) - Wait, 512 threads = 16 warps. Each warp handles 1 query. So block handles 16 queries. - Total blocks = 32 * 32 * (512/16) = 32*32*32 = 32768 blocks. This is a lot of blocks. 32768 / 132 = 248 blocks per SM. Each block does 65536 FMAs per thread. 512 threads * 65536 = 33.5M FMAs per block. Total FMAs = 1.1T? Wait, total work should be 549B FMAs. 32768 blocks * 512 threads * 65536 FMAs = 1.1T FMAs. That's twice as much. Oh wait, I miscalculated. For one query, the total FMAs is 512 keys * 1024 dims * 2 (score + V) = 1,048,576 FMAs. With chunk size 16: 512 keys * 64 chunks * 16 dims * 2 = same. Per thread: for one key, chunk loop does 16 iterations. 16 FMAs for score + 16 for V = 32 FMAs. For 512 keys: 512 * 32 = 16384 FMAs per thread. Total threads = 32768 blocks * 512 threads = 16.7M threads. Total FMAs = 16.7M * 16384 = 274B. That's only half of 549B. Why? Because 16384 FMAs per thread * 16.7M threads = 274B. But we need 549B. Ah, I see. The score computation alone is 512 keys * 1024 = 524k FMAs. But if 512 threads process 512 keys... no, in my design, one thread processes ONE query, not one key. So for one query with 512 keys and D=1024: - Score: 512 * 1024 = 524k FMAs. - V: 512 * 1024 = 524k FMAs. - Total: ~1M FMAs. With chunk size 16: - Score: 512 keys * 64 chunks * 16 FMAs = 524k FMAs. - V: same. Per warp (1 query): total FMAs = ~1M. Per block (16 warps/queries): 16M FMAs. Total blocks = 32*32*32 = 32768. Total FMAs = 32768 * 16M = 524B. Yes, that matches. Per thread FMAs: 1M FMAs / 32 threads per warp = 32768 FMAs per thread. (Actually, 1M / 32 = 32k) 16.7M threads * 32k FMAs = 535B. Close enough. OK so D chunk size doesn't change total work. It just changes the loop structure. Now, with shared memory for K and V: - Block = 16 warps = 512 threads, processing 16 queries. - K tile: BK=16, D=1024, so smem_k = 16*1024*4 = 64KB. - V tile: same 64KB. - Total: 128KB. - With 227KB limit per block, and 64K registers per SM, we might fit 1 block per SM (128KB smem, ~40K registers). - Or if we use fewer threads per block... With 512 threads per block and 128KB smem: - Max concurrent per SM: limited by smem to 1 block (128KB) or 2 blocks (256KB > 227KB, no). So 1 block per SM. - Total concurrent blocks = 132. - Waves = 32768 / 132 = 248 waves. - That's a lot of waves! If each block takes a short time, the overhead of launching all those waves adds up. What if we reduce the number of blocks by increasing queries per block? Block = 64 warps = 2048 threads? Exceeds max 1024. Not possible. Block = 32 warps = 1024 threads. Then block handles 32 queries. Total blocks = 32*32*16 = 16384. Waves = 16384/132 = 124. Still many waves, but better. Block = 64 queries? Not possible with 1 thread per query (max 1024 threads). Unless we have 1 block per 64 queries with 16 threads per query. But 64*16=1024 threads. Then each query is handled by 16 threads cooperatively. That changes things. Actually, the cooperative approach with 1 warp per query and multiple warps per block seems clean. Let's commit to: - 1 block = 32 warps = 1024 threads - Block handles 32 queries - Total blocks = 32 * 32 * 16 = 16384 - Uses shared memory for K and V tiles if we want. But register pressure with 1024 threads is a concern. If each thread needs 64 registers, total = 65536. An SM has 65536 registers. So exactly 1 block per SM. This works! With 1 block per SM = 1024 threads per SM. H100 SMs can handle up to 2048 threads, so occupancy is 50% in terms of threads, but 100% in terms of blocks (1 block is the max due to smem). Actually, 50% thread occupancy is fine. 16384 blocks / 132 SMs = 124 waves. Each wave of 132 blocks. At per-block time of, say, 0.5ms, total = 62ms. That's too long. Wait, 0.5ms per block is way too high. Let's estimate: Each warp does ~1M FMAs. 1024 threads = 32M FMAs per block. At 512 GFLOPS per SM: 32M / 512G = 62 microseconds. 124 waves * 62us = 7.7ms. That seems reasonable! Plus memory loads. 7.7ms is competitive. But can we reduce waves? With 1024 threads per block processing 32 queries, it's 16384 blocks. If we could process 64 queries per block, total blocks = 8192, waves = 62. But we'd need 2048 threads, which exceeds the 1024 limit. Alternatively, process 32 queries with 512 threads (16 warps). Then 2 blocks per SM concurrently if registers allow. If 512 threads * 64 regs = 32K registers, then 2 blocks per SM = 64K registers. OK. And smem per block = 128KB. 2 blocks = 256KB > 227KB. Not OK. So we'd need to reduce smem. If we don't use shared memory and instead rely on L2 cache, we can have 2 blocks per SM with 512 threads each. Total concurrent = 264. Waves = 32768 / 264 = 124. Same issue. If we don't use shared memory and use 256 threads (8 warps) processing 8 queries: - Total blocks = 32*32*64 = 65536. - Waves = 65536 / (132 * 8) ??? With 256 threads * 64 regs = 16K regs. Up to 4 blocks per SM = 64K regs. So 4 blocks = 1024 threads. - Concurrent blocks = 132 * 4 = 528. - Waves = 65536 / 528 = 124. It seems like regardless of block size, if we process all queries with 1 thread per query, we get ~124 waves. The key is that total parallel work is 524288 queries, and SMs can run ~1000-2000 threads. With 132 SMs * ~1500 threads = 198k threads concurrent. 524k threads / 198k = 2.6 waves. Wait, I think my wave calculation was wrong. Total threads = 524288. Concurrent threads across all SMs = 132 SMs * threads_per_SM. If threads_per_SM = 1024 (1 block of 1024), total concurrent = 132 * 1024 = 135k. Waves = 524k / 135k = 3.88 waves. If threads_per_SM = 2048 (2 blocks of 1024 or 4 blocks of 512), total concurrent = 132 * 2048 = 270k. Waves = 524k / 270k = 1.94 waves. So we're looking at 2-4 waves. That's very good! I was dividing blocks by SMs before, but it should be total threads by concurrent threads. For shared memory limited to 1 block per SM (1024 threads), waves = 524k / (132*1024) = 3.88. For 2 blocks per SM (if possible), waves = 2. For 512 threads with 2 blocks = same 1024 threads per SM. Waves = 3.88. Actually 512 threads * 2 blocks = 1024 threads per SM. Same as 1 block of 1024. To get to 2048 threads per SM: - 2 blocks of 1024: each block would need < 113KB smem. With BK=16 and K+V=128KB, we can't. - 4 blocks of 256: each block needs < 56KB smem. We can if we don't use smem or use very small tiles. - 4 blocks of 512: need < 56KB smem, and 512*128 regs? 512*64=32K, 4 blocks=128K > 64K. Not possible. So max concurrency is likely 1 block of 1024 threads, or 2 blocks of 512 with no shared memory. Either way, ~1024-2048 threads per SM. Let's say 1024 for safety. Waves = 3.88. Total runtime ≈ 4 * block_time. Block_time for 32 warps processing 32 queries: - Each warp: 1M FMAs. At 50% peak (due to memory), ~0.1ms. - 4 waves * 0.1ms = 0.4ms. That's very fast! But this assumes perfect utilization. Let's be more conservative: 2ms. Now, how does the reference do? The reference with cuBLAS + softmax might take 5-15ms depending on backend. So if I can write a clean fused kernel with no shared memory and good register usage, I might actually beat the reference. Let me think about whether to use shared memory or not. - Without shared memory: relies on L2 cache for K/V reuse among warps. Clean, no sync overhead. - With shared memory: explicit reuse, but limits occupancy. For N=512, D=1024, K and V for BK=16 keys = 128KB shared memory. This is fine if we use 512 threads per block and only 1 block per SM. Actually, I just realized: with shared memory for K and V, we MUST synchronize threads. But each warp handles one query, and the score computation for each query needs ALL K dimensions. So we can load a chunk of K[0:BK][0:DT] into smem, then all warps compute their partial scores using this chunk. But partial scores must be accumulated across chunks. This means we need a per-query score accumulator. But the score is per query per key, not per query per chunk. Wait, no: the score is the dot product of Q and K, which is sum over D. If we compute partial dot products for a chunk of D, we accumulate them. Actually, for BK=16 and D chunk = 64: - Load K[0:16][0:64] into smem: 16*64 = 1024 floats = 4KB. - All 32 warps load their Q[q][0:64] into registers. - Warp q computes partial scores for q and keys 0..15 using chunk 0..63. - Then load next chunk of K and Q. - After all chunks, we have accumulated the full scores. This doesn't require a score array. Each warp just accumulates s_k for each key k: ```cpp float s[16]; // registers for key_tile: for d_chunk: load K chunk for k=0..15: s[k] += q[d] * K[k][d]; ``` But `s[16]` per warp = 16 registers. And we need 32 lanes computing. If each lane handles one key, then lane k handles s[k] for ONE query. So each lane needs only 1 register per query. But we have 32 lanes and 16 keys. Some lanes are idle. If each warp handles 1 query and BK=16: - 32 lanes, 16 keys. Each lane handles 2 keys? Or one key with 2 lanes? - Actually, 16 keys and 32 lanes means each key gets 2 lanes. 2 lanes can split D. - Lane 0,1 handle key 0, splitting D in half. - This uses all threads and parallelizes dot product! But then we need shuffle to reduce across the 2 lanes for each key. This is getting complicated. The no-shared-memory approach with warp-level reductions and 32 dims per thread is much simpler and likely performance-sufficient. Let me finalize on this approach: one-warp-per-query, no shared memory. Wait, I think I should try both the no-shared-memory version and a shared-memory version. But I need to write the code first. Actually, let me reconsider: with the no-shared-memory version, each warp loads K independently. L2 bandwidth is used for every warp. L2 bandwidth on H100 is around 10-20 TB/s. HBM is 3.35 TB/s. If K and V don't fit in L2 for reuse (because many different heads compete), traffic goes to HBM and we're very slow. With shared memory, K is loaded explicitly once per block and reused by all 32 warps. This guarantees we read K from HBM only once per block, and then all 32 warps use it from shared memory. For V, same thing. So shared memory is safer and gives explicit control over memory traffic. I'll design a kernel with: - 1 warp per query - Block = 32 warps = 1024 threads - Block processes 32 queries - Shared memory for K and V tiles - BK=16 keys per tile - D processed in chunks, but wait... Actually, for the dot product with shared memory, let me reconsider the structure: If BK=16 and each warp handles 1 query: - We want to compute 16 scores per query per tile. - A warp has 32 lanes. We can have each lane compute 1 score using half the lanes, or have lanes cooperate (2 lanes per key). With 2 lanes per key and BK=16: 16*2 = 32 lanes. Perfect! - Lane pair (2k, 2k+1) computes score for key k. - They split D in half. Lane 2k handles D[0:512], lane 2k+1 handles D[512:1024]. - After partial dot products, shuffle reduce lane 2k and 2k+1. - Now all 32 lanes have all 16 scores? No, lane 2k has score for key k. To have all lanes know all scores for softmax, we can shuffle. But with only 32 lanes and 16 keys, we can store scores in 16 lanes and have the other 16 lanes idle during softmax? Or we can distribute scores and do warp-level operations. Actually, with 16 scores, we can have: - lane 0..15 store scores for keys 0..15. - lanes 16..31 are idle during score storage. Then for softmax: - Max over 16 lanes: `__shfl_xor` can reduce 16 values if they're in lanes 0..15. But `__shfl_xor` operates on all 32 lanes. - We can do: lanes 0..15 have scores, lanes 16..31 have -inf. - Warp max = fmaxf across all 32 lanes = max of the 16 scores. - Similarly for sum. Then for V accumulation: all 32 lanes can participate, each handling 1024/32 = 32 dims of output. - For key k (0..15), we need p_k for that key. - p_k is stored in lane k. - Lane k broadcasts p_k to all lanes. - Then all lanes update their 32 dims for this key. This works! And all lanes are used efficiently for V accumulation. Register usage per lane: - Q: 1024/2 = 512 floats if we split D... wait, we don't need to keep all Q in registers. We can process D in chunks of DT (say 64), and load Q and K for each chunk. ```cpp float q_local[32]; // DT=64, split across 2 lanes = 32 elements per lane. float s0 = 0; // score for key handled by this lane ``` Actually, with 2 lanes per key, DT could be the full 64 of a chunk. - For D chunk 0..63: - All lanes load Q[0..63]? No, each lane loads only its portion. - Wait, Q for the warp's query needs to be accessible to all lanes. - Since we split D across 2 lanes, each lane loads a sub-chunk. For chunk dt (64 elements): - Lane 2k loads Q[dt:dt+32] - Lane 2k+1 loads Q[dt+32:dt+64] - K[k][dt:dt+32] and K[k][dt+32:dt+64] loaded from shared memory. Actually, K is in shared memory, so all lanes can access it. We'll load K[0:16][dt:dt+64] into smem. For simplicity, let me process D in chunks of 32 to match the warp size: - DT = 32. - Each lane loads exactly one float of Q and one float of K per chunk? No, we want each thread to process a contiguous chunk. Wait, instead of this complexity, let me just go with the no-shared-memory warp design I described earlier. It's simpler and avoids all these tile complications. And with online softmax, it's correctly fused. Final design: one-warp-per-query, no shared memory. - Block = 32 warps = 1024 threads. - Grid = (32*32, 16) where x = batch*head, y = query_group (512/32 = 16). - Each warp handles one query. 32 warps handle 32 queries in a block. Actually, in CUDA, blocks are limited to 1024 threads, and grids are up to 2^31-1 in x, 65535 in y and z. I can use grid = (32*32*16, 1) = 16384 blocks, each with 32 warps. Or to be cleaner: grid = (32, 32, 16), block = (32, 32, 1)? Wait, block can have 3D thread indexing. blockDim.z is up to 64. blockDim.x * blockDim.y * blockDim.z <= 1024. If blockDim = (32, 32, 1) = 1024 threads. Then grid = (32, 32, 16). Within block: - threadIdx.x = lane_id (0..31) - threadIdx.y = warp_id (0..31), representing query offset within the block So query_id = blockIdx.z * 32 + threadIdx.y. head_id = blockIdx.y. batch_id = blockIdx.x. Wait, lane_id = threadIdx.x, and warp_id = threadIdx.y. But threadIdx.x=32 and threadIdx.y=32 means it's 2D. CUDA will create 32 warps, each of 32 threads. But the warp grouping is by threadIdx.x being the fastest varying dimension, so a warp consists of threads (0..31, y) for fixed y. Yes, that's correct. So within warp y, all threads have threadIdx.x = 0..31. This is one warp. And blockIdx.z tells which group of 32 queries in the sequence. In this configuration: - threadIdx.x determines which slice of D the thread handles. - Each thread handles D/32 = 32 dimensions (since D=1024 and 32 threads). For accessing Q: Q[b][h][q][d] - d start = threadIdx.x * 32. - Thread loads Q[q][d_start + i] for i=0..31. Coalesced. For K and V: same pattern. For dot product and accumulation, we use warp-level shuffle for the score. Let me write this kernel out carefully. I need to be careful with numerical precision and correctness. ```cpp __global__ void attention_kernel( const float* __restrict__ Q, const float* __restrict__ K, const float* __restrict__ V, float* __restrict__ O, int B, int H, int N, int D) { int batch = blockIdx.x; int head = blockIdx.y; int q_group = blockIdx.z; int lane = threadIdx.x; int warp = threadIdx.y; int q = q_group * 32 + warp; if (q >= N) return; int head_offset = ((batch * H + head) * N) * D; int q_offset = head_offset + q * D; float scale = 1.0f / sqrtf((float)D); // Load Q: 32 floats per thread float q_local[32]; #pragma unroll for (int i = 0; i < 32; ++i) { q_local[i] = Q[q_offset + lane * 32 + i]; } float o_local[32] = {0.0f}; float m = -1e20f; float l = 0.0f; // Iterate over all keys for (int k = 0; k < N; ++k) { int k_offset = head_offset + k * D; // Compute dot product Q·K float dot = 0.0f; #pragma unroll for (int i = 0; i < 32; ++i) { float k_val = K[k_offset + lane * 32 + i]; dot += q_local[i] * k_val; } // Warp-level reduction #pragma unroll for (int offset = 16; offset > 0; offset /= 2) { dot += __shfl_xor_sync(0xFFFFFFFF, dot, offset); } dot *= scale; // Online softmax update float new_m = fmaxf(m, dot); float corr = expf(m - new_m); float p = expf(dot - new_m); // Load V and update output #pragma unroll for (int i = 0; i < 32; ++i) { float v_val = V[k_offset + lane * 32 + i]; o_local[i] = o_local[i] * corr + p * v_val; } l = l * corr + p; m = new_m; } // Normalize and write output float norm = 1.0f / l; #pragma unroll for (int i = 0; i < 32; ++i) { O[q_offset + lane * 32 + i] = o_local[i] * norm; } } ``` Wait, I need to check something critical about this kernel. The issue: `l = l * corr + p`. This maintains the sum of exp(s_i - m_current). At the end, m_current = max over all keys, so l = sum exp(s_i - max). And we normalize by l. This is correct. But what about `corr = exp(m - new_m)`? If m = -inf initially, and new_m = dot for the first key: - corr = exp(-inf - dot) = 0 - Then o_local = o_local * 0 + p * v = p * v. - l = 0 + p = exp(dot - dot) = 1. This is correct. For the second key, if dot < m: - new_m = m (old max) - corr = exp(m - m) = 1 - p = exp(dot - m) - o_local = o_local + p * v - l = l + p Correct. If dot > m: - new_m = dot - corr = exp(m - dot) - p = exp(dot - dot) = 1 - o_local = o_local * corr + 1 * v - l = l * corr + 1 Correct. However, there's a subtle precision issue. When m updates, we re-weight all previous output by `corr`. This is done elementwise with `o_local[i] = o_local[i] * corr + p * v`. This is the standard online softmax formulation. It should be correct. Now, the interesting point: `l` is only used at the end for normalization. But we could avoid accumulating l by keeping a running sum. Actually, l = sum exp(s_i - m_current). Yes. Is there any issue with `__shfl_xor_sync` and the loop unrolling? It should be fine. One concern: warp divergence due to `if (q >= N) return;`. With N=512 and q_group=0..15, q = q_group*32 + warp. For q_group=15, q = 480..511. All valid since N=512. So no divergence for this problem size. Another concern: memory alignment. Q, K, V are allocated by torch.rand, which uses aligned memory. q_offset + lane*32 + i: lane*32 is 128-byte aligned for lane=0, 128-byte aligned for lane=1? No, lane*32 + i for i=0 gives q_offset + lane*32. Consecutive lanes access q_offset + 0, 32, 64, 96, ... which are 128 bytes apart. This is NOT coalesced for 4-byte loads. OH NO! I made a critical mistake. For coalesced access, thread t should access address `base + t` (or `base + t * stride` where stride is contiguous). Here, lane 0 accesses offset 0, lane 1 accesses offset 32. If each thread loads a single float, lane 0 accesses byte 0, lane 1 accesses byte 128. These are 128 bytes apart. A warp of 32 threads issues 32 separate 128-byte cache line requests. That's terrible! To be coalesced, thread t should access `base + t`. Then we'd load Q[q][lane + i*32] instead of Q[q][lane*32 + i]. Let me think: if I want each thread to load 32 contiguous floats, the layout should be: Thread t loads: `Q[offset + t], Q[offset + t + 32], Q[offset + t + 64], ...` Wait, no. For coalescing, at a given instruction, all threads in a warp should access consecutive memory addresses OR the same address. If I do: ```cpp for (int i = 0; i < 32; ++i) { float val = Q[offset + lane + i * 32]; } ``` At i=0: threads access offset+0, offset+1, offset+2, ..., offset+31. Perfectly coalesced (one 128-byte transaction). At i=1: threads access offset+32, offset+33, ..., offset+63. Perfectly coalesced. ... At i=31: threads access offset+992, ..., offset+1023. Coalesced. This is perfect! And each thread ends up with 32 floats, but they're not contiguous in D; they're strided by 32. Does this matter for the dot product? ```cpp dot += q_local[i] * k_local[i]; ``` where q_local[i] = Q[q][lane + i*32] and k_local[i] = K[k][lane + i*32]. This computes: sum over i of Q[q][lane + i*32] * K[k][lane + i*32]. After warp shuffle reduce, we get sum over all lanes and all i of Q[q][lane + i*32] * K[k][lane + i*32] = sum over all d of Q[q][d] * K[k][d]. Yes, this is correct! The order doesn't matter for the dot product. And for the output accumulation with V: ```cpp o_local[i] += p * v_local[i]; ``` where v_local[i] = V[k][lane + i*32]. This updates O[q][lane + i*32]. At the end, we write O[q][lane + i*32] back. Again, the writes are coalesced because at each i, consecutive threads write consecutive addresses. Perfect! This is the correct memory pattern. Let me rewrite the kernel with this pattern. The performance should be much better. One more optimization: float4 loads and stores. Since each warp load is perfectly coalesced and covers 128 bytes (32 floats), we can use float4 to load 4 floats per instruction. ```cpp float4 q_local[8]; // 8 * 4 = 32 floats for (int i = 0; i < 8; ++i) { q_local[i] = reinterpret_cast<const float4*>(Q + q_offset)[lane + i * 8]; } ``` Wait: `reinterpret_cast<const float4*>(Q + q_offset)[lane + i * 8]` means: - Base pointer: Q + q_offset, cast to float4* - Index: lane + i * 8 - float4 at address = base + (lane + i*8) * sizeof(float4) = base + (lane + i*8) * 16 bytes. - For lane=0..31 and i=0..7: - i=0: offsets 0..31 in float4 units = bytes 0..496 (32 * 16) - i=1: offsets 32..63 = bytes 512..1008 - ... - i=7: offsets 224..255 = bytes 3584..4088 - Total: 256 float4s = 1024 floats = 4096 bytes. - Consecutive lanes access consecutive float4s at each i. Perfectly coalesced. This is equivalent to thread t accessing floats at indices t, t+32, t+64, ..., t+992. Because float4 index j corresponds to floats 4j, 4j+1, 4j+2, 4j+3. If lane loads float4 at index `lane + i*8`, that covers floats: 4*(lane + i*8) through 4*(lane + i*8)+3 = 4*lane + 32*i through 4*lane + 32*i + 3 Wait, thread lane loads a float4 that contains 4 floats, but we don't want that if each thread should handle individual floats for the dot product. We'd need to unpack the float4. Actually, using float4 is fine if we unpack: ```cpp float4 q4 = ...; q_local[i*4+0] = q4.x; q_local[i*4+1] = q4.y; q_local[i*4+2] = q4.z; q_local[i*4+3] = q4.w; ``` But this creates extra move instructions. It could still be beneficial because memory bandwidth improves. Alternatively, we can restructure the dot product to use float4 for FMAs too. But then we'd need to do 4 FMAs per float4, which is fine. Actually, let me just keep it simple with float loads first, and switch to float4 if needed. The compiler might vectorize automatically. Wait, D=1024 and we process 32 chunks of 32 floats per thread. That's 32 iterations. Using float would be 32 memory instructions. Using float4 reduces to 8 memory instructions. This is a significant reduction. Let me use float4: ```cpp #define DIM_PER_THREAD 32 #define FLOAT4_PER_THREAD (DIM_PER_THREAD / 4) float4 q4[FLOAT4_PER_THREAD]; #pragma unroll for (int i = 0; i < FLOAT4_PER_THREAD; ++i) { q4[i] = reinterpret_cast<const float4*>(Q + q_offset)[lane + i * 32]; } float q_local[32]; #pragma unroll for (int i = 0; i < FLOAT4_PER_THREAD; ++i) { q_local[i*4+0] = q4[i].x; q_local[i*4+1] = q4[i].y; q_local[i*4+2] = q4[i].z; q_local[i*4+3] = q4[i].w; } ``` Wait, `lane + i * 32` for float4 indexing: - `lane` ranges from 0 to 31. - `i` ranges from 0 to 7. - Index = lane + i*32. - For i=0: float4 indices 0..31. These are consecutive float4s. - For i=1: float4 indices 32..63. - ... - Total float4s accessed per warp = 32 * 8 = 256 float4s = 1024 floats. - Memory layout: Q is float array. Cast to float4*. Accessing index j gets bytes j*16 through j*16+15. - Thread lane accesses index lane + i*32. Since 32 threads access indices 0..31 for i=0, they access bytes 0..511. Perfectly coalesced. - For i=1, indices 32..63, bytes 512..1023. Coalesced. - And so on. This is correct. And it loads 8 float4s per thread. Now for the inner loop over K and V: ```cpp for (int k = 0; k < N; ++k) { float4 dot4 = make_float4(0,0,0,0); // We need to compute partial dot product using float4 #pragma unroll for (int i = 0; i < FLOAT4_PER_THREAD; ++i) { float4 k4 = reinterpret_cast<const float4*>(K + k_offset)[lane + i * 32]; dot4.x += q_local[i*4+0] * k4.x; dot4.y += q_local[i*4+1] * k4.y; dot4.z += q_local[i*4+2] * k4.z; dot4.w += q_local[i*4+3] * k4.w; } float dot = dot4.x + dot4.y + dot4.z + dot4.w; // warp reduce... ``` Wait, I can keep q in float array, but load K as float4 and unpack. Or load both as float4 and keep in float4 for dot product if we reorganize. Actually, we can avoid unpacking Q and store it as float4 too, but then the dot product needs element-wise multiply. With float4, we can't directly do `q4[i] * k4[i]` because CUDA doesn't have built-in float4 multiply-add. We'd need to do `.x * .x + .y * .y + ...` manually, which is what I wrote. Alternatively, we can keep everything as floats and let the compiler handle vector loads. Modern compilers (nvcc with -O3) often automatically vectorize loads when they see coalesced float loads in a loop with known bounds. Let me try a simpler version first with float loads, and then consider float4 if performance is needed. Actually, for correctness and clarity, I'll use float4 for memory loads but process as floats. Let me also think about whether I should use shared memory. With float4 loads, memory bandwidth is already good. And using no shared memory means more threads per SM. I think the no-shared-memory version is worth trying. The kernel is extremely simple. One more optimization: the K and V loads in the inner loop could be vectorized with float4. Let me write the simplified kernel: Block shape: (32, 32, 1) = 1024 threads. Grid shape: (32, 32, 16) = 16384 blocks. Wait, gridDim.z max is 65535. 16 is fine. But `blockIdx.z` is used for query groups. However, I need to map addresses correctly. The input layout is [B, H, N, D] with default strides. In PyTorch, default stride for [B, H, N, D] is [H*N*D, N*D, D, 1]. So offset = ((b * H + h) * N + n) * D + d. Let me write the full kernel code. Actually, I should also consider using `__ldg` for read-only data to help caching. And for the output write, we can just use regular store. Let me also think about numerical precision. The reference uses fp32. My kernel uses fp32 FMAs and `expf`. `torch.nn.functional.scaled_dot_product_attention` might use a different softmax algorithm or different precision for intermediate accumulations. But online softmax in fp32 should be accurate enough. One more thought: `expf` can be expensive. With 512 keys and 32 warps, each warp does 512 `expf` calls. `expf` takes dozens of cycles. This might be a bottleneck. We have 512 keys * 16384 warps = 8.4M `expf` calls. Each `expf` might take ~20-30 cycles. At 2GHz, that's about 10-15ms across all SMs. Not too bad, but significant. Can we use fast `__expf`? The `--use_fast_math` flag should enable it. But it might sacrifice precision. Let's try it and see if correctness passes. Actually, the reference uses standard `exp`. With fast math, the error might be too large. But fp32 tolerance is 1e-4. `__expf` vs `expf` difference is usually around 1e-5 or so. It might be fine. But let's first try with `expf` and see if we pass correctness. If yes, we can try `__expf` for speed. Wait, the first kernel already passed correctness. This new kernel uses a different algorithm but should give the same result. Wait, there's a subtle difference. In standard attention: - S = Q @ K^T - P = softmax(S / sqrt(D)) - O = P @ V In my online softmax formulation: - I compute dot products and scale. - Then I apply online softmax to get normalized weights. - Then O = weighted sum of V. This should give exactly the same result as the standard formulation (up to floating point associativity differences). But does PyTorch's `scaled_dot_product_attention` use exactly the same online softmax? No, but with fp32 tolerance of 1e-4, minor differences from associativity should be acceptable. Let me think about register pressure. With `q_local[32]`, `o_local[32]`, and the loop variables, it's about 68 registers. But the compiler might keep `q_local` in registers. With unrolled loops, it definitely will. However, 68 registers + some stratch = ~80. For 1024 threads, that's 80K > 64K. The compiler might have to spill. But we can avoid `q_local[32]` by loading Q on the fly: ```cpp float q_val = Q[q_offset + lane + i * 32]; dot += q_val * k_val; ``` This way, we only need a scalar `q_val` per iteration, not an array. The compiler might not keep it in a register between iterations, or it might reload from L1 cache. Actually, with warp-level code and 32 iterations, the compiler will likely keep the Q values in registers or at least keep around. But there's no guarantee. If Q gets spilled to local memory, that's bad. To help the compiler, we can load Q into a smaller array and explicitly unroll: ```cpp float q00, q01, q02, ..., q31; // load // use in loop // or use float4 float4 q4[8]; #pragma unroll for (int i = 0; i < 8; ++i) q4[i] = loadQ(...); // in K loop: float4 k4[8]; #pragma unroll for (int i = 0; i < 8; ++i) { k4[i] = loadK(...); dot += q4[i].x * k4[i].x; dot += q4[i].y * k4[i].y; dot += q4[i].z * k4[i].z; dot += q4[i].w * k4[i].w; } ``` With float4, the compiler has fewer variables to track. But the register usage for 8 float4s = 32 registers. For the V accumulation, we need `o_local` as float or float4: ```cpp float4 o4[8] = {0}; #pragma unroll for (int i = 0; i < 8; ++i) { float4 v4 = loadV(...); o4[i] = o4[i] * corr + v4 * p; // but float4 doesn't support * for scalars directly // we need: o4[i].x = o4[i].x * corr + v4.x * p; } ``` Hmm, float4 arithmetic with scalars isn't built-in. I'd have to write helper functions or manual arithmetic. Maybe I should stick with plain float arrays. Let the compiler optimize. Actually, another option: use `__half2` or vector types? No, we're doing fp32. I think the simplest correct approach is: ```cpp float q_reg[32]; for (int i = 0; i < 32; ++i) q_reg[i] = ...; float o_reg[32] = {0}; for (int k = 0; k < N; ++k) { float dot = 0; for (int i = 0; i < 32; ++i) { dot += q_reg[i] * K[k_offset + lane + i * 32]; } // reduce // softmax for (int i = 0; i < 32; ++i) { o_reg[i] = o_reg[i] * corr + p * V[k_offset + lane + i * 32]; } } ``` NVCC will likely keep `q_reg` in registers. With 32 registers and `o_reg` with 32 registers, plus loop temps, that's about 70 registers. It might spill a little, but with unrolled loops it should be OK. Let me write this kernel and compile it. One more thing: `blockIdx.z` is the third dimension of the grid. `gridDim.z` must be <= 65535. 16 is fine. Wait, can I use `blockIdx.z` in `load_inline`? Yes, standard CUDA. Let me write the complete Python code with the kernel. Actually, I realize something. With 1024 threads per block and only `threadIdx.x` and `threadIdx.y`, blockDim = (32, 32, 1). A warp consists of threads with the same `threadIdx.y` and consecutive `threadIdx.x` values. The hardware groups threads by `threadIdx.x` first. So threads (0..31, y) form a warp. This is correct. But note: within a warp, consecutive threads are consecutive in `threadIdx.x`. My memory access uses `lane = threadIdx.x`. So lane 0..31 are consecutive threads in a warp. Memory accesses at the same `i` value, address = base + lane + i*32, are perfectly coalesced. Good. OK, let me write the kernel. Due to the complexity of float4 with scalar multiply, I'll use plain float and let nvcc handle vector loads. But I'll hint vectorization through pointer casting or pragma unroll. Actually, I can try float4 to save instructions. Let me write helper macros: ```cpp inline __device__ float4 operator*(const float4& a, float b) { return make_float4(a.x * b, a.y * b, a.z * b, a.w * b); } inline __device__ float4 operator+(const float4& a, const float4& b) { return make_float4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w); } inline __device__ void operator+=(float4& a, const float4& b) { a.x += b.x; a.y += b.y; a.z += b.z; a.w += b.w; } inline __device__ float4 operator*(float b, const float4& a) { return a * b; } ``` However, operators on float4 might already be defined in some CUDA versions. Actually, no, CUDA provides float4 struct but not arithmetic operators by default (except in some headers for specific purposes). I can define them as above. But `corr` is a scalar float. So `o4[i] = o4[i] * corr + v4 * p;` works with these operators. For dot product with float4: ```cpp float dot = 0; #pragma unroll for (int i = 0; i < 8; ++i) { float4 q4 = q_local[i]; float4 k4 = reinterpret_cast<const float4*>(K + k_offset)[lane + i * 32]; dot += q4.x * k4.x + q4.y * k4.y + q4.z * k4.z + q4.w * k4.w; } ``` This is cleaner and should reduce instruction count. For the output with float4: ```cpp float4 o_local[8] = {}; // ... #pragma unroll for (int i = 0; i < 8; ++i) { float4 v4 = reinterpret_cast<const float4*>(V + k_offset)[lane + i * 32]; float4 p4 = make_float4(p, p, p, p); float4 corr4 = make_float4(corr, corr, corr, corr); o_local[i] = o_local[i] * corr4 + v4 * p4; } ``` Actually `make_float4(p,p,p,p)` for each iteration creates extra instructions. I can just use a helper function: ```cpp __device__ float4 fma4(float4 o, float corr, float4 v, float p) { return make_float4(o.x * corr + v.x * p, ...); } ``` Or even simpler, just use 4 scalar FMAs: ```cpp o_local[i].x = o_local[i].x * corr + v4.x * p; o_local[i].y = o_local[i].y * corr + v4.y * p; ... ``` The compiler will optimize this. With 8 iterations and 4 statements each, plus unrolling, it should generate good code. Actually, wait. With `q_local` as float4, I can't easily write `o_local[i].x = ...` if o_local is float4 array. It works, but it's verbose. Let me try the float4 approach with manual element access, and see if it compiles. Now, let me think about one potential issue: when `lane + i * 32` reaches 256 (for i=8), that's out of bounds for a single float4 row of 1024/4 = 256 elements. Wait: - float4 indices for one row of 1024 floats: 0 to 255. - lane ranges 0..31. - i ranges 0..7. - index = lane + i * 32. - i=0: 0..31 - i=1: 32..63 - ... - i=7: 224..255 - Max index = 255. Correct! And 256 * 4 = 1024 floats. Good, no out-of-bounds. Now, what about the alignment requirement for float4? The base pointer Q + q_offset must be 16-byte aligned. Since torch tensors are aligned to at least 256 bytes typically, and q_offset is a multiple of D*4 = 4096 bytes (since q ranges over integers), the alignment is satisfied. Wait, for K + k_offset, k_offset is also a multiple of D*4 = 4096. So it's also aligned. OK let me write the code now. I'll use: - Block: (32, 32, 1) = 1024 threads - Grid: (32, 32, 16) = 16384 blocks - q = blockIdx.z * 32 + threadIdx.y Wait, grid dimensions: (32, 32, 16) means: - blockIdx.x: 0..31 (batch) - blockIdx.y: 0..31 (head) - blockIdx.z: 0..15 (query group) But there are only 32 batches and 32 heads. The total output is [B, H, N, D]. For block (b, h, g), it processes queries g*32 .. g*32+31 for batch b and head h. Now, I should double check: the grid has 32*32*16 = 16384 blocks. Each block has 32 warps, each handling 1 query. Total queries handled = 16384 * 32 = 524288. Total queries in tensor = 32*32*512 = 524288. Perfect match. Wait, I need to make sure the kernel handles N properly. N=512, groups=16. group=15 handles queries 480..511. All valid. Edge case: what if N is not divisible by 32? We have `if (q >= N) return;`. In this specific problem N=512, so it's exactly divisible and the check is never true. But it's good practice. Let me write the complete kernel code now. Actually, one issue with using `blockIdx.z`: the maximum is quite large, but only 16. Fine. But I realize something: in my grid=(32,32,16), the launch configuration would be: ```cpp dim3 grid(32, 32, 16); dim3 block(32, 32, 1); ``` Wait, block has 32*32 = 1024 threads. Grid has 32*32*16 = 16384 blocks. Total threads = 16.7M. The kernel code: ```cpp __global__ void fused_attention( const float* __restrict__ Q, const float* __restrict__ K, const float* __restrict__ V, float* __restrict__ O, int B, int H, int N, int D) { int batch = blockIdx.x; int head = blockIdx.y; int q_group = blockIdx.z; int lane = threadIdx.x; int warp = threadIdx.y; int q = q_group * 32 + warp; if (q >= N) return; int stride = H * N * D; int head_stride = N * D; int q_offset = batch * stride + head * head_stride + q * D; float scale = 1.0f / sqrtf((float)D); // Load Q float4 q4[8]; #pragma unroll for (int i = 0; i < 8; ++i) { q4[i] = reinterpret_cast<const float4*>(Q + q_offset)[lane + i * 32]; } float o_reg[32]; #pragma unroll for (int i = 0; i < 32; ++i) o_reg[i] = 0.0f; float m = -1e20f; float l = 0.0f; for (int k = 0; k < N; ++k) { int k_offset = batch * stride + head * head_stride + k * D; // Compute dot product float4 k4[8]; float dot = 0.0f; #pragma unroll for (int i = 0; i < 8; ++i) { k4[i] = reinterpret_cast<const float4*>(K + k_offset)[lane + i * 32]; dot += q4[i].x * k4[i].x; dot += q4[i].y * k4[i].y; dot += q4[i].z * k4[i].z; dot += q4[i].w * k4[i].w; } // Warp reduce to get full dot product #pragma unroll for (int offset = 16; offset > 0; offset /= 2) { dot += __shfl_xor_sync(0xFFFFFFFF, dot, offset); } dot *= scale; // Online softmax float new_m = fmaxf(m, dot); float corr = expf(m - new_m); float p = expf(dot - new_m); // Load V and update output #pragma unroll for (int i = 0; i < 8; ++i) { float4 v4 = reinterpret_cast<const float4*>(V + k_offset)[lane + i * 32]; o_reg[i*4+0] = o_reg[i*4+0] * corr + p * v4.x; o_reg[i*4+1] = o_reg[i*4+1] * corr + p * v4.y; o_reg[i*4+2] = o_reg[i*4+2] * corr + p * v4.z; o_reg[i*4+3] = o_reg[i*4+3] * corr + p * v4.w; } l = l * corr + p; m = new_m; } float norm = 1.0f / l; #pragma unroll for (int i = 0; i < 8; ++i) { float4 o4; o4.x = o_reg[i*4+0] * norm; o4.y = o_reg[i*4+1] * norm; o4.z = o_reg[i*4+2] * norm; o4.w = o_reg[i*4+3] * norm; reinterpret_cast<float4*>(O + q_offset)[lane + i * 32] = o4; } } ``` Wait, I need to check: does `q4[i]` register array consume 8*4 = 32 registers? Yes. `o_reg[32]` = 32 registers. `k4[8]` = 32 registers (temporarily). But with loop unrolling, the compiler might optimize the lifetimes. Actually, `k4` is only needed inside the K loop. But the loop has N=512 iterations, so K is loaded repeatedly. The compiler might keep `q4` in registers throughout, and `o_reg` too. `k4` might be allocated as temporaries inside the loop body. Total register pressure: q4(32) + o_reg(32) + m(1) + l(1) + corr(1) + p(1) + dot(1) + k4(32) + temps = ~100 registers. That's too much. Hmm, 100 registers per thread * 1024 threads = 102400 registers. But an SM only has 65536 registers. So with 1024 threads, we'd be limited to 64 registers per thread. The compiler will definitely spill. Can we optimize registers? Alternative: don't load all Q into registers. Load Q on the fly: ```cpp // No q4 array for (int k = 0; k < N; ++k) { float dot = 0.0f; #pragma unroll for (int i = 0; i < 8; ++i) { float4 q4 = reinterpret_cast<const float4*>(Q + q_offset)[lane + i * 32]; float4 k4 = reinterpret_cast<const float4*>(K + k_offset)[lane + i * 32]; dot += q4.x * k4.x + q4.y * k4.y + q4.z * k4.z + q4.w * k4.w; } // ... } ``` Now q4 is a temporary float4, not an array. It doesn't consume persistent registers. Same for k4 in this loop. For the output accumulation, we need to load V too: ```cpp for (int k = 0; k < N; ++k) { // compute dot... // update o_reg #pragma unroll for (int i = 0; i < 8; ++i) { float4 v4 = reinterpret_cast<const float4*>(V + k_offset)[lane + i * 32]; o_reg[i*4+0] = ...; } } ``` q4 is loaded twice (once in dot product, but wait - if we don't cache Q, we load it 512 times! That's bad. Q for this query has 1024 floats = 4KB. Loading it 512 times = 2MB per query. Total Q traffic = 524k * 2MB = 1TB. That's terrible. So we DO need to keep Q in registers or shared memory. What if we load Q into shared memory? Each block has 32 queries. Q for 32 queries = 32*1024*4 = 128KB. Shared memory can hold this, but it limits occupancy. Alternatively, each warp loads Q into SHARED MEMORY (not registers) for its one query: - smem_q[32][32] (1024 floats for one query). - But shared memory is per block, so all 32 warps share it. - Wait, we'd need 32 query worths = 128KB. That's too much. What if a block only has 8 warps = 256 threads, processing 8 queries? - Q smem = 8*1024*4 = 32KB. - Can have 4 blocks per SM = 128KB. But 4*256 = 1024 threads per SM. OK. - But grid now has more blocks. Actually, with q4 array of 8 float4s = 32 floats = 32 registers, it's not terrible. The issue is combining q4, o_reg, and k4. What if o_reg is stored in shared memory too? - O accumulator for 32 queries: 32*1024*4 = 128KB. Too much for shared. What if we process outputs in chunks? - Compute all scores first (using cached Q), then accumulate V in chunks. - But we said storing scores for 512 keys is 512 floats per query = 2KB per query. - For 32 warps = 64KB. This fits in shared memory! - Actually no, 32 queries * 512 scores = 16K floats = 64KB. That's reasonable. - We could store scores in shared memory, then process V in chunks and update global O directly. Algorithm: 1. Load Q into registers (8 float4s = 32 regs). 2. For each key, compute score and store in shared memory: score[q][k]. 3. Wait, but we need 32*512 = 16K floats = 64KB. BlockDim = 32 warps. If each warp computes 16 scores (one per key for some keys)... This is getting complicated again. Let me try a simpler register approach with lower block size. What if block = 8 warps = 256 threads, processing 8 queries. Grid: (32, 32, 64) = 65536 blocks. (8 queries per block, 512/8 = 64 groups). Each warp handles 1 query. 8 warps = 8 queries. Total blocks = 32*32*64 = 65536. Total threads = 65536 * 256 = 16.7M. Same total threads. With 256 threads per block: - Register limit = 65536 / 256 = 256 registers per thread! Much more room. - q4[8] = 32 registers. - o_reg[32] = 32 registers. - k4[8] temporary = 32 registers (compiler manages lifetime). - dot, m, l, corr, p = 5. - Total peak ~100 registers. Under 256 limit. Plenty of room. - Also, smem is not used, so occupancy isn't limited by shared memory. - Concurrent blocks per SM: 65536 / 128 = 512 threads. 2048 / 512 = 4 blocks per SM. - Actually 2048 threads max / 256 threads per block = 8 blocks per SM. - But register limit: 8 * 256 * 100 = 204800 > 65536. So max blocks = 65536 / (256 * 100) = 2.5, i.e. 2 blocks. - 2 blocks * 256 = 512 threads per SM. That's fine. - Total concurrent blocks = 132 * 2 = 264. - Waves = 65536 / 264 = 248 waves? No, total threads = 16.7M. - Concurrent threads = 132 SMs * 512 threads = 67.5k threads. - Waves = 16.7M / 67.5k = 248 waves. Wait, that's a lot of waves. Each wave takes some time. Block time: each block processes 8 queries * 512 keys. Per warp: 512 keys * 1024 FMAs = 524k FMAs. At 50% peak, maybe ~15 microseconds per warp. 8 warps execute concurrently within a block (they're in the same SM). So block time ~15 microseconds. 248 waves * 15 microseconds = 3.7ms. That's not bad! Actually, with 512 threads per SM and 2 blocks, the 2 blocks from the same SM might execute interleaved. But the wave model still roughly applies. For 1024 threads per block: - Concurrent threads per SM: 65536 / (1024 * 100) = 0.64 blocks due to registers. So 1 block = 1024 threads. - Total concurrent threads = 132 * 1024 = 135k. - Waves = 16.7M / 135k = 124 waves. - Block time: 32 warps * 512 keys. But they're all on the same SM. Latency hiding between warps. Time might still be ~15us (or a bit more due to less parallelism per block? No, 32 warps is plenty). - Total = 124 * 15us = 1.86ms. 1024-thread blocks seem better for overall runtime because fewer waves. But register pressure is the issue. Can we reduce registers? With q4[8] and o_reg[32], that's 64 persistent registers. If we can reduce o_reg... What if instead of 32 floats for o_reg, we process V in chunks and keep only 4 or 8 outputs at a time? Wait, no. For online softmax, we need ALL output dimensions to be updated at once because the `corr` factor applies to all dims. We could process chunks: ```cpp for (int d_chunk = 0; d_chunk < D; d_chunk += 32) { float o_local[32] = {0}; m = -inf; l = 0; for (int k = 0; k < N; ++k) { float dot = 0; // compute full dot product using all Q // load V[k][d_chunk:d_chunk+32] // update o_local } // write o_local / l } ``` This requires recomputing the dot product for each d_chunk! That's 32x more compute. Not worth it. Alternative: compute and store scores in shared memory, then do V accumulation in a separate pass. - Load Q. - For each key, compute score and store in smem_score[warp][k] (512 floats per warp = 2KB, times 32 warps = 64KB total). - Then find max and compute softmax weights in a separate step. - Then for each key, broadcast weight, load V, accumulate to o_reg[32]. - o_reg[32] still needs 32 registers. But q4 is freed after score computation. This doesn't save registers. Another idea: use half the registers for q4 by loading Q in two passes. For D=1024, process i=0..3 first, then i=4..7. But then we load Q twice. Q is small (4KB per query), loading twice is 8KB. From L1 cache, it's fast and doesn't increase HBM traffic. And it halves Q register usage! ```cpp for (int k = 0; k < N; ++k) { float dot = 0.0f; // First half of Q float4 q4[4]; #pragma unroll for (int i = 0; i < 4; ++i) { q4[i] = reinterpret_cast<const float4*>(Q + q_offset)[lane + i * 32]; } #pragma unroll for (int i = 0; i < 4; ++i) { float4 k4 = reinterpret_cast<const float4*>(K + k_offset)[lane + i * 32]; dot += q4[i].x * k4.x + q4[i].y * k4.y + q4[i].z * k4.z + q4[i].w * k4.w; } // Second half of Q #pragma unroll for (int i = 4; i < 8; ++i) { q4[i-4] = reinterpret_cast<const float4*>(Q + q_offset)[lane + i * 32]; } #pragma unroll for (int i = 4; i < 8; ++i) { float4 k4 = reinterpret_cast<const float4*>(K + k_offset)[lane + i * 32]; dot += q4[i-4].x * k4.x + ...; } ``` This uses only 4 float4s for Q at a time = 16 floats persistent. But `q4` is local to the loop, the compiler might keep it in registers or reload from L1. If it reloads from L1, the extra Q loads are fast (L1 cache hit). Actually, since we're loading the same Q over and over across 512 iterations, Q will definitely be in L1 cache. Re-loading is cheap. So we don't need to keep Q in registers at all! ```cpp for (int k = 0; k < N; ++k) { float dot = 0.0f; #pragma unroll for (int i = 0; i < 8; ++i) { float4 q4 = reinterpret_cast<const float4*>(Q + q_offset)[lane + i * 32]; float4 k4 = reinterpret_cast<const float4*>(K + k_offset)[lane + i * 32]; dot += q4.x * k4.x + q4.y * k4.y + q4.z * k4.z + q4.w * k4.w; } // ... } ``` Now `q4` is just one float4 per iteration. The compiler uses 4 registers for it, reusing across iterations. `k4` similarly. After the dot product, we compute softmax update, then load V and update `o_reg[32]`. `o_reg` still needs 32 registers. But 32 registers + a few temporaries is very manageable! Total ~50 registers. Wait, this is great! Q is loaded fresh each key iteration. But since all warps in the block repeatedly access the SAME Q addresses (different queries, but each warp's Q is fixed), and L1 cache on H100 is per-SM, all threads in the block will hit in L1 for Q loads after the first miss. Actually, each warp loads Q for its own query. Query data is 4KB. L1 cache per SM is 128KB or 256KB depending on configuration. 32 warps * 4KB = 128KB. This might just barely fit in L1, or some of it might spill to L2. But even L2 bandwidth is high. And the dot product loads Q and K. K is 4KB per key? No, one key is 4096 bytes. For 512 keys = 2MB. K doesn't stay in L1. So per key iteration, we load: - Q: 4KB (for one query) from L1 or L2 - K: 4KB (for one key) from L1 or L2 - V: 4KB for accumulation Total per block per key = 32 queries * 4KB Q + 1 key * 4KB K + 1 key * 4KB V? Actually, all 32 warps in a block are loading Q for different queries simultaneously. So total concurrent load = 32 * 4KB = 128KB of Q. Plus 32 * 4KB K? No, all warps load the same K[k][0:1023]. So K is broadcast across all warps. The hardware fetches one 4KB cache line set for K and broadcasts it. Wait, is K broadcast across warps? L1 cache serves requests from all warps in the SM. If 32 warps all request K[k][0:1023], and the data is in L1, it can be broadcast. But L1 might not have enough read ports to serve 32 warps simultaneously. In practice, warps have some serialization in accessing L1. But with 32 warps on an SM, and each warp issuing loads every few cycles, it should be OK. The L1 cache has high bandwidth. The key insight is that we only need 32 persistent registers for `o_reg`, plus a few temporaries. This allows us to use 1024 threads per block comfortably without register spilling. Let me build this kernel: ```cpp __global__ void fused_attention_kernel( const float* __restrict__ Q, const float* __restrict__ K, const float* __restrict__ V, float* __restrict__ O, int B, int H, int N, int D) { int batch = blockIdx.x; int head = blockIdx.y; int q_group = blockIdx.z; int lane = threadIdx.x; int warp = threadIdx.y; int q = q_group * 32 + warp; if (q >= N) return; int stride = H * N * D; int head_stride = N * D; int q_offset = batch * stride + head * head_stride + q * D; float scale = 1.0f / sqrtf((float)D); float o_reg[32]; #pragma unroll for (int i = 0; i < 32; ++i) o_reg[i] = 0.0f; float m = -1e20f; float l = 0.0f; for (int k = 0; k < N; ++k) { int k_offset = batch * stride + head * head_stride + k * D; float dot = 0.0f; #pragma unroll for (int i = 0; i < 8; ++i) { float4 q4 = reinterpret_cast<const float4*>(Q + q_offset)[lane + i * 32]; float4 k4 = reinterpret_cast<const float4*>(K + k_offset)[lane + i * 32]; dot += q4.x * k4.x; dot += q4.y * k4.y; dot += q4.z * k4.z; dot += q4.w * k4.w; } #pragma unroll for (int offset = 16; offset > 0; offset /= 2) { dot += __shfl_xor_sync(0xFFFFFFFF, dot, offset); } dot *= scale; float new_m = fmaxf(m, dot); float corr = expf(m - new_m); float p = expf(dot - new_m); #pragma unroll for (int i = 0; i < 8; ++i) { float4 v4 = reinterpret_cast<const float4*>(V + k_offset)[lane + i * 32]; o_reg[i*4+0] = o_reg[i*4+0] * corr + p * v4.x; o_reg[i*4+1] = o_reg[i*4+1] * corr + p * v4.y; o_reg[i*4+2] = o_reg[i*4+2] * corr + p * v4.z; o_reg[i*4+3] = o_reg[i*4+3] * corr + p * v4.w; } l = l * corr + p; m = new_m; } float inv_l = 1.0f / l; #pragma unroll for (int i = 0; i < 8; ++i) { float4 o4; o4.x = o_reg[i*4+0] * inv_l; o4.y = o_reg[i*4+1] * inv_l; o4.z = o_reg[i*4+2] * inv_l; o4.w = o_reg[i*4+3] * inv_l; reinterpret_cast<float4*>(O + q_offset)[lane + i * 32] = o4; } } ``` Wait, I realize there's a crucial issue: Q is loaded inside the K loop. The loads `reinterpret_cast<const float4*>(Q + q_offset)[lane + i * 32]` will be executed for every key iteration (512 times). This is 512 * 4KB = 2MB per query. But Q is only 4KB. So we're reading Q 512x more than necessary. However, with L1 caching, subsequent reads of Q will be L1 hits. The extra reads don't go to HBM, but they consume L1 bandwidth and instructions. Can the compiler optimize this and keep Q in registers? It can't keep all of it, but with the loop unrolling and 8 float4 loads, it could keep each float4 in registers for the duration of the inner loop body, then reload in the next key iteration. The compiler will NOT cache Q across outer loop iterations. But since Q is invariant across the K loop, shouldn't the compiler optimize this? Actually, the compiler is not allowed to move loads out of loops in general if there could be aliasing. But with `__restrict__` and const, it might optimize. However, in my experience, nvcc doesn't always move global loads out of loops. To help the compiler, we could add `__ldg` (load global with cache) to hint that the data is readonly and cacheable. Or we could manually move Q loads before the loop. But to manually move Q loads before the loop, we need at least 32 registers. That causes the pressure issue again. Alternatively, use shared memory to explicitly cache Q. But then we're back to shared memory limits. Wait, what about L1 automatic caching? On modern GPUs, global loads go through L1. When a warp loads Q in iteration k=0, the data goes into L1. In iteration k=1, the same addresses are loaded from L1. The latency is just a few cycles. The instruction overhead is still there, but it's fast. With 512 iterations and 8 float4 loads per iteration = 4096 load instructions per warp. At 1 load per cycle per SM, this is negligible compared to the compute. Actually, the real issue is whether the loads are redundant with respect to the instruction pipeline. Issuing redundant load instructions consumes issue slots. A warp executing the K loop does: - 8 loads of Q - 8 loads of K - 32 FMAs - 8 loads of V - 8*4 = 32 FMAs for V update - shuffle reduce - 2 expf - Total: ~16 loads + ~64 FMAs + 2 expf + some ALU. The loads can be overlapped with FMAs due to latency hiding. Instruction issue is probably not the bottleneck. Actually, let me calculate more carefully. A warp issues about 1-2 instructions per cycle to the execution units. For 512 iterations, if each iteration takes ~40 cycles, total is ~20000 cycles = 10 microseconds. For 32 warps in a block, blocks take ~10 microseconds. Total waves = 124. Total time = 1.24ms. This is really fast! Much faster than I expected. And much faster than cuBLAS approach. But I realize my estimate might be too optimistic. Let me be more conservative. Say each key iteration takes 100 cycles due to expf latency and memory. 512 * 100 = 51200 cycles = 25.6 microseconds per warp. 32 warps with good latency hiding = 25.6 microseconds per block. 124 waves = 3.2ms. If the reference math backend takes ~15ms, then 3.2ms is much faster. Even if it's 5-6ms, we might still win. One concern: `expf` is expensive. We call it twice per key per warp. For N=512, that's 1024 expf calls per warp. Each expf takes ~20-30 cycles. Two expf = 40-60 cycles per iteration. At 2GHz, that's 20-30 nanoseconds per iteration. 512 * 25ns = 12.8 microseconds. For 32 warps executing concurrently with instruction-level parallelism, blocks might take ~15 microseconds. 124 waves = 1.86ms. Actually, if `expf` is so dominant, the whole kernel might be faster than I thought. But wait - is the expf throughput really the bottleneck? H100 has special function units (SFU) for transcendental ops. Each SM might have limited SFU throughput. SFUs on Ampere/Hopper: each SM has 4 SFUs (I think). Each SFU can compute 1 `expf` per clock. So an SM can do 4 expf per cycle. With 32 warps on an SM (1024 threads), and each warp needing 2 expf per iteration... well, warps issue expf only 512 times over the entire kernel. But the latency means the SM is not fully utilized on SFUs. If each warp issues 512 * 2 = 1024 expf instructions, and each takes say 20 cycles to retire. A block has 32 warps, so 32768 expf instructions total. If an SM can retire 4 expf per cycle, it takes 8192 cycles = 4.1 microseconds for expfs on one SM. That's not too bad. But we're running multiple SMs. With 132 SMs, total expf capacity = 132 * 4 = 528 per cycle. Total expfs = 16384 blocks * 32 warps * 1024 expfs = 536M expfs. 536M / 528 = 1M cycles = 0.5ms. So expf is not a major bottleneck at the chip level. OK, I think the kernel is promising. Let me try compiling it and checking correctness. But wait, I just thought of another issue. The loads `reinterpret_cast<const float4*>(Q + q_offset)[lane + i * 32]` where q_offset = ... + q*D. If q_offset is passed in correctly, and Q is a float pointer, then Q + q_offset is float-aligned. But for float4, we need 16-byte alignment. q_offset might not be 16-byte aligned if q*D is not divisible by 4. But D=1024, and q_offset = batch*stride + head*head_stride + q*D. Since D=1024 is divisible by 4, q*D is divisible by 4. stride and head_stride are also multiples of D and thus divisible by 4. So the alignment is actually 4-byte for float, which means for float4 we need 16-byte alignment. Since everything is divisible by 1024, and 1024 / 4 = 256, so offsets are multiples of 256 float4s, which is definitely 16-byte aligned. Good. One more edge case: when total D is not divisible by 128 (bytes per warp load), but D=1024 is divisible by 128. Good. Now let me think about whether I should use float4 for loads. On Hopper, 128-byte loads are very efficient. Using float4 issues 128-bit loads, and 8 float4 loads cover the full 1024 bytes. Alternatively, I could use `__ldcs` or `__ldg` for uncached/cached loads. For Q which is reused across iterations, L1 caching is ideal. The default `ld.global` goes through L1 on Hopper (unlike Ampere where it bypasses L1 by default, but `ld.global.L1::128` is available). Actually, on Ampere and later, `ld.global` can use L1. The compiler usually does the right thing. I'll use `__ldg` (load through read-only cache) for Q, K, V. `__ldg` is a hint for frequent reuse. For Q: since it's invarient, `__ldg` is perfect. For K and V: they are sequential and each element is read by all warps in a block (though with 1024 threads, it's 32 warps). `__ldg` might help because the read-only cache can broadcast. ```cpp float4 q4 = __ldg(reinterpret_cast<const float4*>(Q + q_offset) + lane + i * 32); ``` Actually, I need to cast properly. `__ldg` takes a pointer. So: ```cpp const float4* q_ptr = reinterpret_cast<const float4*>(Q + q_offset); float4 q4 = __ldg(q_ptr + lane + i * 32); ``` Let me check if `__ldg` works on `float4`. Yes, it works on any type. But wait, when using `__ldg`, the address must be 16-byte aligned. Our addresses are. Good. One more concern: `__ldg` might bypass L1 on some architectures and use the texture cache / read-only cache. For data that's reused within a warp/thread, this might not be the best. But for data read by all warps, the L1 cache and texture cache both handle it. Actually, for K which is read by all 32 warps in a block, any cache (L1 or read-only) is good. Q is read repeatedly by one warp. L1 is ideal because Q stays hot in L1. I'll just use default loads for now. If I want, I can use `__ldg` for all. The compiler will generate `LDG` instructions. Now let me write the python module and test. Wait, one more thing. I should check if the problem's `torch.nn.functional.scaled_dot_product_attention` might have different behavior for edge cases. The input tensors are `torch.rand(batch_size, num_heads, sequence_length, embedding_dimension)`. So shape is [32, 32, 512, 1024]. In my kernel, I use the standard formula: scale = 1/sqrt(D). PyTorch's SDPA does the same: `softmax(Q @ K^T / sqrt(D)) @ V`. Actually, I need to confirm: PyTorch fp32 SDPA uses `1/sqrt(E)` where E is the embedding dimension per head. Here D=1024. My kernel does the same. Let me now implement and test this kernel. I'll use the compile_kernel and run_correctness functions. One more optimization thought: can I handle the dot product and V update more efficiently? For the dot product, I do 8 float4 loads and 4 multiply-adds per float4. That's 32 operations. I can combine into: ```cpp dot += q4.x * k4.x + q4.y * k4.y + q4.z * k4.z + q4.w * k4.w; ``` This is a multiply-add chain. The compiler will generate 4 FMADs (actually FMAs). For V update: ```cpp o_reg[i*4+0] = fmaf(p, v4.x, o_reg[i*4+0] * corr); ``` Wait, `o * corr + p * v` = `fmaf(p, v, o * corr)`. The compiler might use FMA if it can. But first it computes `o * corr`, then `fma(p, v, o*corr)`. This uses 2 multiplies and 1 add per element. An FMA instruction does 1 multiply-add. So ideally we'd want `o * corr + p * v` to be a single FMA if `corr` is 1, but it's not. Actually, `fma(p, v, o * corr)` still requires a separate multiply for `o * corr`. But we can compute `o_reg[i] *= corr;` first, then `o_reg[i] += p * v;` using FMA for the second step. This is actually 2 instructions vs a combined instruction that might be the same. Let's just write it clearly and let nvcc optimize. Actually, for corr and p being scalars, if we do: ```cpp #pragma unroll for (int i = 0; i < 32; ++i) { o_reg[i] *= corr; } #pragma unroll for (int i = 0; i < 32; ++i) { o_reg[i] += p * v_local[i]; } ``` This might allow better instruction scheduling because the multiply and the FMA are separated. But it requires 64 instructions. If we fuse them, it's still 64 instructions. However, with floats, the second form uses FMA: `o += p * v` (after separate multiply). But `o *= corr` is just MUL. The first form: `o = o * corr + p * v`. If `corr` is not 1, this is not a pure FMA. It requires two multiplications. However, on modern GPUs, there's a `DFMA` or similar? No, standard FMA is `a * b + c`. Here we have `a * b + c * d`, which requires 2 multiplies and 1 add. So there's no way to reduce it below 2 multiplies + 1 add per output element. The separated form might be better for instruction pipelining. Alright, let me just write the kernel. I'll keep the code as above: one statement per element. One more thing: should I use `rpowf` or `fdiv_rn` for `1.0f / sqrtf(D)`? `sqrtf` and division are done once per kernel. Negligible. OK, let me code this up. Actually, I realize there might be an simpler issue: the CUDA warp shuffle with `__shfl_xor_sync(0xFFFFFFFF, dot, offset)` works across the warp. But my threads are arranged as `threadIdx.x=lane, threadIdx.y=warp`. Within a warp, `threadIdx.y` is constant and `threadIdx.x` varies. The warp consists of threads with consecutive `threadIdx.x`. So the shuffle operates correctly within each warp. Wait, is that right? In CUDA, threads in a block are linearized as: `linear_tid = threadIdx.x + threadIdx.y * blockDim.x + threadIdx.z * blockDim.x * blockDim.y` Warps consist of 32 consecutive threads in this linearization. So threads (0..31, constant y, constant z) are in the same warp. Yes! But the warp shuffle operates on threads in the same warp. And within a warp, the lane ID is `linear_tid % 32 = threadIdx.x % 32`. Since `threadIdx.x` ranges 0..31, lane ID is exactly `threadIdx.x`. So `__shfl_xor_sync(0xFFFFFFFF, dot, offset)` correctly exchanges between threads with different `threadIdx.x` values within the same warp. And since `dot` is computed identically by all threads in the warp (they load different q and k but end up with different partial sums), the shuffle adds them up correctly. Perfect. But I need to make sure that threads in different warps don't participate in each other's shuffles. The mask `0xFFFFFFFF` means all threads in the warp participate. And since shuffle is warp-local, threads in warp 0 don't see warp 1's values. Good. Now, let me also think about whether blockDim=(32,32,1) could cause issues with the number of warps. blockDim.x=32, blockDim.y=32, blockDim.z=1. Total threads=1024. Maximum allowed. This is fine. I'll write the code and compile. Wait, one more concern: `reinterpret_cast<const float4*>(Q + q_offset)`. Q is `const float*`, so Q + q_offset is `const float*`. Casting to `const float4*` is fine if aligned. But C++ syntax: `reinterpret_cast<const float4*>(Q + q_offset)`. When indexing: `[lane + i * 32]`. The base pointer is float4*, so this accesses the (lane + i*32)th float4. That's correct. Let me compile this. Actually, wait. `threadIdx.y` ranges from 0 to 31. And `warp` = `threadIdx.y`. So 32 warps. But CUDA limits to max 32 warps per block. 32*32=1024 threads. This is exactly at the limit. OK. What about blockIdx.z? The grid dimension is dim3(32, 32, 16). But maximum gridDim.z is 65535. 16 is fine. Wait, dim3 constructor: `dim3 grid(32, 32, 16)` or `dim3 grid(X, Y, Z)` where Z defaults to 1. So `dim3 grid(32, 32, 16)` is fine. Actually, I realize that `load_inline` might have limitations on what CUDA features work, but it should compile standard CUDA code. Let me write the full python module and compile it.

The kernel compilation timed out, likely because the inline PTX / C++ compilation was very slow with float4 operations. Let me try a simpler and more reliable approach that uses shared memory tiling and is more carefully constructed. I'll try a simpler tile-based approach. Let me also make sure it avoids float4 which can cause compilation complexity issues. Actually, the first kernel compiled and was correct. Let me try to optimize it further or use a simpler approach. Let me think about what gives good speedup for this problem. 32batch * 32heads, each with 512x1024 attention. This is 32*32 = 1024 heads, 512 seq len, 1024 dim. The PyTorch reference is using `scaled_dot_product_attention` which on H100 is likely extremely optimized (using FlashAttention-2/3 via cuDNN or native impl). Getting speedup over this is very hard. Let me think about what could be faster than PyTorch's native implementation. The native implementation likely already uses FlashAttention-3 on H100 which is highly optimized. For custom CUDA to beat it, we'd need to leverage specific hardware features or reduce overhead. Perhaps a simple kernel that: 1. Uses single loop (not tile-based, just direct) with vector loads 2. Or uses very large tiles with minimal synchronization Actually, with seq len 512 and dim 1024, each head does 512 * 512 * 1024 = 268M ops, times 1024 heads = ~274G flops. At 989 TFLOPS, this would take ~276 microseconds compute bound, but it's actually memory bound because Q/K/V need to be read. Wait, let me compute memory: 3 * 32 * 32 * 512 * 1024 * 4 bytes = 6.29 GB. At 3.35 TB/s, that's ~1.88 ms. But flash attention significantly reduces memory bandwidth by recomputing and fusing. The reference likely runs very fast. Let me try a simpler approach - maybe a direct unrolled attention that's optimized for exactly 512x1024. Given the dimensions are fixed, we can hardcode them. But wait - looking at the docstring more carefully: "Replace the forward computation of the PyTorch model below with a custom CUDA kernel". The model might not use the fixed dimensions; those are just for the benchmark. We should handle arbitrary dimensions. Given reference is `scaled_dot_product_attention` on H100, this will use extremely optimized fused kernels. Let me try to: 1. Use a simple but well-structured tile-based approach 2. Avoid complex template metaprogramming that causes slow compilation 3. Focus on getting correctness first, then speedup Actually, looking at the problem again - pyTorch's scaled_dot_product_attention may have some overhead from the Python/dispatch layer. A custom kernel called directly might have less overhead. But the actual GPU computation will be hard to beat. I'll try another approach: using a well-tuned flash-attention-like kernel but with very specific tile sizes for this workload. Let me try with smaller shared memory usage and simpler structure to avoid compile timeouts. Actually, one thing I could try is using half-precision internally even though input/output is fp32 (cast down, compute in fp16/bf16, cast up). On H100 this would give much higher compute throughput. But would the accuracy still pass? Let me check - accuracy is fp32 atol=rtol=1e-4. Using related formats internally might be borderline but should still pass for random data. Actually, wait - the reference implementation is fp32. The inputs are fp32. Let me look at the tolerance: fp32 atol=rtol=1e-4. Using TF32 (which H100 tensor cores use by default) should be fine. But for simplicity, let me stick to FP32 compute for now. Let me try a simpler, cleaner approach with the first kernel structure but optimizing the parallelization. The issue with the first kernel was likely it had shared memory bank conflicts and unoptimal thread usage. Let me try a new kernel design: - Each block handles 8 queries x 32 dims at a time - Each warp handles one query - Block size = 256 (8 warps) - For each key tile of 32 keys: - Load Q (8 x 64) to smem - Load K (32 x 64) to smem - Each warp computes 1 query x 32 keys partial dots - Use warp shuffle for reduction - Online softmax update - Compute O update Let me write a cleaner version. Actually, let me think differently. For 512x1024 attention, if I use: - BQ=16, BK=64, BD=64 - Each block handles 16 queries (2 warps per query, or 1 warp per query with 16 warps = 512 threads) Actually on H100, maximum threads per block is 1024, but occupancy matters. Let me try a very simple approach where each block handles fewer queries but with more parallelism. Or even simpler: each thread block handles exactly 1 query, and all 128/256/512 threads collaborate on that 1 query against all keys. For N=512, D=1024: - Block handles 1 query - 256 threads - All 256 threads load K (512 keys x 1024 dims) - that's 512K floats = 2MB, too much for smem - Need tiling Let me go with a BQ=8, BK=64, BD=64 tile size: - Block handles 8 queries - Block size = 256 threads (8 warps) - Each warp handles 1 query - For key tile: 64 keys, each warp computes scores for its query vs 64 keys - For dot product, each warp loads Q[query][d:d+64] in registers (shared within warp via shfl) - And loads K[key][d:d+64] from smem Actually, let me use the following design: - B_q = 8 queries per block - B_k = 64 keys per tile - B_d = 64 dims per tile - Block = 256 threads = 8 warps - Each warp handles 1 query × 8 keys (warp has 32 threads, use them along dim) Wait, let me simplify. Each block = 256 threads, 8 warps. - Warp 0 handles query 0, computing scores with all 64 keys (8 keys per thread using vectorized loads...) Actually, let me just use a simpler design: - B_q = 4, B_k = 64 - Block = 256 threads - Each of 4 warps handles 1 query - Within a warp, thread i computes score for key i (only 32 threads, so need 2 warps per query for 64 keys...) OK let me try: - B_q = 8 queries - B_k = 64 keys - 8 warps, each warp handles 1 query - For key tile of 64: each warp computes 64 scores but warp only has 32 threads - So each thread in warp computes 2 scores For dim=1024, we tile by BD=64: - Each query needs 64 dims loaded, that's 64 floats = 256 bytes - Q_smem = 8 × 64 floats - K_smem = 64 × 64 floats - Total smem = 8*64 + 64*64 = 4608 floats = ~18KB, very reasonable Let me write this kernel carefully. Wait, let me just go back to the first kernel concept but fix issues and make it compile faster. The first kernel compiled fine and was correct. Its issue was probably performance: it used a suboptimal configuration. Let me tune it. Actually, the first kernel had BQ=8, BK=16, DT=64. With 256 threads: - 8 warps (warp_id 0..7) - Each warp handles 1 query - lane_id 0..15 handles one key each (BK=16), and lane_id 16..31 handles the same keys for the second half of D - This means for D=1024, we'd need 1024/64 = 16 iterations - Each block does 8 queries at a time, so 512/8 = 64 iterations for Q - Total SMs = 132 or so on H100 For block count: B=32, H=32, so 1024 blocks. Each SM can run multiple blocks. But the first kernel was correct and compiled. Let me try to improve its performance: 1. Use larger BK to reduce loop iterations 2. Better load patterns 3. Vectorized loads if possible Actually, maybe I should try using the CUTLASS-style approach or just try a different parallelization: one thread block per (batch, head) with all threads cooperating. For seq=512, dim=1024, one thread block per head: - 512 threads per block - Each thread handles 1 query - All 512 threads cooperate on loading keys and values Actually thread blocks with 512 threads are fine. Let me try: block = (128 threads for D parallel) handling a subset of queries. Hmm, let me just try a cleaner version of the tile-based approach. Given that PyTorch's implementation is very optimized, my best bet is to: 1. Write a clean, simple kernel that minimizes overhead 2. Use appropriate tile sizes for the workload 3. Try different configurations to find what works Let me write a new kernel with: - B_q = 8 (queries per block) - B_k = 128 (keys per tile) - B_d = 128 (dims per tile) - Block size = 256 threads Smem: Q = 8*128 = 1024 floats, K/V = 128*128 = 16384 floats, P = 8*128 = 1024 floats, O = 8*1024 = 8192 floats. Total = ~107KB, which might be too large (limit is ~227KB but reduces occupancy). Let me try: B_q=8, B_k=64, B_d=64: - Q_smem = 8*64 = 512 floats - K_smem = 64*64 = 4096 floats - P_smem = 8*64 = 512 floats - O_smem = 8*1024 = 8192 floats - m/l = 8 + 8 = 16 floats Total = ~53KB. This is good for occupancy. Let me write this kernel with proper online softmax. Actually, I realize the first kernel already had similar structure but had a bug or performance issue. Let me think about if I can just improve that kernel. Wait - looking at the first kernel output pass, the max_abs_err=8.34e-07 which is very good. But was the kernel fast? We don't know since we only ran correctness. Let me just try the first kernel at timing by submitting it, or I should try to optimize further. Since we're limited to 30 tool calls, let me be more efficient. Let me write a more optimized version and try to compile it. Let me try a new approach: For each (batch, head), launch blocks where each block handles a range of queries. Use 256 threads, handle 8 queries per block, and tile over keys and dims. Actually, for this specific shape (512, 1024), let me try: - B_q = 4 (queries per block) - Block size = 512 threads (16 warps) - Each warp handles partial computations - Actually, 512 threads is too many for occupancy Let me try 128 threads per block, B_q = 4: - 4 warps - Warp 0 handles query 0, warp 1 handles query 1, etc. - Each warp has 32 threads - For key tile of 32: thread i computes score query vs key i - For dim=1024, tile by 64: 16 iterations This seems clean. Let me write it. Actually lane_id 0..31 directly maps to 32 keys if BK=32. But then for the O update, we need 32 keys x 1024 dims which is a lot of computation. Let me think about the O update. After computing P (B_q x B_k), we need to compute O += P @ V_tile. For B_q=4, B_k=32, B_d=64: - Each query needs to compute dot product with 32 keys over 64 dims - For query q, o_d = sum_k p[q][k] * v[k][d] - We can parallelize over d, having each thread handle some d's If block=128 threads, B_q=4: - We have 4 warps - For O update, we could have each warp handle 1 query, and the 32 threads in warp handle d in parallel (but D=1024, not 32) - Or each warp handles partial d range This gets complicated. Let me try a simpler assignment. For block=256 threads, B_q=8: - 8 warps, each handles 1 query - During score computation: each warp computes 8 scores (32 threads / need more) - Actually if BK=64, each thread computes 2 scores - We use warp shuffle to reduce across dim Let me write a cleaner kernel: For B_q=8, BK=64, BD=64, block=256: - 8 warps, warp id 0..7 handles query id 0..7 within the block - For score computation: - Each thread loads some part of Q and K - We need each thread to compute its assigned score Actually, let me assign: within a warp, thread i computes score for query q vs key k = base_k + i, where i = 0..31, so we handle 32 keys per warp. But BK=64, so we need 2 iterations per warp or use 64 threads per warp (impossible, warp is 32). So if BK=64, we could handle it in 2 sub-tiles of 32, or use a different distribution. Let me try BK=32 which fits in a warp. Then for N=512, we need 512/32 = 16 key tiles. For B_q=8, BK=32, BD=64: - Q_smem = 8 × 64 = 512 floats - K_smem = 32 × 64 = 2048 floats - P_smem = 8 × 32 = 256 floats - O_smem = 8 × 1024 = 8192 floats - m/l = 16 floats Total = ~41KB. Good. For block=256 threads = 8 warps: - Each warp handles 1 query - During score: warp computes 32 scores (1 per thread for 32 keys) - But each thread needs the dot product over 1024 dims - We tile over BD=64, so 16 iterations - In each iteration, thread accumulates partial dot For the dot product: thread k needs Q[q][0:1024] and K[k][0:1024]. - Q[q] is loaded to smem once per query (shared across warp) - K[k] is loaded to smem once per key tile - Dot product is done in tiles of 64 So in each BD iteration: - Load Q[q][d:d+64] to smem_q (done cooperatively by block) - Load K[base_k + k][d:d+64] to smem_k (done cooperatively) - Thread k reads smem_q[q][0:64] and smem_k[k][0:64], computes dot But thread k only needs smem_q for one specific q (its warp's q). So all 32 threads in warp 0 read the same smem_q[0][:] but different smem_k[k][:]. Wait, that's 32 threads reading the same smem_q row - bank conflicts! We need to avoid this. To avoid bank conflicts: have threads in a warp read different elements and use warp shuffle to broadcast. Or transpose the data. Actually for smem_q[8][64+pad]: if we use pad=1, then row q starts at q*(65). Threads reading column c from same row: addresses are q*65 + c, which for different c are different banks (bank = (q*65 + c) % 32). Since c differs by thread, no conflicts! Good. Wait no - bank conflicts happen when multiple threads read different addresses that map to the same bank. Bank = address_in_bytes / 4 % 32. So for smem_q[q][c], address = q*(64+pad) + c. If pad=1, address = q*65 + c. For thread k reading c=k, bank = (q*65 + k) % 32 = (q*1 + k) % 32 = (q+k) % 32. Since each thread has different k, banks are different. No conflict! For K: thread k reads smem_k[k][c] where c varies. Address = k*(64+pad) + c. Bank = (k*65 + c) % 32 = (k + c) % 32. For fixed k and varying c across threads? Wait, no - in our case, thread k reads all c in 0..63 (loop) for its dot product, not in parallel with other threads. Actually each thread computes its own dot product independently: ``` for c in 0..63: acc += smem_q[q][c] * smem_k[k][c] ``` This is done independently by each thread. No bank conflicts because it's a single thread reading sequentially! So the access pattern is: - smem_q: all threads read different columns from the same row (or different rows) - NO conflict with pad=1 - smem_k: each thread reads its own row sequentially - NO conflict This is great. Let me write this kernel. For O update: after computing P[8][32], we need to compute O[8][1024] += P[8][32] @ V[32][1024]. - For each query q, we do: o[d] = sum_k p[q][k] * v[k][d] - We can parallelize over (q, d) - With 256 threads and 8 queries x 1024 dims = 8192 elements to update - Each thread handles 8192/256 = 32 elements For d splitting: 1024 / 32 = 32 groups... Actually let's do: - thread idx = q * 128 + d where d handles 128 dims - No, block is 256 threads. For 8 queries, each query gets 32 threads. - Each of 32 threads handles 1024/32 = 32 dims. So thread assignment for O update: - warp_id = 0..7 (for 8 queries) - lane_id = 0..31 (for 32 dim groups) - Each warp handles 1 query, each lane handles D/32 = 32 dims For the O update with key tile of 32 keys: - For each key k in 0..31: - p = P[q][k] (scalar, same for all threads in warp) - For each d in thread's dim range: - o[d] += p * V[k][d] To load V efficiently: V is [32][1024]. We can load it to smem_v[32][64] for tiles of 64 dims. - Each key tile needs V[0:32][0:1024] - We process V in tiles of BD=64 dims - For each V tile: - Load V[k][d:d+64] to smem_v for all k - Each warp computes partial O update for its query Actually, we can also just load V from global memory directly without smem for the O update, but smem is faster. Let me write this kernel. I'll compute P first, do online softmax, then do O updates. Actually, doing online softmax correctly with tiled processing is tricky. Let me use the standard approach from FlashAttention: For each query q: ``` m = -inf, l = 0, o = 0 for each key tile j: s_j = q @ k_j^T # scores for this tile m_j = max(s_j) m_new = max(m, m_j) p_j = exp(s_j - m_new) corr = exp(m - m_new) # correction factor for previous accumulator o = o * corr + p_j @ v_j l = l * corr + sum(p_j) m = m_new end o = o / l ``` The update needs to happen for each key tile. Let me structure this as: 1. Compute s_j for the current tile (all 8 warps do this in parallel for their queries) 2. Within each warp, compute m_j (max of 32 scores), l_j (sum of exp scores) 3. Update running m, l, o For step 3, we need warp-level communication. Actually we can do the update inside each warp independently since each warp handles one query. Let me write the kernel properly. Actually wait - there's a subtlety. If we do the update per tile, we need: - After computing s_j, we need p_j = exp(s_j - m_new) for each key in the tile. We can store p_j in smem_p or in registers if we have enough. Each thread in warp has one score (since warp has 32 threads and key tile is 32). So each thread stores its own exp-adjusted score in a register. But for the O update, all threads in the warp need all p_j values (i.e., p_j for all 32 keys) because each O element is a weighted sum over all keys. So either: 1. Store p_j in smem_p[8][32], then all threads can read 2. Or do the O update in a way where threads accumulate independently Option 1 is cleaner. smem_p = 8 * 32 * 4 = 1024 bytes. Very small. But for O update, threads need to read p_j[k] for all k, which means 32 reads per thread. If each thread in warp handles 32 dims, that's 32 dims * 32 keys = 1024 multiply-adds. We can do this. Actually, if warp handles query q and lane i handles dims D_start:D_end: ``` for d in D_start to D_end: sum = 0 for k in 0 to 31: sum += P[q][k] * V[k][d] O[q][d] = corr_prev * O[q][d] + sum ``` This requires reading P[q][k] 32 times per d, which is 32*32 = 1024 reads from smem_p per thread. Since P[q][k] is the same for all threads in warp working on same q, we can have each thread in warp read a different P[q][k] and broadcast. But that's complex. Actually, with 8 warps and 32 threads per warp: - For each key k (0..31): - All threads in warp k? No... - Actually we can have the warp read P[q][0..31] into registers, one per thread. Thread i holds P[q][i]. - Then for computing O update for each d, we need all P[q][k]. We can use warp shuffle to get them: ``` float acc = 0; for (int k = 0; k < 32; ++k) { float p = __shfl_sync(0xFFFFFFFF, p_reg[k_owner], k); acc += p * V[k][d]; } ``` But k_owner is just thread k, so `p_reg[k_owner]` means the register value of thread k. We can use: ``` for (int k = 0; k < 32; ++k) { float p = __shfl_sync(0xFFFFFFFF, p_local, k); // p_local is this thread's p // Wait, __shfl_sync broadcasts from lane k to all lanes float p = __shfl_sync(0xFFFFFFFF, p_value_of_thread_k, k); ``` Actually __shfl_sync with srcLane=k gives the value from lane k. If each thread stores its own p_score, then: ``` float p_broadcast = __shfl_sync(0xFFFFFFFF, my_p_score, k); // gets lane k's value ``` This gives us p_broadcast = p[q][k] for all threads. Then each thread multiplies by V[k][d] for its own d. This is elegant! No smem_p needed for the O update. But we need V[k][d] for the correct k and d. Since each thread handles different d's, we all read different V addresses. For V access: if we load V tile to smem_v[32][64], then for each key k, thread reads smem_v[k][d_in_tile]. Each thread works on 32 dims, perhaps 16 dims per V tile of 64... Wait, each thread handles 32 dims out of 1024. We process V in tiles of BD=64. For each V tile, each thread handles 32 dims within that tile... but 32 > 64? No, if each tile is 64 dims and each thread handles 32, then only 2 threads per warp would be used for O update within one V tile. That's bad. Let me redesign: each warp handles 1 query, and for O update, each lane handles 1 dim across the full 1024 dims. But warp only has 32 threads, so each lane handles 1024/32 = 32 dims. For V tile of BD=64: each lane handles up to 32 dims, and each V tile is 64 dims. So 2 lanes cover 1 V tile... no wait, 32 lanes * 32 dims = 1024 dims. If V tile is 64, then all 32 lanes together cover the tile, with each lane handling 2 dims per tile? Actually 32 lanes * X dims = 64, so X = 2 dims per lane. So in each V tile of 64 dims, lane i handles dims i*2 and i*2+1. That's 2 multiply-adds per key. For 32 keys and 2 dims: 32 * 2 = 64 multiply-adds. With 32 threads in warp, each does 64 ops. Plus 32 shfl operations. This seems reasonable. Wait but in the previous part (score computation), lane i handled key i. Now in O update, lane i handles dims i*2..i*2+1. The warp role changes between phases. That's fine, we just reassign. Wait, I was thinking of this wrong. Let me restructure: Phase 1 (score computation): block-size cooperation - 8 warps, 32 threads each - Each warp handles 1 query - Thread i in warp computes score for key i (i < 32) - After reduction over BD tiles, we have s[q][i] for each thread Phase 2 (online softmax and O update): - Thread i still in warp for query q - It stores s[q][i] as its p_score after softmax adjustment - But O update needs p[q][k] for all k to compute all dims - We can do the O update across all d's with all threads, where each thread accumulates contributions for ALL dims it owns, iterating over keys and using shuffle to get p[q][k] Actually, to avoid the complexity, let me just store the full P matrix in smem. It's only 8 * 32 * 4 = 1KB. Very small. Then for O update: - For each query q, we iterate over V tiles - Load V tile to smem_v - Each warp handles 1 query - Each thread in warp handles 32 dims - For each key k, read p[q][k] from smem and multiply with V[k][d] This is simple and doesn't need shuffle. Let me go with this. So for V tile of BD=64: - warp for query q - thread i handles dims [i*32, i*32+31] ? No, with 32 threads and D=1024, each thread handles 32 dims. - But V tile is only 64 dims. So in one V tile, 2 threads are active? That's terrible utilization. Ah I see the issue. If block has 256 threads and D=1024, each thread naturally handles 4 dims (1024/256 = 4). So let me use this assignment globally: thread t handles dims t*4, t*4+1, t*4+2, t*4+3. For score computation, we still need threads organized by keys. So the thread role changes between score and O phases. That's fine. For score phase: warp 0-7 each handles 1 query, thread 0-31 in warp handles keys 0-31 For O phase: all 256 threads cooperate, thread t handles dims (t*4):(t*4+4) for all queries? Or per query? Actually let's think: if thread t handles o[q][t*4:d*4+4] for some q, then we need to iterate over all 8 queries, having threads work on their assigned dims for each query. For 256 threads and 8 queries * 1024 dims = 8192 elements, each thread handles 8192/256 = 32 elements = 8 float4s. Actually 8192 / 256 = 32. So each thread handles 32 dims. That's what I had before. Let me assign: thread t handles a group of 32 consecutive dims. t = 0 handles dims 0-31 for all queries? No, that doesn't make sense. Better: the 256 threads are divided into 8 groups of 32 (the warps). Group w handles query w. Within group, thread i handles dims i*32 to i*32+31. So during O update: - Each warp loads its p[q][0:32] from smem or registers - Each thread in warp reads p[q][k] for all k and multiplies with V[k][d] for d in [i*32, i*32+31] This is 32 reads from p and 32 reads from V per d, but d ranges over 32 values, so total 1024 ops. Wait, for each key k, thread reads V[k][d0:d0+32] which is 32 floats. For 32 keys, that's 32*32 = 1024 reads from V. Plus 32 reads from p. Each read does 32 multiply-adds. Total FMAs: 32 * 32 = 1024. Read bytes: 1024*4 + 32*4 = ~4KB from smem. This is fine. But here's the issue: during O update, we'd need to process V in tiles because smem_v can't hold 32*1024 floats = 128KB. We need to tile over D. If we process V in tiles of BD=64: - smem_v[32][64+pad] = 32 * 68 * 4 = ~8.5KB - Each thread handles 32 dims, so it needs 2 V tiles to cover its dims - Actually if each thread handles 32 dims and BD=64, then 2 threads cover one tile. The other 30 threads in warp cover... wait, 32 threads * 32 dims = 1024 dims. If tile is 64, then each thread handles 0.5 tiles worth. That's awkward. Let me use BD=32 for V tiles: - smem_v[32][32+pad] = 32 * 36 * 4 = ~4.5KB - Each thread handles 32 dims = 1 full tile - 32 threads cover 32 tiles simultaneously... no wait, all threads together handle all dims of one query. If each thread handles 32 dims and there are 32 threads, that's 1024 dims total. But they work on the same query in parallel. Actually I think I'm overcomplicating. Let me simplify: - For V tile of BD=32, we load 32 keys x 32 dims - Then each thread updates its 32 dims (which might span multiple tiles) - For a thread handling dims d_start:d_start+31, if d_start is aligned to 32, then it uses exactly one tile So with thread i handling dims [i*32, i*32+31]: - For V tile t (dims t*32 to t*32+31): - Thread i participates only if i == t - So only 1 thread per warp is active per V tile! This is terrible. We need all threads to be active. Alternative: assign threads to handle interleaved dims or non-consecutive dims. For example, thread i in warp handles dims i, i+32, i+64, ..., i+992 (stride-32 access). Then for V tile of 64 dims (say 0-63), all 32 threads are active, each handling 2 dims. Yes! This works better. Let me use: - Thread i handles dims [i, i+32, i+64, ..., i+992] = 32 dims total, spaced by 32 - For V tile of BD=64 (dims dt to dt+63): - Each thread handles dims i and i+32 within the tile (2 dims per thread) - All 32 threads are active This gives good parallelism. And bank conflicts for V: if smem_v[32][68] with pad=4, then addresses are k*68 + d. Bank = (k*68 + d) % 32 = (k*4 + d) % 32. For thread i reading d=i and d=i+32: bank1 = (k*4 + i) % 32, bank2 = (k*4 + i + 32) % 32 = (k*4 + i) % 32. Same bank! Conflict! To avoid this, use larger pad or different access pattern. If smem_v[32][64+8] = [32][72], address = k*72 + d. Bank = (k*8 + d) % 32. For d=i, bank = (k*8 + i) % 32. For d=i+32, bank = (k*8 + i + 32) % 32 = (k*8 + i) % 32. Still same! The issue is that i and i+32 differ by 32, so they map to the same bank. To avoid this, we could have each thread read a contiguous chunk of 2 dims, but then threads need to handle different offset pairs. Actually, if thread i reads dims 2i and 2i+1, then for a V tile of 64, all threads read contiguous pairs. Banks: (k*72 + 2i) % 32 and (k*72 + 2i+1) % 32. These differ by 1, so different banks. And across different i, no conflicts (2i and 2j differ by multiple of 2, and since 2 and 32 share factor 2, conflicts occur when i and j differ by 16). Actually (2i) % 32 = (2j) % 32 when i-j=16. So threads 0 and 16, 1 and 17, etc. conflict. With 32 threads in warp, and bank size = 32, if 2 threads read same bank, they conflict. So threads i and i+16 both read bank 2i (mod 32). This is a 2-way conflict. A 2-way bank conflict is not terrible. It doubles the access time. With good occupancy, this might be acceptable. Alternatively, we can pad more: smem_v[32][80], address = k*80 + d. Bank = (k*16 + d) % 32. For d=2i, bank = (k*16 + 2i) % 32. Threads i and i+16: (k*16 + 2i) % 32 vs (k*16 + 2i + 32) % 32 = same. Still conflict! Actually, the fundamental issue is that we have 32 threads and 32 banks, and we're reading 64 consecutive elements. Any assignment of 32 threads to 64 elements with each thread reading 2 elements will have at least 2 threads hitting the same bank, because 64 elements span exactly 2 complete bank cycles (64/32=2), but with 32 threads, by pigeonhole, some bank gets 2 accesses. Actually, with 32 threads each reading 2 elements, total 64 element-reads. There are 32 banks. Ideal is 2 reads per bank. So 2-way conflict is expected and unavoidable. But we can at least make it uniform (all banks have 2 accesses) rather than some having more. To achieve uniform 2-way: if thread i reads elements i and i+32, then bank(i) = i%32 = i, bank(i+32) = (i+32)%32 = i. Both go to bank i. So bank i gets exactly 2 accesses (from thread i). This is perfectly uniform! But wait, if all threads in warp access the same bank simultaneously, that's a 32-way conflict! Because bank i is accessed by thread i (for element i) and also thread i accesses bank i again (for element i+32). Actually no - each thread accesses 2 elements in sequence, not simultaneously. So first all threads access element i (different banks: 0,1,2,...,31), then all access element i+32 (same banks: 0,1,2,...,31). Each bank gets 1 access in each step. No conflict within a step! Ah, if the two reads are sequential (not in the same instruction), then there's no conflict in either step. Step 1: thread i reads bank i. Step 2: thread i reads bank i. Each step is conflict-free! So my original plan works fine. The compiler should interleave or sequence the reads properly. Now let me refine the plan: For each query group (8 queries per block): 1. Initialize O[8][1024] = 0, m[8] = -inf, l[8] = 0 2. For each key tile j (32 keys per tile, 16 tiles total for N=512): a. For each dim tile dt (64 dims per tile, 16 tiles total for D=1024): - Load Q[q][dt:dt+64] to smem_q for q=0..7 - Load K[k][dt:dt+64] to smem_k for k=0..31 - Each warp computes scores for its query vs all 32 keys - Thread i in warp computes score for key i - Accumulates dot product over dt tiles b. After all dt, we have s[q][i] in thread i's register c. Warp-level softmax: - m_j = warp_max(s[q][i]) over i=0..31 - Compute p_i = exp(s[q][i] - m_new) where m_new = max(m_prev, m_j) - l_j = warp_sum(p_i) - Update: corr = exp(m_prev - m_new), o = o * corr, l = l * corr + l_j, m = m_new d. Store p_i to smem_p[q][i] (or keep in register if possible) e. For each dim tile dt of V (64 dims): - Load V[k][dt:dt+64] to smem_v - Update O[q][d] += p[q][k] * V[k][d] for all d in thread's range 3. After all key tiles, normalize O: O[q][d] /= l[q] 4. Write O to global memory For step 2e, if we keep p_i in registers, we don't need smem_p. Each thread has its own p_i, and during O update, we need all p[q][k]. We can use warp shuffle or just read from smem_p. If using smem_p: after softmax, all threads write their p_i to smem_p[q][thread_id]. Then for O update, threads read p[q][k] for k=0..31. Since threads in warp read the same p[q][k] but for different d, and p is in smem, there might be conflicts if all threads read the same address. But reading from the same smem address across a warp is broadcast, not conflicting! CUDA handles this efficiently. Actually, for compute capability >= 7.x, reading the same address from smem by all threads in a warp is a broadcast and takes the same time as a single read (it doesn't serialize). This is great! So the plan: - Store p[q][0:32] in smem_p[q][0:32] - Each thread reads p[q][k] sequentially for k=0..31 - Each thread handles 32 dims (spaced by 32) - For V tile of 64 dims, each thread updates 2 dims For a V tile of 64 dims [dt:dt+64]: - Thread i in warp handles dims dt+i and dt+i+32 - For each key k in 0..31: - p = smem_p[q][k] (broadcast, fast) - v0 = smem_v[k][i] - v1 = smem_v[k][i+32] - o[i] += p * v0 - o[i+32] += p * v1 Wait, but if dt=0, thread i handles dims i and i+32. If dt=64, thread i handles dims 64+i and 64+i+32. So the O accumulator should be organized per query, with indices based on absolute d, not relative to tile. Each thread needs 32 float registers for its O values. That's fine (32 * 4 = 128 bytes per thread, 32 threads = 4KB per warp, well within register limits). Wait, but I said each thread handles 32 dims total, not per tile. So yes, 32 registers for O per thread. Actually, let me reconsider the thread assignment. With 8 warps and 32 threads each = 256 threads. For O update across 8 queries and 1024 dims: - Option A: warp handles 1 query, thread handles 32 dims. Total coverage: 8 warps * 32 threads * 32 dims = 8192 = 8*1024. Perfect. - Option B: threads not organized by warp, just flat 256 threads. Each handles 32 dims but for which query? We'd need 32 queries to use all threads, but we only have 8. So Option A is correct. Warp w handles query w, thread i handles dims i*32 to i*32+31? No, that's the issue with V tiles. Let me use: Warp w handles query w, thread i handles dims with step 32: i, i+32, i+64, ..., i+992. Total 32 dims. For V tile [dt, dt+64): - Thread i handles dt+i and dt+i+32, which are 2 of its 32 dims - If dt = 0, handles i and i+32 - If dt = 64, handles 64+i and 64+i+32 - etc. Good, this works. Now let me think about registers. Each thread needs: - O accumulator: 32 floats - Score accumulation: 1 float - p_value: 1 float - Maybe some temporaries Total: ~34 floats = 136 bytes. At 255 registers max, that's ~34 registers. Very comfortable. For smem: - smem_q[8][68] (pad=4): 8*68*4 = 2176 bytes - smem_k[32][68]: 32*68*4 = 8704 bytes - smem_v[32][68]: 32*68*4 = 8704 bytes - smem_p[8][32]: 8*32*4 = 1024 bytes - smem_m[8], smem_l[8]: 64 bytes Total: ~20.6 KB. Excellent for occupancy. Now let me write this kernel carefully. Wait, I realize a problem. In step 2a, I said load Q for q=0..7 and K for k=0..31. But Q is [8][1024] and we load it in tiles of 64 dims. For each dt, we load a [8][64] tile. Similarly K is [32][64]. The block has 256 threads. Loading Q tile: 8*64 = 512 floats. Each thread loads 2 floats. Loading K tile: 32*64 = 2048 floats. Each thread loads 8 floats. This is fine, we can loop. Actually, let me structure loads as float4 when possible for vectorized memory access. For Q [batch][head][N][D], the memory layout is N-major then D-major (row-major for the last two dims). So Q[q][d:d+4] are contiguous and can be loaded as float4. For Q tile of [8][64]: - Thread 0 loads Q[0][0:4], Q[0][4:8], etc. (since it needs to load 2 floats? Actually with 256 threads and 512 floats, each loads 2 floats, which is half a float4. Not optimal.) Let me change: have each thread load 1 float4 = 4 floats. Then 512/4 = 128 threads needed for Q. And 2048/4 = 512 threads needed for K. But we only have 256 threads. So we load Q with 128 threads and K with 256 threads... but all 256 need to participate. Actually, we can have all 256 threads participate in loading both Q and K sequentially: - First, 256 threads load 256*2 = 512 floats for Q (each thread loads 2 floats, or better, 128 threads load 1 float4 each and the other 128 do nothing... no, we want all threads busy for coalescing). Let me have each thread load float2 (2 floats). 256 threads load 512 floats. This is a single coalesced operation if the float2s are contiguous across thread idx. For Q: thead t loads float2 at Q + base + t*2. This is perfectly coalesced. For K: we need 2048 floats. We can do 4 rounds of 512 floats each, or have threads load more. Each thread can load float4 = 4 floats, so 256 threads load 1024 floats in one go. We need 2 such loads for K (2048 floats). Actually, to keep it simple, I'll use float4 where aligned and possible. For K tile of [32][64], each row is 64 contiguous floats. With float4, that's 16 float4s per row. Total: 32*16 = 512 float4s. We have 256 threads, so each thread loads 2 float4s. For Q tile of [8][64], each row is 16 float4s. Total: 8*16 = 128 float4s. We have 256 threads, so 128 threads load 1 float4 each. We can have threads 0..127 load Q, threads 128..255 wait or help with K. This is getting complicated. Let me keep it simple: each thread loads a fixed number of elements with a loop. Actually, the simplest is to use a uniform load pattern: ``` for (int i = threadIdx.x; i < total_elements; i += blockDim.x) { smem[i] = gmem[i]; } ``` This is simple and works, but might not be perfectly coalesced if the mapping from linear i to 2D isn't ideal. Smem layouts are 2D but gmem is also effectively 2D. Wait, if we flatten smem_q to 1D: smem_q[q*68 + d], and gmem is Q[head_offset + q*N*D + d], then: ``` int i = threadIdx.x; while (i < 8 * 64) { int q = i / 64; int d = i % 64; smem_q[q*68 + d] = Q[head_offset + (qg+q)*D + dt + d]; i += 256; } ``` This reads Q in a somewhat strided pattern across q. For q=0, thread 0 reads d=0, thread 1 reads d=1, etc. This is coalesced within a warp (threads read consecutive addresses). When we wrap around to q=1 at i=64, thread 0 reads Q[(qg+1)*D + dt + 0], which is far from Q[qg*D + dt + 0]. The memory address jumps by D=1024 floats = 4KB. Still, within each warp's 32 threads, if they all read from the same q, it's coalesced. But if a warp spans across q values, e.g., threads 0-31 reading i=0..31 (q=0, d=0..31) and i=32..63 (q=0, d=32..63) and i=64..95 (q=1, d=0..31), then all threads are still somewhat coalesced since consecutive threads read addresses that are either +1 or jump by ~4KB. Actually, better is to have warp-coordinated loads. With 256 threads = 8 warps, we can assign warp w to load row w of Q (8 rows total, 8 warps). Each warp loads 64 floats. 32 threads in warp load 2 floats each (or 1 float4 + 1 float4 for 16 threads, etc.) Let me do: - Warp w loads Q row w: 64 floats - Each thread in warp loads 2 floats (float2) - Thread lane i loads Q[w][2i:2i+2] for i=0..31, but wait warp has 32 lanes and 64 elements, so yes, each thread loads 2 contiguous floats as float2 - Coalescing: thread i reads address base + 2i, so addresses are 0,2,4,6,... which is coalesced (consecutive threads read consecutive addresses) For K: - 32 rows, 8 warps. Each warp loads 4 rows. - Warp w loads rows 4w, 4w+1, 4w+2, 4w+3 - Each row has 64 floats. 32 threads can load 2 floats each. - So warp loads 4 rows * 64 floats = 256 float loads = 128 float2 loads. - Each thread loads 128/32 = 4 float2s. - Thread i loads from row r = 4w + (i/8), offset = 2*(i%8) within that row... actually let me think. For 4 rows and 32 threads, we can assign: - Thread i: row = 4w + (i / 8), offset = 2 * (i % 8) - i=0: row 4w, offset 0 - i=1: row 4w, offset 2 - ... - i=7: row 4w, offset 14 - i=8: row 4w+1, offset 0 - etc. This is coalesced: threads 0..7 read row 4w offsets 0,2,4,...,14 - consecutive addresses! Threads 8..15 read row 4w+1 offsets 0,2,...,14 - also consecutive but 64 floats apart from previous row. Actually there's a gap of 64 floats between row 4w and 4w+1. This means the warp has two coalesced segments. This is still fine. The hardware can handle it. OK, now let me also consider the V loading. Same pattern as K. For the O update, I want to use float4 stores for writing results. Thread i in warp w handles query q=w, dims i, i+32, i+64, ..., i+992. That's 32 dims, not a multiple of 4 necessarily. Actually 32 is a multiple of 4, so 8 float4s per thread. But the dims are strided by 32, not contiguous. So we can't use float4 for O writing since O[q][i], O[q][i+32] are spaced 128 bytes apart (32 floats). For the final write to global memory, we just do 32 scalar stores per thread. That's fine, it's not on the critical path. Alternatively, we can store O in smem first, then do a coalesced write. But smem O is [8][1024]. Actually, we can write O directly from registers to gmem. Each thread writes 32 floats. 256 threads write 8192 floats simultaneously. The warp writes 32 threads * 32 floats = 1024 floats. But these are scattered (thread i writes to address base + i, base + i + 32, etc.). For warp 0 (query 0): thread i writes O[0][i], O[0][i+32], etc. Within warp, consecutive threads write to consecutive addresses only if we process one segment at a time. E.g., if all threads in warp first write their d=i element (for i=0..31), those addresses are base+0, base+1, ..., base+31. That's coalesced! Then we write d=i+32 elements: base+32, ..., base+63. Also coalesced. So if we loop over the 32 dim-groups and have all threads write one element each per group, we get perfect coalescing. For thread i handling dims {i, i+32, i+64, ..., i+992}: ``` for (int g = 0; g < 32; ++g) { int d = i + g * 32; O[q*D + d] = o_reg[g]; } ``` Warp writes: thread i writes address base + i + g*32. For fixed g, addresses are base+g*32, base+g*32+1, ..., base+g*32+31. Consecutive! Perfect coalescing. Great, we'll use this pattern. Now let me write the complete kernel. I'll also carefully handle the scale factor (1/sqrt(D)). Let me also think about whether to fuse the O update with the score computation tile-by-tile, or keep them as separate phases. In the FLashAttention algorithm, we update O after each key tile. This reduces register pressure and memory usage. Actually, for our tiled approach: - For each key tile j (32 keys): 1. Compute scores S_j for this tile (over all dim tiles) 2. Do online softmax to get P_j 3. Update O accumulator: O = O * corr + P_j @ V_j This means we need to keep O accumulator across key tiles. O is 8*1024 floats in smem = 32KB, or in registers. Actually with our design, O is in registers per thread (32 floats each). Total register storage: 256 threads * 32 floats = 8192 floats = 32KB. This is fine. Wait, but does each thread keep its own 32 O values? Yes. And P_j and V_j are used to update these. Let me now write the kernel step by step. Global setup: ``` int batch = blockIdx.x; int head = blockIdx.y; int qg = 0; // We handle all queries in one block? No, block handles 8 queries. If N=512, we need 64 blocks per head. ``` Wait, I only have grid(B, H). That's 32*32 = 1024 blocks. Each block needs to handle 512/8 = 64 query groups? No, I only have one block per (batch, head). So I need to either: 1. Have grid(B, H, 64) for 64 query groups 2. Or have block loop over query groups Option 1 is better for parallelism. Let me use grid(B, H, (N+7)/8). For N=512, that's 64. Total blocks = 32*32*64 = 65536. Each SM can handle multiple blocks. Block size = 256 threads. Block handles qg = blockIdx.z * 8 to qg+7. Let me update the grid: ``` dim3 grid(B, H, (N + BQ - 1) / BQ); ``` With BQ=8, grid.z = 64. Now for the kernel: ``` __global__ void attention_kernel(...) { int b = blockIdx.x; int h = blockIdx.y; int qg = blockIdx.z * BQ; // query group start int tid = threadIdx.x; int warp_id = tid / 32; int lane_id = tid % 32; int q_local = warp_id; // 0..7 int q = qg + q_local; if (q >= N) return; // head offset in global memory int base = ((b * H + h) * N) * D; int q_offset = base + q * D; // registers float o_reg[32]; #pragma unroll for (int i = 0; i < 32; ++i) o_reg[i] = 0.0f; float m = -1e20f; float l = 0.0f; float scale = 1.0f / sqrtf((float)D); __shared__ float smem_q[BQ][64 + 4]; __shared__ float smem_k[BK][64 + 4]; __shared__ float smem_v[BK][64 + 4]; __shared__ float smem_p[BQ][BK]; __shared__ float smem_m[BQ]; __shared__ float smem_l[BQ]; // if first warp, init smem_m and smem_l if (warp_id == 0) { if (lane_id < BQ) { smem_m[lane_id] = -1e20f; smem_l[lane_id] = 0.0f; } } __syncthreads(); // Pre-load Q to smem (or just process tile by tile) // Actually we need Q in registers per thread, not smem, since each warp handles one query // But we load Q to smem for cooperative loading then threads read for (int k_tile = 0; k_tile < N; k_tile += BK) { // Step 1: Compute scores for this key tile float score = 0.0f; for (int d_tile = 0; d_tile < D; d_tile += 64) { // Load Q tile: [8][64] // Warp w loads row w { int q_load = warp_id; int row_offset = qg + q_load; if (row_offset < N) { int gmem_q_base = base + row_offset * D + d_tile; int smem_q_row = q_load; // Each thread loads 2 floats (float2) for (int j = lane_id; j < 32; j += 32) { // j is the float2 index within the 64-float row // Actually we want each thread to load 2 floats, so 32 threads cover 64 floats float2 q2 = reinterpret_cast<const float2*>(Q + gmem_q_base)[j]; smem_q[smem_q_row][j*2 + 0] = q2.x; smem_q[smem_q_row][j*2 + 1] = q2.y; } } else { // zero fill for (int j = lane_id; j < 32; j += 32) { smem_q[q_load][j*2 + 0] = 0.0f; smem_q[q_load][j*2 + 1] = 0.0f; } } } // Load K tile: [32][64] // 8 warps, each warp loads 4 rows { int rows_per_warp = BK / 8; // 4 for (int r = 0; r < rows_per_warp; ++r) { int k_load = warp_id * rows_per_warp + r; if (k_tile + k_load < N) { int gmem_k_base = base + (k_tile + k_load) * D + d_tile; int smem_k_row = k_load; for (int j = lane_id; j < 32; j += 32) { float2 k2 = reinterpret_cast<const float2*>(K + gmem_k_base)[j]; smem_k[smem_k_row][j*2 + 0] = k2.x; smem_k[smem_k_row][j*2 + 1] = k2.y; } } else { for (int j = lane_id; j < 32; j += 32) { smem_k[k_load][j*2 + 0] = 0.0f; smem_k[k_load][j*2 + 1] = 0.0f; } } } } __syncthreads(); // Compute dot product: my query vs my key int my_k = lane_id; // key within this tile if (k_tile + my_k < N) { #pragma unroll for (int d = 0; d < 64; ++d) { score += smem_q[q_local][d] * smem_k[my_k][d]; } } __syncthreads(); } score *= scale; // Step 2: Online softmax for this tile float my_s = score; float m_tile = my_s; // warp max #pragma unroll for (int offset = 16; offset > 0; offset /= 2) { m_tile = fmaxf(m_tile, __shfl_xor_sync(0xFFFFFFFF, m_tile, offset)); } // load previous m, l from smem float m_prev, l_prev; if (lane_id == 0) { m_prev = smem_m[q_local]; l_prev = smem_l[q_local]; } m_prev = __shfl_sync(0xFFFFFFFF, m_prev, 0); l_prev = __shfl_sync(0xFFFFFFFF, l_prev, 0); float m_new = fmaxf(m_prev, m_tile); float corr_prev = expf(m_prev - m_new); float corr_tile = expf(m_tile - m_new); float p_val = expf(my_s - m_new); float l_tile = p_val; #pragma unroll for (int offset = 16; offset > 0; offset /= 2) { l_tile += __shfl_xor_sync(0xFFFFFFFF, l_tile, offset); } float l_new = l_prev * corr_prev + l_tile * corr_tile; if (lane_id == 0) { smem_m[q_local] = m_new; smem_l[q_local] = l_new; } // Store p to smem if (k_tile + lane_id < N) { smem_p[q_local][lane_id] = p_val; } else { smem_p[q_local][lane_id] = 0.0f; } __syncthreads(); // Step 3: Update O accumulator // O[q][d] = O[q][d] * corr_prev + sum_k(p[q][k] * V[k][d]) for (int d_tile = 0; d_tile < D; d_tile += 64) { // Load V tile { int rows_per_warp = BK / 8; for (int r = 0; r < rows_per_warp; ++r) { int k_load = warp_id * rows_per_warp + r; if (k_tile + k_load < N) { int gmem_v_base = base + (k_tile + k_load) * D + d_tile; int smem_v_row = k_load; for (int j = lane_id; j < 32; j += 32) { float2 v2 = reinterpret_cast<const float2*>(V + gmem_v_base)[j]; smem_v[smem_v_row][j*2 + 0] = v2.x; smem_v[smem_v_row][j*2 + 1] = v2.y; } } else { for (int j = lane_id; j < 32; j += 32) { smem_v[k_load][j*2 + 0] = 0.0f; smem_v[k_load][j*2 + 1] = 0.0f; } } } } __syncthreads(); // Update O for my query and my dims // Thread i in warp handles dims: d_tile + i and d_tile + i + 32 int d0 = d_tile + lane_id; int d1 = d_tile + lane_id + 32; float sum0 = 0.0f; float sum1 = 0.0f; #pragma unroll for (int k = 0; k < BK; ++k) { if (k_tile + k < N) { float p = smem_p[q_local][k]; sum0 += p * smem_v[k][lane_id]; sum1 += p * smem_v[k][lane_id + 32]; } } // Find which register slots these dims belong to // Thread lane_id handles absolute dims: lane_id, lane_id+32, lane_id+64, ..., lane_id+992 // d0 = d_tile + lane_id belongs to slot d_tile / 64 * 2 + 0? Not exactly. // Actually: d0 = d_tile + lane_id. Since lane_id is fixed, d0 takes values: lane_id, 64+lane_id, 128+lane_id, ... // These are slots: 0, 2, 4, ... (even slots) since each stride is 64 and we have 2 dims per tile. // Wait: for lane_id=0: dims are 0, 32, 64, 96, ... // d0 for tile 0: 0 (slot 0) // d0 for tile 1: 64 (slot 2, since 64/32=2) // So slot = d0 / 32 = (d_tile + lane_id) / 32. int slot0 = d0 / 32; int slot1 = d1 / 32; o_reg[slot0] = o_reg[slot0] * corr_prev + sum0; o_reg[slot1] = o_reg[slot1] * corr_prev + sum1; __syncthreads(); } } // Write out and normalize float l_final = smem_l[q_local]; float inv_l = 1.0f / l_final; #pragma unroll for (int g = 0; g < 32; ++g) { int d = lane_id + g * 32; O[q_offset + d] = o_reg[g] * inv_l; } } ``` Hmm wait, there's an issue with how I'm computing slot indices. Let me trace: - lane_id = 0: handles dims {0, 32, 64, 96, 128, 160, 192, 224, 256, 288, 320, 352, 384, 416, 448, 480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800, 832, 864, 896, 928, 960, 992} - For d_tile=0: d0=0, slot0=0; d1=32, slot1=1 - For d_tile=64: d0=64, slot0=2; d1=96, slot1=3 - For d_tile=128: d0=128, slot0=4; d1=160, slot1=5 So slots used by lane 0: 0, 1, 2, 3, 4, 5, ..., 31. That's all 32 slots. For lane_id=1: handles dims {1, 33, 65, ..., 993} - d_tile=0: d0=1, slot0=0; d1=33, slot1=1 - Wait, slot0 = 1/32 = 0, slot1 = 33/32 = 1. Hmm, but lane 0 and lane 1 both write to slot 0! o_reg[0] for lane 0 is dim 0, but o_reg[0] for lane 1 is dim 1. That's fine because each thread has its own o_reg array. Actually yes! Each thread has its own private o_reg[32]. Lane 0's o_reg[0] corresponds to dim 0, lane 1's o_reg[0] corresponds to dim 1, etc. They're independent. So the slot computation is: - slot0 = (d_tile + lane_id) / 32 - For lane_id=1, d_tile=0: slot0 = 1/32 = 0. Correct, dim 1 is the first element. - For lane_id=1, d_tile=64: slot0 = 65/32 = 2. Hmm wait, 65/32 = 2 (integer division). But dim 65 is the 3rd element for lane 1: {1, 33, 65, ...}. So it should be slot 2. Yes, correct. Wait let me recheck: lane 1's elements: e[0]=1, e[1]=33, e[2]=65. For d_tile=64, d0=65. 65/32 = 2 (integer division). Yes, that's e[2]. Correct. For d1=97: 97/32 = 3. e[3]=97. Correct. So slot0 and slot1 are correct. Now let me check the update: o_reg[slot0] = o_reg[slot0] * corr_prev + sum0 For tile 0 (d_tile=0): - corr_prev = exp(m_prev - m_new) for the first key tile - For the very first tile, m_prev = -inf, so corr_prev should be 0. But in our code, for the first tile, what is m_prev? We initialize smem_m[q_local] = -1e20f. Then: - m_prev = -1e20f - m_tile = max score in tile - m_new = max(-1e20f, m_tile) = m_tile - corr_prev = exp(-1e20f - m_tile) ≈ 0 - p_val = exp(score - m_tile) And o_reg starts at 0. So: o_reg[slot] = 0 * 0 + sum = sum For l_prev = 0, l_new = 0 * 0 + l_tile = l_tile. Correct. For the second tile: - m_prev = m_tile1 - m_tile = max score in tile2 - m_new = max(m_tile1, m_tile2) - corr_prev = exp(m_tile1 - m_new) - p_val = exp(score - m_new) - l_new = l_tile1 * corr_prev + l_tile2 * corr_tile And O update: o_reg = o_tile1 * corr_prev + sum_tile2 This is the correct online softmax update. Wait, but there's a subtlety! In step 3, after the first key tile, O accumulates using corr_prev. For the second tile, O_prev was computed with p1 * V1. We multiply by corr_prev = exp(m1 - m_new). This is correct because the "true" P1 weights should have been exp(S1 - m_new), not exp(S1 - m1). So we retroactively adjust O_prev by multiplying by exp(m1 - m_new). This is the standard FlashAttention online softmax approach. However, there is a correctness issue in my code! In the O update loop, I'm using corr_prev which is computed for the current tile. But corr_prev changes for each key tile. So for tile 2, I multiply the accumulated O (which includes tile 1's contribution) by corr_prev_2. This is correct. But wait, what about the P values stored in smem_p? They're p_val = exp(my_s - m_new). For tile 2, these use m_new for tile 2. So the V update for tile 2 uses correctly normalized P. And the old O is adjusted by corr_prev. This is correct. One concern: the loop variable `corr_prev` is the same for all threads in a warp because m_prev and m_new are the same for all threads in the warp (within a query). But between warps (different queries), corr_prev is different. In my code, `corr_prev` is computed per-warp because m_prev and m_new are warp-level quantities. This is correct! Thread i in warp w uses corr_prev for query q=w. But there's an issue: `smem_p` is shared across all warps, but each warp only reads its own row (q_local = warp_id). So no conflict. Another issue: in the V loading, I use: ``` for (int r = 0; r < rows_per_warp; ++r) { int k_load = warp_id * rows_per_warp + r; ``` This loads K rows in a strided pattern. For warp_id 0: rows 0,1,2,3. Warp_id 1: rows 4,5,6,7. etc. This is fine for correctness. For the score computation, Q is loaded similarly. Wait, I realize there might be an issue with the Q loading. In score computation, each warp loads Q for its corresponding query. But Q needs to be loaded for all 8 queries since all warps use all 8 Q rows (actually each warp only uses its own Q row). Wait no - warp w only needs Q for q_local = w. So warp w only needs to load Q row w. But in my code, I load all 8 Q rows cooperatively. Actually I have: ``` int q_load = warp_id; int row_offset = qg + q_load; ``` So warp w loads Q row w. But then other warps also load their own rows. After __syncthreads, all 8 rows are in smem_q. Warps then read smem_q[q_local] where q_local = warp_id. Warp w reads the row that warp w loaded. This is correct but somewhat redundant - each warp could just load its own Q directly to registers and avoid smem_q entirely! Let me think: if each warp loads Q for its query directly into a register array (64 floats shared across warp via shfl or just kept in smem), we could skip smem_q. But with 64 floats and each thread needing all 64 for the dot product, we'd need to broadcast. Smem is actually perfect for this: load once, read many times (each thread reads all 64 values). But each warp only reads one Q row. And there are 8 warps. With smem_q, all warps load their row, then __syncthreads, then each warp reads its own row. This is efficient enough. But we could eliminate the __syncthreads after Q load by having each warp load its own Q and immediately use it before the K load... no, because K is also needed. Actually, for the score computation, we need Q[q][d] and K[k][d] for the same d range. We can load both to smem and then compute. The __syncthreads after both loads ensures visibility. But we can optimize: each warp can load Q for its query into its own registers (broadcast via shfl) or into a dedicated smem area. Actually, since smem_q is only read by one warp (the warp that loaded it), there's no real need for __syncthreads on smem_q. The __syncthreads is needed for smem_k because all warps read all K rows. Wait, in the score computation, warp w computes dot products for its query vs all 32 keys. It needs K for all 32 keys, not just its own subset. So smem_k must be visible to all warps. smem_q is only needed by warp w for its own query, so as long as warp w sees its own write, it's fine. However, without __syncthreads between Q load and K load, there's no guarantee. We need __syncthreads after K load at least. Actually, for simplicity and safety, I'll keep a __syncthreads after both loads. Now, looking more carefully, there's a correctness issue in how I handle the score accumulation. Let me trace for a complete example: N=512, D=1024, BK=32, BD=64. Key tile 0: k = 0..31. dim tile 0: d = 0..63. - Load Q[0..7][0..63] to smem_q - Load K[0..31][0..63] to smem_k - __syncthreads - Each thread in warp w computes: score += sum_d smem_q[w][d] * smem_k[thread_id][d] This computes Q[w] · K[thread_id] for d in 0..63. After all 16 dim tiles, score = Q[w] · K[thread_id] over all 1024 dims. Correct. Now, lane_id 0 in warp 0 computes score = Q[0] · K[0]. lane_id 1 computes Q[0] · K[1]. etc. Warp-level max of these 32 scores gives the max score for query 0 vs keys 0..31. Correct. Then p_val = exp(score - m_new). Correct. For O update with V[0..31][0..1024]: We process V in dim tiles of 64. For dim tile 0: V[0..31][0..63]. - Thread 0 in warp 0: d0=0, d1=32. sum0 = sum_k p[0][k] * V[k][0] sum1 = sum_k p[0][k] * V[k][32] - These are correct. o_reg[0] gets sum0 (slot0 = 0/32 = 0). o_reg[1] gets sum1 (slot1 = 32/32 = 1). For dim tile 1: V[0..31][64..127]. - Thread 0: d0=64, d1=96. sum0 = sum_k p[0][k] * V[k][64] sum1 = sum_k p[0][k] * V[k][96] - o_reg[2] += corr_prev * old + sum0 (slot0 = 64/32 = 2) - o_reg[3] += corr_prev * old + sum1 (slot1 = 96/32 = 3) Wait! The update formula is: o_reg = o_reg * corr_prev + sum. For the first key tile, corr_prev = exp(m_prev - m_new). For the first tile, m_prev = -1e20f, so corr_prev is essentially 0. So o_reg becomes sum. That's correct. For the second key tile, we have accumulated O from tile 1 in o_reg. Now for tile 2, we compute corr_prev = exp(m_tile1 - m_new), where m_new = max(m_tile1, m_tile2). We then do: o_reg = o_reg * corr_prev + sum_tile2 This is correct - we adjust the old O by the correction factor and add the new contribution. After all 16 key tiles, o_reg contains the unnormalized O. Then we multiply by 1/l_final. Correct. Wait, there's another issue! The P values are stored as p_val = exp(score - m_new), where m_new is the running maximum after seeing the current tile. But for the O update, we use p_val directly without any additional correction. However, FlashAttention stores: - p = exp(S - m_new) (using the new max) - The old O is multiplied by exp(m_prev - m_new) - The new contribution uses p directly This is correct. After normalization by l_new at the very end, everything is consistent. Actually wait - there's another subtle point. In FlashAttention, after each tile: - l_new = l_prev * exp(m_prev - m_new) + sum(exp(S_j - m_new)) - O_new = O_prev * exp(m_prev - m_new) + P_j @ V_j, where P_j = exp(S_j - m_new) This is exactly what my code does. l_new is stored in smem_l, and O_new is in o_reg. But at the end, O_final = O_new / l_new. Correct. One more issue: in my code, I initialize smem_m and smem_l, but also have local variables m and l that are initialized but never used. Let me clean that up. Also, for the condition `if (k_tile + my_k < N)`, since N=512 and BK=32, and we iterate k_tile from 0 to 480 in steps of 32, the condition is always true for N=512. But to be general, I should keep it. However, since my_k = lane_id (0..31) and k_tile+31 < 512 when k_tile <= 480, the condition is always true. For generality with arbitrary N, it's good to keep. Actually, the issue is that the score computation loop has: ``` if (k_tile + my_k < N) { for (d = 0; d < 64; ++d) { score += smem_q[q_local][d] * smem_k[my_k][d]; } } ``` But my_k = lane_id, which is always < 32. And k_tile + lane_id < k_tile + 32. For the last tile where k_tile+32 might exceed N, some threads would skip the computation. But we should probably have all threads always compute since we pad smem_k with zeros. Actually no - if k_tile + my_k >= N, then that key doesn't exist. The score should be 0, and p_val should be 0. But the warp max would still be correct because other keys have real scores. The skipped contribution to l_tile would be 0 (since p_val=0 for out-of-bounds keys). And for O update, smem_p is already set to 0 for out-of-bounds. Wait, but if some threads don't participate in the dot product, the score variable for those threads remains at whatever it accumulated from previous tiles (which should be 0 since they also skipped). Let me make sure score is only accumulated for valid keys. Actually, looking at the code structure: for each key tile, we initialize score = 0.0f, then accumulate over dim tiles. So if a thread skips the inner loop, score remains 0. That's correct. But wait! I just noticed: score is initialized to 0 OUTSIDE the d_tile loop, but INSIDE the k_tile loop. So for each key tile, score starts at 0. Good. For key tiles that are partially out of bounds, some threads get score=0 and others get real scores. The warp max correctly computes the max. The out-of-bounds threads contribute p_val=exp(0 - m_new) if m_new < 0, or whatever. Wait, if m_new is from a real key with score=5, then p_val for out-of-bounds key with score=0 would be exp(0 - 5) = 0.0067, which is non-zero! This is wrong. Ah! This is a bug. For out-of-bounds keys, the score should be -inf, not 0. Because they shouldn't contribute to the softmax at all. Let me fix this: for threads where k_tile + lane_id >= N, set the score to a very negative number (like -1e20f) so it doesn't affect the max, and p_val = 0. Actually, in the original kernel I wrote, I had: ```cpp if (k_tile + my_k < N) { // compute dot } else { score = -1e20f; } ``` Or better yet, initialize score to -1e20f for invalid threads before the d_tile loop. Actually, it should be set after all d tiles. But if we set it to -1e20f at the start, it won't change. Wait, but if the thread is valid, we want score = real dot product. If invalid, we want score = -inf. So: ```cpp float score = (k_tile + lane_id < N) ? 0.0f : -1e20f; for (d_tile ...) { if (k_tile + lane_id < N) { score += smem_q[q_local][d] * smem_k[lane_id][d]; } } score *= scale; ``` Wait, multiplying -1e20f by scale (~0.03) gives -3e19, which is still very negative and should not affect max. But for safety, let me do: ```cpp float score; if (k_tile + lane_id < N) { score = 0.0f; for (d_tile ...) { score += ... } score *= scale; } else { score = -1e20f; } ``` Or actually, let me just use `-1e20f` for invalid keys. But what if all keys in the tile are invalid? Then m_tile = -1e20f, m_new = max(m_prev, -1e20f) = m_prev. And corr_prev = exp(m_prev - m_prev) = 1, corr_tile = exp(-1e20f - m_prev) = 0. And l_tile = 0 (since all p_vals are exp(very negative)). So no update. This is correct. But the issue is that for a partially valid tile, invalid keys have score=-1e20f, so they don't affect m_tile. p_val = exp(-1e20f - m_new) = 0. So they don't contribute. Correct. Let me add this to the code. Also, for the p storage: ```cpp if (k_tile + lane_id < N) { smem_p[q_local][lane_id] = p_val; } else { smem_p[q_local][lane_id] = 0.0f; } ``` Actually this was already in my code. But the issue is that p_val = exp(my_s - m_new), and if my_s = -1e20f, then p_val = 0. So we don't even need the if-check for smem_p. However, to avoid numerical issues, it's good to be explicit. OK let me write a clean version of this kernel now. Hmm, I'm worried about compilation time again. Let me try to keep the code relatively simple and avoid very deep unrolling if that causes compilation issues. The previous simple kernel compiled fine. Let me make sure this one isn't too complex. Also, I should consider: the first kernel was correct and compiled quickly. Maybe I can just try that one with timing, or make small optimizations. But I believe the new approach (with proper tile sizes and register usage) will be much faster. Let me write the clean kernel. I'll be careful about the syntax and try to compile it. Actually, I need to think about one more thing: the first kernel used a 2D grid and had the block loop over qg. Here I'm using a 3D grid. Both are fine. Let me also think about occupancy. Block size = 256 threads, smem ~24KB. On H100, each SM can have up to 32 blocks and 2048 threads. With 256 threads per block, max blocks per SM = 2048/256 = 8. With 24KB smem and 228KB total, max blocks = 228/24 ≈ 9. So occupancy is limited by threads to 8 blocks per SM, or 2048/132 SMs = about 15.5 blocks per SM on average across the GPU. Wait, with 65536 blocks and 132 SMs, that's ~496 blocks per SM. But each SM can only run 8 at a time (due to thread limit), so total concurrent blocks = 132 * 8 = 1056. The rest are queued. This is fine; the GPU scheduler handles it. Actually, with 65536 blocks and 1056 concurrent blocks, each block runs about 65536/1056 = 62 "waves". This might lead to some tail effect but should be OK. Wait, but each block only does 8 queries out of 512. So 64 blocks per head. With 1024 heads, that's 65536 blocks. Yes. Actually, I wonder if using fewer blocks with more work per block would be better. For example, each block handles 64 queries (all queries for one head). Then we'd have 1024 blocks total (one per head). With 1024 blocks and 1056 concurrent capacity, most blocks run simultaneously. But can a block handle 64 queries? That would require 64 warps = 2048 threads. At the block limit of 1024 threads (on most GPUs), we'd need to handle 32 queries with 1024 threads. Or handle 64 queries in a loop with 32 warps across multiple iterations. Actually, H100 allows up to 1024 threads per block in practice (sometimes 1536 or 2048 depending on config, but typically 1024). With 1024 threads = 32 warps, each warp handles 1 query, we can handle 32 queries per block. For N=512, we'd need 16 blocks per head. Total blocks = 32*32*16 = 16384. This is still plenty. With 1024 threads per block and 32 warps, each block handles 32 queries. Let me see if this is better. Block = 1024 threads, B_q = 32, BK = 32 or 64. Smem for Q = 32 * 64 = 2048 floats, K = 64 * 64 = 4096 floats, V = 64 * 64 = 4096 floats, P = 32 * 64 = 2048 floats. Total ~50KB. OK for occupancy. With 1024 threads = 32 warps, each warp handles 1 query. 32 queries per block. For N=512, blockIdx.z = 512/32 = 16. Total blocks = 32 * 32 * 16 = 16384. Concurrent capacity: 132 SMs * (1024 threads limit: 2048/1024 = 2 blocks per SM, or smem 228/50 = 4 blocks per SM). So limited by threads: 2 blocks per SM * 132 = 264 concurrent blocks. 16384 / 264 = 62 waves. Same number of waves but fewer blocks overall. Probably not much difference. With 256 threads per block, 65536 blocks, 1056 concurrent, 62 waves. With 1024 threads per block, 16384 blocks, 264 concurrent, 62 waves. The 256-thread version has better occupancy per block (each block finishes faster and frees up resources sooner). Actually, with smaller blocks, the GPU has more flexibility to schedule. 256 threads is a common sweet spot. Let me stick with 256 threads and B_q=8. One more concern: using 3D grid with z dimension might have some overhead compared to 2D grid with loop. Let me use 2D grid and loop over qg within block. This is what the first kernel did. But with 64 iterations, that's extra loop overhead. Hmm, but it's just a simple for loop. Actually, I'll use 1 block per (batch, head, query_group) = 3D grid since it's cleaner. PyTorch handles 3D grids fine. Wait, CUDA grid can only be up to 65535 in z dimension for some older compute capabilities, but for sm_90 (H100), it's 65535 as well? Let me check. Actually, for compute capability >= 3.0, the z dimension limit is 65535. And we need 64, which is well within. So 3D grid is fine. Let me now write the actual code. I'll be very careful. Actually, given my analysis and the complexity, let me try a simpler approach first: just use the first kernel but with a few optimizations. The first kernel was correct and compiled. Let me see if I can make it faster by tuning tile sizes and avoiding unnecessary shared memory. Wait, I realize the first kernel had BQ=8, BK=16, DT=64, block=256. I can try BK=64, which reduces the number of outer loops and might improve performance. But let me write a carefully optimized version instead. Let me try a much simpler approach: each block handles 1 query, and all threads in the block cooperate on that 1 query against all keys. This is very straightforward. For 1 query per block (batch, head, query_index): - Block handles 1 query - Block size = 512 threads (or 256) - We need to compute softmax(Q[k] @ K^T) @ V for one query With 512 threads: - Thread i handles key i for score computation? But N=512, so exactly 1 key per thread. - For dim=1024, we tile by 64 or 128. Score phase: - Load Q for this query (1024 floats). We can load to smem or keep in registers. - For each key k: - Thread k computes dot product Q · K[k] over 1024 dims - Since 512 threads, thread k handles key k and key k+256? Or just one key per thread. With 512 threads and 512 keys, each thread handles 1 key. Dot product over 1024 dims: - We can load K in tiles. For each dim tile of 64: - Load Q[0:64] to smem_q - Load K[k][0:64] to smem_k for all k... but that's 512*64 = 32768 floats, too large for smem. So we need to tile over keys too. For example, process 64 keys at a time: - Load Q to smem_q (1024 floats, tiled by 64) - For each group of 64 keys: - Load K for these 64 keys (64*64 = 4096 floats per dim tile) - 8 threads per key? Hmm. This gets complicated with 1 query per block. The multi-query-per-block approach seems better. OK, I'll write the 8-queries-per-block kernel now. But let me simplify one thing: I'll avoid float2 loads and just use simple float loads with proper coalescing. The compiler should handle this reasonably well. Actually, for coalescing, float4 is important for bandwidth. But since the dot product is compute-bound, load bandwidth might not be the bottleneck. And with shared memory tiling, we load each element once and reuse it many times. Let me write the kernel without float4 first to ensure correctness, then optimize. Wait, I already know the correctness approach works. Let me write a reasonably optimized version. One more thing: I realize I can avoid some shared memory by not using smem_p. Instead, I can compute the O update directly from the p values stored in registers, using warp shuffle. If each thread stores p_val in a register, and we need to read p[q][k] during the O update, we can use __shfl_sync to broadcast from the thread that owns that p value. For warp w (query q): - Thread k stores p_k = smem_p[q][k] (or directly computed) - During O update, for each key k in 0..BK-1: - p = __shfl_sync(0xFFFFFFFF, p_k, k) // broadcasts thread k's p to all threads - Each thread multiplies p by V[k][d] for their d Wait, __shfl_sync(0xFFFFFFFF, source_val, srcLane) gets the value from srcLane. If each thread has my_p in a register, then: ``` for (int k = 0; k < BK; ++k) { float p = __shfl_sync(0xFFFFFFFF, my_p, k); // p is now p[q][k] for all threads // each thread uses p with its own V[k][d] } ``` This broadcasts p from lane k to all lanes. It's very efficient! And we save smem_p. But for this to work, all threads in the warp must have computed their p_val. And for the O update, we iterate k from 0 to BK-1, broadcasting each p. This is a bit of overhead but saves smem. Let me use this approach. No smem_p needed. For the O update with V in smem_v[BK][BD+pad]: ``` for (int k = 0; k < BK; ++k) { float p = __shfl_sync(0xFFFFFFFF, my_p, k); if (d0 < D) o_reg[slot0] += p * smem_v[k][lane_id]; if (d1 < D) o_reg[slot1] += p * smem_v[k][lane_id + 32]; } ``` Wait, but this accumulates without the corr_prev factor. We need to apply corr_prev after each key tile! In FlashAttention, O is updated as: ``` O_new = O_old * corr_prev + P_tile @ V_tile ``` Where P_tile @ V_tile is the matrix product for the current tile. We compute this as a block and add to O_old * corr_prev. So the update is: ``` for each d: sum = 0 for k in tile: sum += p[q][k] * V[k][d] O[q][d] = O[q][d] * corr_prev + sum ``` We can pre-multiply O by corr_prev before adding sum: ``` for each register slot: o_reg[slot] *= corr_prev // then add sum ``` Or we can do it in the loop: ``` o_reg[slot] = o_reg[slot] * corr_prev + sum ``` Either way. But we need to multiply ALL slots by corr_prev before adding the new sum. We can do: ``` #pragma unroll for (int g = 0; g < 32; ++g) { o_reg[g] *= corr_prev; } ``` before the V loop, or inline during the V loop (but that would multiply by corr_prev for each k, which is wrong). So we must pre-multiply or post-multiply correctly. Pre-multiply approach: ``` // After computing p values and before V update for (g = 0; g < 32; ++g) o_reg[g] *= corr_prev; // Then accumulate the new contribution for (k = 0; k < BK; ++k) { float p = __shfl_sync(..., k); o_reg[slot0] += p * smem_v[k][lane_id]; o_reg[slot1] += p * smem_v[k][lane_id + 32]; } ``` This is correct and clean. But wait, what about d1 = d_tile + lane_id + 32 exceeding D? We should check. For D=1024, d_tile ranges 0..960 in steps of 64. d1 = d_tile + lane_id + 32. Max d1 = 960 + 31 + 32 = 1023. For lane_id=31, d1 = 1023 which is valid (last index). But for a generic D that's not a multiple of 64, we might need padding. Let's assume D is a multiple of 64 (which 1024 is). Also, for N not a multiple of 32, we'd need the out-of-bounds handling. OK, let me now write the actual kernel code. I'll keep it clean and well-commented. Hmm, after all this analysis, let me just try something simpler and see if it compiles. The key insight is that PyTorch's sdpa is already heavily optimized, so my custom kernel might not beat it easily. But let's try. Actually, let me also consider using Tensor Cores! On H100, we have FP32/BF16/FP16 Tensor Cores with massive throughput. If I can express attention using mma instructions or WMMA, it could be much faster. But writing raw tensor core code is very complex. Alternatively, I could try using CUTLASS through PyTorch's custom op. But that might be too complex for this format. Let me try my tiled kernel with careful shared memory usage and see how it performs. One last idea: since the PyTorch reference uses `torch.nn.functional.scaled_dot_product_attention`, which on H100 dispatches to FlashAttention-2 or FlashAttention-3, it's very hard to beat. A potentially faster approach is to use the same algorithm but with less overhead or better-optimized tiles for exactly this problem size. Or, I could just try to call cuDNN or cublas directly if they expose attention primitives. But cuDNN attention API is complex. Let me write the kernel and submit. I'll go with the 8-queries-per-block approach. Wait, I notice an important thing about the first kernel that passed correctness: the grid was (B, H) and block was 256. But the kernel had hardcoded BQ=8 and BK=16 with the outer loop over qg. So it was actually handling 8 queries per iteration, but within 2D grid... wait, let me re-examine. Looking at first kernel: ```cpp for (int qg = 0; qg < N; qg += BQ) { ``` This means each block loops over all query groups! With BQ=8 and N=512, that's 64 iterations. And grid(B, H) = 1024 blocks. Each block handles ALL queries for one (batch, head). So total work is 1024 blocks, each doing a lot of work. This is actually very good! Because 1024 blocks is less than the concurrent capacity. With 132 SMs and each running multiple blocks (say 8 with 256 threads), concurrent capacity is ~1056. So 1024 blocks can all run at once, each handling one head. So actually the first kernel's structure was already quite good in terms of parallelism. The issue was that BK=16 was too small, causing many iterations over key tiles (512/16=32 tiles). And DT=64 meant 16 iterations per tile for D=1024. If I increase BK to 64, key tiles reduce to 512/64=8. This could be much faster. Let me try a kernel with similar structure but BK=64 and optimized further. Also, looking at the first kernel's O update, it used smem_o[BQ][1024+PAD] and processed V in tiles. The O smem was 8*1028*4 = ~33KB. Total smem was probably much larger. Let me write a new kernel with: - Grid(B, H), block(256) - Each block handles ALL queries for one (batch, head) - B_q = 8 (per iteration) - B_k = 64 (keys per tile) - B_d = 64 (dims per tile) - Similar structure but optimized Actually wait - with block(B, H) and each block handling all queries, we have exactly 1024 blocks. For H100 with 132 SMs, each SM gets ~7-8 blocks. With 256 threads each, that's ~1792-2048 threads per SM, well within limits. And we have plenty of work per block. Let me write this kernel now, being very careful. Actually, I want to try something: use float4 for memory access where possible, and simplify the kernel to minimize register usage. Let me write the kernel: ```cpp #include <torch/extension.h> #include <cuda_runtime.h> #include <math.h> #define BQ 8 #define BK 64 #define BD 64 #define WARPS 8 #define THREADS 256 __global__ void attention_kernel( const float* __restrict__ Q, const float* __restrict__ K, const float* __restrict__ V, float* __restrict__ O, int B, int H, int N, int D) { int b = blockIdx.x; int h = blockIdx.y; int tid = threadIdx.x; int warp_id = tid / 32; int lane_id = tid % 32; int head_offset = ((b * H + h) * N) * D; float scale = 1.0f / sqrtf((float)D); __shared__ float smem_q[BQ][BD + 4]; __shared__ float smem_k[BK][BD + 4]; __shared__ float smem_v[BK][BD + 4]; __shared__ float smem_m[BQ]; __shared__ float smem_l[BQ]; for (int qg = 0; qg < N; qg += BQ) { // Initialize O, m, l for this query group float o_reg[32]; #pragma unroll for (int i = 0; i < 32; ++i) o_reg[i] = 0.0f; float m_reg = -1e20f; float l_reg = 0.0f; __syncthreads(); if (tid < BQ) { smem_m[tid] = -1e20f; smem_l[tid] = 0.0f; } __syncthreads(); for (int kt = 0; kt < N; kt += BK) { // Compute scores for this key tile float score = (qg + warp_id < N && kt + lane_id < N) ? 0.0f : -1e20f; for (int dt = 0; dt < D; dt += BD) { // Load Q tile for (int i = tid; i < BQ * BD; i += THREADS) { int q = i / BD; int d = i % BD; int q_idx = qg + q; smem_q[q][d] = (q_idx < N) ? Q[head_offset + q_idx * D + dt + d] : 0.0f; } // Load K tile for (int i = tid; i < BK * BD; i += THREADS) { int k = i / BD; int d = i % BD; int k_idx = kt + k; smem_k[k][d] = (k_idx < N) ? K[head_offset + k_idx * D + dt + d] : 0.0f; } __syncthreads(); // Compute partial score if (qg + warp_id < N && kt + lane_id < N) { #pragma unroll for (int d = 0; d < BD; ++d) { score += smem_q[warp_id][d] * smem_k[lane_id][d]; } } __syncthreads(); } score *= scale; // Online softmax float my_s = score; float m_tile = my_s; #pragma unroll for (int offset = 16; offset > 0; offset /= 2) { m_tile = fmaxf(m_tile, __shfl_xor_sync(0xFFFFFFFF, m_tile, offset)); } float m_prev = m_reg; float m_new = fmaxf(m_prev, m_tile); float corr_prev = expf(m_prev - m_new); float corr_tile = expf(m_tile - m_new); float p_val = expf(my_s - m_new); float l_tile = p_val; #pragma unroll for (int offset = 16; offset > 0; offset /= 2) { l_tile += __shfl_xor_sync(0xFFFFFFFF, l_tile, offset); } float l_new = l_reg * corr_prev + l_tile * corr_tile; m_reg = m_new; l_reg = l_new; // Pre-multiply O by corr_prev #pragma unroll for (int i = 0; i < 32; ++i) { o_reg[i] *= corr_prev; } // Update O with P_tile @ V_tile for (int dt = 0; dt < D; dt += BD) { // Load V tile for (int i = tid; i < BK * BD; i += THREADS) { int k = i / BD; int d = i % BD; int k_idx = kt + k; smem_v[k][d] = (k_idx < N) ? V[head_offset + k_idx * D + dt + d] : 0.0f; } __syncthreads(); int d0 = dt + lane_id; int d1 = dt + lane_id + 32; int slot0 = d0 / 32; int slot1 = d1 / 32; float sum0 = 0.0f; float sum1 = 0.0f; #pragma unroll for (int k = 0; k < BK; ++k) { float p = __shfl_sync(0xFFFFFFFF, p_val, k); if (kt + k < N) { sum0 += p * smem_v[k][lane_id]; sum1 += p * smem_v[k][lane_id + 32]; } } if (qg + warp_id < N) { if (d0 < D) o_reg[slot0] += sum0; if (d1 < D) o_reg[slot1] += sum1; } __syncthreads(); } } // Write out int q_idx = qg + warp_id; if (q_idx < N) { float inv_l = 1.0f / l_reg; #pragma unroll for (int g = 0; g < 32; ++g) { int d = lane_id + g * 32; if (d < D) { O[head_offset + q_idx * D + d] = o_reg[g] * inv_l; } } } __syncthreads(); } } ``` Wait, I see an issue: I'm using m_reg and l_reg as per-thread registers, but these need to be per-query (per-warp), not per-thread. Since all threads in a warp compute the same query, and they all compute the same m_new and l_new, having each thread keep its own m_reg and l_reg is fine as long as they all compute the same values. And they do, because the warp-level reductions ensure all lanes get the same result. But wait - m_reg and l_reg are initialized per-thread. For all threads in warp w, m_reg starts at -1e20f and l_reg at 0. After processing each tile, all threads in the warp update to the same values (because __shfl gives the same result to all lanes). So this is correct. However, the initialization inside the qg loop: ```cpp if (tid < BQ) { smem_m[tid] = -1e20f; smem_l[tid] = 0.0f; } ``` sets smem_m/l which I never use! I can remove smem_m and smem_l entirely since I'm using m_reg and l_reg in registers. Actually, I don't need smem_m or smem_l at all. Let me remove them to save shared memory. Let me also think about the O update more carefully. I have: ```cpp for (int dt = 0; dt < D; dt += BD) { ... int d0 = dt + lane_id; int d1 = dt + lane_id + 32; int slot0 = d0 / 32; int slot1 = d1 / 32; // compute sum0, sum1 if (qg + warp_id < N) { if (d0 < D) o_reg[slot0] += sum0; if (d1 < D) o_reg[slot1] += sum1; } } ``` For D=1024 and dt from 0 to 960 step 64: - dt=0: d0=lane_id, slot0=0; d1=lane_id+32, slot1=1 - dt=64: d0=64+lane_id, slot0=2; d1=96+lane_id, slot1=3 - dt=128: d0=128+lane_id, slot0=4; d1=160+lane_id, slot1=5 Wait! slot0 = (64 + lane_id) / 32 = 2 + lane_id/32 = 2 (since lane_id < 32). slot1 = (96 + lane_id) / 32 = 3 + lane_id/32 = 3. So dt=0 uses slots 0,1. dt=64 uses slots 2,3. dt=128 uses 4,5. This covers all 32 slots over 16 iterations. Correct! But wait, d0/32 with d0=64+lane_id gives 2 + lane_id/32 = 2 for lane_id in [0,31]. Correct. However, I think there's a flaw. Let me check for lane_id=0: - dt=0: d0=0, slot0=0; d1=32, slot1=1 - dt=64: d0=64, slot0=2; d1=96, slot1=3 - dt=128: d0=128, slot0=4; d1=160, slot1=5 Yes, correct. For lane_id=31: - dt=0: d0=31, slot0=0; d1=63, slot1=1 - dt=64: d0=95, slot0=2; d1=127, slot1=3 Yes, 95/32=2 (integer), 127/32=3. Correct. Great, this works. One concern: the `__shfl_sync` for p_val. The source lane is `k`. But `k` ranges 0..63 (BK=64), and warp only has 32 lanes. So __shfl_sync(0xFFFFFFFF, p_val, k) for k >= 32 will get the value from lane k-32 in the same warp? No, __shfl_sync reads from srcLane which must be in [0,31]. If k >= 32, this is out of bounds! Wait, BK=64! A warp only has 32 threads, but BK=64 keys. So __shfl_sync with srcLane >= 32 is invalid! I need to handle BK=64 with 32-thread warps. Options: 1. Use two shuffles: for k in 0..31 and k in 32..63, but lanes 32..63 don't exist. 2. Use shared memory for P: smem_p[64] where key k's score is stored by some thread. 3. Reduce BK to 32. If BK=32, then __shfl_sync(0xFFFFFFFF, p_val, k) for k in 0..31 works perfectly. But we'd need 512/32 = 16 key tiles instead of 8. That's more tiles but simpler. Let me use BK=32. With BK=32: - Key tiles: 512/32 = 16 - B_k = 32, B_d = 64 - smem_k[32][68] = 32*68*4 = ~8704 bytes - smem_v[32][68] = same - smem_q[8][68] = 8*68*4 = ~2176 bytes - Total smem = ~20KB And __shfl_sync with srcLane in 0..31 works. For BK=32, the p_val distribution: - lane_id 0..31 each computes one score and one p_val. - __shfl_sync(0xFFFFFFFF, p_val, k) gets p_val from lane k. This is perfect. But wait, can we use __shfl_sync for p_val when my_s = -1e20f? The p_val will be exp(-1e20f - m_new) which underflows to 0. So it's fine. However, all threads in the warp need to participate in __shfl_sync for it to work. Threads with out-of-bounds queries (qg + warp_id >= N) still need to call __shfl_sync or the warp will hang. Actually, I have the condition `if (qg + warp_id < N && kt + lane_id < N)` for score computation, which means some threads might not reach the __shfl_sync if the condition is false for them. But __shfl_sync requires a mask of participating threads. If some threads don't participate (because they skipped the computation), we should still call __shfl_sync with all threads active. Let me restructure: always compute score (0 or -1e20f based on validity), always do warp reductions with all threads, always do __shfl_sync with all threads. The only place we need conditionals is for loading from global memory and writing to global memory. Actually, since N=512 is a multiple of all our tile sizes, and we'll call the kernel with dimensions that are multiples, we might not need many conditionals. But to be general, let's keep some. For the score computation, if kt + lane_id >= N, set score = -1e20f. All threads still participate in __shfl_xor and __shfl_sync. For the O update, if kt + k < N, use p * smem_v[k][...], else add 0. Actually since p_val = 0 for out-of-bounds, and smem_v is 0 for out-of-bounds (from the load), the product is 0 regardless. So we might not even need the conditional in the inner loop. Wait, but if kt + lane_id >= N, the p_val is -1e20f, and exp(-1e20f - m_new) = 0. But if m_new is also -1e20f (all keys in tile invalid), then exp(-1e20f - (-1e20f)) = exp(0) = 1. That's wrong! So if all keys in a tile are invalid, we need m_tile = -1e20f, and the tile should contribute nothing. Let me trace: - All scores = -1e20f - m_tile = -1e20f (from reduction) - m_new = max(m_prev, -1e20f) = m_prev - corr_prev = exp(m_prev - m_prev) = 1 - corr_tile = exp(-1e20f - m_prev) = 0 - p_val = exp(-1e20f - m_prev) = 0 (for all threads) - l_tile = 0 - l_new = l_prev * 1 + 0 = l_prev Then O update: - o_reg *= corr_prev = o_reg * 1 = o_reg - sum over k: p_val * V[k][d] = 0 * V = 0 - o_reg += 0 So nothing changes. This is correct! But for p_val computation: if my_s = -1e20f and m_new = m_prev (e.g., 5), then p_val = exp(-1e20f - 5) which definitely underflows to 0. In practice, with float, exp(-1e20f) is 0. So p_val = 0. Is there a case where p_val becomes NaN? If m_new is also -1e20f and my_s is -1e20f, then -1e20f - (-1e20f) = NaN? Actually, -1e20f is a representable float. -1e20f - (-1e20f) = 0. So exp(0) = 1. But we saw that if m_new = m_prev = -1e20f, then corr_prev = 1, corr_tile = 0, and the update does nothing. And p_val = exp(-1e20f - (-1e20f)) = 1? Wait no: If m_prev = -1e20f and m_tile = -1e20f: - m_new = max(-1e20f, -1e20f) = -1e20f - corr_prev = exp(-1e20f - (-1e20f)) = exp(0) = 1 - corr_tile = exp(-1e20f - (-1e20f)) = 1 - p_val = exp(my_s - m_new) = exp(-1e20f - (-1e20f)) = 1 for all threads! Then l_tile = 32 (all 32 threads contribute 1) l_new = 0 * 1 + 32 * 1 = 32 And O update: o_reg *= 1, then sum += 1 * V[k][d] for all k. This adds V for non-existent keys! This is a bug! If all keys in a tile are invalid, we should add nothing. The issue is that when all keys are invalid, my_s = -1e20f, but we still compute p_val = 1 because of the cancellation. This is wrong. To fix this, we need to explicitly handle out-of-bounds keys by setting p_val = 0 when kt + lane_id >= N. We should mask the p_val. Let me add: ```cpp float p_val = expf(my_s - m_new); if (kt + lane_id >= N) p_val = 0.0f; ``` This ensures invalid keys contribute 0. And when all are invalid, l_tile = 0, l_new = l_prev * corr_prev + 0 * corr_tile = l_prev * 1 = l_prev. No contribution. Correct. Similarly, I need to make sure the O update doesn't add anything. With p_val = 0 for invalid keys, sum is 0. Good. Let me also handle V loading: for k where kt + k >= N, smem_v[k][...] should be 0 (or we just ensure p_val is 0). Either way works. One more thing: when kt + lane_id >= N but we still call __shfl_sync(0xFFFFFFFF, p_val, k), what happens? p_val is in a register. Even if lane_id is out of bounds, it still has a p_val (which we set to 0). And __shfl_sync reads p_val from the source lane. So even for k where kt+k >= N, we get p_val = 0 (because that lane set it to 0). This is correct! So the key fix is: always set p_val = 0 for invalid keys before the __shfl_sync. OK, let me write this kernel now. Actually wait - I realize __shfl_sync's srcLane parameter is just the lane ID (0-31). It works regardless of whether that lane computed a valid score. So we're good. Also, I need to handle the case where qg + warp_id >= N. In that case, the query doesn't exist. All the computation for that warp should produce 0 output. And importantly, the __shfl_xor and __shfl_sync calls must still happen to avoid warp divergence issues. If qg + warp_id >= N: - Score should be -1e20f - m_new and l_new computations should result in no change - O should remain 0 - At write time, skip writing So: ```cpp bool valid_q = (qg + warp_id < N); float score = (valid_q && kt + lane_id < N) ? 0.0f : -1e20f; ``` Wait, but if q is invalid and k is valid, score = -1e20f. This gives p_val = 0 after softmax. m_tile = -1e20f. m_new = max(m_prev, -1e20f) = m_prev. So no change. O stays 0. This is correct. And __shfl_xor still works since all lanes have valid float values. For the O update, valid_q needs to be checked before writing o_reg. Actually, since o_reg stays 0 and l_reg stays 0 (if no valid keys ever processed), we'd divide by 0 at the end. Wait! If qg + warp_id >= N, then no valid query exists. m_reg stays at -1e20f, l_reg stays at 0. At the end, inv_l = 1/0 = inf. And we'd try to write inf * 0 = NaN. We need to handle this. For invalid queries, just write 0 (or skip writing). So at write time: ```cpp if (valid_q) { float inv_l = 1.0f / l_reg; // write } else { // write 0 or skip int q_idx = qg + warp_id; for (g = 0; g < 32; ++g) { int d = lane_id + g * 32; if (d < D) O[head_offset + q_idx * D + d] = 0.0f; } } ``` Actually, we can skip the O write for invalid queries since O is already 0-initialized by torch.empty_like? No, empty_like doesn't initialize to 0; it has garbage. We need to write 0 for invalid queries. Hmm, for simplicity, since the benchmark uses N=512 and all our tile sizes divide evenly, and our grid covers exactly N queries (no padding needed), we might never hit invalid queries. But to be safe, I'll handle it. Let me also consider: do I even need the `valid_q` check for score? If qg + warp_id < N is always true (because qg ranges in multiples of BQ and N is a multiple of BQ), then no check needed for N=512, BQ=8. But the kernel should handle arbitrary N. Let me put the check in and see if the compiler optimizes it away when it can. Actually, for performance, having conditionals inside the hot loop is bad. But `valid_q` is the same for all lanes in a warp, so it's a uniform branch and should be efficient. For `kt + lane_id < N`, this is also uniform in the sense that for a given kt, either all lanes pass or some pass and some fail (only at the last tile). For the common case where kt + 31 < N, all lanes pass, and the compiler can optimize. For the last tile, some lanes take the else branch, but warp divergence is limited to one tile out of many. So performance impact is small. Given all this analysis, let me write the kernel now. Oh wait, one more thing: `__syncthreads()` inside `if (valid_q)` would be wrong because it needs to be called by all threads. I should not put __syncthreads inside conditional blocks. Let me restructure to avoid this. For score computation: ```cpp for (int kt = 0; kt < N; kt += BK) { bool valid_k = (kt + lane_id < N); bool valid_qk = valid_q && valid_k; float score = valid_qk ? 0.0f : -1e20f; for (int dt = 0; dt < D; dt += BD) { // Load Q and K (all threads participate) // ... __syncthreads(); if (valid_qk) { #pragma unroll for (int d = 0; d < BD; ++d) { score += smem_q[warp_id][d] * smem_k[lane_id][d]; } } __syncthreads(); } score *= scale; // softmax (all threads participate) // ... } ``` Wait, but if valid_q is false, valid_qk is false, and the threads skip the inner score accumulation. But they still participate in __syncthreads and the warp shuffle operations. This is correct. For the warp shuffle operations, all threads must participate. Since we're calling __shfl_xor_sync(0xFFFFFFFF, ...), all 32 lanes are active. And if some lanes have score=-1e20f, the reduction still works correctly. However, __shfl_xor_sync uses a mask. If some threads are inactive in the warp due to divergence, we need to be careful. But in our case, all threads reach __shfl_xor_sync. The `if (valid_qk)` only affects the inner score accumulation, not the control flow reaching the shuffle. Wait, the `if (valid_qk)` is inside the dt loop. The shuffle operations are outside the dt loop. So all threads reach the shuffle regardless of valid_qk. This is correct. For O update with __shfl_sync, same thing - all threads in warp reach it. Good. Let me write this. Actually, to make it even cleaner, I could avoid the `if (valid_qk)` around the inner loop by having all threads accumulate 0 if invalid: ```cpp float qk = valid_qk ? smem_q[warp_id][d] * smem_k[lane_id][d] : 0.0f; score += qk; ``` But this introduces a conditional inside the innermost loop which might hurt performance. Better is to have divergence only once per dt tile: ```cpp if (valid_qk) { #pragma unroll for (int d = 0; d < BD; ++d) { score += smem_q[warp_id][d] * smem_k[lane_id][d]; } } ``` Since valid_qk is uniform across the warp (for valid_q true, valid_k is uniform for a full tile), this branch is either all-take or all-dont-take for the warp. So no warp divergence! Wait, valid_k = kt + lane_id < N. For a given kt, this is true for lanes 0..(N-kt-1) and false for the rest. So for the last tile, it's not uniform. But for all other tiles, it is uniform. For the last tile, half the warp might diverge. That's still OK. And valid_q is the same for all lanes in a warp (depends only on warp_id). So `valid_qk` is uniform if valid_q is false (all lanes false) or if valid_k is uniform (most tiles). The divergence is minimal. Let me write the final kernel now. I'll compile it and see.