Skip to content

Commit dc6394e

Browse files
ouqingliangouqingliang
authored
[fix]: fix moe hpp bug. (kvcache-ai#1780)
fix moe hpp init bug.
1 parent ad7674a commit dc6394e

File tree

5 files changed

+58
-51
lines changed

5 files changed

+58
-51
lines changed

‎kt-kernel/operators/amx/awq-moe.hpp‎

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -31,16 +31,16 @@ class AMX_AWQ_MOE_TP : public AMX_MOE_BASE<T, AMX_AWQ_MOE_TP<T>> {
3131
private:
3232
using Base = AMX_MOE_BASE<T, AMX_AWQ_MOE_TP<T>>;
3333
using Base::config_;
34-
using Base::tp_part_idx;
35-
using Base::gate_bb_;
36-
using Base::up_bb_;
37-
using Base::down_bb_;
38-
using Base::gate_up_ba_;
39-
using Base::gate_bc_;
40-
using Base::up_bc_;
4134
using Base::down_ba_;
35+
using Base::down_bb_;
4236
using Base::down_bc_;
37+
using Base::gate_bb_;
38+
using Base::gate_bc_;
39+
using Base::gate_up_ba_;
4340
using Base::m_local_num_;
41+
using Base::tp_part_idx;
42+
using Base::up_bb_;
43+
using Base::up_bc_;
4444

4545
std::filesystem::path prefix;
4646

@@ -265,7 +265,7 @@ class AMX_AWQ_MOE_TP : public AMX_MOE_BASE<T, AMX_AWQ_MOE_TP<T>> {
265265
config_.quant_config.group_size)) != 0) {
266266
printf("verify error\n");
267267
for (size_t i = 0; i < T::BufferB::required_size(config_.hidden_size, config_.intermediate_size,
268-
config_.quant_config.group_size);
268+
config_.quant_config.group_size);
269269
++i) {
270270
if (verify_bb[i] != check_bb[i]) {
271271
printf("Difference at byte %zu: verify_bb_%d[%zu] = %02x, check_bb[%zu] = %02x\n", i, compare_expers, i,
@@ -393,19 +393,21 @@ class AMX_AWQ_MOE_TP : public AMX_MOE_BASE<T, AMX_AWQ_MOE_TP<T>> {
393393

394394
AMX_AWQ_MOE_TP() = default;
395395

396-
AMX_AWQ_MOE_TP(GeneralMOEConfig config, int tp_part_idx_ = 0) : Base(config, tp_part_idx_) {
396+
AMX_AWQ_MOE_TP(GeneralMOEConfig config, int tp_part_idx_ = 0) : Base(config, tp_part_idx_) {}
397+
398+
void derived_init() {
397399
auto& quant_config = config_.quant_config;
398400
if (quant_config.group_size == 0 || !quant_config.zero_point) {
399401
throw std::runtime_error("AWQ-Quantization AMX MoE only support KGroup Int4_1");
400402
}
401403

402-
printf("Creating AMX_AWQ_MOE_TP %d at numa %d\n", tp_part_idx_, numa_node_of_cpu(sched_getcpu()));
404+
printf("Creating AMX_AWQ_MOE_TP %d at numa %d\n", tp_part_idx, numa_node_of_cpu(sched_getcpu()));
403405

404406
auto& load = config_.load;
405407
auto& save = config_.save;
406408

407409
prefix = config_.path;
408-
prefix = prefix / ("_layer_" + std::to_string(config_.layer_idx)) / ("_numa_" + std::to_string(tp_part_idx_));
410+
prefix = prefix / ("_layer_" + std::to_string(config_.layer_idx)) / ("_numa_" + std::to_string(tp_part_idx));
409411
if (save) {
410412
std::cout << "Creating " << prefix << std::endl;
411413
std::filesystem::create_directories(prefix);
@@ -431,9 +433,7 @@ class AMX_AWQ_MOE_TP : public AMX_MOE_BASE<T, AMX_AWQ_MOE_TP<T>> {
431433
size_t buffer_b_required_size_impl(size_t n, size_t k) const {
432434
return T::BufferB::required_size(n, k, config_.quant_config.group_size);
433435
}
434-
size_t buffer_c_required_size_impl(size_t m, size_t n) const {
435-
return T::BufferC::required_size(m, n);
436-
}
436+
size_t buffer_c_required_size_impl(size_t m, size_t n) const { return T::BufferC::required_size(m, n); }
437437

438438
std::shared_ptr<typename T::BufferA> make_buffer_a_impl(size_t m, size_t k, void* data) const {
439439
return std::make_shared<typename T::BufferA>(m, k, config_.quant_config.group_size, data);

‎kt-kernel/operators/amx/fp8-moe.hpp‎

Lines changed: 6 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -13,16 +13,6 @@
1313

1414
// #define DEBUG_FP8_MOE
1515

16-
#include <immintrin.h>
17-
18-
#include <algorithm>
19-
#include <cstddef>
20-
#include <cstdint>
21-
#include <cstring>
22-
#include <memory>
23-
#include <string>
24-
#include <vector>
25-
2616
#include "la/amx_raw_buffers.hpp"
2717
#include "la/amx_raw_kernels.hpp"
2818
#include "moe_base.hpp"
@@ -57,12 +47,15 @@ class AMX_FP8_MOE_TP : public AMX_MOE_BASE<T, AMX_FP8_MOE_TP<T>> {
5747
AMX_FP8_MOE_TP() = default;
5848

5949
AMX_FP8_MOE_TP(GeneralMOEConfig config, int tp_part_idx_ = 0) : Base(config, tp_part_idx_) {
50+
// Initialization now happens in derived_init() which is called by base constructor
51+
}
52+
53+
void derived_init() {
6054
auto& quant_config = config_.quant_config;
6155
if (quant_config.group_size == 0 || quant_config.zero_point) {
62-
throw std::runtime_error("KT-Kernel fp8 MoE only support block-wise FP8. group_size = %d, zero_point = %d",
63-
quant_config.group_size, quant_config.zero_point);
56+
throw std::runtime_error("KT-Kernel fp8 MoE only support block-wise FP8");
6457
}
65-
printf("Created AMX_FP8_MOE_TP %d at numa %d\n", tp_part_idx_, numa_node_of_cpu(sched_getcpu()));
58+
printf("Created AMX_FP8_MOE_TP %d at numa %d\n", tp_part_idx, numa_node_of_cpu(sched_getcpu()));
6659
}
6760

6861
~AMX_FP8_MOE_TP() = default;

‎kt-kernel/operators/amx/k2-moe.hpp‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,12 +45,14 @@ class AMX_K2_MOE_TP : public AMX_MOE_BASE<T, AMX_K2_MOE_TP<T>> {
4545

4646
AMX_K2_MOE_TP() = default;
4747

48-
AMX_K2_MOE_TP(GeneralMOEConfig config, int tp_part_idx_ = 0) : Base(config, tp_part_idx_) {
48+
AMX_K2_MOE_TP(GeneralMOEConfig config, int tp_part_idx_ = 0) : Base(config, tp_part_idx_) {}
49+
50+
void derived_init() {
4951
auto& quant_config = config_.quant_config;
5052
if (quant_config.group_size == 0 || quant_config.zero_point) {
5153
throw std::runtime_error("Kimi-K2 MoE only support KGroup Int4");
5254
}
53-
printf("Creating AMX_K2_MOE_TP %d at numa %d\n", tp_part_idx_, numa_node_of_cpu(sched_getcpu()));
55+
printf("Creating AMX_K2_MOE_TP %d at numa %d\n", tp_part_idx, numa_node_of_cpu(sched_getcpu()));
5456
}
5557

5658
~AMX_K2_MOE_TP() = default;

‎kt-kernel/operators/amx/moe.hpp‎

Lines changed: 20 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -21,18 +21,16 @@ class AMX_MOE_TP : public AMX_MOE_BASE<T, AMX_MOE_TP<T>> {
2121
private:
2222
using Base = AMX_MOE_BASE<T, AMX_MOE_TP<T>>;
2323
using Base::config_;
24-
using Base::tp_part_idx;
25-
using Base::gate_bb_;
26-
using Base::up_bb_;
27-
using Base::down_bb_;
28-
using Base::gate_up_ba_;
29-
using Base::gate_bc_;
30-
using Base::up_bc_;
3124
using Base::down_ba_;
25+
using Base::down_bb_;
3226
using Base::down_bc_;
27+
using Base::gate_bb_;
28+
using Base::gate_bc_;
29+
using Base::gate_up_ba_;
3330
using Base::m_local_num_;
34-
35-
std::filesystem::path prefix;
31+
using Base::tp_part_idx;
32+
using Base::up_bb_;
33+
using Base::up_bc_;
3634

3735
void* gate_proj_; // [expert_num * intermediate_size * hidden_size ( /32 if
3836
// quantized)]
@@ -140,11 +138,15 @@ class AMX_MOE_TP : public AMX_MOE_BASE<T, AMX_MOE_TP<T>> {
140138
AMX_MOE_TP() = default;
141139

142140
AMX_MOE_TP(GeneralMOEConfig config, int tp_part_idx = 0) : Base(config, tp_part_idx) {
141+
// Initialization now happens in derived_init() which is called by base constructor
142+
}
143+
144+
void derived_init() {
143145
printf("Creating AMX_MOE_TP %d at numa %d\n", tp_part_idx, numa_node_of_cpu(sched_getcpu()));
144146
auto& load = config_.load;
145147
auto& save = config_.save;
146148

147-
prefix = config_.path;
149+
std::filesystem::path prefix = config_.path;
148150
prefix = prefix / ("_layer_" + std::to_string(config_.layer_idx)) / ("_numa_" + std::to_string(tp_part_idx));
149151
if (save) {
150152
std::cout << "Creating " << prefix << std::endl;
@@ -169,15 +171,9 @@ class AMX_MOE_TP : public AMX_MOE_BASE<T, AMX_MOE_TP<T>> {
169171
// CRTP buffer creation - no group_size
170172
// ============================================================================
171173

172-
size_t buffer_a_required_size_impl(size_t m, size_t k) const {
173-
return T::BufferA::required_size(m, k);
174-
}
175-
size_t buffer_b_required_size_impl(size_t n, size_t k) const {
176-
return T::BufferB::required_size(n, k);
177-
}
178-
size_t buffer_c_required_size_impl(size_t m, size_t n) const {
179-
return T::BufferC::required_size(m, n);
180-
}
174+
size_t buffer_a_required_size_impl(size_t m, size_t k) const { return T::BufferA::required_size(m, k); }
175+
size_t buffer_b_required_size_impl(size_t n, size_t k) const { return T::BufferB::required_size(n, k); }
176+
size_t buffer_c_required_size_impl(size_t m, size_t n) const { return T::BufferC::required_size(m, n); }
181177

182178
std::shared_ptr<typename T::BufferA> make_buffer_a_impl(size_t m, size_t k, void* data) const {
183179
return std::make_shared<typename T::BufferA>(m, k, data);
@@ -260,6 +256,9 @@ class AMX_MOE_TP : public AMX_MOE_BASE<T, AMX_MOE_TP<T>> {
260256
} else {
261257
int nth = T::recommended_nth(config_.intermediate_size);
262258
static uint8_t mat_type_all = 3, mat_split = 1;
259+
std::filesystem::path prefix = config_.path;
260+
prefix = prefix / ("_layer_" + std::to_string(config_.layer_idx)) / ("_numa_" + std::to_string(tp_part_idx));
261+
263262
if (config_.load) {
264263
std::cout << "Loading from " << prefix << std::endl;
265264
for (int task_id = 0; task_id < config_.expert_num * mat_type_all * mat_split; task_id++) {
@@ -335,7 +334,7 @@ class AMX_MOE_TP : public AMX_MOE_BASE<T, AMX_MOE_TP<T>> {
335334
if (config_.save) {
336335
pool->do_work_stealing_job(
337336
config_.expert_num * mat_type_all, nullptr,
338-
[this, physical_to_logical_map](int task_id) {
337+
[this, physical_to_logical_map, prefix](int task_id) {
339338
int64_t expert_idx = task_id / mat_type_all;
340339
expert_idx = expert_map(physical_to_logical_map, expert_idx);
341340
uint8_t mat_class = task_id % mat_type_all;
@@ -426,7 +425,7 @@ class TP_MOE<AMX_MOE_TP<K>> : public TP_MOE<AMX_MOE_BASE<K, AMX_MOE_TP<K>>> {
426425

427426
this->weights_loaded = true;
428427
} else if (config.path != "") {
429-
printf("TP Load from file\n");
428+
printf("TP Load from file %s\n", config.path.c_str());
430429
DO_TPS_LOAD_WEIGHTS(pool);
431430
this->weights_loaded = true;
432431
} else {

‎kt-kernel/operators/amx/moe_base.hpp‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
#include <cstring>
2424
#include <filesystem>
2525
#include <fstream>
26+
#include <memory>
2627
#include <string>
2728
#include <utility>
2829
#include <vector>
@@ -78,7 +79,10 @@ class AMX_MOE_BASE {
7879
using output_t = float;
7980
static constexpr double ELEMENT_SIZE = T::ELEMENT_SIZE;
8081

81-
AMX_MOE_BASE(GeneralMOEConfig config, int tp_part_idx_) : tp_part_idx(tp_part_idx_), config_(config) { init(); }
82+
AMX_MOE_BASE(GeneralMOEConfig config, int tp_part_idx_) : tp_part_idx(tp_part_idx_), config_(config) {
83+
init();
84+
derived()->derived_init();
85+
}
8286

8387
void init() {
8488
if (config_.load && config_.path == "") {
@@ -639,6 +643,15 @@ class AMX_MOE_BASE {
639643
Derived* derived() { return static_cast<Derived*>(this); }
640644
const Derived* derived_const() const { return static_cast<const Derived*>(this); }
641645

646+
// ============================================================================
647+
// Derived class initialization hook
648+
// Called after base class init() completes, allows derived classes to perform
649+
// their own initialization that depends on base class being fully initialized
650+
// ============================================================================
651+
void derived_init() {
652+
// Default implementation does nothing - derived classes can override
653+
}
654+
642655
// ============================================================================
643656
// Virtual points for buffer creation and size calculation
644657
// Default implementations use group_size (for KGroup quantization like K2)

0 commit comments

Comments
 (0)