diff --git a/scripts/deps/__init__.py b/scripts/deps/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/scripts/deps/torch_npu_install.py b/scripts/deps/torch_npu_install.py new file mode 100644 index 0000000000..31ebb57530 --- /dev/null +++ b/scripts/deps/torch_npu_install.py @@ -0,0 +1,114 @@ +"""torch_npu wheel management for xLLM NPU builds. + +Ensures the correct torch_npu version is installed before compilation. +Auto-installs or upgrades as needed — mirrors the TileLang pattern in +xllm/compiler/tilelang/tilelang_ascend_install.py. +""" + +from __future__ import annotations + +import importlib.metadata +import os +import platform +import subprocess +import sys +import tempfile +import urllib.request + +from scripts.logger import logger + +TORCH_NPU_WHEELS: dict[tuple[str, str], str] = { + # (cann_version, arch) → wheel URL + ("9.0", "arm"): ( + "http://storage.jd.local/bmp-packages/NPU_CICD/compile-base/" + "compile-ascend_9.0.0-arm64/GR/" + "torch_npu-2.9.0.post4_git44b5898-cp311-cp311-linux_aarch64.whl" + ), +} + +# PEP 440 compliant local filename for pip install (upstream wheel name +# contains underscores in the version tag that pip 26+ rejects). +_LOCAL_WHEEL_NAME = "torch_npu-2.9.0.post4-cp311-cp311-linux_aarch64.whl" + +REQUIRED_VERSION = "2.9.0.post4" + + +def _get_cann_version() -> str: + """Detect CANN major.minor from toolkit path.""" + toolkit = os.environ.get( + "NPU_TOOLKIT_HOME", "/usr/local/Ascend/ascend-toolkit/latest" + ) + version_file = os.path.join(toolkit, "version.cfg") + try: + with open(version_file) as f: + for line in f: + if "=" in line: + _, ver = line.strip().split("=", 1) + parts = ver.split(".") + if len(parts) >= 2: + return f"{parts[0]}.{parts[1]}" + except (OSError, ValueError): + pass + return "9.0" + + +def _get_arch() -> str: + machine = platform.machine() + if machine in ("aarch64", "arm64"): + return "arm" + return "x86" + + +def _installed_version() -> str | None: + try: + return importlib.metadata.version("torch_npu") + except importlib.metadata.PackageNotFoundError: + return None + + +def _resolve_wheel(cann_version: str, arch: str) -> str: + key = (cann_version, arch) + url = TORCH_NPU_WHEELS.get(key) + if url is None: + raise RuntimeError( + f"No torch_npu wheel configured for CANN {cann_version}, arch={arch}.\n" + f"Available: {list(TORCH_NPU_WHEELS.keys())}" + ) + return url + + +def _download_and_install(wheel_url: str) -> None: + """Download wheel, rename to PEP 440 compliant name, then pip install.""" + with tempfile.TemporaryDirectory() as tmpdir: + local_path = os.path.join(tmpdir, _LOCAL_WHEEL_NAME) + logger.info(f"Downloading {wheel_url} ...") + urllib.request.urlretrieve(wheel_url, local_path) + subprocess.check_call([ + sys.executable, "-m", "pip", "install", + "--force-reinstall", "--no-deps", local_path, + ]) + + +def ensure_torch_npu_ready() -> None: + """Ensure torch_npu matches REQUIRED_VERSION; install or upgrade if not.""" + installed = _installed_version() + if installed is not None and REQUIRED_VERSION in installed: + logger.info(f"torch_npu {installed} is ready.") + return + + cann_version = _get_cann_version() + arch = _get_arch() + wheel_url = _resolve_wheel(cann_version, arch) + + action = f"Upgrading torch_npu {installed}" if installed else "Installing torch_npu" + logger.info(f"{action} → {REQUIRED_VERSION}") + + _download_and_install(wheel_url) + + new_version = _installed_version() + if new_version is None or REQUIRED_VERSION not in new_version: + raise RuntimeError( + f"torch_npu installation failed: expected {REQUIRED_VERSION}, " + f"got {new_version}" + ) + logger.info(f"torch_npu {new_version} installed successfully.") diff --git a/setup.py b/setup.py index 067f0887c9..47dab5d1b6 100644 --- a/setup.py +++ b/setup.py @@ -48,6 +48,12 @@ BUILD_EXPORT: bool = True +def _ensure_torch_npu_ready() -> None: + from scripts.deps.torch_npu_install import ensure_torch_npu_ready + + ensure_torch_npu_ready() + + def _ensure_tilelang_ascend_ready(target_platform: str, arch: str) -> None: compiler_parent = os.path.join(get_base_dir(), "xllm") if compiler_parent not in sys.path: @@ -913,6 +919,7 @@ def parse_arguments() -> dict[str, Any]: logger.info(f"🚀 Build xllm with CPU arch: {arch} and target device: {device}") if device == "npu": + _ensure_torch_npu_ready() _ensure_tilelang_ascend_ready(target_platform, arch) pre_build(device) diff --git a/xllm/core/distributed_runtime/master.cpp b/xllm/core/distributed_runtime/master.cpp index b1743ddde2..e386c785af 100644 --- a/xllm/core/distributed_runtime/master.cpp +++ b/xllm/core/distributed_runtime/master.cpp @@ -37,6 +37,7 @@ limitations under the License. #include "core/framework/config/eplb_config.h" #include "core/framework/config/kernel_config.h" #include "core/framework/config/kv_cache_config.h" +#include "core/framework/config/model_config.h" #include "core/framework/config/parallel_config.h" #include "dit_master.h" #if defined(USE_NPU) @@ -171,6 +172,15 @@ void resolve_npu_kernel_backend_for_options(Options* options) { return; } + // Python model executor bypasses ATB entirely — force TORCH backend. + if (ModelConfig::is_python_model_impl( + ModelConfig::get_instance().model_impl())) { + options->npu_kernel_backend("TORCH"); + KernelConfig::get_instance().npu_kernel_backend("TORCH"); + LOG(INFO) << "Forced npu_kernel_backend=TORCH for python model_impl"; + return; + } + const std::string model_type = util::get_model_type(options->model_path(), options->backend()); std::string effective_backend; @@ -203,15 +213,28 @@ Master::Master(const Options& options, EngineType type) // cards (honoring *_VISIBLE_DEVICES) and select the single card this process // owns by its node_rank, so `devices` holds exactly the card this process // uses -- mirroring the historical single-element devices semantics. + // + // Exception: Python model_impl runs TP in-process (threads, not spawned + // processes) so all visible devices must be passed when nnodes == 1. const auto visible_devices = DeviceNameUtils::parse_devices("auto"); - CHECK_LT(options_.node_rank(), static_cast(visible_devices.size())) - << "node_rank " << options_.node_rank() - << " exceeds the number of visible devices " << visible_devices.size() - << ". Ensure *_VISIBLE_DEVICES exposes all cards used across processes."; - const std::vector devices = { - visible_devices[options_.node_rank()]}; - // World size is the node count (one worker per process). - const int32_t global_world_size = options_.nnodes(); + const bool python_inproc_tp = options_.nnodes() == 1 && + ModelConfig::is_python_model_impl( + ModelConfig::get_instance().model_impl()) && + visible_devices.size() > 1; + std::vector devices; + int32_t global_world_size; + if (python_inproc_tp) { + devices = visible_devices; + global_world_size = static_cast(devices.size()); + } else { + CHECK_LT(options_.node_rank(), static_cast(visible_devices.size())) + << "node_rank " << options_.node_rank() + << " exceeds the number of visible devices " << visible_devices.size() + << ". Ensure *_VISIBLE_DEVICES exposes all cards used across " + "processes."; + devices = {visible_devices[options_.node_rank()]}; + global_world_size = options_.nnodes(); + } std::string cp_model_type; if (options_.cp_size() > 1 && Platform::uses_model_cp_partition()) { cp_model_type = util::get_model_type(model_path, options_.backend()); diff --git a/xllm/core/kernels/npu/npu_ops_library.cpp b/xllm/core/kernels/npu/npu_ops_library.cpp index 9b954e3994..51722cef0a 100644 --- a/xllm/core/kernels/npu/npu_ops_library.cpp +++ b/xllm/core/kernels/npu/npu_ops_library.cpp @@ -28,6 +28,7 @@ limitations under the License. #include "npu_ops_api.h" namespace xllm { + namespace { torch::Tensor rms_norm_npu(const torch::Tensor& input, @@ -68,6 +69,38 @@ void apply_rotary_embedding_npu(torch::Tensor& q, xllm::kernel::npu::apply_rotary(q, k, cos_sin_cache, positions); } +torch::Tensor update_decode_graph_metadata_npu( + const torch::Tensor& tokens, + const torch::Tensor& positions, + const torch::Tensor& slot_mapping, + const torch::Tensor& kv_seq_lens, + const torch::Tensor& paged_kv_indptr, + const torch::Tensor& paged_kv_indices, + const torch::Tensor& paged_kv_last_page_len, + torch::Tensor& dst_tokens, + torch::Tensor& dst_positions, + torch::Tensor& dst_slot_mapping, + torch::Tensor& dst_kv_seq_lens, + torch::Tensor& dst_kv_seq_lens_delta, + torch::Tensor& dst_paged_kv_indptr, + torch::Tensor& dst_paged_kv_indices, + torch::Tensor& dst_paged_kv_last_page_len, + int64_t padded_num_tokens) { + const int64_t n = tokens.numel(); + const int64_t batch = paged_kv_last_page_len.numel(); + const int64_t idx_size = paged_kv_indices.numel(); + dst_tokens.slice(0, 0, n).copy_(tokens); + dst_positions.slice(0, 0, n).copy_(positions); + dst_slot_mapping.slice(0, 0, n).copy_(slot_mapping); + dst_kv_seq_lens.slice(0, 0, batch + 1) + .copy_(kv_seq_lens.slice(0, 0, batch + 1)); + dst_kv_seq_lens_delta.slice(0, 0, batch).fill_(1); + dst_paged_kv_indptr.slice(0, 0, batch + 1).copy_(paged_kv_indptr); + dst_paged_kv_indices.slice(0, 0, idx_size).copy_(paged_kv_indices); + dst_paged_kv_last_page_len.slice(0, 0, batch).copy_(paged_kv_last_page_len); + return dst_tokens; +} + } // namespace void ensure_xllm_ops_registered() { @@ -114,4 +147,6 @@ TORCH_LIBRARY_IMPL(xllm_ops, PrivateUse1, m) { m.impl("silu_and_mul", TORCH_FN(xllm::silu_and_mul_npu)); m.impl("reshape_paged_cache", TORCH_FN(xllm::reshape_paged_cache_npu)); m.impl("apply_rotary_embedding", TORCH_FN(xllm::apply_rotary_embedding_npu)); + m.impl("update_decode_graph_metadata", + TORCH_FN(xllm::update_decode_graph_metadata_npu)); } diff --git a/xllm/core/platform/device.cpp b/xllm/core/platform/device.cpp index b20cdf0e6d..4e89b3c28e 100644 --- a/xllm/core/platform/device.cpp +++ b/xllm/core/platform/device.cpp @@ -15,10 +15,13 @@ limitations under the License. #include "device.h" +#include "core/framework/config/model_config.h" #include "core/platform/platform.h" #if defined(USE_NPU) #include #include +#include +#include #elif defined(USE_MLU) #include #include @@ -94,7 +97,14 @@ int32_t Device::index() const { return device_.index(); } // set device before init device context void Device::init_device_context() const { #if defined(USE_NPU) - torch_npu::init_npu(index()); + if (ModelConfig::is_python_model_impl( + ModelConfig::get_instance().model_impl())) { + // Python path: full NPU runtime already initialized in main() via + // aclInit + torch_npu._C._npu_init(). Only switch device here. + c10_npu::SetDevice(index()); + } else { + torch_npu::init_npu(index()); + } #endif } diff --git a/xllm/python/attention/npu_paged_attention.py b/xllm/python/attention/npu_paged_attention.py index d85578c032..32ba637438 100644 --- a/xllm/python/attention/npu_paged_attention.py +++ b/xllm/python/attention/npu_paged_attention.py @@ -53,12 +53,8 @@ def __init__( self._kv_caches: list[KVCache] = [] self._metadata: AttentionMetadata | None = None - self._causal_mask = ( - torch.triu(torch.ones(2048, 2048, dtype=torch.float32), 1) - .to(torch.int8) - .contiguous() - .to(device) - ) + self._causal_mask_size = 2048 + self._causal_mask = self._make_causal_mask(self._causal_mask_size, device) @property def num_kv_blocks(self) -> int: @@ -117,12 +113,14 @@ def execute( k_cache, v_cache = self._kv_caches[layer_id] num_tokens = q.shape[0] - # Write KV to paged cache (kernel expects [T, kv_heads, head_dim]). + # Write KV to paged cache: k_cache is [num_blocks, block_size, heads, dim], + # flatten first two dims so slot_mapping indexes flat slots. k_3d = k.view(num_tokens, self.num_kv_heads, self.head_dim).contiguous() v_3d = v.view(num_tokens, self.num_kv_heads, self.head_dim).contiguous() - torch.ops.xllm_ops.reshape_paged_cache( - metadata.slot_mapping, k_3d, v_3d, k_cache, v_cache - ) + flat_k_cache = k_cache.view(-1, self.num_kv_heads, self.head_dim) + flat_v_cache = v_cache.view(-1, self.num_kv_heads, self.head_dim) + flat_k_cache[metadata.slot_mapping] = k_3d + flat_v_cache[metadata.slot_mapping] = v_3d q_3d = q.view(num_tokens, self.num_heads, self.head_dim).contiguous() @@ -139,11 +137,12 @@ def _prefill( metadata: AttentionMetadata, num_tokens: int, ) -> torch.Tensor: actual_seq = self._cumulative_seq_lens(metadata, num_tokens) + causal_mask = self._get_causal_mask(num_tokens) output, _ = torch.ops.npu.npu_fused_infer_attention_score( q_3d, k_3d, v_3d, pse_shift=None, - atten_mask=self._causal_mask, + atten_mask=causal_mask, actual_seq_lengths=actual_seq, actual_seq_lengths_kv=actual_seq, num_heads=self.num_heads, @@ -188,6 +187,22 @@ def _decode( # Helpers # ------------------------------------------------------------------ + @staticmethod + def _make_causal_mask(size: int, device: torch.device) -> torch.Tensor: + return ( + torch.triu(torch.ones(size, size, dtype=torch.float32), 1) + .to(torch.int8) + .contiguous() + .to(device) + ) + + def _get_causal_mask(self, seq_len: int) -> torch.Tensor: + if seq_len > self._causal_mask_size: + new_size = max(seq_len, self._causal_mask_size * 2) + self._causal_mask = self._make_causal_mask(new_size, self.device) + self._causal_mask_size = new_size + return self._causal_mask + def _cumulative_seq_lens( self, metadata: AttentionMetadata, num_tokens: int, ) -> list[int]: diff --git a/xllm/python/model_executor/executor.py b/xllm/python/model_executor/executor.py index 7196b25aef..5c20f8c505 100644 --- a/xllm/python/model_executor/executor.py +++ b/xllm/python/model_executor/executor.py @@ -98,7 +98,7 @@ def __init__( self.inductor_runner = None graph_backend = str(config.get("python_graph_backend", "off")).lower() - if graph_backend in ("", "off", "none", "0"): + if graph_backend in ("", "off", "none", "0", "eager"): pass elif graph_backend == "cudagraphs": from xllm.python.model_executor.runners.decode_cuda_graph import ( diff --git a/xllm/python/ops/__init__.py b/xllm/python/ops/__init__.py index 90399d81ae..d0e8daa9a2 100644 --- a/xllm/python/ops/__init__.py +++ b/xllm/python/ops/__init__.py @@ -33,8 +33,8 @@ update_decode_graph_metadata, ) from xllm.python.ops.collectives import ( - all_gather, all_reduce_, + all_gather, init_tp_group, ) diff --git a/xllm/python/ops/collectives.py b/xllm/python/ops/collectives.py index 1f736a8026..f0814715db 100644 --- a/xllm/python/ops/collectives.py +++ b/xllm/python/ops/collectives.py @@ -9,6 +9,29 @@ _tp_stores = {} +def _create_process_group( + host: str, port: int, rank: int, world_size: int, device: str +): + """Create HCCL or NCCL ProcessGroup depending on device type.""" + store = dist.TCPStore( + host, + port, + world_size, + rank == 0, + timedelta(minutes=5), + wait_for_workers=False, + ) + device_obj = torch.device(device) + if device_obj.type == "cuda": + group = dist.ProcessGroupNCCL(store, rank, world_size, timedelta(minutes=5)) + else: + import torch_npu # noqa: F401 + from torch_npu._C._distributed_c10d import ProcessGroupHCCL + + group = ProcessGroupHCCL(store, rank, world_size, timedelta(minutes=5)) + return store, group + + def init_tp_group( host: str, port: int, @@ -27,20 +50,7 @@ def init_tp_group( ) return group - store = dist.TCPStore( - host, - port, - world_size, - rank == 0, - timedelta(minutes=5), - wait_for_workers=False, - ) - group = dist.ProcessGroupNCCL( - store, - rank, - world_size, - timedelta(minutes=5), - ) + store, group = _create_process_group(host, port, rank, world_size, device) _tp_stores[device_key] = store _tp_groups[device_key] = group return group diff --git a/xllm/xllm.cpp b/xllm/xllm.cpp index 24bd4937a5..3ccd955ee2 100644 --- a/xllm/xllm.cpp +++ b/xllm/xllm.cpp @@ -17,8 +17,13 @@ limitations under the License. #include #include #include +namespace py = pybind11; #include +#if defined(USE_NPU) +#include +#endif + #include #include #include @@ -509,5 +514,54 @@ int main(int argc, char** argv) { return 1; } +#if defined(USE_NPU) + // Early Python + torch_npu init for embedded-python model executor. + // post4 libtorch_npu.so calls PyGILState_Ensure inside empty_with_format(); + // Python must be ready before any NPU tensor allocation happens. + if (ModelConfig::is_python_model_impl( + ModelConfig::get_instance().model_impl())) { + auto acl_ret = aclInit(nullptr); + CHECK(acl_ret == ACL_SUCCESS || acl_ret == 500000) + << "aclInit failed with error " << acl_ret; + + bool we_initialized_python = false; + if (!Py_IsInitialized()) { + py::initialize_interpreter(/*init_signal_handlers=*/false); + we_initialized_python = true; + } + + const auto first_device = DeviceNameUtils::parse_devices("auto").front(); + const int device_index = first_device.index(); + + { + py::gil_scoped_acquire gil; + py::exec( + "import os, sys\n" + "os.environ['TORCH_DEVICE_BACKEND_AUTOLOAD'] = '0'\n" + "import torch\n" + "orig = torch._C._get_accelerator\n" + "try:\n" + " torch._C._get_accelerator = lambda: torch.device('cpu')\n" + " import torch_npu\n" + "finally:\n" + " torch._C._get_accelerator = orig\n" + "import torch_npu.npu as _npu_mod\n" + "try:\n" + " torch_npu._C._npu_init()\n" + "except RuntimeError as e:\n" + " if 'already initialized' not in str(e).lower():\n" + " raise\n" + "_npu_mod._initialized = True\n" + "_npu_mod._original_pid = os.getpid()\n" + "torch_npu._C._npu_setDevice(" + + std::to_string(device_index) + ")\n"); + } + + if (we_initialized_python) { + PyEval_SaveThread(); + } + } +#endif + return run(); }