Skip to content

Commit 56b3173

Browse files
kmorennvkmorennv
committed
Add MMQ fusion test cases with weight-scale epilogue support
1 parent 66fcc8a commit 56b3173

1 file changed

Lines changed: 218 additions & 4 deletions

File tree

‎tests/test-backend-ops.cpp‎

Lines changed: 218 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,35 @@
4545
#include <vector>
4646
#include <unordered_map>
4747

48+
#ifdef _WIN32
49+
static void set_environment_variable(const char * name, const char * value) {
50+
_putenv_s(name, value ? value : "");
51+
}
52+
#else
53+
static void set_environment_variable(const char * name, const char * value) {
54+
if (value) {
55+
setenv(name, value, 1);
56+
} else {
57+
unsetenv(name);
58+
}
59+
}
60+
#endif
61+
62+
struct scoped_environment_variable {
63+
const char * name;
64+
bool had_value;
65+
std::string old_value;
66+
67+
scoped_environment_variable(const char * name, const char * value)
68+
: name(name), had_value(getenv(name) != nullptr), old_value(had_value ? getenv(name) : "") {
69+
set_environment_variable(name, value);
70+
}
71+
72+
~scoped_environment_variable() {
73+
set_environment_variable(name, had_value ? old_value.c_str() : nullptr);
74+
}
75+
};
76+
4877
#ifdef __EMSCRIPTEN__
4978
# define N_THREADS 1
5079
#else
@@ -1313,10 +1342,10 @@ struct test_case {
13131342
}
13141343
}
13151344

