[feat](kt-kernel): Ascend NPU backend for single-card DeepSeek-V4-Flash offload

.gitignore

@@ -17,6 +17,9 @@ compile_commands.json
 ktransformers/server/local_store/
 ktransformers/server_test1.db
 *.patch
+# ...but kt-kernel's third-party patch series is tracked on purpose: the build
+# applies it to third_party/llama.cpp at configure time.
+!kt-kernel/third_party_patches/**/*.patch
 img/
 tmp*.txt
 test.txt

kt-kernel/CMakeLists.txt

@@ -18,7 +18,13 @@ option(KTRANSFORMERS_USE_MUSA "ktransformers: use MUSA" OFF)
 option(KTRANSFORMERS_USE_ROCM "ktransformers: use ROCM" OFF)
 option(KTRANSFORMERS_USE_MACA "ktransformers: use MACA" OFF)
 option(KTRANSFORMERS_USE_SYCL "ktransformers: use SYCL GPTQ INT4 MoE" OFF)
+option(KTRANSFORMERS_USE_ASCEND_NPU "ktransformers: use Ascend NPU (CANN)" OFF)
 option(KTRANSFORMERS_CUDA_STATIC_RUNTIME "ktransformers: statically link CUDA runtime" ON)
+option(LLAMA_ARM_DOTPROD "llama: enable ARM NEON SDOT/UDOT" ON)
+option(LLAMA_ARM_FP16 "llama: enable ARM NEON FP16" ON)
+option(LLAMA_ARM_SVE "llama: enable ARM SVE" OFF)
+option(LLAMA_ARM_BF16 "llama: enable ARM BF16" OFF)
+option(LLAMA_ARM_I8MM "llama: enable ARM I8MM" OFF)
 option(KTRANSFORMERS_CPU_USE_KML "ktransformers: CPU use KML" OFF)
 option(KTRANSFORMERS_CPU_USE_AMX_AVX512 "ktransformers: CPU use AMX or AVX512" OFF)
 option(KTRANSFORMERS_CPU_USE_AMX "ktransformers: CPU use AMX" OFF)
@@ -38,13 +44,14 @@ foreach(_KT_GPU_BACKEND IN ITEMS
         KTRANSFORMERS_USE_SYCL
         KTRANSFORMERS_USE_ROCM
         KTRANSFORMERS_USE_MUSA
-        KTRANSFORMERS_USE_MACA)
+        KTRANSFORMERS_USE_MACA
+        KTRANSFORMERS_USE_ASCEND_NPU)
     if(${_KT_GPU_BACKEND})
         math(EXPR _KT_GPU_BACKEND_COUNT "${_KT_GPU_BACKEND_COUNT} + 1")
     endif()
 endforeach()
 if(_KT_GPU_BACKEND_COUNT GREATER 1)
-    message(FATAL_ERROR "CUDA, SYCL, ROCm, MUSA, and MACA backends are mutually exclusive")
+    message(FATAL_ERROR "CUDA, SYCL, ROCm, MUSA, MACA, and Ascend backends are mutually exclusive")
 endif()
 
 # Choose compilers BEFORE project() so CMake honors them
