118 lines
3.3 KiB
Python
Executable File
118 lines
3.3 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Apply the CoreX vLLM patches needed by Llama 4 Scout / Maverick.
|
|
|
|
The patches in ./patches were generated against the exact base image
|
|
`registry.iluvatar.com.cn:10443/customer/sz/vllm0.17.0-4.4.0-x86:v7`
|
|
(vLLM 0.17.0+corex.20260612070235).
|
|
|
|
Scout needs only the chunked-attention patch. Maverick additionally needs the
|
|
FP8 quantization / FP8 MoE patches and the `corexx_fp8_moe` fallback module.
|
|
One image contains the union of both patch sets; the FP8 paths are only used
|
|
when the mounted model is FP8.
|
|
|
|
The script is idempotent: it checks for a marker in each target file and skips
|
|
patches that are already applied.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import ast
|
|
import os
|
|
import pathlib
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
|
|
PATCH_DIR = pathlib.Path(
|
|
os.environ.get("LLAMA4_PATCH_DIR", "/opt/llama4/patches")
|
|
)
|
|
|
|
# LLAMA4_PKG_ROOT is only needed for local testing against a copy of the
|
|
# vendor files. In the image it is derived from the installed vllm package.
|
|
_pkg_env = os.environ.get("LLAMA4_PKG_ROOT")
|
|
if _pkg_env:
|
|
PKG = pathlib.Path(_pkg_env)
|
|
else:
|
|
import vllm
|
|
|
|
PKG = pathlib.Path(vllm.__file__).resolve().parent.parent
|
|
|
|
# patch file -> (target file relative to the package root, applied marker)
|
|
TARGETS = {
|
|
"01_fp8_quant.patch": (
|
|
"vllm/_custom_ops.py",
|
|
"COREXX_FP8_QUANT_PATCH",
|
|
),
|
|
"02_fp8_moe.patch": (
|
|
"vllm/model_executor/layers/fused_moe/fused_moe.py",
|
|
"COREXX_FP8_MOE_PATCH",
|
|
),
|
|
"03_llama4_chunked_attn.patch": (
|
|
"vllm/v1/attention/backends/utils.py",
|
|
"seq_lens_np=seqlens_k_local",
|
|
),
|
|
}
|
|
|
|
|
|
def apply_one(patch_name: str, target_rel: str, marker: str) -> None:
|
|
target = PKG / target_rel
|
|
if not target.exists():
|
|
sys.exit(f"target file not found: {target}")
|
|
if marker in target.read_text():
|
|
print(f"[skip] {target} (already patched)")
|
|
return
|
|
|
|
patch_file = PATCH_DIR / patch_name
|
|
if not patch_file.exists():
|
|
sys.exit(f"missing patch file: {patch_file}")
|
|
|
|
result = subprocess.run(
|
|
[
|
|
"patch",
|
|
"-p1",
|
|
"--forward",
|
|
"--batch",
|
|
"--no-backup-if-mismatch",
|
|
"-d",
|
|
str(PKG),
|
|
"-i",
|
|
str(patch_file),
|
|
],
|
|
text=True,
|
|
capture_output=True,
|
|
)
|
|
if result.stdout:
|
|
print(result.stdout, end="")
|
|
if result.returncode != 0:
|
|
if result.stderr:
|
|
print(result.stderr, file=sys.stderr)
|
|
sys.exit(f"patch failed: {patch_name}")
|
|
if marker not in target.read_text():
|
|
sys.exit(f"patch {patch_name} applied but marker not found in {target}")
|
|
print(f"[ok] {target}")
|
|
|
|
|
|
def main() -> None:
|
|
for patch_name, (target_rel, marker) in TARGETS.items():
|
|
apply_one(patch_name, target_rel, marker)
|
|
|
|
src = PATCH_DIR / "corexx_fp8_moe.py"
|
|
dst = PKG / "vllm/model_executor/layers/fused_moe/corexx_fp8_moe.py"
|
|
if not src.exists():
|
|
sys.exit(f"missing module: {src}")
|
|
shutil.copy2(src, dst)
|
|
print(f"[copy] {src} -> {dst}")
|
|
|
|
for path in (
|
|
PKG / "vllm/_custom_ops.py",
|
|
PKG / "vllm/model_executor/layers/fused_moe/fused_moe.py",
|
|
dst,
|
|
PKG / "vllm/v1/attention/backends/utils.py",
|
|
):
|
|
ast.parse(path.read_text())
|
|
print("all patches applied; syntax OK")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|