[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.

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

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
+The K/V gather pattern is adapted from prefix_prefill.py (lines 100-170):
-  grid = (num_seqs, num_heads, num_partitions)
+  bn = tl.load(block_tables + seq * stride + (token // block_size) * stride)
-  Each program instance computes attention for one (seq, head, partition).
+  off_k = bn * stride_kc_b + kv_head * stride_kc_h + (d // x) * stride_kc_dx + ...
-  
+  k = tl.load(key_cache + off_k, mask=...)
-  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 
+For decode (BLOCK_M=1), the Q tile is just one vector [HEAD_DIM].
-  grid = (num_seqs, num_heads)
+The inner loop iterates over BLOCK_N KV tokens per step.
-  Each program instance reduces across partitions for one (seq, head).
+Online softmax accumulates (max, sum, weighted_V) across steps.
-  
+
-  Algorithm:
+After all steps in a partition, we have:
-    1. Load max_logits[seq, head, :num_parts] → find global_max
+  max_logits[seq, head, part]: running max
-    2. Rescale: weights[p] = exp(max[p] - global_max) * sum[p]
+  exp_sums[seq, head, part]: running exp sum
-    3. Normalize and weighted sum of tmp_output
+  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
+  Phase 1: K tile [BLOCK_N, HEAD_DIM] loaded via gather (no explicit SMEM tile)
-    At BLOCK_N=64, head_dim=128, fp16: 64*128*2*2 = 32KB ≤ 48KB ✓
+           Triton manages register allocation for tl.load + tl.dot
-  Phase 2: No SMEM needed (max_partitions ≈ 200, fits in registers)
+           At BLOCK_N=32, HEAD_DIM=256: 32×256 fp16 values in registers = 16KB
-
+  Phase 2: No SMEM needed (partitions ≈ 200, all 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
@@ -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]
+    tmp_output_ptr,     # [num_seqs, num_heads, max_num_parts, head_size]
-    exp_sums_ptr,      # [num_seqs, num_heads, max_num_parts]
+    exp_sums_ptr,       # [num_seqs, num_heads, max_num_parts]
-    max_logits_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]
+    query_ptr,          # [num_seqs, num_heads, head_size]
-    key_cache_ptr,     # [num_blocks, num_kv_heads, head_size/x, block_size, x]
+    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]
+    value_cache_ptr,    # [num_blocks, num_kv_heads, head_size, block_size]
-    block_tables_ptr,  # [num_seqs, max_blocks_per_seq]
+    block_tables_ptr,   # [num_seqs, max_blocks_per_seq]
-    seq_lens_ptr,      # [num_seqs]
+    seq_lens_ptr,       # [num_seqs]
     # Scalars
-    scale,
+    scale: tl.float32,
-    num_kv_heads,
+    num_queries_per_kv: tl.int32,
-    block_size,
+    block_size: tl.int32,
-    max_blocks_per_seq,
+    x_pack: tl.int32,             # key_cache packing factor: 16 // sizeof(dtype)
-    max_num_parts,
+    # Strides: query [S, H, D]
-    # Strides
+    stride_qs: tl.int32, stride_qh: tl.int32, stride_qd: tl.int32,
-    stride_qt_s, stride_qt_h, stride_qt_d,
+    # Strides: key_cache [B, KH, D/X, BS, X]
-    stride_kc_b, stride_kc_h, stride_kc_dx, stride_kc_bs, stride_kc_x,
+    stride_kc_b: tl.int32, stride_kc_h: tl.int32,
-    stride_vc_b, stride_vc_h, stride_vc_d, stride_vc_bs,
+    stride_kc_dx: tl.int32, stride_kc_bs: tl.int32, stride_kc_x: tl.int32,
-    stride_bt_s, stride_bt_b,
+    # Strides: value_cache [B, KH, D, BS]
-    stride_to_s, stride_to_h, stride_to_p, stride_to_d,
+    stride_vc_b: tl.int32, stride_vc_h: tl.int32,
-    stride_es_s, stride_es_h, stride_es_p,
+    stride_vc_d: tl.int32, stride_vc_bs: tl.int32,
-    # Constants
+    # 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
+    BLOCK_N: tl.constexpr,
-    X_PACK: tl.constexpr,        # key cache packing factor (16 // element_size)
 ):
-    """Phase 1: Per-partition attention computation.
+    """Phase 1: Per-partition paged attention for decode (BLOCK_M=1).
-    
+
-    Each program computes attention for one (seq, head, partition).
+    Grid: (num_seqs, num_heads, max_num_partitions)
-    Iterates over BLOCK_N tokens at a time within the partition.
+    Each program instance processes one (seq, head, partition) triple.
-    Uses online softmax (Flash Attention style) to compute max, sum, and weighted V.
+
+    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
+    # GQA: map query head → KV head
-    num_queries_per_kv = (tl.program_id(1) + 1)  # placeholder — need actual num_heads/num_kv_heads
+    kv_head_idx = head_idx // num_queries_per_kv
-    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 vector: [HEAD_DIM]
-    # Load query: [HEAD_DIM]
+    offs_d = tl.arange(0, 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 + seq_idx * stride_qs + head_idx * stride_qh + offs_d * stride_qd).to(tl.float32)
-    q = tl.load(query_ptr + q_offsets).to(tl.float32)
+
-    
     # Online softmax state
-    m_i = float('-inf')  # running max
+    m_i = float('-inf')       # running max
-    l_i = 0.0           # running sum of exp
+    l_i = 0.0                 # running exp sum
-    # Accumulator for weighted V: [HEAD_DIM]
+    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)  # weighted V accumulator
-    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)
+
-    
+    # KV token offsets within each BLOCK_N chunk
-    # Iterate over KV tokens in this partition, BLOCK_N at a time
+    offs_n = tl.arange(0, BLOCK_N)
-    for token_start in range(part_start, part_end, BLOCK_N):
+
-        token_end = tl.minimum(token_start + BLOCK_N, part_end)
+    # Iterate over BLOCK_N KV tokens at a time
-        n_tokens = token_end - token_start
+    for start_n in range(part_start, part_end, BLOCK_N):
-        
+        # Token positions in the sequence
-        # For each token, find its physical block and offset
+        token_ids = start_n + offs_n
-        token_offsets = tl.arange(0, BLOCK_N)
+        valid_mask = token_ids < part_end
-        valid_mask = token_offsets < n_tokens
+
-        
+        # === Paged K gather (from prefix_prefill.py pattern) ===
-        global_token_ids = token_start + token_offsets
+        # Look up physical block numbers from block_tables
-        block_indices = global_token_ids // block_size
+        block_indices = token_ids // block_size
-        within_block_offsets = global_token_ids % block_size
+        within_block = token_ids % block_size
-        
+
-        # Look up physical block numbers from block_table
+        # bn: physical block ids [BLOCK_N]
-        bt_offsets = seq_idx * stride_bt_s + block_indices * stride_bt_b
+        bn = tl.load(
-        physical_blocks = tl.load(block_tables_ptr + bt_offsets, mask=valid_mask, other=0)
+            block_tables_ptr + seq_idx * stride_bt_s + block_indices * stride_bt_b,
-        
+            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]
+        # K offsets: key_cache[bn, kv_head, d//x, within_block, d%x]
-        # For each token, load K[physical_block, kv_head, :, within_block_offset, :]
+        # Layout: [num_blocks, num_kv_heads, head_size/x, block_size, x]
-        # → [BLOCK_N, HEAD_DIM]
+        # off_k: [HEAD_DIM, BLOCK_N] — each column is one token's K vector
-        
+        off_k = (bn[None, :] * stride_kc_b +
-        # Compute QK^T scores for this chunk
+                 kv_head_idx * stride_kc_h +
-        # scores[n] = sum_d(q[d] * k[n, d]) * scale
+                 (offs_d[:, None] // x_pack) * stride_kc_dx +
-        # This requires loading K values — which is complex with paged layout
+                 within_block[None, :] * stride_kc_bs +
-        # TODO: implement the actual paged K gather in Triton
+                 (offs_d[:, None] % x_pack) * stride_kc_x)
-        # For now, this is a skeleton showing the algorithm structure
+
-        
+        k = tl.load(key_cache_ptr + off_k, mask=valid_mask[None, :], other=0.0)  # [D, N]
-        # --- Placeholder: scores computation ---
+
-        # In a full implementation, we would:
+        # Scores: q @ k = [1, D] @ [D, N] → [N]
-        # 1. For each token n in [0, BLOCK_N):
+        # For BLOCK_M=1: this is a dot product per KV token
-        #    a. physical_block = block_tables[seq, global_token_ids[n] // block_size]
+        scores = tl.sum(q[:, None] * k, axis=0) * scale  # [BLOCK_N]
-        #    b. offset = global_token_ids[n] % block_size
+        scores = tl.where(valid_mask, scores, float('-inf'))
-        #    c. k[n, :] = key_cache[physical_block, kv_head, :, offset, :].reshape(HEAD_DIM)
+
-        # 2. scores = q @ k.T * scale
+        # Online softmax update
-        # 3. Online softmax update
+        m_ij = tl.max(scores, axis=0)  # scalar: max of this chunk
-        # 4. Load V similarly, accumulate weighted V
+        m_i_new = tl.maximum(m_i, m_ij)
-        pass
+
-    
+        alpha = tl.exp(m_i - m_i_new)
-    # Store results
+        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
+    # Store accumulated output: [HEAD_DIM]
-    out_offsets = (seq_idx * stride_to_s + head_idx * stride_to_h + 
+    out_base = seq_idx * stride_to_s + head_idx * stride_to_h + part_idx * stride_to_p
-                   part_idx * stride_to_p + tl.arange(0, HEAD_DIM) * stride_to_d)
+    tl.store(tmp_output_ptr + out_base + offs_d * stride_to_d, acc.to(tmp_output_ptr.dtype.element_ty))
-    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  
+    # 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_out_s: tl.int32, stride_out_h: tl.int32, stride_out_d: tl.int32,
-    stride_to_s, stride_to_h, stride_to_p, stride_to_d,
+    stride_to_s: tl.int32, stride_to_h: tl.int32,
-    stride_es_s, stride_es_h, stride_es_p,
+    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.
+    """Phase 2: Cross-partition log-sum-exp reduction.
-    
+
-    Each program reduces across partitions for one (seq, head).
+    Grid: (num_seqs, num_heads)
-    Numerically stable log-sum-exp combination.
+    Combines partition results using CCCL summary_statistics pattern.
-    
-    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
+    # 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
+    es_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,
+    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
+    # Weighted sum of partition outputs
-    # For each dimension d in HEAD_DIM:
+    offs_d = tl.arange(0, HEAD_DIM)
-    #   output[d] = sum_p(weights[p] * tmp_output[seq, head, p, d])
+    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)
-    for d in range(HEAD_DIM):
+
-        to_base = seq_idx * stride_to_s + head_idx * stride_to_h + d * stride_to_d
+    for p in range(MAX_NUM_PARTS):
-        part_vals = tl.load(tmp_output_ptr + to_base + part_offsets * stride_to_p,
+        if p < num_parts:
-                           mask=valid_mask, other=0.0)
+            w = tl.load(max_logits_ptr + es_base + p * stride_es_p)  # reload for weight
-        val = tl.sum(weights * part_vals, axis=0)
+            w_rescaled = tl.exp(w - global_max) * tl.load(exp_sums_ptr + es_base + p * stride_es_p) / total
-        tl.store(output_ptr + seq_idx * stride_out_s + head_idx * stride_out_h + d * stride_out_d,
+
-                 val)
+            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.
+    """Launch Triton V2 kernels."""
-    
-    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
+    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
+    # 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)
-    # --- Phase 1: Use PyTorch for now (Triton K/V gather skeleton above) ---
+    # head_dim=128: BLOCK_N=64 → 64×128×2 = 16KB per tile
-    # TODO: Complete the Triton Phase 1 kernel with proper paged K/V gather
+    BLOCK_N = 32 if head_size > 128 else 64
-    from paged_attention_v2_pytorch import paged_attention_v2_pytorch
+
-    paged_attention_v2_pytorch(
+    # Phase 1: partition attention
-        output, exp_sums, max_logits, tmp_output,
+    num_partitions = (max_seq_len + PARTITION_SIZE - 1) // PARTITION_SIZE
-        query, key_cache, value_cache,
+    grid_phase1 = (num_seqs, num_heads, num_partitions)
-        num_kv_heads, scale, block_tables, seq_lens,
+
-        block_size, max_seq_len, alibi_slopes,
+    _paged_attn_v2_partition_kernel[grid_phase1](
-        kv_cache_dtype, k_scale, v_scale,
+        tmp_output, exp_sums, max_logits,
-    )
+        query, key_cache, value_cache, block_tables, seq_lens,
-    # Phase 1 writes tmp_output, exp_sums, max_logits
+        scale, num_queries_per_kv, block_size, x_pack,
-    # Phase 2 below will re-reduce them (redundant but correct)
+        # query strides
-    
+        query.stride(0), query.stride(1), query.stride(2),
-    # --- Phase 2: Triton reduction kernel ---
+        # key_cache strides
-    # This replaces the Python einsum reduction with a single Triton launch
+        key_cache.stride(0), key_cache.stride(1), key_cache.stride(2),
-    MAX_NUM_PARTS_CONST = triton.next_power_of_2(max_num_parts)
+        key_cache.stride(3), key_cache.stride(4),
-    if MAX_NUM_PARTS_CONST > 1024:
+        # value_cache strides
-        MAX_NUM_PARTS_CONST = 1024  # Safety cap
+        value_cache.stride(0), value_cache.stride(1), value_cache.stride(2),
-    
+        value_cache.stride(3),
-    grid_reduce = (num_seqs, num_heads)
+        # block_tables strides
-    _paged_attn_v2_reduce_kernel[grid_reduce](
+        block_tables.stride(0), block_tables.stride(1),
-        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 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,
     )

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,