""" paged_attention_v2_triton.py — Triton kernel for PagedAttention V2 on BI-V100 ================================================================================ Replaces the Python partition loop with a single Triton kernel launch. Phase 1 kernel: paged_attn_v2_partition grid = (num_seqs, num_heads, num_partitions) Each program instance computes attention for one (seq, head, partition). Algorithm per instance: 1. Load Q vector for this (seq, head): [head_dim] 2. Load K/V from paged cache for this partition's token range 3. Compute QK^T scores, online softmax max + sum 4. Compute weighted V output 5. Store: tmp_output[seq, head, part, :], exp_sums[seq, head, part], max_logits[seq, head, part] Phase 2 kernel: paged_attn_v2_reduce grid = (num_seqs, num_heads) Each program instance reduces across partitions for one (seq, head). Algorithm: 1. Load max_logits[seq, head, :num_parts] → find global_max 2. Rescale: weights[p] = exp(max[p] - global_max) * sum[p] 3. Normalize and weighted sum of tmp_output SMEM analysis: Phase 1: K tile [BLOCK_N, head_dim] + V tile [BLOCK_N, head_dim] in SMEM At BLOCK_N=64, head_dim=128, fp16: 64*128*2*2 = 32KB ≤ 48KB ✓ Phase 2: No SMEM needed (max_partitions ≈ 200, fits in registers) Deploy: This kernel requires Triton to be functional on BI-V100. patch_enable_triton.py already enables Triton with try/fallback. If Triton works, this kernel replaces the Python V2 for decode. If Triton doesn't work, fall back to paged_attention_v2_pytorch.py. """ import torch import triton import triton.language as tl from typing import Optional @triton.jit def _paged_attn_v2_partition_kernel( # Outputs tmp_output_ptr, # [num_seqs, num_heads, max_num_parts, head_size] exp_sums_ptr, # [num_seqs, num_heads, max_num_parts] max_logits_ptr, # [num_seqs, num_heads, max_num_parts] # Inputs query_ptr, # [num_seqs, num_heads, head_size] key_cache_ptr, # [num_blocks, num_kv_heads, head_size/x, block_size, x] value_cache_ptr, # [num_blocks, num_kv_heads, head_size, block_size] block_tables_ptr, # [num_seqs, max_blocks_per_seq] seq_lens_ptr, # [num_seqs] # Scalars scale, num_kv_heads, block_size, max_blocks_per_seq, max_num_parts, # Strides stride_qt_s, stride_qt_h, stride_qt_d, stride_kc_b, stride_kc_h, stride_kc_dx, stride_kc_bs, stride_kc_x, stride_vc_b, stride_vc_h, stride_vc_d, stride_vc_bs, stride_bt_s, stride_bt_b, stride_to_s, stride_to_h, stride_to_p, stride_to_d, stride_es_s, stride_es_h, stride_es_p, # Constants PARTITION_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr, BLOCK_N: tl.constexpr, # KV tokens processed per inner loop iteration X_PACK: tl.constexpr, # key cache packing factor (16 // element_size) ): """Phase 1: Per-partition attention computation. Each program computes attention for one (seq, head, partition). Iterates over BLOCK_N tokens at a time within the partition. Uses online softmax (Flash Attention style) to compute max, sum, and weighted V. """ seq_idx = tl.program_id(0) head_idx = tl.program_id(1) part_idx = tl.program_id(2) seq_len = tl.load(seq_lens_ptr + seq_idx) # This partition's token range part_start = part_idx * PARTITION_SIZE part_end = tl.minimum(part_start + PARTITION_SIZE, seq_len) if part_start >= seq_len: # This partition is beyond the sequence length — write -inf/0 tl.store(max_logits_ptr + seq_idx * stride_es_s + head_idx * stride_es_h + part_idx * stride_es_p, float('-inf')) tl.store(exp_sums_ptr + seq_idx * stride_es_s + head_idx * stride_es_h + part_idx * stride_es_p, 0.0) return # GQA: map head_idx to kv_head_idx num_queries_per_kv = (tl.program_id(1) + 1) # placeholder — need actual num_heads/num_kv_heads kv_head_idx = head_idx // (stride_qt_h // stride_kc_h) if stride_kc_h > 0 else head_idx # TODO: fix GQA mapping # Load query: [HEAD_DIM] q_offsets = seq_idx * stride_qt_s + head_idx * stride_qt_h + tl.arange(0, HEAD_DIM) * stride_qt_d q = tl.load(query_ptr + q_offsets).to(tl.float32) # Online softmax state m_i = float('-inf') # running max l_i = 0.0 # running sum of exp # Accumulator for weighted V: [HEAD_DIM] acc = tl.zeros([HEAD_DIM], dtype=tl.float32) # Iterate over KV tokens in this partition, BLOCK_N at a time for token_start in range(part_start, part_end, BLOCK_N): token_end = tl.minimum(token_start + BLOCK_N, part_end) n_tokens = token_end - token_start # For each token, find its physical block and offset token_offsets = tl.arange(0, BLOCK_N) valid_mask = token_offsets < n_tokens global_token_ids = token_start + token_offsets block_indices = global_token_ids // block_size within_block_offsets = global_token_ids % block_size # Look up physical block numbers from block_table bt_offsets = seq_idx * stride_bt_s + block_indices * stride_bt_b physical_blocks = tl.load(block_tables_ptr + bt_offsets, mask=valid_mask, other=0) # Load K for these tokens: need to gather from paged cache # K shape: [num_blocks, num_kv_heads, head_size/x, block_size, x] # For each token, load K[physical_block, kv_head, :, within_block_offset, :] # → [BLOCK_N, HEAD_DIM] # Compute QK^T scores for this chunk # scores[n] = sum_d(q[d] * k[n, d]) * scale # This requires loading K values — which is complex with paged layout # TODO: implement the actual paged K gather in Triton # For now, this is a skeleton showing the algorithm structure # --- Placeholder: scores computation --- # In a full implementation, we would: # 1. For each token n in [0, BLOCK_N): # a. physical_block = block_tables[seq, global_token_ids[n] // block_size] # b. offset = global_token_ids[n] % block_size # c. k[n, :] = key_cache[physical_block, kv_head, :, offset, :].reshape(HEAD_DIM) # 2. scores = q @ k.T * scale # 3. Online softmax update # 4. Load V similarly, accumulate weighted V pass # Store results tl.store(max_logits_ptr + seq_idx * stride_es_s + head_idx * stride_es_h + part_idx * stride_es_p, m_i) tl.store(exp_sums_ptr + seq_idx * stride_es_s + head_idx * stride_es_h + part_idx * stride_es_p, l_i) # Store accumulated output out_offsets = (seq_idx * stride_to_s + head_idx * stride_to_h + part_idx * stride_to_p + tl.arange(0, HEAD_DIM) * stride_to_d) tl.store(tmp_output_ptr + out_offsets, acc.to(tmp_output_ptr.dtype.element_ty)) @triton.jit def _paged_attn_v2_reduce_kernel( # Output output_ptr, # [num_seqs, num_heads, head_size] # Inputs tmp_output_ptr, # [num_seqs, num_heads, max_num_parts, head_size] exp_sums_ptr, # [num_seqs, num_heads, max_num_parts] max_logits_ptr, # [num_seqs, num_heads, max_num_parts] seq_lens_ptr, # [num_seqs] # Scalars max_num_parts, # Strides stride_out_s, stride_out_h, stride_out_d, stride_to_s, stride_to_h, stride_to_p, stride_to_d, stride_es_s, stride_es_h, stride_es_p, # Constants PARTITION_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr, MAX_NUM_PARTS: tl.constexpr, ): """Phase 2: Cross-partition reduction. Each program reduces across partitions for one (seq, head). Numerically stable log-sum-exp combination. This corresponds to CCCL's summary_statistics binary_op pattern: combining partial statistics from independent segments. """ seq_idx = tl.program_id(0) head_idx = tl.program_id(1) seq_len = tl.load(seq_lens_ptr + seq_idx) num_parts = (seq_len + PARTITION_SIZE - 1) // PARTITION_SIZE # Load all partition max_logits and exp_sums part_offsets = tl.arange(0, MAX_NUM_PARTS) valid_mask = part_offsets < num_parts ml_base = seq_idx * stride_es_s + head_idx * stride_es_h part_max = tl.load(max_logits_ptr + ml_base + part_offsets * stride_es_p, mask=valid_mask, other=float('-inf')) part_sum = tl.load(exp_sums_ptr + ml_base + part_offsets * stride_es_p, mask=valid_mask, other=0.0) # Global max across partitions global_max = tl.max(part_max, axis=0) # Rescale: weights[p] = exp(max[p] - global_max) * sum[p] rescale = tl.exp(part_max - global_max) * part_sum total = tl.sum(rescale, axis=0) weights = rescale / total # [MAX_NUM_PARTS] # Weighted combination of partition outputs # For each dimension d in HEAD_DIM: # output[d] = sum_p(weights[p] * tmp_output[seq, head, p, d]) for d in range(HEAD_DIM): to_base = seq_idx * stride_to_s + head_idx * stride_to_h + d * stride_to_d part_vals = tl.load(tmp_output_ptr + to_base + part_offsets * stride_to_p, mask=valid_mask, other=0.0) val = tl.sum(weights * part_vals, axis=0) tl.store(output_ptr + seq_idx * stride_out_s + head_idx * stride_out_h + d * stride_out_d, val) def paged_attention_v2_triton( output: torch.Tensor, exp_sums: torch.Tensor, max_logits: torch.Tensor, tmp_output: torch.Tensor, query: torch.Tensor, key_cache: torch.Tensor, value_cache: torch.Tensor, num_kv_heads: int, scale: float, block_tables: torch.Tensor, seq_lens: torch.Tensor, block_size: int, max_seq_len: int, alibi_slopes: Optional[torch.Tensor], kv_cache_dtype: str = "auto", k_scale: float = 1.0, v_scale: float = 1.0, **kwargs, ) -> None: """Triton-based PagedAttention V2. NOTE: The Phase 1 kernel's K/V gather from paged cache is a skeleton. The paged cache layout (key_cache: [blocks, kv_heads, head_dim/x, block_size, x]) requires indirect memory access (gather via block_tables) which is complex in Triton. The Phase 2 reduction kernel is complete. Current status: Phase 1: SKELETON — falls back to PyTorch partition loop Phase 2: COMPLETE — Triton reduction kernel When Phase 1 is complete, this will be a single-launch V2: grid = (num_seqs, num_heads, max_num_partitions) for Phase 1 grid = (num_seqs, num_heads) for Phase 2 """ num_seqs, num_heads, head_size = query.shape max_num_parts = tmp_output.shape[2] PARTITION_SIZE = 512 BLOCK_N = 64 # Must fit in SMEM: BLOCK_N * head_dim * 2B * 2 ≤ 48KB # --- Phase 1: Use PyTorch for now (Triton K/V gather skeleton above) --- # TODO: Complete the Triton Phase 1 kernel with proper paged K/V gather from paged_attention_v2_pytorch import paged_attention_v2_pytorch paged_attention_v2_pytorch( output, exp_sums, max_logits, tmp_output, query, key_cache, value_cache, num_kv_heads, scale, block_tables, seq_lens, block_size, max_seq_len, alibi_slopes, kv_cache_dtype, k_scale, v_scale, ) # Phase 1 writes tmp_output, exp_sums, max_logits # Phase 2 below will re-reduce them (redundant but correct) # --- Phase 2: Triton reduction kernel --- # This replaces the Python einsum reduction with a single Triton launch MAX_NUM_PARTS_CONST = triton.next_power_of_2(max_num_parts) if MAX_NUM_PARTS_CONST > 1024: MAX_NUM_PARTS_CONST = 1024 # Safety cap grid_reduce = (num_seqs, num_heads) _paged_attn_v2_reduce_kernel[grid_reduce]( output, tmp_output, exp_sums, max_logits, seq_lens, max_num_parts, # output strides output.stride(0), output.stride(1), output.stride(2), # tmp_output strides tmp_output.stride(0), tmp_output.stride(1), tmp_output.stride(2), tmp_output.stride(3), # exp_sums strides (same layout as max_logits) exp_sums.stride(0), exp_sums.stride(1), exp_sums.stride(2), # Constants PARTITION_SIZE=PARTITION_SIZE, HEAD_DIM=head_size, MAX_NUM_PARTS=MAX_NUM_PARTS_CONST, )