From 821c59500d481c835f1142b07daa42a3e56ab627 Mon Sep 17 00:00:00 2001 From: dylanyunlon Date: Wed, 5 Aug 2026 08:39:54 +0000 Subject: [PATCH] =?UTF-8?q?[CLEANUP]=20Remove=2013=20dead=20patch=20script?= =?UTF-8?q?s=20=E2=80=94=20only=201=20remains=20(transformers=20registrati?= =?UTF-8?q?on)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Removed (replaced by full-file cp in patch_ops.sh): - patch_model_runner.py → replaced by model_runner.py (1932 lines) - patch_xformers_sdpa_seq.py → replaced by xformers.py (901 lines) - patch_xformers_sdpa_seq_kernel.py → was unused - patch_xformers_sdpa_batch.py → was unused - patch_xformers_sdpa_batch_kernel.py → was unused - patch_vllm_qwen3_5.py → replaced by registry.py (455 lines) - patch_vllm_tool_parser.py → replaced by tool_parsers_init.py - patch_enable_triton.py → was unused - patch_head256_triton.py → was unused - patch_ixformer_native.py → was unused - patch_paged_attention_v2.py → was unused - patch_triton_tuning.py → was unused - patch_vectorized_decode.py → was unused Remaining: patch_transformers_qwen3_5.py (1 script, unavoidable — modifies pip-installed transformers which is version-specific) Architecture: 13 blind string-replace scripts → 0. All base modifications are now full-file replacements with complete source context. --- qwen3_6_scripts/patch_enable_triton.py | 195 ----------- qwen3_6_scripts/patch_head256_triton.py | 182 ---------- qwen3_6_scripts/patch_ixformer_native.py | 304 ----------------- qwen3_6_scripts/patch_model_runner.py | 78 ----- qwen3_6_scripts/patch_paged_attention_v2.py | 192 ----------- qwen3_6_scripts/patch_triton_tuning.py | 91 ----- qwen3_6_scripts/patch_vectorized_decode.py | 72 ---- qwen3_6_scripts/patch_vllm_qwen3_5.py | 76 ----- qwen3_6_scripts/patch_vllm_tool_parser.py | 79 ----- qwen3_6_scripts/patch_xformers_sdpa_batch.py | 192 ----------- .../patch_xformers_sdpa_batch_kernel.py | 191 ----------- qwen3_6_scripts/patch_xformers_sdpa_seq.py | 321 ------------------ .../patch_xformers_sdpa_seq_kernel.py | 181 ---------- 13 files changed, 2154 deletions(-) delete mode 100644 qwen3_6_scripts/patch_enable_triton.py delete mode 100644 qwen3_6_scripts/patch_head256_triton.py delete mode 100644 qwen3_6_scripts/patch_ixformer_native.py delete mode 100644 qwen3_6_scripts/patch_model_runner.py delete mode 100644 qwen3_6_scripts/patch_paged_attention_v2.py delete mode 100644 qwen3_6_scripts/patch_triton_tuning.py delete mode 100644 qwen3_6_scripts/patch_vectorized_decode.py delete mode 100644 qwen3_6_scripts/patch_vllm_qwen3_5.py delete mode 100644 qwen3_6_scripts/patch_vllm_tool_parser.py delete mode 100644 qwen3_6_scripts/patch_xformers_sdpa_batch.py delete mode 100644 qwen3_6_scripts/patch_xformers_sdpa_batch_kernel.py delete mode 100644 qwen3_6_scripts/patch_xformers_sdpa_seq.py delete mode 100644 qwen3_6_scripts/patch_xformers_sdpa_seq_kernel.py diff --git a/qwen3_6_scripts/patch_enable_triton.py b/qwen3_6_scripts/patch_enable_triton.py deleted file mode 100644 index ea64e5c1..00000000 --- a/qwen3_6_scripts/patch_enable_triton.py +++ /dev/null @@ -1,195 +0,0 @@ -""" -patch_enable_triton.py — Enable Triton kernels on BI-V100 with safety fallback -================================================================================ - -The baseline disables Triton entirely (HAS_TRITON = False) because the default -kernel configuration hangs BI-V100. But Triton 2.3.1 IS installed in the image. - -Strategy: - 1. Set HAS_TRITON = True so prefix_prefill.py is imported - 2. Patch prefix_prefill.py with conservative tile sizes (BLOCK=64, NUM_WARPS=4) - 3. Add a timeout-protected first-call test in forward_prefix: - - Try Triton kernel with 1-second timeout - - If it hangs or errors, permanently fall back to PyTorch path - - Log the result so we know which path is active - -This is the key performance unlock: - PyTorch fallback: Python for-loop, ~20 tokens/sec on prefill - Triton kernel: GPU-parallel Flash Attention, potentially 10-50x faster - -Risk mitigation: - - If Triton still hangs at BLOCK=64/NUM_WARPS=4, the timeout catches it - - All functional tests still pass (same math, different implementation) - - The fallback is the exact same _forward_prefix_pytorch from baseline - -Deploy: python3 qwen3_6_scripts/patch_enable_triton.py - Must run AFTER patch_ops.sh (which deploys paged_attn.py) - Must run AFTER patch_triton_tuning.py (which sets BLOCK=64, NUM_WARPS=4) -""" - -import os - -# --- 1. Enable HAS_TRITON --- - -TRITON_IMPORT_PATH = "/usr/local/corex/lib/python3/dist-packages/vllm/triton_utils/importing.py" -TRITON_IMPORT_PATHS = [ - TRITON_IMPORT_PATH, - "/usr/local/corex/lib64/python3/dist-packages/vllm/triton_utils/importing.py", -] - -OLD_TRITON = "HAS_TRITON = False" -NEW_TRITON = """\ -# BI-V100: Triton 2.3.1 is present. Enable it with conservative tile sizes. -# If Triton kernels hang, the timeout in paged_attn.py will catch it. -try: - import triton - HAS_TRITON = True -except ImportError: - HAS_TRITON = False""" - - -def patch_triton_import(): - for path in TRITON_IMPORT_PATHS: - if not os.path.exists(path): - continue - with open(path, "r") as f: - content = f.read() - if "HAS_TRITON = True" in content: - print(f" [skip] {path}: HAS_TRITON already True") - return True - if OLD_TRITON in content: - content = content.replace(OLD_TRITON, NEW_TRITON, 1) - with open(path, "w") as f: - f.write(content) - print(f" [ok] {path}: HAS_TRITON = False → True (with import guard)") - return True - print(" [error] importing.py not found") - return False - - -# --- 2. Patch paged_attn.py forward_prefix to try Triton with fallback --- - -PAGED_ATTN_PATH = "/usr/local/corex/lib/python3/dist-packages/vllm/attention/ops/paged_attn.py" - -# The patched paged_attn.py (from patch_ops.sh) has: -# def forward_prefix(...): -# return PagedAttention._forward_prefix_pytorch(...) -# -# We replace it with a try-Triton-first version: - -OLD_FORWARD_PREFIX = """\ - @staticmethod - def forward_prefix( - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - kv_cache_dtype: str, - key_cache: torch.Tensor, - value_cache: torch.Tensor, - block_tables: torch.Tensor, - query_start_loc: torch.Tensor, - seq_lens_tensor: torch.Tensor, - context_lens: torch.Tensor, - max_query_len: int, - alibi_slopes: Optional[torch.Tensor], - sliding_window: Optional[int], - k_scale: float, - v_scale: float, - ) -> torch.Tensor: - # NOTE: The Triton context_attention_fwd kernel hangs on Iluvatar - # BI-V100 hardware (same class of issue as cudnnFlashAttnForward). - # Use a pure-PyTorch fallback that reads the paged KV cache directly. - return PagedAttention._forward_prefix_pytorch( - query, key, value, - key_cache, value_cache, - block_tables, query_start_loc, - seq_lens_tensor, context_lens, - )""" - -NEW_FORWARD_PREFIX = """\ - # Triton prefill: try once, fall back permanently if it fails - _triton_prefill_ok = None # None=untested, True=works, False=failed - - @staticmethod - def forward_prefix( - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - kv_cache_dtype: str, - key_cache: torch.Tensor, - value_cache: torch.Tensor, - block_tables: torch.Tensor, - query_start_loc: torch.Tensor, - seq_lens_tensor: torch.Tensor, - context_lens: torch.Tensor, - max_query_len: int, - alibi_slopes: Optional[torch.Tensor], - sliding_window: Optional[int], - k_scale: float, - v_scale: float, - ) -> torch.Tensor: - # Try Triton kernel if available and not known to fail - if PagedAttention._triton_prefill_ok is not False: - try: - from vllm.triton_utils import HAS_TRITON - if HAS_TRITON: - from vllm.attention.ops.prefix_prefill import context_attention_fwd - output = torch.empty_like(query) - context_attention_fwd( - query, key, value, output, kv_cache_dtype, - key_cache, value_cache, block_tables, - query_start_loc[:-1], seq_lens_tensor, context_lens, - max_query_len, k_scale, v_scale, - alibi_slopes, sliding_window, - ) - if PagedAttention._triton_prefill_ok is None: - print("[paged_attn] Triton prefill kernel: SUCCESS", flush=True) - PagedAttention._triton_prefill_ok = True - return output - except Exception as e: - print(f"[paged_attn] Triton prefill failed: {type(e).__name__}: {e}", - flush=True) - print("[paged_attn] Falling back to PyTorch prefill permanently", flush=True) - PagedAttention._triton_prefill_ok = False - - # PyTorch fallback (same as baseline) - return PagedAttention._forward_prefix_pytorch( - query, key, value, - key_cache, value_cache, - block_tables, query_start_loc, - seq_lens_tensor, context_lens, - )""" - - -def patch_paged_attn(): - if not os.path.exists(PAGED_ATTN_PATH): - print(f" [error] {PAGED_ATTN_PATH} not found") - return False - with open(PAGED_ATTN_PATH, "r") as f: - content = f.read() - if "_triton_prefill_ok" in content: - print(f" [skip] {PAGED_ATTN_PATH}: already has Triton try/fallback") - return True - if OLD_FORWARD_PREFIX in content: - content = content.replace(OLD_FORWARD_PREFIX, NEW_FORWARD_PREFIX, 1) - with open(PAGED_ATTN_PATH, "w") as f: - f.write(content) - print(f" [ok] {PAGED_ATTN_PATH}: added Triton try/fallback in forward_prefix") - return True - print(f" [warn] {PAGED_ATTN_PATH}: forward_prefix anchor not found") - return False - - -def main(): - print("=== patch_enable_triton: Enable Triton with safety fallback ===") - print("\n--- Step 1: Enable HAS_TRITON ---") - patch_triton_import() - print("\n--- Step 2: Triton try/fallback in forward_prefix ---") - patch_paged_attn() - print("\nDone. On first prefill request:") - print(" - If Triton works at BLOCK=64/NUM_WARPS=4 → 10-50x prefill speedup") - print(" - If Triton hangs/errors → auto-fallback to PyTorch (same as baseline)") - - -if __name__ == "__main__": - main() diff --git a/qwen3_6_scripts/patch_head256_triton.py b/qwen3_6_scripts/patch_head256_triton.py deleted file mode 100644 index 4123b3a6..00000000 --- a/qwen3_6_scripts/patch_head256_triton.py +++ /dev/null @@ -1,182 +0,0 @@ -""" -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() diff --git a/qwen3_6_scripts/patch_ixformer_native.py b/qwen3_6_scripts/patch_ixformer_native.py deleted file mode 100644 index 9c523a78..00000000 --- a/qwen3_6_scripts/patch_ixformer_native.py +++ /dev/null @@ -1,304 +0,0 @@ -""" -patch_ixformer_native.py — Enable ixformer native kernels on BI-V100 -===================================================================== - -Hardware-verified fixes (2026-07-31): - -1. V1 paged_attention: head_mapping must be Tensor, not int. - Verified: V1 with Tensor head_mapping matches manual attention (diff < 0.001). - Performance: 0.034ms (256 tok), 0.272ms (8192 tok). - -2. V2 paged_attention: native kernel EXISTS (vllm_single_query_cached_kv_attention_v2) - but produces INCORRECT output (diff=1.28 vs V1, norm mismatch). - The native V2 kernel expects different cache layout [B,H,bs,d] and even with - correct conversion, the output doesn't match V1 on the same data. - STATUS: Keep Python V2 fallback (paged_attention_v2_pytorch.py) for seq > 8192. - TODO: Investigate V2 native kernel parameter semantics. - -3. flash_attn_func: WORKS with head_dim=256, GQA. - ixf_F.flash_attn_func(q, k, v, causal=True) produces correct output. - This should replace the Python _run_sdpa_fallback for prefill. - -4. Triton: installed but vllm can't find it (path mismatch). - Fix: symlink + sys.path insertion. -""" - -import os -import sys -import shutil - -VLLM_ROOT = "/usr/local/corex/lib64/python3/dist-packages/vllm" -CUSTOM_OPS_PATH = os.path.join(VLLM_ROOT, "_custom_ops.py") - - -def patch_v1_head_mapping(): - """Fix V1: convert head_mapping from int to Tensor.""" - with open(CUSTOM_OPS_PATH, "r") as f: - content = f.read() - - old_v1 = '''def paged_attention_v1( - output, - query, - key_cache, - value_cache, - head_mapping, - scale, - block_tables, - context_lens, - block_size, - max_context_len, - alibi_slopes=None, - kv_cache_dtype=None, -): - return ixf_F.vllm_single_query_cached_kv_attention( - output, - query, - key_cache, - value_cache, - head_mapping, - scale, - block_tables, - context_lens, - block_size, - max_context_len, - alibi_slopes, - )''' - - new_v1 = '''def paged_attention_v1( - output, - query, - key_cache, - value_cache, - head_mapping, - scale, - block_tables, - context_lens, - block_size, - max_context_len, - alibi_slopes=None, - kv_cache_dtype=None, -): - # BI-V100: ixformer requires head_mapping as Tensor, not int. - # Verified: V1 with Tensor matches manual attention (max diff < 0.001). - if isinstance(head_mapping, int): - num_kv_heads = head_mapping - num_heads = query.shape[1] - num_queries_per_kv = num_heads // num_kv_heads - head_mapping = torch.repeat_interleave( - torch.arange(num_kv_heads, dtype=torch.int32, device=query.device), - num_queries_per_kv) - return ixf_F.vllm_single_query_cached_kv_attention( - output, - query, - key_cache, - value_cache, - head_mapping, - scale, - block_tables, - context_lens, - block_size, - max_context_len, - alibi_slopes, - )''' - - if "isinstance(head_mapping, int)" in content: - print(" [skip] V1 head_mapping fix already applied") - return True - if old_v1 in content: - content = content.replace(old_v1, new_v1, 1) - with open(CUSTOM_OPS_PATH, "w") as f: - f.write(content) - print(" [ok] V1: Added int→Tensor conversion for head_mapping") - return True - else: - print(" [warn] V1 function body not found — check manually") - return False - - -def patch_v2_python_fallback(): - """V2: Replace NotImplementedError with Python V2 fallback. - - The native V2 kernel exists but produces incorrect output. - Use paged_attention_v2_pytorch.py instead. - """ - with open(CUSTOM_OPS_PATH, "r") as f: - content = f.read() - - # Check if V2 is still NotImplementedError - if "raise NotImplementedError()" not in content: - print(" [skip] V2 NotImplementedError already replaced") - return True - - # Add import for Python V2 - import_line = "from vllm.paged_attention_v2_pytorch import paged_attention_v2_pytorch" - if import_line not in content: - anchor = "import ixformer.functions as ixf_F" - if anchor in content: - content = content.replace(anchor, anchor + "\n" + import_line, 1) - print(" [ok] Added Python V2 import") - - # Replace NotImplementedError with Python V2 call - old_v2_end = """ blocksparse_block_size: int = 64, - blocksparse_head_sliding_step: int = 0, -) -> None: - raise NotImplementedError()""" - - new_v2_end = """ blocksparse_block_size: int = 64, - blocksparse_head_sliding_step: int = 0, -) -> None: - # BI-V100: Native V2 kernel exists but has correctness issues. - # Using Python V2 (single-bmm + GQA broadcast) as fallback. - paged_attention_v2_pytorch( - out, exp_sum, max_logits, tmp_out, - 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, - )""" - - if old_v2_end in content: - content = content.replace(old_v2_end, new_v2_end, 1) - with open(CUSTOM_OPS_PATH, "w") as f: - f.write(content) - print(" [ok] V2: Replaced NotImplementedError with Python V2 fallback") - return True - else: - print(" [warn] V2 NotImplementedError block not found") - return False - - -def deploy_v2_module(): - """Copy Python V2 module into vllm package.""" - src = "/workspace/paged_attention_v2_pytorch.py" - dst = os.path.join(VLLM_ROOT, "paged_attention_v2_pytorch.py") - if os.path.exists(dst): - print(f" [skip] {dst} already exists") - return True - if os.path.exists(src): - shutil.copy2(src, dst) - print(f" [ok] Copied paged_attention_v2_pytorch.py → vllm/") - return True - else: - print(f" [warn] {src} not found — V2 fallback won't work") - return False - - -def patch_triton_path(): - """Fix Triton import path.""" - triton_src = "/usr/local/lib/python3.10/site-packages/triton" - triton_dst = "/usr/local/corex/lib64/python3/dist-packages/triton" - if os.path.exists(triton_src) and not os.path.exists(triton_dst): - try: - os.symlink(triton_src, triton_dst) - print(f" [ok] Symlinked triton → corex dist-packages") - except Exception as e: - print(f" [warn] Symlink failed: {e}") - else: - print(" [skip] Triton symlink already exists or source not found") - - # Also symlink triton's dependencies - for dep in ["triton"]: - src = f"/usr/local/lib/python3.10/site-packages/{dep}" - dst = f"/usr/local/corex/lib64/python3/dist-packages/{dep}" - if os.path.exists(src) and not os.path.exists(dst): - try: - os.symlink(src, dst) - except: - pass - - -def patch_flash_attn_prefill(): - """Enable ixformer flash_attn for prefill instead of Python fallback. - - The xformers backend's _run_sdpa_fallback is used when head_dim > 128. - With ixformer.flash_attn_func confirmed working at head_dim=256, - we can replace the fallback with a call to the native kernel. - """ - xformers_path = os.path.join(VLLM_ROOT, "attention/backends/xformers.py") - if not os.path.exists(xformers_path): - print(" [warn] xformers.py not found") - return False - - with open(xformers_path, "r") as f: - content = f.read() - - if "ixf_F.flash_attn_func" in content: - print(" [skip] flash_attn already patched into xformers.py") - return True - - # Find the _run_sdpa_fallback method and add flash_attn as first attempt - marker = "def _run_sdpa_fallback" - if marker not in content: - print(" [warn] _run_sdpa_fallback not found in xformers.py") - return False - - # Add import at top - if "import ixformer.functions as ixf_F" not in content: - content = "import ixformer.functions as ixf_F\n" + content - - # Insert flash_attn attempt at the start of _run_sdpa_fallback - old_def = " def _run_sdpa_fallback(" - new_def = """ def _run_sdpa_flash_attn( - self, - output: torch.Tensor, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - seq_lens: List[int], - is_prefill: bool, - ) -> torch.Tensor: - \"\"\"Try ixformer flash_attn_func first (native kernel, head_dim=256 OK).\"\"\" - try: - # flash_attn expects [batch, seqlen, nheads, headdim] - # Our inputs are [num_tokens, num_heads, head_size] - # Need to reshape per sequence - if is_prefill and len(seq_lens) == 1: - sq = seq_lens[0] - q = query[:sq].unsqueeze(0).transpose(1, 2) # [1, sq, H, d] - # Wait — flash_attn expects [B, S, H, D] not [B, H, S, D] - # query is [num_tokens, num_heads, head_size] - q = query[:sq].unsqueeze(0) # [1, sq, H, d] - k = key[:sq].unsqueeze(0) # [1, sq, kv_H, d] - v = value[:sq].unsqueeze(0) # [1, sq, kv_H, d] - out = ixf_F.flash_attn_func(q, k, v, causal=True) - output[:sq] = out.squeeze(0) - return output - except Exception: - pass - return self._run_sdpa_fallback(output, query, key, value, seq_lens, is_prefill) - - def _run_sdpa_fallback(""" - - content = content.replace(old_def, new_def, 1) - - with open(xformers_path, "w") as f: - f.write(content) - print(" [ok] Added flash_attn prefill path before _run_sdpa_fallback") - return True - - -def main(): - print("=== patch_ixformer_native: Hardware-verified kernel fixes ===\n") - - print("--- 1. V1 head_mapping int→Tensor ---") - patch_v1_head_mapping() - - print("\n--- 2. V2 Python fallback (native V2 has correctness issues) ---") - deploy_v2_module() - patch_v2_python_fallback() - - print("\n--- 3. Triton path fix ---") - patch_triton_path() - - print("\n--- 4. flash_attn for prefill ---") - patch_flash_attn_prefill() - - print("\n=== Summary ===") - print(" V1 decode (seq ≤ 8192): ixformer native kernel ✓ (0.03-0.27ms)") - print(" V2 decode (seq > 8192): Python V2 fallback (native V2 incorrect)") - print(" Prefill: ixformer flash_attn_func ✓ (head_dim=256 confirmed)") - print(" Triton: symlinked for import resolution") - - -if __name__ == "__main__": - main() diff --git a/qwen3_6_scripts/patch_model_runner.py b/qwen3_6_scripts/patch_model_runner.py deleted file mode 100644 index e10ad271..00000000 --- a/qwen3_6_scripts/patch_model_runner.py +++ /dev/null @@ -1,78 +0,0 @@ -""" -Fix: prefix_cache_hit stays True for chunked-prefill chunk 2+ even when past cache. - -Root cause: - model_runner.py _compute_for_prefix_cache_hit has three cases: - Case 1: prefix_cache_len <= context_len → "already past cache, do normal" - Case 2: context_len < prefix_cache_len < seq_len → partial hit, correct - Case 3: seq_len <= prefix_cache_len → full hit, reduce to 1 token - - Case 1 does nothing (leaves prefix_cache_hit = True). Then in utils.py: - if inter_data.prefix_cache_hit: - block_table = computed_block_nums ← ONLY the original prefix blocks! - - But context_len > prefix_cache_len means chunk 1 tokens (between prefix_cache_len - and context_len) are ALSO in KV cache and need to be in block_table. - block_table = computed_block_nums misses all chunk-1 blocks. - - In _forward_prefix_pytorch: - num_ctx_blocks = ceil(context_len / block_size) # e.g. 268 - block_tables.shape[1] = len(computed_block_nums) # e.g. 12 <-- too small! - At tile_blk >= 12: blk_ids is empty → k_t shape [..., 0] → amax crash. - -Fix: - Set prefix_cache_hit = False for Case 1, so utils.py falls through to: - elif chunked_prefill_enabled: - block_table = block_tables[seq_id] ← full block table (prefix + chunk1) -""" - -import re -import sys - -CANDIDATE_PATHS = [ - "/usr/local/corex/lib64/python3/dist-packages/vllm/worker/model_runner.py", - "/usr/local/corex/lib/python3/dist-packages/vllm/worker/model_runner.py", -] - -OLD_BLOCK = """\ - if prefix_cache_len <= context_len: - # We already passed the cache hit region, - # so do normal computation. - pass""" - -NEW_BLOCK = """\ - if prefix_cache_len <= context_len: - # We already passed the cache hit region, - # so do normal computation. - # Must clear prefix_cache_hit so _add_seq_group uses the full - # block_tables (prefix + previous-chunk blocks) instead of only - # computed_block_nums (prefix only). Without this, block_tables - # passed to _forward_prefix_pytorch is too narrow for context_len, - # causing an empty blk_ids slice and a zero-dim amax() crash. - inter_data.prefix_cache_hit = False""" - -import os - -patched = False -for path in CANDIDATE_PATHS: - if not os.path.exists(path): - continue - with open(path, "r") as f: - src = f.read() - if OLD_BLOCK not in src: - if NEW_BLOCK in src: - print(f"[patch_model_runner] already patched: {path}") - patched = True - break - print(f"[patch_model_runner] WARNING: expected block not found in {path}, skipping") - continue - patched_src = src.replace(OLD_BLOCK, NEW_BLOCK, 1) - with open(path, "w") as f: - f.write(patched_src) - print(f"[patch_model_runner] patched Case-1 prefix_cache_hit fix in: {path}") - patched = True - break - -if not patched: - print("[patch_model_runner] ERROR: could not find model_runner.py at any known path", file=sys.stderr) - sys.exit(1) diff --git a/qwen3_6_scripts/patch_paged_attention_v2.py b/qwen3_6_scripts/patch_paged_attention_v2.py deleted file mode 100644 index 4c5b47ef..00000000 --- a/qwen3_6_scripts/patch_paged_attention_v2.py +++ /dev/null @@ -1,192 +0,0 @@ -""" -patch_paged_attention_v2.py — Enable PagedAttention V2 on BI-V100 -================================================================== - -The baseline has paged_attention_v2 = raise NotImplementedError(). -paged_attn.py hardcodes use_v1=True to avoid calling it. - -This patch: - 1. Copies paged_attention_v2_pytorch.py into the vllm package - 2. Patches _custom_ops.py to call the PyTorch V2 implementation - 3. Patches paged_attn.py to enable V2 for long sequences (>8192 tokens) - -Performance impact: - V1 processes the entire KV sequence in a single kernel launch per (seq, head). - When seq_len > 8192, the single-block V1 kernel is memory-bandwidth-limited. - V2 splits the sequence into PARTITION_SIZE=512 chunks, processes them in - parallel, then reduces. For seq_len=100K: 195 parallel partitions vs 1. - - Expected improvement: 30-50% on Output TPS for long-context decode. - This matches the competition's advanced (30%) and special (50%) award tiers. - -Deploy: - cp paged_attention_v2_pytorch.py /usr/local/corex/lib/python3/dist-packages/vllm/ - python3 qwen3_6_scripts/patch_paged_attention_v2.py -""" - -import os -import shutil - -VLLM_ROOTS = [ - "/usr/local/corex/lib/python3/dist-packages/vllm", - "/usr/local/corex/lib64/python3/dist-packages/vllm", -] - -V2_MODULE_PYTORCH = "paged_attention_v2_pytorch.py" -V2_MODULE_TRITON = "paged_attention_v2_triton.py" - - -def find_vllm_root(): - for root in VLLM_ROOTS: - if os.path.exists(os.path.join(root, "_custom_ops.py")): - return root - return None - - -def patch_custom_ops(vllm_root): - """Replace paged_attention_v2 NotImplementedError with PyTorch implementation.""" - path = os.path.join(vllm_root, "_custom_ops.py") - - with open(path, "r") as f: - content = f.read() - - # Add import at the top (after existing imports) - 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: - # Insert after the last import line - anchor = "from vllm.platforms import current_platform" - if anchor in content: - content = content.replace( - anchor, - anchor + "\n" + import_line, - 1 - ) - print(" [ok] Added V2 import") - else: - print(" [warn] Import anchor not found") - return False - - # Replace the NotImplementedError body - old_v2 = ''' blocksparse_block_size: int = 64, - blocksparse_head_sliding_step: int = 0, -) -> None: - raise NotImplementedError()''' - - new_v2 = ''' blocksparse_block_size: int = 64, - blocksparse_head_sliding_step: int = 0, -) -> None: - # BI-V100: PyTorch V2 implementation (replaces NotImplementedError) - _v2_impl( - out, exp_sum, max_logits, tmp_out, - 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, - tp_rank, blocksparse_local_blocks, - blocksparse_vert_stride, blocksparse_block_size, - blocksparse_head_sliding_step, - )''' - - if "paged_attention_v2_pytorch(" in content: - print(" [skip] V2 body already patched") - elif old_v2 in content: - content = content.replace(old_v2, new_v2, 1) - print(" [ok] Replaced V2 NotImplementedError with PyTorch implementation") - else: - print(" [warn] V2 function body not found as expected") - return False - - with open(path, "w") as f: - f.write(content) - print(f" Written: {path}") - return True - - -def patch_paged_attn(vllm_root): - """Enable V2 for long sequences instead of forcing V1.""" - path = os.path.join(vllm_root, "attention/ops/paged_attn.py") - - with open(path, "r") as f: - content = f.read() - - # The baseline has: - # use_v1 = (max_seq_len <= 8192 and ...) - # use_v1 = True # <-- hardcoded override - # We want to remove the hardcoded override so V2 is used for long sequences. - - old_heuristic = " use_v1 = True" - new_heuristic = " # use_v1 = True # Removed: V2 now works on BI-V100 (paged_attention_v2_pytorch)" - - if "V2 now works" in content: - print(" [skip] V1 override already removed") - elif old_heuristic in content: - content = content.replace(old_heuristic, new_heuristic, 1) - print(" [ok] Removed use_v1=True hardcode — V2 enabled for seq_len > 8192") - else: - print(" [warn] use_v1=True line not found") - return False - - with open(path, "w") as f: - f.write(content) - print(f" Written: {path}") - return True - - -def deploy_v2_module(vllm_root): - """Copy the V2 PyTorch module into the vllm package.""" - src = os.path.join(os.path.dirname(__file__), "..", V2_MODULE_PYTORCH) - if not os.path.exists(src): - src = os.path.join("/workspace", V2_MODULE_PYTORCH) - if not os.path.exists(src): - # Try relative to this script - src = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", V2_MODULE_PYTORCH) - - dst = os.path.join(vllm_root, V2_MODULE_PYTORCH) - - if os.path.exists(dst): - print(f" [skip] V2 module {dst} already exists") - return True - - if not os.path.exists(src): - print(f" [error] V2 module not found at {src}") - return False - - shutil.copy2(src, dst) - print(f" [ok] Copied {V2_MODULE_PYTORCH} → {dst}") - return True - - -def main(): - print("=== patch_paged_attention_v2: Enable V2 on BI-V100 ===\n") - - vllm_root = find_vllm_root() - if not vllm_root: - print("[error] vllm package not found") - return - - print(f"vllm root: {vllm_root}\n") - - print("Step 1: Deploy V2 PyTorch module") - deploy_v2_module(vllm_root) - - print("\nStep 2: Patch _custom_ops.py") - patch_custom_ops(vllm_root) - - print("\nStep 3: Patch paged_attn.py (enable V2 for long sequences)") - patch_paged_attn(vllm_root) - - print("\nDone. V2 is now enabled for seq_len > 8192.") - - -if __name__ == "__main__": - main() diff --git a/qwen3_6_scripts/patch_triton_tuning.py b/qwen3_6_scripts/patch_triton_tuning.py deleted file mode 100644 index d7bf838c..00000000 --- a/qwen3_6_scripts/patch_triton_tuning.py +++ /dev/null @@ -1,91 +0,0 @@ -""" -patch_triton_tuning.py — BI-V100 Triton kernel parameter optimization -====================================================================== - -Patches prefix_prefill.py to use BI-V100-optimal BLOCK and NUM_WARPS values. - -Hardware derivation: - BI-V100 SMEM = 48KB. Triton Flash Attention needs K+V tiles in SMEM: - SMEM = BLOCK_N × head_dim × sizeof(fp16) × 2 - BLOCK_N=64, head_dim=128 → 32KB ≤ 48KB ✓ (current, correct) - BLOCK_N=128, head_dim=128 → 64KB > 48KB ✗ (would crash) - → BLOCK must stay at 64. - - NUM_WARPS derivation: - At BLOCK=64, each block does 64 query positions. - 8 warps = 256 threads → each thread handles 32 elements from Q tile. - 4 warps = 128 threads → each thread handles 64 elements. - - With 16 SMs (confirmed) and typical grid of 37K+ blocks: - At 8 warps + 32KB SMEM: 1 block per SM (SMEM-limited) - At 4 warps + 32KB SMEM: potentially 2 blocks per SM - - BI-V100 is bandwidth-limited (900 GB/s), not latency-limited. - Fewer warps hiding latency matters less; more blocks = better. - → NUM_WARPS = 4 - -Deploy: python3 qwen3_6_scripts/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", -] - -# Original line (baseline): -OLD_BLOCK = " BLOCK = 128 if current_platform.has_device_capability(80) else 64\n NUM_WARPS = 8" - -# Optimized for BI-V100: -NEW_BLOCK = """\ - # BI-V100 optimization (patch_triton_tuning.py): - # BLOCK=64: SMEM constraint — BLOCK_N=128 overflows 48KB at head_dim=128 - # NUM_WARPS=4: bandwidth-limited GPU benefits from more blocks/SM over more warps - # Derivation: 4 warps at BLOCK=64 allows 2 concurrent blocks per SM, - # doubling occupancy vs 8 warps (which is SMEM-limited to 1 block/SM). - BLOCK = 64 - NUM_WARPS = 4""" - - -def patch(): - for path in PREFIX_PREFILL_PATHS: - if not os.path.exists(path): - continue - - with open(path, "r") as f: - content = f.read() - - if "NUM_WARPS = 4" in content: - print(f" [skip] {path}: already patched") - return - - if OLD_BLOCK not in content: - # Try the alternative: maybe it's already using BLOCK=64 hardcoded - alt_old = " BLOCK = 64\n NUM_WARPS = 8" - if alt_old in content: - content = content.replace(alt_old, NEW_BLOCK, 1) - with open(path, "w") as f: - f.write(content) - print(f" [ok] {path}: patched NUM_WARPS 8→4 (BLOCK was already 64)") - return - print(f" [warn] {path}: original block not found, manual check needed") - return - - content = content.replace(OLD_BLOCK, NEW_BLOCK, 1) - with open(path, "w") as f: - f.write(content) - print(f" [ok] {path}: patched BLOCK=64, NUM_WARPS=4") - return - - print(" [error] prefix_prefill.py not found at any expected path") - - -def main(): - print("=== patch_triton_tuning: BI-V100 Triton kernel optimization ===") - patch() - print("Done.") - - -if __name__ == "__main__": - main() diff --git a/qwen3_6_scripts/patch_vectorized_decode.py b/qwen3_6_scripts/patch_vectorized_decode.py deleted file mode 100644 index 27b4b945..00000000 --- a/qwen3_6_scripts/patch_vectorized_decode.py +++ /dev/null @@ -1,72 +0,0 @@ -""" -patch_vectorized_decode.py — Vectorize the decode PyTorch fallback -=================================================================== - -Current _forward_decode_pytorch (used when seq_len > 32768): - for i in range(num_seqs): ← Python for-loop - k_t = key_cache[blk_ids]... ← per-sequence gather - attn_w = torch.matmul(q, k_t) ← per-sequence matmul - output[i] = ... - -Problem: When num_seqs=1 (competition config), this loop runs once. - But the inner operations do seq_len worth of gather+matmul in Python. - The real bottleneck is the .permute().contiguous().view() chain on K/V, - which creates multiple intermediate tensors. - -Optimization: Fuse the gather and reduce steps: - 1. Use torch.index_select instead of fancy indexing for K/V gather - 2. Pre-compute the scale factor into Q - 3. Avoid the .float() → .to(orig_dtype) round-trip where possible - 4. Use torch.baddbmm for fused scale+matmul - -This won't change the asymptotic complexity, but reduces Python overhead -and intermediate tensor allocations. The real fix is making paged_attention_v1 -work at seq_len > 32768 (raise the threshold or fix the kernel). - -Deploy: python3 qwen3_6_scripts/patch_vectorized_decode.py -""" - -import os - -PAGED_ATTN_PATH = "/usr/local/corex/lib/python3/dist-packages/vllm/attention/ops/paged_attn.py" - -# Raise the threshold: try letting ixf_F.paged_attention_v1 handle longer sequences -# The baseline sets it to 32768 because v1 "fails for long contexts" -# But this might be a conservative limit — let's try 65536 first -# If it crashes, the user can lower it back - -OLD_THRESHOLD = " _PYTORCH_DECODE_THRESHOLD = 32768" -NEW_THRESHOLD = """\ - # BI-V100: Try higher threshold for compiled v1 kernel. - # Baseline: 32768 (conservative). We try 65536 — the v1 kernel is - # orders of magnitude faster than the Python fallback. - # If v1 crashes at higher seq_lens, lower this back to 32768. - _PYTORCH_DECODE_THRESHOLD = 65536""" - - -def patch(): - if not os.path.exists(PAGED_ATTN_PATH): - print(f" [error] {PAGED_ATTN_PATH} not found") - return - with open(PAGED_ATTN_PATH, "r") as f: - content = f.read() - if "PYTORCH_DECODE_THRESHOLD = 65536" in content: - print(f" [skip] already patched to 65536") - return - if OLD_THRESHOLD in content: - content = content.replace(OLD_THRESHOLD, NEW_THRESHOLD, 1) - with open(PAGED_ATTN_PATH, "w") as f: - f.write(content) - print(f" [ok] _PYTORCH_DECODE_THRESHOLD: 32768 → 65536") - else: - print(f" [warn] threshold anchor not found") - - -def main(): - print("=== patch_vectorized_decode: raise decode threshold ===") - patch() - print("Done.") - - -if __name__ == "__main__": - main() diff --git a/qwen3_6_scripts/patch_vllm_qwen3_5.py b/qwen3_6_scripts/patch_vllm_qwen3_5.py deleted file mode 100644 index a8e091a4..00000000 --- a/qwen3_6_scripts/patch_vllm_qwen3_5.py +++ /dev/null @@ -1,76 +0,0 @@ -""" -Patches the vLLM model registry and deploys the Qwen3_5 model file. - -Deploy steps on the remote machine: - 1. cp modified_scripts/qwen3_5.py \ - /usr/local/corex/lib64/python3/dist-packages/vllm/model_executor/models/qwen3_5.py - 2. python3 modified_scripts/patch_vllm_qwen3_5.py - -Also edit your model config.json to set: - "architectures": ["Qwen3_5ForCausalLM"] - -Target: vLLM at /usr/local/corex/lib64/python3/dist-packages/vllm/ -""" - -VLLM_ROOT = "/usr/local/corex/lib64/python3/dist-packages/vllm" -REGISTRY = f"{VLLM_ROOT}/model_executor/models/registry.py" - - -def patch_file(path, replacements): - with open(path, "r") as f: - content = f.read() - - patched = False - for old, new in replacements: - if new in content: - print(f" [skip] already patched: {repr(new[:70])}") - continue - if old not in content: - print(f" [warn] anchor not found: {repr(old[:70])}") - continue - content = content.replace(old, new, 1) - patched = True - print(f" [ok] patched after: {repr(old[:70])}") - - if patched: - with open(path, "w") as f: - f.write(content) - - -def main(): - print(f"=== Patching {REGISTRY} ===") - patch_file(REGISTRY, [ - ( - ' "Qwen3ForCausalLM": ("qwen3", "Qwen3ForCausalLM"),\n' - ' "Qwen3MoeForCausalLM": ("qwen3_moe", "Qwen3MoeForCausalLM"),', - ' "Qwen3ForCausalLM": ("qwen3", "Qwen3ForCausalLM"),\n' - ' "Qwen3MoeForCausalLM": ("qwen3_moe", "Qwen3MoeForCausalLM"),\n' - ' "Qwen3_5ForCausalLM": ("qwen3_5", "Qwen3_5ForCausalLM"),\n' - ' "Qwen3_5MoeForCausalLM": ("qwen3_5", "Qwen3_5MoeForCausalLM"),', - ), - ]) - - print("\n=== Verification ===") - try: - import importlib.util - spec = importlib.util.spec_from_file_location( - "qwen3_5", - f"{VLLM_ROOT}/model_executor/models/qwen3_5.py", - ) - mod = importlib.util.module_from_spec(spec) - # Quick check: does the class exist? - spec.loader.exec_module(mod) - cls = mod.Qwen3_5ForCausalLM - print(f" Qwen3_5ForCausalLM found: {cls}") - cls_moe = mod.Qwen3_5MoeForCausalLM - print(f" Qwen3_5MoeForCausalLM found: {cls_moe}") - except Exception as e: - print(f" [warn] verification failed (may be OK at runtime): {e}") - - print("\nDone. Remember to:") - print(" 1. Set config.json 'architectures': ['Qwen3_5ForCausalLM'] or ['Qwen3_5MoEForCausalLM']") - print(" 2. Run patch_transformers_qwen3_5.py if not already done") - - -if __name__ == "__main__": - main() diff --git a/qwen3_6_scripts/patch_vllm_tool_parser.py b/qwen3_6_scripts/patch_vllm_tool_parser.py deleted file mode 100644 index f2575ba9..00000000 --- a/qwen3_6_scripts/patch_vllm_tool_parser.py +++ /dev/null @@ -1,79 +0,0 @@ -""" -Patches vLLM 0.6.3 to register Qwen3CoderToolParser under the name "qwen3_coder". - -Deploy steps on the remote machine (already called by patch_ops.sh): - 1. cp qwen3coder_tool_parser.py \ - /usr/local/corex/lib/python3/dist-packages/vllm/entrypoints/openai/tool_parsers/ - 2. python3 patch_vllm_tool_parser.py - -Usage after patching: - --tool-call-parser qwen3_coder --enable-auto-tool-choice -""" - -import os - -VLLM_ROOT = "/usr/local/corex/lib/python3/dist-packages/vllm" -TOOL_PARSERS_DIR = f"{VLLM_ROOT}/entrypoints/openai/tool_parsers" -INIT_FILE = f"{TOOL_PARSERS_DIR}/__init__.py" - - -def patch_file(path, replacements): - with open(path, "r") as f: - content = f.read() - - patched = False - for old, new in replacements: - if new in content: - print(f" [skip] already patched: {repr(new[:70])}") - continue - if old not in content: - print(f" [warn] anchor not found: {repr(old[:70])}") - continue - content = content.replace(old, new, 1) - patched = True - print(f" [ok] patched: {repr(old[:50])} -> {repr(new[:50])}") - - if patched: - with open(path, "w") as f: - f.write(content) - - -def main(): - if not os.path.isdir(TOOL_PARSERS_DIR): - raise FileNotFoundError( - f"Tool parsers directory not found: {TOOL_PARSERS_DIR}\n" - "Verify the vLLM installation path.") - - print(f"=== Patching {INIT_FILE} ===") - patch_file(INIT_FILE, [ - ( - "from .mistral_tool_parser import MistralToolParser", - "from .mistral_tool_parser import MistralToolParser\n" - "from .qwen3coder_tool_parser import Qwen3CoderToolParser", - ), - ( - '"MistralToolParser", "Internlm2ToolParser", "Llama3JsonToolParser"\n]', - '"MistralToolParser", "Internlm2ToolParser", "Llama3JsonToolParser",\n' - ' "Qwen3CoderToolParser"\n]', - ), - ]) - - print("\n=== Verification ===") - try: - import importlib.util - spec = importlib.util.spec_from_file_location( - "qwen3coder_tool_parser", - f"{TOOL_PARSERS_DIR}/qwen3coder_tool_parser.py", - ) - mod = importlib.util.module_from_spec(spec) - print(f" Module spec loaded: {spec.name}") - print(" (full import requires torch/vllm runtime — skipping exec)") - except Exception as e: - print(f" [warn] spec check failed: {e}") - - print("\nDone. Start vLLM server with:") - print(" --tool-call-parser qwen3_coder --enable-auto-tool-choice") - - -if __name__ == "__main__": - main() diff --git a/qwen3_6_scripts/patch_xformers_sdpa_batch.py b/qwen3_6_scripts/patch_xformers_sdpa_batch.py deleted file mode 100644 index a585b4d0..00000000 --- a/qwen3_6_scripts/patch_xformers_sdpa_batch.py +++ /dev/null @@ -1,192 +0,0 @@ -""" -策略：批量（block-diagonal）fallback — 纯 PyTorch 数学实现 -============================================================= -构建块对角 causal mask，对整批序列一次 matmul + softmax， -完全绕开所有硬件 flash attention kernel。 - -背景： - ixformer flshattF: head_dim > 128 报错拒绝 - cudnnFlashAttnForward: 接受 head_dim=256，但数值结果错误（输出全"!"） - 两者大概率是同一硬件单元，ixformer 提前拦截了硬件不支持的配置。 - 纯 matmul 路径完全绕开硬件 flash attention，数值正确。 - -优点： - 数值正确。 - 并发请求 prefill attention 在 GPU 上真正并行（一次大 matmul）。 - -缺点： - 峰值显存 = total_tokens² × H × dtype_size - total_tokens 受 --max-num-batched-tokens 控制，max-model-len 控制不住。 - -内存参考（fp16，H_local=6，--max-num-batched-tokens=T）： - T=2048 → 峰值 ~50 MB - T=4096 → 峰值 ~200 MB - T=8192 → 峰值 ~800 MB - T=16384 → 峰值 ~3.2 GB - -Deploy: - python3 modified_scripts/patch_xformers_sdpa_batch.py -""" - -XFORMERS_PATH = ( - "/usr/local/corex/lib64/python3/dist-packages/" - "vllm/attention/backends/xformers.py" -) - -FALLBACK_METHOD = ''' - def _run_sdpa_fallback( - self, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attn_metadata: "XFormersMetadata", - ) -> torch.Tensor: - """批量纯数学 attention fallback。 - - 构建块对角 causal mask（等价于 ixformer BlockDiagonalCausalMask）， - 对整批序列一次 matmul + softmax，GPU 并行处理所有序列。 - - 块对角 mask 结构（seq1 len=3，seq2 len=2）： - s1,0 s1,1 s1,2 s2,0 s2,1 - s1,0 [ 0 -inf -inf -inf -inf ] - s1,1 [ 0 0 -inf -inf -inf ] - s1,2 [ 0 0 0 -inf -inf ] - s2,0 [-inf -inf -inf 0 -inf ] - s2,1 [-inf -inf -inf 0 0 ] - - softmax 在 float32 下计算防止 float16 溢出，结果转回原始 dtype。 - - Args: - query : [1, total_prefill_tokens, num_heads, head_dim] - key : [1, total_prefill_tokens, num_kv_heads, head_dim] - value : [1, total_prefill_tokens, num_kv_heads, head_dim] - Returns: - [1, total_prefill_tokens, num_heads, head_dim] - """ - assert attn_metadata.seq_lens is not None - orig_dtype = query.dtype - total_tokens = query.shape[1] - - # ── 构建块对角 causal mask [T, T] ──────────────────────────────── - # 全部初始化为 -inf，再对每条序列的对角块填入下三角 0 - mask = torch.full( - (total_tokens, total_tokens), - float("-inf"), - dtype=torch.float32, - device=query.device, - ) - start = 0 - for seq_len in attn_metadata.seq_lens: - end = start + seq_len - mask[start:end, start:end] = torch.tril( - torch.zeros(seq_len, seq_len, - dtype=torch.float32, device=query.device) - ) - start = end - - # ── [1, H, T, D]，.contiguous() ────────────────────────────────── - q_all = query.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0) - k_all = key.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0) - v_all = value.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0) - - # ── GQA：展开 KV heads ──────────────────────────────────────────── - if k_all.shape[1] != q_all.shape[1]: - n = q_all.shape[1] // k_all.shape[1] - k_all = k_all.repeat_interleave(n, dim=1).contiguous() - v_all = v_all.repeat_interleave(n, dim=1).contiguous() - - # ── 纯数学 attention（float32 防溢出）──────────────────────────── - # [1, H, T, T] - attn_w = torch.matmul(q_all.float(), k_all.float().transpose(-2, -1)) - attn_w = attn_w * self.scale - attn_w = attn_w + mask # 加法广播：mask [T,T] → [1, H, T, T] - attn_w = torch.softmax(attn_w, dim=-1) - - out = torch.matmul(attn_w, v_all.float()).to(orig_dtype) - # [1, H, T, D] → [1, T, H, D] - return out.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0) - -''' - -OLD_XFORMER_BLOCK = """\ - self.attn_op = xops.fmha.flash.FwOp() - if self.alibi_slopes is None: - # Add the batch dimension. - query = query.unsqueeze(0) - key = key.unsqueeze(0) - value = value.unsqueeze(0) - out = xops.memory_efficient_attention_forward( - query, - key, - value, - attn_bias=attn_bias[0], - p=0.0, - scale=self.scale, - op = self.attn_op - ) - return out.view_as(original_query)\ -""" - -NEW_XFORMER_BLOCK = """\ - self.attn_op = xops.fmha.flash.FwOp() - if self.alibi_slopes is None: - # Add the batch dimension. - query = query.unsqueeze(0) - key = key.unsqueeze(0) - value = value.unsqueeze(0) - if self.head_size > 128: - out = self._run_sdpa_fallback(query, key, value, attn_metadata) - else: - out = xops.memory_efficient_attention_forward( - query, - key, - value, - attn_bias=attn_bias[0], - p=0.0, - scale=self.scale, - op=self.attn_op, - ) - return out.view_as(original_query)\ -""" - -INJECT_ANCHOR = " def _run_memory_efficient_xformers_forward(" - - -def patch_file(path): - with open(path, "r") as f: - content = f.read() - changed = False - - if "_run_sdpa_fallback" in content: - print(" [skip] _run_sdpa_fallback already present") - elif INJECT_ANCHOR not in content: - print(" [warn] inject anchor not found") - else: - content = content.replace(INJECT_ANCHOR, FALLBACK_METHOD + INJECT_ANCHOR, 1) - print(" [ok] injected _run_sdpa_fallback (batch, pure-math)") - changed = True - - if NEW_XFORMER_BLOCK in content: - print(" [skip] dispatch block already patched") - elif OLD_XFORMER_BLOCK in content: - content = content.replace(OLD_XFORMER_BLOCK, NEW_XFORMER_BLOCK, 1) - print(" [ok] patched dispatch block") - changed = True - else: - print(" [warn] dispatch block anchor not found") - - if changed: - with open(path, "w") as f: - f.write(content) - print(f" Written: {path}") - - -def main(): - print("=== patch_xformers_sdpa_batch (batch, pure-math) ===") - print(f"Target: {XFORMERS_PATH}") - patch_file(XFORMERS_PATH) - print("\nDone.") - - -if __name__ == "__main__": - main() diff --git a/qwen3_6_scripts/patch_xformers_sdpa_batch_kernel.py b/qwen3_6_scripts/patch_xformers_sdpa_batch_kernel.py deleted file mode 100644 index e7f647ff..00000000 --- a/qwen3_6_scripts/patch_xformers_sdpa_batch_kernel.py +++ /dev/null @@ -1,191 +0,0 @@ -""" -策略：批量（block-diagonal）— F.scaled_dot_product_attention，可走硬件 kernel -============================================================================= -构建块对角 causal mask，对整批序列一次 F.scaled_dot_product_attention。 -与 patch_xformers_sdpa_batch.py（纯 matmul）的区别： - SDPA 会根据 PyTorch/驱动能力分发到最优 kernel（Flash Attention / - mem-efficient attention / math fallback），而不是固定走 cublas matmul。 - -历史说明： - 该方案最早因输出全"!"而被弃用，后续排查确认"!"由 mamba_cache.py bug - 引起，与 attention 实现无关。当前恢复此方案用于性能对比测试。 - -已知硬件限制（BI-V100）： - cudnnFlashAttnForward 不支持 is_causal=True（报错）。 - 本实现使用 is_causal=False + 显式块对角 additive mask 规避此限制。 - 若 SDPA 仍分发到有问题的 kernel，回退到 patch_xformers_sdpa_batch.py。 - -优点（vs 纯 matmul）： - SDPA 可分发到 Flash Attention kernel → O(L) 显存、更快的 CUDA kernel。 - -缺点： - 依赖硬件 kernel 行为，若 kernel 有 bug 则数值错误（需与 matmul 版对比验证）。 - -Deploy: - python3 modified_scripts/patch_xformers_sdpa_batch_kernel.py -""" - -XFORMERS_PATH = ( - "/usr/local/corex/lib64/python3/dist-packages/" - "vllm/attention/backends/xformers.py" -) - -FALLBACK_METHOD = ''' - def _run_sdpa_fallback( - self, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attn_metadata: "XFormersMetadata", - ) -> torch.Tensor: - """批量 F.scaled_dot_product_attention fallback（可走硬件 kernel）。 - - 构建块对角 causal mask，对整批序列一次 SDPA 调用。 - SDPA 可分发到 Flash Attention / mem-efficient attention kernel。 - is_causal=False + 显式 additive mask，规避 cudnnFlashAttnForward - 不支持 is_causal=True 的限制。 - - 块对角 mask（seq1 len=3，seq2 len=2）： - s1,0 s1,1 s1,2 s2,0 s2,1 - s1,0 [ 0 -inf -inf -inf -inf ] - s1,1 [ 0 0 -inf -inf -inf ] - s1,2 [ 0 0 0 -inf -inf ] - s2,0 [-inf -inf -inf 0 -inf ] - s2,1 [-inf -inf -inf 0 0 ] - - Args: - query : [1, total_prefill_tokens, num_heads, head_dim] - key : [1, total_prefill_tokens, num_kv_heads, head_dim] - value : [1, total_prefill_tokens, num_kv_heads, head_dim] - Returns: - [1, total_prefill_tokens, num_heads, head_dim] - """ - import torch.nn.functional as F - - assert attn_metadata.seq_lens is not None - orig_dtype = query.dtype - total_tokens = query.shape[1] - - # ── 块对角 causal mask [T, T] ───────────────────────────────────── - mask = torch.full( - (total_tokens, total_tokens), - float("-inf"), - dtype=orig_dtype, - device=query.device, - ) - start = 0 - for seq_len in attn_metadata.seq_lens: - end = start + seq_len - mask[start:end, start:end] = torch.tril( - torch.zeros(seq_len, seq_len, dtype=orig_dtype, device=query.device) - ) - start = end - - # ── [1, H, T, D] ────────────────────────────────────────────────── - q_all = query.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0) - k_all = key.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0) - v_all = value.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0) - - # ── GQA：展开 KV heads ──────────────────────────────────────────── - if k_all.shape[1] != q_all.shape[1]: - n = q_all.shape[1] // k_all.shape[1] - k_all = k_all.repeat_interleave(n, dim=1).contiguous() - v_all = v_all.repeat_interleave(n, dim=1).contiguous() - - # ── F.scaled_dot_product_attention（可走硬件 kernel）───────────── - # is_causal=False：避免 cudnnFlashAttnForward "not support causal mode" - # attn_mask 传 additive float mask（非 bool），SDPA 选择 math/kernel 路径 - out = F.scaled_dot_product_attention( - q_all, k_all, v_all, - attn_mask=mask, - dropout_p=0.0, - is_causal=False, - scale=self.scale, - ) - # [1, H, T, D] → [1, T, H, D] - return out.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0) - -''' - -OLD_XFORMER_BLOCK = """\ - self.attn_op = xops.fmha.flash.FwOp() - if self.alibi_slopes is None: - # Add the batch dimension. - query = query.unsqueeze(0) - key = key.unsqueeze(0) - value = value.unsqueeze(0) - out = xops.memory_efficient_attention_forward( - query, - key, - value, - attn_bias=attn_bias[0], - p=0.0, - scale=self.scale, - op = self.attn_op - ) - return out.view_as(original_query)\ -""" - -NEW_XFORMER_BLOCK = """\ - self.attn_op = xops.fmha.flash.FwOp() - if self.alibi_slopes is None: - # Add the batch dimension. - query = query.unsqueeze(0) - key = key.unsqueeze(0) - value = value.unsqueeze(0) - if self.head_size > 128: - out = self._run_sdpa_fallback(query, key, value, attn_metadata) - else: - out = xops.memory_efficient_attention_forward( - query, - key, - value, - attn_bias=attn_bias[0], - p=0.0, - scale=self.scale, - op=self.attn_op, - ) - return out.view_as(original_query)\ -""" - -INJECT_ANCHOR = " def _run_memory_efficient_xformers_forward(" - - -def patch_file(path): - with open(path, "r") as f: - content = f.read() - changed = False - - if "_run_sdpa_fallback" in content: - print(" [skip] _run_sdpa_fallback already present") - elif INJECT_ANCHOR not in content: - print(" [warn] inject anchor not found") - else: - content = content.replace(INJECT_ANCHOR, FALLBACK_METHOD + INJECT_ANCHOR, 1) - print(" [ok] injected _run_sdpa_fallback (batch, F.sdpa kernel)") - changed = True - - if NEW_XFORMER_BLOCK in content: - print(" [skip] dispatch block already patched") - elif OLD_XFORMER_BLOCK in content: - content = content.replace(OLD_XFORMER_BLOCK, NEW_XFORMER_BLOCK, 1) - print(" [ok] patched dispatch block") - changed = True - else: - print(" [warn] dispatch block anchor not found") - - if changed: - with open(path, "w") as f: - f.write(content) - print(f" Written: {path}") - - -def main(): - print("=== patch_xformers_sdpa_batch_kernel (batch, F.sdpa + kernel dispatch) ===") - print(f"Target: {XFORMERS_PATH}") - patch_file(XFORMERS_PATH) - print("\nDone.") - - -if __name__ == "__main__": - main() diff --git a/qwen3_6_scripts/patch_xformers_sdpa_seq.py b/qwen3_6_scripts/patch_xformers_sdpa_seq.py deleted file mode 100644 index 496abc1f..00000000 --- a/qwen3_6_scripts/patch_xformers_sdpa_seq.py +++ /dev/null @@ -1,321 +0,0 @@ -""" -策略：顺序（per-sequence）fallback — 纯 PyTorch 数学实现 -========================================================== -逐条序列用 matmul + softmax 手写 attention，完全绕开所有硬件 -flash attention kernel（ixformer / cudnnFlashAttnForward）。 - -背景： - Iluvatar cudnnFlashAttnForward 存在两个已知问题： - 1. 不支持 is_causal=True（报错） - 2. 使用 attn_mask 路径时数值结果不正确（静默错误，输出全为"!"） - 与华为昇腾 910B4 上 llama.cpp --flash-attn off 修复同类问题的原理相同。 - 纯数学路径（matmul + softmax）在任何 PyTorch 后端上结果都正确。 - -优点： - 数值正确，不依赖任何硬件特定 attention kernel。 - 峰值显存 = max(seq_len)² × H × dtype_size，由 --max-model-len 控制。 - -缺点： - 并发请求的 prefill attention 串行执行。 - O(L²) 显存（无 flash attention 的 O(L) 优化）。 - -内存参考（fp16，H_local=6）： - max-model-len=4096 → 峰值 ~200 MB - max-model-len=8192 → 峰值 ~800 MB - max-model-len=16384 → 峰值 ~3.2 GB - -额外 patch（arg_utils.py）： - vllm 0.6.3 在 max_model_len > 32K 时会自动开启 chunked prefill（无命令行 - 关闭选项），原意是防止 profiling OOM。但 _run_sdpa_fallback 已通过 Q-tiling - 解决了该问题，chunked prefill 反而会把推理路径从 _run_sdpa_fallback 切换到 - _forward_prefix_pytorch，属于不必要的行为变更，因此一并禁用该自动逻辑。 - -Deploy: - python3 modified_scripts/patch_xformers_sdpa_seq.py -""" - -XFORMERS_PATH = ( - "/usr/local/corex/lib64/python3/dist-packages/" - "vllm/attention/backends/xformers.py" -) - -ARG_UTILS_PATH = ( - "/usr/local/corex/lib64/python3/dist-packages/" - "vllm/engine/arg_utils.py" -) - -LOGITS_PROC_PATH = ( - "/usr/local/corex/lib64/python3/dist-packages/" - "vllm/model_executor/layers/logits_processor.py" -) - -# _apply_logits_processors crashes when seq_groups is None (intermediate -# chunked-prefill chunks on the driver rank). Add an early-return guard. -_LP_OLD_BLOCK = """\ -def _apply_logits_processors( - logits: torch.Tensor, - sampling_metadata: SamplingMetadata, -) -> torch.Tensor: - found_logits_processors = False\ -""" - -_LP_NEW_BLOCK = """\ -def _apply_logits_processors( - logits: torch.Tensor, - sampling_metadata: SamplingMetadata, -) -> torch.Tensor: - if sampling_metadata.seq_groups is None: # intermediate chunked-prefill chunk - return logits - found_logits_processors = False\ -""" - -# vllm 0.6.3 自动开启 chunked prefill 的原始块 -_ARG_OLD_BLOCK = """\ - if (is_gpu and not use_sliding_window and not use_spec_decode - and not self.enable_lora - and not self.enable_prompt_adapter): - self.enable_chunked_prefill = True - logger.warning( - "Chunked prefill is enabled by default for models with " - "max_model_len > 32K. Currently, chunked prefill might " - "not work with some features or models. If you " - "encounter any issues, please disable chunked prefill " - "by setting --enable-chunked-prefill=False.")\ -""" - -_ARG_NEW_BLOCK = """\ - if (is_gpu and not use_sliding_window and not use_spec_decode - and not self.enable_lora - and not self.enable_prompt_adapter): - pass # skip auto-enable: Q-tiling in _run_sdpa_fallback - # handles long-context memory without chunked prefill\ -""" - -FALLBACK_METHOD = ''' - def _run_sdpa_fallback( - self, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attn_metadata: "XFormersMetadata", - ) -> torch.Tensor: - """纯数学 causal attention fallback，带 Q-tiling 内存优化。 - - 调用时机：kv_cache.numel()==0（profiling 阶段）。 - 此路径无 KV 缓存前缀，KV 长度 == query 长度。 - - 内存优化（Q-tiling，与 Flash Attention 同思路）： - 将 Q 分成 _Q_CHUNK 大小的子块逐块计算，每块峰值内存 - O(_Q_CHUNK × q_len) 而非 O(q_len²)。 - profiling 阶段序列可能达到 max_model_len（如 20K tokens）， - 不加 Q-tiling 会产生 9.6 GB 矩阵直接 OOM。 - - softmax 在 float32 下计算以防止 float16 溢出，结果转回原始 dtype。 - - Args: - query : [1, total_query_tokens, num_heads, head_dim] - key : [1, total_query_tokens, num_kv_heads, head_dim] - value : [1, total_query_tokens, num_kv_heads, head_dim] - Returns: - [1, total_query_tokens, num_heads, head_dim] - """ - _Q_CHUNK = 256 # 与 _forward_prefix_pytorch 的 _ATTN_Q_CHUNK 保持一致 - - assert attn_metadata.seq_lens is not None - orig_dtype = query.dtype - num_seqs = len(attn_metadata.seq_lens) - - # 推导每条序列的实际 query 长度。 - # 正常 prefill 时 q_len == seq_len；如果将来遇到 chunked 场景， - # query_start_loc 记录的是真实 query token 数（非全序列长度）。 - if (attn_metadata.query_start_loc is not None - and len(attn_metadata.query_start_loc) == num_seqs + 1): - q_lens = [ - int(attn_metadata.query_start_loc[i + 1].item()) - - int(attn_metadata.query_start_loc[i].item()) - for i in range(num_seqs) - ] - else: - q_lens = list(attn_metadata.seq_lens) - - q_flat = query.squeeze(0) # [T, H, D] - k_flat = key.squeeze(0) # [T, Hkv, D] - v_flat = value.squeeze(0) - - output = torch.empty_like(q_flat) - seq_start = 0 - for q_len in q_lens: - seq_end = seq_start + q_len - - # 当前序列的完整 K/V（此路径无前缀，KV == Q） - k_s = k_flat[seq_start:seq_end].permute(1, 0, 2).float() # [Hkv, q_len, D] - v_s = v_flat[seq_start:seq_end].permute(1, 0, 2).float() # [Hkv, q_len, D] - - # GQA：展开 KV heads 至与 query heads 一致 - if k_s.shape[0] != self.num_heads: - n = self.num_heads // k_s.shape[0] - k_s = k_s.repeat_interleave(n, dim=0).contiguous() - v_s = v_s.repeat_interleave(n, dim=0).contiguous() - - # k_pos 用于因果掩码 - k_pos = torch.arange(q_len, device=query.device) - - # Q-tiling：分块处理 query，峰值内存 O(_Q_CHUNK × q_len) - for qc_start in range(0, q_len, _Q_CHUNK): - qc_end = min(qc_start + _Q_CHUNK, q_len) - - # [H, qc, D] - q_c = q_flat[seq_start + qc_start:seq_start + qc_end] \ - .permute(1, 0, 2).float() - - # [H, qc, q_len] - attn_w = torch.matmul(q_c, k_s.transpose(-2, -1)) * self.scale - - # 因果掩码：q_c 里位置 j 只能看 k_pos <= j（相对位置） - qc_q_pos = torch.arange(qc_start, qc_end, device=query.device) - mask = k_pos.unsqueeze(0) > qc_q_pos.unsqueeze(1) - attn_w = attn_w.masked_fill(mask.unsqueeze(0), float("-inf")) - - attn_w = torch.softmax(attn_w, dim=-1) - out_c = torch.matmul(attn_w, v_s).to(orig_dtype) # [H, qc, D] - - output[seq_start + qc_start:seq_start + qc_end] = ( - out_c.permute(1, 0, 2)) - - seq_start = seq_end - - return output.unsqueeze(0) # [1, T, H, D] - -''' - -OLD_XFORMER_BLOCK = """\ - self.attn_op = xops.fmha.flash.FwOp() - if self.alibi_slopes is None: - # Add the batch dimension. - query = query.unsqueeze(0) - key = key.unsqueeze(0) - value = value.unsqueeze(0) - out = xops.memory_efficient_attention_forward( - query, - key, - value, - attn_bias=attn_bias[0], - p=0.0, - scale=self.scale, - op = self.attn_op - ) - return out.view_as(original_query)\ -""" - -NEW_XFORMER_BLOCK = """\ - self.attn_op = xops.fmha.flash.FwOp() - if self.alibi_slopes is None: - # Add the batch dimension. - query = query.unsqueeze(0) - key = key.unsqueeze(0) - value = value.unsqueeze(0) - if self.head_size > 128: - out = self._run_sdpa_fallback(query, key, value, attn_metadata) - else: - out = xops.memory_efficient_attention_forward( - query, - key, - value, - attn_bias=attn_bias[0], - p=0.0, - scale=self.scale, - op=self.attn_op, - ) - return out.view_as(original_query)\ -""" - -INJECT_ANCHOR = " def _run_memory_efficient_xformers_forward(" - - -def patch_file(path): - with open(path, "r") as f: - content = f.read() - changed = False - - if "_run_sdpa_fallback" in content: - print(" [skip] _run_sdpa_fallback already present") - elif INJECT_ANCHOR not in content: - print(" [warn] inject anchor not found") - else: - content = content.replace(INJECT_ANCHOR, FALLBACK_METHOD + INJECT_ANCHOR, 1) - print(" [ok] injected _run_sdpa_fallback (sequential, pure-math)") - changed = True - - if NEW_XFORMER_BLOCK in content: - print(" [skip] dispatch block already patched") - elif OLD_XFORMER_BLOCK in content: - content = content.replace(OLD_XFORMER_BLOCK, NEW_XFORMER_BLOCK, 1) - print(" [ok] patched dispatch block") - changed = True - else: - print(" [warn] dispatch block anchor not found") - - if changed: - with open(path, "w") as f: - f.write(content) - print(f" Written: {path}") - - -def patch_arg_utils(path): - with open(path, "r") as f: - content = f.read() - changed = False - - if "skip auto-enable: Q-tiling" in content: - print(" [skip] chunked-prefill auto-enable already disabled") - elif _ARG_OLD_BLOCK in content: - content = content.replace(_ARG_OLD_BLOCK, _ARG_NEW_BLOCK, 1) - print(" [ok] disabled chunked-prefill auto-enable for 32K+") - changed = True - else: - print(" [warn] target block not found — check arg_utils.py version") - - if changed: - with open(path, "w") as f: - f.write(content) - print(f" Written: {path}") - - -def patch_logits_processor(path): - with open(path, "r") as f: - content = f.read() - changed = False - - if "intermediate chunked-prefill chunk" in content: - print(" [skip] seq_groups=None guard already present") - elif _LP_OLD_BLOCK in content: - content = content.replace(_LP_OLD_BLOCK, _LP_NEW_BLOCK, 1) - print(" [ok] added seq_groups=None guard in _apply_logits_processors") - changed = True - else: - print(" [warn] target block not found — check logits_processor.py version") - - if changed: - with open(path, "w") as f: - f.write(content) - print(f" Written: {path}") - - -def main(): - print("=== patch_xformers_sdpa_seq (sequential, pure-math) ===") - print(f"Target: {XFORMERS_PATH}") - patch_file(XFORMERS_PATH) - - print("\n=== patch_arg_utils (disable chunked-prefill auto-enable) ===") - print(f"Target: {ARG_UTILS_PATH}") - patch_arg_utils(ARG_UTILS_PATH) - - print("\n=== patch_logits_processor (seq_groups=None guard for chunked prefill) ===") - print(f"Target: {LOGITS_PROC_PATH}") - patch_logits_processor(LOGITS_PROC_PATH) - - print("\nDone.") - - -if __name__ == "__main__": - main() diff --git a/qwen3_6_scripts/patch_xformers_sdpa_seq_kernel.py b/qwen3_6_scripts/patch_xformers_sdpa_seq_kernel.py deleted file mode 100644 index 82df8d09..00000000 --- a/qwen3_6_scripts/patch_xformers_sdpa_seq_kernel.py +++ /dev/null @@ -1,181 +0,0 @@ -""" -策略：顺序（per-sequence）— F.scaled_dot_product_attention，可走硬件 kernel -============================================================================= -逐条序列调用 F.scaled_dot_product_attention，is_causal=False + 显式因果 mask。 -与 patch_xformers_sdpa_seq.py（纯 matmul）的区别： - SDPA 可分发到 Flash Attention / mem-efficient attention kernel， - 而纯 matmul 固定走 cublas。 - -硬件限制（BI-V100）： - cudnnFlashAttnForward 不支持 is_causal=True（直接报错）。 - 必须使用 is_causal=False + 显式 additive causal mask。 - 每条序列单独构造上三角 -inf mask，peak 显存 = max(seq_len)² × dtype， - 比 batch 版的 total_tokens² 小得多。 - -与 batch_kernel 的对比： - seq_kernel: 显存小，peak = max_single_seq²；并发 prefill 串行排队 - batch_kernel: 显存大，peak = total_tokens²；并发 prefill 一次并行处理， - 通过 --max-num-batched-tokens 控制 total_tokens 上限 - -Deploy: - python3 modified_scripts/patch_xformers_sdpa_seq_kernel.py -""" - -XFORMERS_PATH = ( - "/usr/local/corex/lib64/python3/dist-packages/" - "vllm/attention/backends/xformers.py" -) - -FALLBACK_METHOD = ''' - def _run_sdpa_fallback( - self, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attn_metadata: "XFormersMetadata", - ) -> torch.Tensor: - """顺序 F.scaled_dot_product_attention fallback（可走硬件 kernel）。 - - 逐条序列调用 SDPA，is_causal=False + 显式上三角 additive mask。 - cudnnFlashAttnForward 不支持 is_causal=True，必须用显式 mask。 - 逐序列构造 mask，peak 显存 = max(seq_len)² × dtype（远小于 batch 版）。 - - Args: - query : [1, total_prefill_tokens, num_heads, head_dim] - key : [1, total_prefill_tokens, num_kv_heads, head_dim] - value : [1, total_prefill_tokens, num_kv_heads, head_dim] - Returns: - [1, total_prefill_tokens, num_heads, head_dim] - """ - import torch.nn.functional as F - - assert attn_metadata.seq_lens is not None - orig_dtype = query.dtype - - q_flat = query.squeeze(0) # [T, H, D] - k_flat = key.squeeze(0) # [T, Hkv, D] - v_flat = value.squeeze(0) - - output = torch.empty_like(q_flat) - start = 0 - for seq_len in attn_metadata.seq_lens: - end = start + seq_len - # [1, H, L, D] - q_s = q_flat[start:end].permute(1, 0, 2).contiguous().unsqueeze(0) - k_s = k_flat[start:end].permute(1, 0, 2).contiguous().unsqueeze(0) - v_s = v_flat[start:end].permute(1, 0, 2).contiguous().unsqueeze(0) - - # GQA：展开 KV heads - if k_s.shape[1] != q_s.shape[1]: - n = q_s.shape[1] // k_s.shape[1] - k_s = k_s.repeat_interleave(n, dim=1).contiguous() - v_s = v_s.repeat_interleave(n, dim=1).contiguous() - - # 逐序列因果 mask [L, L]，上三角 -inf - causal_mask = torch.tril( - torch.zeros(seq_len, seq_len, dtype=orig_dtype, device=q_s.device) - ) - causal_mask = causal_mask.masked_fill( - torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, - device=q_s.device), diagonal=1), - float("-inf"), - ) - - # is_causal=False + 显式 mask，规避 cudnnFlashAttnForward 不支持 is_causal=True - out_s = F.scaled_dot_product_attention( - q_s, k_s, v_s, - attn_mask=causal_mask, - dropout_p=0.0, - is_causal=False, - scale=self.scale, - ) - # [1, H, L, D] → [L, H, D] - output[start:end] = out_s.squeeze(0).permute(1, 0, 2).to(orig_dtype) - start = end - - return output.unsqueeze(0) # [1, T, H, D] - -''' - -OLD_XFORMER_BLOCK = """\ - self.attn_op = xops.fmha.flash.FwOp() - if self.alibi_slopes is None: - # Add the batch dimension. - query = query.unsqueeze(0) - key = key.unsqueeze(0) - value = value.unsqueeze(0) - out = xops.memory_efficient_attention_forward( - query, - key, - value, - attn_bias=attn_bias[0], - p=0.0, - scale=self.scale, - op = self.attn_op - ) - return out.view_as(original_query)\ -""" - -NEW_XFORMER_BLOCK = """\ - self.attn_op = xops.fmha.flash.FwOp() - if self.alibi_slopes is None: - # Add the batch dimension. - query = query.unsqueeze(0) - key = key.unsqueeze(0) - value = value.unsqueeze(0) - if self.head_size > 128: - out = self._run_sdpa_fallback(query, key, value, attn_metadata) - else: - out = xops.memory_efficient_attention_forward( - query, - key, - value, - attn_bias=attn_bias[0], - p=0.0, - scale=self.scale, - op=self.attn_op, - ) - return out.view_as(original_query)\ -""" - -INJECT_ANCHOR = " def _run_memory_efficient_xformers_forward(" - - -def patch_file(path): - with open(path, "r") as f: - content = f.read() - changed = False - - if "_run_sdpa_fallback" in content: - print(" [skip] _run_sdpa_fallback already present") - elif INJECT_ANCHOR not in content: - print(" [warn] inject anchor not found") - else: - content = content.replace(INJECT_ANCHOR, FALLBACK_METHOD + INJECT_ANCHOR, 1) - print(" [ok] injected _run_sdpa_fallback (seq, F.sdpa kernel)") - changed = True - - if NEW_XFORMER_BLOCK in content: - print(" [skip] dispatch block already patched") - elif OLD_XFORMER_BLOCK in content: - content = content.replace(OLD_XFORMER_BLOCK, NEW_XFORMER_BLOCK, 1) - print(" [ok] patched dispatch block") - changed = True - else: - print(" [warn] dispatch block anchor not found") - - if changed: - with open(path, "w") as f: - f.write(content) - print(f" Written: {path}") - - -def main(): - print("=== patch_xformers_sdpa_seq_kernel (seq, F.sdpa + kernel dispatch) ===") - print(f"Target: {XFORMERS_PATH}") - patch_file(XFORMERS_PATH) - print("\nDone.") - - -if __name__ == "__main__": - main()