@@ -239,18 +246,26 @@ if(CMAKE_OSX_ARCHITECTURES STREQUAL "arm64" OR CMAKE_GENERATOR_PLATFORM_LWR STRE
             # Raspberry Pi 3, 4, Zero 2 (32-bit)
             list(APPEND ARCH_FLAGS -mno-unaligned-access)
         endif()
-        # add_compile_definitions(__ARM_NEON)
-        # list(APPEND ARCH_FLAGS -march=armv8.2-a+fp16+dotprod)
-        # add_compile_definitions(__ARM_FEATURE_DOTPROD)
-        # add_compile_definitions(__aarch64__)
-
-        # add_compile_definitions(__ARM_NEON)
-        list(APPEND ARCH_FLAGS -march=armv8.2-a+fp16+dotprod+sve+bf16)
-        # list(APPEND ARCH_FLAGS -march=armv8-a+dotprod+sha3+sm4+fp16fml+sve+rng+sb+ssbs+i8mm+bf16+flagm+pauth)
-        # add_compile_definitions(__ARM_FEATURE_DOTPROD)
-        # add_compile_definitions(__ARM_FEATURE_SVE)
-        # add_compile_definitions(__ARM_FEATURE_MATMUL_INT8)
-        # add_compile_definitions(__aarch64__)
+        # ARM extensions are selected dynamically. K920 / Cortex-A76 = armv8.2-a + fp16 + dotprod
+        # (NEON only). K930+ / Neoverse-V1+ add SVE/BF16/I8MM. Toggle via -DLLAMA_ARM_SVE=ON etc.
+        set(_kt_arm_arch "armv8.2-a")
+        if(LLAMA_ARM_FP16)
+            set(_kt_arm_arch "${_kt_arm_arch}+fp16")
+        endif()
+        if(LLAMA_ARM_DOTPROD)
+            set(_kt_arm_arch "${_kt_arm_arch}+dotprod")
+        endif()
+        if(LLAMA_ARM_SVE)
+            set(_kt_arm_arch "${_kt_arm_arch}+sve")
+        endif()
+        if(LLAMA_ARM_BF16)
+            set(_kt_arm_arch "${_kt_arm_arch}+bf16")
+        endif()
+        if(LLAMA_ARM_I8MM)
+            set(_kt_arm_arch "${_kt_arm_arch}+i8mm")
+        endif()
+        list(APPEND ARCH_FLAGS "-march=${_kt_arm_arch}")
+        message(STATUS "ARM target: -march=${_kt_arm_arch}")
     endif()
 elseif(CMAKE_OSX_ARCHITECTURES STREQUAL "x86_64" OR CMAKE_GENERATOR_PLATFORM_LWR MATCHES "^(x86_64|i686|amd64|x64|win32)$" OR
     (NOT CMAKE_OSX_ARCHITECTURES AND NOT CMAKE_GENERATOR_PLATFORM_LWR AND
@@ -474,6 +489,70 @@ add_compile_options("$<$<COMPILE_LANGUAGE:CXX>:${ARCH_FLAGS}>")
 add_compile_options("$<$<COMPILE_LANGUAGE:C>:${ARCH_FLAGS}>")
 
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/../third_party/pybind11 ${CMAKE_CURRENT_BINARY_DIR}/third_party/pybind11)
+
+# ---------------------------------------------------------------------------
+# third_party/llama.cpp patches
+#
+# The llama.cpp submodule is pinned to an upstream commit (tag b3173) so that
+# `git clone --recursive` works for everyone.  That tag predates MXFP4, which
+# DeepSeek-V4 GGUF experts require, so the delta is kept as a patch series in
+# kt-kernel/third_party_patches/llama.cpp and applied here, before the
+# add_subdirectory() below configures llama.cpp.
+#
+# Idempotency without a marker file (a marker inside the submodule would show
+# up as untracked forever): an already-applied patch fails `git apply --check`
+# but succeeds `git apply --reverse --check`.
+# ---------------------------------------------------------------------------
+set(KT_LLAMA_CPP_DIR ${CMAKE_CURRENT_SOURCE_DIR}/../third_party/llama.cpp)
+set(KT_LLAMA_CPP_PATCH_DIR ${CMAKE_CURRENT_SOURCE_DIR}/third_party_patches/llama.cpp)
+file(GLOB KT_LLAMA_CPP_PATCHES "${KT_LLAMA_CPP_PATCH_DIR}/*.patch")
+list(SORT KT_LLAMA_CPP_PATCHES)
+
+if(KT_LLAMA_CPP_PATCHES)
+    find_package(Git QUIET)
+    if(NOT GIT_EXECUTABLE)
+        set(GIT_EXECUTABLE git)
+    endif()
+endif()
+
+foreach(KT_PATCH IN LISTS KT_LLAMA_CPP_PATCHES)
+    execute_process(
+        COMMAND ${GIT_EXECUTABLE} apply --check "${KT_PATCH}"
+        WORKING_DIRECTORY ${KT_LLAMA_CPP_DIR}
+        RESULT_VARIABLE KT_PATCH_CHECK
+        OUTPUT_QUIET
+        ERROR_QUIET)
+    if(KT_PATCH_CHECK EQUAL 0)
+        execute_process(
+            COMMAND ${GIT_EXECUTABLE} apply "${KT_PATCH}"
+            WORKING_DIRECTORY ${KT_LLAMA_CPP_DIR}
+            RESULT_VARIABLE KT_PATCH_APPLY
+            ERROR_VARIABLE KT_PATCH_APPLY_ERR)
+        if(NOT KT_PATCH_APPLY EQUAL 0)
+            message(FATAL_ERROR
+                "Failed to apply llama.cpp patch ${KT_PATCH}:\n${KT_PATCH_APPLY_ERR}")
+        endif()
+        message(STATUS "llama.cpp: applied patch ${KT_PATCH}")
+    else()
+        execute_process(
+            COMMAND ${GIT_EXECUTABLE} apply --reverse --check "${KT_PATCH}"
+            WORKING_DIRECTORY ${KT_LLAMA_CPP_DIR}
+            RESULT_VARIABLE KT_PATCH_REVERSE_CHECK
+            OUTPUT_QUIET
+            ERROR_QUIET)
+        if(KT_PATCH_REVERSE_CHECK EQUAL 0)
+            message(STATUS "llama.cpp: patch already applied, skipping ${KT_PATCH}")
+        else()
+            message(FATAL_ERROR
+                "llama.cpp patch ${KT_PATCH} neither applies nor is already applied.\n"
+                "third_party/llama.cpp is probably not at the pinned commit or has local edits.\n"
+                "Reset it and retry:\n"
+                "  git submodule update --init --force third_party/llama.cpp\n"
+                "  cd third_party/llama.cpp && git apply ${KT_PATCH}")
+        endif()
+    endif()
+endforeach()
+
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/../third_party/llama.cpp ${CMAKE_CURRENT_BINARY_DIR}/third_party/llama.cpp)
 
 include_directories(${CMAKE_CURRENT_SOURCE_DIR}/../third_party)
@@ -562,6 +641,27 @@ elseif(KTRANSFORMERS_USE_MACA)
     add_compile_definitions(KTRANSFORMERS_USE_MACA=1)
 elseif(KTRANSFORMERS_USE_SYCL)
     message(STATUS "SYCL GPTQ INT4 support enabled")
+elseif(KTRANSFORMERS_USE_ASCEND_NPU)
+    message(STATUS "Ascend NPU (CANN) backend selected")
+    add_compile_definitions(KTRANSFORMERS_USE_ASCEND_NPU=1)
+    add_compile_definitions(USE_ASCEND_NPU=1)
+    if(DEFINED ENV{ASCEND_TOOLKIT_HOME})
+        set(_kt_cann_root "$ENV{ASCEND_TOOLKIT_HOME}")
+    elseif(DEFINED ENV{CANN_HOME})
+        set(_kt_cann_root "$ENV{CANN_HOME}")
+    else()
+        set(_kt_cann_root "/usr/local/Ascend/ascend-toolkit/latest")
+    endif()
+    find_path(ACL_INCLUDE_DIR acl/acl_rt.h
+        HINTS "${_kt_cann_root}/include"
+        REQUIRED)
+    find_library(ASCEND_CL_LIBRARY NAMES ascendcl
+        HINTS "${_kt_cann_root}/lib64" "${_kt_cann_root}/runtime/lib64"
+        REQUIRED)
+    message(STATUS "CANN root: ${_kt_cann_root}")
+    message(STATUS "Ascend CL include: ${ACL_INCLUDE_DIR}")
+    message(STATUS "Ascend CL library: ${ASCEND_CL_LIBRARY}")
+    include_directories(${ACL_INCLUDE_DIR})
 elseif(KTRANSFORMERS_CPU_USE_KML)
     message(STATUS "KML CPU detected")
 else()
@@ -795,6 +895,10 @@ if(NOT HOST_IS_X86 AND KTRANSFORMERS_CPU_USE_KML)
     target_compile_definitions(${PROJECT_NAME} PRIVATE CPU_USE_KML)
 endif()
 target_link_libraries(${PROJECT_NAME} PRIVATE llama PkgConfig::HWLOC OpenMP::OpenMP_CXX)
+if(KTRANSFORMERS_USE_ASCEND_NPU AND ASCEND_CL_LIBRARY)
+    target_link_libraries(${PROJECT_NAME} PRIVATE ${ASCEND_CL_LIBRARY})
+    target_include_directories(${PROJECT_NAME} PRIVATE ${ACL_INCLUDE_DIR})
+endif()
 if(NOT HOST_IS_X86 AND KTRANSFORMERS_CPU_USE_KML)
     if(KTRANSFORMERS_CPU_DEBUG)
         # add_executable(convert-test ${CMAKE_CURRENT_SOURCE_DIR}/operators/kml/convert-test.cpp)

kt-kernel/cpu_backend/ascend_callback_worker.cpp

@@ -0,0 +1,153 @@
+#if defined(KTRANSFORMERS_USE_ASCEND_NPU)
+
+#include "ascend_callback_worker.h"
+
+#include <atomic>
+#include <cstdio>
+#include <mutex>
+#include <set>
+#include <thread>
+#include <vector>
+
+namespace kt::ascend {
+
+namespace {
+
+constexpr int kProcessReportTimeoutMs = 100;
+
+std::mutex g_mu;
+std::set<aclrtStream> g_subscribed_streams;
+aclrtContext g_context = nullptr;
+std::thread g_worker;
+std::atomic<bool> g_stop{false};
+std::atomic<bool> g_started{false};
+// Set by the worker thread once it has a usable ACL context and is about to
+// enter the aclrtProcessReport() loop.  aclrtSubscribeReport() must not run
+// before that, otherwise the very first callbacks can be dropped.
+std::atomic<bool> g_worker_ready{false};
+uint64_t g_worker_thread_id = 0;
+
+void worker_main(aclrtContext ctx) {
+  if (ctx != nullptr) {
+    aclError err = aclrtSetCurrentContext(ctx);
+    if (err != ACL_SUCCESS) {
+      std::fprintf(stderr,
+                   "[kt-kernel] ascend_callback_worker: aclrtSetCurrentContext failed (%d)\n",
+                   static_cast<int>(err));
+    }
+  }
+  g_worker_ready.store(true, std::memory_order_release);
+
+  while (!g_stop.load(std::memory_order_acquire)) {
+    (void)aclrtProcessReport(kProcessReportTimeoutMs);
+  }
+
+  // Drain callbacks that were enqueued between the last loop iteration and the
+  // stop request; without this they are silently dropped and their waiters hang.
+  (void)aclrtProcessReport(kProcessReportTimeoutMs);
+
+  // ACL requires the *subscribing* thread to unsubscribe.  Doing it from
+  // shutdown_callback_worker() after join() would target a dead thread id and
+  // silently do nothing.
+  for (aclrtStream stream : g_subscribed_streams) {
+    (void)aclrtUnSubscribeReport(g_worker_thread_id, stream);
+  }
+}
+
+void start_worker_locked(aclrtContext ctx) {
+  if (g_started.load(std::memory_order_acquire)) {
+    return;
+  }
+  g_context = ctx;
+  g_stop.store(false, std::memory_order_release);
+  g_worker_ready.store(false, std::memory_order_release);
+  g_worker = std::thread([ctx]() { worker_main(ctx); });
+  g_worker_thread_id = static_cast<uint64_t>(g_worker.native_handle());
+  g_started.store(true, std::memory_order_release);
+
+  // Barrier: no aclrtSubscribeReport() before the worker is in its report loop.
+  while (!g_worker_ready.load(std::memory_order_acquire)) {
+    std::this_thread::yield();
+  }
+
+  for (aclrtStream stream : g_subscribed_streams) {
+    aclError err = aclrtSubscribeReport(g_worker_thread_id, stream);
+    if (err != ACL_SUCCESS) {
+      std::fprintf(stderr,
+                   "[kt-kernel] ascend_callback_worker: aclrtSubscribeReport failed (%d)\n",
+                   static_cast<int>(err));
+    }
+  }
+}
+
+void subscribe_stream_locked(aclrtStream stream) {
+  if (stream == nullptr) {
+    return;
+  }
+  if (g_subscribed_streams.count(stream) != 0) {
+    return;
+  }
+  g_subscribed_streams.insert(stream);
+  if (g_started.load(std::memory_order_acquire)) {
+    aclError err = aclrtSubscribeReport(g_worker_thread_id, stream);
+    if (err != ACL_SUCCESS) {
+      std::fprintf(stderr,
+                   "[kt-kernel] ascend_callback_worker: aclrtSubscribeReport failed (%d)\n",
+                   static_cast<int>(err));
+    }
+  }
+}
+
+}  // namespace
+
+void ensure_callback_worker(aclrtContext ctx) {
+  std::lock_guard<std::mutex> lock(g_mu);
+  if (g_started.load(std::memory_order_acquire)) {
+    return;
+  }
+  aclrtContext use_ctx = ctx;
+  if (use_ctx == nullptr) {
+    aclError err = aclrtGetCurrentContext(&use_ctx);
+    if (err != ACL_SUCCESS || use_ctx == nullptr) {
+      std::fprintf(stderr,
+                   "[kt-kernel] ascend_callback_worker: no ACL context; call after torch.npu init\n");
+      return;
+    }
+  }
+  start_worker_locked(use_ctx);
+}
+
+void ensure_stream_subscribed(aclrtStream stream) {
+  std::lock_guard<std::mutex> lock(g_mu);
+  if (!g_started.load(std::memory_order_acquire)) {
+    aclrtContext ctx = nullptr;
+    (void)aclrtGetCurrentContext(&ctx);
+    start_worker_locked(ctx);
+  }
+  subscribe_stream_locked(stream);
+}
+
+void shutdown_callback_worker() {
+  std::lock_guard<std::mutex> lock(g_mu);
+  if (!g_started.load(std::memory_order_acquire)) {
+    return;
+  }
+  g_stop.store(true, std::memory_order_release);
+  if (g_worker.joinable()) {
+    // worker_main() drains the report queue and unsubscribes before returning.
+    g_worker.join();
+  }
+  g_subscribed_streams.clear();
+  g_started.store(false, std::memory_order_release);
+  g_worker_ready.store(false, std::memory_order_release);
+  g_worker_thread_id = 0;
+}
+
+bool callback_worker_running() {
+  std::lock_guard<std::mutex> lock(g_mu);
+  return g_started.load(std::memory_order_acquire) && g_worker_ready.load(std::memory_order_acquire);
+}
+
+}  // namespace kt::ascend
+
+#endif  // KTRANSFORMERS_USE_ASCEND_NPU

kt-kernel/cpu_backend/ascend_callback_worker.h

@@ -0,0 +1,34 @@
+#pragma once
+
+// Ascend ACL stream callback subscriber for kt-kernel.
+//
+// CANN dispatches aclrtLaunchCallback tasks only on a host thread that has called
+// aclrtSubscribeReport(threadId, stream) and is running aclrtProcessReport() in a
+// loop.  This worker mirrors torch_npu's pattern (see NPUGraph.cpp / Graph.cpp).
+
+#if defined(KTRANSFORMERS_USE_ASCEND_NPU)
+
+#include <acl/acl_rt.h>
+
+#include <cstdint>
+
+namespace kt::ascend {
+
+// Start the global callback worker (idempotent).  Call after ACL/torch.npu init.
+// If ctx is null, uses aclrtGetCurrentContext().
+void ensure_callback_worker(aclrtContext ctx = nullptr);
+
+// Register ``stream`` with the worker so enqueued callbacks are dispatched.
+void ensure_stream_subscribed(aclrtStream stream);
+
+// Optional shutdown (process exit).
+void shutdown_callback_worker();
+
+// True iff the worker thread was started and has entered its aclrtProcessReport
+// loop.  When this is false, stream callbacks are never dispatched and callers
+// must fall back to the synchronous submit/sync path.
+bool callback_worker_running();
+
+}  // namespace kt::ascend
+
+#endif  // KTRANSFORMERS_USE_ASCEND_NPU