From 6f100c96095f32f31db5ed4d717815ecc923de77 Mon Sep 17 00:00:00 2001 From: Pan-Boyi Date: Thu, 13 Aug 2026 20:05:30 +0800 Subject: [PATCH 01/13] kt-kernel: support DeepSeek V4 offload on Ascend Signed-off-by: Pan-Boyi --- kt-kernel/CMakeLists.txt | 68 ++- .../cpu_backend/ascend_callback_worker.cpp | 127 +++++ .../cpu_backend/ascend_callback_worker.h | 29 ++ kt-kernel/cpu_backend/cpuinfer.h | 8 +- kt-kernel/cpu_backend/vendors/ascend_npu.h | 64 +++ kt-kernel/cpu_backend/vendors/vendor.h | 2 + kt-kernel/ext_bindings.cpp | 18 + kt-kernel/install.sh | 92 +++- kt-kernel/operators/llamafile/moe.hpp | 359 ++++++++----- kt-kernel/pyproject.toml | 4 +- kt-kernel/python/experts.py | 14 +- kt-kernel/python/experts_base.py | 471 ++++++++++++++++-- kt-kernel/python/utils/llamafile.py | 76 ++- kt-kernel/python/utils/loader.py | 109 ++-- kt-kernel/setup.py | 82 ++- third_party/llama.cpp | 2 +- 16 files changed, 1289 insertions(+), 236 deletions(-) create mode 100644 kt-kernel/cpu_backend/ascend_callback_worker.cpp create mode 100644 kt-kernel/cpu_backend/ascend_callback_worker.h create mode 100644 kt-kernel/cpu_backend/vendors/ascend_npu.h diff --git a/kt-kernel/CMakeLists.txt b/kt-kernel/CMakeLists.txt index 66ad54973..bb035f152 100644 --- a/kt-kernel/CMakeLists.txt +++ b/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) @@ -36,13 +42,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 @@ -237,18 +244,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 @@ -544,6 +559,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() @@ -750,6 +786,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) diff --git a/kt-kernel/cpu_backend/ascend_callback_worker.cpp b/kt-kernel/cpu_backend/ascend_callback_worker.cpp new file mode 100644 index 000000000..6f69d92cf --- /dev/null +++ b/kt-kernel/cpu_backend/ascend_callback_worker.cpp @@ -0,0 +1,127 @@ +#if defined(KTRANSFORMERS_USE_ASCEND_NPU) + +#include "ascend_callback_worker.h" + +#include +#include +#include +#include +#include +#include + +namespace kt::ascend { + +namespace { + +constexpr int kProcessReportTimeoutMs = 100; + +std::mutex g_mu; +std::set g_subscribed_streams; +aclrtContext g_context = nullptr; +std::thread g_worker; +std::atomic g_stop{false}; +std::atomic g_started{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(err)); + } + } + + while (!g_stop.load(std::memory_order_acquire)) { + (void)aclrtProcessReport(kProcessReportTimeoutMs); + } +} + +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 = std::thread([ctx]() { worker_main(ctx); }); + g_worker_thread_id = static_cast(g_worker.native_handle()); + g_started.store(true, std::memory_order_release); + + 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(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(err)); + } + } +} + +} // namespace + +void ensure_callback_worker(aclrtContext ctx) { + std::lock_guard 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 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 lock(g_mu); + if (!g_started.load(std::memory_order_acquire)) { + return; + } + g_stop.store(true, std::memory_order_release); + if (g_worker.joinable()) { + g_worker.join(); + } + for (aclrtStream stream : g_subscribed_streams) { + (void)aclrtUnSubscribeReport(g_worker_thread_id, stream); + } + g_subscribed_streams.clear(); + g_started.store(false, std::memory_order_release); + g_worker_thread_id = 0; +} + +} // namespace kt::ascend + +#endif // KTRANSFORMERS_USE_ASCEND_NPU diff --git a/kt-kernel/cpu_backend/ascend_callback_worker.h b/kt-kernel/cpu_backend/ascend_callback_worker.h new file mode 100644 index 000000000..8719a4860 --- /dev/null +++ b/kt-kernel/cpu_backend/ascend_callback_worker.h @@ -0,0 +1,29 @@ +#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 + +#include + +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(); + +} // namespace kt::ascend + +#endif // KTRANSFORMERS_USE_ASCEND_NPU diff --git a/kt-kernel/cpu_backend/cpuinfer.h b/kt-kernel/cpu_backend/cpuinfer.h index 8366d1ce8..227187d68 100644 --- a/kt-kernel/cpu_backend/cpuinfer.h +++ b/kt-kernel/cpu_backend/cpuinfer.h @@ -26,6 +26,8 @@ #include "vendors/hip.h" #elif KTRANSFORMERS_USE_MACA #include "vendors/maca.h" +#elif KTRANSFORMERS_USE_ASCEND_NPU +#include "vendors/ascend_npu.h" #endif #include "./vendors/vendor.h" @@ -86,7 +88,8 @@ class CPUInfer { #ifndef KTRANSFORMERS_CPU_ONLY void submit_with_cuda_stream(intptr_t user_cuda_stream, std::pair params) { #if defined(KTRANSFORMERS_USE_CUDA) || defined(KTRANSFORMERS_USE_CUDA_HOST_CALLBACKS) || \ - defined(KTRANSFORMERS_USE_MUSA) || defined(KTRANSFORMERS_USE_ROCM) || defined(KTRANSFORMERS_USE_MACA) + defined(KTRANSFORMERS_USE_MUSA) || defined(KTRANSFORMERS_USE_ROCM) || defined(KTRANSFORMERS_USE_MACA) || \ + defined(KTRANSFORMERS_USE_ASCEND_NPU) void (*func)(void*) = (void (*)(void*))params.first; void* args = (void*)params.second; *((CPUInfer**)args) = this; @@ -112,7 +115,8 @@ class CPUInfer { #ifndef KTRANSFORMERS_CPU_ONLY void sync_with_cuda_stream(intptr_t user_cuda_stream, size_t allow_n_pending = 0) { #if defined(KTRANSFORMERS_USE_CUDA) || defined(KTRANSFORMERS_USE_CUDA_HOST_CALLBACKS) || \ - defined(KTRANSFORMERS_USE_MUSA) || defined(KTRANSFORMERS_USE_ROCM) || defined(KTRANSFORMERS_USE_MACA) + defined(KTRANSFORMERS_USE_MUSA) || defined(KTRANSFORMERS_USE_ROCM) || defined(KTRANSFORMERS_USE_MACA) || \ + defined(KTRANSFORMERS_USE_ASCEND_NPU) SyncArgs* args = new SyncArgs{this, allow_n_pending}; cudaLaunchHostFunc((cudaStream_t)user_cuda_stream, (cudaHostFn_t)&sync_, (void*)args); #endif diff --git a/kt-kernel/cpu_backend/vendors/ascend_npu.h b/kt-kernel/cpu_backend/vendors/ascend_npu.h new file mode 100644 index 000000000..259dcd6ac --- /dev/null +++ b/kt-kernel/cpu_backend/vendors/ascend_npu.h @@ -0,0 +1,64 @@ +#pragma once + +// ============================================================================ +// Ascend NPU vendor adapter for cpu_backend. +// +// Provides CUDA-shaped wrappers around CANN aclrt* runtime APIs so the rest of +// cpu_backend (cpuinfer.h, ext_bindings.cpp etc.) can keep using cudaStream_t / +// cudaHostFn_t / cudaLaunchHostFunc names. Mirrors the pattern used by +// vendors/hip.h and vendors/musa.h. +// +// IMPORTANT semantic difference from CUDA: +// * aclrtLaunchCallback() inserts a callback into the stream's report queue. +// The callback is dispatched by a dedicated host "callback thread" which +// must (a) have been registered with aclrtSubscribeReport(threadId, stream) +// and (b) be continuously calling aclrtProcessReport(timeout) in a loop. +// * If no such subscriber thread exists, queued callbacks NEVER fire and any +// code waiting on the callback (e.g. CPUInfer::sync_with_cuda_stream) will +// hang forever. +// * The submit_with_cuda_stream / sync_with_cuda_stream path therefore +// requires the host process to also spin up a callback worker (see +// cpu_backend/ascend_callback_worker.cpp). The synchronous CPUInfer::submit() +// / sync() path works without it +// and is what the Phase 1 PoC uses. +// ============================================================================ + +#include +#include + +#include + +#if defined(KTRANSFORMERS_USE_ASCEND_NPU) +#include "../ascend_callback_worker.h" +#endif + +// ---- types ----------------------------------------------------------------- +using cudaStream_t = aclrtStream; +using cudaError_t = aclError; +using cudaHostFn_t = aclrtCallback; // both are void (*)(void *) + +// ACL_SUCCESS is `static const int = 0;` in acl_base_rt.h. Re-expose with the +// CUDA name. Using an inline constexpr avoids ODR issues across TUs. +inline constexpr cudaError_t cudaSuccess = ACL_SUCCESS; + +// ---- callbacks ------------------------------------------------------------- +// CUDA: cudaLaunchHostFunc(stream, fn, userData) +// ACL : aclrtLaunchCallback(fn, userData, blockType, stream) +// +// We pick ACL_CALLBACK_NO_BLOCK to match CUDA's fire-and-forget enqueue +// semantics. (ACL_CALLBACK_BLOCK can block the host if the device queue +// is full.) See header note above re: the required subscriber thread. +static inline cudaError_t cudaLaunchHostFunc(cudaStream_t stream, cudaHostFn_t fn, void* userData) { +#if defined(KTRANSFORMERS_USE_ASCEND_NPU) + kt::ascend::ensure_stream_subscribed(stream); +#endif + return aclrtLaunchCallback(fn, userData, ACL_CALLBACK_NO_BLOCK, stream); +} + +// ---- error reporting ------------------------------------------------------- +// CUDA gives a static string per error code; ACL only exposes the *most recent* +// error message on the current thread. Best-effort emulation. +static inline const char* cudaGetErrorString(cudaError_t /*err*/) { + const char* m = aclGetRecentErrMsg(); + return (m && *m) ? m : "ACL error"; +} diff --git a/kt-kernel/cpu_backend/vendors/vendor.h b/kt-kernel/cpu_backend/vendors/vendor.h index dfc8d37c8..c460c945b 100644 --- a/kt-kernel/cpu_backend/vendors/vendor.h +++ b/kt-kernel/cpu_backend/vendors/vendor.h @@ -10,6 +10,8 @@ #include "musa.h" #elif USE_MACA #include "maca.h" +#elif USE_ASCEND_NPU +#include "ascend_npu.h" #endif #endif // CPUINFER_VENDOR_VENDOR_H diff --git a/kt-kernel/ext_bindings.cpp b/kt-kernel/ext_bindings.cpp index 355ef8285..adaa5d866 100644 --- a/kt-kernel/ext_bindings.cpp +++ b/kt-kernel/ext_bindings.cpp @@ -21,6 +21,9 @@ #include "cpu_backend/cpuinfer.h" #include "cpu_backend/worker_pool.h" +#if defined(KTRANSFORMERS_USE_ASCEND_NPU) +#include "cpu_backend/ascend_callback_worker.h" +#endif #include "operators/common.hpp" #if defined(USE_MOE_KERNEL) @@ -514,6 +517,21 @@ void bind_moe_module(py::module_& moe_module, const char* name) { } PYBIND11_MODULE(kt_kernel_ext, m) { +#if defined(KTRANSFORMERS_USE_ASCEND_NPU) + m.def( + "init_ascend_callback_worker", + []() { kt::ascend::ensure_callback_worker(nullptr); }, + "Start ACL aclrtProcessReport worker for stream callbacks (Ascend NPU)."); + m.def( + "subscribe_ascend_stream", + [](intptr_t stream_handle) { + kt::ascend::ensure_stream_subscribed(reinterpret_cast(stream_handle)); + }, + py::arg("stream_handle"), + "Subscribe an aclrtStream with the global callback worker."); + m.def("shutdown_ascend_callback_worker", &kt::ascend::shutdown_callback_worker, + "Stop the ACL callback worker thread."); +#endif py::class_(m, "WorkerPool").def(py::init()); py::class_(m, "WorkerPoolConfig") .def(py::init<>()) diff --git a/kt-kernel/install.sh b/kt-kernel/install.sh index 06a7d8c91..4c79787ab 100755 --- a/kt-kernel/install.sh +++ b/kt-kernel/install.sh @@ -76,6 +76,7 @@ Optional variables (with defaults): CPUINFER_ENABLE_AVX512_BF16=ON/OFF Override BF16 detection (auto if unset) CPUINFER_ENABLE_AVX512_VBMI=ON/OFF Override VBMI detection (auto if unset) CPUINFER_ENABLE_CPPTRACE=ON/OFF Enable native crash tracing (default OFF) + CPUINFER_PIP_NO_DEPS=0/1 Install the built wheel without changing the image's Python stack Software Fallback Support: ✓ If VNNI not available: Uses AVX512BW fallback (2-3x slower but works) @@ -146,6 +147,33 @@ install_dependencies() { esac } +# Function to detect ARM (aarch64) features from /proc/cpuinfo "Features:" line. +# Returns: "has_dotprod has_fp16 has_sve has_bf16 has_i8mm" (space-separated 0/1 values). +detect_arm_features() { + local has_dotprod=0 has_fp16=0 has_sve=0 has_bf16=0 has_i8mm=0 + if [ -f /proc/cpuinfo ]; then + local feats + feats=$(grep -m1 -E "^Features\s*:" /proc/cpuinfo | tr ' ' '\n') + echo "$feats" | grep -qE "^asimddp$" && has_dotprod=1 + echo "$feats" | grep -qE "^(asimdhp|fphp)$" && has_fp16=1 + echo "$feats" | grep -qE "^sve$" && has_sve=1 + echo "$feats" | grep -qE "^bf16$" && has_bf16=1 + echo "$feats" | grep -qE "^i8mm$" && has_i8mm=1 + fi + echo "$has_dotprod $has_fp16 $has_sve $has_bf16 $has_i8mm" +} + +# Detect Ascend CANN install. Echoes the toolkit root if found, else empty. +detect_cann_root() { + for cand in "${ASCEND_TOOLKIT_HOME:-}" "${CANN_HOME:-}" "/usr/local/Ascend/ascend-toolkit/latest"; do + if [ -n "$cand" ] && [ -f "$cand/include/acl/acl_rt.h" ]; then + echo "$cand" + return 0 + fi + done + echo "" +} + # Function to detect CPU features # Returns: "has_amx has_avx512f has_avx512_vnni has_avx512_bf16 has_avx512_vbmi" (space-separated 0/1 values) detect_cpu_features() { @@ -227,6 +255,59 @@ build_step() { echo "==========================================" echo "" + HOST_ARCH="$(uname -m)" + echo "Host arch: $HOST_ARCH" + + if [ "$HOST_ARCH" = "aarch64" ] || [ "$HOST_ARCH" = "arm64" ]; then + # ARM (aarch64) auto-detect path: Kunpeng / Neoverse / Apple Silicon. + # Returns "dotprod fp16 sve bf16 i8mm" + ARM_FEATURES=$(detect_arm_features) + HAS_DOTPROD=$(echo "$ARM_FEATURES" | cut -d' ' -f1) + HAS_FP16=$(echo "$ARM_FEATURES" | cut -d' ' -f2) + HAS_SVE=$(echo "$ARM_FEATURES" | cut -d' ' -f3) + HAS_BF16=$(echo "$ARM_FEATURES" | cut -d' ' -f4) + HAS_I8MM=$(echo "$ARM_FEATURES" | cut -d' ' -f5) + + echo "ARM features: DOTPROD=$HAS_DOTPROD FP16=$HAS_FP16 SVE=$HAS_SVE BF16=$HAS_BF16 I8MM=$HAS_I8MM" + + export CPUINFER_CPU_INSTRUCT=NATIVE + export CPUINFER_ENABLE_AMX=OFF + [ "$HAS_DOTPROD" = "1" ] && export CPUINFER_ARM_DOTPROD=ON || export CPUINFER_ARM_DOTPROD=OFF + [ "$HAS_FP16" = "1" ] && export CPUINFER_ARM_FP16=ON || export CPUINFER_ARM_FP16=OFF + [ "$HAS_SVE" = "1" ] && export CPUINFER_ARM_SVE=ON || export CPUINFER_ARM_SVE=OFF + [ "$HAS_BF16" = "1" ] && export CPUINFER_ARM_BF16=ON || export CPUINFER_ARM_BF16=OFF + [ "$HAS_I8MM" = "1" ] && export CPUINFER_ARM_I8MM=ON || export CPUINFER_ARM_I8MM=OFF + + # CANN auto-detection. If found, enable Ascend NPU backend; otherwise + # build pure CPU and rely on the LLAMA_MOE_TP / llamafile NEON+SDOT path. + CANN_ROOT="$(detect_cann_root)" + if [ -n "$CANN_ROOT" ]; then + echo "✓ CANN detected at: $CANN_ROOT" + export ASCEND_TOOLKIT_HOME="$CANN_ROOT" + export CPUINFER_USE_ASCEND_NPU=1 + else + echo "ℹ CANN not detected; building CPU-only (no NPU host callback support)" + export CPUINFER_USE_ASCEND_NPU=0 + fi + + # K920 / Cortex-A76 cannot run the SVE micro-kernels in kt-kernel's KML path. + # Keep KML / BLIS off unless the user explicitly opts in. + : "${CPUINFER_ENABLE_KML:=OFF}" + : "${CPUINFER_ENABLE_BLIS:=OFF}" + export CPUINFER_ENABLE_KML CPUINFER_ENABLE_BLIS + + echo "" + echo "Configuration (aarch64):" + echo " CPUINFER_USE_ASCEND_NPU = $CPUINFER_USE_ASCEND_NPU" + echo " CPUINFER_ARM_SVE = $CPUINFER_ARM_SVE" + echo " CPUINFER_ARM_BF16 = $CPUINFER_ARM_BF16" + echo " CPUINFER_ARM_I8MM = $CPUINFER_ARM_I8MM" + echo " CPUINFER_ENABLE_KML = $CPUINFER_ENABLE_KML (forced OFF unless SVE present)" + echo "" + # Skip the x86-only AMX/AVX512 detection below. + : + else + # detect_cpu_features returns "has_amx has_avx512f has_avx512_vnni has_avx512_bf16 has_avx512_vbmi" CPU_FEATURES=$(detect_cpu_features) HAS_AMX=$(echo "$CPU_FEATURES" | cut -d' ' -f1) @@ -317,6 +398,7 @@ build_step() { echo "" echo "To use manual configuration instead, run: $0 build --manual" echo "" + fi # end aarch64-vs-x86 branch else # Manual mode - validate user configuration (no exports) if [ -z "$CPUINFER_CPU_INSTRUCT" ] || [ -z "$CPUINFER_ENABLE_AMX" ]; then @@ -399,10 +481,16 @@ echo " CPUINFER_PARALLEL = ${CPUINFER_PARALLEL:-AUTO}" echo " CPUINFER_VERBOSE = ${CPUINFER_VERBOSE:-1}" echo "" +PIP_DEP_ARGS=() +if [ "${CPUINFER_PIP_NO_DEPS:-0}" = "1" ]; then + PIP_DEP_ARGS+=(--no-deps) + echo " pip dependency resolution disabled (CPUINFER_PIP_NO_DEPS=1)" +fi + if [ ${CPUINFER_VERBOSE:-1} = "0" ]; then - python3 -m pip install . + python3 -m pip install "${PIP_DEP_ARGS[@]}" . else - python3 -m pip install . -v + python3 -m pip install "${PIP_DEP_ARGS[@]}" . -v fi } diff --git a/kt-kernel/operators/llamafile/moe.hpp b/kt-kernel/operators/llamafile/moe.hpp index 3a0d76197..809bfaab7 100644 --- a/kt-kernel/operators/llamafile/moe.hpp +++ b/kt-kernel/operators/llamafile/moe.hpp @@ -21,6 +21,27 @@ #include "llama.cpp/ggml-quants.h" #include "llama.cpp/ggml.h" #include "llamafile/sgemm.h" +#if defined(__aarch64__) +#include +#endif + +// --------------------------------------------------------------------------- +// KT_MOE_PHASE_TIMING=1 — Session D Phase-0 测量：decode (forward_one) 三段 +// (input量化 / gate+up job / down job) + TP 层 merge 的累计耗时，定位固定开销 F +// 的构成。env 关闭时仅一次 getenv + 分支，零扰动。每 tp 每 4300 次调用 +// (=43层×100 token) 打一行均值到 stderr。 +// 线程安全性：同一 tp_part_idx 的 forward_one 在 decode 流内串行（逐层）， +// 不同 tp 并行但写不同槽位 → 无需原子。 +// --------------------------------------------------------------------------- +static inline bool kt_phase_timing_on() { + static const bool on = std::getenv("KT_MOE_PHASE_TIMING") != nullptr; + return on; +} +struct KtPhaseAcc { + uint64_t calls = 0; + uint64_t quant_ns = 0, gateup_ns = 0, down_ns = 0; +}; +inline KtPhaseAcc g_kt_phase_acc[16]; inline void debug_quant(void* input, ggml_type type) { std::vector output(ggml_blck_size(type)); @@ -31,6 +52,37 @@ inline void debug_quant(void* input, ggml_type type) { printf("\n"); } +// --------------------------------------------------------------------------- +// kt_effective_vec_dot_type +// 解决 aarch64 (Kunpeng K920 / Cortex-A76, **no SVE / no i8mm**) 上 +// kt-kernel llamafile sgemm 的 BF16/Q8_0 路径不健全问题： +// +// * BF16 weight + BF16 input：tinyblas_cpu_sgemm.inc ARM_NEON path 要求 +// `Btype == GGML_TYPE_F32`（line 209: `if (Btype != F32) return NOT_SUPPORTED;`）， +// 但 ggml type_traits 对 BF16 给出 `vec_dot_type = BF16`， +// 上层 forward_one/forward_many 默认把 BF16 input 喂给 sgemm +// → llamafile_sgemm 返回 false → `throw "llamafile not supported"`。 +// +// * 解法：在 aarch64-without-SVE 平台上，把 BF16 weight 的有效 vec_dot_type +// 声明为 F32。input 路径会因此走 to_float(bf16 → fp32) + memcpy(F32→F32 buffer) +// （`from_float()` 已在 conversion.h 对 F32 short-circuit 成 memcpy）， +// buffer 大小自动按 fp32 (4 bytes/elem) 分配 —— 比原 BF16 (2 bytes/elem) +// 大一倍，足够装 fp32 数据。sgemm 则走 ARM_NEON 已支持的 +// `Atype=BF16, Btype=F32, Ctype=F32` 路径（tinyblas_cpu_sgemm.inc line 125-133）。 +// +// * SVE 机器走原 BF16-BF16 path（vec_dot_type 不改），不打扰原性能优化。 +// * 其他 weight type（Q8_0/Q4_K/…）走原 vec_dot_type，不动； +// 若 Q8_0 NaN 仍存在，单独在 sgemm 内部修复，而非这里。 +// --------------------------------------------------------------------------- +static inline ggml_type kt_effective_vec_dot_type(ggml_type weight_type) { +#if defined(__aarch64__) && !defined(__ARM_FEATURE_SVE) + if (weight_type == GGML_TYPE_BF16) { + return GGML_TYPE_F32; + } +#endif + return ggml_internal_get_type_traits(weight_type).vec_dot_type; +} + class LLAMA_MOE_TP { private: GeneralMOEConfig config_; @@ -95,12 +147,12 @@ class LLAMA_MOE_TP { mem_requests.append_pointer(&s_input_fp32_, sizeof(float) * config_.hidden_size); mem_requests.append_pointer( &s_gate_input_, config_.hidden_size * - ggml_type_size(ggml_internal_get_type_traits((ggml_type)config_.gate_type).vec_dot_type) / - ggml_blck_size(ggml_internal_get_type_traits((ggml_type)config_.gate_type).vec_dot_type)); + ggml_type_size(kt_effective_vec_dot_type((ggml_type)config_.gate_type)) / + ggml_blck_size(kt_effective_vec_dot_type((ggml_type)config_.gate_type))); mem_requests.append_pointer( &s_up_input_, config_.hidden_size * - ggml_type_size(ggml_internal_get_type_traits((ggml_type)config_.up_type).vec_dot_type) / - ggml_blck_size(ggml_internal_get_type_traits((ggml_type)config_.up_type).vec_dot_type)); + ggml_type_size(kt_effective_vec_dot_type((ggml_type)config_.up_type)) / + ggml_blck_size(kt_effective_vec_dot_type((ggml_type)config_.up_type))); s_gate_output_.resize(config_.num_experts_per_tok); s_up_output_.resize(config_.num_experts_per_tok); s_intermediate_fp32_.resize(config_.num_experts_per_tok); @@ -113,8 +165,8 @@ class LLAMA_MOE_TP { mem_requests.append_pointer( &s_down_input_[i], config_.intermediate_size * - ggml_type_size(ggml_internal_get_type_traits((ggml_type)config_.down_type).vec_dot_type) / - ggml_blck_size(ggml_internal_get_type_traits((ggml_type)config_.down_type).vec_dot_type)); + ggml_type_size(kt_effective_vec_dot_type((ggml_type)config_.down_type)) / + ggml_blck_size(kt_effective_vec_dot_type((ggml_type)config_.down_type))); mem_requests.append_pointer(&s_down_output_[i], sizeof(float) * config_.hidden_size); } mem_requests.append_pointer(&s_output_fp32_, sizeof(float) * config_.hidden_size); @@ -129,22 +181,22 @@ class LLAMA_MOE_TP { mem_requests.append_pointer( &m_gate_input_[i], config_.hidden_size * - ggml_type_size(ggml_internal_get_type_traits((ggml_type)config_.gate_type).vec_dot_type) / - ggml_blck_size(ggml_internal_get_type_traits((ggml_type)config_.gate_type).vec_dot_type)); + ggml_type_size(kt_effective_vec_dot_type((ggml_type)config_.gate_type)) / + ggml_blck_size(kt_effective_vec_dot_type((ggml_type)config_.gate_type))); mem_requests.append_pointer( &m_up_input_[i], config_.hidden_size * - ggml_type_size(ggml_internal_get_type_traits((ggml_type)config_.up_type).vec_dot_type) / - ggml_blck_size(ggml_internal_get_type_traits((ggml_type)config_.up_type).vec_dot_type)); + ggml_type_size(kt_effective_vec_dot_type((ggml_type)config_.up_type)) / + ggml_blck_size(kt_effective_vec_dot_type((ggml_type)config_.up_type))); } mem_requests.append_pointer( &m_local_gate_input_, config_.num_experts_per_tok * config_.group_max_len * config_.hidden_size * - ggml_type_size(ggml_internal_get_type_traits((ggml_type)config_.gate_type).vec_dot_type) / - ggml_blck_size(ggml_internal_get_type_traits((ggml_type)config_.gate_type).vec_dot_type)); + ggml_type_size(kt_effective_vec_dot_type((ggml_type)config_.gate_type)) / + ggml_blck_size(kt_effective_vec_dot_type((ggml_type)config_.gate_type))); mem_requests.append_pointer( &m_local_up_input_, config_.num_experts_per_tok * config_.group_max_len * config_.hidden_size * - ggml_type_size(ggml_internal_get_type_traits((ggml_type)config_.up_type).vec_dot_type) / - ggml_blck_size(ggml_internal_get_type_traits((ggml_type)config_.up_type).vec_dot_type)); + ggml_type_size(kt_effective_vec_dot_type((ggml_type)config_.up_type)) / + ggml_blck_size(kt_effective_vec_dot_type((ggml_type)config_.up_type))); mem_requests.append_pointer(&m_local_gate_output_, sizeof(float) * config_.num_experts_per_tok * config_.group_max_len * config_.intermediate_size); mem_requests.append_pointer(&m_local_up_output_, sizeof(float) * config_.num_experts_per_tok * @@ -154,8 +206,8 @@ class LLAMA_MOE_TP { mem_requests.append_pointer( &m_local_down_input_, config_.num_experts_per_tok * config_.group_max_len * config_.intermediate_size * - ggml_type_size(ggml_internal_get_type_traits((ggml_type)config_.down_type).vec_dot_type) / - ggml_blck_size(ggml_internal_get_type_traits((ggml_type)config_.down_type).vec_dot_type)); + ggml_type_size(kt_effective_vec_dot_type((ggml_type)config_.down_type)) / + ggml_blck_size(kt_effective_vec_dot_type((ggml_type)config_.down_type))); mem_requests.append_pointer(&m_local_down_output_, sizeof(float) * config_.num_experts_per_tok * config_.group_max_len * config_.hidden_size); m_output_fp32_.resize(config_.group_max_len); @@ -190,9 +242,6 @@ class LLAMA_MOE_TP { } void load_weights(int complete_intermediate_size, int offset) { - auto local_gate_proj = m_local_gate_proj_; - auto local_up_proj = m_local_up_proj_; - auto local_down_proj = m_local_down_proj_; auto& config = config_; // printf("gate load weights:"); // debug_quant(config.gate_proj, (ggml_type)config.gate_type); @@ -221,31 +270,52 @@ class LLAMA_MOE_TP { uint8_t* down_proj = (uint8_t*)config.down_proj + offset * ggml_type_size((ggml_type)config.down_type) / ggml_blck_size((ggml_type)config.down_type); - for (int i = 0; i < config.expert_num; ++i) { - memcpy(local_gate_proj, gate_proj, - config.intermediate_size * config.hidden_size * ggml_type_size((ggml_type)config.gate_type) / - ggml_blck_size((ggml_type)config.gate_type)); - memcpy(local_up_proj, up_proj, - config.intermediate_size * config.hidden_size * ggml_type_size((ggml_type)config.up_type) / - ggml_blck_size((ggml_type)config.up_type)); + // Per-expert byte strides. The source tensors are laid out with the FULL + // intermediate_size (complete_intermediate_size); this TP only owns the + // [offset, offset+intermediate_size) block — hence the base-pointer offset + // above (src strides) and the smaller local destination strides below. + const size_t gate_dst_stride = (size_t)config.intermediate_size * config.hidden_size * + ggml_type_size((ggml_type)config.gate_type) / + ggml_blck_size((ggml_type)config.gate_type); + const size_t gate_src_stride = (size_t)complete_intermediate_size * config.hidden_size * + ggml_type_size((ggml_type)config.gate_type) / + ggml_blck_size((ggml_type)config.gate_type); + const size_t up_dst_stride = (size_t)config.intermediate_size * config.hidden_size * + ggml_type_size((ggml_type)config.up_type) / ggml_blck_size((ggml_type)config.up_type); + const size_t up_src_stride = (size_t)complete_intermediate_size * config.hidden_size * + ggml_type_size((ggml_type)config.up_type) / ggml_blck_size((ggml_type)config.up_type); + const size_t down_dst_row = (size_t)config.intermediate_size * ggml_type_size((ggml_type)config.down_type) / + ggml_blck_size((ggml_type)config.down_type); + const size_t down_src_row = (size_t)complete_intermediate_size * ggml_type_size((ggml_type)config.down_type) / + ggml_blck_size((ggml_type)config.down_type); + const size_t down_dst_stride = (size_t)config.hidden_size * down_dst_row; + const size_t down_src_stride = (size_t)config.hidden_size * down_src_row; + + uint8_t* const local_gate_base = m_local_gate_proj_; + uint8_t* const local_up_base = m_local_up_proj_; + uint8_t* const local_down_base = m_local_down_proj_; + + // Copy one expert's gate/up/down into the (disjoint) local buffers. Experts + // write non-overlapping destination regions and read disjoint source spans, + // so this is embarrassingly parallel across i. + auto copy_expert = [&](int i) { + memcpy(local_gate_base + (size_t)i * gate_dst_stride, gate_proj + (size_t)i * gate_src_stride, gate_dst_stride); + memcpy(local_up_base + (size_t)i * up_dst_stride, up_proj + (size_t)i * up_src_stride, up_dst_stride); + uint8_t* ld = local_down_base + (size_t)i * down_dst_stride; + uint8_t* sd = down_proj + (size_t)i * down_src_stride; for (int j = 0; j < config.hidden_size; ++j) { - memcpy(local_down_proj, down_proj, - config.intermediate_size * ggml_type_size((ggml_type)config.down_type) / - ggml_blck_size((ggml_type)config.down_type)); - local_down_proj += config.intermediate_size * ggml_type_size((ggml_type)config.down_type) / - ggml_blck_size((ggml_type)config.down_type); - down_proj += complete_intermediate_size * ggml_type_size((ggml_type)config.down_type) / - ggml_blck_size((ggml_type)config.down_type); + memcpy(ld, sd, down_dst_row); + ld += down_dst_row; + sd += down_src_row; } - local_gate_proj += config.intermediate_size * config.hidden_size * ggml_type_size((ggml_type)config.gate_type) / - ggml_blck_size((ggml_type)config.gate_type); - local_up_proj += config.intermediate_size * config.hidden_size * ggml_type_size((ggml_type)config.up_type) / - ggml_blck_size((ggml_type)config.up_type); - gate_proj += complete_intermediate_size * config.hidden_size * ggml_type_size((ggml_type)config.gate_type) / - ggml_blck_size((ggml_type)config.gate_type); - up_proj += complete_intermediate_size * config.hidden_size * ggml_type_size((ggml_type)config.up_type) / - ggml_blck_size((ggml_type)config.up_type); - } + }; + + // Parallelize the per-expert reshuffle across this NUMA subpool's worker + // threads. The legacy serial loop left each TP's load 1-wide (8-wide overall + // via do_numa_job) on a 192-core box. Mirrors forward()'s + // get_subpool(tp_part_idx)->do_work_stealing_job nesting inside do_numa_job. + config_.pool->get_subpool(tp_part_idx)->do_work_stealing_job(config.expert_num, + [&](int i) { copy_expert(i); }); } void warm_up() { @@ -264,37 +334,46 @@ class LLAMA_MOE_TP { } } - static float act_fn(float x) { return x / (1.0f + expf(-x)); } + static float act_fn(float gate, float up, float swiglu_limit) { + if (swiglu_limit > 0.0f) { + gate = fminf(gate, swiglu_limit); + up = fmaxf(-swiglu_limit, fminf(up, swiglu_limit)); + } + return gate / (1.0f + expf(-gate)) * up; + } void forward_one(int k, const int64_t* expert_ids, const float* weights, const void* input, float* output) { auto pool = config_.pool->get_subpool(tp_part_idx); #ifdef FORWARD_TIME_PROFILE auto t0 = std::chrono::high_resolution_clock::now(); #endif + const bool kt_pt = kt_phase_timing_on(); + std::chrono::high_resolution_clock::time_point kt_pt0, kt_pt1, kt_pt2; + if (kt_pt) kt_pt0 = std::chrono::high_resolution_clock::now(); const void* gate_input_ptr; const void* up_input_ptr; - if ((ggml_type)config_.hidden_type == ggml_internal_get_type_traits((ggml_type)config_.gate_type).vec_dot_type && - (ggml_type)config_.hidden_type == ggml_internal_get_type_traits((ggml_type)config_.up_type).vec_dot_type) { + if ((ggml_type)config_.hidden_type == kt_effective_vec_dot_type((ggml_type)config_.gate_type) && + (ggml_type)config_.hidden_type == kt_effective_vec_dot_type((ggml_type)config_.up_type)) { gate_input_ptr = up_input_ptr = input; } else { to_float(input, s_input_fp32_, config_.hidden_size, (ggml_type)config_.hidden_type); - if (ggml_internal_get_type_traits((ggml_type)config_.gate_type).vec_dot_type == - ggml_internal_get_type_traits((ggml_type)config_.up_type).vec_dot_type) { + if (kt_effective_vec_dot_type((ggml_type)config_.gate_type) == + kt_effective_vec_dot_type((ggml_type)config_.up_type)) { from_float(s_input_fp32_, s_gate_input_, config_.hidden_size, - ggml_internal_get_type_traits((ggml_type)config_.gate_type).vec_dot_type); + kt_effective_vec_dot_type((ggml_type)config_.gate_type)); gate_input_ptr = up_input_ptr = s_gate_input_; } else { if ((ggml_type)config_.hidden_type != - ggml_internal_get_type_traits((ggml_type)config_.gate_type).vec_dot_type) { + kt_effective_vec_dot_type((ggml_type)config_.gate_type)) { from_float(s_input_fp32_, s_gate_input_, config_.hidden_size, - ggml_internal_get_type_traits((ggml_type)config_.gate_type).vec_dot_type); + kt_effective_vec_dot_type((ggml_type)config_.gate_type)); gate_input_ptr = s_gate_input_; } else { gate_input_ptr = input; } - if ((ggml_type)config_.hidden_type != ggml_internal_get_type_traits((ggml_type)config_.up_type).vec_dot_type) { + if ((ggml_type)config_.hidden_type != kt_effective_vec_dot_type((ggml_type)config_.up_type)) { from_float(s_input_fp32_, s_up_input_, config_.hidden_size, - ggml_internal_get_type_traits((ggml_type)config_.up_type).vec_dot_type); + kt_effective_vec_dot_type((ggml_type)config_.up_type)); up_input_ptr = s_up_input_; } else { up_input_ptr = input; @@ -305,16 +384,18 @@ class LLAMA_MOE_TP { #ifdef FORWARD_TIME_PROFILE // printf("gate_input: "); // debug_quant(const_cast(gate_input_ptr), - // ggml_internal_get_type_traits((ggml_type)config_.gate_type).vec_dot_type); + // kt_effective_vec_dot_type((ggml_type)config_.gate_type)); // printf("up_input: "); // debug_quant(const_cast(up_input_ptr), - // ggml_internal_get_type_traits((ggml_type)config_.up_type).vec_dot_type); + // kt_effective_vec_dot_type((ggml_type)config_.up_type)); auto t1 = std::chrono::high_resolution_clock::now(); fmt::print("numa_node: {}, convert time: {}\n", tp_part_idx, std::chrono::duration_cast(t1 - t0).count()); #endif + if (kt_pt) kt_pt1 = std::chrono::high_resolution_clock::now(); + int activated_expert = 0; for (int i = 0; i < k; i++) { if (config_.should_skip_expert(expert_ids[i])) { @@ -349,7 +430,7 @@ class LLAMA_MOE_TP { config_.hidden_size / ggml_blck_size((ggml_type)config_.gate_type), gate_input_ptr, config_.hidden_size / ggml_blck_size((ggml_type)config_.gate_type), gate_output_ptr, config_.m_block, 0, 1, GGML_TASK_TYPE_COMPUTE, (ggml_type)config_.gate_type, - ggml_internal_get_type_traits((ggml_type)config_.gate_type).vec_dot_type, GGML_TYPE_F32, + kt_effective_vec_dot_type((ggml_type)config_.gate_type), GGML_TYPE_F32, GGML_PREC_DEFAULT); if (ok == false) [[unlikely]] { throw std::runtime_error("llamafile not supported"); @@ -365,33 +446,34 @@ class LLAMA_MOE_TP { up_proj_ptr, config_.hidden_size / ggml_blck_size((ggml_type)config_.up_type), up_input_ptr, config_.hidden_size / ggml_blck_size((ggml_type)config_.up_type), up_output_ptr, config_.m_block, 0, 1, GGML_TASK_TYPE_COMPUTE, (ggml_type)config_.up_type, - ggml_internal_get_type_traits((ggml_type)config_.up_type).vec_dot_type, GGML_TYPE_F32, + kt_effective_vec_dot_type((ggml_type)config_.up_type), GGML_TYPE_F32, GGML_PREC_DEFAULT); for (int i = ith * config_.m_block; i < (ith + 1) * config_.m_block; i++) { - s_intermediate_fp32_[act_idx][i] = act_fn(s_gate_output_[act_idx][i]) * s_up_output_[act_idx][i]; + s_intermediate_fp32_[act_idx][i] = + act_fn(s_gate_output_[act_idx][i], s_up_output_[act_idx][i], config_.swiglu_limit); } if (config_.m_block % - ggml_blck_size(ggml_internal_get_type_traits((ggml_type)config_.down_type).vec_dot_type) == + ggml_blck_size(kt_effective_vec_dot_type((ggml_type)config_.down_type)) == 0) { float* intermediate_fp32_ptr = s_intermediate_fp32_[act_idx] + ith * config_.m_block; void* down_input_ptr = s_down_input_[act_idx] + ith * config_.m_block * - ggml_type_size(ggml_internal_get_type_traits((ggml_type)config_.down_type).vec_dot_type) / - ggml_blck_size(ggml_internal_get_type_traits((ggml_type)config_.down_type).vec_dot_type); + ggml_type_size(kt_effective_vec_dot_type((ggml_type)config_.down_type)) / + ggml_blck_size(kt_effective_vec_dot_type((ggml_type)config_.down_type)); from_float(intermediate_fp32_ptr, down_input_ptr, config_.m_block, - ggml_internal_get_type_traits((ggml_type)config_.down_type).vec_dot_type); + kt_effective_vec_dot_type((ggml_type)config_.down_type)); } }, nullptr); } - if (config_.m_block % ggml_blck_size(ggml_internal_get_type_traits((ggml_type)config_.down_type).vec_dot_type) != + if (config_.m_block % ggml_blck_size(kt_effective_vec_dot_type((ggml_type)config_.down_type)) != 0) { for (int i = 0; i < activated_expert; i++) { from_float(s_intermediate_fp32_[i], s_down_input_[i], config_.intermediate_size, - ggml_internal_get_type_traits((ggml_type)config_.down_type).vec_dot_type); + kt_effective_vec_dot_type((ggml_type)config_.down_type)); } } @@ -402,6 +484,7 @@ class LLAMA_MOE_TP { fmt::print("numa_node: {}, gate/up time: {}\n", tp_part_idx, std::chrono::duration_cast(t2 - t1).count()); #endif + if (kt_pt) kt_pt2 = std::chrono::high_resolution_clock::now(); nth = config_.hidden_size / config_.m_block; pool->do_work_stealing_job( @@ -429,7 +512,7 @@ class LLAMA_MOE_TP { down_proj_ptr, config_.intermediate_size / ggml_blck_size((ggml_type)config_.down_type), s_down_input_[expert_idx], config_.intermediate_size / ggml_blck_size((ggml_type)config_.down_type), down_output_ptr, config_.m_block, 0, 1, GGML_TASK_TYPE_COMPUTE, (ggml_type)config_.down_type, - ggml_internal_get_type_traits((ggml_type)config_.down_type).vec_dot_type, GGML_TYPE_F32, + kt_effective_vec_dot_type((ggml_type)config_.down_type), GGML_TYPE_F32, GGML_PREC_DEFAULT); float expert_weight = 0.0f; @@ -454,6 +537,22 @@ class LLAMA_MOE_TP { fmt::print("numa_node: {}, total time: {}\n", tp_part_idx, std::chrono::duration_cast(t3 - t0).count()); #endif + if (kt_pt) { + auto kt_pt3 = std::chrono::high_resolution_clock::now(); + auto ns = [](auto a, auto b) { + return (uint64_t)std::chrono::duration_cast(b - a).count(); + }; + auto& acc = g_kt_phase_acc[tp_part_idx & 15]; + acc.calls++; + acc.quant_ns += ns(kt_pt0, kt_pt1); + acc.gateup_ns += ns(kt_pt1, kt_pt2); + acc.down_ns += ns(kt_pt2, kt_pt3); + if (acc.calls % 4300 == 0) { + fprintf(stderr, "[KT_PHASE tp%d] n=%llu avg/layer-call: quant=%.1fus gateup=%.1fus down=%.1fus\n", + tp_part_idx, (unsigned long long)acc.calls, acc.quant_ns / 1e3 / acc.calls, + acc.gateup_ns / 1e3 / acc.calls, acc.down_ns / 1e3 / acc.calls); + } + } } void forward_many(int qlen, int k, const int64_t* expert_ids, const float* weights, const void* input, @@ -485,21 +584,21 @@ class LLAMA_MOE_TP { m_local_gate_input_ptr_[i] = m_local_gate_input_ + offset * config_.hidden_size * - ggml_type_size(ggml_internal_get_type_traits((ggml_type)config_.gate_type).vec_dot_type) / - ggml_blck_size(ggml_internal_get_type_traits((ggml_type)config_.gate_type).vec_dot_type); + ggml_type_size(kt_effective_vec_dot_type((ggml_type)config_.gate_type)) / + ggml_blck_size(kt_effective_vec_dot_type((ggml_type)config_.gate_type)); m_local_up_input_ptr_[i] = m_local_up_input_ + offset * config_.hidden_size * - ggml_type_size(ggml_internal_get_type_traits((ggml_type)config_.up_type).vec_dot_type) / - ggml_blck_size(ggml_internal_get_type_traits((ggml_type)config_.up_type).vec_dot_type); + ggml_type_size(kt_effective_vec_dot_type((ggml_type)config_.up_type)) / + ggml_blck_size(kt_effective_vec_dot_type((ggml_type)config_.up_type)); m_local_gate_output_ptr_[i] = m_local_gate_output_ + offset * config_.intermediate_size; m_local_up_output_ptr_[i] = m_local_up_output_ + offset * config_.intermediate_size; m_local_intermediate_fp32_ptr_[i] = m_local_intermediate_fp32_ + offset * config_.intermediate_size; m_local_down_input_ptr_[i] = m_local_down_input_ + offset * config_.intermediate_size * - ggml_type_size(ggml_internal_get_type_traits((ggml_type)config_.down_type).vec_dot_type) / - ggml_blck_size(ggml_internal_get_type_traits((ggml_type)config_.down_type).vec_dot_type); + ggml_type_size(kt_effective_vec_dot_type((ggml_type)config_.down_type)) / + ggml_blck_size(kt_effective_vec_dot_type((ggml_type)config_.down_type)); m_local_down_output_ptr_[i] = m_local_down_output_ + offset * config_.hidden_size; offset += m_local_num_[i]; if (m_local_num_[i] > 0) { @@ -525,9 +624,9 @@ class LLAMA_MOE_TP { const void* gate_input_ptr; const void* up_input_ptr; if ((ggml_type)config_.hidden_type == - ggml_internal_get_type_traits((ggml_type)config_.gate_type).vec_dot_type && + kt_effective_vec_dot_type((ggml_type)config_.gate_type) && (ggml_type)config_.hidden_type == - ggml_internal_get_type_traits((ggml_type)config_.up_type).vec_dot_type) { + kt_effective_vec_dot_type((ggml_type)config_.up_type)) { gate_input_ptr = up_input_ptr = (uint8_t*)input + i * config_.hidden_size * ggml_type_size((ggml_type)config_.hidden_type) / ggml_blck_size((ggml_type)config_.hidden_type); @@ -535,16 +634,16 @@ class LLAMA_MOE_TP { to_float((uint8_t*)input + i * config_.hidden_size * ggml_type_size((ggml_type)config_.hidden_type) / ggml_blck_size((ggml_type)config_.hidden_type), m_input_fp32_[i], config_.hidden_size, (ggml_type)config_.hidden_type); - if (ggml_internal_get_type_traits((ggml_type)config_.gate_type).vec_dot_type == - ggml_internal_get_type_traits((ggml_type)config_.up_type).vec_dot_type) { + if (kt_effective_vec_dot_type((ggml_type)config_.gate_type) == + kt_effective_vec_dot_type((ggml_type)config_.up_type)) { from_float(m_input_fp32_[i], m_gate_input_[i], config_.hidden_size, - ggml_internal_get_type_traits((ggml_type)config_.gate_type).vec_dot_type); + kt_effective_vec_dot_type((ggml_type)config_.gate_type)); gate_input_ptr = up_input_ptr = m_gate_input_[i]; } else { if ((ggml_type)config_.hidden_type != - ggml_internal_get_type_traits((ggml_type)config_.gate_type).vec_dot_type) { + kt_effective_vec_dot_type((ggml_type)config_.gate_type)) { from_float(m_input_fp32_[i], m_gate_input_[i], config_.hidden_size, - ggml_internal_get_type_traits((ggml_type)config_.gate_type).vec_dot_type); + kt_effective_vec_dot_type((ggml_type)config_.gate_type)); gate_input_ptr = m_gate_input_[i]; } else { gate_input_ptr = (uint8_t*)input + i * config_.hidden_size * @@ -552,9 +651,9 @@ class LLAMA_MOE_TP { ggml_blck_size((ggml_type)config_.hidden_type); } if ((ggml_type)config_.hidden_type != - ggml_internal_get_type_traits((ggml_type)config_.up_type).vec_dot_type) { + kt_effective_vec_dot_type((ggml_type)config_.up_type)) { from_float(m_input_fp32_[i], m_up_input_[i], config_.hidden_size, - ggml_internal_get_type_traits((ggml_type)config_.up_type).vec_dot_type); + kt_effective_vec_dot_type((ggml_type)config_.up_type)); up_input_ptr = m_up_input_[i]; } else { up_input_ptr = (uint8_t*)input + i * config_.hidden_size * @@ -569,20 +668,20 @@ class LLAMA_MOE_TP { } memcpy(m_local_gate_input_ptr_[expert_ids[i * k + j]] + m_local_pos_[i][j] * config_.hidden_size * - ggml_type_size(ggml_internal_get_type_traits((ggml_type)config_.gate_type).vec_dot_type) / - ggml_blck_size(ggml_internal_get_type_traits((ggml_type)config_.gate_type).vec_dot_type), + ggml_type_size(kt_effective_vec_dot_type((ggml_type)config_.gate_type)) / + ggml_blck_size(kt_effective_vec_dot_type((ggml_type)config_.gate_type)), gate_input_ptr, config_.hidden_size * - ggml_type_size(ggml_internal_get_type_traits((ggml_type)config_.gate_type).vec_dot_type) / - ggml_blck_size(ggml_internal_get_type_traits((ggml_type)config_.gate_type).vec_dot_type)); + ggml_type_size(kt_effective_vec_dot_type((ggml_type)config_.gate_type)) / + ggml_blck_size(kt_effective_vec_dot_type((ggml_type)config_.gate_type))); memcpy(m_local_up_input_ptr_[expert_ids[i * k + j]] + m_local_pos_[i][j] * config_.hidden_size * - ggml_type_size(ggml_internal_get_type_traits((ggml_type)config_.up_type).vec_dot_type) / - ggml_blck_size(ggml_internal_get_type_traits((ggml_type)config_.up_type).vec_dot_type), + ggml_type_size(kt_effective_vec_dot_type((ggml_type)config_.up_type)) / + ggml_blck_size(kt_effective_vec_dot_type((ggml_type)config_.up_type)), up_input_ptr, config_.hidden_size * - ggml_type_size(ggml_internal_get_type_traits((ggml_type)config_.up_type).vec_dot_type) / - ggml_blck_size(ggml_internal_get_type_traits((ggml_type)config_.up_type).vec_dot_type)); + ggml_type_size(kt_effective_vec_dot_type((ggml_type)config_.up_type)) / + ggml_blck_size(kt_effective_vec_dot_type((ggml_type)config_.up_type))); } }, nullptr); @@ -623,7 +722,7 @@ class LLAMA_MOE_TP { config_.hidden_size / ggml_blck_size((ggml_type)config_.gate_type), gate_input_ptr, config_.hidden_size / ggml_blck_size((ggml_type)config_.gate_type), gate_output_ptr, config_.intermediate_size, 0, 1, GGML_TASK_TYPE_COMPUTE, (ggml_type)config_.gate_type, - ggml_internal_get_type_traits((ggml_type)config_.gate_type).vec_dot_type, GGML_TYPE_F32, + kt_effective_vec_dot_type((ggml_type)config_.gate_type), GGML_TYPE_F32, GGML_PREC_DEFAULT); void* up_input_ptr = m_local_up_input_ptr_[expert_idx]; @@ -638,25 +737,26 @@ class LLAMA_MOE_TP { up_proj_ptr, config_.hidden_size / ggml_blck_size((ggml_type)config_.up_type), up_input_ptr, config_.hidden_size / ggml_blck_size((ggml_type)config_.up_type), up_output_ptr, config_.intermediate_size, 0, 1, GGML_TASK_TYPE_COMPUTE, (ggml_type)config_.up_type, - ggml_internal_get_type_traits((ggml_type)config_.up_type).vec_dot_type, GGML_TYPE_F32, GGML_PREC_DEFAULT); + kt_effective_vec_dot_type((ggml_type)config_.up_type), GGML_TYPE_F32, GGML_PREC_DEFAULT); for (int i = 0; i < m_local_num_[expert_idx]; i++) { for (int j = ith * m_block; j < (ith + 1) * m_block; j++) { m_local_intermediate_fp32_ptr_[expert_idx][i * config_.intermediate_size + j] = - act_fn(m_local_gate_output_ptr_[expert_idx][i * config_.intermediate_size + j]) * - m_local_up_output_ptr_[expert_idx][i * config_.intermediate_size + j]; + act_fn(m_local_gate_output_ptr_[expert_idx][i * config_.intermediate_size + j], + m_local_up_output_ptr_[expert_idx][i * config_.intermediate_size + j], + config_.swiglu_limit); } float* intermediate_fp32_ptr = m_local_intermediate_fp32_ptr_[expert_idx] + i * config_.intermediate_size + ith * m_block; void* down_input_ptr = m_local_down_input_ptr_[expert_idx] + i * config_.intermediate_size * - ggml_type_size(ggml_internal_get_type_traits((ggml_type)config_.down_type).vec_dot_type) / - ggml_blck_size(ggml_internal_get_type_traits((ggml_type)config_.down_type).vec_dot_type) + + ggml_type_size(kt_effective_vec_dot_type((ggml_type)config_.down_type)) / + ggml_blck_size(kt_effective_vec_dot_type((ggml_type)config_.down_type)) + ith * m_block * - ggml_type_size(ggml_internal_get_type_traits((ggml_type)config_.down_type).vec_dot_type) / - ggml_blck_size(ggml_internal_get_type_traits((ggml_type)config_.down_type).vec_dot_type); + ggml_type_size(kt_effective_vec_dot_type((ggml_type)config_.down_type)) / + ggml_blck_size(kt_effective_vec_dot_type((ggml_type)config_.down_type)); from_float(intermediate_fp32_ptr, down_input_ptr, m_block, - ggml_internal_get_type_traits((ggml_type)config_.down_type).vec_dot_type); + kt_effective_vec_dot_type((ggml_type)config_.down_type)); } }, nullptr); @@ -691,7 +791,7 @@ class LLAMA_MOE_TP { config_.intermediate_size / ggml_blck_size((ggml_type)config_.down_type), down_input_ptr, config_.intermediate_size / ggml_blck_size((ggml_type)config_.down_type), down_output_ptr, config_.hidden_size, 0, 1, GGML_TASK_TYPE_COMPUTE, (ggml_type)config_.down_type, - ggml_internal_get_type_traits((ggml_type)config_.down_type).vec_dot_type, GGML_TYPE_F32, + kt_effective_vec_dot_type((ggml_type)config_.down_type), GGML_TYPE_F32, GGML_PREC_DEFAULT); }, nullptr); @@ -788,31 +888,58 @@ class TP_MOE : public TP_MOE_Common { void merge_results(int qlen, void* output, bool incremental) { auto pool = this->config.pool; + const bool kt_pt = kt_phase_timing_on(); + std::chrono::high_resolution_clock::time_point kt_m0; + if (kt_pt) kt_m0 = std::chrono::high_resolution_clock::now(); + // Tile over (token, hidden-chunk) so decode (qlen=1) still spreads across the + // whole pool instead of running the 8-NUMA reduce + from_float on a single core. + // F-opt Phase 1 (Session D): merge was ~98us/layer single-core at qlen=1 + // (~11% of cpu_moe_wall). Only chunk when hidden_type is unblocked (BF16/F16/F32, + // blck==1) so an arbitrary element boundary is always valid; block-quant hidden + // types fall back to per-token (num_chunks=1), preserving the original behavior. + const int H = config.hidden_size; + const ggml_type htype = (ggml_type)config.hidden_type; + const size_t hsz = ggml_type_size(htype); + const int hblck = ggml_blck_size(htype); + int num_chunks = 1; + if (hblck == 1) { + num_chunks = H / 256; + if (num_chunks < 1) num_chunks = 1; + if (num_chunks > 32) num_chunks = 32; + } + const int chunk = (H + num_chunks - 1) / num_chunks; pool->do_work_stealing_job( - qlen, nullptr, - [this, output, incremental](int token_nth) { + qlen * num_chunks, nullptr, + [this, output, incremental, H, htype, hsz, hblck, num_chunks, chunk](int task_id) { + const int token_nth = task_id / num_chunks; + const int c = task_id % num_chunks; + const int e0 = c * chunk; + const int e1 = std::min(H, e0 + chunk); + if (e0 >= e1) return; + float* base0 = local_output_numa[0] + token_nth * H; if (incremental) { - to_float((uint8_t*)output + token_nth * config.hidden_size * ggml_type_size((ggml_type)config.hidden_type) / - ggml_blck_size((ggml_type)config.hidden_type), - local_output + token_nth * config.hidden_size, config.hidden_size, (ggml_type)config.hidden_type); - for (int e = 0; e < config.hidden_size; e++) { - local_output_numa[0][token_nth * config.hidden_size + e] += - local_output[token_nth * config.hidden_size + e]; - } + to_float((uint8_t*)output + (size_t)(token_nth * H + e0) * hsz / hblck, + local_output + token_nth * H + e0, e1 - e0, htype); + for (int e = e0; e < e1; e++) base0[e] += local_output[token_nth * H + e]; } - auto& tp_count = this->tp_count; - for (int i = 1; i < tp_count; i++) { - for (int e = 0; e < config.hidden_size; e++) { - local_output_numa[0][token_nth * config.hidden_size + e] += - local_output_numa[i][token_nth * config.hidden_size + e]; - } + for (int i = 1; i < this->tp_count; i++) { + const float* basei = local_output_numa[i] + token_nth * H; + for (int e = e0; e < e1; e++) base0[e] += basei[e]; } - from_float(local_output_numa[0] + token_nth * config.hidden_size, - (uint8_t*)output + token_nth * config.hidden_size * ggml_type_size((ggml_type)config.hidden_type) / - ggml_blck_size((ggml_type)config.hidden_type), - config.hidden_size, (ggml_type)config.hidden_type); + from_float(base0 + e0, (uint8_t*)output + (size_t)(token_nth * H + e0) * hsz / hblck, e1 - e0, htype); }, nullptr); + if (kt_pt) { + auto kt_m1 = std::chrono::high_resolution_clock::now(); + // 槽位 15 专用于 merge（forward_one 只用 tp 0..7） + auto& acc = g_kt_phase_acc[15]; + acc.calls++; + acc.quant_ns += (uint64_t)std::chrono::duration_cast(kt_m1 - kt_m0).count(); + if (acc.calls % 4300 == 0) { + fprintf(stderr, "[KT_PHASE merge] n=%llu avg/layer-call: merge=%.1fus (qlen=%d)\n", + (unsigned long long)acc.calls, acc.quant_ns / 1e3 / acc.calls, qlen); + } + } } }; #endif diff --git a/kt-kernel/pyproject.toml b/kt-kernel/pyproject.toml index f877ddae1..1e13cb7b0 100644 --- a/kt-kernel/pyproject.toml +++ b/kt-kernel/pyproject.toml @@ -21,7 +21,9 @@ classifiers = [ requires-python = ">=3.11" dependencies = [ # Core dependencies - "torch==2.9.1", + # Keep the kernel usable with both the 2.9 reference stack and CANN 9 images + # whose torch-npu package requires torch 2.10 exactly. + "torch>=2.9.1,<2.11", "safetensors>=0.4.0", "numpy>=1.24.0", "triton>=2.0.0", diff --git a/kt-kernel/python/experts.py b/kt-kernel/python/experts.py index fcb1bc1bb..5c18e2eb8 100644 --- a/kt-kernel/python/experts.py +++ b/kt-kernel/python/experts.py @@ -359,19 +359,17 @@ def _create_inference_wrapper( raise NotImplementedError(f"Unsupported inference method: {method}") # Create and return backend instance. - # `swiglu_limit != 0` is meaningful only on the MXFP4 path. NativeMoEWrapper - # also serves RAWINT4 / FP8 / BF16 / FP8_PERCHANNEL / GPTQ_INT4, so a - # `backend_cls is NativeMoEWrapper` test would silently forward a stale - # 10.0 (e.g., from a leftover SGLANG_DSV4_2604_SUBMODE=2604B in the env) - # into a non-MXFP4 backend; act_fn would then clamp gate/up to ±10 with - # no warning. Gate strictly on method instead. Origin: kt-sglang 耦合. + # `swiglu_limit != 0` is meaningful only on native MXFP4/MXFP8 or the + # LLAMAFILE backend. LLAMAFILE validates the GGUF tensor types while loading + # and applies the clamp in its scalar/NEON-independent activation step. extra_kwargs = {} - if method in ("MXFP4", "MXFP8"): + if method in ("MXFP4", "MXFP8", "LLAMAFILE"): extra_kwargs["swiglu_limit"] = swiglu_limit extra_kwargs["swiglu_alpha"] = swiglu_alpha elif swiglu_limit != 0.0: raise ValueError( - f"swiglu_limit={swiglu_limit} is only supported on method='MXFP4'/'MXFP8', " + f"swiglu_limit={swiglu_limit} is only supported on " + "method='MXFP4'/'MXFP8'/'LLAMAFILE', " f"got method={method!r} (backend={backend_cls.__name__}). This " f"usually means SGLANG_DSV4_2604_SUBMODE=2604B is set in the " f"environment while the current launch does not actually use " diff --git a/kt-kernel/python/experts_base.py b/kt-kernel/python/experts_base.py index bc58b98f0..eb8d073b6 100644 --- a/kt-kernel/python/experts_base.py +++ b/kt-kernel/python/experts_base.py @@ -18,6 +18,104 @@ from kt_kernel import kt_kernel_ext +# ----------------------------------------------------------------------------- +# NPU stream-callback bypass. +# +# On Ascend NPU, `CPUInfer::submit_with_cuda_stream` calls `aclrtLaunchCallback`, +# which inserts the function into a per-stream **callback report queue**. ACL +# requires a dedicated subscriber thread (registered via `aclrtSubscribeReport` +# and continuously running `aclrtProcessReport`) to dispatch those callbacks. +# +# kt-kernel does NOT currently start such a worker (see +# `cpu_backend/vendors/ascend_npu.h` — marked TODO Phase 3). Without a +# subscriber, queued callbacks **silently never fire**, so the CPU forward +# task never runs and `output_cpu` stays all-zero. `sync_with_cuda_stream` +# likewise schedules its sync_ via the same callback queue and silently +# completes without actually syncing anything. +# +# When the ACL callback worker is active (ascend_callback_worker.cpp, started +# via ``init_ascend_callback_worker``), ``submit_with_cuda_stream`` callbacks +# are dispatched and CPU/NPU overlap works. Set ``KT_FORCE_SYNC_SUBMIT=1`` to +# force the legacy synchronous path for debugging. +# ----------------------------------------------------------------------------- + + +def _should_bypass_stream_callback(device: torch.device) -> bool: + """Return True iff we must use the synchronous submit/sync path.""" + if os.environ.get("KT_FORCE_SYNC_SUBMIT", "") == "1": + return True + return False + + +def _uses_external_npu_report_subscriber() -> bool: + return os.environ.get("KT_EXTERNAL_NPU_REPORT_SUBSCRIBER", "") == "1" + + +def _ensure_ascend_callback_worker() -> None: + """Start kt-kernel ACL callback worker (idempotent).""" + if _uses_external_npu_report_subscriber(): + return + if not hasattr(kt_kernel_ext, "init_ascend_callback_worker"): + return + if getattr(_ensure_ascend_callback_worker, "_done", False): + return + kt_kernel_ext.init_ascend_callback_worker() + _ensure_ascend_callback_worker._done = True # type: ignore[attr-defined] + if hasattr(kt_kernel_ext, "shutdown_ascend_callback_worker"): + import atexit + + atexit.register(kt_kernel_ext.shutdown_ascend_callback_worker) + + +def _sglang_is_capture_mode() -> bool: + """True when sglang is inside ``model_capture_mode()`` (graph capture). + + kt-kernel must remain importable standalone, so the sglang dependency is + optional and any failure is treated as "not capturing". + """ + try: + from sglang.srt.model_executor.runner import get_is_capture_mode + + return bool(get_is_capture_mode()) + except Exception: + return False + + +def _wait_device(device: torch.device) -> None: + """Block until pending async copies on `device`'s current stream finish. + + NOTE on graph capture: torch.{cuda,npu}.synchronize() raises during cuda / + NPU graph capture (NPU returns ERR 107027 "stream is captured"). Skip sync + while capturing; graph MoE uses ``_launch_host_func`` + pinned buffers + (see ``kt_ep_wrapper``). + """ + if device.type == "npu": + try: + if torch.npu.is_current_stream_capturing(): + return + except Exception: + pass + # Defensive fallback: torch.npu.is_current_stream_capturing() reliability + # during torch_npu graph capture is unconfirmed; if it returns False (or + # raises) while capturing, the synchronize() below would attempt a + # stream sync on a captured stream and crash (107027/107030). Mirror the + # capture detection used by kt_ep_wrapper._npu_use_graph_host_callback by + # also consulting sglang's global capture flag, which model_capture_mode() + # sets reliably around the whole capture loop. + if _sglang_is_capture_mode(): + return + torch.npu.synchronize(device) + elif device.type == "cuda": + try: + if torch.cuda.is_current_stream_capturing(): + return + except Exception: + pass + if _sglang_is_capture_mode(): + return + torch.cuda.synchronize(device) + + def generate_gpu_experts_masks( activation_freq: torch.Tensor, num_gpu_experts: int, @@ -174,6 +272,11 @@ def _get_cpu_infer( CPUInfer singleton instance """ if cls._cpu_infer_instance is None: + try: + if torch.npu.is_available(): # type: ignore[attr-defined] + _ensure_ascend_callback_worker() + except Exception: + pass worker_config = kt_kernel_ext.WorkerPoolConfig() if numa_nodes is not None: @@ -374,24 +477,35 @@ def select_deferred_experts( return immediate_ids, deferred_ids - def submit_forward( + def _check_qlen_fits_cpp_buffers(self, hidden_states: torch.Tensor) -> None: + """Fail loudly when qlen would overrun the C++ MoE output buffer. + + ``moe-tp.hpp`` sizes ``local_output_numa[i]`` by ``max_possible_qlen()`` = + ``max(max_len, group_max_len)``, and both are set to ``chunked_prefill_size`` + (``utils/llamafile.py``). ``TP::forward`` then hands the *full* qlen to + ``MOE::forward``, whose recursion only splits the internal scratch — the + caller-supplied output pointer still advances across ``qlen * hidden_size``. + So ``qlen > chunked_prefill_size`` writes past the allocation and corrupts the + heap, surfacing later as an unrelated ``malloc(): unaligned tcache chunk`` + abort. Raise here instead, mirroring the SFT path (``sft/base.py``). + """ + qlen = hidden_states.numel() // hidden_states.shape[-1] + if qlen > self.chunked_prefill_size: + raise ValueError( + f"qlen ({qlen}) exceeds chunked_prefill_size ({self.chunked_prefill_size}); " + "the C++ MoE output buffer is sized by chunked_prefill_size and would be " + "overrun. Raise --chunked-prefill-size or reduce the prefill chunk." + ) + + def _prepare_forward_cpu_buffers( self, hidden_states: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor, - cuda_stream, - ): - """ - Submit forward inference task to CPU (non-blocking). - - Args: - hidden_states: Input hidden states [batch_size, hidden_size] - topk_ids: Top-k expert IDs [batch_size, num_experts_per_tok] - topk_weights: Top-k expert weights [batch_size, num_experts_per_tok] - cuda_stream: CUDA stream for synchronization - """ + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], tuple, int, int]: + """D2H copy into pinned CPU buffers; return deferred ids and buffer handles.""" + self._check_qlen_fits_cpp_buffers(hidden_states) flat_hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) - batch_size = flat_hidden_states.shape[0] ( input_tensor_cpu, @@ -405,16 +519,14 @@ def submit_forward( current_slot = self.layer_idx % KExpertsCPUBuffer.buffer_depth next_slot = (current_slot + 1) % KExpertsCPUBuffer.buffer_depth - bsz_slot_tensor = bsz_tensor_cpu[current_slot] topk_ids_long = topk_ids.to(torch.long) - immediate_ids: torch.Tensor - deferred_ids: Optional[torch.Tensor] if self.max_deferred_experts_per_token > 0: protected_k = self.num_experts_per_tok - self.max_deferred_experts_per_token - - immediate_ids, deferred_ids = self.select_deferred_experts(topk_ids_long, topk_weights, protected_k) + immediate_ids, deferred_ids = self.select_deferred_experts( + topk_ids_long, topk_weights, protected_k + ) else: immediate_ids = topk_ids_long deferred_ids = None @@ -422,38 +534,274 @@ def submit_forward( input_tensor_cpu[current_slot].copy_(flat_hidden_states, non_blocking=True) weights_cpu[current_slot].copy_(topk_weights, non_blocking=True) immediate_experts_ids_cpu[current_slot].copy_(immediate_ids, non_blocking=True) + if deferred_ids is not None: + deferred_experts_ids_cpu[current_slot].copy_(deferred_ids, non_blocking=True) + + buffers = ( + input_tensor_cpu, + immediate_experts_ids_cpu, + deferred_experts_ids_cpu, + weights_cpu, + output_cpu, + bsz_tensor_cpu, + _output_gpu, + ) + return immediate_ids, deferred_ids, buffers, current_slot, next_slot + def copy_inputs_to_cpu_buffers( + self, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + ) -> None: + """Copy MoE inputs to pinned CPU buffers (for NPU graph host callbacks).""" + self._prepare_forward_cpu_buffers(hidden_states, topk_ids, topk_weights) + + def forward_on_pinned_buffers( + self, + hidden_states: torch.Tensor, + cuda_stream, + ) -> None: + """Run CPU MoE on buffers already filled (sync or stream callback).""" + self._check_qlen_fits_cpp_buffers(hidden_states) + flat_hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + ( + input_tensor_cpu, + immediate_experts_ids_cpu, + deferred_experts_ids_cpu, + weights_cpu, + output_cpu, + bsz_tensor_cpu, + _output_gpu, + ) = KExpertsCPUBuffer.get_buffer(flat_hidden_states, self.num_experts_per_tok) + + current_slot = self.layer_idx % KExpertsCPUBuffer.buffer_depth + next_slot = (current_slot + 1) % KExpertsCPUBuffer.buffer_depth + bsz_slot_tensor = bsz_tensor_cpu[current_slot] + + bypass = _should_bypass_stream_callback(hidden_states.device) incremental = BaseMoEWrapper._layer_has_pending_deferred.get(self.layer_idx - 1, False) - self.cpu_infer.submit_with_cuda_stream( - cuda_stream, - self.moe.forward_task( + immediate_task = self.moe.forward_task( + bsz_slot_tensor.data_ptr(), + immediate_experts_ids_cpu[current_slot].size(-1), + immediate_experts_ids_cpu[current_slot].data_ptr(), + weights_cpu[current_slot].data_ptr(), + input_tensor_cpu[current_slot].data_ptr(), + output_cpu[current_slot].data_ptr(), + incremental, + ) + if bypass: + self.cpu_infer.submit(immediate_task) + else: + # Correct + fast async path. ``submit_with_cuda_stream`` enqueues the + # CPU-MoE via an ACL host callback whose firing is not host-observable + # (NO_BLOCK) -> a later host-side drain can race ahead of it (empty + # queue -> stale output_cpu read by the H2D; nondeterministic on heavy + # prefill). Enqueue synchronously on this host thread instead (work is + # guaranteed submitted; it still runs on the WorkerPool and overlaps the + # GPU experts queued after). Keep subscribe_ascend_stream — that stream + # registration is what keeps decode's host-callback dispatch fast; the + # async submit itself is not required for it. + if ( + hidden_states.device.type == "npu" + and not _uses_external_npu_report_subscriber() + and hasattr( + kt_kernel_ext, "subscribe_ascend_stream" + ) + ): + kt_kernel_ext.subscribe_ascend_stream(int(cuda_stream)) + self.cpu_infer.submit(immediate_task) + + BaseMoEWrapper._layer_has_pending_deferred[self.layer_idx] = False + has_deferred = ( + self.max_deferred_experts_per_token > 0 + and (deferred_experts_ids_cpu[current_slot] >= 0).any().item() + ) + if has_deferred: + deferred_task = self.moe.forward_task( bsz_slot_tensor.data_ptr(), - immediate_experts_ids_cpu[current_slot].size(-1), - immediate_experts_ids_cpu[current_slot].data_ptr(), + deferred_experts_ids_cpu[current_slot].size(-1), + deferred_experts_ids_cpu[current_slot].data_ptr(), weights_cpu[current_slot].data_ptr(), input_tensor_cpu[current_slot].data_ptr(), - output_cpu[current_slot].data_ptr(), - incremental, - ), + output_cpu[next_slot].data_ptr(), + False, + ) + if bypass: + self.cpu_infer.submit(deferred_task) + else: + self.cpu_infer.submit_with_cuda_stream(cuda_stream, deferred_task) + BaseMoEWrapper._layer_has_pending_deferred[self.layer_idx] = True + + def run_pinned_forward_sync( + self, + hidden_states: torch.Tensor, + cuda_stream, + ) -> None: + """Submit + sync CPU MoE on pre-filled buffers (NPU graph host callback). + + Called from ``aclrtLaunchCallback`` / ``_launch_host_func``; must not enqueue + nested stream callbacks. + """ + del cuda_stream # unused — sync path only + self._check_qlen_fits_cpp_buffers(hidden_states) + flat_hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + ( + input_tensor_cpu, + immediate_experts_ids_cpu, + deferred_experts_ids_cpu, + weights_cpu, + output_cpu, + bsz_tensor_cpu, + _output_gpu, + ) = KExpertsCPUBuffer.get_buffer(flat_hidden_states, self.num_experts_per_tok) + + current_slot = self.layer_idx % KExpertsCPUBuffer.buffer_depth + next_slot = (current_slot + 1) % KExpertsCPUBuffer.buffer_depth + bsz_slot_tensor = bsz_tensor_cpu[current_slot] + + incremental = BaseMoEWrapper._layer_has_pending_deferred.get(self.layer_idx - 1, False) + immediate_task = self.moe.forward_task( + bsz_slot_tensor.data_ptr(), + immediate_experts_ids_cpu[current_slot].size(-1), + immediate_experts_ids_cpu[current_slot].data_ptr(), + weights_cpu[current_slot].data_ptr(), + input_tensor_cpu[current_slot].data_ptr(), + output_cpu[current_slot].data_ptr(), + incremental, + ) + self.cpu_infer.submit(immediate_task) + BaseMoEWrapper._layer_has_pending_deferred[self.layer_idx] = False + has_deferred = ( + self.max_deferred_experts_per_token > 0 + and (deferred_experts_ids_cpu[current_slot] >= 0).any().item() + ) + if has_deferred: + deferred_task = self.moe.forward_task( + bsz_slot_tensor.data_ptr(), + deferred_experts_ids_cpu[current_slot].size(-1), + deferred_experts_ids_cpu[current_slot].data_ptr(), + weights_cpu[current_slot].data_ptr(), + input_tensor_cpu[current_slot].data_ptr(), + output_cpu[next_slot].data_ptr(), + False, + ) + self.cpu_infer.submit(deferred_task) + BaseMoEWrapper._layer_has_pending_deferred[self.layer_idx] = True + allow_pending = 1 if BaseMoEWrapper._layer_has_pending_deferred.get(self.layer_idx, False) else 0 + self.cpu_infer.sync(allow_pending) + + def submit_forward( + self, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + cuda_stream, + ): + """ + Submit forward inference task to CPU (non-blocking). + + Args: + hidden_states: Input hidden states [batch_size, hidden_size] + topk_ids: Top-k expert IDs [batch_size, num_experts_per_tok] + topk_weights: Top-k expert weights [batch_size, num_experts_per_tok] + cuda_stream: CUDA stream for synchronization + """ + _immediate_ids, deferred_ids, _buffers, current_slot, next_slot = ( + self._prepare_forward_cpu_buffers(hidden_states, topk_ids, topk_weights) + ) + ( + input_tensor_cpu, + immediate_experts_ids_cpu, + deferred_experts_ids_cpu, + weights_cpu, + output_cpu, + bsz_tensor_cpu, + _output_gpu, + ) = _buffers + bsz_slot_tensor = bsz_tensor_cpu[current_slot] + + bypass = _should_bypass_stream_callback(hidden_states.device) + # Both paths now submit the CPU-MoE synchronously on this host thread, which + # reads input_tensor_cpu immediately -> the input D2H (queued async on the + # stream by copy_inputs_to_cpu_buffers) MUST be finished first, else the CPU + # MoE reads a half-copied input. Wait unconditionally (bypass already did). + _wait_device(hidden_states.device) + + incremental = BaseMoEWrapper._layer_has_pending_deferred.get(self.layer_idx - 1, False) + immediate_task = self.moe.forward_task( + bsz_slot_tensor.data_ptr(), + immediate_experts_ids_cpu[current_slot].size(-1), + immediate_experts_ids_cpu[current_slot].data_ptr(), + weights_cpu[current_slot].data_ptr(), + input_tensor_cpu[current_slot].data_ptr(), + output_cpu[current_slot].data_ptr(), + incremental, ) + if bypass: + self.cpu_infer.submit(immediate_task) + else: + # Correct + fast async path. ``submit_with_cuda_stream`` enqueues the + # CPU-MoE via an ACL host callback whose firing is not host-observable + # (NO_BLOCK) -> a later host-side drain can race ahead of it (empty + # queue -> stale output_cpu read by the H2D; nondeterministic on heavy + # prefill). Enqueue synchronously on this host thread instead (work is + # guaranteed submitted; it still runs on the WorkerPool and overlaps the + # GPU experts queued after). Keep subscribe_ascend_stream — that stream + # registration is what keeps decode's host-callback dispatch fast; the + # async submit itself is not required for it. + if ( + hidden_states.device.type == "npu" + and not _uses_external_npu_report_subscriber() + and hasattr( + kt_kernel_ext, "subscribe_ascend_stream" + ) + ): + kt_kernel_ext.subscribe_ascend_stream(int(cuda_stream)) + self.cpu_infer.submit(immediate_task) BaseMoEWrapper._layer_has_pending_deferred[self.layer_idx] = False if deferred_ids is not None: - deferred_experts_ids_cpu[current_slot].copy_(deferred_ids, non_blocking=True) - self.cpu_infer.submit_with_cuda_stream( - cuda_stream, - self.moe.forward_task( - bsz_slot_tensor.data_ptr(), - deferred_experts_ids_cpu[current_slot].size(-1), - deferred_experts_ids_cpu[current_slot].data_ptr(), - weights_cpu[current_slot].data_ptr(), - input_tensor_cpu[current_slot].data_ptr(), - output_cpu[next_slot].data_ptr(), - False, - ), + _wait_device(hidden_states.device) + deferred_task = self.moe.forward_task( + bsz_slot_tensor.data_ptr(), + deferred_experts_ids_cpu[current_slot].size(-1), + deferred_experts_ids_cpu[current_slot].data_ptr(), + weights_cpu[current_slot].data_ptr(), + input_tensor_cpu[current_slot].data_ptr(), + output_cpu[next_slot].data_ptr(), + False, ) + if bypass: + self.cpu_infer.submit(deferred_task) + else: + if ( + hidden_states.device.type == "npu" + and not _uses_external_npu_report_subscriber() + and hasattr( + kt_kernel_ext, "subscribe_ascend_stream" + ) + ): + kt_kernel_ext.subscribe_ascend_stream(int(cuda_stream)) + self.cpu_infer.submit(deferred_task) BaseMoEWrapper._layer_has_pending_deferred[self.layer_idx] = True + def copy_forward_output_to_device(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Copy pinned CPU output to the device tensor (CPU work already finished).""" + flat_hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + ( + _input_tensor_cpu, + _immediate_experts_ids_cpu, + _deferred_experts_ids_cpu, + _weights_cpu, + output_cpu, + _bsz_tensor_cpu, + output_gpu, + ) = KExpertsCPUBuffer.get_buffer(flat_hidden_states, self.num_experts_per_tok) + current_slot = self.layer_idx % KExpertsCPUBuffer.buffer_depth + output_gpu[current_slot].copy_(output_cpu[current_slot], non_blocking=True) + return output_gpu[current_slot] + def sync_forward(self, hidden_states: torch.Tensor, cuda_stream) -> torch.Tensor: """ Synchronize and retrieve forward inference results. @@ -478,9 +826,49 @@ def sync_forward(self, hidden_states: torch.Tensor, cuda_stream) -> torch.Tensor current_slot = self.layer_idx % KExpertsCPUBuffer.buffer_depth allow_pending = 1 if BaseMoEWrapper._layer_has_pending_deferred.get(self.layer_idx, False) else 0 - self.cpu_infer.sync_with_cuda_stream(cuda_stream, allow_pending) - output_gpu[current_slot].copy_(output_cpu[current_slot], non_blocking=True) - return output_gpu[current_slot] + bypass = _should_bypass_stream_callback(hidden_states.device) + if bypass: + self.cpu_infer.sync(allow_pending) + else: + if ( + hidden_states.device.type == "npu" + and not _uses_external_npu_report_subscriber() + and hasattr( + kt_kernel_ext, "subscribe_ascend_stream" + ) + ): + kt_kernel_ext.subscribe_ascend_stream(int(cuda_stream)) + if hidden_states.device.type == "npu": + # ACL host callbacks do not order the following H2D copy behind + # the WorkerPool drain. Wait for the accelerator work that + # overlapped CPU MoE, then drain on the host before copying the + # completed pinned output back to the NPU. + _wait_device(hidden_states.device) + self.cpu_infer.sync(allow_pending) + else: + self.cpu_infer.sync_with_cuda_stream(cuda_stream, allow_pending) + + if os.environ.get("KT_DEBUG_MOE_OUT", "") == "1": + oc = output_cpu[current_slot] + ic = _input_tensor_cpu[current_slot] + wc = _weights_cpu[current_slot] + ec = _immediate_experts_ids_cpu[current_slot] + finite_oc = torch.isfinite(oc).all().item() + finite_ic = torch.isfinite(ic).all().item() + finite_wc = torch.isfinite(wc).all().item() + nan_count = int(torch.isnan(oc).sum().item()) + inf_count = int(torch.isinf(oc).sum().item()) + print( + f"[KT_DEBUG] layer={self.layer_idx} slot={current_slot} bypass={bypass} " + f"output_cpu: finite={finite_oc} nan={nan_count}/{oc.numel()} inf={inf_count} " + f"first16={oc.float().flatten()[:16].tolist()} | " + f"input_cpu.finite={finite_ic} first8={ic.float().flatten()[:8].tolist()} | " + f"weights_cpu.finite={finite_wc} first={wc.flatten()[:8].tolist()} | " + f"expert_ids={ec.flatten()[:12].tolist()}", + flush=True, + ) + + return self.copy_forward_output_to_device(hidden_states) def forward( self, @@ -541,4 +929,3 @@ def clear_buffer_cache(): KExpertsCPUBuffer.capture_buffers.clear() KExpertsCPUBuffer.temp_bs = 0 KExpertsCPUBuffer.temp_buffer = tuple() - diff --git a/kt-kernel/python/utils/llamafile.py b/kt-kernel/python/utils/llamafile.py index 0f1c4403c..9a105d817 100644 --- a/kt-kernel/python/utils/llamafile.py +++ b/kt-kernel/python/utils/llamafile.py @@ -1,6 +1,9 @@ -import torch -from typing import List, Optional +from __future__ import annotations + import os +from typing import Dict, List, Optional + +import torch # Use relative imports for package structure from ..experts_base import BaseMoEWrapper @@ -22,9 +25,13 @@ class LlamafileMoEWrapper(BaseMoEWrapper): """ Llamafile-based MoE wrapper implementation. Supports GGUF quantized weights with llamafile backend. + + GGUFLoader is cached **per resolved weight path** (file or directory): multiple MoE layers + that share one merged GGUF reuse a single mmap; **per-layer split GGUFs** each get their own + loader (Phase 2-B). """ - _gguf_loader_instance = None # Singleton GGUFLoader + _gguf_loaders_by_path: Dict[str, GGUFLoader] = {} def __init__( self, @@ -42,6 +49,8 @@ def __init__( max_deferred_experts_per_token: Optional[int] = None, method: str = "LLAMAFILE", numa_nodes: Optional[List[int]] = None, + swiglu_limit: float = 0.0, + swiglu_alpha: float = 0.0, ): """ Initialize Llamafile MoE Wrapper. @@ -73,10 +82,10 @@ def __init__( if not os.path.exists(weight_path): raise FileNotFoundError(f"GGUF weight path not found: {weight_path}") - # Initialize GGUF loader (singleton) - if LlamafileMoEWrapper._gguf_loader_instance is None: - LlamafileMoEWrapper._gguf_loader_instance = GGUFLoader(weight_path) - self.gguf_loader = LlamafileMoEWrapper._gguf_loader_instance + cache_key = os.path.realpath(weight_path) + if cache_key not in LlamafileMoEWrapper._gguf_loaders_by_path: + LlamafileMoEWrapper._gguf_loaders_by_path[cache_key] = GGUFLoader(weight_path) + self.gguf_loader = LlamafileMoEWrapper._gguf_loaders_by_path[cache_key] # Validate TP configuration with QK_K alignment QK_K = 256 @@ -119,6 +128,8 @@ def __init__( print(f" TP {tp_id}: size={tp_size}, offset={current_offset}, blocks={tp_blocks}") current_offset += tp_size + self._swiglu_alpha = float(swiglu_alpha) + # Initialize base class super().__init__( layer_idx=layer_idx, @@ -135,6 +146,7 @@ def __init__( max_deferred_experts_per_token=max_deferred_experts_per_token, method=method, numa_nodes=numa_nodes, + swiglu_limit=swiglu_limit, ) self.weights_to_keep = None @@ -176,12 +188,28 @@ def load_weights(self, physical_to_logical_map_cpu: Optional[torch.Tensor] = Non base_key = f"blk.{self.layer_idx}" - # Load quantized tensors from GGUF - gate_data, gate_type = self.gguf_loader.get_undequanted_tensor_and_ggml_type(f"{base_key}.ffn_gate_exps.weight") - - up_data, up_type = self.gguf_loader.get_undequanted_tensor_and_ggml_type(f"{base_key}.ffn_up_exps.weight") - - down_data, down_type = self.gguf_loader.get_undequanted_tensor_and_ggml_type(f"{base_key}.ffn_down_exps.weight") + # KT_DUMMY_CPU_WEIGHTS: skip the multi-GB GGUF disk read and fabricate + # same-layout zero buffers so the C++ MOE / load_weights_task path runs + # identically (buffer sizing, kernel selection, capture) but loads much + # faster. Debug-only: weights are meaningless — never use for accuracy. + if os.environ.get("KT_DUMMY_CPU_WEIGHTS"): + if self.layer_idx == 0: + print( + "[LlamafileMoEWrapper] KT_DUMMY_CPU_WEIGHTS=1: fabricating zero " + "CPU MoE weights (skips GGUF read). Outputs are MEANINGLESS — " + "for capture iteration only, not accuracy validation.", + flush=True, + ) + _load_tensor = self.gguf_loader.get_dummy_tensor_and_ggml_type + else: + _load_tensor = self.gguf_loader.get_undequanted_tensor_and_ggml_type + + # Load quantized tensors from GGUF (or fabricate dummies) + gate_data, gate_type = _load_tensor(f"{base_key}.ffn_gate_exps.weight") + + up_data, up_type = _load_tensor(f"{base_key}.ffn_up_exps.weight") + + down_data, down_type = _load_tensor(f"{base_key}.ffn_down_exps.weight") # Keep tensors alive self.weights_to_keep = (gate_data, up_data, down_data) @@ -202,8 +230,26 @@ def load_weights(self, physical_to_logical_map_cpu: Optional[torch.Tensor] = Non # Llamafile-specific configuration moe_config.m_block = 32 # Parallel block size moe_config.group_min_len = 10 # Use forward_one when qlen < 10 - moe_config.max_len = self.chunked_prefill_size - moe_config.group_max_len = max(1, int(self.chunked_prefill_size)) + # Defensive fallback: chunked_prefill_size <= 0 (e.g. -1 meaning "disabled" in sglang + # baseline) would otherwise let C++ compute max_possible_qlen() = max(max_len=-1, + # group_max_len=max(1,-1)=1) = 1, sizing per-NUMA fp32 output buffer + # (moe-tp.hpp:130 local_output_numa[i]) to a single token. The very first prefill + # with qlen > 1 then overruns the buffer and corrupts glibc tcache metadata + # (observed as "malloc(): unaligned tcache chunk detected" Fatal Python error). + # Clamp to a safe positive value so KT can always alloc a fp32 buffer at least as + # large as the per-call qlen the caller will pass. + _effective_chunk = int(self.chunked_prefill_size) + if _effective_chunk <= 0: + _effective_chunk = 2048 + print( + f"[LlamafileMoEWrapper] chunked_prefill_size={self.chunked_prefill_size} " + f"<= 0 is unsafe for KT MoE C++ buffer sizing; falling back to " + f"{_effective_chunk}." + ) + moe_config.max_len = _effective_chunk + moe_config.group_max_len = _effective_chunk + moe_config.swiglu_limit = self.swiglu_limit + moe_config.swiglu_alpha = self._swiglu_alpha # Set weight pointers moe_config.gate_proj = gate_data.data_ptr() diff --git a/kt-kernel/python/utils/loader.py b/kt-kernel/python/utils/loader.py index c24f78937..aa6e71fb4 100644 --- a/kt-kernel/python/utils/loader.py +++ b/kt-kernel/python/utils/loader.py @@ -48,6 +48,44 @@ class GGMLQuantizationType(IntEnum): F64 = 28 IQ1_M = 29 BF16 = 30 + MXFP4 = 39 # OCP microscaling FP4 (E2M1 + ue8m0), id aligned with upstream ggml + + +# (block_size, type_size) per GGML quant type — bytes per `block_size` elements. +# Module-level so both the real loader and the dummy-weight fabrication path +# (KT_DUMMY_CPU_WEIGHTS) share one source of truth. +GGML_QUANT_SIZES = { + GGMLQuantizationType.F32: (1, 4), + GGMLQuantizationType.F16: (1, 2), + GGMLQuantizationType.BF16: (1, 2), + GGMLQuantizationType.Q4_0: (32, 2 + 16), + GGMLQuantizationType.Q4_1: (32, 2 + 2 + 16), + GGMLQuantizationType.Q5_0: (32, 2 + 4 + 16), + GGMLQuantizationType.Q5_1: (32, 2 + 2 + 4 + 16), + GGMLQuantizationType.Q8_0: (32, 2 + 32), + GGMLQuantizationType.Q8_1: (32, 4 + 4 + 32), + GGMLQuantizationType.Q2_K: (256, 2 + 2 + 256 // 16 + 256 // 4), + GGMLQuantizationType.Q3_K: (256, 2 + 256 // 4 + 256 // 8 + 12), + GGMLQuantizationType.Q4_K: (256, 2 + 2 + 256 // 2 + 12), + GGMLQuantizationType.Q5_K: (256, 2 + 2 + 256 // 2 + 256 // 8 + 12), + GGMLQuantizationType.Q6_K: (256, 2 + 256 // 2 + 256 // 4 + 256 // 16), + GGMLQuantizationType.Q8_K: (256, 4 + 256 + 256 // 8), + GGMLQuantizationType.IQ2_XXS: (256, 2 + 256 // 4), + GGMLQuantizationType.IQ2_XS: (256, 2 + 256 // 4 + 256 // 32), + GGMLQuantizationType.IQ3_XXS: (256, 2 + 256 // 4 + 256 // 8), + GGMLQuantizationType.IQ1_S: (256, 2 + 256 // 8 + 256 // 16), + GGMLQuantizationType.IQ4_NL: (32, 2 + 16), + GGMLQuantizationType.IQ3_S: (256, 2 + 256 // 4 + 256 // 8 + 256 // 32 + 4), + GGMLQuantizationType.IQ2_S: (256, 2 + 256 // 4 + 256 // 16), + GGMLQuantizationType.IQ4_XS: (256, 2 + 2 + 256 // 2 + 256 // 64), + GGMLQuantizationType.I8: (1, 1), + GGMLQuantizationType.I16: (1, 2), + GGMLQuantizationType.I32: (1, 4), + GGMLQuantizationType.I64: (1, 8), + GGMLQuantizationType.F64: (1, 8), + GGMLQuantizationType.IQ1_M: (256, 256 // 8 + 256 // 16 + 256 // 32), + GGMLQuantizationType.MXFP4: (32, 1 + 16), # 1B e8m0 scale + 16B nibble-packed E2M1 +} def translate_name_to_gguf(name): @@ -1024,44 +1062,49 @@ def get_undequanted_tensor_and_ggml_type(self, name: str): n_elements = info["n_elements"] ggml_type = info["dtype"] - GGML_QUANT_SIZES = { - GGMLQuantizationType.F32: (1, 4), - GGMLQuantizationType.F16: (1, 2), - GGMLQuantizationType.BF16: (1, 2), - GGMLQuantizationType.Q4_0: (32, 2 + 16), - GGMLQuantizationType.Q4_1: (32, 2 + 2 + 16), - GGMLQuantizationType.Q5_0: (32, 2 + 4 + 16), - GGMLQuantizationType.Q5_1: (32, 2 + 2 + 4 + 16), - GGMLQuantizationType.Q8_0: (32, 2 + 32), - GGMLQuantizationType.Q8_1: (32, 4 + 4 + 32), - GGMLQuantizationType.Q2_K: (256, 2 + 2 + 256 // 16 + 256 // 4), - GGMLQuantizationType.Q3_K: (256, 2 + 256 // 4 + 256 // 8 + 12), - GGMLQuantizationType.Q4_K: (256, 2 + 2 + 256 // 2 + 12), - GGMLQuantizationType.Q5_K: (256, 2 + 2 + 256 // 2 + 256 // 8 + 12), - GGMLQuantizationType.Q6_K: (256, 2 + 256 // 2 + 256 // 4 + 256 // 16), - GGMLQuantizationType.Q8_K: (256, 4 + 256 + 256 // 8), - GGMLQuantizationType.IQ2_XXS: (256, 2 + 256 // 4), - GGMLQuantizationType.IQ2_XS: (256, 2 + 256 // 4 + 256 // 32), - GGMLQuantizationType.IQ3_XXS: (256, 2 + 256 // 4 + 256 // 8), - GGMLQuantizationType.IQ1_S: (256, 2 + 256 // 8 + 256 // 16), - GGMLQuantizationType.IQ4_NL: (32, 2 + 16), - GGMLQuantizationType.IQ3_S: (256, 2 + 256 // 4 + 256 // 8 + 256 // 32 + 4), - GGMLQuantizationType.IQ2_S: (256, 2 + 256 // 4 + 256 // 16), - GGMLQuantizationType.IQ4_XS: (256, 2 + 2 + 256 // 2 + 256 // 64), - GGMLQuantizationType.I8: (1, 1), - GGMLQuantizationType.I16: (1, 2), - GGMLQuantizationType.I32: (1, 4), - GGMLQuantizationType.I64: (1, 8), - GGMLQuantizationType.F64: (1, 8), - GGMLQuantizationType.IQ1_M: (256, 256 // 8 + 256 // 16 + 256 // 32), - } - block_size, type_size = GGML_QUANT_SIZES[ggml_type] n_bytes = n_elements * type_size // block_size data_bytes = mmap_data[offset : offset + n_bytes] - data = torch.from_numpy(np.frombuffer(data_bytes, dtype=np.uint8).copy()) + # Zero-copy: the returned tensor is only ever consumed as a *read-only source* by the + # C++ MoE load_weights reshuffle (llamafile/moe.hpp memcpy SOURCE), which makes the + # single necessary copy into its NUMA-local buffers. So view the mmap directly instead + # of copying every layer's ~6.85GB single-threaded into anonymous RAM. The mmap stays + # alive in GGUFLoader.file_data_map; the wrapper pins the tensors in + # self.weights_to_keep until cpu_infer.sync() returns. + # (np.frombuffer over the memmap slice = read-only view, no copy; torch.from_numpy + # shares that buffer so data.data_ptr() points into the mmap. Use from_numpy, NOT + # torch.frombuffer — the latter is redirected to NPU by torch_npu's transfer_to_npu.) + data = torch.from_numpy(np.frombuffer(data_bytes, dtype=np.uint8)) + + return data, ggml_type + + def get_dummy_tensor_and_ggml_type(self, name: str): + """Fabricate a zero buffer matching a tensor's on-disk byte layout. + + Used only by the ``KT_DUMMY_CPU_WEIGHTS`` fast-iteration path: it reads + tensor metadata (shape / quant type / element count) from the cheap + ``tensor_info`` map and allocates a same-sized ``uint8`` buffer WITHOUT + touching the mmap'd weight bytes. This skips the multi-GB GGUF disk read + so graph capture can be iterated quickly; the returned values are + meaningless and must NOT be used for accuracy validation. + + Returns ``(data, ggml_type)`` with the same contract as + ``get_undequanted_tensor_and_ggml_type``. + """ + name = translate_name_to_gguf(name) + + if name not in self.tensor_info: + raise KeyError(f"Tensor '{name}' not found in GGUF files") + + info = self.tensor_info[name] + n_elements = info["n_elements"] + ggml_type = info["dtype"] + + block_size, type_size = GGML_QUANT_SIZES[ggml_type] + n_bytes = n_elements * type_size // block_size + data = torch.zeros(n_bytes, dtype=torch.uint8) return data, ggml_type diff --git a/kt-kernel/setup.py b/kt-kernel/setup.py index bff9307bf..964912826 100644 --- a/kt-kernel/setup.py +++ b/kt-kernel/setup.py @@ -34,14 +34,22 @@ CPUINFER_NATIVE=ON (override LLAMA_NATIVE) -GPU backends: +GPU/NPU backends: CPUINFER_USE_CUDA=0/1 -DKTRANSFORMERS_USE_CUDA CPUINFER_USE_SYCL=0/1 -DKTRANSFORMERS_USE_SYCL (GPTQ INT4 MoE) CPUINFER_USE_ROCM=0/1 -DKTRANSFORMERS_USE_ROCM CPUINFER_USE_MUSA=0/1 -DKTRANSFORMERS_USE_MUSA CPUINFER_USE_MACA=0/1 -DKTRANSFORMERS_USE_MACA + CPUINFER_USE_ASCEND_NPU=0/1 -DKTRANSFORMERS_USE_ASCEND_NPU MACA_PATH=/opt/maca MACA SDK root +ARM aarch64 feature toggles: + CPUINFER_ARM_DOTPROD=ON/OFF -DLLAMA_ARM_DOTPROD + CPUINFER_ARM_FP16=ON/OFF -DLLAMA_ARM_FP16 + CPUINFER_ARM_SVE=ON/OFF -DLLAMA_ARM_SVE + CPUINFER_ARM_BF16=ON/OFF -DLLAMA_ARM_BF16 + CPUINFER_ARM_I8MM=ON/OFF -DLLAMA_ARM_I8MM + Usage: pip install . Or build wheel: @@ -207,6 +215,20 @@ def detect_cpu_info(self) -> dict: flags.update(m.group(1).lower().split()) info["raw"]["flags"] = flags + # ARM feature summary (Kunpeng / Neoverse / Apple Silicon). + # /proc/cpuinfo on Linux/aarch64 lists these in "Features:". + if info["vendor"] == "arm" or "aarch64" in info["arch"]: + if "asimddp" in flags: + info["features"].add("ARM_DOTPROD") + if "asimdhp" in flags or "fphp" in flags: + info["features"].add("ARM_FP16") + if "sve" in flags: + info["features"].add("ARM_SVE") + if "bf16" in flags: + info["features"].add("ARM_BF16") + if "i8mm" in flags: + info["features"].add("ARM_I8MM") + # feature summary if any(f in flags or f in low for f in ["avx512f", "avx512bw", "avx512dq", "avx512vl"]): info["features"].add("AVX512") @@ -479,6 +501,20 @@ def detect_cuda_toolkit() -> bool: return True return False + # Auto-detect Ascend CANN toolkit if user did not explicitly set CPUINFER_USE_ASCEND_NPU. + # We check $ASCEND_TOOLKIT_HOME, $CANN_HOME, and the default install + # prefix /usr/local/Ascend/ascend-toolkit/latest; treat the toolkit as + # available iff include/acl/acl_rt.h exists under that root. + def detect_cann_toolkit() -> str | None: + for env_name in ("ASCEND_TOOLKIT_HOME", "CANN_HOME"): + p = os.environ.get(env_name) + if p and (Path(p) / "include" / "acl" / "acl_rt.h").exists(): + return p + default_root = Path("/usr/local/Ascend/ascend-toolkit/latest") + if (default_root / "include" / "acl" / "acl_rt.h").exists(): + return str(default_root) + return None + # Locate nvcc executable (without forcing user to set -DCMAKE_CUDA_COMPILER) def find_nvcc_path() -> str | None: cuda_home = os.environ.get("CUDA_HOME") @@ -522,7 +558,13 @@ def find_maca_path() -> str | None: if cuda_env is None: requested_non_cuda_gpu = any( _env_get_bool(name, False) - for name in ("CPUINFER_USE_SYCL", "CPUINFER_USE_ROCM", "CPUINFER_USE_MUSA", "CPUINFER_USE_MACA") + for name in ( + "CPUINFER_USE_SYCL", + "CPUINFER_USE_ROCM", + "CPUINFER_USE_MUSA", + "CPUINFER_USE_MACA", + "CPUINFER_USE_ASCEND_NPU", + ) ) if requested_non_cuda_gpu: os.environ["CPUINFER_USE_CUDA"] = "0" @@ -532,6 +574,20 @@ def find_maca_path() -> str | None: os.environ["CPUINFER_USE_CUDA"] = "1" if auto_cuda else "0" print(f"-- CPUINFER_USE_CUDA not set; auto-detected CUDA toolkit: {'YES' if auto_cuda else 'NO'}") + npu_env = _env_get_bool("CPUINFER_USE_ASCEND_NPU", None) + if npu_env is None: + cann_root = detect_cann_toolkit() + host_is_arm = platform.machine().lower() in ("aarch64", "arm64") + cuda_active = os.environ.get("CPUINFER_USE_CUDA") == "1" + auto_npu = cann_root is not None and host_is_arm and not cuda_active + os.environ["CPUINFER_USE_ASCEND_NPU"] = "1" if auto_npu else "0" + if auto_npu: + os.environ.setdefault("ASCEND_TOOLKIT_HOME", cann_root) + print( + "-- CPUINFER_USE_ASCEND_NPU not set; auto-detected CANN toolkit: " + + (f"YES ({cann_root})" if auto_npu else "NO") + ) + enabled_gpu_backends = [ name for name, env_name in ( @@ -540,6 +596,7 @@ def find_maca_path() -> str | None: ("ROCM", "CPUINFER_USE_ROCM"), ("MUSA", "CPUINFER_USE_MUSA"), ("MACA", "CPUINFER_USE_MACA"), + ("ASCEND", "CPUINFER_USE_ASCEND_NPU"), ) if _env_get_bool(env_name, False) ] @@ -701,6 +758,27 @@ def find_maca_path() -> str | None: if maca_path and not os.environ.get("MACA_PATH"): cmake_args.append(f"-DMACA_PATH={maca_path}") print("-- Enabling MACA backend (-DKTRANSFORMERS_USE_MACA=ON)") + if _env_get_bool("CPUINFER_USE_ASCEND_NPU", False): + cmake_args.append("-DKTRANSFORMERS_USE_ASCEND_NPU=ON") + print("-- Enabling Ascend NPU backend (-DKTRANSFORMERS_USE_ASCEND_NPU=ON)") + if not _env_get_bool("CPUINFER_ENABLE_KML", None): + os.environ.setdefault("CPUINFER_ENABLE_KML", "OFF") + if not _env_get_bool("CPUINFER_ENABLE_BLIS", None): + os.environ.setdefault("CPUINFER_ENABLE_BLIS", "OFF") + + if platform.machine().lower() in ("aarch64", "arm64"): + arm_feature_map = ( + ("CPUINFER_ARM_DOTPROD", "LLAMA_ARM_DOTPROD", "ARM_DOTPROD"), + ("CPUINFER_ARM_FP16", "LLAMA_ARM_FP16", "ARM_FP16"), + ("CPUINFER_ARM_SVE", "LLAMA_ARM_SVE", "ARM_SVE"), + ("CPUINFER_ARM_BF16", "LLAMA_ARM_BF16", "ARM_BF16"), + ("CPUINFER_ARM_I8MM", "LLAMA_ARM_I8MM", "ARM_I8MM"), + ) + for env_name, cmake_flag, feature in arm_feature_map: + if _forward_bool_env(cmake_args, env_name, cmake_flag): + continue + value = "ON" if feature in d["features"] else "OFF" + cmake_args.append(f"-D{cmake_flag}={value}") # Respect user extra CMAKE_ARGS (space separated) extra = os.environ.get("CMAKE_ARGS") diff --git a/third_party/llama.cpp b/third_party/llama.cpp index a94e6ff87..ca9486437 160000 --- a/third_party/llama.cpp +++ b/third_party/llama.cpp @@ -1 +1 @@ -Subproject commit a94e6ff8774b7c9f950d9545baf0ce35e8d1ed2f +Subproject commit ca948643712e7bb11f543b404d0ce15fd53c0d2f From ebf6405623ff39e6ce3b231785541c392016c83e Mon Sep 17 00:00:00 2001 From: Pan-Boyi Date: Thu, 13 Aug 2026 21:53:02 +0800 Subject: [PATCH 02/13] fix: support MXFP4 GGUF in llamafile sgemm Signed-off-by: Pan-Boyi --- third_party/llamafile/tinyblas_cpu_sgemm.inc | 33 +++++++++++++++++++- 1 file changed, 32 insertions(+), 1 deletion(-) diff --git a/third_party/llamafile/tinyblas_cpu_sgemm.inc b/third_party/llamafile/tinyblas_cpu_sgemm.inc index e83636b61..9663f35e4 100644 --- a/third_party/llamafile/tinyblas_cpu_sgemm.inc +++ b/third_party/llamafile/tinyblas_cpu_sgemm.inc @@ -257,6 +257,37 @@ bool llamafile_sgemm_impl(long m, long n, long k, const void* A, long lda, const #endif } + case GGML_TYPE_MXFP4: { + if (Btype != GGML_TYPE_Q8_0) + return NOT_SUPPORTED; + if (task != GGML_TASK_TYPE_COMPUTE) + return true; + + // MXFP4 is an upstream GGML type, but llamafile has no native + // tinyBLAS tile for it yet. Use the GGML vec-dot implementation + // (NEON accelerated on AArch64) so LLAMAFILE MoE can consume + // MXFP4 GGUF weights instead of rejecting them at first token. + const block_mxfp4* a = static_cast(A); + const block_q8_0* b = static_cast(B); + const long jobs = m * n; + for (long job = ith; job < jobs; job += nth) { + const long i = job % m; + const long j = job / m; + float value = 0.0f; + ggml_vec_dot_mxfp4_q8_0( + k * QK_MXFP4, + &value, + 0, + a + i * lda, + 0, + b + j * ldb, + 0, + 1); + C[j * ldc + i] = static_cast(value); + } + return true; + } + default: return NOT_SUPPORTED; } @@ -358,4 +389,4 @@ bool llamafile_sgemm(long m, long n, long k, const void* A, long lda, const void default: return NOT_SUPPORTED; } -} \ No newline at end of file +} From 26a1e73cc8fc1743528909492a16f8266500c7c4 Mon Sep 17 00:00:00 2001 From: Pan-Boyi Date: Fri, 14 Aug 2026 11:21:31 +0800 Subject: [PATCH 03/13] [chore](kt-kernel): remove internal markers from Ascend offload sources Replace internal phase/session markers and Chinese comments in the Ascend NPU offload sources with English prose, fix a stale cross-reference that claimed kt-kernel does not start an ACL callback worker (this commit series adds one), make the phase-timing report interval configurable via KT_MOE_PHASE_TIMING_INTERVAL instead of a hardcoded model-specific value, and apply black to the Python sources touched here. No functional change. Signed-off-by: Pan-Boyi --- kt-kernel/cpu_backend/vendors/ascend_npu.h | 4 +- kt-kernel/operators/llamafile/moe.hpp | 72 ++++++++++++++-------- kt-kernel/python/experts_base.py | 54 ++++++---------- kt-kernel/python/utils/llamafile.py | 2 +- 4 files changed, 68 insertions(+), 64 deletions(-) diff --git a/kt-kernel/cpu_backend/vendors/ascend_npu.h b/kt-kernel/cpu_backend/vendors/ascend_npu.h index 259dcd6ac..da763144a 100644 --- a/kt-kernel/cpu_backend/vendors/ascend_npu.h +++ b/kt-kernel/cpu_backend/vendors/ascend_npu.h @@ -19,8 +19,8 @@ // * The submit_with_cuda_stream / sync_with_cuda_stream path therefore // requires the host process to also spin up a callback worker (see // cpu_backend/ascend_callback_worker.cpp). The synchronous CPUInfer::submit() -// / sync() path works without it -// and is what the Phase 1 PoC uses. +// / sync() path works without it and is used as the fallback when no +// callback worker is available. // ============================================================================ #include diff --git a/kt-kernel/operators/llamafile/moe.hpp b/kt-kernel/operators/llamafile/moe.hpp index 809bfaab7..8beecf145 100644 --- a/kt-kernel/operators/llamafile/moe.hpp +++ b/kt-kernel/operators/llamafile/moe.hpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include @@ -26,17 +27,31 @@ #endif // --------------------------------------------------------------------------- -// KT_MOE_PHASE_TIMING=1 — Session D Phase-0 测量：decode (forward_one) 三段 -// (input量化 / gate+up job / down job) + TP 层 merge 的累计耗时，定位固定开销 F -// 的构成。env 关闭时仅一次 getenv + 分支，零扰动。每 tp 每 4300 次调用 -// (=43层×100 token) 打一行均值到 stderr。 -// 线程安全性：同一 tp_part_idx 的 forward_one 在 decode 流内串行（逐层）， -// 不同 tp 并行但写不同槽位 → 无需原子。 +// KT_MOE_PHASE_TIMING=1 accumulates the per-stage cost of decode +// (forward_one): input quantization, the gate+up job, the down job, and the +// TP-level merge. It is meant for locating the fixed per-call overhead. +// When the env var is unset the cost is a single getenv plus a branch, so the +// instrumentation is non-perturbing. One averaged line is printed to stderr +// every kt_phase_report_interval() calls per TP partition. +// Thread safety: forward_one for a given tp_part_idx runs serially within the +// decode stream (layer by layer), and different TP partitions write to +// different slots, so no atomics are required. // --------------------------------------------------------------------------- static inline bool kt_phase_timing_on() { static const bool on = std::getenv("KT_MOE_PHASE_TIMING") != nullptr; return on; } +// How many accumulated calls to average over before printing one report line. +// Override with KT_MOE_PHASE_TIMING_INTERVAL; the default is only a convenient +// order of magnitude, not a model-specific constant. +static inline uint64_t kt_phase_report_interval() { + static const uint64_t interval = [] { + const char* raw = std::getenv("KT_MOE_PHASE_TIMING_INTERVAL"); + const uint64_t parsed = raw ? std::strtoull(raw, nullptr, 10) : 0; + return parsed ? parsed : 4096; + }(); + return interval; +} struct KtPhaseAcc { uint64_t calls = 0; uint64_t quant_ns = 0, gateup_ns = 0, down_ns = 0; @@ -54,25 +69,27 @@ inline void debug_quant(void* input, ggml_type type) { // --------------------------------------------------------------------------- // kt_effective_vec_dot_type -// 解决 aarch64 (Kunpeng K920 / Cortex-A76, **no SVE / no i8mm**) 上 -// kt-kernel llamafile sgemm 的 BF16/Q8_0 路径不健全问题： +// Works around the incomplete BF16 path of the llamafile sgemm on aarch64 +// cores without SVE and without i8mm (armv8.2-a + fp16 + dotprod only): // -// * BF16 weight + BF16 input：tinyblas_cpu_sgemm.inc ARM_NEON path 要求 -// `Btype == GGML_TYPE_F32`（line 209: `if (Btype != F32) return NOT_SUPPORTED;`）， -// 但 ggml type_traits 对 BF16 给出 `vec_dot_type = BF16`， -// 上层 forward_one/forward_many 默认把 BF16 input 喂给 sgemm -// → llamafile_sgemm 返回 false → `throw "llamafile not supported"`。 +// * BF16 weight with BF16 input: the ARM_NEON path in +// tinyblas_cpu_sgemm.inc requires `Btype == GGML_TYPE_F32` and otherwise +// returns NOT_SUPPORTED, but the ggml type traits report +// `vec_dot_type = BF16` for BF16. forward_one/forward_many therefore feed +// a BF16 input to the sgemm, llamafile_sgemm returns false and the caller +// throws "llamafile not supported". // -// * 解法：在 aarch64-without-SVE 平台上，把 BF16 weight 的有效 vec_dot_type -// 声明为 F32。input 路径会因此走 to_float(bf16 → fp32) + memcpy(F32→F32 buffer) -// （`from_float()` 已在 conversion.h 对 F32 short-circuit 成 memcpy）， -// buffer 大小自动按 fp32 (4 bytes/elem) 分配 —— 比原 BF16 (2 bytes/elem) -// 大一倍，足够装 fp32 数据。sgemm 则走 ARM_NEON 已支持的 -// `Atype=BF16, Btype=F32, Ctype=F32` 路径（tinyblas_cpu_sgemm.inc line 125-133）。 +// * Fix: on aarch64 without SVE, declare the effective vec_dot_type of a +// BF16 weight to be F32. The input path then runs to_float(bf16 -> fp32) +// followed by a memcpy into the F32 buffer (from_float() already +// short-circuits to a memcpy for F32 in conversion.h), and the buffer is +// sized for fp32 (4 bytes/elem), i.e. twice the BF16 size, which is enough +// to hold the fp32 data. The sgemm then takes the +// `Atype=BF16, Btype=F32, Ctype=F32` path that ARM_NEON already supports. // -// * SVE 机器走原 BF16-BF16 path（vec_dot_type 不改），不打扰原性能优化。 -// * 其他 weight type（Q8_0/Q4_K/…）走原 vec_dot_type，不动； -// 若 Q8_0 NaN 仍存在，单独在 sgemm 内部修复，而非这里。 +// * Cores with SVE keep the original BF16-BF16 path so their existing +// performance tuning is untouched. +// * Other weight types (Q8_0, Q4_K, ...) keep their original vec_dot_type. // --------------------------------------------------------------------------- static inline ggml_type kt_effective_vec_dot_type(ggml_type weight_type) { #if defined(__aarch64__) && !defined(__ARM_FEATURE_SVE) @@ -547,7 +564,7 @@ class LLAMA_MOE_TP { acc.quant_ns += ns(kt_pt0, kt_pt1); acc.gateup_ns += ns(kt_pt1, kt_pt2); acc.down_ns += ns(kt_pt2, kt_pt3); - if (acc.calls % 4300 == 0) { + if (acc.calls % kt_phase_report_interval() == 0) { fprintf(stderr, "[KT_PHASE tp%d] n=%llu avg/layer-call: quant=%.1fus gateup=%.1fus down=%.1fus\n", tp_part_idx, (unsigned long long)acc.calls, acc.quant_ns / 1e3 / acc.calls, acc.gateup_ns / 1e3 / acc.calls, acc.down_ns / 1e3 / acc.calls); @@ -893,8 +910,9 @@ class TP_MOE : public TP_MOE_Common { if (kt_pt) kt_m0 = std::chrono::high_resolution_clock::now(); // Tile over (token, hidden-chunk) so decode (qlen=1) still spreads across the // whole pool instead of running the 8-NUMA reduce + from_float on a single core. - // F-opt Phase 1 (Session D): merge was ~98us/layer single-core at qlen=1 - // (~11% of cpu_moe_wall). Only chunk when hidden_type is unblocked (BF16/F16/F32, + // Before this change the merge ran single-core and was a significant part + // of the fixed per-layer cost at qlen=1. + // Only chunk when hidden_type is unblocked (BF16/F16/F32, // blck==1) so an arbitrary element boundary is always valid; block-quant hidden // types fall back to per-token (num_chunks=1), preserving the original behavior. const int H = config.hidden_size; @@ -931,11 +949,11 @@ class TP_MOE : public TP_MOE_Common { nullptr); if (kt_pt) { auto kt_m1 = std::chrono::high_resolution_clock::now(); - // 槽位 15 专用于 merge（forward_one 只用 tp 0..7） + // Slot 15 is reserved for the merge; forward_one only uses slots 0..7. auto& acc = g_kt_phase_acc[15]; acc.calls++; acc.quant_ns += (uint64_t)std::chrono::duration_cast(kt_m1 - kt_m0).count(); - if (acc.calls % 4300 == 0) { + if (acc.calls % kt_phase_report_interval() == 0) { fprintf(stderr, "[KT_PHASE merge] n=%llu avg/layer-call: merge=%.1fus (qlen=%d)\n", (unsigned long long)acc.calls, acc.quant_ns / 1e3 / acc.calls, qlen); } diff --git a/kt-kernel/python/experts_base.py b/kt-kernel/python/experts_base.py index eb8d073b6..71a12f790 100644 --- a/kt-kernel/python/experts_base.py +++ b/kt-kernel/python/experts_base.py @@ -17,7 +17,6 @@ from kt_kernel import kt_kernel_ext - # ----------------------------------------------------------------------------- # NPU stream-callback bypass. # @@ -26,17 +25,17 @@ # requires a dedicated subscriber thread (registered via `aclrtSubscribeReport` # and continuously running `aclrtProcessReport`) to dispatch those callbacks. # -# kt-kernel does NOT currently start such a worker (see -# `cpu_backend/vendors/ascend_npu.h` — marked TODO Phase 3). Without a -# subscriber, queued callbacks **silently never fire**, so the CPU forward -# task never runs and `output_cpu` stays all-zero. `sync_with_cuda_stream` -# likewise schedules its sync_ via the same callback queue and silently -# completes without actually syncing anything. +# Without such a subscriber, queued callbacks **silently never fire**, so the +# CPU forward task never runs and `output_cpu` stays all-zero. +# `sync_with_cuda_stream` likewise schedules its sync_ via the same callback +# queue and silently completes without actually syncing anything. # -# When the ACL callback worker is active (ascend_callback_worker.cpp, started -# via ``init_ascend_callback_worker``), ``submit_with_cuda_stream`` callbacks -# are dispatched and CPU/NPU overlap works. Set ``KT_FORCE_SYNC_SUBMIT=1`` to -# force the legacy synchronous path for debugging. +# kt-kernel therefore starts the subscriber itself (see +# `cpu_backend/ascend_callback_worker.cpp`, entered via +# ``init_ascend_callback_worker``); ``subscribe_ascend_stream`` registers each +# stream with it, after which ``submit_with_cuda_stream`` callbacks are +# dispatched and CPU/NPU overlap works. Set ``KT_FORCE_SYNC_SUBMIT=1`` to fall +# back to the synchronous submit/sync path for debugging. # ----------------------------------------------------------------------------- @@ -282,8 +281,7 @@ def _get_cpu_infer( if numa_nodes is not None: if len(numa_nodes) != threadpool_count: raise ValueError( - f"numa_nodes length ({len(numa_nodes)}) must match " - f"threadpool_count ({threadpool_count})" + f"numa_nodes length ({len(numa_nodes)}) must match " f"threadpool_count ({threadpool_count})" ) subpool_numa_map = list(numa_nodes) else: @@ -524,9 +522,7 @@ def _prepare_forward_cpu_buffers( topk_ids_long = topk_ids.to(torch.long) if self.max_deferred_experts_per_token > 0: protected_k = self.num_experts_per_tok - self.max_deferred_experts_per_token - immediate_ids, deferred_ids = self.select_deferred_experts( - topk_ids_long, topk_weights, protected_k - ) + immediate_ids, deferred_ids = self.select_deferred_experts(topk_ids_long, topk_weights, protected_k) else: immediate_ids = topk_ids_long deferred_ids = None @@ -605,17 +601,14 @@ def forward_on_pinned_buffers( if ( hidden_states.device.type == "npu" and not _uses_external_npu_report_subscriber() - and hasattr( - kt_kernel_ext, "subscribe_ascend_stream" - ) + and hasattr(kt_kernel_ext, "subscribe_ascend_stream") ): kt_kernel_ext.subscribe_ascend_stream(int(cuda_stream)) self.cpu_infer.submit(immediate_task) BaseMoEWrapper._layer_has_pending_deferred[self.layer_idx] = False has_deferred = ( - self.max_deferred_experts_per_token > 0 - and (deferred_experts_ids_cpu[current_slot] >= 0).any().item() + self.max_deferred_experts_per_token > 0 and (deferred_experts_ids_cpu[current_slot] >= 0).any().item() ) if has_deferred: deferred_task = self.moe.forward_task( @@ -673,8 +666,7 @@ def run_pinned_forward_sync( self.cpu_infer.submit(immediate_task) BaseMoEWrapper._layer_has_pending_deferred[self.layer_idx] = False has_deferred = ( - self.max_deferred_experts_per_token > 0 - and (deferred_experts_ids_cpu[current_slot] >= 0).any().item() + self.max_deferred_experts_per_token > 0 and (deferred_experts_ids_cpu[current_slot] >= 0).any().item() ) if has_deferred: deferred_task = self.moe.forward_task( @@ -707,8 +699,8 @@ def submit_forward( topk_weights: Top-k expert weights [batch_size, num_experts_per_tok] cuda_stream: CUDA stream for synchronization """ - _immediate_ids, deferred_ids, _buffers, current_slot, next_slot = ( - self._prepare_forward_cpu_buffers(hidden_states, topk_ids, topk_weights) + _immediate_ids, deferred_ids, _buffers, current_slot, next_slot = self._prepare_forward_cpu_buffers( + hidden_states, topk_ids, topk_weights ) ( input_tensor_cpu, @@ -753,9 +745,7 @@ def submit_forward( if ( hidden_states.device.type == "npu" and not _uses_external_npu_report_subscriber() - and hasattr( - kt_kernel_ext, "subscribe_ascend_stream" - ) + and hasattr(kt_kernel_ext, "subscribe_ascend_stream") ): kt_kernel_ext.subscribe_ascend_stream(int(cuda_stream)) self.cpu_infer.submit(immediate_task) @@ -778,9 +768,7 @@ def submit_forward( if ( hidden_states.device.type == "npu" and not _uses_external_npu_report_subscriber() - and hasattr( - kt_kernel_ext, "subscribe_ascend_stream" - ) + and hasattr(kt_kernel_ext, "subscribe_ascend_stream") ): kt_kernel_ext.subscribe_ascend_stream(int(cuda_stream)) self.cpu_infer.submit(deferred_task) @@ -833,9 +821,7 @@ def sync_forward(self, hidden_states: torch.Tensor, cuda_stream) -> torch.Tensor if ( hidden_states.device.type == "npu" and not _uses_external_npu_report_subscriber() - and hasattr( - kt_kernel_ext, "subscribe_ascend_stream" - ) + and hasattr(kt_kernel_ext, "subscribe_ascend_stream") ): kt_kernel_ext.subscribe_ascend_stream(int(cuda_stream)) if hidden_states.device.type == "npu": diff --git a/kt-kernel/python/utils/llamafile.py b/kt-kernel/python/utils/llamafile.py index 9a105d817..6c9b6d26d 100644 --- a/kt-kernel/python/utils/llamafile.py +++ b/kt-kernel/python/utils/llamafile.py @@ -28,7 +28,7 @@ class LlamafileMoEWrapper(BaseMoEWrapper): GGUFLoader is cached **per resolved weight path** (file or directory): multiple MoE layers that share one merged GGUF reuse a single mmap; **per-layer split GGUFs** each get their own - loader (Phase 2-B). + loader. """ _gguf_loaders_by_path: Dict[str, GGUFLoader] = {} From b5bf065a6556f399acb093e1211ebc87243191c0 Mon Sep 17 00:00:00 2001 From: Pan-Boyi Date: Fri, 14 Aug 2026 15:05:22 +0800 Subject: [PATCH 04/13] [build](kt-kernel): pin llama.cpp to upstream b3173 and apply MXFP4 as a patch The llama.cpp gitlink pointed at a fork commit that exists only in the local clone, so `git clone --recursive` of this repository failed for everyone else. Restore the pin to a94e6ff8 (refs/tags/b3173, resolvable upstream) and keep the MXFP4 delta as a patch series under kt-kernel/third_party_patches/llama.cpp/. kt-kernel/CMakeLists.txt applies the series at configure time, right before add_subdirectory() pulls llama.cpp in. Idempotency needs no marker file: an already-applied patch fails `git apply --check` but succeeds `git apply --reverse --check`, and only when both fail does configure abort with the manual recovery commands. Signed-off-by: Pan-Boyi --- .gitignore | 3 + kt-kernel/CMakeLists.txt | 64 ++++++ .../llama.cpp/0001-ggml-mxfp4-type.patch | 211 ++++++++++++++++++ .../third_party_patches/llama.cpp/README.md | 57 +++++ third_party/llama.cpp | 2 +- 5 files changed, 336 insertions(+), 1 deletion(-) create mode 100644 kt-kernel/third_party_patches/llama.cpp/0001-ggml-mxfp4-type.patch create mode 100644 kt-kernel/third_party_patches/llama.cpp/README.md diff --git a/.gitignore b/.gitignore index f44a0ab87..d5f3f568a 100644 --- a/.gitignore +++ b/.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 diff --git a/kt-kernel/CMakeLists.txt b/kt-kernel/CMakeLists.txt index bb035f152..9726c2085 100644 --- a/kt-kernel/CMakeLists.txt +++ b/kt-kernel/CMakeLists.txt @@ -471,6 +471,70 @@ add_compile_options("$<$:${ARCH_FLAGS}>") add_compile_options("$<$:${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) diff --git a/kt-kernel/third_party_patches/llama.cpp/0001-ggml-mxfp4-type.patch b/kt-kernel/third_party_patches/llama.cpp/0001-ggml-mxfp4-type.patch new file mode 100644 index 000000000..a0c31961b --- /dev/null +++ b/kt-kernel/third_party_patches/llama.cpp/0001-ggml-mxfp4-type.patch @@ -0,0 +1,211 @@ +diff --git a/ggml-common.h b/ggml-common.h +index e8efceb76..c0aebeb74 100644 +--- a/ggml-common.h ++++ b/ggml-common.h +@@ -186,6 +186,17 @@ typedef struct { + } block_q8_0; + static_assert(sizeof(block_q8_0) == sizeof(ggml_half) + QK8_0, "wrong q8_0 block size/padding"); + ++// MXFP4 (OCP microscaling FP4 E2M1, per-32 group, ue8m0 scale). ++// One byte E8M0 group exponent + 16 nibble-packed E2M1 codes (32 weights / block). ++// Block layout matches upstream llama.cpp: qs[j] low nibble = element j, ++// high nibble = element j+16 (half-block interleave). ++#define QK_MXFP4 32 ++typedef struct { ++ uint8_t e; ++ uint8_t qs[QK_MXFP4/2]; ++} block_mxfp4; ++static_assert(sizeof(block_mxfp4) == sizeof(uint8_t) + QK_MXFP4/2, "wrong mxfp4 block size/padding"); ++ + #define QK8_1 32 + typedef struct { + union { +diff --git a/ggml-quants.c b/ggml-quants.c +index 0b346c11e..a447da35d 100644 +--- a/ggml-quants.c ++++ b/ggml-quants.c +@@ -3565,6 +3565,38 @@ void dequantize_row_iq4_nl(const block_iq4_nl * restrict x, float * restrict y, + } + } + ++// MXFP4 (OCP microscaling FP4): E2M1 LUT scaled by two so the factor of one ++// half can be folded into the E8M0 scale conversion below. ++static const int8_t kvalues_mxfp4[16] = {0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12}; ++ ++static inline float ggml_e8m0_to_fp32_half(uint8_t x) { ++ uint32_t bits; ++ if (x < 2) { ++ bits = (uint32_t) 0x00200000 << x; ++ } else { ++ bits = (uint32_t) (x - 1) << 23; ++ } ++ float result; ++ memcpy(&result, &bits, sizeof(float)); ++ return result; ++} ++#define GGML_E8M0_TO_FP32_HALF(x) ggml_e8m0_to_fp32_half(x) ++ ++void dequantize_row_mxfp4(const block_mxfp4 * restrict x, float * restrict y, int64_t k) { ++ assert(k % QK_MXFP4 == 0); ++ const int64_t nb = k / QK_MXFP4; ++ ++ for (int i = 0; i < nb; i++) { ++ const uint8_t * qs = x[i].qs; ++ const float d = GGML_E8M0_TO_FP32_HALF(x[i].e); ++ for (int j = 0; j < QK_MXFP4/2; ++j) { ++ y[j] = d * kvalues_mxfp4[qs[j] & 0xf]; ++ y[j + QK_MXFP4/2] = d * kvalues_mxfp4[qs[j] >> 4]; ++ } ++ y += QK_MXFP4; ++ } ++} ++ + void dequantize_row_iq4_xs(const block_iq4_xs * restrict x, float * restrict y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; +@@ -11289,6 +11321,71 @@ void ggml_vec_dot_iq4_nl_q8_0(int n, float * restrict s, size_t bs, const void * + #endif + } + ++// MXFP4 (E2M1 + ue8m0) weights multiplied by Q8_0 activations. The AArch64 ++// path uses NEON dot products and deliberately avoids SVE/i8mm requirements so ++// it can run on Kunpeng 920. Other architectures use the scalar fallback. ++void ggml_vec_dot_mxfp4_q8_0(int n, float * restrict s, size_t bs, const void * restrict vx, size_t bx, const void * restrict vy, size_t by, int nrc) { ++ assert(nrc == 1); ++ UNUSED(nrc); ++ UNUSED(bx); ++ UNUSED(by); ++ UNUSED(bs); ++ assert(n % QK_MXFP4 == 0); ++ static_assert(QK_MXFP4 == QK8_0, "QK_MXFP4 and QK8_0 must be the same"); ++ ++ const block_mxfp4 * restrict x = vx; ++ const block_q8_0 * restrict y = vy; ++ const int nb = n / QK_MXFP4; ++ int ib = 0; ++ float sumf = 0; ++ ++#if defined __ARM_NEON ++ const int8x16_t values = vld1q_s8(kvalues_mxfp4); ++ const uint8x16_t m4b = vdupq_n_u8(0x0f); ++ uint8x16x2_t q4bits; ++ int8x16x4_t q4b; ++ int8x16x4_t q8b; ++ int32x4_t prod_1, prod_2; ++ float32x4_t sumv0 = vdupq_n_f32(0.0f); ++ float32x4_t sumv1 = vdupq_n_f32(0.0f); ++ ++ for (; ib + 1 < nb; ib += 2) { ++ __builtin_prefetch((const char *) &x[ib] + 512, 0, 0); ++ q4bits.val[0] = vld1q_u8(x[ib + 0].qs); ++ q4bits.val[1] = vld1q_u8(x[ib + 1].qs); ++ q8b.val[0] = vld1q_s8(y[ib + 0].qs); ++ q8b.val[1] = vld1q_s8(y[ib + 0].qs + 16); ++ q8b.val[2] = vld1q_s8(y[ib + 1].qs); ++ q8b.val[3] = vld1q_s8(y[ib + 1].qs + 16); ++ ++ q4b.val[0] = ggml_vqtbl1q_s8(values, vandq_u8(q4bits.val[0], m4b)); ++ q4b.val[1] = ggml_vqtbl1q_s8(values, vshrq_n_u8(q4bits.val[0], 4)); ++ q4b.val[2] = ggml_vqtbl1q_s8(values, vandq_u8(q4bits.val[1], m4b)); ++ q4b.val[3] = ggml_vqtbl1q_s8(values, vshrq_n_u8(q4bits.val[1], 4)); ++ ++ prod_1 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), q4b.val[0], q8b.val[0]), q4b.val[1], q8b.val[1]); ++ prod_2 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), q4b.val[2], q8b.val[2]), q4b.val[3], q8b.val[3]); ++ ++ sumv0 = vfmaq_n_f32(sumv0, vcvtq_f32_s32(prod_1), ++ GGML_E8M0_TO_FP32_HALF(x[ib + 0].e) * GGML_FP16_TO_FP32(y[ib + 0].d)); ++ sumv1 = vfmaq_n_f32(sumv1, vcvtq_f32_s32(prod_2), ++ GGML_E8M0_TO_FP32_HALF(x[ib + 1].e) * GGML_FP16_TO_FP32(y[ib + 1].d)); ++ } ++ sumf = vaddvq_f32(vaddq_f32(sumv0, sumv1)); ++#endif ++ for (; ib < nb; ++ib) { ++ const float d = GGML_FP16_TO_FP32(y[ib].d) * GGML_E8M0_TO_FP32_HALF(x[ib].e); ++ int sumi1 = 0; ++ int sumi2 = 0; ++ for (int j = 0; j < QK_MXFP4/2; ++j) { ++ sumi1 += y[ib].qs[j] * kvalues_mxfp4[x[ib].qs[j] & 0xf]; ++ sumi2 += y[ib].qs[j + QK_MXFP4/2] * kvalues_mxfp4[x[ib].qs[j] >> 4]; ++ } ++ sumf += d * (sumi1 + sumi2); ++ } ++ *s = sumf; ++} ++ + void ggml_vec_dot_iq4_xs_q8_K(int n, float * restrict s, size_t bs, const void * restrict vx, size_t bx, const void * restrict vy, size_t by, int nrc) { + assert(nrc == 1); + UNUSED(nrc); +diff --git a/ggml-quants.h b/ggml-quants.h +index 4d436a8f0..a9a272fdc 100644 +--- a/ggml-quants.h ++++ b/ggml-quants.h +@@ -76,6 +76,7 @@ void dequantize_row_iq1_m (const block_iq1_m * GGML_RESTRICT x, float * GGML_ + void dequantize_row_iq4_nl (const block_iq4_nl * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); + void dequantize_row_iq4_xs (const block_iq4_xs * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); + void dequantize_row_iq3_s (const block_iq3_s * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); ++void dequantize_row_mxfp4 (const block_mxfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); + + // Dot product + void ggml_vec_dot_q4_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +@@ -97,6 +98,7 @@ void ggml_vec_dot_iq3_xxs_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const + void ggml_vec_dot_iq1_s_q8_K (int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + void ggml_vec_dot_iq1_m_q8_K (int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + void ggml_vec_dot_iq4_nl_q8_0 (int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); ++void ggml_vec_dot_mxfp4_q8_0 (int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + void ggml_vec_dot_iq4_xs_q8_K (int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + void ggml_vec_dot_iq3_s_q8_K (int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +diff --git a/ggml.c b/ggml.c +index d5d33c2ba..6fa04fd3c 100644 +--- a/ggml.c ++++ b/ggml.c +@@ -902,6 +902,17 @@ static const ggml_type_traits_t type_traits[GGML_TYPE_COUNT] = { + .vec_dot = (ggml_vec_dot_t) ggml_vec_dot_bf16, + .vec_dot_type = GGML_TYPE_BF16, + .nrows = 1, ++ }, ++ [GGML_TYPE_MXFP4] = { ++ // OCP microscaling FP4 (E2M1 + ue8m0), weight-only checkpoint format. ++ .type_name = "mxfp4", ++ .blck_size = QK_MXFP4, ++ .type_size = sizeof(block_mxfp4), ++ .is_quantized = true, ++ .to_float = (ggml_to_float_t) dequantize_row_mxfp4, ++ .vec_dot = ggml_vec_dot_mxfp4_q8_0, ++ .vec_dot_type = GGML_TYPE_Q8_0, ++ .nrows = 1, + } + }; + +diff --git a/ggml.h b/ggml.h +index 13502a362..574b528f1 100644 +--- a/ggml.h ++++ b/ggml.h +@@ -377,6 +377,7 @@ extern "C" { + GGML_TYPE_F64 = 28, + GGML_TYPE_IQ1_M = 29, + GGML_TYPE_BF16 = 30, ++ GGML_TYPE_MXFP4 = 39, // OCP microscaling FP4 (E2M1 + ue8m0), id aligned with upstream + GGML_TYPE_COUNT, + }; + +diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py +index fb20cfabb..3c33158be 100644 +--- a/gguf-py/gguf/constants.py ++++ b/gguf-py/gguf/constants.py +@@ -903,6 +903,7 @@ class GGMLQuantizationType(IntEnum): + F64 = 28 + IQ1_M = 29 + BF16 = 30 ++ MXFP4 = 39 + + + # TODO: add GGMLFileType from ggml_ftype in ggml.h +@@ -1017,6 +1018,7 @@ GGML_QUANT_SIZES: dict[GGMLQuantizationType, tuple[int, int]] = { + GGMLQuantizationType.F64: (1, 8), + GGMLQuantizationType.IQ1_M: (256, QK_K // 8 + QK_K // 16 + QK_K // 32), + GGMLQuantizationType.BF16: (1, 2), ++ GGMLQuantizationType.MXFP4: (32, 1 + 16), + } + + diff --git a/kt-kernel/third_party_patches/llama.cpp/README.md b/kt-kernel/third_party_patches/llama.cpp/README.md new file mode 100644 index 000000000..60fe25e94 --- /dev/null +++ b/kt-kernel/third_party_patches/llama.cpp/README.md @@ -0,0 +1,57 @@ +# llama.cpp patches + +`third_party/llama.cpp` is pinned to the upstream commit +`a94e6ff8774b7c9f950d9545baf0ce35e8d1ed2f`, which is `refs/tags/b3173` in +. The pin deliberately points at a +commit that exists upstream, so `git clone --recursive` of this repository +works for everyone. Local-only fork commits must **not** be pinned here. + +b3173 predates MXFP4, so the deltas kt-kernel needs on top of it live in this +directory as a patch series and are applied at configure time by +`kt-kernel/CMakeLists.txt` (right before `add_subdirectory(... llama.cpp)`). + +## Patches + +### `0001-ggml-mxfp4-type.patch` + +Adds the OCP microscaling FP4 type (`GGML_TYPE_MXFP4`, id 39 — the same id +upstream llama.cpp later assigned) to ggml: + +* `ggml-common.h` — `block_mxfp4` (one `ue8m0` byte + 16 nibble-packed `E2M1` + codes per 32 weights), matching the upstream half-block interleave. +* `ggml-quants.{c,h}` — `dequantize_row_mxfp4` and `ggml_vec_dot_mxfp4_q8_0` + (NEON dot-product path that avoids SVE/i8mm so it runs on Kunpeng 920, plus a + scalar fallback for every other architecture). +* `ggml.c` / `ggml.h` — the enum value and the `type_traits` entry wiring the + dequantizer and the `vec_dot` (`vec_dot_type = GGML_TYPE_Q8_0`). +* `gguf-py/gguf/constants.py` — `GGMLQuantizationType.MXFP4` and its + `GGML_QUANT_SIZES` entry. + +This is required because the DeepSeek-V4 GGUF stores its CPU-offloaded expert +tensors as MXFP4; without it ggml rejects the tensor type and kt-kernel cannot +load the experts. + +## Applying by hand + +The build applies these automatically. To do it manually (e.g. when debugging): + +```bash +cd third_party/llama.cpp +for p in ../../kt-kernel/third_party_patches/llama.cpp/*.patch; do + git apply "$p" +done +``` + +To revert: + +```bash +cd third_party/llama.cpp +for p in $(ls -r ../../kt-kernel/third_party_patches/llama.cpp/*.patch); do + git apply --reverse "$p" +done +``` + +The CMake step is idempotent and needs no marker file: for each patch it first +tries `git apply --check`; if that fails it tries `git apply --reverse --check`, +and a success there means the patch is already applied, so it is skipped. Only +when both checks fail does the configure step abort. diff --git a/third_party/llama.cpp b/third_party/llama.cpp index ca9486437..a94e6ff87 160000 --- a/third_party/llama.cpp +++ b/third_party/llama.cpp @@ -1 +1 @@ -Subproject commit ca948643712e7bb11f543b404d0ce15fd53c0d2f +Subproject commit a94e6ff8774b7c9f950d9545baf0ce35e8d1ed2f From 2fd5ff8b22813f0fc121f16cd4837462dc2629b3 Mon Sep 17 00:00:00 2001 From: Pan-Boyi Date: Fri, 14 Aug 2026 15:05:35 +0800 Subject: [PATCH 05/13] [build](kt-kernel): make dependencies self-consistent for CANN 9 / aarch64 - requirements.txt pinned torch==2.9.1 while pyproject already allowed >=2.9.1,<2.11; torch_npu for CANN 9.0.0 needs torch 2.10, so the hard pin made the two files contradict each other. Widen the requirements pin to match. - triton has no aarch64 wheels and the Ascend path never uses it. Guard it with a platform_machine marker in both files so an aarch64 install does not fail resolving it. - CMakeLists.txt does find_library(NUMA_LIBRARY numa) with FATAL_ERROR on failure and moe.hpp includes directly, but install.sh only asked for hwloc and relied on its dependency chain to drag libnuma in. Install libnuma-dev / numactl-devel explicitly. Signed-off-by: Pan-Boyi --- kt-kernel/install.sh | 18 +++++++++--------- kt-kernel/pyproject.toml | 3 ++- kt-kernel/requirements.txt | 4 ++-- 3 files changed, 13 insertions(+), 12 deletions(-) diff --git a/kt-kernel/install.sh b/kt-kernel/install.sh index 4c79787ab..214348b69 100755 --- a/kt-kernel/install.sh +++ b/kt-kernel/install.sh @@ -125,24 +125,24 @@ install_dependencies() { # Install dependencies based on OS case "$OS" in debian|ubuntu|linuxmint|pop) - echo "Detected Debian-based system. Installing libhwloc-dev and pkg-config..." + echo "Detected Debian-based system. Installing libhwloc-dev, libnuma-dev and pkg-config..." $SUDO apt update - $SUDO apt install -y libhwloc-dev pkg-config + $SUDO apt install -y libhwloc-dev libnuma-dev pkg-config ;; fedora|rhel|centos|rocky|almalinux) - echo "Detected Red Hat-based system. Installing hwloc-devel and pkgconfig..." - $SUDO dnf install -y hwloc-devel pkgconfig || $SUDO yum install -y hwloc-devel pkgconfig + echo "Detected Red Hat-based system. Installing hwloc-devel, numactl-devel and pkgconfig..." + $SUDO dnf install -y hwloc-devel numactl-devel pkgconfig || $SUDO yum install -y hwloc-devel numactl-devel pkgconfig ;; arch|manjaro) - echo "Detected Arch-based system. Installing hwloc and pkgconf..." - $SUDO pacman -S --noconfirm hwloc pkgconf + echo "Detected Arch-based system. Installing hwloc, numactl and pkgconf..." + $SUDO pacman -S --noconfirm hwloc numactl pkgconf ;; opensuse*|sles) - echo "Detected openSUSE-based system. Installing hwloc-devel and pkg-config..." - $SUDO zypper install -y hwloc-devel pkg-config + echo "Detected openSUSE-based system. Installing hwloc-devel, libnuma-devel and pkg-config..." + $SUDO zypper install -y hwloc-devel libnuma-devel pkg-config ;; *) - echo "Warning: Unsupported OS '$OS'. Please manually install libhwloc-dev and pkg-config." + echo "Warning: Unsupported OS '$OS'. Please manually install libhwloc-dev, libnuma-dev and pkg-config." ;; esac } diff --git a/kt-kernel/pyproject.toml b/kt-kernel/pyproject.toml index 1e13cb7b0..733c3b060 100644 --- a/kt-kernel/pyproject.toml +++ b/kt-kernel/pyproject.toml @@ -26,7 +26,8 @@ dependencies = [ "torch>=2.9.1,<2.11", "safetensors>=0.4.0", "numpy>=1.24.0", - "triton>=2.0.0", + # triton has no aarch64 wheels; the Ascend build does not use it. + "triton>=2.0.0; platform_machine == 'x86_64' or platform_machine == 'AMD64'", "gguf>=0.17.0", # CLI dependencies "typer>=0.9.0", diff --git a/kt-kernel/requirements.txt b/kt-kernel/requirements.txt index d2b31e2e1..6070f43f5 100644 --- a/kt-kernel/requirements.txt +++ b/kt-kernel/requirements.txt @@ -3,8 +3,8 @@ # You can skip this file if you already have these packages installed # Core dependencies (minimum versions) -torch==2.9.1 +torch>=2.9.1,<2.11 safetensors>=0.4.0 numpy>=1.24.0 -triton>=2.0.0 +triton>=2.0.0; platform_machine == "x86_64" gguf>=0.17.0 From 77df4caf39c4d4b8839eabc195806b43f860f211 Mon Sep 17 00:00:00 2001 From: Pan-Boyi Date: Fri, 14 Aug 2026 15:05:50 +0800 Subject: [PATCH 06/13] [fix](kt-kernel): correct Ascend ACL callback worker lifecycle Three defects in the ACL report-subscriber worker: - aclrtUnSubscribeReport() ran in shutdown_callback_worker() *after* g_worker.join(), when the thread is already dead and g_worker_thread_id no longer names a live thread, so it was a no-op. ACL also requires the subscribing thread to unsubscribe. Move it to the end of worker_main(). - start_worker_locked() called aclrtSubscribeReport() immediately after spawning the thread, so subscriptions could race ahead of the worker reaching its aclrtProcessReport() loop and the first callbacks could be dropped. Add a g_worker_ready release/acquire barrier. - Callbacks enqueued between the last loop iteration and the stop request were silently discarded, hanging their waiters. Drain once more after the loop. Also add callback_worker_running() plus its is_ascend_callback_worker_running() binding, and use it from Python: warn if the worker did not come up, and treat "worker not running" as a reason to take the synchronous submit/sync path instead of letting submit_with_cuda_stream hang forever. The probe is memoized because it sits on the per-layer forward path, and is looked up with getattr() so older extension builds still work. While here, submit_forward() no longer degrades the CUDA path: the synchronous submit plus unconditional _wait_device() were only needed for the ACL host callback ordering problem on NPU, but they cost x86/CUDA an extra torch.cuda.synchronize() per layer and all CPU/GPU overlap. Gate them on device.type == "npu" and leave CUDA on the upstream submit_with_cuda_stream path. Drop the KT_DEBUG_MOE_OUT block, which printed real activations. Signed-off-by: Pan-Boyi --- .../cpu_backend/ascend_callback_worker.cpp | 32 ++++- .../cpu_backend/ascend_callback_worker.h | 5 + kt-kernel/ext_bindings.cpp | 2 + kt-kernel/python/experts_base.py | 119 +++++++++++------- 4 files changed, 107 insertions(+), 51 deletions(-) diff --git a/kt-kernel/cpu_backend/ascend_callback_worker.cpp b/kt-kernel/cpu_backend/ascend_callback_worker.cpp index 6f69d92cf..bbcbf445a 100644 --- a/kt-kernel/cpu_backend/ascend_callback_worker.cpp +++ b/kt-kernel/cpu_backend/ascend_callback_worker.cpp @@ -21,6 +21,10 @@ aclrtContext g_context = nullptr; std::thread g_worker; std::atomic g_stop{false}; std::atomic 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 g_worker_ready{false}; uint64_t g_worker_thread_id = 0; void worker_main(aclrtContext ctx) { @@ -32,10 +36,22 @@ void worker_main(aclrtContext ctx) { static_cast(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) { @@ -44,10 +60,16 @@ void start_worker_locked(aclrtContext ctx) { } 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(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) { @@ -112,16 +134,20 @@ void shutdown_callback_worker() { } 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(); } - for (aclrtStream stream : g_subscribed_streams) { - (void)aclrtUnSubscribeReport(g_worker_thread_id, stream); - } 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 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 diff --git a/kt-kernel/cpu_backend/ascend_callback_worker.h b/kt-kernel/cpu_backend/ascend_callback_worker.h index 8719a4860..01894399f 100644 --- a/kt-kernel/cpu_backend/ascend_callback_worker.h +++ b/kt-kernel/cpu_backend/ascend_callback_worker.h @@ -24,6 +24,11 @@ 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 diff --git a/kt-kernel/ext_bindings.cpp b/kt-kernel/ext_bindings.cpp index adaa5d866..4cc1d3bc2 100644 --- a/kt-kernel/ext_bindings.cpp +++ b/kt-kernel/ext_bindings.cpp @@ -531,6 +531,8 @@ PYBIND11_MODULE(kt_kernel_ext, m) { "Subscribe an aclrtStream with the global callback worker."); m.def("shutdown_ascend_callback_worker", &kt::ascend::shutdown_callback_worker, "Stop the ACL callback worker thread."); + m.def("is_ascend_callback_worker_running", &kt::ascend::callback_worker_running, + "True iff the ACL callback worker is started and dispatching reports."); #endif py::class_(m, "WorkerPool").def(py::init()); py::class_(m, "WorkerPoolConfig") diff --git a/kt-kernel/python/experts_base.py b/kt-kernel/python/experts_base.py index 71a12f790..647ab2948 100644 --- a/kt-kernel/python/experts_base.py +++ b/kt-kernel/python/experts_base.py @@ -12,11 +12,14 @@ import torch from typing import Dict, List, Optional, Tuple from abc import ABC, abstractmethod -import os import ctypes +import logging +import os from kt_kernel import kt_kernel_ext +logger = logging.getLogger(__name__) + # ----------------------------------------------------------------------------- # NPU stream-callback bypass. # @@ -39,10 +42,43 @@ # ----------------------------------------------------------------------------- +# Memoized result of the C++ worker probe. None = not probed yet. +_ascend_worker_degraded: Optional[bool] = None + + +def _ascend_callback_worker_running() -> Optional[bool]: + """Probe the C++ callback worker; None when the build predates the probe.""" + probe = getattr(kt_kernel_ext, "is_ascend_callback_worker_running", None) + if probe is None: + return None + try: + return bool(probe()) + except Exception: + return None + + +def _ascend_callback_worker_degraded() -> bool: + """True iff the ACL callback worker is known NOT to be dispatching reports. + + Probed once and memoized: the worker's state only changes at start-up and at + interpreter shutdown, and this sits on the per-layer forward path. Builds + without the probe are treated as healthy (backwards compatible). + """ + global _ascend_worker_degraded + if _ascend_worker_degraded is None: + _ascend_worker_degraded = _ascend_callback_worker_running() is False + return _ascend_worker_degraded + + def _should_bypass_stream_callback(device: torch.device) -> bool: """Return True iff we must use the synchronous submit/sync path.""" if os.environ.get("KT_FORCE_SYNC_SUBMIT", "") == "1": return True + if device.type == "npu" and not _uses_external_npu_report_subscriber(): + # Without a dispatcher thread, aclrtLaunchCallback tasks never fire and + # submit_with_cuda_stream would hang forever. Degrade to sync instead. + if _ascend_callback_worker_degraded(): + return True return False @@ -60,6 +96,14 @@ def _ensure_ascend_callback_worker() -> None: return kt_kernel_ext.init_ascend_callback_worker() _ensure_ascend_callback_worker._done = True # type: ignore[attr-defined] + + global _ascend_worker_degraded + _ascend_worker_degraded = _ascend_callback_worker_running() is False + if _ascend_worker_degraded: + logger.warning( + "kt-kernel ACL callback worker failed to start (no ACL context?); " + "falling back to the synchronous CPU-MoE submit/sync path." + ) if hasattr(kt_kernel_ext, "shutdown_ascend_callback_worker"): import atexit @@ -714,11 +758,19 @@ def submit_forward( bsz_slot_tensor = bsz_tensor_cpu[current_slot] bypass = _should_bypass_stream_callback(hidden_states.device) - # Both paths now submit the CPU-MoE synchronously on this host thread, which - # reads input_tensor_cpu immediately -> the input D2H (queued async on the - # stream by copy_inputs_to_cpu_buffers) MUST be finished first, else the CPU - # MoE reads a half-copied input. Wait unconditionally (bypass already did). - _wait_device(hidden_states.device) + # NPU only: submit the CPU-MoE synchronously on this host thread. + # ``submit_with_cuda_stream`` enqueues it via an ACL host callback whose + # firing is not host-observable (NO_BLOCK) -> a later host-side drain can + # race ahead of it (empty queue -> stale output_cpu read by the H2D; + # nondeterministic on heavy prefill). The synchronous submit still runs on + # the WorkerPool and overlaps the device experts queued after it. + # CUDA keeps the upstream async submit_with_cuda_stream path unchanged. + sync_submit = bypass or hidden_states.device.type == "npu" + if sync_submit: + # The synchronous submit reads input_tensor_cpu immediately -> the input + # D2H queued async on the stream by _prepare_forward_cpu_buffers MUST be + # finished first, else the CPU MoE reads a half-copied input. + _wait_device(hidden_states.device) incremental = BaseMoEWrapper._layer_has_pending_deferred.get(self.layer_idx - 1, False) immediate_task = self.moe.forward_task( @@ -730,29 +782,25 @@ def submit_forward( output_cpu[current_slot].data_ptr(), incremental, ) - if bypass: - self.cpu_infer.submit(immediate_task) - else: - # Correct + fast async path. ``submit_with_cuda_stream`` enqueues the - # CPU-MoE via an ACL host callback whose firing is not host-observable - # (NO_BLOCK) -> a later host-side drain can race ahead of it (empty - # queue -> stale output_cpu read by the H2D; nondeterministic on heavy - # prefill). Enqueue synchronously on this host thread instead (work is - # guaranteed submitted; it still runs on the WorkerPool and overlaps the - # GPU experts queued after). Keep subscribe_ascend_stream — that stream - # registration is what keeps decode's host-callback dispatch fast; the - # async submit itself is not required for it. + if sync_submit: + # Keep subscribe_ascend_stream — that stream registration is what keeps + # decode's host-callback dispatch (sync_forward) fast; the async submit + # itself is not required for it. if ( - hidden_states.device.type == "npu" + not bypass + and hidden_states.device.type == "npu" and not _uses_external_npu_report_subscriber() and hasattr(kt_kernel_ext, "subscribe_ascend_stream") ): kt_kernel_ext.subscribe_ascend_stream(int(cuda_stream)) self.cpu_infer.submit(immediate_task) + else: + self.cpu_infer.submit_with_cuda_stream(cuda_stream, immediate_task) BaseMoEWrapper._layer_has_pending_deferred[self.layer_idx] = False if deferred_ids is not None: - _wait_device(hidden_states.device) + if sync_submit: + _wait_device(hidden_states.device) deferred_task = self.moe.forward_task( bsz_slot_tensor.data_ptr(), deferred_experts_ids_cpu[current_slot].size(-1), @@ -762,16 +810,10 @@ def submit_forward( output_cpu[next_slot].data_ptr(), False, ) - if bypass: + if sync_submit: self.cpu_infer.submit(deferred_task) else: - if ( - hidden_states.device.type == "npu" - and not _uses_external_npu_report_subscriber() - and hasattr(kt_kernel_ext, "subscribe_ascend_stream") - ): - kt_kernel_ext.subscribe_ascend_stream(int(cuda_stream)) - self.cpu_infer.submit(deferred_task) + self.cpu_infer.submit_with_cuda_stream(cuda_stream, deferred_task) BaseMoEWrapper._layer_has_pending_deferred[self.layer_idx] = True def copy_forward_output_to_device(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -834,26 +876,6 @@ def sync_forward(self, hidden_states: torch.Tensor, cuda_stream) -> torch.Tensor else: self.cpu_infer.sync_with_cuda_stream(cuda_stream, allow_pending) - if os.environ.get("KT_DEBUG_MOE_OUT", "") == "1": - oc = output_cpu[current_slot] - ic = _input_tensor_cpu[current_slot] - wc = _weights_cpu[current_slot] - ec = _immediate_experts_ids_cpu[current_slot] - finite_oc = torch.isfinite(oc).all().item() - finite_ic = torch.isfinite(ic).all().item() - finite_wc = torch.isfinite(wc).all().item() - nan_count = int(torch.isnan(oc).sum().item()) - inf_count = int(torch.isinf(oc).sum().item()) - print( - f"[KT_DEBUG] layer={self.layer_idx} slot={current_slot} bypass={bypass} " - f"output_cpu: finite={finite_oc} nan={nan_count}/{oc.numel()} inf={inf_count} " - f"first16={oc.float().flatten()[:16].tolist()} | " - f"input_cpu.finite={finite_ic} first8={ic.float().flatten()[:8].tolist()} | " - f"weights_cpu.finite={finite_wc} first={wc.flatten()[:8].tolist()} | " - f"expert_ids={ec.flatten()[:12].tolist()}", - flush=True, - ) - return self.copy_forward_output_to_device(hidden_states) def forward( @@ -915,3 +937,4 @@ def clear_buffer_cache(): KExpertsCPUBuffer.capture_buffers.clear() KExpertsCPUBuffer.temp_bs = 0 KExpertsCPUBuffer.temp_buffer = tuple() + From b9f41b74360d13e025b288b0834842441bb0fe8c Mon Sep 17 00:00:00 2001 From: Pan-Boyi Date: Fri, 14 Aug 2026 15:05:59 +0800 Subject: [PATCH 07/13] [chore](kt-kernel): drop KT_DUMMY_CPU_WEIGHTS dev scaffolding The dummy-weight path fabricated zero expert buffers to skip the multi-GB GGUF read during graph-capture iteration. It has served its purpose and only adds a way to silently run a model that produces meaningless output. GGMLQuantizationType.MXFP4 and the module-level GGML_QUANT_SIZES table stay: the real loader needs both. Signed-off-by: Pan-Boyi --- kt-kernel/python/utils/llamafile.py | 28 ++++++--------------------- kt-kernel/python/utils/loader.py | 30 ----------------------------- 2 files changed, 6 insertions(+), 52 deletions(-) diff --git a/kt-kernel/python/utils/llamafile.py b/kt-kernel/python/utils/llamafile.py index 6c9b6d26d..bf79ae7e4 100644 --- a/kt-kernel/python/utils/llamafile.py +++ b/kt-kernel/python/utils/llamafile.py @@ -188,28 +188,12 @@ def load_weights(self, physical_to_logical_map_cpu: Optional[torch.Tensor] = Non base_key = f"blk.{self.layer_idx}" - # KT_DUMMY_CPU_WEIGHTS: skip the multi-GB GGUF disk read and fabricate - # same-layout zero buffers so the C++ MOE / load_weights_task path runs - # identically (buffer sizing, kernel selection, capture) but loads much - # faster. Debug-only: weights are meaningless — never use for accuracy. - if os.environ.get("KT_DUMMY_CPU_WEIGHTS"): - if self.layer_idx == 0: - print( - "[LlamafileMoEWrapper] KT_DUMMY_CPU_WEIGHTS=1: fabricating zero " - "CPU MoE weights (skips GGUF read). Outputs are MEANINGLESS — " - "for capture iteration only, not accuracy validation.", - flush=True, - ) - _load_tensor = self.gguf_loader.get_dummy_tensor_and_ggml_type - else: - _load_tensor = self.gguf_loader.get_undequanted_tensor_and_ggml_type - - # Load quantized tensors from GGUF (or fabricate dummies) - gate_data, gate_type = _load_tensor(f"{base_key}.ffn_gate_exps.weight") - - up_data, up_type = _load_tensor(f"{base_key}.ffn_up_exps.weight") - - down_data, down_type = _load_tensor(f"{base_key}.ffn_down_exps.weight") + # Load quantized tensors from GGUF + gate_data, gate_type = self.gguf_loader.get_undequanted_tensor_and_ggml_type(f"{base_key}.ffn_gate_exps.weight") + + up_data, up_type = self.gguf_loader.get_undequanted_tensor_and_ggml_type(f"{base_key}.ffn_up_exps.weight") + + down_data, down_type = self.gguf_loader.get_undequanted_tensor_and_ggml_type(f"{base_key}.ffn_down_exps.weight") # Keep tensors alive self.weights_to_keep = (gate_data, up_data, down_data) diff --git a/kt-kernel/python/utils/loader.py b/kt-kernel/python/utils/loader.py index aa6e71fb4..6d250e83b 100644 --- a/kt-kernel/python/utils/loader.py +++ b/kt-kernel/python/utils/loader.py @@ -52,8 +52,6 @@ class GGMLQuantizationType(IntEnum): # (block_size, type_size) per GGML quant type — bytes per `block_size` elements. -# Module-level so both the real loader and the dummy-weight fabrication path -# (KT_DUMMY_CPU_WEIGHTS) share one source of truth. GGML_QUANT_SIZES = { GGMLQuantizationType.F32: (1, 4), GGMLQuantizationType.F16: (1, 2), @@ -1079,34 +1077,6 @@ def get_undequanted_tensor_and_ggml_type(self, name: str): return data, ggml_type - def get_dummy_tensor_and_ggml_type(self, name: str): - """Fabricate a zero buffer matching a tensor's on-disk byte layout. - - Used only by the ``KT_DUMMY_CPU_WEIGHTS`` fast-iteration path: it reads - tensor metadata (shape / quant type / element count) from the cheap - ``tensor_info`` map and allocates a same-sized ``uint8`` buffer WITHOUT - touching the mmap'd weight bytes. This skips the multi-GB GGUF disk read - so graph capture can be iterated quickly; the returned values are - meaningless and must NOT be used for accuracy validation. - - Returns ``(data, ggml_type)`` with the same contract as - ``get_undequanted_tensor_and_ggml_type``. - """ - name = translate_name_to_gguf(name) - - if name not in self.tensor_info: - raise KeyError(f"Tensor '{name}' not found in GGUF files") - - info = self.tensor_info[name] - n_elements = info["n_elements"] - ggml_type = info["dtype"] - - block_size, type_size = GGML_QUANT_SIZES[ggml_type] - n_bytes = n_elements * type_size // block_size - - data = torch.zeros(n_bytes, dtype=torch.uint8) - return data, ggml_type - class GPTQSafeTensorLoader(FP8SafeTensorLoader): """Loader for symmetric GPTQ-Int4 expert weights (qweight + scales, no qzeros). From 9597c8a3dfe48427817fdec15880ff82883a2ec6 Mon Sep 17 00:00:00 2001 From: Pan-Boyi Date: Fri, 14 Aug 2026 18:40:32 +0800 Subject: [PATCH 08/13] [feat](kt-kernel): add the fused AscendC MXFP4->W8A8 dequant kernel The DeepSeek-V4 streaming-prefill "depool" path keeps the MoE experts on the host as native MXFP4 (~137GB) instead of a resident W8A8 NZ pool (~277GB), and converts one layer at a time on the device. That conversion is this operator: a single AscendC pass that reads the MXFP4 codes + e8m0 block scales once and emits the int8 weight planes plus a per-output-channel bf16 scale, i.e. exactly the tensors the NPU grouped-matmul MoE path consumes. mxfp4_fused_kernel.cpp exposes two entry points: launch_mxfp4_fused for consecutive-nibble codes (native safetensors) and launch_mxfp4_fused_blk, which de-interleaves GGUF block_mxfp4 (scale|codes per 17-byte block) in UB via Gather so no host-side strided de-interleave is needed. mxfp4_fused_op.py builds the .so with bisheng on first use, caches it, and wraps both entry points in the chunked ND->NZ post-step. Relative to the source drop: - the .so is no longer forced into the source directory, which is read-only under some container mounts; it falls back to a user cache dir and can be redirected with KT_MXFP4_SO_DIR, - a failing bisheng invocation now reports the command, ASCEND_TOOLKIT_HOME and the compiler stderr instead of a bare CalledProcessError, - the ASCEND_TOOLKIT_HOME fallback is the version-independent install symlink rather than a pinned CANN release. The kernel is standalone: it is loaded by path via KT_MXFP4_OP_DIR and is not imported by the kt-kernel build or by any kt_kernel module. Signed-off-by: Pan-Boyi --- .../ascendc_mxfp4/mxfp4_fused_kernel.cpp | 348 ++++++++++++++++++ .../tools/ascendc_mxfp4/mxfp4_fused_op.py | 332 +++++++++++++++++ 2 files changed, 680 insertions(+) create mode 100644 kt-kernel/tools/ascendc_mxfp4/mxfp4_fused_kernel.cpp create mode 100644 kt-kernel/tools/ascendc_mxfp4/mxfp4_fused_op.py diff --git a/kt-kernel/tools/ascendc_mxfp4/mxfp4_fused_kernel.cpp b/kt-kernel/tools/ascendc_mxfp4/mxfp4_fused_kernel.cpp new file mode 100644 index 000000000..4743461c3 --- /dev/null +++ b/kt-kernel/tools/ascendc_mxfp4/mxfp4_fused_kernel.cpp @@ -0,0 +1,348 @@ +// AscendC FUSED MXFP4 -> W8A8: int8 weight + per-output-channel oscale in ONE pass (reads MXFP4 +// once). Block-partitioned: each core owns a contiguous, ACC-aligned row range; per row it does +// decode/scale/reduce once, emits int8 (two contiguous planes [lo|hi]) AND accumulates oscale into +// a UB block; each full block is flushed as ONE large contiguous DataCopy (the idiom that makes +// the small per-channel scale store survive alongside loads + int8 stores). +#include "kernel_operator.h" +using namespace AscendC; + +constexpr int32_t HALF_MAX = 2048; +constexpr int32_t IN_MAX = 4096; +constexpr int32_t NB_MAX = 128; +constexpr int32_t ACC = 512; // oscale flush block (floats), 8-aligned + +extern "C" __global__ __aicore__ void mxfp4_fused( + GM_ADDR codes, GM_ADDR scaleg, GM_ADDR outg, GM_ADDR oscaleg, + GM_ADDR lutLoG, GM_ADDR lutHiG, GM_ADDR lutE8G, GM_ADDR scOffG, + uint32_t R, uint32_t HALF, uint32_t NB, uint32_t IN) +{ + const int32_t blkid = GetBlockIdx(); + const int32_t nblk = GetBlockNum(); + + GlobalTensor gCodes, gScale, gOut; + GlobalTensor gOscale, gLutLo, gLutHi, gLutE8; + GlobalTensor gScOff; + gCodes.SetGlobalBuffer((__gm__ uint8_t *)codes); + gScale.SetGlobalBuffer((__gm__ uint8_t *)scaleg); + gOut.SetGlobalBuffer((__gm__ uint8_t *)outg); + gOscale.SetGlobalBuffer((__gm__ float *)oscaleg); + gLutLo.SetGlobalBuffer((__gm__ float *)lutLoG); + gLutHi.SetGlobalBuffer((__gm__ float *)lutHiG); + gLutE8.SetGlobalBuffer((__gm__ float *)lutE8G); + gScOff.SetGlobalBuffer((__gm__ uint32_t *)scOffG); + + TPipe pipe; + TQue qCodes, qScale; + TQue qOut; + pipe.InitBuffer(qCodes, 1, HALF_MAX * sizeof(uint8_t)); + pipe.InitBuffer(qScale, 1, (NB_MAX + 32) * sizeof(uint8_t)); + pipe.InitBuffer(qOut, 1, IN_MAX * sizeof(uint8_t)); + TBuf tLutLo, tLutHi, tLutE8, tScOff; + TBuf tComb, tOff, tOffH, tScI, tScF, tScHalf, tAbs, tWork, tAcc; + pipe.InitBuffer(tLutLo, 256 * sizeof(float)); + pipe.InitBuffer(tLutHi, 256 * sizeof(float)); + pipe.InitBuffer(tLutE8, 256 * sizeof(float)); + pipe.InitBuffer(tScOff, HALF_MAX * sizeof(uint32_t)); + pipe.InitBuffer(tComb, 2 * HALF_MAX * sizeof(float)); + pipe.InitBuffer(tOff, HALF_MAX * sizeof(int32_t)); + pipe.InitBuffer(tOffH, HALF_MAX * sizeof(half)); + pipe.InitBuffer(tScI, NB_MAX * sizeof(int32_t)); + pipe.InitBuffer(tScF, NB_MAX * sizeof(float)); + pipe.InitBuffer(tScHalf, HALF_MAX * sizeof(float)); + pipe.InitBuffer(tAbs, HALF_MAX * sizeof(float)); + pipe.InitBuffer(tWork, HALF_MAX * sizeof(float)); + pipe.InitBuffer(tAcc, ACC * sizeof(float)); + + LocalTensor lutLo = tLutLo.Get(); + LocalTensor lutHi = tLutHi.Get(); + LocalTensor lutE8 = tLutE8.Get(); + LocalTensor scOff = tScOff.Get(); + LocalTensor comb = tComb.Get(); + LocalTensor off = tOff.Get(); + LocalTensor offH = tOffH.Get(); + LocalTensor scI = tScI.Get(); + LocalTensor scF = tScF.Get(); + LocalTensor scHalf = tScHalf.Get(); + LocalTensor absb = tAbs.Get(); + LocalTensor work = tWork.Get(); + LocalTensor acc = tAcc.Get(); + + DataCopy(lutLo, gLutLo, 256); + DataCopy(lutHi, gLutHi, 256); + DataCopy(lutE8, gLutE8, 256); + DataCopy(scOff, gScOff, HALF); + PipeBarrier(); + + const uint32_t chunk = ((R + nblk - 1) / nblk + (ACC - 1)) / ACC * ACC; + const uint32_t rStart = (uint32_t)blkid * chunk; + uint32_t rEnd = rStart + chunk; + if (rEnd > R) rEnd = R; + const uint32_t scLoad = (NB + 31) / 32 * 32; + + for (uint32_t base = rStart; base < rEnd; base += ACC) { + uint32_t bend = base + ACC; + if (bend > rEnd) bend = rEnd; + for (uint32_t r = base; r < bend; r++) { + LocalTensor vlo = comb; + LocalTensor vhi = comb[HALF]; + + LocalTensor cu = qCodes.AllocTensor(); + DataCopy(cu, gCodes[(uint64_t)r * HALF], HALF); + qCodes.EnQue(cu); + LocalTensor cuU = qCodes.DeQue(); + LocalTensor su = qScale.AllocTensor(); + DataCopy(su, gScale[(uint64_t)r * NB], scLoad); + qScale.EnQue(su); + LocalTensor suU = qScale.DeQue(); + + Cast(offH, cuU, RoundMode::CAST_NONE, HALF); + Muls(offH, offH, (half)4.0, HALF); + Cast(off, offH, RoundMode::CAST_RINT, HALF); + LocalTensor offU = off.ReinterpretCast(); + Gather(vlo, lutLo, offU, (uint32_t)0, HALF); + Gather(vhi, lutHi, offU, (uint32_t)0, HALF); + qCodes.FreeTensor(cuU); + + Cast(offH, suU, RoundMode::CAST_NONE, NB); + Muls(offH, offH, (half)4.0, NB); + Cast(scI, offH, RoundMode::CAST_RINT, NB); + Gather(scF, lutE8, scI.ReinterpretCast(), (uint32_t)0, NB); + qScale.FreeTensor(suU); + PipeBarrier(); + Gather(scHalf, scF, scOff, (uint32_t)0, HALF); + Mul(vlo, vlo, scHalf, HALF); + Mul(vhi, vhi, scHalf, HALF); + PipeBarrier(); + + Abs(absb, vlo, HALF); + Abs(work, vhi, HALF); + Max(scHalf, absb, work, HALF); + PipeBarrier(); + LocalTensor fa = scHalf, fb = absb; + for (uint32_t h = HALF >> 1; h >= 8; h >>= 1) { + Max(fb, fa, fa[h], h); + PipeBarrier(); + LocalTensor tmp = fa; fa = fb; fb = tmp; + } + PipeBarrier(); + float amax = fa.GetValue(0); + for (int i = 1; i < 8; i++) { float v = fa.GetValue(i); if (v > amax) amax = v; } + if (amax < 1e-8f) amax = 1e-8f; + acc.SetValue(r - base, amax / 127.0f); // accumulate oscale (flushed per block) + float inv = 127.0f / amax; + + PipeBarrier(); // scalar inv -> vector + Muls(vlo, vlo, inv, HALF); + Muls(vhi, vhi, inv, HALF); + PipeBarrier(); + Mins(vlo, vlo, 127.0f, HALF); Maxs(vlo, vlo, -127.0f, HALF); + Mins(vhi, vhi, 127.0f, HALF); Maxs(vhi, vhi, -127.0f, HALF); + PipeBarrier(); + + LocalTensor outrow = qOut.AllocTensor(); + LocalTensor outI = outrow.ReinterpretCast(); + Cast(offH, vlo, RoundMode::CAST_NONE, HALF); + PipeBarrier(); + Cast(outI, offH, RoundMode::CAST_RINT, HALF); + PipeBarrier(); + Cast(offH, vhi, RoundMode::CAST_NONE, HALF); + PipeBarrier(); + Cast(outI[HALF], offH, RoundMode::CAST_RINT, HALF); + PipeBarrier(); + qOut.EnQue(outrow); + LocalTensor outU = qOut.DeQue(); + DataCopy(gOut[(uint64_t)r * IN], outU, IN); + qOut.FreeTensor(outU); + } + // flush the oscale block as one large contiguous DataCopy (8-aligned base) + PipeBarrier(); + DataCopy(gOscale[base], acc, ACC); + PipeBarrier(); + } +} + +extern "C" void launch_mxfp4_fused(void *stream, uint32_t blockdim, + uint8_t *codes, uint8_t *scale, uint8_t *out, uint8_t *oscale, + uint8_t *lutLo, uint8_t *lutHi, uint8_t *lutE8, uint8_t *scOff, + uint32_t R, uint32_t HALF, uint32_t NB, uint32_t IN) +{ + mxfp4_fused<<>>( + (GM_ADDR)codes, (GM_ADDR)scale, (GM_ADDR)out, (GM_ADDR)oscale, + (GM_ADDR)lutLo, (GM_ADDR)lutHi, (GM_ADDR)lutE8, (GM_ADDR)scOff, + R, HALF, NB, IN); +} + +// ---- block-input variant: reads raw GGUF block_mxfp4 ([nb*17] per row = nb x (1 e8m0 scale + 16 +// codes)) and de-interleaves IN UB via Gather (same gather-from-UB-by-offset op the base kernel +// already uses for scHalf), so the host does NO de-interleave (the slow 16-of-17 strided int8 copy). +// codeOff[j] = byte offset of code j in the half-cast block buffer = ((j/16)*17 + 1 + j%16)*2; +// scaleOff[b] = (b*17)*2. Everything after the input load is byte-identical to mxfp4_fused. +constexpr int32_t BLK_MAX = (HALF_MAX / 16) * 17; // 2176 + +extern "C" __global__ __aicore__ void mxfp4_fused_blk( + GM_ADDR blocks, GM_ADDR outg, GM_ADDR oscaleg, + GM_ADDR lutLoG, GM_ADDR lutHiG, GM_ADDR lutE8G, GM_ADDR scOffG, + GM_ADDR codeOffG, GM_ADDR scaleOffG, + uint32_t R, uint32_t HALF, uint32_t NB, uint32_t IN) +{ + const int32_t blkid = GetBlockIdx(); + const int32_t nblk = GetBlockNum(); + + GlobalTensor gBlocks, gOut; + GlobalTensor gOscale, gLutLo, gLutHi, gLutE8; + GlobalTensor gScOff, gCodeOff, gScaleOff; + gBlocks.SetGlobalBuffer((__gm__ uint8_t *)blocks); + gOut.SetGlobalBuffer((__gm__ uint8_t *)outg); + gOscale.SetGlobalBuffer((__gm__ float *)oscaleg); + gLutLo.SetGlobalBuffer((__gm__ float *)lutLoG); + gLutHi.SetGlobalBuffer((__gm__ float *)lutHiG); + gLutE8.SetGlobalBuffer((__gm__ float *)lutE8G); + gScOff.SetGlobalBuffer((__gm__ uint32_t *)scOffG); + gCodeOff.SetGlobalBuffer((__gm__ uint32_t *)codeOffG); + gScaleOff.SetGlobalBuffer((__gm__ uint32_t *)scaleOffG); + + TPipe pipe; + TQue qBlk; + TQue qOut; + pipe.InitBuffer(qBlk, 1, BLK_MAX * sizeof(uint8_t)); + pipe.InitBuffer(qOut, 1, IN_MAX * sizeof(uint8_t)); + TBuf tLutLo, tLutHi, tLutE8, tScOff, tCodeOff, tScaleOff, tBlkH; + TBuf tComb, tOff, tOffH, tScI, tScF, tScHalf, tAbs, tWork, tAcc; + pipe.InitBuffer(tLutLo, 256 * sizeof(float)); + pipe.InitBuffer(tLutHi, 256 * sizeof(float)); + pipe.InitBuffer(tLutE8, 256 * sizeof(float)); + pipe.InitBuffer(tScOff, HALF_MAX * sizeof(uint32_t)); + pipe.InitBuffer(tCodeOff, HALF_MAX * sizeof(uint32_t)); + pipe.InitBuffer(tScaleOff, NB_MAX * sizeof(uint32_t)); + pipe.InitBuffer(tBlkH, BLK_MAX * sizeof(half)); + pipe.InitBuffer(tComb, 2 * HALF_MAX * sizeof(float)); + pipe.InitBuffer(tOff, HALF_MAX * sizeof(int32_t)); + pipe.InitBuffer(tOffH, HALF_MAX * sizeof(half)); + pipe.InitBuffer(tScI, NB_MAX * sizeof(int32_t)); + pipe.InitBuffer(tScF, NB_MAX * sizeof(float)); + pipe.InitBuffer(tScHalf, HALF_MAX * sizeof(float)); + pipe.InitBuffer(tAbs, HALF_MAX * sizeof(float)); + pipe.InitBuffer(tWork, HALF_MAX * sizeof(float)); + pipe.InitBuffer(tAcc, ACC * sizeof(float)); + + LocalTensor lutLo = tLutLo.Get(); + LocalTensor lutHi = tLutHi.Get(); + LocalTensor lutE8 = tLutE8.Get(); + LocalTensor scOff = tScOff.Get(); + LocalTensor codeOff = tCodeOff.Get(); + LocalTensor scaleOff = tScaleOff.Get(); + LocalTensor blkH = tBlkH.Get(); + LocalTensor comb = tComb.Get(); + LocalTensor off = tOff.Get(); + LocalTensor offH = tOffH.Get(); + LocalTensor scI = tScI.Get(); + LocalTensor scF = tScF.Get(); + LocalTensor scHalf = tScHalf.Get(); + LocalTensor absb = tAbs.Get(); + LocalTensor work = tWork.Get(); + LocalTensor acc = tAcc.Get(); + + DataCopy(lutLo, gLutLo, 256); + DataCopy(lutHi, gLutHi, 256); + DataCopy(lutE8, gLutE8, 256); + DataCopy(scOff, gScOff, HALF); + DataCopy(codeOff, gCodeOff, HALF); + DataCopy(scaleOff, gScaleOff, NB); + PipeBarrier(); + + const uint32_t nb17 = (HALF / 16) * 17; + const uint32_t chunk = ((R + nblk - 1) / nblk + (ACC - 1)) / ACC * ACC; + const uint32_t rStart = (uint32_t)blkid * chunk; + uint32_t rEnd = rStart + chunk; + if (rEnd > R) rEnd = R; + + for (uint32_t base = rStart; base < rEnd; base += ACC) { + uint32_t bend = base + ACC; + if (bend > rEnd) bend = rEnd; + for (uint32_t r = base; r < bend; r++) { + LocalTensor vlo = comb; + LocalTensor vhi = comb[HALF]; + + LocalTensor bu = qBlk.AllocTensor(); + DataCopy(bu, gBlocks[(uint64_t)r * nb17], (nb17 + 31) / 32 * 32); + qBlk.EnQue(bu); + LocalTensor buU = qBlk.DeQue(); + Cast(blkH, buU, RoundMode::CAST_NONE, nb17); // blocks -> half + qBlk.FreeTensor(buU); + PipeBarrier(); + + Gather(offH, blkH, codeOff, (uint32_t)0, HALF); // de-interleave codes -> half + Muls(offH, offH, (half)4.0, HALF); + Cast(off, offH, RoundMode::CAST_RINT, HALF); + LocalTensor offU = off.ReinterpretCast(); + Gather(vlo, lutLo, offU, (uint32_t)0, HALF); + Gather(vhi, lutHi, offU, (uint32_t)0, HALF); + + Gather(offH, blkH, scaleOff, (uint32_t)0, NB); // de-interleave scale -> half + Muls(offH, offH, (half)4.0, NB); + Cast(scI, offH, RoundMode::CAST_RINT, NB); + Gather(scF, lutE8, scI.ReinterpretCast(), (uint32_t)0, NB); + PipeBarrier(); + Gather(scHalf, scF, scOff, (uint32_t)0, HALF); + Mul(vlo, vlo, scHalf, HALF); + Mul(vhi, vhi, scHalf, HALF); + PipeBarrier(); + + Abs(absb, vlo, HALF); + Abs(work, vhi, HALF); + Max(scHalf, absb, work, HALF); + PipeBarrier(); + LocalTensor fa = scHalf, fb = absb; + for (uint32_t h = HALF >> 1; h >= 8; h >>= 1) { + Max(fb, fa, fa[h], h); + PipeBarrier(); + LocalTensor tmp = fa; fa = fb; fb = tmp; + } + PipeBarrier(); + float amax = fa.GetValue(0); + for (int i = 1; i < 8; i++) { float v = fa.GetValue(i); if (v > amax) amax = v; } + if (amax < 1e-8f) amax = 1e-8f; + acc.SetValue(r - base, amax / 127.0f); + float inv = 127.0f / amax; + + PipeBarrier(); + Muls(vlo, vlo, inv, HALF); + Muls(vhi, vhi, inv, HALF); + PipeBarrier(); + Mins(vlo, vlo, 127.0f, HALF); Maxs(vlo, vlo, -127.0f, HALF); + Mins(vhi, vhi, 127.0f, HALF); Maxs(vhi, vhi, -127.0f, HALF); + PipeBarrier(); + + LocalTensor outrow = qOut.AllocTensor(); + LocalTensor outI = outrow.ReinterpretCast(); + Cast(offH, vlo, RoundMode::CAST_NONE, HALF); + PipeBarrier(); + Cast(outI, offH, RoundMode::CAST_RINT, HALF); + PipeBarrier(); + Cast(offH, vhi, RoundMode::CAST_NONE, HALF); + PipeBarrier(); + Cast(outI[HALF], offH, RoundMode::CAST_RINT, HALF); + PipeBarrier(); + qOut.EnQue(outrow); + LocalTensor outU = qOut.DeQue(); + DataCopy(gOut[(uint64_t)r * IN], outU, IN); + qOut.FreeTensor(outU); + } + PipeBarrier(); + DataCopy(gOscale[base], acc, ACC); + PipeBarrier(); + } +} + +extern "C" void launch_mxfp4_fused_blk(void *stream, uint32_t blockdim, + uint8_t *blocks, uint8_t *out, uint8_t *oscale, + uint8_t *lutLo, uint8_t *lutHi, uint8_t *lutE8, uint8_t *scOff, + uint8_t *codeOff, uint8_t *scaleOff, + uint32_t R, uint32_t HALF, uint32_t NB, uint32_t IN) +{ + mxfp4_fused_blk<<>>( + (GM_ADDR)blocks, (GM_ADDR)out, (GM_ADDR)oscale, + (GM_ADDR)lutLo, (GM_ADDR)lutHi, (GM_ADDR)lutE8, (GM_ADDR)scOff, + (GM_ADDR)codeOff, (GM_ADDR)scaleOff, + R, HALF, NB, IN); +} diff --git a/kt-kernel/tools/ascendc_mxfp4/mxfp4_fused_op.py b/kt-kernel/tools/ascendc_mxfp4/mxfp4_fused_op.py new file mode 100644 index 000000000..d1112988e --- /dev/null +++ b/kt-kernel/tools/ascendc_mxfp4/mxfp4_fused_op.py @@ -0,0 +1,332 @@ +"""Runtime wrapper for the fused AscendC MXFP4->W8A8 kernel, for kt_stream_prefill depool. + +Builds libmxfp4fused.so on first use (bisheng), loads it via ctypes, and exposes: + + mxfp4_layer_to_nz_slots(c13, s13, c2, s2, H, I, blockdim=40) + -> (w13_nz, s13b, w2_nz, s2b) # exactly the slot tensors npu_fused_experts consumes + # w*_nz: FRACTAL_NZ int8 [E,IN,OUT]; s*b: bf16 [E,OUT] + +Inputs are this layer's combined MXFP4 (device uint8): + c13/s13: w13 = cat(w1,w3) codes [E,2I,H/2] + e8m0 scale [E,2I,H/32] + c2/s2 : w2 codes [E,H,I/2] + e8m0 scale [E,H,I/32] + +Reads MXFP4 once; one kernel pass per projection. Validated end-to-end (cos 0.99999976 vs fp32 +golden through the NPU grouped-matmul MoE path). + +Point sglang's ``kt_stream_prefill`` at this directory with ``KT_MXFP4_OP_DIR``. +""" + +import ctypes +import os +import subprocess +import threading +from pathlib import Path + +import numpy as np +import torch + +_HERE = Path(__file__).resolve().parent +_SRC = _HERE / "mxfp4_fused_kernel.cpp" +_NZ = 29 +_ACC = 512 +_FP4 = np.array([0, 0.5, 1, 1.5, 2, 3, 4, 6, 0, -0.5, -1, -1.5, -2, -3, -4, -6], np.float32) + +# The fused kernel is launched via a raw ctypes <<>> call. With TASK_QUEUE_ENABLE=1 torch +# dispatches its ops through an async queue that is NOT ordered against that direct launch, so the +# post-step racing `out`/`osc` needs an explicit per-chunk host sync (correctness; but it stalls the +# host and serializes the convert with the rest of the forward -> slow prefill). With +# TASK_QUEUE_ENABLE=0 torch ops go straight to the stream, so kernel+post-step are FIFO-ordered and +# NO sync is needed (validated: deterministic + byte-equal, runs async -> fast). So only sync when +# the task queue is on. +_TQ_SYNC = os.environ.get("TASK_QUEUE_ENABLE", "1") != "0" + +_lib = None +_lock = threading.Lock() +_consts_cache = {} + + +def _cann_home(): + # ASCEND_TOOLKIT_HOME is exported by the CANN set_env.sh; the fallback is the + # version-independent install symlink so no CANN release is hardcoded here. + return os.environ.get("ASCEND_TOOLKIT_HOME", "/usr/local/Ascend/ascend-toolkit/latest") + + +def _so_path(): + """Where to build/cache libmxfp4fused.so. + + The source directory is preferred (keeps the .so next to the kernel, survives + across runs), but it is read-only in some container mounts and in a pip-installed + tree, so fall back to a user cache directory. Override with KT_MXFP4_SO_DIR. + """ + override = os.environ.get("KT_MXFP4_SO_DIR") + if override: + d = Path(override) + elif os.access(_HERE, os.W_OK): + d = _HERE + else: + d = ( + Path(os.environ.get("SGLANG_CACHE_DIR") or os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache")) + / "ascendc_mxfp4" + ) + d.mkdir(parents=True, exist_ok=True) + return d / "libmxfp4fused.so" + + +def _build(so: Path): + cann = _cann_home() + tk = f"{cann}/aarch64-linux/tikcpp" + inc = [ + f"{tk}/tikcfw", + f"{tk}/tikcfw/impl", + f"{tk}/tikcfw/interface", + f"{tk}/tikcfw/lib", + f"{cann}/aarch64-linux/include", + ] + cmd = [ + "bisheng", + "-x", + "asc", + "--cce-aicore-arch=dav-c220", + "-O2", + "-std=c++17", + "-fPIC", + "-shared", + *[f"-I{p}" for p in inc], + str(_SRC), + "-o", + str(so), + f"-L{cann}/aarch64-linux/lib64", + "-lruntime", + "-lascendcl", + ] + proc = subprocess.run(cmd, capture_output=True) + if proc.returncode != 0: + raise RuntimeError( + f"bisheng failed to build {_SRC.name} (exit {proc.returncode}).\n" + f" command: {' '.join(cmd)}\n" + f" ASCEND_TOOLKIT_HOME={cann}\n" + f" stderr: {proc.stderr.decode('utf-8', 'replace')[-2000:]}" + ) + + +def get_lib(): + """Build (if needed) and load the fused kernel .so. Thread-safe, idempotent.""" + global _lib + if _lib is not None: + return _lib + with _lock: + if _lib is None: + so = _so_path() + if not so.exists() or so.stat().st_mtime < _SRC.stat().st_mtime: + _build(so) + lib = ctypes.CDLL(str(so)) + lib.launch_mxfp4_fused.restype = None + lib.launch_mxfp4_fused.argtypes = ( + [ctypes.c_void_p, ctypes.c_uint32] + [ctypes.c_void_p] * 8 + [ctypes.c_uint32] * 4 + ) + lib.launch_mxfp4_fused_blk.restype = None + lib.launch_mxfp4_fused_blk.argtypes = ( + [ctypes.c_void_p, ctypes.c_uint32] + [ctypes.c_void_p] * 9 + [ctypes.c_uint32] * 4 + ) + _lib = lib + return _lib + + +def _consts(HALF, NB, dev): + key = (HALF, NB, str(dev)) + if key in _consts_cache: + return _consts_cache[key] + b = np.arange(256, dtype=np.int64) + lutLo = _FP4[b & 0xF].astype(np.float32) + lutHi = _FP4[(b >> 4) & 0xF].astype(np.float32) + lutE8 = ((b.astype(np.uint32)) << 23).view(np.float32).astype(np.float32) + j = np.arange(HALF, dtype=np.int64) + scOff = ((j >> 4) * 4).astype(np.int32) + out = tuple(torch.from_numpy(a).to(dev) for a in (lutLo, lutHi, lutE8, scOff)) + _consts_cache[key] = out + return out + + +_blk_consts_cache = {} + + +def _blk_consts(HALF, NB, dev): + """codeOff/scaleOff: byte offsets of code j / scale block b in the half-cast GGUF block buffer + ([nb,17] per row). Used by mxfp4_fused_blk to de-interleave in UB via Gather.""" + key = (HALF, NB, str(dev)) + if key in _blk_consts_cache: + return _blk_consts_cache[key] + j = np.arange(HALF, dtype=np.int64) + codeOff = (((j // 16) * 17 + 1 + (j % 16)) * 2).astype(np.uint32) + b = np.arange(NB, dtype=np.int64) + scaleOff = ((b * 17) * 2).astype(np.uint32) + out = (torch.from_numpy(codeOff).to(dev), torch.from_numpy(scaleOff).to(dev)) + _blk_consts_cache[key] = out + return out + + +_NZ_CHUNK = int(os.environ.get("KT_MXFP4_NZ_CHUNK", "32")) # experts/chunk -> bounds HBM transient + + +def convert_proj(codes_dev, scale_dev, IN, blockdim=40, packing="consecutive", out_nz=None): + """One projection: MXFP4 codes/scale [E,OUT,*] -> (q_nz [E,IN,OUT] FRACTAL_NZ, oscale bf16 [E,OUT]). + + Chunked over experts so the transient (int8 planes + de-interleave + NZ cast) stays small — + only the final [E,IN,OUT] NZ output is full-size (HBM-bounded like the W8A8 slot). + + out_nz: optional pre-allocated FRACTAL_NZ [E,IN,OUT] int8 buffer to write into (the reserved + streaming slot). When given, no per-call ~GBs output allocation happens — the layer's NZ is + produced straight into the reused slot (HBM budgeted once at load). Must already be NZ-format + with matching shape. When None, a fresh buffer is allocated (back-compat). + + packing: nibble layout of the code bytes — how the kernel's lo/hi planes map back to K-positions. + "consecutive" (native safetensors): byte j -> Kpos 2j (lo), 2j+1 (hi) -> interleave. + "halfblock" (GGUF block_mxfp4): byte j -> Kpos g*32+jl (lo), +16 (hi) within its 32-group + -> per-group [lo0..15 | hi0..15] concat. + Both decode the SAME K-ordered weights bit-for-bit; the kernel and scale->block mapping (scOff) + are packing-agnostic, so only this post-step rearrange differs (no .so change).""" + import torch_npu + + lib = get_lib() + dev = codes_dev.device + E, OUT, HALF = codes_dev.shape + NB = scale_dev.shape[2] + HALFp = IN // 2 + lutLo, lutHi, lutE8, scOff = _consts(HALF, NB, dev) + st = torch.npu.current_stream().npu_stream + P = lambda t: ctypes.c_void_p(t.data_ptr()) + + oscale = torch.empty((E, OUT), dtype=torch.bfloat16, device=dev) + for c in range(0, E, _NZ_CHUNK): + ce = min(c + _NZ_CHUNK, E) + Ec = ce - c + Rc = Ec * OUT + cd = codes_dev[c:ce].reshape(Rc, HALF).contiguous() + sd = scale_dev[c:ce].reshape(Rc, NB).contiguous() + out = torch.empty((Rc, IN), dtype=torch.int8, device=dev) # two planes [lo|hi] + Rp = (Rc + _ACC - 1) // _ACC * _ACC + osc = torch.empty((Rp,), dtype=torch.float32, device=dev) + lib.launch_mxfp4_fused( + ctypes.c_void_p(st), + blockdim, + P(cd), + P(sd), + P(out), + P(osc), + P(lutLo), + P(lutHi), + P(lutE8), + P(scOff), + Rc, + HALF, + NB, + IN, + ) + # See _TQ_SYNC: only needed when the task queue is on (then the raw ctypes launch is not + # ordered against the torch post-step reading `out`/`osc`). With it off, stream FIFO orders + # them and we skip the host stall -> async convert, fast prefill. + if _TQ_SYNC: + torch.npu.synchronize() + # De-interleave the [lo|hi] planes (contiguous stack) then transpose OUT<->IN. The old depool + # hot spot was (a) a strided 1-byte de-interleave scatter (~2.4s/layer) and (b) an int8 + # transpose that degenerates to a 1-byte gather (~0.6s, ~20GB/s). (a) is gone via the + # contiguous stack; (b) is killed by transposing in fp16 (vectorized) and round-tripping + # int8->fp16->int8 — exact because |q|<=127. Net post-step ~3s -> ~0.13s. The .contiguous() + # is mandatory: feeding a transposed view to format_cast lays down WRONG NZ bytes on device + # (looks fine via .cpu() which de-formats, but grouped_matmul reads garbage). + lo, hi = out[:, :HALFp], out[:, HALFp:] + if packing == "halfblock": + nb = HALFp // 16 + q = torch.cat([lo.reshape(Rc, nb, 16), hi.reshape(Rc, nb, 16)], dim=2).reshape(Ec, OUT, IN) + else: + q = torch.stack([lo, hi], dim=2).reshape(Ec, OUT, IN) # consecutive interleave [E,OUT,IN] + nd = q.to(torch.float16).transpose(1, 2).contiguous().to(torch.int8) # [E,IN,OUT] + nz = torch_npu.npu_format_cast(nd, _NZ) + if out_nz is None: + out_nz = torch.empty((E,) + tuple(nz.shape[1:]), dtype=torch.int8, device=dev) + out_nz[c:ce].copy_(nz) + oscale[c:ce] = osc[:Rc].reshape(Ec, OUT).to(torch.bfloat16) + # Second sync (task-queue-on only): let the osc read finish before the next chunk reuses it. + if _TQ_SYNC: + torch.npu.synchronize() + del out, q, nd, nz, osc, cd, sd + return out_nz, oscale + + +def mxfp4_layer_to_nz_slots(c13, s13, c2, s2, H, I, blockdim=40, packing="consecutive", out_w13=None, out_w2=None): + """Full layer depool conversion -> (w13_nz, s13b, w2_nz, s2b), the exact tensors the streaming + slot + npu_fused_experts consume (replacing the resident W8A8 pool). packing: see convert_proj + ("consecutive" for native safetensors codes, "halfblock" for GGUF block_mxfp4 codes). + out_w13/out_w2: optional pre-reserved NZ slots to convert into (no per-layer output alloc).""" + w13_nz, s13b = convert_proj(c13, s13, H, blockdim, packing, out_nz=out_w13) + w2_nz, s2b = convert_proj(c2, s2, I, blockdim, packing, out_nz=out_w2) + return w13_nz, s13b, w2_nz, s2b + + +def convert_proj_blk(blocks_dev, IN, blockdim=40, out_nz=None): + """One projection from RAW GGUF block_mxfp4 [E,OUT,nb*17] -> (q_nz [E,IN,OUT] FRACTAL_NZ, oscale + bf16 [E,OUT]). The de-interleave (scale|codes per 17B block) is done IN-KERNEL (mxfp4_fused_blk, + UB Gather) -- no host/device de-interleave (the slow 16-of-17 strided int8 copy). The kernel + output `out` (two [lo|hi] planes) is byte-identical to the de-interleaved path, so the post-step + is the same half-block rearrange. out_nz: optional pre-reserved NZ buffer.""" + import torch_npu + + lib = get_lib() + dev = blocks_dev.device + E, OUT, NB17 = blocks_dev.shape + nb = NB17 // 17 + HALF = nb * 16 + HALFp = IN // 2 + lutLo, lutHi, lutE8, scOff = _consts(HALF, nb, dev) + codeOff, scaleOff = _blk_consts(HALF, nb, dev) + st = torch.npu.current_stream().npu_stream + P = lambda t: ctypes.c_void_p(t.data_ptr()) + oscale = torch.empty((E, OUT), dtype=torch.bfloat16, device=dev) + for c in range(0, E, _NZ_CHUNK): + ce = min(c + _NZ_CHUNK, E) + Ec = ce - c + Rc = Ec * OUT + bd = blocks_dev[c:ce].reshape(Rc, NB17).contiguous() + out = torch.empty((Rc, IN), dtype=torch.int8, device=dev) + Rp = (Rc + _ACC - 1) // _ACC * _ACC + osc = torch.empty((Rp,), dtype=torch.float32, device=dev) + lib.launch_mxfp4_fused_blk( + ctypes.c_void_p(st), + blockdim, + P(bd), + P(out), + P(osc), + P(lutLo), + P(lutHi), + P(lutE8), + P(scOff), + P(codeOff), + P(scaleOff), + Rc, + HALF, + nb, + IN, + ) + if _TQ_SYNC: + torch.npu.synchronize() + lo, hi = out[:, :HALFp], out[:, HALFp:] + nbb = HALFp // 16 + q = torch.cat([lo.reshape(Rc, nbb, 16), hi.reshape(Rc, nbb, 16)], dim=2).reshape(Ec, OUT, IN) + nd = q.to(torch.float16).transpose(1, 2).contiguous().to(torch.int8) + nz = torch_npu.npu_format_cast(nd, _NZ) + if out_nz is None: + out_nz = torch.empty((E,) + tuple(nz.shape[1:]), dtype=torch.int8, device=dev) + out_nz[c:ce].copy_(nz) + oscale[c:ce] = osc[:Rc].reshape(Ec, OUT).to(torch.bfloat16) + if _TQ_SYNC: + torch.npu.synchronize() + del out, q, nd, nz, osc, bd + return out_nz, oscale + + +def mxfp4_layer_to_nz_slots_blk(blk13, blk2, H, I, blockdim=40, out_w13=None, out_w2=None): + """Full layer conversion from RAW GGUF blocks (in-kernel de-interleave) -> slot tensors. + blk13 = cat(gate,up) blocks [E,2I,nbH*17]; blk2 = down blocks [E,H,nbI*17].""" + w13_nz, s13b = convert_proj_blk(blk13, H, blockdim, out_nz=out_w13) + w2_nz, s2b = convert_proj_blk(blk2, I, blockdim, out_nz=out_w2) + return w13_nz, s13b, w2_nz, s2b From 84aa637cc4978ef7f430fc6a2deec73da228ea03 Mon Sep 17 00:00:00 2001 From: Pan-Boyi Date: Fri, 14 Aug 2026 20:07:01 +0800 Subject: [PATCH 09/13] [docs](kt-kernel): record why LLAMA_MOE_TP weights cannot alias the GGUF mmap load_weights copies each expert into node-local buffers, and the obvious optimisation is to skip that and point m_local_*_proj_ straight at the GGUF mmap the caller already holds. That is only correct when this TP owns the whole tensor, and TP_MOE::load_weights hands each of the threadpool_count subpools a distinct intermediate_size / threadpool_count slice. Every supported deployment runs one subpool per NUMA node, so the condition is unreachable and the branch would be dead code. Comment only; no behaviour change. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Q7su5SAw6wskXDpzuPLMb4 Signed-off-by: Pan-Boyi --- kt-kernel/operators/llamafile/moe.hpp | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/kt-kernel/operators/llamafile/moe.hpp b/kt-kernel/operators/llamafile/moe.hpp index 8beecf145..da19b8c04 100644 --- a/kt-kernel/operators/llamafile/moe.hpp +++ b/kt-kernel/operators/llamafile/moe.hpp @@ -287,6 +287,15 @@ class LLAMA_MOE_TP { uint8_t* down_proj = (uint8_t*)config.down_proj + offset * ggml_type_size((ggml_type)config.down_type) / ggml_blck_size((ggml_type)config.down_type); + // The reshuffle below always copies. It degenerates to an identity copy only when this TP owns + // the whole tensor (offset == 0 && complete_intermediate_size == config.intermediate_size), + // which is the sole case where the three m_local_*_proj_ pointers could instead alias the + // caller's GGUF mmap and save the ~137GB of node-local buffers. TP_MOE::load_weights splits the + // tensor across threadpool_count subpools, so that case needs threadpool_count == 1; every + // supported deployment runs one subpool per NUMA node (8), where each TP gets + // intermediate_size / 8 and no alias is possible. Do not add the alias branch without also + // making a single-subpool configuration real -- it would be dead code. + // Per-expert byte strides. The source tensors are laid out with the FULL // intermediate_size (complete_intermediate_size); this TP only owns the // [offset, offset+intermediate_size) block — hence the base-pointer offset From 1a297be27b0eb2122563f773c6ec91cc470fd0d9 Mon Sep 17 00:00:00 2001 From: Pan-Boyi Date: Sat, 15 Aug 2026 11:21:53 +0800 Subject: [PATCH 10/13] [build](kt-kernel): carry the NumPy 2 fix for the GGUF reader kt-kernel/python/utils/loader.py reads the per-layer MXFP4 expert files through gguf.gguf_reader.GGUFReader. That reader calls ndarray.newbyteorder, which NumPy 2.0 removed, so on NumPy >= 2.0 loading the experts fails with AttributeError before the first layer is read. llama.cpp is pinned to b3173, which predates the upstream fix, so carry the delta the same way as the MXFP4 type: an in-repo pure diff applied at configure time. CMakeLists.txt already globs the patch directory and applies each patch idempotently, so no build change is needed. NumPy 1.x environments are unaffected: the helper rewrites the call to dtype.newbyteorder, which behaves identically there. Signed-off-by: Pan-Boyi --- .../0002-gguf-numpy2-byteorder.patch | 44 +++++++++++++++++++ .../third_party_patches/llama.cpp/README.md | 11 +++++ 2 files changed, 55 insertions(+) create mode 100644 kt-kernel/third_party_patches/llama.cpp/0002-gguf-numpy2-byteorder.patch diff --git a/kt-kernel/third_party_patches/llama.cpp/0002-gguf-numpy2-byteorder.patch b/kt-kernel/third_party_patches/llama.cpp/0002-gguf-numpy2-byteorder.patch new file mode 100644 index 000000000..3b355b1cb --- /dev/null +++ b/kt-kernel/third_party_patches/llama.cpp/0002-gguf-numpy2-byteorder.patch @@ -0,0 +1,44 @@ +diff --git a/gguf-py/gguf/gguf_reader.py b/gguf-py/gguf/gguf_reader.py +index e48bc00c3..be7a994a4 100644 +--- a/gguf-py/gguf/gguf_reader.py ++++ b/gguf-py/gguf/gguf_reader.py +@@ -35,6 +35,16 @@ logger = logging.getLogger(__name__) + READER_SUPPORTED_VERSIONS = [2, GGUF_VERSION] + + ++def _np_apply_byteorder( ++ arr: npt.NDArray[Any], ++ order: Literal["I"] | Literal["S"] | Literal["<"] | Literal[">"] | None, ++) -> npt.NDArray[Any]: ++ """NumPy 2+ compatible replacement for ndarray.newbyteorder (removed in NumPy 2.0).""" ++ if order is None or order == "I": ++ return arr ++ return arr.view(arr.dtype.newbyteorder(order)) ++ ++ + class ReaderField(NamedTuple): + # Offset to start of this field. + offset: int +@@ -96,7 +106,7 @@ class GGUFReader: + # If we get 0 here that means it's (probably) a GGUF file created for + # the opposite byte order of the machine this script is running on. + self.byte_order = 'S' +- temp_version = temp_version.newbyteorder(self.byte_order) ++ temp_version = _np_apply_byteorder(temp_version, self.byte_order) + version = temp_version[0] + if version not in READER_SUPPORTED_VERSIONS: + raise ValueError(f'Sorry, file appears to be version {version} which we cannot handle') +@@ -135,10 +145,9 @@ class GGUFReader: + count = int(count) + itemsize = int(np.empty([], dtype = dtype).itemsize) + end_offs = offset + itemsize * count +- return ( +- self.data[offset:end_offs] +- .view(dtype = dtype)[:count] +- .newbyteorder(override_order or self.byte_order) ++ return _np_apply_byteorder( ++ self.data[offset:end_offs].view(dtype=dtype)[:count], ++ override_order or self.byte_order, + ) + + def _push_field(self, field: ReaderField, skip_sum: bool = False) -> int: diff --git a/kt-kernel/third_party_patches/llama.cpp/README.md b/kt-kernel/third_party_patches/llama.cpp/README.md index 60fe25e94..871d0c7dc 100644 --- a/kt-kernel/third_party_patches/llama.cpp/README.md +++ b/kt-kernel/third_party_patches/llama.cpp/README.md @@ -31,6 +31,17 @@ This is required because the DeepSeek-V4 GGUF stores its CPU-offloaded expert tensors as MXFP4; without it ggml rejects the tensor type and kt-kernel cannot load the experts. +### `0002-gguf-numpy2-byteorder.patch` + +`gguf-py/gguf/gguf_reader.py` — routes the two byte-order conversions through a +`_np_apply_byteorder` helper that calls `dtype.newbyteorder` instead of +`ndarray.newbyteorder`, which NumPy 2.0 removed. + +`kt-kernel/python/utils/loader.py` reads the per-layer expert GGUFs through +`gguf.gguf_reader.GGUFReader`, so on NumPy >= 2.0 the unpatched reader raises +`AttributeError` before any expert is loaded. Environments still on NumPy 1.x are +unaffected either way — the helper is a no-op rewrite there. + ## Applying by hand The build applies these automatically. To do it manually (e.g. when debugging): From 96b2a07325961e4a72db493867b7b77784538c67 Mon Sep 17 00:00:00 2001 From: pan-boyi Date: Tue, 18 Aug 2026 11:36:41 +0800 Subject: [PATCH 11/13] [perf](kt-kernel): alias the GGUF mmap when one subpool owns the whole tensor LLAMA_MOE_TP::load_weights reshuffles each NUMA subpool's slice of the expert tensors into node-local buffers. With --kt-threadpool-count 1 the subpool owns the whole tensor, so offset == 0 and intermediate_size == complete_intermediate_size: the reshuffle is an identity copy, and its only effect is a second, byte-identical copy of the expert set in anonymous memory. A preceding commit documented this case and declined to handle it, on the grounds that "every supported deployment runs one subpool per NUMA node (8)". That premise does not hold on a single-NUMA host, where one subpool is the only valid setting -- and one subpool is what the reference launcher defaults to. The cost is not just the memory. On a single-NUMA Atlas A3 (229 GiB RAM, 137 GiB MXFP4 GGUF set) the redundant copy left 62 GiB of page cache for a 137 GiB working set. Streaming prefill (KT_PREFILL_STREAM=1 with KT_MXFP4_GGUF_DEDUP=1) sweeps the entire expert set through that mmap on every request, so it thrashed and re-read 135.6 GiB from a 350 MB/s disk each time. Take the identity case: free the reservations and point the three m_local_*_proj_ pointers at the caller's mmap. The mmap then becomes the single shared copy, so the CPU MoE and the streaming-prefill reader hit the same page cache. Multi-NUMA has strided per-node slices and still falls through to the copy. The memcpy also warmed the page cache as a side effect; MADV_WILLNEED restores that warming without the copy, so the first long prefill does not pay a cold read. Measured on Atlas A3, single die, one NUMA node, --kt-cpuinfer 32 --kt-threadpool-count 1 --mem-fraction-static 0.81, all five feature switches on, 1 warmup + 3 measured iterations: scheduler RssAnon 140.2 GiB -> 2.9 GiB disk read per prefill 135.6 GiB -> 0 prefill @ 801 tokens 397.1 s -> 15.6 s prefill @ 7823 tokens 395.9 s -> 16.5 s decode @ 7823 tokens 19.2 tok/s -> 19.9 tok/s Co-Authored-By: Claude Opus 5 --- kt-kernel/operators/llamafile/moe.hpp | 55 +++++++++++++++++++++++---- 1 file changed, 47 insertions(+), 8 deletions(-) diff --git a/kt-kernel/operators/llamafile/moe.hpp b/kt-kernel/operators/llamafile/moe.hpp index da19b8c04..1bc4f89be 100644 --- a/kt-kernel/operators/llamafile/moe.hpp +++ b/kt-kernel/operators/llamafile/moe.hpp @@ -5,6 +5,8 @@ #endif #include #include +#include // madvise / MADV_WILLNEED (warm the aliased GGUF mmap into page cache) +#include // sysconf(_SC_PAGESIZE) #include #include @@ -108,6 +110,10 @@ class LLAMA_MOE_TP { uint8_t* m_local_gate_proj_; // [expert_num * intermediate_size * hidden_size ( /32 if quantized)] uint8_t* m_local_up_proj_; // [expert_num * intermediate_size * hidden_size ( /32 if quantized)] uint8_t* m_local_down_proj_; // [expert_num * hidden_size * intermediate_size ( /32 if quantized)] + // Single-NUMA (tp_count==1): the three proj pointers ALIAS the source GGUF mmap instead of + // owning a memcpy'd copy (see load_weights identity short-circuit). When true, the ctor's + // lazy new[] reservations were freed and these must NOT be delete[]'d (they point into mmap). + bool m_weights_aliased_ = false; float* s_input_fp32_; // [hidden_size] uint8_t* s_gate_input_; // [hidden_size * ggml_type_size(ggml_internal_get_type_traits(gate_type).vec_dot_type) / @@ -287,14 +293,47 @@ class LLAMA_MOE_TP { uint8_t* down_proj = (uint8_t*)config.down_proj + offset * ggml_type_size((ggml_type)config.down_type) / ggml_blck_size((ggml_type)config.down_type); - // The reshuffle below always copies. It degenerates to an identity copy only when this TP owns - // the whole tensor (offset == 0 && complete_intermediate_size == config.intermediate_size), - // which is the sole case where the three m_local_*_proj_ pointers could instead alias the - // caller's GGUF mmap and save the ~137GB of node-local buffers. TP_MOE::load_weights splits the - // tensor across threadpool_count subpools, so that case needs threadpool_count == 1; every - // supported deployment runs one subpool per NUMA node (8), where each TP gets - // intermediate_size / 8 and no alias is possible. Do not add the alias branch without also - // making a single-subpool configuration real -- it would be dead code. + // Single-NUMA fast path (tp_count==1 => offset==0 && intermediate==complete): the per-expert + // reshuffle below degenerates to an IDENTITY copy (dst layout == src layout). Instead of + // duplicating the whole expert set into anonymous RAM, ALIAS the source (GGUF mmap) directly: + // * ~138GB anonymous RssAnon -> ~0 (weights are read-only after load; forward only reads). + // * the mmap becomes the SINGLE shared copy (CPU MoE compute + streaming-prefill dedup share + // one page cache) => streaming-prefill stops re-reading cold GGUF from disk. + // Multi-NUMA (tp_count>1) has strided per-node slices and cannot alias -> falls through to copy. + if (offset == 0 && complete_intermediate_size == config.intermediate_size) { + delete[] m_local_gate_proj_; + delete[] m_local_up_proj_; + delete[] m_local_down_proj_; + m_local_gate_proj_ = gate_proj; // == (uint8_t*)config.gate_proj (offset==0) + m_local_up_proj_ = up_proj; + m_local_down_proj_ = down_proj; + m_weights_aliased_ = true; + + // Aliasing skips the memcpy that used to (as a side effect) read the whole GGUF into page + // cache at load. Restore that warming WITHOUT copying: MADV_WILLNEED kicks off async + // readahead of the aliased expert tensors, so by the time traffic arrives the CPU-MoE / + // streaming-prefill shared page cache is hot (else the first long prefill pays a one-time + // cold read). Non-blocking, best-effort (page-align the start; ignore errors). + const size_t gate_bytes = (size_t)config.expert_num * config.intermediate_size * config.hidden_size * + ggml_type_size((ggml_type)config.gate_type) / + ggml_blck_size((ggml_type)config.gate_type); + const size_t up_bytes = (size_t)config.expert_num * config.intermediate_size * config.hidden_size * + ggml_type_size((ggml_type)config.up_type) / ggml_blck_size((ggml_type)config.up_type); + const size_t down_bytes = (size_t)config.expert_num * config.hidden_size * config.intermediate_size * + ggml_type_size((ggml_type)config.down_type) / + ggml_blck_size((ggml_type)config.down_type); + const long pg = sysconf(_SC_PAGESIZE); + auto warm = [pg](void* p, size_t n) { + if (!p || !n || pg <= 0) return; + uintptr_t a = (uintptr_t)p, start = a & ~((uintptr_t)pg - 1); + madvise((void*)start, n + (a - start), MADV_WILLNEED); // best-effort readahead + }; + warm(m_local_gate_proj_, gate_bytes); + warm(m_local_up_proj_, up_bytes); + warm(m_local_down_proj_, down_bytes); + return; + } + // Per-expert byte strides. The source tensors are laid out with the FULL // intermediate_size (complete_intermediate_size); this TP only owns the From 067292633a1bd826b812777e6ad98928bbf881c3 Mon Sep 17 00:00:00 2001 From: Pan-Boyi Date: Fri, 21 Aug 2026 11:22:11 +0800 Subject: [PATCH 12/13] [build](kt-kernel): pin torch back to 2.9.1 Reverts the relaxed `torch>=2.9.1,<2.11` bound from review discussion: KT is not free to move independently here. ktransformers, sglang-kt, accelerate-kt and transformers-kt converged on 2.9.1 and the fine-tuning path needs all of them to agree, so relaxing the bound is a risk for that side even though inference alone would tolerate 2.10. Verified 2.9.1 on two Ascend hosts (isolated venvs, system environments untouched): torch 2.9.1 + torch_npu 2.9.1 import, `torch.npu.is_available()`, a real bf16 matmul on device, server start-up and question answering all pass on CANN 9.0.0. Deployment and answering only -- no accuracy or throughput comparison was run. The comment now records the consequence for vendor images: a CANN 9 image that ships torch 2.10 has to rebuild its torch-linked extensions for 2.9.1 rather than have this bound widened. kt-kernel's own extension links no libtorch symbols at all, so it is unaffected either way. Co-Authored-By: Claude Opus 5 --- kt-kernel/pyproject.toml | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/kt-kernel/pyproject.toml b/kt-kernel/pyproject.toml index 733c3b060..b34f6f59b 100644 --- a/kt-kernel/pyproject.toml +++ b/kt-kernel/pyproject.toml @@ -21,9 +21,11 @@ classifiers = [ requires-python = ">=3.11" dependencies = [ # Core dependencies - # Keep the kernel usable with both the 2.9 reference stack and CANN 9 images - # whose torch-npu package requires torch 2.10 exactly. - "torch>=2.9.1,<2.11", + # Pinned, not a range: ktransformers, sglang-kt, accelerate-kt and transformers-kt + # converged on 2.9.1, and the fine-tuning path depends on all of them agreeing. A CANN 9 + # image that ships torch 2.10 therefore needs its torch-linked extensions rebuilt for + # 2.9.1 rather than this bound relaxed -- see the Ascend tutorial. + "torch==2.9.1", "safetensors>=0.4.0", "numpy>=1.24.0", # triton has no aarch64 wheels; the Ascend build does not use it. From 247f0ee3cc0d714e7c94825ca3ad76c9a215e664 Mon Sep 17 00:00:00 2001 From: pan-boyi Date: Wed, 26 Aug 2026 14:44:18 +0800 Subject: [PATCH 13/13] [build](kt-kernel): pin torch to 2.9.1 in requirements.txt too Applies the review suggestion on kt-kernel/requirements.txt:6. pyproject.toml was pinned to torch==2.9.1 but requirements.txt still carried the old torch>=2.9.1,<2.11 range, so the two files disagreed about the one version the fine-tuning stack needs everyone to agree on. The section heading said "minimum versions", which is no longer what this is. Co-Authored-By: Claude Opus 5 --- kt-kernel/requirements.txt | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/kt-kernel/requirements.txt b/kt-kernel/requirements.txt index 6070f43f5..9087878c6 100644 --- a/kt-kernel/requirements.txt +++ b/kt-kernel/requirements.txt @@ -2,8 +2,10 @@ # These dependencies will be automatically installed when running `pip install .` # You can skip this file if you already have these packages installed -# Core dependencies (minimum versions) -torch>=2.9.1,<2.11 +# Core dependencies +# torch is pinned, not a range, and must stay in step with pyproject.toml -- see the +# comment there. +torch==2.9.1 safetensors>=0.4.0 numpy>=1.24.0 triton>=2.0.0; platform_machine == "x86_64"