mirror of
https://github.com/kvcache-ai/ktransformers.git
synced 2026-10-16 21:01:14 +08:00
losewayy
a02efe36b5
* [feat](kt-kernel): add AVX2 RAWINT4 routed-expert LoRA SFT backend * [fix](kt-kernel): fix RAWINT4 SFT wrapper and clamp gradients Allow the RAWINT4 production wrapper on AVX2 runtimes and reject shared backward pools before construction. Preserve FP32 clamp decisions across BF16 caches, remove unused configuration state, and cover rectangular dimensions, multiblock dX, clamp boundaries, and production wrapper tasks with tighter LoRA gradient checks. * [fix](kt-kernel): plumb swiglu_limit through the RAWINT4 SFT python path The production wrapper hard-zeroed swiglu_limit for every format except MXFP4, and both SFT python gates rejected non-zero values for RAWINT4 even though MOESFTConfig carries the field and the AVX2 kernel implements the clamp in forward plus the gradient mask in backward. RAWINT4_SFT now accepts a finite non-negative swiglu_limit end to end, while the AMX K2 part (which computes plain SiLU) fails fast instead of silently dropping the clamp. The production-wrapper contract test is parametrized over swiglu_limit 0 and 2, a model-config plumb test covers plain and nested text_config shapes, the __main__ entry exercises both limits, and bench_rawint4_sft.py gains --swiglu-limit.
1217 lines
59 KiB
C++
1217 lines
59 KiB
C++
/**
|
||
* @Description :
|
||
* @Author : chenht2022, Jianwei Dong
|
||
* @Date : 2024-07-22 02:03:22
|
||
* @Version : 1.0.0
|
||
* @LastEditors : Jianwei Dong
|
||
* @LastEditTime : 2024-08-26 22:47:06
|
||
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||
**/
|
||
// Python bindings
|
||
#include <sys/types.h>
|
||
#include <sys/wait.h>
|
||
#include <unistd.h>
|
||
|
||
#if defined(KTRANSFORMERS_ENABLE_CPPTRACE)
|
||
#include <cpptrace/cpptrace.hpp>
|
||
#endif
|
||
#include <csignal>
|
||
#include <cstddef>
|
||
#include <cstring>
|
||
|
||
#include "bf16_sumsq.h"
|
||
#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)
|
||
#include "operators/moe_kernel/la/kernel.hpp"
|
||
#include "operators/moe_kernel/moe.hpp"
|
||
#endif
|
||
|
||
#if defined(__aarch64__) && defined(CPU_USE_KML)
|
||
#if defined(KTRANSFORMERS_CPU_MLA)
|
||
#include "operators/kml/deepseekv3.hpp"
|
||
#include "operators/kml/gate.hpp"
|
||
#include "operators/kml/mla.hpp"
|
||
#include "operators/kml/mla_int8.hpp"
|
||
#endif
|
||
#include "operators/kml/moe.hpp"
|
||
static const bool _is_plain_ = true;
|
||
#else
|
||
static const bool _is_plain_ = false;
|
||
#endif
|
||
|
||
#if defined(__x86_64__) && defined(USE_AMX_AVX_KERNEL)
|
||
#include "operators/amx/awq-moe.hpp"
|
||
#include "operators/amx/bf16-moe.hpp" // Native BF16 MoE using CRTP pattern, with fallback for AVX512F
|
||
#include "operators/amx/fp4-moe.hpp" // MXFP4 MoE: FP4 E2M1 weights × BF16 activations
|
||
#include "operators/amx/fp8-moe.hpp" // FP8 MoE requires AVX512 BF16 support, with fallback for AVX512F+BW
|
||
#include "operators/amx/fp8-perchannel-moe.hpp" // FP8 Per-Channel MoE for GLM-4.7-FP8
|
||
#include "operators/amx/k2-moe.hpp"
|
||
#include "operators/amx/la/amx_kernels.hpp"
|
||
#include "operators/amx/moe.hpp"
|
||
#include "operators/amx/mxfp8-moe.hpp" // MXFP8 MoE: FP8 E4M3fn weights × BF16 activations (MiniMax M3)
|
||
#include "operators/amx/sft-k2-moe.hpp"
|
||
#include "operators/amx/sft_moe.hpp"
|
||
#include "operators/moe-sft-tp.hpp"
|
||
#endif
|
||
// AVX2 backends — always available on x86_64 (no AMX/AVX512 dependency)
|
||
#if defined(__x86_64__)
|
||
#include "operators/avx2/bf16-moe.hpp"
|
||
#include "operators/avx2/fp8-moe.hpp"
|
||
#include "operators/avx2/gptq_int4-moe.hpp"
|
||
#include "operators/avx2/gptq_int4_avxvnni-moe.hpp"
|
||
#include "operators/avx2/gptq_int4_avxvnni_packed-moe.hpp"
|
||
#include "operators/avx2/mxfp4-moe.hpp"
|
||
#include "operators/avx2/sft_mxfp4-moe.hpp"
|
||
#include "operators/avx2/sft_rawint4-moe.hpp"
|
||
#if !defined(USE_AMX_AVX_KERNEL)
|
||
#include "operators/moe-sft-tp.hpp"
|
||
#endif
|
||
#include "operators/avx2/mxfp8-moe.hpp"
|
||
#include "operators/avx2/rawint4-moe.hpp"
|
||
#include "operators/avx2/rawint4_avxvnni-moe.hpp"
|
||
#endif
|
||
#if defined(USE_SYCL)
|
||
#include "operators/sycl/gptq_int4_sycl-moe.hpp"
|
||
#endif
|
||
|
||
#include <pybind11/stl.h> // std::vector/std::pair/std::string conversions
|
||
|
||
#include <cstdint>
|
||
#include <memory>
|
||
#include <type_traits>
|
||
|
||
#include "fp8_layerwise_transport.hpp"
|
||
#include "operators/kvcache/kvcache.h"
|
||
#include "operators/llamafile/linear.h"
|
||
#include "operators/llamafile/mla.hpp"
|
||
#include "operators/llamafile/mlp.h"
|
||
#include "operators/llamafile/moe.hpp"
|
||
#include "pybind11/pybind11.h"
|
||
|
||
namespace py = pybind11;
|
||
using namespace pybind11::literals;
|
||
|
||
py::object to_float_ptr(uintptr_t input_ptr, int size, ggml_type type) {
|
||
if (type < 0 || type >= GGML_TYPE_COUNT) {
|
||
PyErr_SetString(PyExc_ValueError, "Invalid ggml_type");
|
||
throw py::error_already_set();
|
||
}
|
||
|
||
py::module torch = py::module::import("torch");
|
||
py::dict kwargs;
|
||
kwargs["dtype"] = torch.attr("float32");
|
||
py::object tensor = torch.attr("empty")(size, **kwargs);
|
||
|
||
uintptr_t output_ptr = tensor.attr("data_ptr")().cast<uintptr_t>();
|
||
float* output_float_ptr = reinterpret_cast<float*>(output_ptr);
|
||
|
||
try {
|
||
to_float(reinterpret_cast<void*>(input_ptr), output_float_ptr, size, type);
|
||
} catch (const std::exception& e) {
|
||
PyErr_SetString(PyExc_RuntimeError, e.what());
|
||
throw py::error_already_set();
|
||
}
|
||
|
||
return tensor;
|
||
}
|
||
|
||
py::object from_float_ptr(uintptr_t input_ptr, int size, ggml_type type) {
|
||
if (type < 0 || type >= GGML_TYPE_COUNT) {
|
||
PyErr_SetString(PyExc_ValueError, "Invalid ggml_type");
|
||
throw py::error_already_set();
|
||
}
|
||
|
||
py::module torch = py::module::import("torch");
|
||
|
||
size_t output_elem_bytes = ggml_type_size(type);
|
||
size_t output_elem_count = (size + ggml_blck_size(type) - 1) / ggml_blck_size(type);
|
||
size_t total_bytes = output_elem_count * output_elem_bytes;
|
||
|
||
py::dict kwargs;
|
||
kwargs["dtype"] = torch.attr("uint8");
|
||
py::object tensor = torch.attr("empty")(total_bytes, **kwargs);
|
||
|
||
uintptr_t output_ptr = tensor.attr("data_ptr")().cast<uintptr_t>();
|
||
void* output_void_ptr = reinterpret_cast<void*>(output_ptr);
|
||
|
||
try {
|
||
from_float(reinterpret_cast<float*>(input_ptr), output_void_ptr, size, type);
|
||
} catch (const std::exception& e) {
|
||
PyErr_SetString(PyExc_RuntimeError, e.what());
|
||
throw py::error_already_set();
|
||
}
|
||
|
||
return tensor;
|
||
}
|
||
|
||
template <typename T>
|
||
std::vector<std::vector<uintptr_t>> void_ptr_nested_to_uint(const std::vector<std::vector<T*>>& input) {
|
||
std::vector<std::vector<uintptr_t>> result;
|
||
for (const auto& row : input) {
|
||
std::vector<uintptr_t> new_row;
|
||
for (auto ptr : row) {
|
||
new_row.push_back(reinterpret_cast<uintptr_t>(ptr));
|
||
}
|
||
result.push_back(std::move(new_row));
|
||
}
|
||
return result;
|
||
}
|
||
|
||
template <typename T>
|
||
std::vector<std::vector<T*>> uint_to_void_ptr_nested(const std::vector<std::vector<uintptr_t>>& input) {
|
||
std::vector<std::vector<T*>> result;
|
||
for (const auto& row : input) {
|
||
std::vector<T*> new_row;
|
||
for (auto val : row) {
|
||
new_row.push_back(reinterpret_cast<T*>(val));
|
||
}
|
||
result.push_back(std::move(new_row));
|
||
}
|
||
return result;
|
||
}
|
||
|
||
#define DEF_PTR_PROPERTY(cls, name) \
|
||
def_property( \
|
||
#name, [](const cls& self) { return reinterpret_cast<uintptr_t>(self.name); }, \
|
||
[](cls& self, uintptr_t val) { self.name = reinterpret_cast<void*>(val); })
|
||
|
||
#define DEF_PTR_2D_PROPERTY(cls, name) \
|
||
def_property( \
|
||
#name, [](const cls& self) { return void_ptr_nested_to_uint<void>(self.name); }, \
|
||
[](cls& self, const std::vector<std::vector<uintptr_t>>& val) { \
|
||
self.name = uint_to_void_ptr_nested<void>(val); \
|
||
})
|
||
|
||
template <class T>
|
||
class MOEBindings {
|
||
public:
|
||
class WarmUpBindings {
|
||
public:
|
||
struct Args {
|
||
CPUInfer* cpuinfer;
|
||
TP_MOE<T>* moe;
|
||
};
|
||
static void inner(void* args) {
|
||
Args* args_ = (Args*)args;
|
||
args_->cpuinfer->enqueue(&TP_MOE<T>::warm_up, args_->moe);
|
||
}
|
||
static std::pair<intptr_t, intptr_t> cpuinfer_interface(std::shared_ptr<TP_MOE<T>> moe) {
|
||
Args* args = new Args{nullptr, moe.get()};
|
||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||
}
|
||
};
|
||
class LoadWeightsBindings {
|
||
public:
|
||
struct Args {
|
||
CPUInfer* cpuinfer;
|
||
TP_MOE<T>* moe;
|
||
};
|
||
static void inner(void* args) {
|
||
Args* args_ = (Args*)args;
|
||
args_->cpuinfer->enqueue(&TP_MOE<T>::load_weights, args_->moe);
|
||
}
|
||
static std::pair<intptr_t, intptr_t> cpuinfer_interface(std::shared_ptr<TP_MOE<T>> moe,
|
||
const uintptr_t physical_to_logical_map = 0) {
|
||
Args* args = new Args{nullptr, moe.get()};
|
||
if (physical_to_logical_map) {
|
||
// printf("debug physical_to_logical_map in arg:%lu\n", physical_to_logical_map);
|
||
moe->config.physical_to_logical_map = reinterpret_cast<void*>(physical_to_logical_map);
|
||
// printf("moe ptr:%p,confirm: moe->config.physical_to_logical_map:%lu\n", reinterpret_cast<void*>(moe.get()),
|
||
// reinterpret_cast<uintptr_t>(moe->config.physical_to_logical_map));
|
||
}
|
||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||
}
|
||
static std::pair<intptr_t, intptr_t> cpuinfer_interface(std::shared_ptr<TP_MOE<T>> moe) {
|
||
return cpuinfer_interface(moe, 0);
|
||
}
|
||
};
|
||
class ForwardBindings {
|
||
public:
|
||
struct Args {
|
||
CPUInfer* cpuinfer;
|
||
TP_MOE<T>* moe;
|
||
intptr_t qlen;
|
||
int k;
|
||
intptr_t expert_ids;
|
||
intptr_t weights;
|
||
intptr_t input;
|
||
intptr_t output;
|
||
bool incremental;
|
||
};
|
||
static void inner(void* args) {
|
||
Args* args_ = (Args*)args;
|
||
args_->cpuinfer->enqueue(&TP_MOE<T>::forward_binding, args_->moe, args_->qlen, args_->k, args_->expert_ids,
|
||
args_->weights, args_->input, args_->output, args_->incremental);
|
||
}
|
||
static std::pair<intptr_t, intptr_t> cpuinfer_interface(std::shared_ptr<TP_MOE<T>> moe, intptr_t qlen, int k,
|
||
intptr_t expert_ids, intptr_t weights, intptr_t input,
|
||
intptr_t output, bool incremental = false) {
|
||
Args* args = new Args{nullptr, moe.get(), qlen, k, expert_ids, weights, input, output, incremental};
|
||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||
}
|
||
static std::pair<intptr_t, intptr_t> cpuinfer_interface(std::shared_ptr<TP_MOE<T>> moe, intptr_t qlen, int k,
|
||
intptr_t expert_ids, intptr_t weights, intptr_t input,
|
||
intptr_t output) {
|
||
return cpuinfer_interface(moe, qlen, k, expert_ids, weights, input, output, false);
|
||
}
|
||
};
|
||
};
|
||
|
||
#if defined(__x86_64__)
|
||
template <class T>
|
||
class MOESFTBindings {
|
||
public:
|
||
class WarmUpBindings {
|
||
public:
|
||
struct Args {
|
||
CPUInfer* cpuinfer;
|
||
TP_MOE_SFT<T>* moe;
|
||
};
|
||
static void inner(void* args) {
|
||
Args* args_ = (Args*)args;
|
||
args_->cpuinfer->enqueue(&TP_MOE_SFT<T>::warm_up, args_->moe);
|
||
}
|
||
static std::pair<intptr_t, intptr_t> cpuinfer_interface(std::shared_ptr<TP_MOE_SFT<T>> moe) {
|
||
Args* args = new Args{nullptr, moe.get()};
|
||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||
}
|
||
};
|
||
|
||
class LoadWeightsBindings {
|
||
public:
|
||
struct Args {
|
||
CPUInfer* cpuinfer;
|
||
TP_MOE_SFT<T>* moe;
|
||
};
|
||
static void inner(void* args) {
|
||
Args* args_ = (Args*)args;
|
||
args_->cpuinfer->enqueue(&TP_MOE_SFT<T>::load_weights, args_->moe);
|
||
}
|
||
static std::pair<intptr_t, intptr_t> cpuinfer_interface(std::shared_ptr<TP_MOE_SFT<T>> moe) {
|
||
Args* args = new Args{nullptr, moe.get()};
|
||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||
}
|
||
};
|
||
|
||
class ForwardSFTBindings {
|
||
public:
|
||
struct Args {
|
||
CPUInfer* cpuinfer;
|
||
TP_MOE_SFT<T>* moe;
|
||
intptr_t qlen;
|
||
int k;
|
||
intptr_t expert_ids;
|
||
intptr_t weights;
|
||
intptr_t input;
|
||
intptr_t output;
|
||
bool save_for_backward;
|
||
};
|
||
static void inner(void* args) {
|
||
Args* args_ = (Args*)args;
|
||
args_->cpuinfer->enqueue(&TP_MOE_SFT<T>::forward_sft_binding, args_->moe, args_->qlen, args_->k,
|
||
args_->expert_ids, args_->weights, args_->input, args_->output,
|
||
args_->save_for_backward);
|
||
}
|
||
static std::pair<intptr_t, intptr_t> cpuinfer_interface(std::shared_ptr<TP_MOE_SFT<T>> moe, intptr_t qlen, int k,
|
||
intptr_t expert_ids, intptr_t weights, intptr_t input,
|
||
intptr_t output, bool save_for_backward) {
|
||
Args* args = new Args{nullptr, moe.get(), qlen, k, expert_ids, weights, input, output, save_for_backward};
|
||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||
}
|
||
};
|
||
|
||
class BackwardBindings {
|
||
public:
|
||
struct Args {
|
||
CPUInfer* cpuinfer;
|
||
TP_MOE_SFT<T>* moe;
|
||
intptr_t grad_output;
|
||
intptr_t grad_input;
|
||
intptr_t grad_gate_lora_a;
|
||
intptr_t grad_gate_lora_b;
|
||
intptr_t grad_up_lora_a;
|
||
intptr_t grad_up_lora_b;
|
||
intptr_t grad_down_lora_a;
|
||
intptr_t grad_down_lora_b;
|
||
intptr_t grad_weights;
|
||
intptr_t grad_gate_proj;
|
||
intptr_t grad_up_proj;
|
||
intptr_t grad_down_proj;
|
||
bool accumulate_optimizer_grads;
|
||
float optimizer_grad_scale;
|
||
};
|
||
static void inner(void* args) {
|
||
Args* args_ = (Args*)args;
|
||
args_->cpuinfer->enqueue(&TP_MOE_SFT<T>::backward_binding, args_->moe, args_->grad_output, args_->grad_input,
|
||
args_->grad_gate_lora_a, args_->grad_gate_lora_b, args_->grad_up_lora_a,
|
||
args_->grad_up_lora_b, args_->grad_down_lora_a, args_->grad_down_lora_b,
|
||
args_->grad_weights, args_->grad_gate_proj, args_->grad_up_proj, args_->grad_down_proj,
|
||
args_->accumulate_optimizer_grads, args_->optimizer_grad_scale);
|
||
}
|
||
static std::pair<intptr_t, intptr_t> cpuinfer_interface(
|
||
std::shared_ptr<TP_MOE_SFT<T>> moe, intptr_t grad_output, intptr_t grad_input, intptr_t grad_gate_lora_a,
|
||
intptr_t grad_gate_lora_b, intptr_t grad_up_lora_a, intptr_t grad_up_lora_b, intptr_t grad_down_lora_a,
|
||
intptr_t grad_down_lora_b, intptr_t grad_weights, intptr_t grad_gate_proj, intptr_t grad_up_proj,
|
||
intptr_t grad_down_proj, bool accumulate_optimizer_grads = false, float optimizer_grad_scale = 1.0f) {
|
||
Args* args = new Args{nullptr,
|
||
moe.get(),
|
||
grad_output,
|
||
grad_input,
|
||
grad_gate_lora_a,
|
||
grad_gate_lora_b,
|
||
grad_up_lora_a,
|
||
grad_up_lora_b,
|
||
grad_down_lora_a,
|
||
grad_down_lora_b,
|
||
grad_weights,
|
||
grad_gate_proj,
|
||
grad_up_proj,
|
||
grad_down_proj,
|
||
accumulate_optimizer_grads,
|
||
optimizer_grad_scale};
|
||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||
}
|
||
};
|
||
|
||
class UpdateLoRAWeightsBindings {
|
||
public:
|
||
struct Args {
|
||
CPUInfer* cpuinfer;
|
||
TP_MOE_SFT<T>* moe;
|
||
intptr_t gate_lora_a;
|
||
intptr_t gate_lora_b;
|
||
intptr_t up_lora_a;
|
||
intptr_t up_lora_b;
|
||
intptr_t down_lora_a;
|
||
intptr_t down_lora_b;
|
||
};
|
||
static void inner(void* args) {
|
||
// Debug code for Bug #18 - commented out after fix verified
|
||
// printf("[DEBUG UpdateLoRAWeightsBindings::inner] called\n");
|
||
Args* args_ = (Args*)args;
|
||
// printf(" moe=%p, gate_lora_a=%p, gate_lora_b=%p\n", (void*)args_->moe, (void*)args_->gate_lora_a,
|
||
// (void*)args_->gate_lora_b); printf(" up_lora_a=%p, up_lora_b=%p\n", (void*)args_->up_lora_a,
|
||
// (void*)args_->up_lora_b); printf(" down_lora_a=%p, down_lora_b=%p\n", (void*)args_->down_lora_a,
|
||
// (void*)args_->down_lora_b);
|
||
args_->cpuinfer->enqueue(&TP_MOE_SFT<T>::update_lora_weights_binding, args_->moe, args_->gate_lora_a,
|
||
args_->gate_lora_b, args_->up_lora_a, args_->up_lora_b, args_->down_lora_a,
|
||
args_->down_lora_b);
|
||
// printf("[DEBUG UpdateLoRAWeightsBindings::inner] enqueue done\n");
|
||
}
|
||
static std::pair<intptr_t, intptr_t> cpuinfer_interface(std::shared_ptr<TP_MOE_SFT<T>> moe, intptr_t gate_lora_a,
|
||
intptr_t gate_lora_b, intptr_t up_lora_a,
|
||
intptr_t up_lora_b, intptr_t down_lora_a,
|
||
intptr_t down_lora_b) {
|
||
Args* args =
|
||
new Args{nullptr, moe.get(), gate_lora_a, gate_lora_b, up_lora_a, up_lora_b, down_lora_a, down_lora_b};
|
||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||
}
|
||
};
|
||
};
|
||
|
||
template <typename MoeSftTP>
|
||
void bind_moe_sft_module(py::module_& moe_module, const char* name) {
|
||
using MoeClass = TP_MOE_SFT<MoeSftTP>;
|
||
using MoeBindings = MOESFTBindings<MoeSftTP>;
|
||
|
||
py::class_<MoeClass, MoE_Interface, std::shared_ptr<MoeClass>>(moe_module, name)
|
||
.def(py::init<MOESFTConfig>())
|
||
.def("warm_up_task", &MoeBindings::WarmUpBindings::cpuinfer_interface)
|
||
.def("load_weights_task", &MoeBindings::LoadWeightsBindings::cpuinfer_interface)
|
||
.def("forward_sft_task", &MoeBindings::ForwardSFTBindings::cpuinfer_interface)
|
||
.def("backward_task", &MoeBindings::BackwardBindings::cpuinfer_interface, py::arg("grad_output"),
|
||
py::arg("grad_input"), py::arg("grad_gate_lora_a"), py::arg("grad_gate_lora_b"), py::arg("grad_up_lora_a"),
|
||
py::arg("grad_up_lora_b"), py::arg("grad_down_lora_a"), py::arg("grad_down_lora_b"), py::arg("grad_weights"),
|
||
py::arg("grad_gate_proj"), py::arg("grad_up_proj"), py::arg("grad_down_proj"),
|
||
py::arg("accumulate_optimizer_grads") = false, py::arg("optimizer_grad_scale") = 1.0f)
|
||
.def("update_lora_weights_task", &MoeBindings::UpdateLoRAWeightsBindings::cpuinfer_interface)
|
||
.def("warm_up", &MoeClass::warm_up)
|
||
.def("load_weights", &MoeClass::load_weights)
|
||
.def("forward_sft", &MoeClass::forward_sft_binding)
|
||
.def("backward", &MoeClass::backward_binding, py::arg("grad_output"), py::arg("grad_input"),
|
||
py::arg("grad_gate_lora_a"), py::arg("grad_gate_lora_b"), py::arg("grad_up_lora_a"),
|
||
py::arg("grad_up_lora_b"), py::arg("grad_down_lora_a"), py::arg("grad_down_lora_b"), py::arg("grad_weights"),
|
||
py::arg("grad_gate_proj"), py::arg("grad_up_proj"), py::arg("grad_down_proj"),
|
||
py::arg("accumulate_optimizer_grads") = false, py::arg("optimizer_grad_scale") = 1.0f)
|
||
.def("update_lora_weights", &MoeClass::update_lora_weights_binding)
|
||
.def("prepare_and_save_bwd",
|
||
[](MoeClass& self, intptr_t gate, intptr_t up, intptr_t down, const std::string& path) {
|
||
self.prepare_and_save_bwd((void*)gate, (void*)up, (void*)down, path);
|
||
})
|
||
.def("submit_backward_repack", &MoeClass::submit_backward_repack)
|
||
.def("wait_backward_repack", &MoeClass::wait_backward_repack)
|
||
.def("get_profile_stats", &MoeClass::get_profile_stats, py::arg("reset") = false)
|
||
.def("reset_profile_stats", &MoeClass::reset_profile_stats)
|
||
// Update base weight BF16 pointers for reload_base_weights (full mode training)
|
||
// After calling this, call load_weights_task() to re-quantize BF16->AMX
|
||
.def("set_base_weight_pointers", [](MoeClass& self, intptr_t gate, intptr_t up, intptr_t down) {
|
||
self.set_base_weight_pointers((void*)gate, (void*)up, (void*)down);
|
||
});
|
||
}
|
||
#endif // defined(__x86_64__)
|
||
|
||
template <typename MoeTP>
|
||
void bind_moe_module(py::module_& moe_module, const char* name) {
|
||
using MoeClass = TP_MOE<MoeTP>;
|
||
using MoeBindings = MOEBindings<MoeTP>;
|
||
|
||
auto moe_cls = py::class_<MoeClass, MoE_Interface, std::shared_ptr<MoeClass>>(moe_module, name);
|
||
|
||
moe_cls.def(py::init<GeneralMOEConfig>())
|
||
.def("warm_up_task", &MoeBindings::WarmUpBindings::cpuinfer_interface)
|
||
.def("load_weights_task",
|
||
py::overload_cast<std::shared_ptr<MoeClass>>(&MoeBindings::LoadWeightsBindings::cpuinfer_interface))
|
||
.def("load_weights_task",
|
||
py::overload_cast<std::shared_ptr<MoeClass>, const uintptr_t>(
|
||
&MoeBindings::LoadWeightsBindings::cpuinfer_interface),
|
||
py::arg("physical_to_logical_map"))
|
||
// .def("forward_task", &MoeBindings::ForwardBindings::cpuinfer_interface)
|
||
.def("forward_task",
|
||
py::overload_cast<std::shared_ptr<MoeClass>, intptr_t, int, intptr_t, intptr_t, intptr_t, intptr_t>(
|
||
&MoeBindings::ForwardBindings::cpuinfer_interface))
|
||
.def("forward_task",
|
||
py::overload_cast<std::shared_ptr<MoeClass>, intptr_t, int, intptr_t, intptr_t, intptr_t, intptr_t, bool>(
|
||
&MoeBindings::ForwardBindings::cpuinfer_interface))
|
||
.def("warm_up", &MoeClass::warm_up)
|
||
.def("load_weights", &MoeClass::load_weights)
|
||
.def("forward", &MoeClass::forward_binding);
|
||
|
||
// Bind write_weight_scale_to_buffer_task for MoE types that support it
|
||
// Uses SFINAE to detect if MoeClass has write_weight_scale_to_buffer method
|
||
if constexpr (requires { &MoeClass::write_weight_scale_to_buffer; }) {
|
||
struct WriteWeightScaleToBufferBindings {
|
||
struct Args {
|
||
CPUInfer* cpuinfer;
|
||
MoeClass* moe;
|
||
int gpu_tp_count;
|
||
int expert_id;
|
||
std::vector<uintptr_t> w13_weight_ptrs;
|
||
std::vector<uintptr_t> w13_scale_ptrs;
|
||
std::vector<uintptr_t> w2_weight_ptrs;
|
||
std::vector<uintptr_t> w2_scale_ptrs;
|
||
};
|
||
|
||
static void inner(void* args) {
|
||
Args* args_ = (Args*)args;
|
||
args_->cpuinfer->enqueue(&MoeClass::write_weight_scale_to_buffer, args_->moe, args_->gpu_tp_count,
|
||
args_->expert_id, args_->w13_weight_ptrs, args_->w13_scale_ptrs, args_->w2_weight_ptrs,
|
||
args_->w2_scale_ptrs);
|
||
}
|
||
|
||
static std::pair<intptr_t, intptr_t> cpuinfer_interface(std::shared_ptr<MoeClass> moe, int gpu_tp_count,
|
||
int expert_id, py::list w13_weight_ptrs,
|
||
py::list w13_scale_ptrs, py::list w2_weight_ptrs,
|
||
py::list w2_scale_ptrs) {
|
||
// Convert Python lists to std::vector<uintptr_t>
|
||
std::vector<uintptr_t> w13_weight_vec, w13_scale_vec, w2_weight_vec, w2_scale_vec;
|
||
|
||
for (auto item : w13_weight_ptrs) w13_weight_vec.push_back(py::cast<uintptr_t>(item));
|
||
for (auto item : w13_scale_ptrs) w13_scale_vec.push_back(py::cast<uintptr_t>(item));
|
||
for (auto item : w2_weight_ptrs) w2_weight_vec.push_back(py::cast<uintptr_t>(item));
|
||
for (auto item : w2_scale_ptrs) w2_scale_vec.push_back(py::cast<uintptr_t>(item));
|
||
|
||
Args* args = new Args{nullptr, moe.get(), gpu_tp_count, expert_id,
|
||
w13_weight_vec, w13_scale_vec, w2_weight_vec, w2_scale_vec};
|
||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||
}
|
||
};
|
||
|
||
moe_cls.def("write_weight_scale_to_buffer_task", &WriteWeightScaleToBufferBindings::cpuinfer_interface,
|
||
py::arg("gpu_tp_count"), py::arg("expert_id"), py::arg("w13_weight_ptrs"), py::arg("w13_scale_ptrs"),
|
||
py::arg("w2_weight_ptrs"), py::arg("w2_scale_ptrs"));
|
||
|
||
moe_cls.def(
|
||
"run_layerwise_fp8_batch",
|
||
[](std::shared_ptr<MoeClass> moe, const std::shared_ptr<kt::layerwise::FP8LayerwiseTransport>& transport,
|
||
std::uint64_t epoch, std::int64_t layer_id, int expert_count) {
|
||
if (!transport) throw std::invalid_argument("FP8 layerwise transport is null");
|
||
const int tp_size = transport->tp_size();
|
||
transport->run_producer(
|
||
epoch, layer_id, expert_count,
|
||
[moe, tp_size](int expert_id, const std::vector<std::uintptr_t>& w13_weight_ptrs,
|
||
const std::vector<std::uintptr_t>& w13_scale_ptrs,
|
||
const std::vector<std::uintptr_t>& w2_weight_ptrs,
|
||
const std::vector<std::uintptr_t>& w2_scale_ptrs) {
|
||
moe->write_weight_scale_to_buffer(tp_size, expert_id, w13_weight_ptrs, w13_scale_ptrs, w2_weight_ptrs,
|
||
w2_scale_ptrs);
|
||
});
|
||
},
|
||
py::arg("transport"), py::arg("epoch"), py::arg("layer_id"), py::arg("expert_count"),
|
||
py::call_guard<py::gil_scoped_release>());
|
||
}
|
||
}
|
||
|
||
PYBIND11_MODULE(kt_kernel_ext, m) {
|
||
m.def("_bf16_sumsq_ptr", &bf16_sumsq_ptr, py::arg("data_ptr"), py::arg("numel"),
|
||
py::call_guard<py::gil_scoped_release>());
|
||
#if defined(HAVE_AMX)
|
||
m.attr("__cpu_variant__") = "amx";
|
||
m.attr("__int8_kernel__") = "amx-int8";
|
||
#elif defined(__AVX512BF16__) && defined(__AVX512VNNI__)
|
||
m.attr("__cpu_variant__") = "avx512_bf16";
|
||
m.attr("__int8_kernel__") = amx::int8_vnni_backend_name();
|
||
#elif defined(__AVX512VBMI__)
|
||
m.attr("__cpu_variant__") = "avx512_vbmi";
|
||
m.attr("__int8_kernel__") = "unsupported";
|
||
#elif defined(__AVX512VNNI__)
|
||
m.attr("__cpu_variant__") = "avx512_vnni";
|
||
m.attr("__int8_kernel__") = "unsupported";
|
||
#elif defined(__AVX512F__)
|
||
m.attr("__cpu_variant__") = "avx512_base";
|
||
m.attr("__int8_kernel__") = "unsupported";
|
||
#else
|
||
m.attr("__cpu_variant__") = "avx2";
|
||
m.attr("__int8_kernel__") = "unsupported";
|
||
#endif
|
||
m.attr("__int8_weight_layout__") = "kt-int8-n32-k64-vnni-v1";
|
||
#if defined(USE_AMX_AVX_KERNEL) && defined(__AVX512BF16__)
|
||
m.attr("__rawint4_kernel__") = "amx-int4-kgroup-g32";
|
||
#elif defined(__x86_64__) && !defined(USE_AMX_AVX_KERNEL)
|
||
m.attr("__rawint4_kernel__") = "avx2-int4-kgroup-g32";
|
||
#else
|
||
m.attr("__rawint4_kernel__") = "unsupported";
|
||
#endif
|
||
m.attr("__rawint4_weight_layout__") = "compressed-tensors-rawint4-g32-v1";
|
||
#if defined(__AVX512BF16__) && defined(__AVX512VBMI__)
|
||
m.attr("__fp8_kernel__") = "avx512-fp8-decode-bf16";
|
||
#else
|
||
m.attr("__fp8_kernel__") = "unsupported";
|
||
#endif
|
||
m.attr("__fp8_weight_layout__") = "block-e4m3-128x128";
|
||
m.attr("FP8_LAYERWISE_CONTROL_BYTES") = kt::layerwise::kFP8LayerwiseControlBytes;
|
||
m.attr("FP8_LAYERWISE_MAX_TP_SIZE") = kt::layerwise::kFP8LayerwiseMaxTPSize;
|
||
|
||
m.def("initialize_fp8_layerwise_control", &kt::layerwise::initialize_fp8_layerwise_control,
|
||
py::arg("control_ptr"), py::arg("control_size"), py::arg("tp_size"));
|
||
|
||
py::class_<kt::layerwise::FP8LayerwiseTransport, std::shared_ptr<kt::layerwise::FP8LayerwiseTransport>>(
|
||
m, "FP8LayerwiseTransport")
|
||
.def(py::init<std::uintptr_t, std::size_t, int, int, int, const std::vector<std::uintptr_t>&,
|
||
const std::vector<std::uintptr_t>&, const std::vector<std::uintptr_t>&,
|
||
const std::vector<std::size_t>&, int, std::uint64_t>(),
|
||
py::arg("control_ptr"), py::arg("control_size"), py::arg("rank"), py::arg("tp_size"),
|
||
py::arg("cuda_device"), py::arg("local_host_ptrs"), py::arg("local_gpu_ptrs"),
|
||
py::arg("all_rank_host_ptrs"), py::arg("expert_nbytes"), py::arg("num_experts"),
|
||
py::arg("timeout_ms") = 60000)
|
||
.def("join", &kt::layerwise::FP8LayerwiseTransport::join, py::arg("epoch"), py::arg("layer_id"),
|
||
py::arg("expert_count"), py::call_guard<py::gil_scoped_release>())
|
||
.def(
|
||
"wait",
|
||
[](kt::layerwise::FP8LayerwiseTransport& transport, std::uint64_t epoch) {
|
||
kt::layerwise::FP8LayerwiseStats stats;
|
||
{
|
||
py::gil_scoped_release release;
|
||
stats = transport.wait(epoch);
|
||
}
|
||
py::dict result;
|
||
result["epoch"] = stats.epoch;
|
||
result["layer_id"] = stats.layer_id;
|
||
result["expert_count"] = stats.expert_count;
|
||
result["rank"] = stats.rank;
|
||
result["writer_ms"] = stats.writer_ms;
|
||
result["slot_wait_ms"] = stats.slot_wait_ms;
|
||
result["h2d_ms"] = stats.h2d_ms;
|
||
result["total_ms"] = stats.total_ms;
|
||
result["bytes"] = stats.bytes;
|
||
result["poisoned"] = stats.poisoned;
|
||
result["error_code"] = stats.error_code;
|
||
result["error_rank"] = stats.error_rank;
|
||
result["error_message"] = stats.error_message;
|
||
return result;
|
||
},
|
||
py::arg("epoch"))
|
||
.def("close", &kt::layerwise::FP8LayerwiseTransport::close, py::call_guard<py::gil_scoped_release>())
|
||
.def_property_readonly("rank", &kt::layerwise::FP8LayerwiseTransport::rank)
|
||
.def_property_readonly("tp_size", &kt::layerwise::FP8LayerwiseTransport::tp_size)
|
||
.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<aclrtStream>(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_<WorkerPool>(m, "WorkerPool").def(py::init<int>());
|
||
py::class_<WorkerPoolConfig>(m, "WorkerPoolConfig")
|
||
.def(py::init<>())
|
||
.def_readwrite("subpool_count", &WorkerPoolConfig::subpool_count)
|
||
.def_readwrite("subpool_numa_map", &WorkerPoolConfig::subpool_numa_map)
|
||
.def_readwrite("subpool_thread_count", &WorkerPoolConfig::subpool_thread_count);
|
||
|
||
py::class_<CPUInfer>(m, "CPUInfer")
|
||
.def(py::init<int>())
|
||
.def(py::init<WorkerPoolConfig>())
|
||
.def("submit", &CPUInfer::submit)
|
||
.def("sync", &CPUInfer::sync, py::arg("allow_n_pending") = 0)
|
||
.def_readwrite("backend_", &CPUInfer::backend_)
|
||
#ifndef KTRANSFORMERS_CPU_ONLY
|
||
.def("sync_with_cuda_stream", &CPUInfer::sync_with_cuda_stream, py::arg("user_cuda_stream"),
|
||
py::arg("allow_n_pending") = 0)
|
||
.def("submit_with_cuda_stream", &CPUInfer::submit_with_cuda_stream)
|
||
#endif
|
||
;
|
||
|
||
auto linear_module = m.def_submodule("linear");
|
||
py::class_<LinearConfig>(linear_module, "LinearConfig")
|
||
.def(py::init([](int hidden_size, int intermediate_size, int stride, int group_max_len, intptr_t proj,
|
||
int proj_type, int hidden_type) {
|
||
return LinearConfig(hidden_size, intermediate_size, stride, group_max_len, (void*)proj, (ggml_type)proj_type,
|
||
(ggml_type)hidden_type);
|
||
}));
|
||
// py::class_<Linear>(linear_module, "Linear")
|
||
// .def(py::init<LinearConfig>())
|
||
// .def("warm_up", &LinearBindings::WarmUpBindings::cpuinfer_interface)
|
||
// .def("forward", &LinearBindings::ForwardBindings::cpuinfer_interface);
|
||
|
||
auto mlp_module = m.def_submodule("mlp");
|
||
py::class_<MLPConfig>(mlp_module, "MLPConfig")
|
||
.def(py::init([](int hidden_size, int intermediate_size, int stride, int group_max_len, intptr_t gate_proj,
|
||
intptr_t up_proj, intptr_t down_proj, int gate_type, int up_type, int down_type,
|
||
int hidden_type) {
|
||
return MLPConfig(hidden_size, intermediate_size, stride, group_max_len, (void*)gate_proj, (void*)up_proj,
|
||
(void*)down_proj, (ggml_type)gate_type, (ggml_type)up_type, (ggml_type)down_type,
|
||
(ggml_type)hidden_type);
|
||
}));
|
||
// py::class_<MLP>(mlp_module, "MLP")
|
||
// .def(py::init<MLPConfig>())
|
||
// .def("warm_up", &MLPBindings::WarmUpBindings::cpuinfer_interface)
|
||
// .def("forward", &MLPBindings::ForwardBindings::cpuinfer_interface);
|
||
|
||
py::class_<GeneralConfig>(m, "GeneralConfig")
|
||
.def(py::init<>())
|
||
.def_readwrite("vocab_size", &GeneralConfig::vocab_size)
|
||
.def_readwrite("hidden_size", &GeneralConfig::hidden_size)
|
||
.def_readwrite("num_experts_per_tok", &GeneralConfig::num_experts_per_tok)
|
||
.def_readwrite("n_routed_experts", &GeneralConfig::n_routed_experts)
|
||
.def_readwrite("n_shared_experts", &GeneralConfig::n_shared_experts)
|
||
.def_readwrite("max_qlen", &GeneralConfig::max_qlen)
|
||
.DEF_PTR_PROPERTY(GeneralConfig, lm_heads_ptr)
|
||
.def_readwrite("lm_heads_type", &GeneralConfig::lm_heads_type)
|
||
.DEF_PTR_PROPERTY(GeneralConfig, norm_weights_ptr)
|
||
.def_readwrite("norm_weights_type", &GeneralConfig::norm_weights_type)
|
||
.DEF_PTR_PROPERTY(GeneralConfig, token_embd_ptr)
|
||
.def_readwrite("token_embd_type", &GeneralConfig::token_embd_type)
|
||
.def_readwrite("pool", &GeneralConfig::pool);
|
||
#if defined(__aarch64__) && defined(CPU_USE_KML) && defined(KTRANSFORMERS_CPU_MLA)
|
||
py::class_<DeepseekV3ForCausalLM, std::shared_ptr<DeepseekV3ForCausalLM>>(m, "DeepseekV3ForCausalLM")
|
||
.def(py::init([](GeneralConfig config) { return std::make_shared<DeepseekV3ForCausalLM>(config); }))
|
||
.def_readwrite("model", &DeepseekV3ForCausalLM::model)
|
||
.def("forward", &DeepseekV3ForCausalLM::forward_binding);
|
||
|
||
py::class_<DeepseekV3Model, std::shared_ptr<DeepseekV3Model>>(m, "DeepseekV3Model")
|
||
.def(py::init([](GeneralConfig config) { return std::make_shared<DeepseekV3Model>(config); }))
|
||
.def_readwrite("layers", &DeepseekV3Model::layers);
|
||
|
||
py::class_<DeepseekV3DecoderLayer, std::shared_ptr<DeepseekV3DecoderLayer>>(m, "DeepseekV3DecoderLayer")
|
||
.def(py::init([](GeneralConfig config, size_t layer_idx) {
|
||
return std::make_shared<DeepseekV3DecoderLayer>(config, layer_idx);
|
||
}))
|
||
.def("load_norm", &DeepseekV3DecoderLayer::load_norm_binding)
|
||
.def_readwrite("self_attn", &DeepseekV3DecoderLayer::self_attn)
|
||
.def_readwrite("gate", &DeepseekV3DecoderLayer::gate)
|
||
.def_readwrite("ffn", &DeepseekV3DecoderLayer::ffn);
|
||
#endif
|
||
auto mla_module = m.def_submodule("mla");
|
||
py::class_<GeneralMLAConfig>(mla_module, "MLAConfig")
|
||
.def(py::init([](size_t hidden_size, size_t q_lora_rank, size_t num_heads, size_t nope_size, size_t rope_size,
|
||
size_t kv_lora_rank) {
|
||
return GeneralMLAConfig(hidden_size, q_lora_rank, num_heads, nope_size, rope_size, kv_lora_rank);
|
||
}))
|
||
.def_readwrite("layer_idx", &GeneralMLAConfig::layer_idx)
|
||
.def_readwrite("pool", &GeneralMLAConfig::pool)
|
||
.def_readwrite("token_count_in_page", &GeneralMLAConfig::token_count_in_page)
|
||
.def_readwrite("max_qlen", &GeneralMLAConfig::max_qlen)
|
||
.def_readwrite("max_kvlen", &GeneralMLAConfig::max_kvlen)
|
||
|
||
.def_readwrite("max_position_embeddings", &GeneralMLAConfig::max_position_embeddings)
|
||
.def_readwrite("rope_scaling_factor", &GeneralMLAConfig::rope_scaling_factor)
|
||
.def_readwrite("rope_theta", &GeneralMLAConfig::rope_theta)
|
||
.def_readwrite("rope_scaling_beta_fast", &GeneralMLAConfig::rope_scaling_beta_fast)
|
||
.def_readwrite("rope_scaling_beta_slow", &GeneralMLAConfig::rope_scaling_beta_slow)
|
||
.def_readwrite("rope_scaling_mscale", &GeneralMLAConfig::rope_scaling_mscale)
|
||
.def_readwrite("rope_scaling_mscale_all_dim", &GeneralMLAConfig::rope_scaling_mscale_all_dim)
|
||
.def_readwrite("rope_scaling_original_max_position_embeddings",
|
||
&GeneralMLAConfig::rope_scaling_original_max_position_embeddings)
|
||
|
||
.DEF_PTR_PROPERTY(GeneralMLAConfig, q_a_proj)
|
||
.DEF_PTR_PROPERTY(GeneralMLAConfig, q_a_norm)
|
||
.DEF_PTR_PROPERTY(GeneralMLAConfig, q_b_proj)
|
||
.DEF_PTR_PROPERTY(GeneralMLAConfig, kv_a_proj_with_mqa)
|
||
.DEF_PTR_PROPERTY(GeneralMLAConfig, kv_a_norm)
|
||
.DEF_PTR_PROPERTY(GeneralMLAConfig, kv_b_proj)
|
||
.DEF_PTR_PROPERTY(GeneralMLAConfig, o_proj)
|
||
|
||
.def_readwrite("q_a_proj_type", &GeneralMLAConfig::q_a_proj_type)
|
||
.def_readwrite("q_a_norm_type", &GeneralMLAConfig::q_a_norm_type)
|
||
.def_readwrite("q_b_proj_type", &GeneralMLAConfig::q_b_proj_type)
|
||
.def_readwrite("kv_a_proj_with_mqa_type", &GeneralMLAConfig::kv_a_proj_with_mqa_type)
|
||
.def_readwrite("kv_a_norm_type", &GeneralMLAConfig::kv_a_norm_type)
|
||
.def_readwrite("kv_b_proj_type", &GeneralMLAConfig::kv_b_proj_type)
|
||
.def_readwrite("w_o_type", &GeneralMLAConfig::w_o_type)
|
||
.def_readwrite("page_count", &GeneralMLAConfig::page_count)
|
||
|
||
;
|
||
py::class_<MLA_Interface, std::shared_ptr<MLA_Interface>>(mla_module, "MLA_Interface");
|
||
#if defined(__aarch64__) && defined(CPU_USE_KML) && defined(KTRANSFORMERS_CPU_MLA)
|
||
py::class_<TP_MLA<KML_MLA_TP<float16_t>>, MLA_Interface, std::shared_ptr<TP_MLA<KML_MLA_TP<float16_t>>>>(mla_module,
|
||
"MLA_F16")
|
||
.def(py::init<GeneralMLAConfig>())
|
||
.def("load_weights", &TP_MLA<KML_MLA_TP<float16_t>>::load_weights)
|
||
.def("forward",
|
||
[](TP_MLA<KML_MLA_TP<float16_t>>& op, std::vector<int> qlens, std::vector<std::vector<int>> page_tables,
|
||
std::vector<int> kvlens, intptr_t input,
|
||
intptr_t output) { op.forward(qlens, page_tables, kvlens, (const void*)input, (void*)output); })
|
||
.def("set_local_pages", &TP_MLA<KML_MLA_TP<float16_t>>::set_local_pages)
|
||
.def("set_pages", [](TP_MLA<KML_MLA_TP<float16_t>>& op, std::vector<std::vector<intptr_t>> nope_pages,
|
||
std::vector<std::vector<intptr_t>> rope_pages) {
|
||
std::vector<std::vector<void*>> nope_pages_ptr;
|
||
std::vector<std::vector<void*>> rope_pages_ptr;
|
||
op.set_pages(nope_pages_ptr, rope_pages_ptr);
|
||
});
|
||
|
||
py::class_<TP_MLA<KML_MLA_TP<float>>, MLA_Interface, std::shared_ptr<TP_MLA<KML_MLA_TP<float>>>>(mla_module,
|
||
"MLA_F32")
|
||
.def(py::init<GeneralMLAConfig>())
|
||
.def("load_weights", &TP_MLA<KML_MLA_TP<float>>::load_weights)
|
||
.def("forward",
|
||
[](TP_MLA<KML_MLA_TP<float>>& op, std::vector<int> qlens, std::vector<std::vector<int>> page_tables,
|
||
std::vector<int> kvlens, intptr_t input,
|
||
intptr_t output) { op.forward(qlens, page_tables, kvlens, (const void*)input, (void*)output); })
|
||
.def("set_local_pages", &TP_MLA<KML_MLA_TP<float>>::set_local_pages)
|
||
.def("set_pages", [](TP_MLA<KML_MLA_TP<float>>& op, std::vector<std::vector<intptr_t>> nope_pages,
|
||
std::vector<std::vector<intptr_t>> rope_pages) {
|
||
std::vector<std::vector<void*>> nope_pages_ptr;
|
||
std::vector<std::vector<void*>> rope_pages_ptr;
|
||
op.set_pages(nope_pages_ptr, rope_pages_ptr);
|
||
});
|
||
py::class_<TP_MLA<KML_MLA_TP_QUAN<float>>, MLA_Interface, std::shared_ptr<TP_MLA<KML_MLA_TP_QUAN<float>>>>(
|
||
mla_module, "MLA_QUAN_F32")
|
||
.def(py::init<GeneralMLAConfig>())
|
||
.def("load_weights", &TP_MLA<KML_MLA_TP_QUAN<float>>::load_weights)
|
||
.def("forward",
|
||
[](TP_MLA<KML_MLA_TP_QUAN<float>>& op, std::vector<int> qlens, std::vector<std::vector<int>> page_tables,
|
||
std::vector<int> kvlens, intptr_t input,
|
||
intptr_t output) { op.forward(qlens, page_tables, kvlens, (const void*)input, (void*)output); })
|
||
.def("set_local_pages", &TP_MLA<KML_MLA_TP_QUAN<float>>::set_local_pages)
|
||
.def("set_pages", [](TP_MLA<KML_MLA_TP_QUAN<float>>& op, std::vector<std::vector<intptr_t>> nope_pages,
|
||
std::vector<std::vector<intptr_t>> rope_pages) {
|
||
std::vector<std::vector<void*>> nope_pages_ptr;
|
||
std::vector<std::vector<void*>> rope_pages_ptr;
|
||
op.set_pages(nope_pages_ptr, rope_pages_ptr);
|
||
});
|
||
|
||
auto gate_module = m.def_submodule("gate");
|
||
py::class_<GeneralGateConfig>(gate_module, "GateConfig")
|
||
.def(py::init([](int hidden_size, int num_experts_per_tok, int n_routed_experts, int n_group, int topk_group) {
|
||
return GeneralGateConfig(hidden_size, num_experts_per_tok, n_routed_experts, n_group, topk_group);
|
||
}))
|
||
.def_readwrite("routed_scaling_factor", &GeneralGateConfig::routed_scaling_factor)
|
||
|
||
.def_readwrite("layer_idx", &GeneralGateConfig::layer_idx)
|
||
.def_readwrite("pool", &GeneralGateConfig::pool)
|
||
.DEF_PTR_PROPERTY(GeneralGateConfig, weight)
|
||
.def_readwrite("weight_type", &GeneralGateConfig::weight_type)
|
||
.DEF_PTR_PROPERTY(GeneralGateConfig, e_score_correction_bias)
|
||
.def_readwrite("e_score_correction_bias_type", &GeneralGateConfig::e_score_correction_bias_type)
|
||
|
||
;
|
||
py::class_<MoEGate, std::shared_ptr<MoEGate>>(gate_module, "MoEGate")
|
||
.def(py::init<GeneralGateConfig>())
|
||
.def("forward", &MoEGate::forward_binding);
|
||
#endif
|
||
|
||
py::class_<QuantConfig>(m, "QuantConfig")
|
||
.def(py::init<>())
|
||
.def_readwrite("quant_method", &QuantConfig::quant_method)
|
||
.def_readwrite("bits", &QuantConfig::bits)
|
||
.def_readwrite("group_size", &QuantConfig::group_size)
|
||
.def_readwrite("zero_point", &QuantConfig::zero_point)
|
||
.def_readwrite("per_channel", &QuantConfig::per_channel);
|
||
|
||
auto moe_module = m.def_submodule("moe");
|
||
|
||
py::class_<GeneralMOEConfig>(moe_module, "MOEConfig")
|
||
.def(py::init([](int expert_num, int routed_expert_num, int hidden_size, int intermediate_size) {
|
||
return GeneralMOEConfig(expert_num, routed_expert_num, hidden_size, intermediate_size);
|
||
}))
|
||
.def(py::init(
|
||
[](int expert_num, int routed_expert_num, int hidden_size, int intermediate_size, int num_gpu_experts) {
|
||
GeneralMOEConfig cfg(expert_num, routed_expert_num, hidden_size, intermediate_size);
|
||
cfg.num_gpu_experts = num_gpu_experts;
|
||
return cfg;
|
||
}))
|
||
.def(py::init([](int expert_num, int routed_expert_num, int hidden_size, int intermediate_size,
|
||
uintptr_t gpu_experts_mask_ptr) {
|
||
GeneralMOEConfig cfg(expert_num, routed_expert_num, hidden_size, intermediate_size);
|
||
cfg.gpu_experts_mask = reinterpret_cast<uint8_t*>(gpu_experts_mask_ptr);
|
||
cfg.compute_num_gpu_experts();
|
||
return cfg;
|
||
}))
|
||
// Core config fields (required for Python access after construction)
|
||
.def_readwrite("expert_num", &GeneralMOEConfig::expert_num)
|
||
.def_readwrite("num_experts_per_tok", &GeneralMOEConfig::num_experts_per_tok)
|
||
.def_readwrite("hidden_size", &GeneralMOEConfig::hidden_size)
|
||
.def_readwrite("intermediate_size", &GeneralMOEConfig::intermediate_size)
|
||
.def_readwrite("layer_idx", &GeneralMOEConfig::layer_idx)
|
||
.def_readwrite("pool", &GeneralMOEConfig::pool)
|
||
|
||
.def_readonly("num_gpu_experts", &GeneralMOEConfig::num_gpu_experts)
|
||
.def_property(
|
||
"gpu_experts_mask",
|
||
[](const GeneralMOEConfig& self) { return reinterpret_cast<uintptr_t>(self.gpu_experts_mask); },
|
||
[](GeneralMOEConfig& self, uintptr_t val) { self.gpu_experts_mask = reinterpret_cast<uint8_t*>(val); })
|
||
.DEF_PTR_PROPERTY(GeneralMOEConfig, physical_to_logical_map)
|
||
|
||
.DEF_PTR_PROPERTY(GeneralMOEConfig, gate_proj)
|
||
.DEF_PTR_PROPERTY(GeneralMOEConfig, up_proj)
|
||
.DEF_PTR_PROPERTY(GeneralMOEConfig, down_proj)
|
||
|
||
.DEF_PTR_PROPERTY(GeneralMOEConfig, gate_scale)
|
||
.DEF_PTR_PROPERTY(GeneralMOEConfig, up_scale)
|
||
.DEF_PTR_PROPERTY(GeneralMOEConfig, down_scale)
|
||
|
||
.DEF_PTR_PROPERTY(GeneralMOEConfig, gate_zero)
|
||
.DEF_PTR_PROPERTY(GeneralMOEConfig, up_zero)
|
||
.DEF_PTR_PROPERTY(GeneralMOEConfig, down_zero)
|
||
|
||
.def_readwrite("quant_config", &GeneralMOEConfig::quant_config)
|
||
|
||
.def_readwrite("max_len", &GeneralMOEConfig::max_len)
|
||
|
||
.DEF_PTR_2D_PROPERTY(GeneralMOEConfig, gate_projs)
|
||
.DEF_PTR_2D_PROPERTY(GeneralMOEConfig, up_projs)
|
||
.DEF_PTR_2D_PROPERTY(GeneralMOEConfig, down_projs)
|
||
|
||
.DEF_PTR_2D_PROPERTY(GeneralMOEConfig, gate_scales)
|
||
.DEF_PTR_2D_PROPERTY(GeneralMOEConfig, up_scales)
|
||
.DEF_PTR_2D_PROPERTY(GeneralMOEConfig, down_scales)
|
||
|
||
.DEF_PTR_2D_PROPERTY(GeneralMOEConfig, gate_zeros)
|
||
.DEF_PTR_2D_PROPERTY(GeneralMOEConfig, up_zeros)
|
||
.DEF_PTR_2D_PROPERTY(GeneralMOEConfig, down_zeros)
|
||
|
||
.DEF_PTR_2D_PROPERTY(GeneralMOEConfig, gate_bwd_projs)
|
||
.DEF_PTR_2D_PROPERTY(GeneralMOEConfig, up_bwd_projs)
|
||
.DEF_PTR_2D_PROPERTY(GeneralMOEConfig, down_bwd_projs)
|
||
.DEF_PTR_2D_PROPERTY(GeneralMOEConfig, gate_bwd_scales)
|
||
.DEF_PTR_2D_PROPERTY(GeneralMOEConfig, up_bwd_scales)
|
||
.DEF_PTR_2D_PROPERTY(GeneralMOEConfig, down_bwd_scales)
|
||
|
||
.def_readwrite("path", &GeneralMOEConfig::path)
|
||
.def_readwrite("save", &GeneralMOEConfig::save)
|
||
.def_readwrite("load", &GeneralMOEConfig::load)
|
||
.def_readwrite("share_backward_bb", &GeneralMOEConfig::share_backward_bb)
|
||
.def_readwrite("share_cache_pool", &GeneralMOEConfig::share_cache_pool)
|
||
.def_readwrite("m_block", &GeneralMOEConfig::m_block)
|
||
.def_readwrite("group_min_len", &GeneralMOEConfig::group_min_len)
|
||
.def_readwrite("group_max_len", &GeneralMOEConfig::group_max_len)
|
||
|
||
.def_readwrite("gate_type", &GeneralMOEConfig::gate_type)
|
||
.def_readwrite("up_type", &GeneralMOEConfig::up_type)
|
||
.def_readwrite("down_type", &GeneralMOEConfig::down_type)
|
||
.def_readwrite("hidden_type", &GeneralMOEConfig::hidden_type)
|
||
.def_readwrite("max_cache_depth", &GeneralMOEConfig::max_cache_depth)
|
||
// V4-Flash 2604B SwiGLU clamp limit (0.0 = disabled). See common.hpp.
|
||
.def_readwrite("swiglu_limit", &GeneralMOEConfig::swiglu_limit)
|
||
.def_readwrite("swiglu_alpha", &GeneralMOEConfig::swiglu_alpha)
|
||
|
||
;
|
||
|
||
// MOESFTConfig - extends GeneralMOEConfig with LoRA support
|
||
py::class_<MOESFTConfig, GeneralMOEConfig>(moe_module, "MOESFTConfig")
|
||
.def(py::init<>())
|
||
.def(py::init([](int expert_num, int routed_expert_num, int hidden_size, int intermediate_size) {
|
||
return MOESFTConfig(expert_num, routed_expert_num, hidden_size, intermediate_size);
|
||
}))
|
||
.def_readwrite("lora_rank", &MOESFTConfig::lora_rank)
|
||
.def_readwrite("lora_alpha", &MOESFTConfig::lora_alpha)
|
||
.def_readwrite("lora_dropout", &MOESFTConfig::lora_dropout)
|
||
.DEF_PTR_PROPERTY(MOESFTConfig, gate_lora_a)
|
||
.DEF_PTR_PROPERTY(MOESFTConfig, gate_lora_b)
|
||
.DEF_PTR_PROPERTY(MOESFTConfig, up_lora_a)
|
||
.DEF_PTR_PROPERTY(MOESFTConfig, up_lora_b)
|
||
.DEF_PTR_PROPERTY(MOESFTConfig, down_lora_a)
|
||
.DEF_PTR_PROPERTY(MOESFTConfig, down_lora_b)
|
||
.def_readwrite("full_weight_grad", &MOESFTConfig::full_weight_grad)
|
||
.def_readwrite("authoritative_optimizer_grads", &MOESFTConfig::authoritative_optimizer_grads)
|
||
.DEF_PTR_PROPERTY(MOESFTConfig, grad_gate_proj)
|
||
.DEF_PTR_PROPERTY(MOESFTConfig, grad_up_proj)
|
||
.DEF_PTR_PROPERTY(MOESFTConfig, grad_down_proj);
|
||
|
||
py::class_<MoE_Interface, std::shared_ptr<MoE_Interface>>(moe_module, "MoE_Interface");
|
||
|
||
bind_moe_module<LLAMA_MOE_TP>(moe_module, "MOE");
|
||
|
||
#if defined(__x86_64__) && defined(USE_AMX_AVX_KERNEL)
|
||
bind_moe_module<AMX_MOE_TP<amx::GemmKernel224Int8>>(moe_module, "AMXInt8_MOE");
|
||
bind_moe_module<AMX_MOE_TP<amx::GemmKernel224Int4>>(moe_module, "AMXInt4_MOE");
|
||
bind_moe_module<AMX_MOE_TP<amx::GemmKernel224Int4_1>>(moe_module, "AMXInt4_1_MOE");
|
||
bind_moe_module<AMX_AWQ_MOE_TP<amx::GemmKernel224Int4_1_LowKGroup>>(moe_module, "AMXInt4_1KGroup_MOE");
|
||
bind_moe_module<AMX_K2_MOE_TP<amx::GemmKernel224Int4SmallKGroup>>(moe_module, "AMXInt4_KGroup_MOE");
|
||
bind_moe_module<AMX_K2_MOE_TP<amx::GemmKernel224Int4SmallKGroupBlocked>>(moe_module, "AMXInt4_KGroupBlocked_MOE");
|
||
#if defined(__AVX512F__)
|
||
bind_moe_module<AMX_BF16_MOE_TP<amx::GemmKernel224BF16>>(moe_module, "AMXBF16_MOE");
|
||
bind_moe_module<AMX_FP8_MOE_TP<amx::GemmKernel224FP8>>(moe_module, "AMXFP8_MOE");
|
||
bind_moe_module<AMX_FP8_PERCHANNEL_MOE_TP<amx::GemmKernel224FP8PerChannel>>(moe_module, "AMXFP8PerChannel_MOE");
|
||
bind_moe_module<AMX_FP4_MOE_TP<amx::GemmKernel224MXFP4SmallKGroup>>(moe_module, "AMXFP4_KGroup_MOE");
|
||
bind_moe_module<AMX_MXFP8_MOE_TP<amx::GemmKernel224MXFP8SmallKGroup>>(moe_module, "AMXMXFP8_KGroup_MOE");
|
||
#endif
|
||
#if defined(__AVX512BF16__)
|
||
// SFT MoE with LoRA support (BF16, INT8, INT4, AWQ, K2)
|
||
bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224BF16, AMX_BF16_MOE_TP>>(moe_module, "AMXBF16_SFT_MOE");
|
||
bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224MXFP4SmallKGroup, AMX_FP4_MOE_TP>>(
|
||
moe_module, "MXFP4_SFT_MOE");
|
||
#if defined(__AVX512VBMI__)
|
||
bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224FP8, AMX_FP8_MOE_TP>>(moe_module, "AMXFP8_SFT_MOE");
|
||
#endif
|
||
bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224Int8>>(moe_module, "AMXInt8_SFT_MOE");
|
||
bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224Int4>>(moe_module, "AMXInt4_SFT_MOE");
|
||
// bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224Int4_1>>(moe_module, "AMXInt4_1_SFT_MOE");
|
||
// bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224Int4_1_LowKGroup, AMX_AWQ_MOE_TP>>(moe_module,
|
||
// "AMXInt4_1KGroup_SFT_MOE");
|
||
bind_moe_sft_module<AMX_K2_SFT_MOE_TP<>>(moe_module, "AMXInt4_KGroup_SFT_MOE");
|
||
// SFT MoE with SkipLoRA=true (skip all LoRA computation in backward, only compute base weight grad_input)
|
||
bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224BF16, AMX_BF16_MOE_TP, true>>(moe_module,
|
||
"AMXBF16_SFT_MOE_SkipLoRA");
|
||
bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224Int8, AMX_MOE_TP, true>>(moe_module, "AMXInt8_SFT_MOE_SkipLoRA");
|
||
bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224Int4, AMX_MOE_TP, true>>(moe_module, "AMXInt4_SFT_MOE_SkipLoRA");
|
||
// bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224Int4_1, AMX_MOE_TP, true>>(moe_module,
|
||
// "AMXInt4_1_SFT_MOE_SkipLoRA");
|
||
// bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224Int4_1_LowKGroup, AMX_AWQ_MOE_TP, true>>(
|
||
// moe_module, "AMXInt4_1KGroup_SFT_MOE_SkipLoRA");
|
||
// bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224Int4SmallKGroup, AMX_K2_MOE_TP, true>>(
|
||
// moe_module, "AMXInt4_KGroup_SFT_MOE_SkipLoRA");
|
||
#endif
|
||
#endif
|
||
// AVX2 backends — available on all x86_64 (no AMX/AVX512 requirement)
|
||
#if defined(__x86_64__)
|
||
bind_moe_module<AVX2_BF16_MOE_TP<avx2::GemmKernelAVX2BF16>>(moe_module, "AVX2BF16_MOE");
|
||
bind_moe_module<AVX2_FP8_MOE_TP<avx2::GemmKernelAVX2FP8>>(moe_module, "AVX2FP8_MOE");
|
||
bind_moe_module<AVX2_GPTQ_INT4_MOE_TP<avx2::GemmKernelAVX2GPTQInt4>>(moe_module, "AVX2GPTQInt4_MOE");
|
||
bind_moe_module<AVX2_RAW_INT4_MOE_TP<avx2::GemmKernelAVX2RawInt4>>(moe_module, "AVX2RawInt4_MOE");
|
||
bind_moe_module<AVX2_MXFP4_MOE_TP<avx2::GemmKernelAVX2MXFP4>>(moe_module, "AVX2MXFP4_MOE");
|
||
#if !defined(USE_AMX_AVX_KERNEL)
|
||
// AVX2 tier of the MXFP4 routed-expert LoRA SFT (same Python name as the AMX build)
|
||
bind_moe_sft_module<avx2::AVX2_SFT_MXFP4_MOE_TP<avx2::GemmKernelAVX2MXFP4>>(moe_module, "MXFP4_SFT_MOE");
|
||
// AVX2 tier of the RAWINT4 routed-expert LoRA SFT (same Python name as the AMX build)
|
||
bind_moe_sft_module<avx2::AVX2_SFT_RAWINT4_MOE_TP<avx2::GemmKernelAVX2RawInt4>>(moe_module,
|
||
"AMXInt4_KGroup_SFT_MOE");
|
||
#endif
|
||
bind_moe_module<AVX2_MXFP8_MOE_TP<avx2::GemmKernelAVX2MXFP8>>(moe_module, "AVX2MXFP8_MOE");
|
||
bind_moe_module<AVXVNNI256_GPTQ_INT4_MOE_TP<avxvnni::GemmKernelAVXVNNI256GPTQInt4>>(moe_module,
|
||
"AVXVNNI256GPTQInt4_MOE");
|
||
bind_moe_module<AVXVNNI256_GPTQ_INT4_PACKED_MOE_TP<>>(moe_module, "AVXVNNI256GPTQInt4Packed_MOE");
|
||
bind_moe_module<AVXVNNI256_RAW_INT4_MOE_TP<avxvnni_rawint4::GemmKernelAVXVNNI256RawInt4>>(moe_module,
|
||
"AVXVNNI256RawInt4_MOE");
|
||
#endif
|
||
#if defined(USE_SYCL)
|
||
bind_moe_module<SYCL_GPTQ_INT4_MOE_TP<sycl_int4::GemmKernelSYCLGPTQInt4>>(moe_module, "SYCLGPTQInt4_MOE");
|
||
#endif
|
||
|
||
#if defined(USE_MOE_KERNEL)
|
||
bind_moe_module<MOE_KERNEL_TP<moe_kernel::GemmKernelInt8, _is_plain_>>(moe_module, "Int8_KERNEL_MOE");
|
||
#if defined(__aarch64__) && defined(CPU_USE_KML)
|
||
// amd have not implemented int4 kernel yet
|
||
bind_moe_module<MOE_KERNEL_TP<moe_kernel::GemmKernelInt4, _is_plain_>>(moe_module, "Int4_KERNEL_MOE");
|
||
#endif
|
||
#endif
|
||
|
||
// Expose kernel tiling/runtime parameters so Python can modify them at runtime
|
||
{
|
||
auto tiling_module = moe_module.def_submodule("tiling");
|
||
#if defined(USE_MOE_KERNEL)
|
||
tiling_module.def(
|
||
"get_int8",
|
||
[]() {
|
||
auto t = moe_kernel::GemmKernelInt8::get_tiling();
|
||
py::dict d;
|
||
d["n_block_up_gate"] = std::get<0>(t);
|
||
d["n_block_down"] = std::get<1>(t);
|
||
d["n_block"] = std::get<2>(t);
|
||
d["m_block"] = std::get<3>(t);
|
||
d["k_block"] = std::get<4>(t);
|
||
d["n_block_up_gate_prefi"] = std::get<5>(t);
|
||
d["n_block_down_prefi"] = std::get<6>(t);
|
||
return d;
|
||
},
|
||
"Get current tiling parameters for INT8 kernel");
|
||
tiling_module.def(
|
||
"set_int8",
|
||
[](int n_block_up_gate, int n_block_down, int n_block, int m_block, int k_block, int n_block_up_gate_prefi,
|
||
int n_block_down_prefi) {
|
||
moe_kernel::GemmKernelInt8::set_tiling(n_block_up_gate, n_block_down, n_block, m_block, k_block,
|
||
n_block_up_gate_prefi, n_block_down_prefi);
|
||
},
|
||
py::arg("n_block_up_gate"), py::arg("n_block_down"), py::arg("n_block"), py::arg("m_block"), py::arg("k_block"),
|
||
py::arg("n_block_up_gate_prefi"), py::arg("n_block_down_prefi"), "Set tiling parameters for INT8 kernel");
|
||
|
||
tiling_module.def(
|
||
"get_int4",
|
||
[]() {
|
||
auto t = moe_kernel::GemmKernelInt4::get_tiling();
|
||
py::dict d;
|
||
d["n_block_up_gate"] = std::get<0>(t);
|
||
d["n_block_down"] = std::get<1>(t);
|
||
d["n_block"] = std::get<2>(t);
|
||
d["m_block"] = std::get<3>(t);
|
||
d["k_block"] = std::get<4>(t);
|
||
d["n_block_up_gate_prefi"] = std::get<5>(t);
|
||
d["n_block_down_prefi"] = std::get<6>(t);
|
||
return d;
|
||
},
|
||
"Get current tiling parameters for INT4 kernel");
|
||
tiling_module.def(
|
||
"set_int4",
|
||
[](int n_block_up_gate, int n_block_down, int n_block, int m_block, int k_block, int n_block_up_gate_prefi,
|
||
int n_block_down_prefi) {
|
||
moe_kernel::GemmKernelInt4::set_tiling(n_block_up_gate, n_block_down, n_block, m_block, k_block,
|
||
n_block_up_gate_prefi, n_block_down_prefi);
|
||
},
|
||
py::arg("n_block_up_gate"), py::arg("n_block_down"), py::arg("n_block"), py::arg("m_block"), py::arg("k_block"),
|
||
py::arg("n_block_up_gate_prefi"), py::arg("n_block_down_prefi"), "Set tiling parameters for INT4 kernel");
|
||
|
||
// Convenience: set both
|
||
tiling_module.def(
|
||
"set_all",
|
||
[](int n_block_up_gate, int n_block_down, int n_block, int m_block, int k_block, int n_block_up_gate_prefi,
|
||
int n_block_down_prefi) {
|
||
moe_kernel::GemmKernelInt8::set_tiling(n_block_up_gate, n_block_down, n_block, m_block, k_block,
|
||
n_block_up_gate_prefi, n_block_down_prefi);
|
||
moe_kernel::GemmKernelInt4::set_tiling(n_block_up_gate, n_block_down, n_block, m_block, k_block,
|
||
n_block_up_gate_prefi, n_block_down_prefi);
|
||
},
|
||
py::arg("n_block_up_gate"), py::arg("n_block_down"), py::arg("n_block"), py::arg("m_block"), py::arg("k_block"),
|
||
py::arg("n_block_up_gate_prefi"), py::arg("n_block_down_prefi"),
|
||
"Set tiling parameters for both INT8 and INT4 kernels");
|
||
#endif
|
||
}
|
||
|
||
auto kvcache_module = m.def_submodule("kvcache");
|
||
|
||
py::enum_<AnchorType>(kvcache_module, "AnchorType")
|
||
.value("FIXED", AnchorType::FIXED_ANCHOR)
|
||
.value("DYNAMIC", AnchorType::DYNAMIC)
|
||
.value("QUEST", AnchorType::QUEST)
|
||
.value("BLOCK_MAX", AnchorType::BLOCK_MAX)
|
||
.value("BLOCK_MEAN", AnchorType::BLOCK_MEAN);
|
||
py::enum_<ggml_type>(kvcache_module, "ggml_type")
|
||
// .value("FP16", ggml_type::GGML_TYPE_F16)
|
||
// .value("FP32", ggml_type::GGML_TYPE_F32)
|
||
// .value("Q4_0", ggml_type::GGML_TYPE_Q4_0)
|
||
// .value("Q8_0", ggml_type::GGML_TYPE_Q8_0)
|
||
.value("FP32", GGML_TYPE_F32)
|
||
.value("FP16", GGML_TYPE_F16)
|
||
.value("Q4_0", GGML_TYPE_Q4_0)
|
||
.value("Q4_1", GGML_TYPE_Q4_1)
|
||
.value("Q5_0", GGML_TYPE_Q5_0)
|
||
.value("Q5_1", GGML_TYPE_Q5_1)
|
||
.value("Q8_0", GGML_TYPE_Q8_0)
|
||
.value("Q8_1", GGML_TYPE_Q8_1)
|
||
.value("Q2_K", GGML_TYPE_Q2_K)
|
||
.value("Q3_K", GGML_TYPE_Q3_K)
|
||
.value("Q4_K", GGML_TYPE_Q4_K)
|
||
.value("Q5_K", GGML_TYPE_Q5_K)
|
||
.value("Q6_K", GGML_TYPE_Q6_K)
|
||
.value("Q8_K", GGML_TYPE_Q8_K)
|
||
.value("IQ2_XXS", GGML_TYPE_IQ2_XXS)
|
||
.value("IQ2_XS", GGML_TYPE_IQ2_XS)
|
||
.value("IQ3_XXS", GGML_TYPE_IQ3_XXS)
|
||
.value("IQ1_S", GGML_TYPE_IQ1_S)
|
||
.value("IQ4_NL", GGML_TYPE_IQ4_NL)
|
||
.value("IQ3_S", GGML_TYPE_IQ3_S)
|
||
.value("IQ2_S", GGML_TYPE_IQ2_S)
|
||
.value("IQ4_XS", GGML_TYPE_IQ4_XS)
|
||
.value("I8", GGML_TYPE_I8)
|
||
.value("I16", GGML_TYPE_I16)
|
||
.value("I32", GGML_TYPE_I32)
|
||
.value("I64", GGML_TYPE_I64)
|
||
.value("F64", GGML_TYPE_F64)
|
||
.value("IQ1_M", GGML_TYPE_IQ1_M)
|
||
.value("BF16", GGML_TYPE_BF16)
|
||
.export_values();
|
||
|
||
py::enum_<RetrievalType>(kvcache_module, "RetrievalType")
|
||
.value("LAYER", RetrievalType::LAYER)
|
||
.value("KVHEAD", RetrievalType::KVHEAD)
|
||
.value("QHEAD", RetrievalType::QHEAD);
|
||
|
||
py::class_<KVCacheConfig>(kvcache_module, "KVCacheConfig")
|
||
.def(py::init<int, int, int, int, int, int, AnchorType, ggml_type, RetrievalType, int, int, int, int, int, int>())
|
||
.def_readwrite("layer_num", &KVCacheConfig::layer_num)
|
||
.def_readwrite("kv_head_num", &KVCacheConfig::kv_head_num)
|
||
.def_readwrite("q_head_num", &KVCacheConfig::q_head_num)
|
||
.def_readwrite("head_dim", &KVCacheConfig::head_dim)
|
||
.def_readwrite("block_len", &KVCacheConfig::block_len)
|
||
.def_readwrite("anchor_num", &KVCacheConfig::anchor_num)
|
||
.def_readwrite("anchor_type", &KVCacheConfig::anchor_type)
|
||
.def_readwrite("kv_type", &KVCacheConfig::kv_type)
|
||
.def_readwrite("retrieval_type", &KVCacheConfig::retrieval_type)
|
||
.def_readwrite("layer_step", &KVCacheConfig::layer_step)
|
||
.def_readwrite("token_step", &KVCacheConfig::token_step)
|
||
.def_readwrite("layer_offset", &KVCacheConfig::layer_offset)
|
||
.def_readwrite("max_block_num", &KVCacheConfig::max_block_num)
|
||
.def_readwrite("max_batch_size", &KVCacheConfig::max_batch_size)
|
||
.def_readwrite("max_thread_num", &KVCacheConfig::max_thread_num);
|
||
py::class_<KVCache>(kvcache_module, "KVCache")
|
||
.def(py::init<KVCacheConfig>())
|
||
.def("get_cache_total_len", &KVCache::get_cache_total_len)
|
||
.def("update_cache_total_len",
|
||
[](KVCache& kvcache, int cache_total_len) { kvcache.update_cache_total_len(cache_total_len); });
|
||
|
||
auto utils = m.def_submodule("utils");
|
||
|
||
// 注册转换函数
|
||
utils.def("to_float", &to_float_ptr, "Convert tensor from any GGML type to float32", py::arg("input"),
|
||
py::arg("size"), py::arg("type"));
|
||
|
||
utils.def("from_float", &from_float_ptr, "Convert tensor from float32 to any GGML type", py::arg("input"),
|
||
py::arg("size"), py::arg("type"));
|
||
}
|
||
|
||
#if defined(KTRANSFORMERS_ENABLE_CPPTRACE)
|
||
static void warmup_cpptrace() {
|
||
// 避免第一次调用触发 lazy-loading（malloc 等） :contentReference[oaicite:7]{index=7}
|
||
cpptrace::frame_ptr buffer[10];
|
||
(void)cpptrace::safe_generate_raw_trace(buffer, 10);
|
||
cpptrace::safe_object_frame frame{};
|
||
cpptrace::get_safe_object_frame(buffer[0], &frame);
|
||
}
|
||
|
||
static void crash_handler(int signo, siginfo_t* /*info*/, void* /*ucontext*/) {
|
||
const char* head = "=== crash: signal received ===\n";
|
||
write(STDERR_FILENO, head, std::strlen(head));
|
||
cpptrace::generate_trace().print();
|
||
_exit(128 + signo);
|
||
}
|
||
|
||
__attribute__((constructor)) static void install_handlers() {
|
||
struct sigaction sa;
|
||
std::memset(&sa, 0, sizeof(sa));
|
||
sa.sa_sigaction = &crash_handler;
|
||
sa.sa_flags = SA_SIGINFO;
|
||
sigemptyset(&sa.sa_mask);
|
||
|
||
sigaction(SIGSEGV, &sa, nullptr);
|
||
sigaction(SIGABRT, &sa, nullptr);
|
||
}
|
||
|
||
#endif
|