[OPT] Complete Triton V2 Phase 1 — paged K/V gather from prefix_prefill.py pattern
Phase 1 kernel (_paged_attn_v2_partition_kernel) now has complete
paged K/V gather implementation, adapted from prefix_prefill.py:
K gather:
bn = tl.load(block_tables + seq*stride + (token//block_size)*stride)
off_k = bn * stride_kc_b + kv_head * stride_kc_h +
(d//x) * stride_kc_dx + (token%block_size) * stride_kc_bs +
(d%x) * stride_kc_x
k = tl.load(key_cache + off_k, mask=valid)
V gather (simpler layout):
off_v = bn * stride_vc_b + kv_head * stride_vc_h +
d * stride_vc_d + (token%block_size) * stride_vc_bs
Online softmax (Flash Attention pattern):
m_i_new = max(m_i, max(scores))
alpha = exp(m_i - m_i_new)
acc = acc * alpha * l_i / l_i_new + (p/l_i_new * beta) @ V
Key difference from prefix_prefill.py:
- BLOCK_M=1 (decode: 1 query token) vs BLOCK_M>1 (prefill)
- q @ k is dot product [D]•[D,N] → [N], not matrix [M,D]@[D,N] → [M,N]
- head_dim=256 support: BLOCK_N=32 (vs 64 for head_dim=128)
32×256×2×2 = 32KB ≤ 48KB SMEM ✓
Integration: Triton V2 tried first, PyTorch V2 as fallback.
If Triton works on BI-V100: single GPU launch for all partitions
(grid = num_seqs × num_heads × num_partitions = 1 × 24 × 200 = 4800 blocks)
vs PyTorch's 3 bmm launches.
This commit is contained in:
Claude
@@ -6,6 +6,7 @@ WORKDIR /workspace/
|
||||
# Copy all scripts and the V2 module
|
||||
COPY ./qwen3_6_scripts /workspace/qwen3_6_scripts
|
||||
COPY ./paged_attention_v2_pytorch.py /workspace/paged_attention_v2_pytorch.py
|
||||
COPY ./paged_attention_v2_triton.py /workspace/paged_attention_v2_triton.py
|
||||
|
||||
# Run baseline patches (model registration, xformers fallback, tool parser, etc.)
|
||||
RUN cd ./qwen3_6_scripts && ./patch_ops.sh
|
||||
|
||||
@@ -1,39 +1,35 @@
|
||||
"""
|
||||
paged_attention_v2_triton.py — Triton kernel for PagedAttention V2 on BI-V100
|
||||
================================================================================
|
||||
paged_attention_v2_triton.py — Triton PagedAttention V2 for BI-V100
|
||||
=====================================================================
|
||||
|
||||
Replaces the Python partition loop with a single Triton kernel launch.
|
||||
Two-kernel V2 implementation using Triton:
|
||||
Phase 1: _paged_attn_v2_partition — per-partition attention (paged K/V gather)
|
||||
Phase 2: _paged_attn_v2_reduce — cross-partition log-sum-exp reduction
|
||||
|
||||
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]
|
||||
The K/V gather pattern is adapted from prefix_prefill.py (lines 100-170):
|
||||
bn = tl.load(block_tables + seq * stride + (token // block_size) * stride)
|
||||
off_k = bn * stride_kc_b + kv_head * stride_kc_h + (d // x) * stride_kc_dx + ...
|
||||
k = tl.load(key_cache + off_k, mask=...)
|
||||
|
||||
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
|
||||
For decode (BLOCK_M=1), the Q tile is just one vector [HEAD_DIM].
|
||||
The inner loop iterates over BLOCK_N KV tokens per step.
|
||||
Online softmax accumulates (max, sum, weighted_V) across steps.
|
||||
|
||||
After all steps in a partition, we have:
|
||||
max_logits[seq, head, part]: running max
|
||||
exp_sums[seq, head, part]: running exp sum
|
||||
tmp_output[seq, head, part, :]: unnormalized weighted V
|
||||
|
||||
Phase 2 combines partitions using the CCCL summary_statistics pattern:
|
||||
global_max = max(part_maxes)
|
||||
rescaled_sum = sum(exp(part_max - global_max) * part_sum)
|
||||
output = sum(weight[p] * part_output[p])
|
||||
|
||||
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.
|
||||
Phase 1: K tile [BLOCK_N, HEAD_DIM] loaded via gather (no explicit SMEM tile)
|
||||
Triton manages register allocation for tl.load + tl.dot
|
||||
At BLOCK_N=32, HEAD_DIM=256: 32×256 fp16 values in registers = 16KB
|
||||
Phase 2: No SMEM needed (partitions ≈ 200, all in registers)
|
||||
"""
|
||||
|
||||
import torch
|
||||
@@ -45,185 +41,221 @@ 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]
|
||||
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]
|
||||
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
|
||||
scale: tl.float32,
|
||||
num_queries_per_kv: tl.int32,
|
||||
block_size: tl.int32,
|
||||
x_pack: tl.int32, # key_cache packing factor: 16 // sizeof(dtype)
|
||||
# Strides: query [S, H, D]
|
||||
stride_qs: tl.int32, stride_qh: tl.int32, stride_qd: tl.int32,
|
||||
# Strides: key_cache [B, KH, D/X, BS, X]
|
||||
stride_kc_b: tl.int32, stride_kc_h: tl.int32,
|
||||
stride_kc_dx: tl.int32, stride_kc_bs: tl.int32, stride_kc_x: tl.int32,
|
||||
# Strides: value_cache [B, KH, D, BS]
|
||||
stride_vc_b: tl.int32, stride_vc_h: tl.int32,
|
||||
stride_vc_d: tl.int32, stride_vc_bs: tl.int32,
|
||||
# Strides: block_tables [S, MAX_BLOCKS]
|
||||
stride_bt_s: tl.int32, stride_bt_b: tl.int32,
|
||||
# Strides: tmp_output [S, H, P, D]
|
||||
stride_to_s: tl.int32, stride_to_h: tl.int32,
|
||||
stride_to_p: tl.int32, stride_to_d: tl.int32,
|
||||
# Strides: exp_sums / max_logits [S, H, P]
|
||||
stride_es_s: tl.int32, stride_es_h: tl.int32, stride_es_p: tl.int32,
|
||||
# Compile-time 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)
|
||||
BLOCK_N: tl.constexpr,
|
||||
):
|
||||
"""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.
|
||||
"""Phase 1: Per-partition paged attention for decode (BLOCK_M=1).
|
||||
|
||||
Grid: (num_seqs, num_heads, max_num_partitions)
|
||||
Each program instance processes one (seq, head, partition) triple.
|
||||
|
||||
Adapted from prefix_prefill.py's paged K/V gather pattern.
|
||||
Key difference: BLOCK_M=1 (decode has 1 query token per head).
|
||||
"""
|
||||
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
|
||||
# Unused partition — write sentinel values
|
||||
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)
|
||||
|
||||
|
||||
# GQA: map query head → KV head
|
||||
kv_head_idx = head_idx // num_queries_per_kv
|
||||
|
||||
# Load query vector: [HEAD_DIM]
|
||||
offs_d = tl.arange(0, HEAD_DIM)
|
||||
q = tl.load(query_ptr + seq_idx * stride_qs + head_idx * stride_qh + offs_d * stride_qd).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
|
||||
m_i = float('-inf') # running max
|
||||
l_i = 0.0 # running exp sum
|
||||
acc = tl.zeros([HEAD_DIM], dtype=tl.float32) # weighted V accumulator
|
||||
|
||||
# KV token offsets within each BLOCK_N chunk
|
||||
offs_n = tl.arange(0, BLOCK_N)
|
||||
|
||||
# Iterate over BLOCK_N KV tokens at a time
|
||||
for start_n in range(part_start, part_end, BLOCK_N):
|
||||
# Token positions in the sequence
|
||||
token_ids = start_n + offs_n
|
||||
valid_mask = token_ids < part_end
|
||||
|
||||
# === Paged K gather (from prefix_prefill.py pattern) ===
|
||||
# Look up physical block numbers from block_tables
|
||||
block_indices = token_ids // block_size
|
||||
within_block = token_ids % block_size
|
||||
|
||||
# bn: physical block ids [BLOCK_N]
|
||||
bn = tl.load(
|
||||
block_tables_ptr + seq_idx * stride_bt_s + block_indices * stride_bt_b,
|
||||
mask=valid_mask, other=0)
|
||||
|
||||
# K offsets: key_cache[bn, kv_head, d//x, within_block, d%x]
|
||||
# Layout: [num_blocks, num_kv_heads, head_size/x, block_size, x]
|
||||
# off_k: [HEAD_DIM, BLOCK_N] — each column is one token's K vector
|
||||
off_k = (bn[None, :] * stride_kc_b +
|
||||
kv_head_idx * stride_kc_h +
|
||||
(offs_d[:, None] // x_pack) * stride_kc_dx +
|
||||
within_block[None, :] * stride_kc_bs +
|
||||
(offs_d[:, None] % x_pack) * stride_kc_x)
|
||||
|
||||
k = tl.load(key_cache_ptr + off_k, mask=valid_mask[None, :], other=0.0) # [D, N]
|
||||
|
||||
# Scores: q @ k = [1, D] @ [D, N] → [N]
|
||||
# For BLOCK_M=1: this is a dot product per KV token
|
||||
scores = tl.sum(q[:, None] * k, axis=0) * scale # [BLOCK_N]
|
||||
scores = tl.where(valid_mask, scores, float('-inf'))
|
||||
|
||||
# Online softmax update
|
||||
m_ij = tl.max(scores, axis=0) # scalar: max of this chunk
|
||||
m_i_new = tl.maximum(m_i, m_ij)
|
||||
|
||||
alpha = tl.exp(m_i - m_i_new)
|
||||
beta = tl.exp(m_ij - m_i_new)
|
||||
|
||||
p = tl.exp(scores - m_i_new) # [BLOCK_N]
|
||||
l_ij = tl.sum(p, axis=0)
|
||||
|
||||
l_i_new = alpha * l_i + beta * l_ij if l_i > 0 else l_ij
|
||||
|
||||
# === Paged V gather ===
|
||||
# V offsets: value_cache[bn, kv_head, d, within_block]
|
||||
# Layout: [num_blocks, num_kv_heads, head_size, block_size]
|
||||
off_v = (bn[:, None] * stride_vc_b +
|
||||
kv_head_idx * stride_vc_h +
|
||||
offs_d[None, :] * stride_vc_d +
|
||||
within_block[:, None] * stride_vc_bs)
|
||||
|
||||
v = tl.load(value_cache_ptr + off_v, mask=valid_mask[:, None], other=0.0) # [N, D]
|
||||
|
||||
# Update accumulator: acc = (acc * alpha * l_i / l_i_new) + (p @ V * beta / l_i_new)
|
||||
if l_i > 0:
|
||||
acc_scale = l_i / l_i_new * alpha
|
||||
acc = acc * acc_scale
|
||||
|
||||
p_scaled = p / l_i_new * beta # [BLOCK_N]
|
||||
acc += tl.sum(p_scaled[:, None] * v, axis=0) # [HEAD_DIM]
|
||||
|
||||
l_i = l_i_new
|
||||
m_i = m_i_new
|
||||
|
||||
# Store partition 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))
|
||||
|
||||
# Store accumulated output: [HEAD_DIM]
|
||||
out_base = seq_idx * stride_to_s + head_idx * stride_to_h + part_idx * stride_to_p
|
||||
tl.store(tmp_output_ptr + out_base + offs_d * stride_to_d, 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
|
||||
# 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,
|
||||
max_num_parts: tl.int32,
|
||||
# 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,
|
||||
stride_out_s: tl.int32, stride_out_h: tl.int32, stride_out_d: tl.int32,
|
||||
stride_to_s: tl.int32, stride_to_h: tl.int32,
|
||||
stride_to_p: tl.int32, stride_to_d: tl.int32,
|
||||
stride_es_s: tl.int32, stride_es_h: tl.int32, stride_es_p: tl.int32,
|
||||
# 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.
|
||||
"""Phase 2: Cross-partition log-sum-exp reduction.
|
||||
|
||||
Grid: (num_seqs, num_heads)
|
||||
Combines partition results using CCCL summary_statistics pattern.
|
||||
"""
|
||||
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
|
||||
|
||||
# Load partition statistics
|
||||
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,
|
||||
|
||||
es_base = seq_idx * stride_es_s + head_idx * stride_es_h
|
||||
part_max = tl.load(max_logits_ptr + es_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,
|
||||
part_sum = tl.load(exp_sums_ptr + es_base + part_offsets * stride_es_p,
|
||||
mask=valid_mask, other=0.0)
|
||||
|
||||
# Global max across partitions
|
||||
|
||||
# Global max
|
||||
global_max = tl.max(part_max, axis=0)
|
||||
|
||||
# Rescale: weights[p] = exp(max[p] - global_max) * sum[p]
|
||||
|
||||
# Rescale and normalize
|
||||
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)
|
||||
|
||||
# Weighted sum of partition outputs
|
||||
offs_d = tl.arange(0, HEAD_DIM)
|
||||
acc = tl.zeros([HEAD_DIM], dtype=tl.float32)
|
||||
|
||||
for p in range(MAX_NUM_PARTS):
|
||||
if p < num_parts:
|
||||
w = tl.load(max_logits_ptr + es_base + p * stride_es_p) # reload for weight
|
||||
w_rescaled = tl.exp(w - global_max) * tl.load(exp_sums_ptr + es_base + p * stride_es_p) / total
|
||||
|
||||
to_base = seq_idx * stride_to_s + head_idx * stride_to_h + p * stride_to_p
|
||||
part_out = tl.load(tmp_output_ptr + to_base + offs_d * stride_to_d)
|
||||
acc += w_rescaled * part_out.to(tl.float32)
|
||||
|
||||
# Store final output
|
||||
out_base = seq_idx * stride_out_s + head_idx * stride_out_h
|
||||
tl.store(output_ptr + out_base + offs_d * stride_out_d, acc.to(output_ptr.dtype.element_ty))
|
||||
|
||||
|
||||
def paged_attention_v2_triton(
|
||||
@@ -246,59 +278,60 @@ def paged_attention_v2_triton(
|
||||
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
|
||||
"""
|
||||
"""Launch Triton V2 kernels."""
|
||||
num_seqs, num_heads, head_size = query.shape
|
||||
num_queries_per_kv = num_heads // num_kv_heads
|
||||
max_num_parts = tmp_output.shape[2]
|
||||
|
||||
x_pack = key_cache.shape[-1] # packing factor
|
||||
|
||||
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),
|
||||
# BLOCK_N: must fit in SMEM. For decode (BLOCK_M=1), SMEM is dominated by K/V gather.
|
||||
# head_dim=256: BLOCK_N=32 → 32×256×2 = 16KB per tile (K or V)
|
||||
# head_dim=128: BLOCK_N=64 → 64×128×2 = 16KB per tile
|
||||
BLOCK_N = 32 if head_size > 128 else 64
|
||||
|
||||
# Phase 1: partition attention
|
||||
num_partitions = (max_seq_len + PARTITION_SIZE - 1) // PARTITION_SIZE
|
||||
grid_phase1 = (num_seqs, num_heads, num_partitions)
|
||||
|
||||
_paged_attn_v2_partition_kernel[grid_phase1](
|
||||
tmp_output, exp_sums, max_logits,
|
||||
query, key_cache, value_cache, block_tables, seq_lens,
|
||||
scale, num_queries_per_kv, block_size, x_pack,
|
||||
# query strides
|
||||
query.stride(0), query.stride(1), query.stride(2),
|
||||
# key_cache strides
|
||||
key_cache.stride(0), key_cache.stride(1), key_cache.stride(2),
|
||||
key_cache.stride(3), key_cache.stride(4),
|
||||
# value_cache strides
|
||||
value_cache.stride(0), value_cache.stride(1), value_cache.stride(2),
|
||||
value_cache.stride(3),
|
||||
# block_tables strides
|
||||
block_tables.stride(0), block_tables.stride(1),
|
||||
# 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 strides
|
||||
exp_sums.stride(0), exp_sums.stride(1), exp_sums.stride(2),
|
||||
# Constants
|
||||
PARTITION_SIZE=PARTITION_SIZE,
|
||||
HEAD_DIM=head_size,
|
||||
BLOCK_N=BLOCK_N,
|
||||
)
|
||||
|
||||
# Phase 2: cross-partition reduction
|
||||
MAX_NUM_PARTS_CONST = triton.next_power_of_2(max_num_parts)
|
||||
if MAX_NUM_PARTS_CONST > 1024:
|
||||
MAX_NUM_PARTS_CONST = 1024
|
||||
|
||||
grid_phase2 = (num_seqs, num_heads)
|
||||
_paged_attn_v2_reduce_kernel[grid_phase2](
|
||||
output,
|
||||
tmp_output, exp_sums, max_logits, seq_lens,
|
||||
max_num_parts,
|
||||
output.stride(0), output.stride(1), output.stride(2),
|
||||
tmp_output.stride(0), tmp_output.stride(1), tmp_output.stride(2), tmp_output.stride(3),
|
||||
exp_sums.stride(0), exp_sums.stride(1), exp_sums.stride(2),
|
||||
PARTITION_SIZE=PARTITION_SIZE,
|
||||
HEAD_DIM=head_size,
|
||||
MAX_NUM_PARTS=MAX_NUM_PARTS_CONST,
|
||||
)
|
||||
|
||||
@@ -32,7 +32,8 @@ VLLM_ROOTS = [
|
||||
"/usr/local/corex/lib64/python3/dist-packages/vllm",
|
||||
]
|
||||
|
||||
V2_MODULE = "paged_attention_v2_pytorch.py"
|
||||
V2_MODULE_PYTORCH = "paged_attention_v2_pytorch.py"
|
||||
V2_MODULE_TRITON = "paged_attention_v2_triton.py"
|
||||
|
||||
|
||||
def find_vllm_root():
|
||||
@@ -50,7 +51,15 @@ def patch_custom_ops(vllm_root):
|
||||
content = f.read()
|
||||
|
||||
# Add import at the top (after existing imports)
|
||||
import_line = "from vllm.paged_attention_v2_pytorch import paged_attention_v2_pytorch"
|
||||
import_line = "# Try Triton V2 (single-launch, GPU-parallel) first; PyTorch V2 as fallback
|
||||
try:
|
||||
from vllm.paged_attention_v2_triton import paged_attention_v2_triton as _v2_impl
|
||||
_V2_BACKEND = "triton"
|
||||
except Exception:
|
||||
from vllm.paged_attention_v2_pytorch import paged_attention_v2_pytorch as _v2_impl
|
||||
_V2_BACKEND = "pytorch"
|
||||
import logging
|
||||
logging.getLogger("vllm").info(f"PagedAttention V2 backend: {_V2_BACKEND}")"
|
||||
if import_line in content:
|
||||
print(" [skip] V2 import already present")
|
||||
else:
|
||||
@@ -77,7 +86,7 @@ def patch_custom_ops(vllm_root):
|
||||
blocksparse_head_sliding_step: int = 0,
|
||||
) -> None:
|
||||
# BI-V100: PyTorch V2 implementation (replaces NotImplementedError)
|
||||
paged_attention_v2_pytorch(
|
||||
_v2_impl(
|
||||
out, exp_sum, max_logits, tmp_out,
|
||||
query, key_cache, value_cache,
|
||||
num_kv_heads, scale, block_tables, seq_lens,
|
||||
|
||||
Reference in New Issue
Block a user