feat: use post4 libtorch_npu with early Python init and native ProcessGroupHCCL.

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.")

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)
 

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<int32_t>(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<torch::Device> 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<torch::Device> devices;
+  int32_t global_world_size;
+  if (python_inproc_tp) {
+    devices = visible_devices;
+    global_world_size = static_cast<int32_t>(devices.size());
+  } else {
+    CHECK_LT(options_.node_rank(), static_cast<int32_t>(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());

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));
 }

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 <torch_npu/csrc/aten/NPUGeneratorImpl.h>
 #include <torch_npu/csrc/core/npu/NPUCachingAllocator.h>
+#include <torch_npu/csrc/core/npu/NPUFunctions.h>
+#include <torch_npu/csrc/libs/init_npu.h>
 #elif defined(USE_MLU)
 #include <cn_api.h>
 #include <framework/core/caching_allocator.h>
@@ -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
 }
 

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]:

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 (

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,
 )