"""Paged attention decode kernel for H100 PCIe (SM90, HBM2e 2.0 TB/s). Single-query decode: each batch element has a query of shape (num_heads, head_dim) and attends over a KV cache of seq_len[b] tokens stored as fixed-size pages in a global pool, with block_table[b] listing which pages belong to batch element b. Key optimisation: processes an entire GQA group (GROUP_SIZE query heads sharing one KV head) per thread block, reusing each K/V load for GROUP_SIZE dot products. Uses 2D tensors throughout for vectorised softmax and accumulation. """ import math import torch import torch.nn as nn import triton import triton.language as tl OP_TYPE = "attention" SUPPORTED_PRECISIONS = ["bf16"] HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"] BATCH = 8 NUM_HEADS = 32 NUM_KV_HEADS = 8 HEAD_DIM = 128 SEQ_LEN = 1024 PAGE_SIZE = 16 @triton.jit def _paged_attention_decode_kernel( Q_ptr, KV_cache_ptr, block_table_ptr, seq_lens_ptr, Out_ptr, scale, stride_q_bt, stride_q_hd, stride_kv_blk, stride_kv_page, stride_kv_head, stride_bt_bt, stride_bt_pg, stride_o_bt, stride_o_hd, HEAD_DIM: tl.constexpr, NUM_KV_HEADS: tl.constexpr, PAGE_SIZE: tl.constexpr, GROUP_SIZE: tl.constexpr, MAX_PAGES: tl.constexpr, ): """Process a GQA group: GROUP_SIZE query heads sharing one KV head. Grid: (batch * num_kv_heads,) Uses 2D tensors (GROUP_SIZE, HEAD_DIM) so that K/V loads are amortised across all heads in the group. """ pid = tl.program_id(0) batch_id = pid // NUM_KV_HEADS kv_head = pid % NUM_KV_HEADS seq_len = tl.load(seq_lens_ptr + batch_id) n_pages = tl.cdiv(seq_len, PAGE_SIZE) d_offs = tl.arange(0, HEAD_DIM) # --- Load Q for all GROUP_SIZE heads as a 2D tensor (GROUP_SIZE, HEAD_DIM) --- q_all = tl.zeros((GROUP_SIZE, HEAD_DIM), dtype=tl.float32) for g in tl.static_range(GROUP_SIZE): head_id = kv_head * GROUP_SIZE + g q_ptrs = Q_ptr + batch_id * stride_q_bt + head_id * stride_q_hd + d_offs q_row = tl.load(q_ptrs).to(tl.float32) # Place row g into the 2D tensor row_mask = (tl.arange(0, GROUP_SIZE) == g)[:, None] q_all = tl.where(row_mask, q_row[None, :], q_all) # --- Online softmax state: (GROUP_SIZE, 1) for m, l; (GROUP_SIZE, HEAD_DIM) for acc --- m_all = tl.full((GROUP_SIZE, 1), -float('inf'), dtype=tl.float32) l_all = tl.zeros((GROUP_SIZE, 1), dtype=tl.float32) acc_all = tl.zeros((GROUP_SIZE, HEAD_DIM), dtype=tl.float32) for page_idx in range(MAX_PAGES): page_valid = page_idx < n_pages page_id = tl.load(block_table_ptr + batch_id * stride_bt_bt + page_idx * stride_bt_pg).to(tl.int64) for tok_idx in range(PAGE_SIZE): global_tok = page_idx * PAGE_SIZE + tok_idx tok_valid = page_valid & (global_tok < seq_len) tok_off = page_id * stride_kv_blk + tok_idx * stride_kv_page + kv_head * stride_kv_head # Load K, V once — shape (HEAD_DIM,) k_ptrs = KV_cache_ptr + tok_off + d_offs k = tl.load(k_ptrs, mask=tok_valid, other=0.0).to(tl.float32) v_ptrs = KV_cache_ptr + tok_off + HEAD_DIM + d_offs v = tl.load(v_ptrs, mask=tok_valid, other=0.0).to(tl.float32) # Dot product for all GROUP_SIZE heads at once # scores: (GROUP_SIZE, 1) scores = tl.sum(q_all * k[None, :], axis=1, keep_dims=True) * scale # Online softmax — vectorised across GROUP_SIZE m_new = tl.maximum(m_all, scores) alpha = tl.exp(m_all - m_new) acc_all = alpha * acc_all # (GROUP_SIZE, 1) * (GROUP_SIZE, HEAD_DIM) -> broadcast l_all = alpha * l_all p = tl.exp(scores - m_new) acc_all = acc_all + p * v[None, :] l_all = l_all + p m_all = m_new # --- Normalise and write output for each head --- acc_all = acc_all / l_all # broadcast: / (GROUP_SIZE, 1) for g in tl.static_range(GROUP_SIZE): head_id = kv_head * GROUP_SIZE + g o_ptrs = Out_ptr + batch_id * stride_o_bt + head_id * stride_o_hd + d_offs # Extract row g from the 2D accumulator row_mask = (tl.arange(0, GROUP_SIZE) == g)[:, None] acc_g = tl.sum(acc_all * row_mask, axis=0) # (HEAD_DIM,) tl.store(o_ptrs, acc_g.to(Out_ptr.dtype.element_ty)) class Model(nn.Module): """Single-query paged attention decode with GQA-grouped Triton kernel.""" def __init__( self, batch: int, num_heads: int, num_kv_heads: int, head_dim: int, seq_len: int, page_size: int, ): super().__init__() assert num_heads % num_kv_heads == 0, "num_heads must be a multiple of num_kv_heads (GQA)" 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: torch.Tensor, kv_cache: torch.Tensor, block_table: torch.Tensor, seq_lens: torch.Tensor, ) -> torch.Tensor: B, H, D = query.shape P = self.page_size G = self.group_size Hkv = self.num_kv_heads out = torch.empty(B, H, D, dtype=query.dtype, device=query.device) max_pages = (self.seq_len + P - 1) // P # Grid: one thread block per (batch, kv_head), each processes G heads grid = (B * Hkv,) _paged_attention_decode_kernel[grid]( query, kv_cache, block_table, seq_lens, out, self.scale, query.stride(0), query.stride(1), kv_cache.stride(0), kv_cache.stride(1), kv_cache.stride(2), block_table.stride(0), block_table.stride(1), out.stride(0), out.stride(1), HEAD_DIM=D, NUM_KV_HEADS=Hkv, PAGE_SIZE=P, GROUP_SIZE=G, MAX_PAGES=max_pages, num_warps=4, num_stages=2, ) return out def get_inputs(): """Build random paged inputs for the current module-level shape knobs.""" 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]