diff --git a/Dockerfile b/Dockerfile index 87f3b43d..b2ea93ac 100644 --- a/Dockerfile +++ b/Dockerfile @@ -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 diff --git a/paged_attention_v2_triton.py b/paged_attention_v2_triton.py index 6645d0a4..36e7fc3c 100644 --- a/paged_attention_v2_triton.py +++ b/paged_attention_v2_triton.py @@ -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, ) diff --git a/qwen3_6_scripts/patch_paged_attention_v2.py b/qwen3_6_scripts/patch_paged_attention_v2.py index 0618628b..d5ed2947 100644 --- a/qwen3_6_scripts/patch_paged_attention_v2.py +++ b/qwen3_6_scripts/patch_paged_attention_v2.py @@ -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,