"""Custom Triton paged-attention decode kernel (memory-bound, single-query). Design ------ Decode is bandwidth-bound: the KV cache must be streamed end-to-end once. We launch ONE program per (batch element, KV head) and walk the pages of sequence b via block_table, reading each KV token exactly once and running a FlashAttention style *online softmax* so K and V are consumed in a single streaming pass. A single (b, hkv) program owns the group of G = num_heads / num_kv_heads query heads that share that KV head, so KV is fetched exactly once regardless of the GQA ratio (matching problem.yaml's minimal bytes_formula). To hide memory latency when the batch is small, each sequence is optionally split across several programs (FlashDecoding): kernel 1 produces, per (b, hkv, split), a *partial* (running max, normalizer, weighted sum) over its page range; kernel 2 combines the partials with the online-softmax combine rule. This keeps the GPU's SMs saturated even for batch=4 / long context. KV cache layout: (num_blocks, page_size, num_kv_heads, 2 * head_dim) with the last dim packing [K | V], so a single gather pulls both halves. Triton notes (learned the hard way on this build): * A tl.dot whose contraction dimension is not stride-1 contiguous in memory HANGS on Hopper. The token dimension is strided (stride 2048) in this KV layout, so the V aggregation is done as a broadcast-multiply + sum over tokens (reads V contiguously per token) instead of a tl.dot. * A combined (row & col) store mask HANGS on this build; a row-only mask is used (BLOCK_D == D always, so no column masking is needed). * The MMA requires the leading (query-head) dimension to be a multiple of 16, so BLOCK_G is rounded up to 16 and padding rows are reset every step. """ import math import torch import torch.nn as nn import triton import triton.language as tl @triton.jit def _decode_split_kernel( Q, KV, BLOCK_TABLE, M_PART, L_PART, A_PART, SEQ_LENS, # kv strides stride_kv_b, stride_kv_p, stride_kv_h, stride_kv_d, # q strides stride_q_b, stride_q_h, stride_q_d, # partial strides: (B, Hkv, num_splits, ...) stride_m_b, stride_m_h, stride_m_s, stride_m_g, stride_l_b, stride_l_h, stride_l_s, stride_l_g, stride_a_b, stride_a_h, stride_a_s, stride_a_g, stride_a_d, stride_bt_b, stride_bt_p, scale, B, HKV, G, P, D, num_splits, pages_per_split, num_pages, BLOCK_D: tl.constexpr, BLOCK_G: tl.constexpr, GROUP: tl.constexpr, PAGE: tl.constexpr, NEG_INF: tl.constexpr, ): pid_b = tl.program_id(0) pid_hkv = tl.program_id(1) pid_split = tl.program_id(2) seqlen = tl.load(SEQ_LENS + pid_b).to(tl.int32) num_pages = num_pages # python int passed in start = pid_split * pages_per_split end = start + pages_per_split if end > num_pages: end = num_pages offs_g = tl.arange(0, BLOCK_G) offs_d = tl.arange(0, BLOCK_D) keep = offs_g < GROUP # Load Q for the G query heads mapped to this KV head. q_ptr = Q + pid_b * stride_q_b + pid_hkv * GROUP * stride_q_h Qg = tl.load(q_ptr + offs_g[:, None] * stride_q_h + offs_d[None, :] * stride_q_d, mask=(offs_g[:, None] < GROUP), other=0.0).to(tl.float32) m = tl.full((BLOCK_G,), NEG_INF, dtype=tl.float32) lsum = tl.zeros((BLOCK_G,), dtype=tl.float32) acc = tl.zeros((BLOCK_G, BLOCK_D), dtype=tl.float32) offs_p = tl.arange(0, PAGE) for page in range(start, end): blk = tl.load(BLOCK_TABLE + pid_b * stride_bt_b + page * stride_bt_p).to(tl.int32) kv_base = KV + blk * stride_kv_b + pid_hkv * stride_kv_h k_ptrs = kv_base + offs_p[None, :] * stride_kv_p + offs_d[:, None] * stride_kv_d KT = tl.load(k_ptrs).to(tl.float32) # (D, PAGE) scores = tl.dot(Qg, KT) * scale # (BLOCK_G, PAGE) tok = page * PAGE + offs_p valid = tok < seqlen scores = tl.where(valid[None, :], scores, NEG_INF) m_new = tl.maximum(m, tl.max(scores, axis=1)) # (BLOCK_G,) p = tl.exp(scores - m_new[:, None]) # (BLOCK_G, PAGE) corr = tl.exp(m - m_new) # (BLOCK_G,) v_ptrs = kv_base + D * stride_kv_d + offs_p[:, None] * stride_kv_p + offs_d[None, :] * stride_kv_d Vp = tl.load(v_ptrs).to(tl.float32) # (PAGE, D) p = tl.where(tok[None, :] < seqlen, p, 0.0) # (BLOCK_G, PAGE) acc_corr = acc * corr[:, None] + tl.sum(p[:, :, None] * Vp[None, :, :], axis=1) acc = acc_corr lsum = corr * lsum + tl.sum(p, axis=1) m = m_new # Reset padding rows every step so they never leak into shared vectors. acc = tl.where(keep[:, None], acc, 0.0) lsum = tl.where(keep, lsum, 0.0) m = tl.where(keep, m, NEG_INF) # Write partials at [b, hkv, split, g]. m_ptr = M_PART + pid_b * stride_m_b + pid_hkv * stride_m_h + pid_split * stride_m_s l_ptr = L_PART + pid_b * stride_l_b + pid_hkv * stride_l_h + pid_split * stride_l_s a_ptr = A_PART + pid_b * stride_a_b + pid_hkv * stride_a_h + pid_split * stride_a_s tl.store(m_ptr + offs_g * stride_m_g, m.to(tl.float32)) tl.store(l_ptr + offs_g * stride_l_g, lsum.to(tl.float32)) tl.store(a_ptr + offs_g[:, None] * stride_a_g + offs_d[None, :] * stride_a_d, acc.to(tl.float32)) @triton.jit def _combine_kernel( M_PART, L_PART, A_PART, OUT, stride_m_b, stride_m_h, stride_m_s, stride_m_g, stride_l_b, stride_l_h, stride_l_s, stride_l_g, stride_a_b, stride_a_h, stride_a_s, stride_a_g, stride_a_d, stride_o_b, stride_o_h, stride_o_d, B, HKV, G, D, num_splits, BLOCK_D: tl.constexpr, BLOCK_G: tl.constexpr, GROUP: tl.constexpr, NEG_INF: tl.constexpr, ): pid_b = tl.program_id(0) pid_hkv = tl.program_id(1) offs_g = tl.arange(0, BLOCK_G) offs_d = tl.arange(0, BLOCK_D) keep = offs_g < GROUP m = tl.full((BLOCK_G,), NEG_INF, dtype=tl.float32) lsum = tl.zeros((BLOCK_G,), dtype=tl.float32) acc = tl.zeros((BLOCK_G, BLOCK_D), dtype=tl.float32) for s in range(0, num_splits): m_s = tl.load(M_PART + pid_b * stride_m_b + pid_hkv * stride_m_h + s * stride_m_s + offs_g * stride_m_g).to(tl.float32) l_s = tl.load(L_PART + pid_b * stride_l_b + pid_hkv * stride_l_h + s * stride_l_s + offs_g * stride_l_g).to(tl.float32) a_s = tl.load(A_PART + pid_b * stride_a_b + pid_hkv * stride_a_h + s * stride_a_s + offs_g[:, None] * stride_a_g + offs_d[None, :] * stride_a_d).to(tl.float32) # Treat empty partials (l_s == 0, i.e. m_s at sentinel) as no-op. m_s = tl.where(l_s > 0, m_s, m) m_new = tl.maximum(m, m_s) w_g = tl.exp(m - m_new) w_s = tl.exp(m_s - m_new) acc = acc * w_g[:, None] + a_s * w_s[:, None] lsum = lsum * w_g + l_s * w_s m = m_new out = acc / lsum[:, None] o_ptr = OUT + pid_b * stride_o_b + pid_hkv * GROUP * stride_o_h tl.store(o_ptr + offs_g[:, None] * stride_o_h + offs_d[None, :] * stride_o_d, out.to(tl.bfloat16), mask=(offs_g[:, None] < GROUP)) class Model(nn.Module): def __init__(self, batch, num_heads, num_kv_heads, head_dim, seq_len, page_size): super().__init__() assert num_heads % num_kv_heads == 0 self.batch = batch self.num_heads = num_heads self.num_kv_heads = num_kv_heads self.head_dim = head_dim self.seq_len = seq_len self.page_size = page_size self.group_size = num_heads // num_kv_heads self.scale = 1.0 / math.sqrt(head_dim) self.register_buffer("_dummy", torch.zeros(1, dtype=torch.bfloat16), persistent=False) def forward(self, query, kv_cache, block_table, seq_lens): B, H, D = query.shape Hkv = self.num_kv_heads G = self.group_size P = self.page_size BLOCK_G = ((G + 15) // 16) * 16 out = torch.empty(B, H, D, dtype=query.dtype, device=query.device) num_pages = int(((seq_lens.max() + P - 1) // P).item()) # FlashDecoding: aim for ~512 concurrent programs to hide memory latency. target = 512 num_splits = max(1, min((target + B * Hkv - 1) // (B * Hkv), num_pages)) pages_per_split = (num_pages + num_splits - 1) // num_splits # Partial buffers (B, Hkv, num_splits, BLOCK_G, [D]). m_part = torch.empty(B, Hkv, num_splits, BLOCK_G, dtype=torch.float32, device=query.device) l_part = torch.empty(B, Hkv, num_splits, BLOCK_G, dtype=torch.float32, device=query.device) a_part = torch.empty(B, Hkv, num_splits, BLOCK_G, D, dtype=torch.float32, device=query.device) _decode_split_kernel[(B, Hkv, num_splits)]( query, kv_cache, block_table, m_part, l_part, a_part, seq_lens, kv_cache.stride(0), kv_cache.stride(1), kv_cache.stride(2), kv_cache.stride(3), query.stride(0), query.stride(1), query.stride(2), m_part.stride(0), m_part.stride(1), m_part.stride(2), m_part.stride(3), l_part.stride(0), l_part.stride(1), l_part.stride(2), l_part.stride(3), a_part.stride(0), a_part.stride(1), a_part.stride(2), a_part.stride(3), a_part.stride(4), block_table.stride(0), block_table.stride(1), self.scale, B, Hkv, G, P, D, num_splits, pages_per_split, num_pages, BLOCK_D=D, BLOCK_G=BLOCK_G, GROUP=G, PAGE=P, NEG_INF=-3.0e4, ) _combine_kernel[(B, Hkv)]( m_part, l_part, a_part, out, m_part.stride(0), m_part.stride(1), m_part.stride(2), m_part.stride(3), l_part.stride(0), l_part.stride(1), l_part.stride(2), l_part.stride(3), a_part.stride(0), a_part.stride(1), a_part.stride(2), a_part.stride(3), a_part.stride(4), out.stride(0), out.stride(1), out.stride(2), B, Hkv, G, D, num_splits, BLOCK_D=D, BLOCK_G=BLOCK_G, GROUP=G, NEG_INF=-3.0e4, ) return out def get_inputs(): B = BATCH H = NUM_HEADS Hkv = NUM_KV_HEADS D = HEAD_DIM L = SEQ_LEN P = PAGE_SIZE pages_per_seq = (L + P - 1) // P total_pages = max(B * pages_per_seq + 8, 64) query = torch.randn(B, H, D, dtype=torch.bfloat16) * 0.1 kv_cache = torch.randn(total_pages, P, Hkv, 2 * D, dtype=torch.bfloat16) * 0.1 perm = torch.randperm(total_pages)[: B * pages_per_seq].reshape(B, pages_per_seq).int() block_table = perm.contiguous() seq_lens = torch.full((B,), L, dtype=torch.int32) return [query, kv_cache, block_table, seq_lens] def get_init_inputs(): return [BATCH, NUM_HEADS, NUM_KV_HEADS, HEAD_DIM, SEQ_LEN, PAGE_SIZE] # --- Shape knobs (overridden by check.py / benchmark.py from shapes.py) ---- BATCH = 8 NUM_HEADS = 32 NUM_KV_HEADS = 8 HEAD_DIM = 128 SEQ_LEN = 1024 PAGE_SIZE = 16