project_6/qwen3_6_scripts/patch_head256_triton.py

"""
patch_head256_triton.py — Enable Triton prefill for head_dim=256
=================================================================

Qwen3.6-35B-A3B uses head_dim=256 (confirmed: text_cfg.head_dim=256,
num_heads=24, num_kv_heads=4, GQA ratio=6).

Current problem:
  prefix_prefill.py: BLOCK = 64 (set by patch_triton_tuning.py)
  SMEM needed: BLOCK_N × head_dim × sizeof(fp16) × 2 (K+V tiles)
  = 64 × 256 × 2 × 2 = 64KB > 48KB (BI-V100 SMEM limit)
  → Triton kernel CANNOT launch at BLOCK=64 for head_dim=256.

  xformers.py: _run_sdpa_fallback triggers for head_size > 128.
  This is a Python for-loop with Q-tiling — orders of magnitude slower.

Fix:
  1. In prefix_prefill.py launcher, use BLOCK based on head_dim:
     head_dim ≤ 128: BLOCK = 64, NUM_WARPS = 4  (32KB SMEM, fits)
     head_dim = 256: BLOCK = 32, NUM_WARPS = 4  (32KB SMEM, fits)
     head_dim > 256: BLOCK = 16, NUM_WARPS = 2  (16KB SMEM, fits)
  
  2. Triton BLOCK=32 means more kernel launches per sequence but
     each launch uses only 32KB SMEM — well within 48KB limit.
     32×256×2×2 = 32KB ≤ 48KB ✓

  3. Also optimize _run_sdpa_fallback _Q_CHUNK:
     Current: 256 (same for all head dims)
     For head_dim=256: memory = _Q_CHUNK × seq_len × H × 4 bytes (float32)
     At Q_CHUNK=256, seq_len=100K, H=24: 256×100K×24×4 ≈ 2.3GB
     Better: _Q_CHUNK=128 for head_dim=256 → 1.2GB (safer for OOM)

Deploy: python3 qwen3_6_scripts/patch_head256_triton.py
  Must run AFTER patch_ops.sh and patch_triton_tuning.py
"""

import os

PREFIX_PREFILL_PATHS = [
    "/usr/local/corex/lib/python3/dist-packages/vllm/attention/ops/prefix_prefill.py",
    "/usr/local/corex/lib64/python3/dist-packages/vllm/attention/ops/prefix_prefill.py",
]

XFORMERS_PATHS = [
    "/usr/local/corex/lib/python3/dist-packages/vllm/attention/backends/xformers.py",
    "/usr/local/corex/lib64/python3/dist-packages/vllm/attention/backends/xformers.py",
]

# --- Patch 1: prefix_prefill.py BLOCK selection based on head_dim ---

# The current patch_triton_tuning.py sets:
#   BLOCK = 64
#   NUM_WARPS = 4
# We need to make BLOCK depend on head_dim:

OLD_BLOCK_SETTING = """\
        # BI-V100 optimization (patch_triton_tuning.py):
        # BLOCK=64: SMEM constrains BLOCK_N≤64 for head_dim=128
        # NUM_WARPS=4: fewer warps → more blocks/SM → better occupancy
        BLOCK = 64
        NUM_WARPS = 4"""

NEW_BLOCK_SETTING = """\
        # BI-V100: BLOCK must fit in 48KB SMEM.
        # SMEM = BLOCK_N × head_dim × sizeof(fp16) × 2 (K+V tiles)
        # head_dim=128: BLOCK=64 → 64×128×2×2 = 32KB ✓
        # head_dim=256: BLOCK=32 → 32×256×2×2 = 32KB ✓ (BLOCK=64 → 64KB overflow!)
        # head_dim>256: BLOCK=16 → fallback
        Lk = q.shape[-1]
        if Lk <= 128:
            BLOCK = 64
            NUM_WARPS = 4
        elif Lk <= 256:
            BLOCK = 32
            NUM_WARPS = 4
        else:
            BLOCK = 16
            NUM_WARPS = 2"""

