Skip to content

Fix moe bug. #1783

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

Merged
ouqingliang merged 4 commits into from
Jan 5, 2026
Merged

Fix moe bug. #1783

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
f2baa0c
[fix]: fix moe.hpp load from file bug.
ouqingliang Dec 29, 2025
f8bd1fa
[fix]: fix all moe hpp init bug.
ouqingliang Jan 4, 2026
8bd51e6
[fix]: fix moe & awq-moe ug.
ouqingliang Jan 5, 2026
942b433
Merge branch 'main' into fix-bug-oql
ouqingliang Jan 5, 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
7 changes: 4 additions & 3 deletions kt-kernel/operators/amx/awq-moe.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,8 +42,6 @@ class AMX_AWQ_MOE_TP : public AMX_MOE_BASE<T, AMX_AWQ_MOE_TP<T>> {
using Base::up_bb_;
using Base::up_bc_;

std::filesystem::path prefix;

#ifdef CHECK
char verify_bb[100000000];
char check_bb[100000000];
Expand Down Expand Up @@ -406,7 +404,7 @@ class AMX_AWQ_MOE_TP : public AMX_MOE_BASE<T, AMX_AWQ_MOE_TP<T>> {
auto& load = config_.load;
auto& save = config_.save;

prefix = config_.path;
std::filesystem::path prefix = config_.path;
prefix = prefix / ("_layer_" + std::to_string(config_.layer_idx)) / ("_numa_" + std::to_string(tp_part_idx));
if (save) {
std::cout << "Creating " << prefix << std::endl;
Expand Down Expand Up @@ -498,6 +496,9 @@ class AMX_AWQ_MOE_TP : public AMX_MOE_BASE<T, AMX_AWQ_MOE_TP<T>> {
throw std::runtime_error("AMX load weights from gate_projs is not supported");
} else {
int nth = T::recommended_nth(config_.intermediate_size);
std::filesystem::path prefix = config_.path;
prefix = prefix / ("_layer_" + std::to_string(config_.layer_idx)) / ("_numa_" + std::to_string(tp_part_idx));
Comment on lines +499 to +500
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Jan 5, 2026

Choose a reason for hiding this comment

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

medium

This logic for constructing the prefix path is a duplicate of the code in derived_init() (lines 407-408). This code duplication can lead to maintenance issues, as any change to the path structure would need to be updated in two places.

To improve this, you could re-introduce prefix as a private member variable, initialize it once in derived_init(), and then reuse it in this load_weights() method. This would centralize the path construction logic.


if (config_.load) {
throw std::runtime_error("AMX load weights from file is not supported");
}
Expand Down
11 changes: 0 additions & 11 deletions kt-kernel/operators/amx/moe.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,13 +32,6 @@ class AMX_MOE_TP : public AMX_MOE_BASE<T, AMX_MOE_TP<T>> {
using Base::up_bb_;
using Base::up_bc_;

void* gate_proj_; // [expert_num * intermediate_size * hidden_size ( /32 if
// quantized)]
void* up_proj_; // [expert_num * intermediate_size * hidden_size ( /32 if
// quantized)]
void* down_proj_; // [expert_num * hidden_size * intermediate_size ( /32 if
// quantized)]

#ifdef CHECK
char verify_bb[100000000];
char check_bb[100000000];
Expand Down Expand Up @@ -159,10 +152,6 @@ class AMX_MOE_TP : public AMX_MOE_BASE<T, AMX_MOE_TP<T>> {
throw std::runtime_error("Path not found: " + prefix.string());
}
}

gate_proj_ = config_.gate_proj;
up_proj_ = config_.up_proj;
down_proj_ = config_.down_proj;
}

~AMX_MOE_TP() = default;
Expand Down
Loading