"""Custom paged-attention decode kernel for B200 (SM100, HBM3e). Flash-decoding: one CTA per (batch, kv_head, seq_partition). Each KV element is read exactly once (GQA query heads in a group share the load). Split-K along the sequence gives enough parallelism to fill all 148 SMs even at batch=4. A second small kernel does the flash reduction across partitions. No SDPA / vLLM / FlashInfer calls -- pure Triton. """ 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"] # --- 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 _NUM_SMS = 148 # B200 # =========================================================================== # Decode kernel: produces per-partition partial (o, m, l) or final output. # Grid: (B, Hkv, num_partitions) # =========================================================================== @triton.autotune( configs=[ triton.Config({"BLOCK_N": 32}, num_warps=4, num_stages=4), triton.Config({"BLOCK_N": 32}, num_warps=8, num_stages=5), triton.Config({"BLOCK_N": 64}, num_warps=4, num_stages=4), triton.Config({"BLOCK_N": 64}, num_warps=8, num_stages=4), triton.Config({"BLOCK_N": 64}, num_warps=8, num_stages=5), triton.Config({"BLOCK_N": 64}, num_warps=8, num_stages=6), triton.Config({"BLOCK_N": 128}, num_warps=8, num_stages=3), triton.Config({"BLOCK_N": 128}, num_warps=8, num_stages=4), triton.Config({"BLOCK_N": 128}, num_warps=8, num_stages=5), triton.Config({"BLOCK_N": 128}, num_warps=16, num_stages=3), triton.Config({"BLOCK_N": 128}, num_warps=16, num_stages=4), ], key=["G", "D", "P", "WRITE_FINAL"], ) @triton.jit def _decode_kernel( Q_ptr, KV_ptr, BT_ptr, SL_ptr, OUT_ptr, PO_ptr, PM_ptr, PL_ptr, q_sb, q_sh, kv_sblk, kv_spg, kv_sh, bt_sb, out_sb, out_sh, po_sb, po_sh, po_sp, pm_sb, pm_sh, pm_sp, pl_sb, pl_sh, pl_sp, sm_scale, PART_SEQ, G: tl.constexpr, D: tl.constexpr, HKV: tl.constexpr, P: tl.constexpr, KVD: tl.constexpr, MQ: tl.constexpr, BLOCK_N: tl.constexpr, WRITE_FINAL: tl.constexpr, ): pid_b = tl.program_id(0) pid_h = tl.program_id(1) pid_p = tl.program_id(2) Lb = tl.load(SL_ptr + pid_b) start = pid_p * PART_SEQ # exclusive end of this partition (clamped to seq len); masks the inner loop # so interior partitions don't read into their neighbour's token range. part_end = tl.minimum((pid_p + 1) * PART_SEQ, Lb) # ---- Q tile (MQ, D) bf16; rows >= G are padded with 0 ---- offs_q = tl.arange(0, MQ) g_mask = offs_q < G q_rows = pid_h * G + offs_q offs_d = tl.arange(0, D) q_ptrs = Q_ptr + pid_b * q_sb + q_rows[:, None] * q_sh + offs_d[None, :] Q = tl.load(q_ptrs, mask=g_mask[:, None], other=0.0) # ---- accumulators ---- m_i = tl.full([MQ], float("-inf"), dtype=tl.float32) l_i = tl.full([MQ], 0.0, dtype=tl.float32) o_i = tl.zeros([MQ, D], dtype=tl.float32) if start < Lb: for start_n in range(0, PART_SEQ, BLOCK_N): n = start + start_n + tl.arange(0, BLOCK_N) mask_n = n < part_end page = n // P off_pg = n % P bid = tl.load(BT_ptr + pid_b * bt_sb + page, mask=mask_n, other=0) phys = bid * P + off_pg base = phys * (HKV * KVD) + pid_h * KVD kv_ptrs = KV_ptr + base[:, None] + offs_d[None, :] K = tl.load(kv_ptrs, mask=mask_n[:, None], other=0.0, eviction_policy="evict_first") qk = tl.dot(Q, tl.trans(K)) * sm_scale # (MQ, BLOCK_N) qk = tl.where(mask_n[None, :], qk, float("-inf")) m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) qk = qk - m_ij[:, None] p = tl.exp(qk) # (MQ, BLOCK_N) alpha = tl.exp(m_i - m_ij) l_ij = tl.sum(p, axis=1) V = tl.load(kv_ptrs + D, mask=mask_n[:, None], other=0.0, eviction_policy="evict_first") o_i = o_i * alpha[:, None] + tl.dot(p.to(tl.bfloat16), V) l_i = l_i * alpha + l_ij m_i = m_ij if WRITE_FINAL: o = o_i / l_i[:, None] out_ptrs = OUT_ptr + pid_b * out_sb + q_rows[:, None] * out_sh + offs_d[None, :] tl.store(out_ptrs, o.to(tl.bfloat16), mask=g_mask[:, None]) else: po_ptrs = PO_ptr + pid_b * po_sb + pid_h * po_sh + pid_p * po_sp + offs_q[:, None] * D + offs_d[None, :] tl.store(po_ptrs, o_i, mask=g_mask[:, None]) pm_ptrs = PM_ptr + pid_b * pm_sb + pid_h * pm_sh + pid_p * pm_sp + offs_q tl.store(pm_ptrs, m_i, mask=g_mask) pl_ptrs = PL_ptr + pid_b * pl_sb + pid_h * pl_sh + pid_p * pl_sp + offs_q tl.store(pl_ptrs, l_i, mask=g_mask) # =========================================================================== # Reduce kernel: flash-reduce per-partition partials -> final output. # Grid: (B, H) -- one CTA per query head for high occupancy on small batches. # =========================================================================== @triton.jit def _reduce_kernel( OUT_ptr, PO_ptr, PM_ptr, PL_ptr, out_sb, out_sh, po_sb, po_sh, po_sp, pm_sb, pm_sh, pm_sp, pl_sb, pl_sh, pl_sp, NP, G: tl.constexpr, D: tl.constexpr, BLOCK_P: tl.constexpr, ): pid_b = tl.program_id(0) pid_h = tl.program_id(1) # query head index in [0, H) hkv = pid_h // G g = pid_h % G offs_d = tl.arange(0, D) offs_p = tl.arange(0, BLOCK_P) p_mask = offs_p < NP # partial max for this head: (BLOCK_P,) pm_ptrs = PM_ptr + pid_b * pm_sb + hkv * pm_sh + offs_p * pm_sp + g pm = tl.load(pm_ptrs, mask=p_mask, other=float("-inf")) m_g = tl.max(pm, axis=0) factor = tl.exp(pm - m_g) # (BLOCK_P,) pl_ptrs = PL_ptr + pid_b * pl_sb + hkv * pl_sh + offs_p * pl_sp + g pl = tl.load(pl_ptrs, mask=p_mask, other=0.0) l_g = tl.sum(factor * pl, axis=0) # partial output for this head: (BLOCK_P, D) po_ptrs = (PO_ptr + pid_b * po_sb + hkv * po_sh + offs_p[:, None] * po_sp + g * D + offs_d[None, :]) po = tl.load(po_ptrs, mask=p_mask[:, None], other=0.0) o_g = tl.sum(factor[:, None] * po, axis=0) # (D,) out = o_g / l_g out_ptrs = OUT_ptr + pid_b * out_sb + pid_h * out_sh + offs_d tl.store(out_ptrs, out.to(tl.bfloat16)) def _next_pow2(x): p = 1 while p < x: p <<= 1 return p 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) import os as _os _mq_override = _os.environ.get("KBH_MQ") if _mq_override: self.MQ = int(_mq_override) else: self.MQ = _next_pow2(max(self.group_size, 1)) # ---- choose split-K parallelism ---- # Decode prefers a modest number of fairly large CTAs (per-CTA softmax # overhead + split-K reduce cost penalize very high partition counts). # Empirically ~1 wave (B*Hkv*npart ~= 128) is the sweet spot across the # shape sweep. Going higher only helps the pure-read ceiling, not decode. import os base = batch * num_kv_heads target_ctas = 128 npart = (target_ctas + base - 1) // base npart = max(npart, 1) min_tokens_per_part = 64 max_npart = max(1, seq_len // min_tokens_per_part) npart = min(npart, max_npart) npart = max(npart, 1) env_np = os.environ.get("KBH_NPART") if env_np: npart = int(env_np) self.npart = npart self.part_seq = (seq_len + npart - 1) // npart self.register_buffer("_dummy", torch.zeros(1, dtype=torch.bfloat16), persistent=False) # Preallocate output and partial buffers (static addresses for CUDA graphs). self._out = torch.empty(batch, num_heads, head_dim, dtype=torch.bfloat16, device="cuda") if npart > 1: self._po = torch.empty(batch, num_kv_heads, npart, self.MQ, head_dim, dtype=torch.float32, device="cuda") self._pm = torch.empty(batch, num_kv_heads, npart, self.MQ, dtype=torch.float32, device="cuda") self._pl = torch.empty(batch, num_kv_heads, npart, self.MQ, dtype=torch.float32, device="cuda") else: self._po = self._pm = self._pl = None self._graph = None self._graph_key = None def _launch(self, out, 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 KVD = 2 * D MQ = self.MQ npart = self.npart part_seq = self.part_seq grid = (B, Hkv, npart) if npart == 1: _decode_kernel[grid]( query, kv_cache, block_table, seq_lens, out, None, None, None, query.stride(0), query.stride(1), kv_cache.stride(0), kv_cache.stride(1), kv_cache.stride(2), block_table.stride(0), out.stride(0), out.stride(1), 0, 0, 0, 0, 0, 0, 0, 0, 0, self.scale, part_seq, G=G, D=D, HKV=Hkv, P=P, KVD=KVD, MQ=MQ, WRITE_FINAL=True, ) else: po, pm, pl = self._po, self._pm, self._pl _decode_kernel[grid]( query, kv_cache, block_table, seq_lens, out, po, pm, pl, query.stride(0), query.stride(1), kv_cache.stride(0), kv_cache.stride(1), kv_cache.stride(2), block_table.stride(0), out.stride(0), out.stride(1), po.stride(0), po.stride(1), po.stride(2), pm.stride(0), pm.stride(1), pm.stride(2), pl.stride(0), pl.stride(1), pl.stride(2), self.scale, part_seq, G=G, D=D, HKV=Hkv, P=P, KVD=KVD, MQ=MQ, WRITE_FINAL=False, ) _reduce_kernel[(B, self.num_heads)]( out, po, pm, pl, out.stride(0), out.stride(1), po.stride(0), po.stride(1), po.stride(2), pm.stride(0), pm.stride(1), pm.stride(2), pl.stride(0), pl.stride(1), pl.stride(2), npart, G=G, D=D, BLOCK_P=_next_pow2(npart), ) def forward(self, query, kv_cache, block_table, seq_lens): out = self._out key = (query.data_ptr(), kv_cache.data_ptr(), block_table.data_ptr(), seq_lens.data_ptr()) # If we're already inside an external CUDA-graph capture, just launch # eagerly so the kernel launches get captured by the outer graph. capturing = torch.cuda.is_current_stream_capturing() if not capturing and self._graph is not None and key == self._graph_key: self._graph.replay() return out # Eager launch (also resolves Triton autotune before capture). self._launch(out, query, kv_cache, block_table, seq_lens) # Try to capture a graph for this set of input addresses so subsequent # calls replay with near-zero launch overhead. Falls back to eager. if not capturing and self._graph is None: try: torch.cuda.synchronize() g = torch.cuda.CUDAGraph() with torch.cuda.graph(g): self._launch(out, query, kv_cache, block_table, seq_lens) self._graph = g self._graph_key = key except Exception: self._graph = None 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]