Jeff Bolz
b365c3ff01
I updated test_topk_moe to more closely match llm_graph_context::build_moe_ffn and added coverage for exp_probs_b and some other missing combinations. This exposed a bug in both CUDA and Vulkan backends where they were assuming the input to argsort and the input to get_rows are the same. I'd like to optimize this graph in another change, but for now just get it functional. CUDA also had a bug where it got n_experts from the wrong place, leading to GGML_ASSERT failures in some of the new tests.
22 lines
1023 B
Plaintext
22 lines
1023 B
Plaintext
#include "common.cuh"
|
|
#include "ggml.h"
|
|
|
|
#include <initializer_list>
|
|
|
|
void ggml_cuda_op_topk_moe(ggml_backend_cuda_context & ctx,
|
|
const ggml_tensor * logits,
|
|
ggml_tensor * weights,
|
|
ggml_tensor * ids,
|
|
const bool with_norm,
|
|
const bool delayed_softmax = false,
|
|
ggml_tensor * weight_clamp = nullptr);
|
|
|
|
bool ggml_cuda_should_use_topk_moe(const ggml_tensor * softmax,
|
|
const ggml_tensor * weights,
|
|
const ggml_tensor * get_rows,
|
|
const ggml_tensor * argsort,
|
|
const ggml_tensor * clamp,
|
|
int n_expert);
|
|
|
|
std::initializer_list<enum ggml_op> ggml_cuda_topk_moe_ops(bool with_norm, bool delayed_softmax = false);
|