Skip to content

[perf](kt-kernel): SFT end2end #2094

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
yyj6666667 wants to merge 20 commits into
base: main
Choose a base branch
from
+4,715 −504
Open

[perf](kt-kernel): SFT end2end #2094

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
20 commits
Select commit Hold shift + click to select a range
e99b5e1
Add full FT development snapshot
Illumination111 Jul 10, 2026
2d81e86
[fix](kt-kernel): fix Full FT TP base weight gradients
Illumination111 Jul 10, 2026
20f645c
[fix]: bug fix of 2d81e86
Illumination111 Jul 13, 2026
5ce0767
[fix](kt-kernel): fix AMX BF16 full-weight gradients
Illumination111 Jul 13, 2026
6aa1862
[docs](kt-kernel): document Full FT fork changes and debug history
Illumination111 Jul 13, 2026
06f0606
[docs](kt-kernel): align fork remote terminology
Illumination111 Jul 13, 2026
01c0c92
[chore](kt-kernel): keep agent notes local
Illumination111 Jul 15, 2026
b6f7a2e
[fix](kt-kernel): configure SFT OpenMP threads
Illumination111 Jul 15, 2026
f209878
[perf](kt-kernel): optimize AMX Full-FT weight gradients
Illumination111 Jul 16, 2026
9dc9d93
[feat](kt-kernel): add staged SFT profiling
yyj6666667 Jul 16, 2026
c6f4211
[fix](kt-kernel): reuse inference BF16 kernel for SFT
yyj6666667 Jul 17, 2026
109b403
[perf](kt-kernel): add fine-grained Full-FT profiling
yyj6666667 Jul 17, 2026
34d2102
[perf](kt-kernel): batch BF16 Full-FT weight gradients
yyj6666667 Jul 17, 2026
ea84e6e
[test](kt-kernel): benchmark BF16 dWeight AMX driver
yyj6666667 Jul 17, 2026
273e670
[fix](kt-kernel): label dWeight store as worker CPU time
yyj6666667 Jul 17, 2026
61e63a8
[docs](sft): record BF16 Full-FT performance
yyj6666667 Jul 17, 2026
1e95053
[docs](sft): remove Qwen3 Full-FT performance report
yyj6666667 Jul 17, 2026
66f5f15
[perf](kt-kernel): reduce BF16 Full-FT checkpoint overhead
yyj6666667 Jul 17, 2026
6eeffa7
[perf](kt-kernel): make SFT optimizer gradients authoritative
yyj6666667 Jul 18, 2026
e26e21f
[perf](kt-kernel): avoid eager Full-FT gradient zeroing
yyj6666667 Jul 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions kt-kernel/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -395,6 +395,7 @@ if(HOST_IS_X86)
# 获取不带扩展名的文件名作为 target 名
get_filename_component(test_name ${test_src} NAME_WE)
add_executable(${test_name} ${test_src} ${CMAKE_CURRENT_SOURCE_DIR}/cpu_backend/shared_mem_buffer.cpp)
target_compile_options(${test_name} PRIVATE ${ARCH_FLAGS})
target_link_libraries(${test_name} llama OpenMP::OpenMP_CXX numa)
endforeach()
endif()
Expand Down
55 changes: 46 additions & 9 deletions kt-kernel/ext_bindings.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -325,22 +325,35 @@ class MOESFTBindings {
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_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_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_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);
}
};
Expand Down Expand Up @@ -391,19 +404,37 @@ void bind_moe_sft_module(py::module_& moe_module, const char* name) {
.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)
.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)
.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("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__) && defined(USE_AMX_AVX_KERNEL)

Expand Down Expand Up @@ -779,7 +810,12 @@ PYBIND11_MODULE(kt_kernel_ext, m) {
.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_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");

Expand All @@ -800,7 +836,7 @@ PYBIND11_MODULE(kt_kernel_ext, m) {
#endif
#if defined(__AVX512BF16__)
// SFT MoE with LoRA support (BF16, INT8, INT4, AWQ, K2)
bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224BF>>(moe_module, "AMXBF16_SFT_MOE");
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::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");
Expand All @@ -809,7 +845,8 @@ PYBIND11_MODULE(kt_kernel_ext, m) {
// bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224Int4SmallKGroup, AMX_K2_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::GemmKernel224BF, AMX_MOE_TP, true>>(moe_module, "AMXBF16_SFT_MOE_SkipLoRA");
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,
Expand Down
2 changes: 2 additions & 0 deletions kt-kernel/operators/amx/bf16-moe.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,8 @@
template <class T = amx::GemmKernel224BF16>
class AMX_BF16_MOE_TP : public AMX_MOE_BASE<T, AMX_BF16_MOE_TP<T>> {
using Base = AMX_MOE_BASE<T, AMX_BF16_MOE_TP<T>>;

protected:
using Base::config_;
using Base::down_ba_;
using Base::down_bb_;
Expand Down
117 changes: 112 additions & 5 deletions kt-kernel/operators/amx/la/amx_raw_buffers.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -133,16 +133,13 @@ struct BufferBBF16Impl {
}
void set_data(void* new_ptr) { b = reinterpret_cast<ggml_bf16_t*>(new_ptr); }

void from_mat(ggml_bf16_t* src, int ith, int nth) {
auto [n_start, n_end] = K::split_range_n(n, ith, nth);
int n_block_begin = n_start;
int n_block_size = n_end - n_block_begin;
void pack_block(ggml_bf16_t* src, int src_stride, int n_block_begin, int n_block_size) {
for (int n_begin = 0; n_begin < n_block_size; n_begin += N_STEP) {
for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K_BLOCK) {
int k_block_size = std::min(K_BLOCK, k - k_block_begin);
for (int k_begin = 0; k_begin < k_block_size; k_begin += K_STEP) {
for (int i = 0; i < N_STEP; i++) {
__m512i* s = (__m512i*)(src + (n_block_begin + n_begin + i) * k + k_block_begin + k_begin);
__m512i* s = (__m512i*)(src + (n_begin + i) * src_stride + k_block_begin + k_begin);
__m512i* d = (__m512i*)(b + n_block_begin * k + k_block_begin * n_block_size + n_begin * k_block_size +
k_begin * N_STEP + i * K_STEP);
avx512_copy_32xbf16(s, d);
Expand All @@ -155,6 +152,116 @@ struct BufferBBF16Impl {
}
}
}

void from_mat(ggml_bf16_t* src, int ith, int nth) {
auto [n_start, n_end] = K::split_range_n(n, ith, nth);
int n_block_begin = n_start;
int n_block_size = n_end - n_block_begin;
pack_block(src + n_block_begin * k, k, n_block_begin, n_block_size);
}

void from_mat_strided(ggml_bf16_t* src, int src_stride, int ith, int nth) {
assert(src_stride >= k);
auto [n_start, n_end] = K::split_range_n(n, ith, nth);
int n_block_begin = n_start;
int n_block_size = n_end - n_block_begin;
if (n_block_size <= 0) return;
pack_block(src + (size_t)n_block_begin * src_stride, src_stride, n_block_begin, n_block_size);
}

void from_mat_transposed(ggml_bf16_t* src, int src_n, int src_k, int ith, int nth) {
assert(n == src_k && k == src_n);
auto [n_start, n_end] = K::split_range_n(n, ith, nth);
int n_block_begin = n_start;
int n_block_size = n_end - n_block_begin;
if (n_block_size <= 0) return;

thread_local std::vector<ggml_bf16_t> strip;
strip.resize((size_t)n_block_size * k);
constexpr int TILE = 32;
for (int c_tile = 0; c_tile < k; c_tile += TILE) {
int c_end = std::min(c_tile + TILE, k);
for (int r_tile = 0; r_tile < n_block_size; r_tile += TILE) {
int r_end = std::min(r_tile + TILE, n_block_size);
for (int c = c_tile; c < c_end; c++) {
for (int r = r_tile; r < r_end; r++) {
strip[(size_t)r * k + c] = src[(size_t)c * src_k + n_block_begin + r];
}
}
}
}
pack_block(strip.data(), k, n_block_begin, n_block_size);
}

void to_mat(ggml_bf16_t* dst, int ith, int nth) const {
auto [n_start, n_end] = K::split_range_n(n, ith, nth);
int n_block_begin = n_start;
int n_block_size = n_end - n_block_begin;
if (n_block_size <= 0) return;

alignas(64) ggml_bf16_t tile_copy[N_STEP * K_STEP];
for (int n_begin = 0; n_begin < n_block_size; n_begin += N_STEP) {
for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K_BLOCK) {
int k_block_size = std::min(K_BLOCK, k - k_block_begin);
for (int k_begin = 0; k_begin < k_block_size; k_begin += K_STEP) {
const ggml_bf16_t* tile_src =
b + n_block_begin * k + k_block_begin * n_block_size + n_begin * k_block_size + k_begin * N_STEP;
memcpy(tile_copy, tile_src, sizeof(tile_copy));
transpose_16x16_32bit((__m512i*)tile_copy);
transpose_16x16_32bit((__m512i*)(tile_copy + TILE_N * K_STEP));
for (int i = 0; i < N_STEP; i++) {
__m512i* s = (__m512i*)(tile_copy + i * K_STEP);
__m512i* d = (__m512i*)(dst + (size_t)(n_block_begin + n_begin + i) * k + k_block_begin + k_begin);
avx512_copy_32xbf16(s, d);
}
}
}
}
}

void from_bb_transposed(const BufferBBF16Impl& src, int ith, int nth) {
assert(n == src.k && k == src.n);
auto [n_start, n_end] = K::split_range_n(n, ith, nth);
int dst_n_block_begin = n_start;
int dst_n_block_size = n_end - dst_n_block_begin;
if (dst_n_block_size <= 0) return;

auto tile_ptr = [](ggml_bf16_t* base, int total_n, int total_k, int abs_n, int abs_k) {
int n_block_begin = abs_n / N_BLOCK * N_BLOCK;
int n_within = abs_n - n_block_begin;
int n_block_size = std::min(N_BLOCK, total_n - n_block_begin);
int k_block_begin = abs_k / K_BLOCK * K_BLOCK;
int k_within = abs_k - k_block_begin;
return base + (size_t)n_block_begin * total_k + (size_t)k_block_begin * n_block_size +
(size_t)n_within * std::min(K_BLOCK, total_k - k_block_begin) + (size_t)k_within * N_STEP;
};

alignas(64) ggml_bf16_t src_tile[N_STEP * K_STEP];
alignas(64) ggml_bf16_t dst_tile[N_STEP * K_STEP];
for (int dst_n = 0; dst_n < dst_n_block_size; dst_n += N_STEP) {
for (int dst_k_block = 0; dst_k_block < k; dst_k_block += K_BLOCK) {
int dst_k_block_size = std::min(K_BLOCK, k - dst_k_block);
for (int dst_k = 0; dst_k < dst_k_block_size; dst_k += K_STEP) {
int abs_dst_n = dst_n_block_begin + dst_n;
int abs_dst_k = dst_k_block + dst_k;
ggml_bf16_t* src_ptr = tile_ptr(src.b, src.n, src.k, abs_dst_k, abs_dst_n);
memcpy(src_tile, src_ptr, sizeof(src_tile));
transpose_16x16_32bit((__m512i*)src_tile);
transpose_16x16_32bit((__m512i*)(src_tile + TILE_N * K_STEP));

for (int i = 0; i < N_STEP; i++) {
for (int j = 0; j < K_STEP; j++) {
dst_tile[j * K_STEP + i] = src_tile[i * K_STEP + j];
}
}
Comment on lines +252 to +256

@gemini-code-assist gemini-code-assist Bot Jul 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

In from_bb_transposed, the destination tile dst_tile has dimensions K_STEP (rows) by N_STEP (columns). Therefore, its row stride is N_STEP. The indexing for dst_tile should use j * N_STEP + i instead of j * K_STEP + i. Although K_STEP and N_STEP are currently both 32, using N_STEP is conceptually correct and prevents potential out-of-bounds or layout corruption bugs if the step sizes ever differ in future configurations.

          for (int i = 0; i < N_STEP; i++) {
            for (int j = 0; j < K_STEP; j++) {
              dst_tile[j * N_STEP + i] = src_tile[i * K_STEP + j];
            }
          }

transpose_16x16_32bit((__m512i*)dst_tile);
transpose_16x16_32bit((__m512i*)(dst_tile + TILE_N * K_STEP));
ggml_bf16_t* dst_ptr = tile_ptr(b, n, k, abs_dst_n, abs_dst_k);
memcpy(dst_ptr, dst_tile, sizeof(dst_tile));
}
}
}
}
ggml_bf16_t* get_submat(int n, int k, int n_begin, int k_begin) {
int n_block_begin = n_begin / N_BLOCK * N_BLOCK;
n_begin -= n_block_begin;
Expand Down
Loading
Loading