1316-
test_status_t eval(ggml_backend_t backend1,
1317-
ggml_backend_t backend2,
1318-
const char * op_names_filter,
1319-
printer * output_printer) {
1345+
virtual test_status_t eval(ggml_backend_t backend1,
1346+
ggml_backend_t backend2,
1347+
const char * op_names_filter,
1348+
printer * output_printer) {
13201349
mode = MODE_TEST;
13211350

13221351
ggml_init_params params = {
@@ -4410,6 +4439,174 @@ static void init_mul_mat_id_tensors(ggml_context * ctx, int n_mats) {
44104439
}
44114440
}
44124441

4442+
// GGML_TYPE_NVFP4 MMQ + weight-scale epilogue
4443+
struct test_mul_mat_mmq_fusion : public test_case {
4444+
const bool use_id;
4445+
const int64_t m;
4446+
const int64_t n;
4447+
const int64_t k;
4448+
const int n_mats;
4449+
const int n_used;
4450+
4451+
test_mul_mat_mmq_fusion(
4452+
bool use_id, int64_t m = 32, int64_t n = 64, int64_t k = 256, int n_mats = 16, int n_used = 8)
4453+
: use_id(use_id), m(m), n(n), k(k), n_mats(n_mats), n_used(n_used) {
4454+
GGML_ASSERT(n_used <= n_mats);
4455+
}
4456+
4457+
std::string vars() override {
4458+
return VARS_TO_STR6(use_id, m, n, k, n_mats, n_used);
4459+
}
4460+
4461+
std::string op_desc(ggml_tensor * t) override {
4462+
GGML_UNUSED(t);
4463+
return "MUL_MAT_MMQ_FUSION";
4464+
}
4465+
4466+
ggml_tensor * build_scale_id(
4467+
ggml_context * ctx, ggml_tensor * scale, ggml_tensor * ids, ggml_tensor * out) {
4468+
ggml_tensor * s = ggml_reshape_3d(ctx, scale, 1, n_mats, 1);
4469+
s = ggml_repeat_4d(ctx, s, 1, n_mats, m, 1);
4470+
s = ggml_get_rows(ctx, s, ids);
4471+
return ggml_mul(ctx, out, s);
4472+
}
4473+
4474+
ggml_tensor * build_graph(ggml_context * ctx) override {
4475+
if (!use_id) {
4476+
ggml_tensor * weights = ggml_new_tensor_2d(ctx, GGML_TYPE_NVFP4, k, n);
4477+
ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, k, m);
4478+
ggml_tensor * scale = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
4479+
4480+
return ggml_mul(ctx, ggml_mul_mat(ctx, weights, input), scale);
4481+
}
4482+
4483+
ggml_tensor * weights = ggml_new_tensor_3d(ctx, GGML_TYPE_NVFP4, k, n, n_mats);
4484+
ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, k, n_used, m);
4485+
ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_mats, m);
4486+
if (n_used != n_mats) {
4487+
ids = ggml_view_2d(ctx, ids, n_used, m, ids->nb[1], 0);
4488+
}
4489+
ggml_tensor * scale = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_mats);
4490+
4491+
return build_scale_id(ctx, scale, ids, ggml_mul_mat_id(ctx, weights, input, ids));
4492+
}
4493+
4494+
void initialize_tensors(ggml_context * ctx) override {
4495+
if (use_id) {
4496+
init_mul_mat_id_tensors(ctx, n_mats);
4497+
} else {
4498+
test_case::initialize_tensors(ctx);
4499+
}
4500+
}
4501+
4502+
test_status_t eval(ggml_backend_t backend1,
4503+
ggml_backend_t backend2,
4504+
const char * op_names_filter,
4505+
printer * output_printer) override {
4506+
GGML_UNUSED(backend2);
4507+
4508+
if (strncmp(ggml_backend_name(backend1), "CUDA", 4) != 0) {
4509+
return test_status_t::NOT_SUPPORTED;
4510+
}
4511+
4512+
mode = MODE_TEST;
4513+
ggml_init_params params = {
4514+
/* .mem_size = */ ggml_tensor_overhead()*128 + ggml_graph_overhead(),
4515+
/* .mem_base = */ nullptr,
4516+
/* .no_alloc = */ true,
4517+
};
4518+
ggml_context * ctx = ggml_init(params);
4519+
GGML_ASSERT(ctx);
4520+
4521+
ggml_cgraph * graph = ggml_new_graph(ctx);
4522+
ggml_tensor * out = build_graph(ctx);
4523+
current_op_name = op_desc(out);
4524+
if (!matches_filter(out, op_names_filter)) {
4525+
ggml_free(ctx);
4526+
return test_status_t::SKIPPED;
4527+
}
4528+
ggml_build_forward_expand(graph, out);
4529+
4530+
ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend1);
4531+
if (!buffer) {
4532+
ggml_free(ctx);
4533+
return test_status_t::FAIL;
4534+
}
4535+
if (use_id) {
4536+
ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
4537+
}
4538+
initialize_tensors(ctx);
4539+
ggml_backend_synchronize(backend1);
4540+
4541+
ggml_backend_t backend_fused = ggml_backend_dev_init(ggml_backend_get_device(backend1), nullptr);
4542+
GGML_ASSERT(backend_fused);
4543+
4544+
ggml_context * ctx_fused = ggml_init(params);
4545+
GGML_ASSERT(ctx_fused);
4546+
ggml_cgraph * graph_fused = ggml_new_graph(ctx_fused);
4547+
ggml_tensor * out_fused = build_graph(ctx_fused);
4548+
ggml_build_forward_expand(graph_fused, out_fused);
4549+
4550+
ggml_backend_buffer_t buffer_fused = ggml_backend_alloc_ctx_tensors(ctx_fused, backend_fused);
4551+
GGML_ASSERT(buffer_fused);
4552+
if (use_id) {
4553+
ggml_backend_buffer_set_usage(buffer_fused, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
4554+
}
4555+
ggml_tensor * tensor_unfused = ggml_get_first_tensor(ctx);
4556+
ggml_tensor * tensor_fused = ggml_get_first_tensor(ctx_fused);
4557+
while (tensor_unfused && tensor_fused) {
4558+
GGML_ASSERT(tensor_unfused->type == tensor_fused->type);
4559+
GGML_ASSERT(ggml_are_same_shape(tensor_unfused, tensor_fused));
4560+
if (!ggml_is_view_op(tensor_unfused->op)) {
4561+
ggml_backend_tensor_copy(tensor_unfused, tensor_fused);
4562+
}
4563+
tensor_unfused = ggml_get_next_tensor(ctx, tensor_unfused);
4564+
tensor_fused = ggml_get_next_tensor(ctx_fused, tensor_fused);
4565+
}
4566+
GGML_ASSERT(!tensor_unfused && !tensor_fused);
4567+
ggml_backend_synchronize(backend_fused);
4568+
4569+
ggml_status status_fused;
4570+
{
4571+
scoped_environment_variable enable_fusion("GGML_CUDA_FUSE_WS", "1");
4572+
scoped_environment_variable disable_fusion("GGML_CUDA_NO_FUSE_WS", nullptr);
4573+
status_fused = ggml_backend_graph_compute(backend_fused, graph_fused);
4574+
}
4575+
4576+
ggml_status status_unfused;
4577+
{
4578+
scoped_environment_variable enable_fusion("GGML_CUDA_FUSE_WS", nullptr);
4579+
scoped_environment_variable disable_fusion("GGML_CUDA_NO_FUSE_WS", "1");
4580+
status_unfused = ggml_backend_graph_compute(backend1, graph);
4581+
}
4582+
4583+
std::vector<float> data_unfused(ggml_nelements(out));
4584+
std::vector<float> data_fused(ggml_nelements(out_fused));
4585+
ggml_backend_tensor_get(out, data_unfused.data(), 0, ggml_nbytes(out));
4586+
ggml_backend_tensor_get(out_fused, data_fused.data(), 0, ggml_nbytes(out_fused));
4587+
4588+
double max_diff = 0.0;
4589+
for (size_t i = 0; i < data_unfused.size(); ++i) {
4590+
max_diff = std::max(max_diff, (double) std::fabs(data_unfused[i] - data_fused[i]));
4591+
}
4592+
const bool bit_exact = memcmp(data_unfused.data(), data_fused.data(), ggml_nbytes(out)) == 0;
4593+
const bool passed = status_unfused == GGML_STATUS_SUCCESS && status_fused == GGML_STATUS_SUCCESS && bit_exact;
4594+
4595+
ggml_backend_buffer_free(buffer_fused);
4596+
ggml_backend_free(backend_fused);
4597+
ggml_free(ctx_fused);
4598+
ggml_backend_buffer_free(buffer);
4599+
ggml_free(ctx);
4600+
4601+
char error_buf[64];
4602+
snprintf(error_buf, sizeof(error_buf), "bit_exact=%d, max_diff=%.9f", bit_exact, max_diff);
4603+
const std::string error = passed ? "" : error_buf;
4604+
print_test_result_locked(output_printer,
4605+
test_result(ggml_backend_name(backend1), current_op_name, vars(), "test", true, passed, error));
4606+
return passed ? test_status_t::OK : test_status_t::FAIL;
4607+
}
4608+
};
4609+
44134610
// GGML_OP_MUL_MAT_ID
44144611
struct test_mul_mat_id : public test_case {
44154612
const ggml_type type_a;
@@ -8016,6 +8213,23 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
80168213
std::vector<std::unique_ptr<test_case>> test_cases;
80178214
std::default_random_engine rng(0);
80188215

8216+
test_cases.emplace_back(new test_mul_mat_mmq_fusion(false));
8217+
test_cases.emplace_back(new test_mul_mat_mmq_fusion(true));
8218+
for (bool use_id : { false, true }) {
8219+
test_cases.emplace_back(new test_mul_mat_mmq_fusion(use_id, 11, 256, 4096));
8220+
test_cases.emplace_back(new test_mul_mat_mmq_fusion(use_id, 11, 4096, 256));
8221+
}
8222+
for (int64_t m : { 2, 4, 12 }) {
8223+
test_cases.emplace_back(new test_mul_mat_mmq_fusion(false, m, 8192, 2048));
8224+
test_cases.emplace_back(new test_mul_mat_mmq_fusion(false, m, 32, 2048));
8225+
test_cases.emplace_back(new test_mul_mat_mmq_fusion(false, m, 4096, 2048));
8226+
test_cases.emplace_back(new test_mul_mat_mmq_fusion(false, m, 2048, 4096));
8227+
test_cases.emplace_back(new test_mul_mat_mmq_fusion(false, m, 512, 2048));
8228+
test_cases.emplace_back(new test_mul_mat_mmq_fusion(false, m, 2048, 512));
8229+
test_cases.emplace_back(new test_mul_mat_mmq_fusion(true, m, 512, 2048, 256, 8));
8230+
test_cases.emplace_back(new test_mul_mat_mmq_fusion(true, m, 2048, 512, 256, 8));
8231+
}
8232+
80198233
// unary ops
80208234
for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32}) {
80218235
for (int v : {0, 1}) {

0 commit comments

Comments
 (0)