78 lines
2.1 KiB
Python
78 lines
2.1 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
"""
|
|||
|
|
修复 CoreX 版 vLLM 0.17.0 的 Llama 4 chunked local attention bug。
|
|||
|
|
|
|||
|
|
症状（第一个真实请求时）：
|
|||
|
|
TypeError: CommonAttentionMetadata.__init__() missing 2 required positional
|
|||
|
|
arguments: 'key_start_loc' and 'seq_lens_np'
|
|||
|
|
|
|||
|
|
原因：
|
|||
|
|
vllm/v1/attention/backends/utils.py 的
|
|||
|
|
make_local_attention_virtual_batches() 在构造 CommonAttentionMetadata 时
|
|||
|
|
没有传新版本 dataclass 要求的 key_start_loc / seq_lens_np。
|
|||
|
|
|
|||
|
|
本脚本幂等：已经打过补丁会直接返回 0。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
import sys
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
import vllm
|
|||
|
|
except ImportError as e:
|
|||
|
|
print(f"ERROR: 无法 import vllm: {e}", file=sys.stderr)
|
|||
|
|
sys.exit(1)
|
|||
|
|
|
|||
|
|
TARGET = os.path.join(
|
|||
|
|
os.path.dirname(vllm.__file__),
|
|||
|
|
"v1",
|
|||
|
|
"attention",
|
|||
|
|
"backends",
|
|||
|
|
"utils.py",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
MARKER = "seq_lens_np=seqlens_k_local"
|
|||
|
|
|
|||
|
|
OLD = """ query_start_loc=query_start_loc_cpu.to(device=device, non_blocking=True),
|
|||
|
|
seq_lens=seq_lens_cpu.to(device=device, non_blocking=True),"""
|
|||
|
|
|
|||
|
|
NEW = """ query_start_loc=query_start_loc_cpu.to(device=device, non_blocking=True),
|
|||
|
|
key_start_loc=torch.from_numpy(
|
|||
|
|
np.concatenate(
|
|||
|
|
(np.zeros(1, dtype=np.int32), np.cumsum(seqlens_k_local, dtype=np.int32))
|
|||
|
|
)
|
|||
|
|
).to(device=device, non_blocking=True),
|
|||
|
|
seq_lens=seq_lens_cpu.to(device=device, non_blocking=True),
|
|||
|
|
seq_lens_np=seqlens_k_local,"""
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main() -> int:
|
|||
|
|
if not os.path.exists(TARGET):
|
|||
|
|
print(f"ERROR: 找不到目标文件: {TARGET}", file=sys.stderr)
|
|||
|
|
return 1
|
|||
|
|
|
|||
|
|
with open(TARGET, "r", encoding="utf-8") as f:
|
|||
|
|
src = f.read()
|
|||
|
|
|
|||
|
|
if MARKER in src:
|
|||
|
|
print(f"patch already applied: {TARGET}")
|
|||
|
|
return 0
|
|||
|
|
|
|||
|
|
if OLD not in src:
|
|||
|
|
print(
|
|||
|
|
"ERROR: 目标代码块未找到，vLLM 版本可能不同。\n"
|
|||
|
|
f"请手动检查: {TARGET}",
|
|||
|
|
file=sys.stderr,
|
|||
|
|
)
|
|||
|
|
return 1
|
|||
|
|
|
|||
|
|
with open(TARGET, "w", encoding="utf-8") as f:
|
|||
|
|
f.write(src.replace(OLD, NEW, 1))
|
|||
|
|
|
|||
|
|
print(f"patched: {TARGET}")
|
|||
|
|
return 0
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
sys.exit(main())
|