File size: 3,954 Bytes
37f197f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | from __future__ import annotations
import argparse
from pathlib import Path
ORIGINAL_MOE_BLOCK = """ # Prefer to use the MarlinMoE kernel when it is supported.\n if (\n not check_moe_marlin_supports_layer(layer, group_size)\n or current_platform.is_rocm()\n ):\n"""
GENERIC_MOE_BLOCK = """ # Prefer the generic WNA16 path on CUDA for now. The Marlin MoE\n # repack op can fail at load time with PTX toolchain mismatches on\n # some environments even when the layer is otherwise supported.\n if True:\n"""
def find_vllm_dir() -> Path:
import vllm # type: ignore
return Path(vllm.__file__).resolve().parent
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Patch the active vLLM install for compressed Sarvam inference."
)
parser.add_argument(
"--moe-kernel",
choices=("generic", "marlin"),
default="marlin",
help="Select the WNA16 MoE kernel path to enable in vLLM.",
)
return parser.parse_args()
def replace_once(path: Path, old: str, new: str, marker: str) -> None:
text = path.read_text(encoding="utf-8")
if marker in text:
print(f"already patched: {path}")
return
if old not in text:
raise RuntimeError(f"expected block not found in {path}")
path.write_text(text.replace(old, new, 1), encoding="utf-8")
print(f"patched {path}")
def replace_either(path: Path, first_old: str, second_old: str, new: str) -> None:
text = path.read_text(encoding="utf-8")
if new in text:
print(f"already patched: {path}")
return
if first_old in text:
path.write_text(text.replace(first_old, new, 1), encoding="utf-8")
print(f"patched {path}")
return
if second_old in text:
path.write_text(text.replace(second_old, new, 1), encoding="utf-8")
print(f"patched {path}")
return
raise RuntimeError(f"expected block not found in {path}")
def patch_compressed_tensors_moe(vllm_dir: Path, moe_kernel: str) -> None:
path = (
vllm_dir
/ "model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py"
)
new = GENERIC_MOE_BLOCK if moe_kernel == "generic" else ORIGINAL_MOE_BLOCK
replace_either(path, ORIGINAL_MOE_BLOCK, GENERIC_MOE_BLOCK, new)
def patch_fused_moe_loader(vllm_dir: Path) -> None:
path = vllm_dir / "model_executor/layers/fused_moe/fused_moe.py"
old = """ # If a configuration has been found, return it\n tuned_config = json.load(f)\n # Delete triton_version from tuned_config\n tuned_config.pop(\"triton_version\", None)\n return {int(key): val for key, val in tuned_config.items()}\n"""
new = """ # If a configuration has been found, return it\n tuned_config = json.load(f)\n if not isinstance(tuned_config, dict):\n logger.warning_once(\n \"Ignoring malformed MoE tuned config at %s with type %s; \"\n \"falling back to the default MoE config.\",\n config_file_path,\n type(tuned_config).__name__,\n scope=\"global\",\n )\n continue\n # Delete triton_version from tuned_config\n tuned_config.pop(\"triton_version\", None)\n return {int(key): val for key, val in tuned_config.items()}\n"""
replace_once(
path, old, new, "Ignoring malformed MoE tuned config at %s with type %s;"
)
def main() -> None:
args = parse_args()
vllm_dir = find_vllm_dir()
patch_compressed_tensors_moe(vllm_dir, args.moe_kernel)
patch_fused_moe_loader(vllm_dir)
if __name__ == "__main__":
main()
|