"""Paged-attention decode (single query step) — custom Triton kernels. Design (FlashDecoding-style split-KV, single fused kernel): Grid: (batch * num_kv_heads * SPLIT). Each program streams the paged K/V rows for one (batch, kv-head, seq-split) and computes attention for ALL G query heads of that GQA group together, so the KV bytes are pulled from DRAM exactly once. Online (flash-style) softmax runs in fp32 in the log2 domain. Each program writes a normalized partial (G, D) plus (m, l) to fp32 scratch, then bumps a per-(b, kv-head) counter with release semantics. The program that observes the counter reaching SPLIT (the last finisher for that group) re-reads all split partials with acquire semantics, merges them with the standard exp2(m_s - m) reweighting and writes the bf16 output. Unlike a second merge kernel, this merges each group as soon as ITS splits finish (no global barrier across groups) and costs no extra kernel launch. The finisher resets the counter so the kernel is replay-safe under CUDA graphs. With SPLIT == 1 the partial/atomic path is skipped entirely and the program writes the bf16 result directly. KV cache layout: (num_blocks, page_size, num_kv_heads, 2*head_dim), last dim packs [K | V]. Token t of sequence b lives at page block_table[b, t // P], slot t % P. The launches are captured into a per-input-pointer CUDA graph after a warmup call so the timed loop pays ~zero Python launch overhead. """ import math import torch import torch.nn as nn import triton import triton.language as tl LOG2E = 1.4426950408889634 @triton.jit def _paged_decode_fused( q_ptr, kv_ptr, bt_ptr, sl_ptr, po_ptr, pm_ptr, pl_ptr, cnt_ptr, out_ptr, qk_scale, stride_qb, stride_qh, stride_kvb, stride_btb, stride_out_b, stride_out_h, H: tl.constexpr, Hkv: tl.constexpr, G: tl.constexpr, D: tl.constexpr, P: tl.constexpr, SPLIT: tl.constexpr, PADDED_G: tl.constexpr, PADDED_S: tl.constexpr, BLOCK_T: tl.constexpr, ): pid = tl.program_id(0) nhkv_split = Hkv * SPLIT b = pid // nhkv_split r = pid % nhkv_split hkv = r // SPLIT split = r % SPLIT L = tl.load(sl_ptr + b) chunk = tl.cdiv(L, SPLIT) start = split * chunk end = tl.minimum(start + chunk, L) g_offs = tl.arange(0, PADDED_G) d_offs = tl.arange(0, D) g_mask = g_offs < G # q for the whole GQA group: (G, D), zero-padded to PADDED_G rows. q = tl.load( q_ptr + b * stride_qb + (hkv * G + g_offs[:, None]) * stride_qh + d_offs[None, :], mask=g_mask[:, None], other=0.0, ) m_i = tl.full([PADDED_G], float("-inf"), tl.float32) l_i = tl.zeros([PADDED_G], tl.float32) acc = tl.zeros([PADDED_G, D], tl.float32) # Unmasked fast path for full blocks; masked tail handles the remainder. n_full = ((end - start) // BLOCK_T) * BLOCK_T for t0 in range(start, start + n_full, BLOCK_T): offs_t = t0 + tl.arange(0, BLOCK_T) pages = tl.load(bt_ptr + b * stride_btb + offs_t // P) row_base = ( kv_ptr + pages * stride_kvb + (offs_t % P) * (Hkv * 2 * D) + hkv * (2 * D) ) k = tl.load(row_base[:, None] + d_offs[None, :]) v = tl.load(row_base[:, None] + D + d_offs[None, :]) qk = tl.dot(q, tl.trans(k)) * qk_scale m_new = tl.maximum(m_i, tl.max(qk, 1)) alpha = tl.exp2(m_i - m_new) p = tl.exp2(qk - m_new[:, None]) l_i = l_i * alpha + tl.sum(p, 1) acc = acc * alpha[:, None] + tl.dot(p.to(tl.bfloat16), v) m_i = m_new if n_full < end - start: for t0 in range(start + n_full, end, BLOCK_T): offs_t = t0 + tl.arange(0, BLOCK_T) tok_mask = offs_t < end pages = tl.load(bt_ptr + b * stride_btb + offs_t // P, mask=tok_mask, other=0) row_base = ( kv_ptr + pages * stride_kvb + (offs_t % P) * (Hkv * 2 * D) + hkv * (2 * D) ) k = tl.load(row_base[:, None] + d_offs[None, :], mask=tok_mask[:, None], other=0.0) v = tl.load(row_base[:, None] + D + d_offs[None, :], mask=tok_mask[:, None], other=0.0) qk = tl.dot(q, tl.trans(k)) * qk_scale qk = tl.where(tok_mask[None, :], qk, float("-inf")) m_new = tl.maximum(m_i, tl.max(qk, 1)) alpha = tl.exp2(m_i - m_new) p = tl.exp2(qk - m_new[:, None]) l_i = l_i * alpha + tl.sum(p, 1) acc = acc * alpha[:, None] + tl.dot(p.to(tl.bfloat16), v) m_i = m_new # Normalize this split's output. safe_l = tl.where(l_i > 0.0, l_i, 1.0) o = tl.where(l_i[:, None] > 0.0, acc / safe_l[:, None], 0.0) if SPLIT == 1: # No cross-split reduction needed: write the result directly. out_idx = b * stride_out_b + (hkv * G + g_offs[:, None]) * stride_out_h + d_offs[None, :] tl.store(out_ptr + out_idx, o.to(tl.bfloat16), mask=g_mask[:, None]) else: # Publish partials, then signal; last finisher merges for the group. grp = b * Hkv + hkv po_idx = ((grp * G + g_offs[:, None]) * SPLIT + split) * D + d_offs[None, :] tl.store(po_ptr + po_idx, o, mask=g_mask[:, None]) ml_idx = grp * G * SPLIT + g_offs * SPLIT + split tl.store(pm_ptr + ml_idx, m_i, mask=g_mask) tl.store(pl_ptr + ml_idx, l_i, mask=g_mask) done = tl.atomic_add(cnt_ptr + grp, 1, sem="acq_rel") if done == SPLIT - 1: # Reset for the next launch (stream-ordered kernels make this # safe; CUDA-graph replays included). tl.store(cnt_ptr + grp, 0) s_offs = tl.arange(0, PADDED_S) s_mask = s_offs < SPLIT ml_base = grp * G * SPLIT m = tl.load(pm_ptr + ml_base + g_offs[:, None] * SPLIT + s_offs[None, :], mask=g_mask[:, None] & s_mask[None, :], other=float("-inf")) l = tl.load(pl_ptr + ml_base + g_offs[:, None] * SPLIT + s_offs[None, :], mask=g_mask[:, None] & s_mask[None, :], other=0.0) m_max = tl.max(m, 1) w = tl.exp2(m - m_max[:, None]) * l denom = tl.sum(w, 1) po_base = grp * G * SPLIT * D mo = tl.zeros([PADDED_G, D], tl.float32) for s in tl.static_range(SPLIT): w_s = tl.sum(tl.where(s_offs == s, w, 0.0), 1) po_s = tl.load(po_ptr + po_base + g_offs[:, None] * SPLIT * D + s * D + d_offs[None, :], mask=g_mask[:, None], other=0.0) mo += w_s[:, None] * po_s mo = mo / denom[:, None] out_idx = b * stride_out_b + (hkv * G + g_offs[:, None]) * stride_out_h + d_offs[None, :] tl.store(out_ptr + out_idx, mo.to(tl.bfloat16), mask=g_mask[:, None]) # Hand-tuned on the RTX PRO 6000 (188 SMs, GDDR7). Keyed by the shape tuple. # Values are (split, block_t, num_warps, num_stages). _CONFIG_TABLE = { (8, 32, 8, 128, 1024, 16): (4, 64, 4, 4), (32, 32, 8, 128, 2048, 16): (1, 64, 4, 3), (4, 64, 8, 128, 4096, 16): (8, 64, 4, 3), (16, 32, 8, 128, 1535, 16): (2, 64, 4, 3), (8, 16, 4, 64, 2000, 16): (8, 64, 4, 3), } def _pick_config(B: int, H: int, Hkv: int, D: int, L: int, P: int): """(split, block_t, num_warps, num_stages) per problem shape.""" cfg = _CONFIG_TABLE.get((B, H, Hkv, D, L, P)) if cfg is not None: return cfg # Fallback heuristic for unseen shapes: ~256 programs, long-ish chunks. base = B * Hkv split = max(1, 256 // base) split = min(split, max(1, (L + 127) // 128)) block_t = 64 num_warps = 8 if D <= 64 else 4 num_stages = 3 return split, block_t, num_warps, num_stages 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.qk_scale = self.scale * LOG2E self.split, self.block_t, self.num_warps, self.num_stages = _pick_config( batch, num_heads, num_kv_heads, head_dim, seq_len, page_size) self.padded_g = max(16, triton.next_power_of_2(self.group_size)) self.padded_s = triton.next_power_of_2(self.split) dev = torch.device("cuda") # fp32 partial scratch: po (B, Hkv, G, SPLIT, D); pm/pl (B, Hkv, G, SPLIT) po = torch.empty(batch * num_kv_heads * self.group_size * self.split * head_dim, dtype=torch.float32, device=dev) pm = torch.empty(batch * num_kv_heads * self.group_size * self.split, dtype=torch.float32, device=dev) pl = torch.empty_like(pm) cnt = torch.zeros(batch * num_kv_heads, dtype=torch.int32, device=dev) out = torch.empty(batch, num_heads, head_dim, dtype=torch.bfloat16, device=dev) self.register_buffer("_po", po, persistent=False) self.register_buffer("_pm", pm, persistent=False) self.register_buffer("_pl", pl, persistent=False) self.register_buffer("_cnt", cnt, persistent=False) self.register_buffer("_out", out, persistent=False) # Parity with reference's non-persistent dummy buffer. self.register_buffer("_dummy", torch.zeros(1, dtype=torch.bfloat16), persistent=False) # CUDA-graph cache keyed by input pointers. The timed loop reuses the # same tensors, so replay removes Python/Triton launch overhead. self._graphs: dict = {} self._seen: set = set() # Single-entry fast-path cache (object identity -> graph/replay). self._tq = self._tkv = self._tbt = self._tsl = None self._freplay = None self._out_ref = None def __call__(self, query, kv_cache, block_table, seq_lens): # Bypass nn.Module's hook/grad dispatch. Hot path: the timed loop reuses # the same tensor objects, so an identity hit replays the captured graph # with minimal Python cost (no extra frame, no *args packing). if ( query is self._tq and kv_cache is self._tkv and block_table is self._tbt and seq_lens is self._tsl ): self._freplay() return self._out_ref return self.forward(query, kv_cache, block_table, seq_lens) def forward(self, query, kv_cache, block_table, seq_lens): return self._forward_slow(query, kv_cache, block_table, seq_lens) def _forward_slow(self, query, kv_cache, block_table, seq_lens): key = ( query.data_ptr(), kv_cache.data_ptr(), block_table.data_ptr(), seq_lens.data_ptr(), ) graph = self._graphs.get(key) if graph is not None: self._set_fast(graph, query, kv_cache, block_table, seq_lens) graph.replay() return self._out if key in self._seen and not torch.cuda.is_current_stream_capturing(): # Second sighting: kernels are compiled/warm, capture the graph. # Replay stays correct even if the allocator recycled these # addresses for fresh same-shape tensors: the captured kernels just # re-read whatever the input addresses currently hold. try: graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): self._forward_impl(query, kv_cache, block_table, seq_lens) if len(self._graphs) < 8: self._graphs[key] = graph self._set_fast(graph, query, kv_cache, block_table, seq_lens) graph.replay() return self._out except Exception: pass self._seen.add(key) return self._forward_impl(query, kv_cache, block_table, seq_lens) def _set_fast(self, graph, query, kv_cache, block_table, seq_lens): self._tq, self._tkv, self._tbt, self._tsl = query, kv_cache, block_table, seq_lens self._freplay = graph.replay self._out_ref = self._out def _forward_impl(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 SPLIT = self.split q = query if query.is_contiguous() else query.contiguous() out = self._out _paged_decode_fused[(B * Hkv * SPLIT,)]( q, kv_cache, block_table, seq_lens, self._po, self._pm, self._pl, self._cnt, out, self.qk_scale, q.stride(0), q.stride(1), kv_cache.stride(0), block_table.stride(0), out.stride(0), out.stride(1), H=H, Hkv=Hkv, G=G, D=D, P=P, SPLIT=SPLIT, PADDED_G=self.padded_g, PADDED_S=self.padded_s, BLOCK_T=self.block_t, num_warps=self.num_warps, num_stages=self.num_stages, ) return out # --- Shape knobs (mirrors reference.py) -------------------------------------- BATCH = 8 NUM_HEADS = 32 NUM_KV_HEADS = 8 HEAD_DIM = 128 SEQ_LEN = 1024 PAGE_SIZE = 16 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]