|
8 | 8 | #include "unary-ops.h" |
9 | 9 | #include "vec.h" |
10 | 10 |
|
| 11 | +// MAD-223 G.4.c — ml8 dequant + block layout for GGML_OP_ML8_MUL_MAT. |
| 12 | +#define GGML_COMMON_DECL_CPP |
| 13 | +#include "ggml-common.h" |
| 14 | +#include "ggml-quants.h" |
| 15 | + |
11 | 16 | #include <algorithm> |
12 | 17 | #include <cfloat> |
13 | 18 | #include <cmath> |
@@ -11405,6 +11410,172 @@ void ggml_compute_forward_opt_step_sgd(const ggml_compute_params * params, ggml_ |
11405 | 11410 | } |
11406 | 11411 | } |
11407 | 11412 |
|
| 11413 | +// ─── ggml_compute_forward_ml8_mul_mat (MAD-223 G.4.c) ───────────────────── |
| 11414 | +// |
| 11415 | +// CPU dispatch for GGML_OP_ML8_MUL_MAT. Same math as the previous custom_4d |
| 11416 | +// callback in ggml-ml8.c — dequantize w block-by-block via |
| 11417 | +// dequantize_row_ml8_4_with_lut into a per-thread fp32 scratch row, then dot |
| 11418 | +// against x. Tiled across N-rows by (ith, nth). |
| 11419 | +// |
| 11420 | +// Tensor layouts (set by ggml_ml8_mul_mat): |
| 11421 | +// dst->src[0] = w (GGML_TYPE_ML8_4 , [K, N]) |
| 11422 | +// dst->src[1] = centroids (GGML_TYPE_F8_E4M3, [16, K/QK_ML8]) |
| 11423 | +// dst->src[2] = x (GGML_TYPE_F32 , [K, M]) |
| 11424 | +// dst = y (GGML_TYPE_F32 , [N, M]) |
| 11425 | +void ggml_compute_forward_ml8_mul_mat(const ggml_compute_params * params, ggml_tensor * dst) { |
| 11426 | + const ggml_tensor * w = dst->src[0]; |
| 11427 | + const ggml_tensor * lut = dst->src[1]; |
| 11428 | + const ggml_tensor * x = dst->src[2]; |
| 11429 | + GGML_ASSERT(w && lut && x); |
| 11430 | + GGML_ASSERT(w->type == GGML_TYPE_ML8_4); |
| 11431 | + GGML_ASSERT(lut->type == GGML_TYPE_F8_E4M3); |
| 11432 | + GGML_ASSERT(x->type == GGML_TYPE_F32); |
| 11433 | + GGML_ASSERT(dst->type == GGML_TYPE_F32); |
| 11434 | + |
| 11435 | + const int64_t K = w->ne[0]; |
| 11436 | + const int64_t N = w->ne[1]; |
| 11437 | + const int64_t M = x->ne[1]; |
| 11438 | + GGML_ASSERT(x->ne[0] == K); |
| 11439 | + GGML_ASSERT(K % QK_ML8 == 0); |
| 11440 | + const int64_t n_groups_k = K / QK_ML8; |
| 11441 | + GGML_ASSERT(lut->ne[0] == 16 && lut->ne[1] == n_groups_k); |
| 11442 | + |
| 11443 | + const block_ml8_4 * w_blocks = (const block_ml8_4 *) w->data; |
| 11444 | + const uint8_t * lut_fp8 = (const uint8_t *) lut->data; |
| 11445 | + const float * x_data = (const float *) x->data; |
| 11446 | + float * y_data = (float *) dst->data; |
| 11447 | + |
| 11448 | + const int ith = params->ith; |
| 11449 | + const int nth = params->nth; |
| 11450 | + const int64_t n_per_thread = (N + nth - 1) / nth; |
| 11451 | + const int64_t n_start = (int64_t) ith * n_per_thread; |
| 11452 | + const int64_t n_end = (n_start + n_per_thread < N) ? (n_start + n_per_thread) : N; |
| 11453 | + |
| 11454 | + // Per-thread fp32 scratch for one row of dequantized W (K floats). |
| 11455 | + // Allocated via plain malloc/free (matches the prior ml8-ml8.c behaviour |
| 11456 | + // which ran cleanly under ggml-cpu's threadpool — replacing this with |
| 11457 | + // std::vector triggered libgomp "Thread identifier invalid" errors on the |
| 11458 | + // second forward pass in llama-cli, likely an interaction with libstdc++ |
| 11459 | + // allocator thread-local state and ggml-cpu's threadpool fork model). |
| 11460 | + float * w_row_fp32 = (float *) malloc((size_t) K * sizeof(float)); |
| 11461 | + if (!w_row_fp32) { |
| 11462 | + GGML_ABORT("ggml_compute_forward_ml8_mul_mat: malloc(%zu) failed", |
| 11463 | + (size_t) K * sizeof(float)); |
| 11464 | + } |
| 11465 | + |
| 11466 | + for (int64_t n = n_start; n < n_end; n++) { |
| 11467 | + const block_ml8_4 * w_row = &w_blocks[n * n_groups_k]; |
| 11468 | + dequantize_row_ml8_4_with_lut(w_row, lut_fp8, w_row_fp32, K); |
| 11469 | + for (int64_t m = 0; m < M; m++) { |
| 11470 | + const float * x_col = &x_data[m * K]; |
| 11471 | + float sum = 0.0f; |
| 11472 | + for (int64_t k = 0; k < K; k++) { |
| 11473 | + sum += w_row_fp32[k] * x_col[k]; |
| 11474 | + } |
| 11475 | + y_data[m * N + n] = sum; |
| 11476 | + } |
| 11477 | + } |
| 11478 | + |
| 11479 | + free(w_row_fp32); |
| 11480 | +} |
| 11481 | + |
| 11482 | +// ─── ggml_compute_forward_ml8_apply_rotation (MAD-223 G.4.g) ────────────── |
| 11483 | +// |
| 11484 | +// Y[:, t] = unflatten(H_a^T @ flatten(X[:, t]) @ H_b, d=a*b) |
| 11485 | +// where H_b is the Sylvester Hadamard of size b_dim, built deterministically. |
| 11486 | +// |
| 11487 | +// op_params: int32_t[0]=a_dim, int32_t[1]=b_dim. |
| 11488 | +// src[0]=x F32 [d=a*b, n_tokens]; src[1]=h_a F32 [a, a]; dst=y F32 [d, n_tokens]. |
| 11489 | +// |
| 11490 | +// Ported from the original ggml_custom_4d implementation in ggml-ml8.c. |
| 11491 | +// Uses malloc/free (NOT std::vector) for scratch; std::vector inside a |
| 11492 | +// compute callback triggers libstdc++ allocator thread-local state issues |
| 11493 | +// with ggml-cpu's OpenMP threadpool (see ggml_compute_forward_ml8_mul_mat |
| 11494 | +// above for the same constraint). |
| 11495 | + |
| 11496 | +static void ml8_build_sylvester_cpu(float * H, int64_t b) { |
| 11497 | + H[0] = 1.0f; |
| 11498 | + for (int64_t n = 1; n < b; n *= 2) { |
| 11499 | + const float inv_sqrt2 = 0.70710678118654752440f; |
| 11500 | + for (int64_t i = n - 1; i >= 0; i--) { |
| 11501 | + for (int64_t j = n - 1; j >= 0; j--) { |
| 11502 | + const float v = H[i * b + j] * inv_sqrt2; |
| 11503 | + H[(i ) * b + (j )] = v; |
| 11504 | + H[(i ) * b + (j + n)] = v; |
| 11505 | + H[(i + n) * b + (j )] = v; |
| 11506 | + H[(i + n) * b + (j + n)] = -v; |
| 11507 | + } |
| 11508 | + } |
| 11509 | + } |
| 11510 | +} |
| 11511 | + |
| 11512 | +void ggml_compute_forward_ml8_apply_rotation(const ggml_compute_params * params, ggml_tensor * dst) { |
| 11513 | + const ggml_tensor * x = dst->src[0]; |
| 11514 | + const ggml_tensor * h_a = dst->src[1]; |
| 11515 | + |
| 11516 | + GGML_ASSERT(x->type == GGML_TYPE_F32); |
| 11517 | + GGML_ASSERT(h_a->type == GGML_TYPE_F32); |
| 11518 | + GGML_ASSERT(dst->type == GGML_TYPE_F32); |
| 11519 | + |
| 11520 | + const int32_t * pp = (const int32_t *) dst->op_params; |
| 11521 | + const int64_t a_dim = (int64_t) pp[0]; |
| 11522 | + const int64_t b_dim = (int64_t) pp[1]; |
| 11523 | + const int64_t d_dim = a_dim * b_dim; |
| 11524 | + GGML_ASSERT(x->ne[0] == d_dim); |
| 11525 | + GGML_ASSERT(h_a->ne[0] == a_dim && h_a->ne[1] == a_dim); |
| 11526 | + |
| 11527 | + const int64_t n_tokens = x->ne[1]; |
| 11528 | + const float * x_data = (const float *) x->data; |
| 11529 | + const float * h_a_data = (const float *) h_a->data; |
| 11530 | + float * y_data = (float *) dst->data; |
| 11531 | + |
| 11532 | + float * h_b = (float *) malloc((size_t) b_dim * (size_t) b_dim * sizeof(float)); |
| 11533 | + if (!h_b) { |
| 11534 | + GGML_ABORT("ml8_apply_rotation: malloc h_b(%zu) failed", |
| 11535 | + (size_t) b_dim * (size_t) b_dim * sizeof(float)); |
| 11536 | + } |
| 11537 | + ml8_build_sylvester_cpu(h_b, b_dim); |
| 11538 | + |
| 11539 | + float * xp = (float *) malloc((size_t) d_dim * sizeof(float)); |
| 11540 | + if (!xp) { |
| 11541 | + free(h_b); |
| 11542 | + GGML_ABORT("ml8_apply_rotation: malloc xp(%zu) failed", (size_t) d_dim * sizeof(float)); |
| 11543 | + } |
| 11544 | + |
| 11545 | + const int64_t per_thread = (n_tokens + params->nth - 1) / params->nth; |
| 11546 | + const int64_t t_start = (int64_t) params->ith * per_thread; |
| 11547 | + const int64_t t_end = (t_start + per_thread < n_tokens) ? (t_start + per_thread) : n_tokens; |
| 11548 | + |
| 11549 | + for (int64_t t = t_start; t < t_end; t++) { |
| 11550 | + const float * xt = x_data + t * d_dim; |
| 11551 | + float * yt = y_data + t * d_dim; |
| 11552 | + |
| 11553 | + // Step 1: xp[k, l] = sum_i H_a[i, k] * X[i, l] |
| 11554 | + for (int64_t k = 0; k < a_dim; k++) { |
| 11555 | + for (int64_t l = 0; l < b_dim; l++) { |
| 11556 | + float s = 0.0f; |
| 11557 | + for (int64_t i = 0; i < a_dim; i++) { |
| 11558 | + s += h_a_data[i * a_dim + k] * xt[i * b_dim + l]; |
| 11559 | + } |
| 11560 | + xp[k * b_dim + l] = s; |
| 11561 | + } |
| 11562 | + } |
| 11563 | + // Step 2: yt[k, l] = sum_j xp[k, j] * H_b[j, l] |
| 11564 | + for (int64_t k = 0; k < a_dim; k++) { |
| 11565 | + for (int64_t l = 0; l < b_dim; l++) { |
| 11566 | + float s = 0.0f; |
| 11567 | + for (int64_t j = 0; j < b_dim; j++) { |
| 11568 | + s += xp[k * b_dim + j] * h_b[j * b_dim + l]; |
| 11569 | + } |
| 11570 | + yt[k * b_dim + l] = s; |
| 11571 | + } |
| 11572 | + } |
| 11573 | + } |
| 11574 | + |
| 11575 | + free(xp); |
| 11576 | + free(h_b); |
| 11577 | +} |
| 11578 | + |
11408 | 11579 | static void ggml_compute_forward_fwht_f32(const ggml_compute_params * params, ggml_tensor * dst) { |
11409 | 11580 | const ggml_tensor * src0 = dst->src[0]; |
11410 | 11581 | const ggml_tensor * src1 = dst->src[1]; |
|
0 commit comments