mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 00:36:20 +08:00
e1f71ecdf8
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LFPG29daR1FhchJBTjiEKr
27 lines
914 B
Python
27 lines
914 B
Python
"""JIT build of the page-copy kernel (torch cpp_extension; cached under ~/.cache/torch_extensions)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from types import ModuleType
|
|
|
|
_module = None
|
|
|
|
|
|
def copy_pages_module() -> ModuleType:
|
|
global _module
|
|
if _module is None:
|
|
from torch.utils.cpp_extension import load
|
|
|
|
# Do not touch TORCH_CUDA_ARCH_LIST here: the engine's own JIT kernels (deep_jit / TileLang) run
|
|
# in this process and must keep their arch-specific (sm_1xxa) targets. torch's loader derives the
|
|
# arch for this trivial kernel from the visible device.
|
|
here = Path(__file__).resolve().parent
|
|
_module = load(
|
|
name="deepinfer_hicache_page_copy",
|
|
sources=[str(here / "page_copy.cpp"), str(here / "page_copy.cu")],
|
|
extra_cuda_cflags=["-O3"],
|
|
verbose=False,
|
|
)
|
|
return _module
|