diff --git a/pyproject.toml b/pyproject.toml index ccc8aaca..71e662d2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,10 +17,22 @@ dev = ["pre-commit"] [project.entry-points."vllm.general_plugins"] gguf = "vllm_gguf_plugin:register" +[tool.setuptools] +include-package-data = true + [tool.setuptools.packages.find] where = ["."] include = ["vllm_gguf_plugin*"] +[tool.setuptools.package-data] +vllm_gguf_plugin = [ + "csrc/*.h", + "csrc/*.cpp", + "csrc/gguf/*.h", + "csrc/gguf/*.cu", + "csrc/gguf/*.cuh", +] + [tool.ruff.lint] select = [ # pycodestyle diff --git a/scripts/build_jit_cache_wheel.sh b/scripts/build_jit_cache_wheel.sh new file mode 100755 index 00000000..e46534a3 --- /dev/null +++ b/scripts/build_jit_cache_wheel.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +# Build the optional jit-cache wheel with CUDA arch coverage that mirrors +# scripts/build_release_wheel.sh for the current nvcc. +set -euo pipefail + +if ! command -v nvcc >/dev/null 2>&1; then + echo "error: nvcc not found on PATH" >&2 + exit 1 +fi + +cuda_release=$(nvcc --version | grep -oE 'release [0-9]+\.[0-9]+' | awk '{print $2}') +cuda_major=${cuda_release%.*} +cuda_minor=${cuda_release#*.} + +if [ "$cuda_major" -ge 13 ]; then + export TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;8.7;8.9;9.0;10.0;11.0;12.0" +elif [ "$cuda_major" -ge 12 ] && [ "$cuda_minor" -ge 8 ]; then + export TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.3;12.0;12.1" +else + export TORCH_CUDA_ARCH_LIST="7.0;7.5;8.0;8.6;8.7;8.9;9.0" +fi + +echo "CUDA $cuda_release; TORCH_CUDA_ARCH_LIST=$TORCH_CUDA_ARCH_LIST" + +exec uv build ./vllm-gguf-plugin-jit-cache --wheel --no-build-isolation "$@" diff --git a/setup.py b/setup.py index 34379a1f..97cfbf4f 100644 --- a/setup.py +++ b/setup.py @@ -1,47 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -import sys - from setuptools import setup - -def _should_build_extension() -> bool: - packaging_commands = {"sdist", "egg_info", "dist_info"} - return not any(command in packaging_commands for command in sys.argv[1:]) - - -setup_kwargs: dict = {} - -if _should_build_extension(): - from torch.utils.cpp_extension import BuildExtension, CUDAExtension - - setup_kwargs.update( - ext_modules=[ - CUDAExtension( - name="vllm_gguf_plugin._C_gguf", - sources=[ - "vllm_gguf_plugin/csrc/torch_bindings.cpp", - "vllm_gguf_plugin/csrc/gguf/gguf_kernel.cu", - ], - include_dirs=[ - "vllm_gguf_plugin/csrc", - "vllm_gguf_plugin/csrc/gguf", - ], - py_limited_api=True, - extra_compile_args={ - "cxx": ["-O3", "-std=c++17"], - "nvcc": [ - "-O3", - "-std=c++17", - "--use_fast_math", - # Exposes aoti_torch_get_current_cuda_stream in the AOTI shim. - "-DUSE_CUDA", - ], - }, - ) - ], - cmdclass={"build_ext": BuildExtension}, - options={"bdist_wheel": {"py_limited_api": "cp310"}}, - ) - -setup(**setup_kwargs) +setup() diff --git a/tests/test_plugin.py b/tests/test_plugin.py index 70fd12d6..ea1d0089 100644 --- a/tests/test_plugin.py +++ b/tests/test_plugin.py @@ -1,5 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 +import importlib.util +import re +import sys +from pathlib import Path + +import pytest import torch import vllm.engine.arg_utils as arg_utils_module import vllm.model_executor.layers.vocab_parallel_embedding as vocab_embedding_module @@ -18,6 +24,7 @@ from vllm.model_executor.model_loader import get_model_loader from vllm.transformers_utils.config import get_config_parser +import vllm_gguf_plugin._jit as jit_module import vllm_gguf_plugin.config_parser as gguf_config_parser_module import vllm_gguf_plugin.quantization as gguf_quantization from vllm_gguf_plugin import OOTGGUFConfig, OOTGGUFModelLoader, register @@ -314,3 +321,145 @@ def test_gguf_linear_preserves_cuda_weight_device(monkeypatch): assert layer.qweight.device.type == "cuda" assert layer.qweight_type.device.type == "cuda" + + +def test_gguf_cuda_extension_uses_jit_loader(monkeypatch): + state = {"loaded": False} + captured = {} + + monkeypatch.setattr(jit_module, "_gguf_ops_available", lambda: state["loaded"]) + monkeypatch.setattr(jit_module, "_precompiled_gguf_library_paths", lambda: []) + monkeypatch.setattr(jit_module.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(jit_module.torch.version, "cuda", "12.9", raising=False) + monkeypatch.setattr(jit_module.cpp_extension, "CUDA_HOME", "/usr/local/cuda") + + def fake_load(**kwargs): + state["loaded"] = True + captured.update(kwargs) + return object() + + monkeypatch.setattr(jit_module.cpp_extension, "load", fake_load) + + jit_module.ensure_gguf_cuda_ops_loaded() + jit_module.ensure_gguf_cuda_ops_loaded() + + assert captured["name"] == "_C_gguf" + assert captured["with_cuda"] is True + assert captured["sources"] == [ + str(jit_module._csrc_root() / "torch_bindings.cpp"), + str(jit_module._csrc_root() / "gguf" / "gguf_kernel.cu"), + ] + assert captured["extra_include_paths"] == [ + str(jit_module._csrc_root()), + str(jit_module._csrc_root() / "gguf"), + ] + assert captured["extra_cuda_cflags"] == [ + "-O3", + "-std=c++17", + "--use_fast_math", + "-DUSE_CUDA", + ] + + +def test_gguf_cuda_extension_prefers_precompiled_library(monkeypatch, tmp_path): + state = {"loaded": False} + library_path = tmp_path / "_C_gguf.so" + library_path.write_bytes(b"") + loaded_paths = [] + + monkeypatch.setattr(jit_module, "_gguf_ops_available", lambda: state["loaded"]) + monkeypatch.setattr( + jit_module, "_precompiled_gguf_library_paths", lambda: [library_path] + ) + monkeypatch.setattr(jit_module.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(jit_module.torch.version, "cuda", "12.9", raising=False) + monkeypatch.setattr(jit_module.cpp_extension, "CUDA_HOME", None) + + def fake_load_library(path: str): + loaded_paths.append(path) + state["loaded"] = True + + monkeypatch.setattr(jit_module.torch.ops, "load_library", fake_load_library) + monkeypatch.setattr( + jit_module.cpp_extension, + "load", + lambda **kwargs: pytest.fail(f"unexpected JIT compile: {kwargs}"), + ) + + jit_module.ensure_gguf_cuda_ops_loaded() + jit_module.ensure_gguf_cuda_ops_loaded() + + assert loaded_paths == [str(library_path)] + + +def test_gguf_cuda_extension_requires_cuda_device(monkeypatch): + monkeypatch.setattr(jit_module, "_gguf_ops_available", lambda: False) + monkeypatch.setattr(jit_module.torch.cuda, "is_available", lambda: False) + + with pytest.raises(RuntimeError, match="available CUDA device"): + jit_module.ensure_gguf_cuda_ops_loaded() + + +def test_gguf_precompiled_artifact_tag(monkeypatch): + monkeypatch.setattr(jit_module.platform, "system", lambda: "Linux") + monkeypatch.setattr(jit_module.platform, "machine", lambda: "x86_64") + monkeypatch.setattr(jit_module.torch, "__version__", "2.11.0+cu129", raising=False) + monkeypatch.setattr(jit_module.torch.version, "cuda", "12.9", raising=False) + + assert jit_module.get_gguf_precompiled_artifact_tag() == ( + "linux-x86_64/torch-2.11.0/cuda-12.9/" + f"python-cp{sys.version_info.major}{sys.version_info.minor}" + ) + + +def test_jit_cache_wheel_backend_writes_build_meta(): + backend_path = ( + Path(__file__).resolve().parent.parent + / "vllm-gguf-plugin-jit-cache" + / "build_backend.py" + ) + spec = importlib.util.spec_from_file_location("jit_cache_backend", backend_path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + root_pyproject = ( + Path(__file__).resolve().parent.parent / "pyproject.toml" + ).read_text(encoding="utf-8") + expected_version = re.search( + r'^version = "([^"]+)"$', root_pyproject, re.MULTILINE + ).group(1) + version = module._write_build_meta() + build_meta_path = ( + Path(__file__).resolve().parent.parent + / "vllm-gguf-plugin-jit-cache" + / "vllm_gguf_plugin_precompiled" + / "_build_meta.py" + ) + + assert version == expected_version + assert ( + build_meta_path.read_text(encoding="utf-8") + .strip() + .endswith(f'__version__ = "{expected_version}"') + ) + + +def test_jit_cache_wheel_backend_artifact_dir_matches_runtime_tag(): + backend_path = ( + Path(__file__).resolve().parent.parent + / "vllm-gguf-plugin-jit-cache" + / "build_backend.py" + ) + spec = importlib.util.spec_from_file_location("jit_cache_backend", backend_path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + assert module._artifact_output_dir() == ( + Path(__file__).resolve().parent.parent + / "vllm-gguf-plugin-jit-cache" + / "vllm_gguf_plugin_precompiled" + / "artifacts" + / jit_module.get_gguf_precompiled_artifact_tag() + ) diff --git a/vllm-gguf-plugin-jit-cache/build_backend.py b/vllm-gguf-plugin-jit-cache/build_backend.py new file mode 100644 index 00000000..94c7c0e2 --- /dev/null +++ b/vllm-gguf-plugin-jit-cache/build_backend.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: Apache-2.0 + +import importlib.util +import re +import shutil +from pathlib import Path + +from setuptools import build_meta as _orig +from wheel.bdist_wheel import bdist_wheel + + +_REPO_ROOT = Path(__file__).resolve().parent.parent +_SUBPROJECT_ROOT = Path(__file__).resolve().parent +_PACKAGE_ROOT = _SUBPROJECT_ROOT / "vllm_gguf_plugin_precompiled" +_BUILD_META_PATH = _PACKAGE_ROOT / "_build_meta.py" + + +def _load_jit_module(): + spec = importlib.util.spec_from_file_location( + "vllm_gguf_plugin_jit", + _REPO_ROOT / "vllm_gguf_plugin" / "_jit.py", + ) + if spec is None or spec.loader is None: + raise RuntimeError("Failed to load JIT helper module for cache wheel build.") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _read_main_package_version() -> str: + pyproject_text = (_REPO_ROOT / "pyproject.toml").read_text(encoding="utf-8") + match = re.search(r'^version = "([^"]+)"$', pyproject_text, re.MULTILINE) + if match is None: + raise RuntimeError("Could not determine main package version from pyproject.toml.") + return match.group(1) + + +def _write_build_meta() -> str: + version = _read_main_package_version() + _BUILD_META_PATH.write_text( + '"""Build metadata for vllm-gguf-plugin-jit-cache."""\n' + f'__version__ = "{version}"\n', + encoding="utf-8", + ) + return version + + +def _artifact_output_dir() -> Path: + jit_module = _load_jit_module() + return _PACKAGE_ROOT / "artifacts" / jit_module.get_gguf_precompiled_artifact_tag() + + +def _build_precompiled_artifact() -> None: + jit_module = _load_jit_module() + if jit_module.torch.version.cuda is None: + raise RuntimeError("A CUDA-enabled PyTorch build is required to build the JIT cache wheel.") + if jit_module.cpp_extension.CUDA_HOME is None: + raise RuntimeError("CUDA toolkit not found. Set CUDA_HOME before building the JIT cache wheel.") + + artifacts_root = _PACKAGE_ROOT / "artifacts" + if artifacts_root.exists(): + shutil.rmtree(artifacts_root) + + artifact_output_dir = _artifact_output_dir() + artifact_output_dir.mkdir(parents=True, exist_ok=True) + + build_directory = _REPO_ROOT / "build" / "jit-cache-wheel" + build_directory.mkdir(parents=True, exist_ok=True) + + module = jit_module.cpp_extension.load( + name=jit_module._JIT_EXTENSION_NAME, + sources=jit_module._extension_sources(), + extra_cflags=["-O3", "-std=c++17"], + extra_cuda_cflags=["-O3", "-std=c++17", "--use_fast_math", "-DUSE_CUDA"], + extra_include_paths=jit_module._include_paths(), + build_directory=str(build_directory), + verbose=False, + with_cuda=True, + ) + + library_path = Path(module.__file__) + if not library_path.is_file(): + raise RuntimeError(f"Expected compiled extension at {library_path}.") + + shutil.copy2(library_path, artifact_output_dir / library_path.name) + + +def _prepare_build() -> None: + _write_build_meta() + _build_precompiled_artifact() + + +class PlatformSpecificBdistWheel(bdist_wheel): + def finalize_options(self): + super().finalize_options() + self.root_is_pure = False + + +class _MonkeyPatchBdistWheel: + def __enter__(self): + from setuptools.command import bdist_wheel as setuptools_bdist_wheel + + self.original_bdist_wheel = setuptools_bdist_wheel.bdist_wheel + setuptools_bdist_wheel.bdist_wheel = PlatformSpecificBdistWheel + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + from setuptools.command import bdist_wheel as setuptools_bdist_wheel + + setuptools_bdist_wheel.bdist_wheel = self.original_bdist_wheel + + +def build_wheel(wheel_directory, config_settings=None, metadata_directory=None): + _prepare_build() + with _MonkeyPatchBdistWheel(): + return _orig.build_wheel(wheel_directory, config_settings, metadata_directory) + + +def build_editable(wheel_directory, config_settings=None, metadata_directory=None): + _prepare_build() + orig_build_editable = getattr(_orig, "build_editable", None) + if orig_build_editable is None: + raise RuntimeError("build_editable not supported by setuptools backend") + return orig_build_editable(wheel_directory, config_settings, metadata_directory) + + +def prepare_metadata_for_build_wheel(metadata_directory, config_settings=None): + _write_build_meta() + with _MonkeyPatchBdistWheel(): + return _orig.prepare_metadata_for_build_wheel( + metadata_directory, config_settings + ) + + +def prepare_metadata_for_build_editable(metadata_directory, config_settings=None): + _write_build_meta() + with _MonkeyPatchBdistWheel(): + return _orig.prepare_metadata_for_build_editable( + metadata_directory, config_settings + ) + + +get_requires_for_build_wheel = _orig.get_requires_for_build_wheel +get_requires_for_build_editable = getattr( + _orig, "get_requires_for_build_editable", None +) diff --git a/vllm-gguf-plugin-jit-cache/pyproject.toml b/vllm-gguf-plugin-jit-cache/pyproject.toml new file mode 100644 index 00000000..5a94ed19 --- /dev/null +++ b/vllm-gguf-plugin-jit-cache/pyproject.toml @@ -0,0 +1,23 @@ +[build-system] +requires = ["setuptools>=77.0.3,<81.0.0", "wheel", "torch>=2.9"] +build-backend = "build_backend" +backend-path = ["."] + +[project] +name = "vllm-gguf-plugin-jit-cache" +dynamic = ["version"] +description = "Optional precompiled JIT cache wheel for vllm-gguf-plugin" +readme = {text = "This package contains a precompiled _C_gguf artifact for a specific runtime tag.", content-type = "text/plain"} +license = "Apache-2.0" +requires-python = ">=3.10" +dependencies = [] + +[tool.setuptools] +packages = ["vllm_gguf_plugin_precompiled"] +include-package-data = true + +[tool.setuptools.dynamic] +version = {attr = "vllm_gguf_plugin_precompiled.__version__"} + +[tool.setuptools.package-data] +vllm_gguf_plugin_precompiled = ["artifacts/**/*.so"] diff --git a/vllm-gguf-plugin-jit-cache/vllm_gguf_plugin_precompiled/__init__.py b/vllm-gguf-plugin-jit-cache/vllm_gguf_plugin_precompiled/__init__.py new file mode 100644 index 00000000..a7d0bf52 --- /dev/null +++ b/vllm-gguf-plugin-jit-cache/vllm_gguf_plugin_precompiled/__init__.py @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 + +from ._build_meta import __version__ + +__all__ = ["__version__"] diff --git a/vllm-gguf-plugin-jit-cache/vllm_gguf_plugin_precompiled/_build_meta.py b/vllm-gguf-plugin-jit-cache/vllm_gguf_plugin_precompiled/_build_meta.py new file mode 100644 index 00000000..429a3f96 --- /dev/null +++ b/vllm-gguf-plugin-jit-cache/vllm_gguf_plugin_precompiled/_build_meta.py @@ -0,0 +1,2 @@ +"""Build metadata for vllm-gguf-plugin-jit-cache.""" +__version__ = "0.1.0" diff --git a/vllm_gguf_plugin/_jit.py b/vllm_gguf_plugin/_jit.py new file mode 100644 index 00000000..fa3cad70 --- /dev/null +++ b/vllm_gguf_plugin/_jit.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 + +import importlib.util +import os +import platform +import re +import sys +from importlib.machinery import EXTENSION_SUFFIXES +from pathlib import Path +from threading import Lock + +import torch +from torch.utils import cpp_extension + +_GGUF_LIBRARY_NAMESPACE = "_C_gguf" +_JIT_EXTENSION_NAME = "_C_gguf" +_PRECOMPILED_ARTIFACTS_PACKAGE = "vllm_gguf_plugin_precompiled" +_BUILD_LOCK = Lock() + + +def _gguf_ops_available() -> bool: + return hasattr(torch.ops, _GGUF_LIBRARY_NAMESPACE) and hasattr( + torch.ops._C_gguf, "ggml_dequantize" + ) + + +def _csrc_root() -> Path: + return Path(__file__).resolve().parent / "csrc" + + +def _extension_sources() -> list[str]: + root = _csrc_root() + return [ + str(root / "torch_bindings.cpp"), + str(root / "gguf" / "gguf_kernel.cu"), + ] + + +def _include_paths() -> list[str]: + root = _csrc_root() + return [str(root), str(root / "gguf")] + + +def _normalize_tag_component(value: str) -> str: + return re.sub(r"[^0-9A-Za-z_.-]+", "-", value).strip("-") + + +def _runtime_platform_tag() -> str: + return "-".join( + [ + _normalize_tag_component(platform.system().lower()), + _normalize_tag_component(platform.machine().lower()), + ] + ) + + +def get_gguf_precompiled_artifact_tag() -> str: + torch_version = _normalize_tag_component(torch.__version__.split("+", 1)[0]) + cuda_version = _normalize_tag_component(torch.version.cuda or "cpu") + python_tag = f"cp{sys.version_info.major}{sys.version_info.minor}" + return ( + f"{_runtime_platform_tag()}/torch-{torch_version}/" + f"cuda-{cuda_version}/python-{python_tag}" + ) + + +def _package_root(package_name: str) -> Path | None: + spec = importlib.util.find_spec(package_name) + if spec is None or spec.submodule_search_locations is None: + return None + return Path(next(iter(spec.submodule_search_locations))) + + +def _precompiled_search_roots() -> list[Path]: + roots: list[Path] = [] + env_root = os.environ.get("VLLM_GGUF_PLUGIN_PRECOMPILED_ROOT") + if env_root: + roots.append(Path(env_root)) + roots.append(Path(__file__).resolve().parent / "precompiled") + package_root = _package_root(_PRECOMPILED_ARTIFACTS_PACKAGE) + if package_root is not None: + roots.append(package_root / "artifacts") + return roots + + +def _precompiled_glob_patterns() -> set[str]: + return {f"{_JIT_EXTENSION_NAME}*{suffix}" for suffix in EXTENSION_SUFFIXES} + + +def _precompiled_gguf_library_paths() -> list[Path]: + explicit_library = os.environ.get("VLLM_GGUF_PLUGIN_PRECOMPILED_LIB") + if explicit_library: + return [Path(explicit_library)] + + matches: list[Path] = [] + tag = get_gguf_precompiled_artifact_tag() + for root in _precompiled_search_roots(): + candidate_dir = root / tag + if not candidate_dir.is_dir(): + continue + for pattern in sorted(_precompiled_glob_patterns()): + matches.extend(sorted(candidate_dir.glob(pattern))) + return matches + + +def _load_precompiled_gguf_library() -> bool: + for library_path in _precompiled_gguf_library_paths(): + torch.ops.load_library(str(library_path)) + if _gguf_ops_available(): + return True + return False + + +def ensure_gguf_cuda_ops_loaded() -> None: + if _gguf_ops_available(): + return + + if not torch.cuda.is_available(): + raise RuntimeError( + "vllm-gguf-plugin CUDA kernels require an available CUDA device." + ) + if torch.version.cuda is None: + raise RuntimeError( + "vllm-gguf-plugin CUDA kernels require a CUDA-enabled PyTorch build." + ) + + with _BUILD_LOCK: + if _gguf_ops_available(): + return + + if _load_precompiled_gguf_library(): + return + + if cpp_extension.CUDA_HOME is None: + raise RuntimeError( + "vllm-gguf-plugin could not find the CUDA toolkit. Set CUDA_HOME " + "before using GGUF CUDA ops or install a matching precompiled " + "artifact." + ) + + build_directory = os.environ.get("VLLM_GGUF_PLUGIN_JIT_BUILD_DIR") + if build_directory: + Path(build_directory).mkdir(parents=True, exist_ok=True) + + cpp_extension.load( + name=_JIT_EXTENSION_NAME, + sources=_extension_sources(), + extra_cflags=["-O3", "-std=c++17"], + extra_cuda_cflags=["-O3", "-std=c++17", "--use_fast_math", "-DUSE_CUDA"], + extra_include_paths=_include_paths(), + build_directory=build_directory, + verbose=os.environ.get("VLLM_GGUF_PLUGIN_JIT_VERBOSE") == "1", + with_cuda=True, + ) diff --git a/vllm_gguf_plugin/ops.py b/vllm_gguf_plugin/ops.py index e90a8297..9431cb4a 100644 --- a/vllm_gguf_plugin/ops.py +++ b/vllm_gguf_plugin/ops.py @@ -2,82 +2,105 @@ import torch +from ._jit import ensure_gguf_cuda_ops_loaded + try: from torch.library import register_fake except ImportError: from torch.library import impl_abstract as register_fake -try: - from . import _C_gguf # noqa: F401 -except ImportError: - _C_gguf = None - - -if hasattr(torch.ops, "_C_gguf") and hasattr(torch.ops._C_gguf, "ggml_dequantize"): - - @register_fake("_C_gguf::ggml_dequantize") - def _ggml_dequantize_fake( - W: torch.Tensor, - quant_type: int, - m: torch.SymInt, - n: torch.SymInt, - dtype: torch.dtype | None = None, - ) -> torch.Tensor: - return torch.empty((m, n), dtype=torch.float16, device=W.device) - - @register_fake("_C_gguf::ggml_mul_mat_vec_a8") - def _ggml_mul_mat_vec_a8_fake( - W: torch.Tensor, - X: torch.Tensor, - quant_type: int, - row: torch.SymInt, - ) -> torch.Tensor: - return torch.empty((X.shape[0], row), dtype=X.dtype, device=W.device) - - @register_fake("_C_gguf::ggml_mul_mat_a8") - def _ggml_mul_mat_a8_fake( - W: torch.Tensor, - X: torch.Tensor, - quant_type: int, - row: torch.SymInt, - ) -> torch.Tensor: - return torch.empty((X.size(0), row), dtype=X.dtype, device=W.device) - - @register_fake("_C_gguf::ggml_moe_a8") - def _ggml_moe_a8_fake( - X: torch.Tensor, - W: torch.Tensor, - sorted_token_ids: torch.Tensor, - expert_ids: torch.Tensor, - num_tokens_post_padded: torch.Tensor, - quant_type: int, - row: torch.SymInt, - top_k: torch.SymInt, - tokens: torch.SymInt, - ) -> torch.Tensor: - return torch.empty( - (X.size(0) * top_k, row), dtype=torch.float16, device=W.device - ) - - -if hasattr(torch.ops, "_C_gguf") and hasattr(torch.ops._C_gguf, "ggml_moe_a8_vec"): - - @register_fake("_C_gguf::ggml_moe_a8_vec") - def _ggml_moe_a8_vec_fake( - X: torch.Tensor, - W: torch.Tensor, - topk_ids: torch.Tensor, - top_k: int, - quant_type: int, - row: torch.SymInt, - tokens: torch.SymInt, - ) -> torch.Tensor: - return torch.empty((X.size(0) * top_k, row), dtype=X.dtype, device=W.device) +_BASE_FAKE_OPS_REGISTERED = False +_VEC_FAKE_OP_REGISTERED = False + + +def _ggml_dequantize_fake( + W: torch.Tensor, + quant_type: int, + m: torch.SymInt, + n: torch.SymInt, + dtype: torch.dtype | None = None, +) -> torch.Tensor: + del quant_type, dtype + return torch.empty((m, n), dtype=torch.float16, device=W.device) + + +def _ggml_mul_mat_vec_a8_fake( + W: torch.Tensor, + X: torch.Tensor, + quant_type: int, + row: torch.SymInt, +) -> torch.Tensor: + del quant_type + return torch.empty((X.shape[0], row), dtype=X.dtype, device=W.device) + + +def _ggml_mul_mat_a8_fake( + W: torch.Tensor, + X: torch.Tensor, + quant_type: int, + row: torch.SymInt, +) -> torch.Tensor: + del quant_type + return torch.empty((X.size(0), row), dtype=X.dtype, device=W.device) + + +def _ggml_moe_a8_fake( + X: torch.Tensor, + W: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + quant_type: int, + row: torch.SymInt, + top_k: torch.SymInt, + tokens: torch.SymInt, +) -> torch.Tensor: + del sorted_token_ids, expert_ids, num_tokens_post_padded, quant_type, tokens + return torch.empty((X.size(0) * top_k, row), dtype=torch.float16, device=W.device) + + +def _ggml_moe_a8_vec_fake( + X: torch.Tensor, + W: torch.Tensor, + topk_ids: torch.Tensor, + top_k: int, + quant_type: int, + row: torch.SymInt, + tokens: torch.SymInt, +) -> torch.Tensor: + del topk_ids, quant_type, tokens + return torch.empty((X.size(0) * top_k, row), dtype=X.dtype, device=W.device) + + +def _maybe_register_fake_ops() -> None: + global _BASE_FAKE_OPS_REGISTERED, _VEC_FAKE_OP_REGISTERED + if not hasattr(torch.ops, "_C_gguf"): + return + + if not _BASE_FAKE_OPS_REGISTERED and hasattr(torch.ops._C_gguf, "ggml_dequantize"): + register_fake("_C_gguf::ggml_dequantize")(_ggml_dequantize_fake) + register_fake("_C_gguf::ggml_mul_mat_vec_a8")(_ggml_mul_mat_vec_a8_fake) + register_fake("_C_gguf::ggml_mul_mat_a8")(_ggml_mul_mat_a8_fake) + register_fake("_C_gguf::ggml_moe_a8")(_ggml_moe_a8_fake) + _BASE_FAKE_OPS_REGISTERED = True + + if not _VEC_FAKE_OP_REGISTERED and hasattr(torch.ops._C_gguf, "ggml_moe_a8_vec"): + register_fake("_C_gguf::ggml_moe_a8_vec")(_ggml_moe_a8_vec_fake) + _VEC_FAKE_OP_REGISTERED = True + + +def _ensure_cuda_ops() -> None: + ensure_gguf_cuda_ops_loaded() + _maybe_register_fake_ops() + + +_maybe_register_fake_ops() def ggml_dequantize( W: torch.Tensor, quant_type: int, m: int, n: int, dtype: torch.dtype | None ) -> torch.Tensor: + _ensure_cuda_ops() return torch.ops._C_gguf.ggml_dequantize(W, quant_type, m, n, dtype) @@ -87,6 +110,7 @@ def ggml_mul_mat_vec_a8( quant_type: int, row: int, ) -> torch.Tensor: + _ensure_cuda_ops() return torch.ops._C_gguf.ggml_mul_mat_vec_a8(W, X, quant_type, row) @@ -96,6 +120,7 @@ def ggml_mul_mat_a8( quant_type: int, row: int, ) -> torch.Tensor: + _ensure_cuda_ops() return torch.ops._C_gguf.ggml_mul_mat_a8(W, X, quant_type, row) @@ -110,6 +135,7 @@ def ggml_moe_a8( top_k: int, tokens: int, ) -> torch.Tensor: + _ensure_cuda_ops() return torch.ops._C_gguf.ggml_moe_a8( X, W, @@ -132,12 +158,14 @@ def ggml_moe_a8_vec( row: int, tokens: int, ) -> torch.Tensor: + _ensure_cuda_ops() return torch.ops._C_gguf.ggml_moe_a8_vec( X, W, topk_ids, top_k, quant_type, row, tokens ) def ggml_moe_get_block_size(quant_type: int) -> int: + _ensure_cuda_ops() return torch.ops._C_gguf.ggml_moe_get_block_size(quant_type)