73 lines
2.7 KiB
Python
73 lines
2.7 KiB
Python
|
"""
|
||
|
|
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()
|