Skip to content

GitLab

Unverified Commit 900a7f7c authored by Atream's avatar Atream Committed by GitHub
Browse files

Merge pull request #1271 from kvcache-ai/fix-AMX 

fix AMX
parents 3f14e311 b22cded8
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
csrc/ktransformers_ext/operators/amx/moe.hpp
View file @ 900a7f7c
...@@ -272,8 +272,8 @@ public: ...@@ -272,8 +272,8 @@ public:
void forward(int qlen, int k, const uint64_t *expert_ids, const float *weights, const void *input, void *output, void forward(int qlen, int k, const uint64_t *expert_ids, const float *weights, const void *input, void *output,
int *batch_size_tensor, Backend *backend) { int *batch_size_tensor, Backend *backend) {
bool use_amx = (qlen > 4 * config_.expert_num / config_.routed_expert_num);
qlen = batch_size_tensor[0]; qlen = batch_size_tensor[0];
bool use_amx = (qlen > 4 * config_.expert_num / config_.routed_expert_num);
int activated_expert = 0; int activated_expert = 0;
for (int i = 0; i < config_.expert_num; i++) { for (int i = 0; i < config_.expert_num; i++) {
m_local_num_[i] = 0; m_local_num_[i] = 0;
...@@ -395,4 +395,4 @@ public: ...@@ -395,4 +395,4 @@ public:
} }
}; };
#endif #endif
\ No newline at end of file
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment