dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

[CLEANUP] Remove 13 dead patch scripts — only 1 remains (transformers registration)

Browse Source 
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.
This commit is contained in:
dylanyunlon
2026-08-05 08:39:54 +00:00
parent b902090fb2
commit 821c59500d
13 changed files with 0 additions and 2154 deletions
Download Patch File Download Diff File Expand all files Collapse all files

195
qwen3_6_scripts/patch_enable_triton.py
View File

@@ -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()

182
qwen3_6_scripts/patch_head256_triton.py
View File

@@ -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()

304
qwen3_6_scripts/patch_ixformer_native.py
View File

@@ -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()

78
qwen3_6_scripts/patch_model_runner.py
View File

@@ -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)

192
qwen3_6_scripts/patch_paged_attention_v2.py
View File

@@ -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()

91
qwen3_6_scripts/patch_triton_tuning.py
View File

@@ -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()

72
qwen3_6_scripts/patch_vectorized_decode.py
View File

@@ -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()

76
qwen3_6_scripts/patch_vllm_qwen3_5.py
View File

@@ -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()

79
qwen3_6_scripts/patch_vllm_tool_parser.py
View File

@@ -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()

192
qwen3_6_scripts/patch_xformers_sdpa_batch.py
View File

@@ -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()

191
qwen3_6_scripts/patch_xformers_sdpa_batch_kernel.py
View File

@@ -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()

321
qwen3_6_scripts/patch_xformers_sdpa_seq.py
View File

@@ -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()

181
qwen3_6_scripts/patch_xformers_sdpa_seq_kernel.py
View File

@@ -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()