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