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 bd2adcec2..12fc70393 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) @@ -38,13 +44,14 @@ foreach(_KT_GPU_BACKEND IN ITEMS KTRANSFORMERS_USE_SYCL KTRANSFORMERS_USE_ROCM KTRANSFORMERS_USE_MUSA - KTRANSFORMERS_USE_MACA) + KTRANSFORMERS_USE_MACA + KTRANSFORMERS_USE_ASCEND_NPU) if(${_KT_GPU_BACKEND}) math(EXPR _KT_GPU_BACKEND_COUNT "${_KT_GPU_BACKEND_COUNT} + 1") endif() endforeach() if(_KT_GPU_BACKEND_COUNT GREATER 1) - message(FATAL_ERROR "CUDA, SYCL, ROCm, MUSA, and MACA backends are mutually exclusive") + message(FATAL_ERROR "CUDA, SYCL, ROCm, MUSA, MACA, and Ascend backends are mutually exclusive") endif() # Choose compilers BEFORE project() so CMake honors them @@ -239,18 +246,26 @@ if(CMAKE_OSX_ARCHITECTURES STREQUAL "arm64" OR CMAKE_GENERATOR_PLATFORM_LWR STRE # Raspberry Pi 3, 4, Zero 2 (32-bit) list(APPEND ARCH_FLAGS -mno-unaligned-access) endif() - # add_compile_definitions(__ARM_NEON) - # list(APPEND ARCH_FLAGS -march=armv8.2-a+fp16+dotprod) - # add_compile_definitions(__ARM_FEATURE_DOTPROD) - # add_compile_definitions(__aarch64__) - - # add_compile_definitions(__ARM_NEON) - list(APPEND ARCH_FLAGS -march=armv8.2-a+fp16+dotprod+sve+bf16) - # list(APPEND ARCH_FLAGS -march=armv8-a+dotprod+sha3+sm4+fp16fml+sve+rng+sb+ssbs+i8mm+bf16+flagm+pauth) - # add_compile_definitions(__ARM_FEATURE_DOTPROD) - # add_compile_definitions(__ARM_FEATURE_SVE) - # add_compile_definitions(__ARM_FEATURE_MATMUL_INT8) - # add_compile_definitions(__aarch64__) + # ARM extensions are selected dynamically. K920 / Cortex-A76 = armv8.2-a + fp16 + dotprod + # (NEON only). K930+ / Neoverse-V1+ add SVE/BF16/I8MM. Toggle via -DLLAMA_ARM_SVE=ON etc. + set(_kt_arm_arch "armv8.2-a") + if(LLAMA_ARM_FP16) + set(_kt_arm_arch "${_kt_arm_arch}+fp16") + endif() + if(LLAMA_ARM_DOTPROD) + set(_kt_arm_arch "${_kt_arm_arch}+dotprod") + endif() + if(LLAMA_ARM_SVE) + set(_kt_arm_arch "${_kt_arm_arch}+sve") + endif() + if(LLAMA_ARM_BF16) + set(_kt_arm_arch "${_kt_arm_arch}+bf16") + endif() + if(LLAMA_ARM_I8MM) + set(_kt_arm_arch "${_kt_arm_arch}+i8mm") + endif() + list(APPEND ARCH_FLAGS "-march=${_kt_arm_arch}") + message(STATUS "ARM target: -march=${_kt_arm_arch}") endif() elseif(CMAKE_OSX_ARCHITECTURES STREQUAL "x86_64" OR CMAKE_GENERATOR_PLATFORM_LWR MATCHES "^(x86_64|i686|amd64|x64|win32)$" OR (NOT CMAKE_OSX_ARCHITECTURES AND NOT CMAKE_GENERATOR_PLATFORM_LWR AND @@ -474,6 +489,70 @@ add_compile_options("$<$:${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) @@ -562,6 +641,27 @@ elseif(KTRANSFORMERS_USE_MACA) add_compile_definitions(KTRANSFORMERS_USE_MACA=1) elseif(KTRANSFORMERS_USE_SYCL) message(STATUS "SYCL GPTQ INT4 support enabled") +elseif(KTRANSFORMERS_USE_ASCEND_NPU) + message(STATUS "Ascend NPU (CANN) backend selected") + add_compile_definitions(KTRANSFORMERS_USE_ASCEND_NPU=1) + add_compile_definitions(USE_ASCEND_NPU=1) + if(DEFINED ENV{ASCEND_TOOLKIT_HOME}) + set(_kt_cann_root "$ENV{ASCEND_TOOLKIT_HOME}") + elseif(DEFINED ENV{CANN_HOME}) + set(_kt_cann_root "$ENV{CANN_HOME}") + else() + set(_kt_cann_root "/usr/local/Ascend/ascend-toolkit/latest") + endif() + find_path(ACL_INCLUDE_DIR acl/acl_rt.h + HINTS "${_kt_cann_root}/include" + REQUIRED) + find_library(ASCEND_CL_LIBRARY NAMES ascendcl + HINTS "${_kt_cann_root}/lib64" "${_kt_cann_root}/runtime/lib64" + REQUIRED) + message(STATUS "CANN root: ${_kt_cann_root}") + message(STATUS "Ascend CL include: ${ACL_INCLUDE_DIR}") + message(STATUS "Ascend CL library: ${ASCEND_CL_LIBRARY}") + include_directories(${ACL_INCLUDE_DIR}) elseif(KTRANSFORMERS_CPU_USE_KML) message(STATUS "KML CPU detected") else() @@ -795,6 +895,10 @@ if(NOT HOST_IS_X86 AND KTRANSFORMERS_CPU_USE_KML) target_compile_definitions(${PROJECT_NAME} PRIVATE CPU_USE_KML) endif() target_link_libraries(${PROJECT_NAME} PRIVATE llama PkgConfig::HWLOC OpenMP::OpenMP_CXX) +if(KTRANSFORMERS_USE_ASCEND_NPU AND ASCEND_CL_LIBRARY) + target_link_libraries(${PROJECT_NAME} PRIVATE ${ASCEND_CL_LIBRARY}) + target_include_directories(${PROJECT_NAME} PRIVATE ${ACL_INCLUDE_DIR}) +endif() if(NOT HOST_IS_X86 AND KTRANSFORMERS_CPU_USE_KML) if(KTRANSFORMERS_CPU_DEBUG) # add_executable(convert-test ${CMAKE_CURRENT_SOURCE_DIR}/operators/kml/convert-test.cpp) 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..bbcbf445a --- /dev/null +++ b/kt-kernel/cpu_backend/ascend_callback_worker.cpp @@ -0,0 +1,153 @@ +#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}; +// 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) { + 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)); + } + } + g_worker_ready.store(true, std::memory_order_release); + + while (!g_stop.load(std::memory_order_acquire)) { + (void)aclrtProcessReport(kProcessReportTimeoutMs); + } + + // Drain callbacks that were enqueued between the last loop iteration and the + // stop request; without this they are silently dropped and their waiters hang. + (void)aclrtProcessReport(kProcessReportTimeoutMs); + + // ACL requires the *subscribing* thread to unsubscribe. Doing it from + // shutdown_callback_worker() after join() would target a dead thread id and + // silently do nothing. + for (aclrtStream stream : g_subscribed_streams) { + (void)aclrtUnSubscribeReport(g_worker_thread_id, stream); + } +} + +void start_worker_locked(aclrtContext ctx) { + if (g_started.load(std::memory_order_acquire)) { + return; + } + g_context = ctx; + g_stop.store(false, std::memory_order_release); + g_worker_ready.store(false, std::memory_order_release); + g_worker = std::thread([ctx]() { worker_main(ctx); }); + g_worker_thread_id = static_cast(g_worker.native_handle()); + g_started.store(true, std::memory_order_release); + + // Barrier: no aclrtSubscribeReport() before the worker is in its report loop. + while (!g_worker_ready.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + + for (aclrtStream stream : g_subscribed_streams) { + aclError err = aclrtSubscribeReport(g_worker_thread_id, stream); + if (err != ACL_SUCCESS) { + std::fprintf(stderr, + "[kt-kernel] ascend_callback_worker: aclrtSubscribeReport failed (%d)\n", + static_cast(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()) { + // worker_main() drains the report queue and unsubscribes before returning. + g_worker.join(); + } + g_subscribed_streams.clear(); + g_started.store(false, std::memory_order_release); + g_worker_ready.store(false, std::memory_order_release); + g_worker_thread_id = 0; +} + +bool callback_worker_running() { + std::lock_guard 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 new file mode 100644 index 000000000..01894399f --- /dev/null +++ b/kt-kernel/cpu_backend/ascend_callback_worker.h @@ -0,0 +1,34 @@ +#pragma once + +// Ascend ACL stream callback subscriber for kt-kernel. +// +// CANN dispatches aclrtLaunchCallback tasks only on a host thread that has called +// aclrtSubscribeReport(threadId, stream) and is running aclrtProcessReport() in a +// loop. This worker mirrors torch_npu's pattern (see NPUGraph.cpp / Graph.cpp). + +#if defined(KTRANSFORMERS_USE_ASCEND_NPU) + +#include + +#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(); + +// 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/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..da763144a --- /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 used as the fallback when no +// callback worker is available. +// ============================================================================ + +#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 bfcbe9548..f60b6ac81 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) @@ -611,6 +614,23 @@ PYBIND11_MODULE(kt_kernel_ext, m) { .def_property_readonly("num_experts", &kt::layerwise::FP8LayerwiseTransport::num_experts) .def_property_readonly("closed", &kt::layerwise::FP8LayerwiseTransport::closed); +#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."); + 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") .def(py::init<>()) diff --git a/kt-kernel/install.sh b/kt-kernel/install.sh index 06a7d8c91..214348b69 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) @@ -124,28 +125,55 @@ 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 } +# 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 4ea91f129..88da8b401 100644 --- a/kt-kernel/operators/llamafile/moe.hpp +++ b/kt-kernel/operators/llamafile/moe.hpp @@ -5,12 +5,15 @@ #endif #include #include +#include // madvise / MADV_WILLNEED (warm the aliased GGUF mmap into page cache) +#include // sysconf(_SC_PAGESIZE) #include #include #include #include #include +#include #include #include @@ -21,6 +24,41 @@ #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 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; +}; +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 +69,39 @@ inline void debug_quant(void* input, ggml_type type) { printf("\n"); } +// --------------------------------------------------------------------------- +// kt_effective_vec_dot_type +// 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 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". +// +// * 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. +// +// * 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) + 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_; @@ -39,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) / @@ -95,12 +170,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 +188,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 +204,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 +229,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); @@ -192,9 +267,6 @@ class LLAMA_MOE_TP { ~LLAMA_MOE_TP() { shared_mem_buffer_numa.dealloc(tp_part_idx, this); } 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); @@ -223,31 +295,94 @@ 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)); + // 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 + // [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() { @@ -266,37 +401,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; @@ -307,16 +451,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])) { @@ -351,7 +497,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"); @@ -367,33 +513,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)); } } @@ -404,6 +551,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( @@ -431,7 +579,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; @@ -456,6 +604,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 % 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); + } + } } void forward_many(int qlen, int k, const int64_t* expert_ids, const float* weights, const void* input, @@ -487,21 +651,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) { @@ -527,9 +691,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); @@ -537,16 +701,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 * @@ -554,9 +718,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 * @@ -571,20 +735,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); @@ -625,7 +789,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]; @@ -640,25 +804,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); @@ -693,7 +858,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); @@ -790,31 +955,59 @@ 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. + // 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; + 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(); + // 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 % 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); + } + } } }; #endif diff --git a/kt-kernel/pyproject.toml b/kt-kernel/pyproject.toml index f877ddae1..b34f6f59b 100644 --- a/kt-kernel/pyproject.toml +++ b/kt-kernel/pyproject.toml @@ -21,10 +21,15 @@ classifiers = [ requires-python = ">=3.11" dependencies = [ # Core dependencies + # 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>=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/python/experts.py b/kt-kernel/python/experts.py index 092146d56..5c3cb0fce 100644 --- a/kt-kernel/python/experts.py +++ b/kt-kernel/python/experts.py @@ -363,24 +363,27 @@ def _create_inference_wrapper( raise NotImplementedError(f"Unsupported inference method: {method}") # Create and return backend instance. - # `swiglu_limit != 0` is validated for block-FP8, MXFP4, and MXFP8. + # `swiglu_limit != 0` is validated for block-FP8, MXFP4, MXFP8 and LLAMAFILE. # NativeMoEWrapper also serves RAWINT4 / 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 耦合. + # LLAMAFILE is on the list for the same reason as the native MXFP4 path: it + # validates the GGUF tensor types while loading and applies the clamp in its + # scalar/NEON-independent activation step. extra_kwargs = {} - if method in ("FP8", "MXFP4", "MXFP8"): + if method in ("FP8", "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='FP8'/'MXFP4'/'MXFP8', " + "method='FP8'/'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 " - "FP8/MXFP4/MXFP8 weights — either unset the env or select a " + "FP8/MXFP4/MXFP8/LLAMAFILE weights — either unset the env or select a " "matching --kt-method." ) return backend_cls( diff --git a/kt-kernel/python/experts_base.py b/kt-kernel/python/experts_base.py index bc58b98f0..647ab2948 100644 --- a/kt-kernel/python/experts_base.py +++ b/kt-kernel/python/experts_base.py @@ -12,11 +12,152 @@ 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. +# +# 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. +# +# 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. +# +# 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. +# ----------------------------------------------------------------------------- + + +# 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 + + +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] + + 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 + + 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, @@ -174,13 +315,17 @@ 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: 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: @@ -374,24 +519,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,15 +561,11 @@ 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) else: immediate_ids = topk_ids_long @@ -422,38 +574,264 @@ 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) + 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(), + 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: + 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) - 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, + ) + 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, + ) + 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) + # 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( + 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 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 ( + 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: - 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, - ), + 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), + 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 sync_submit: + 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 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 +856,27 @@ 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) + + return self.copy_forward_output_to_device(hidden_states) def forward( self, diff --git a/kt-kernel/python/utils/llamafile.py b/kt-kernel/python/utils/llamafile.py index 0f1c4403c..bf79ae7e4 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. """ - _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 @@ -202,8 +214,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..6d250e83b 100644 --- a/kt-kernel/python/utils/loader.py +++ b/kt-kernel/python/utils/loader.py @@ -48,6 +48,42 @@ 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. +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,43 +1060,20 @@ 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 diff --git a/kt-kernel/requirements.txt b/kt-kernel/requirements.txt index d2b31e2e1..9087878c6 100644 --- a/kt-kernel/requirements.txt +++ b/kt-kernel/requirements.txt @@ -2,9 +2,11 @@ # 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) +# 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 +triton>=2.0.0; platform_machine == "x86_64" gguf>=0.17.0 diff --git a/kt-kernel/setup.py b/kt-kernel/setup.py index 034b03c9f..64718b4e3 100644 --- a/kt-kernel/setup.py +++ b/kt-kernel/setup.py @@ -36,14 +36,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: @@ -211,6 +219,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") @@ -492,6 +514,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") @@ -535,7 +571,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" @@ -545,6 +587,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 ( @@ -553,6 +609,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) ] @@ -716,6 +773,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/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/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 new file mode 100644 index 000000000..871d0c7dc --- /dev/null +++ b/kt-kernel/third_party_patches/llama.cpp/README.md @@ -0,0 +1,68 @@ +# 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. + +### `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): + +```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/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 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 +}