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

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,