# Alternative: if patch_triton_tuning.py hasn't run yet, patch the original
OLD_BLOCK_ORIGINAL = """\
        BLOCK = 128 if current_platform.has_device_capability(80) else 64
        NUM_WARPS = 8"""

NEW_BLOCK_FROM_ORIGINAL = NEW_BLOCK_SETTING


# --- Patch 2: _run_sdpa_fallback Q_CHUNK for head_dim=256 ---

OLD_Q_CHUNK = "        _Q_CHUNK = 256"
NEW_Q_CHUNK = """\
        # Adapt Q chunk size to head_dim to control memory:
        # head_dim=128: 256 × seq_len × H × 4B → manageable
        # head_dim=256: halve to 128 to avoid OOM on long sequences
        _Q_CHUNK = 128 if self.head_size > 128 else 256"""


def patch_prefix_prefill():
    for path in PREFIX_PREFILL_PATHS:
        if not os.path.exists(path):
            continue
        with open(path, "r") as f:
            content = f.read()
        
        changed = False
        if "head_dim=256: BLOCK=32" in content:
            print(f"  [skip] {path}: head_dim-aware BLOCK already present")
            return True
        
        if OLD_BLOCK_SETTING in content:
            content = content.replace(OLD_BLOCK_SETTING, NEW_BLOCK_SETTING, 1)
            changed = True
            print(f"  [ok] Replaced fixed BLOCK=64 with head_dim-dependent selection")
        elif OLD_BLOCK_ORIGINAL in content:
            content = content.replace(OLD_BLOCK_ORIGINAL, NEW_BLOCK_FROM_ORIGINAL, 1)
            changed = True
            print(f"  [ok] Replaced original BLOCK selection with head_dim-dependent version")
        else:
            print(f"  [warn] Neither BLOCK anchor found in {path}")
        
        # Also need to move Lk computation before BLOCK selection
        # Currently Lk is computed AFTER BLOCK is set (line ~750)
        # We need it before. Check if Lk is already available:
        if "Lk = q.shape[-1]" in content and "Lk, Lk, Lv" in content:
            # Lk is computed later — we duplicate the computation for BLOCK selection
            # This is safe because q.shape[-1] doesn't change
            print(f"  [note] Lk computed early for BLOCK selection + later for kernel args")
        
        if changed:
            with open(path, "w") as f:
                f.write(content)
            print(f"  Written: {path}")
        return changed
    
    print("  [error] prefix_prefill.py not found")
    return False


def patch_xformers_sdpa():
    for path in XFORMERS_PATHS:
        if not os.path.exists(path):
            continue
        with open(path, "r") as f:
            content = f.read()
        
        if "head_size > 128 else 256" in content:
            print(f"  [skip] {path}: adaptive Q_CHUNK already present")
            return True
        
        if OLD_Q_CHUNK in content:
            content = content.replace(OLD_Q_CHUNK, NEW_Q_CHUNK, 1)
            with open(path, "w") as f:
                f.write(content)
            print(f"  [ok] {path}: _Q_CHUNK now adapts to head_dim")
            return True
        else:
            print(f"  [warn] _Q_CHUNK anchor not found in {path}")
    
    print("  [error] xformers.py not found")
    return False


def main():
    print("=== patch_head256_triton: Enable Triton for head_dim=256 ===")
    print(f"  Qwen3.6: head_dim=256, num_heads=24, num_kv_heads=4, GQA=6")
    print()
    
    print("--- Patch 1: prefix_prefill.py BLOCK selection ---")
    patch_prefix_prefill()
    
    print("\n--- Patch 2: xformers Q_CHUNK for head_dim=256 ---")
    patch_xformers_sdpa()
    
    print("\nSMEM budget at BLOCK=32, head_dim=256:")
    print(f"  K tile: 32 × 256 × 2B = 16KB")
    print(f"  V tile: 32 × 256 × 2B = 16KB")
    print(f"  Total: 32KB ≤ 48KB ✓")
    print("\nDone.")


if __name__ == "__main__":
    main()