diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 62758d1dc2..49649d24b3 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -1278,7 +1278,7 @@ static void ggml_compute_forward_mul_mat_scale( ggml_barrier(params->threadpool); - const float * sd = (const float *) scale->data; + const char * sd = (const char *) scale->data; const int64_t sne0 = scale->ne[0], sne1 = scale->ne[1], sne2 = scale->ne[2], sne3 = scale->ne[3]; const int64_t nr = ne1 * ne2 * ne3; for (int64_t ir = ith; ir < nr; ir += nth) { @@ -1286,9 +1286,9 @@ static void ggml_compute_forward_mul_mat_scale( const int64_t i2 = (ir / ne1) % ne2; const int64_t i3 = ir / (ne1 * ne2); float * dp = (float *) ((char *) dst->data + i1*nb1 + i2*nb2 + i3*nb3); - const int64_t so = (i1 % sne1)*sne0 + (i2 % sne2)*sne0*sne1 + (i3 % sne3)*sne0*sne1*sne2; + const size_t so = (i1 % sne1)*scale->nb[1] + (i2 % sne2)*scale->nb[2] + (i3 % sne3)*scale->nb[3]; for (int64_t i0 = 0; i0 < ne0; i0++) { - dp[i0] *= sd[(i0 % sne0) + so]; + dp[i0] *= *(const float *) (sd + (i0 % sne0)*scale->nb[0] + so); } } } @@ -1791,7 +1791,7 @@ static void ggml_compute_forward_mul_mat_id( if (scale) { ggml_barrier(params->threadpool); - const float * sd = (const float *) scale->data; + const char * sd = (const char *) scale->data; const bool per_ch_exp = scale->ne[0] == ne0 && scale->ne[1] == ne02; const bool per_expert = ggml_nelements(scale) == ne02; const int64_t nr = ne1 * ne2; @@ -1801,12 +1801,12 @@ static void ggml_compute_forward_mul_mat_id( const int32_t expert = *(const int32_t *) ((const char *) ids->data + i1*ids->nb[0] + i2*ids->nb[1]); float * dp = (float *) ((char *) dst->data + i1*nb1 + i2*nb2); if (per_ch_exp) { - const float * sp = sd + (int64_t) expert*ne0; + const char * sp = sd + expert*scale->nb[1]; for (int64_t i0 = 0; i0 < ne0; i0++) { - dp[i0] *= sp[i0]; + dp[i0] *= *(const float *) (sp + i0*scale->nb[0]); } } else { - const float s = sd[per_expert ? expert : 0]; + const float s = *(const float *) (sd + (per_expert ? expert*scale->nb[0] : 0)); for (int64_t i0 = 0; i0 < ne0; i0++) { dp[i0] *= s; } diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 28189f5601..8406bf24f8 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -2164,6 +2164,7 @@ static bool ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * // mmvq can fold the per-expert weight scale (src[3]) into its epilogue for NVFP4 const bool scale_fold = dst->src[3] != NULL && src0->type == GGML_TYPE_NVFP4 && ggml_nelements(dst->src[3]) == src0->ne[2] && + ggml_is_contiguous(dst->src[3]) && ggml_cuda_should_fuse_mul_mat_vec_q(dst); ggml_cuda_mm_fusion_args_host fusion = {}; if (scale_fold) { @@ -2327,7 +2328,8 @@ static bool ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * #define CUDA_MUL_MAT_SCALE_BLOCK_SIZE 256 static __global__ void mul_mat_scale_dense(float * dst, const float * scale, int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, - int64_t sne0, int64_t sne1, int64_t sne2, int64_t sne3) { + int64_t sne0, int64_t sne1, int64_t sne2, int64_t sne3, + size_t snb0, size_t snb1, size_t snb2, size_t snb3) { const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; if (i >= ne0*ne1*ne2*ne3) { return; @@ -2336,13 +2338,14 @@ static __global__ void mul_mat_scale_dense(float * dst, const float * scale, const int64_t i1 = (i / ne0) % ne1; const int64_t i2 = (i / (ne0*ne1)) % ne2; const int64_t i3 = i / (ne0*ne1*ne2); - const int64_t si = (i0 % sne0) + (i1 % sne1)*sne0 + (i2 % sne2)*sne0*sne1 + (i3 % sne3)*sne0*sne1*sne2; - dst[i] *= scale[si]; + const size_t si = (i0 % sne0)*snb0 + (i1 % sne1)*snb1 + (i2 % sne2)*snb2 + (i3 % sne3)*snb3; + dst[i] *= *(const float *) ((const char *) scale + si); } -// sne0 == 0: scalar per-expert scale[expert]; sne0 > 0: per-channel-per-expert scale[i0 + expert*sne0] +// strides are in bytes; channel_stride is zero for per-expert scalar scales static __global__ void mul_mat_scale_id(float * dst, const float * scale, const int32_t * ids, - int64_t ne0, int64_t ne1, int64_t ne2, int64_t ids_s0, int64_t ids_s1, int64_t sne0) { + int64_t ne0, int64_t ne1, int64_t ne2, int64_t ids_s0, int64_t ids_s1, + size_t channel_stride, size_t expert_stride) { const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; if (i >= ne0*ne1*ne2) { return; @@ -2351,13 +2354,13 @@ static __global__ void mul_mat_scale_id(float * dst, const float * scale, const const int64_t i1 = (i / ne0) % ne1; const int64_t i2 = i / (ne0*ne1); const int64_t expert = ids[i1*ids_s0 + i2*ids_s1]; - dst[i] *= scale[sne0 > 0 ? i0 + expert*sne0 : expert]; + dst[i] *= *(const float *) ((const char *) scale + i0*channel_stride + expert*expert_stride); } static void ggml_cuda_op_mul_mat_scale(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const bool is_id = dst->op == GGML_OP_MUL_MAT_ID; const ggml_tensor * scale = is_id ? dst->src[3] : dst->src[2]; - GGML_ASSERT(scale != NULL && scale->type == GGML_TYPE_F32 && ggml_is_contiguous(scale)); + GGML_ASSERT(scale != NULL && scale->type == GGML_TYPE_F32); GGML_ASSERT(ggml_is_contiguous(dst)); float * dst_d = (float *) dst->data; @@ -2377,13 +2380,15 @@ static void ggml_cuda_op_mul_mat_scale(ggml_backend_cuda_context & ctx, ggml_ten const int64_t blocks = (total + CUDA_MUL_MAT_SCALE_BLOCK_SIZE - 1) / CUDA_MUL_MAT_SCALE_BLOCK_SIZE; mul_mat_scale_id<<>>( dst_d, scale_d, (const int32_t *) ids->data, ne0, ne1, ne2, - ids->nb[0]/sizeof(int32_t), ids->nb[1]/sizeof(int32_t), per_channel_expert ? scale->ne[0] : 0); + ids->nb[0]/sizeof(int32_t), ids->nb[1]/sizeof(int32_t), + per_channel_expert ? scale->nb[0] : 0, per_channel_expert ? scale->nb[1] : scale->nb[0]); } else { const int64_t total = ne0*ne1*ne2*ne3; const int64_t blocks = (total + CUDA_MUL_MAT_SCALE_BLOCK_SIZE - 1) / CUDA_MUL_MAT_SCALE_BLOCK_SIZE; mul_mat_scale_dense<<>>( dst_d, scale_d, ne0, ne1, ne2, ne3, - scale->ne[0], scale->ne[1], scale->ne[2], scale->ne[3]); + scale->ne[0], scale->ne[1], scale->ne[2], scale->ne[3], + scale->nb[0], scale->nb[1], scale->nb[2], scale->nb[3]); } CUDA_CHECK(cudaGetLastError()); } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 0ababc0d03..0087a4a21b 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -5397,9 +5397,10 @@ struct test_mul_mat : public test_case { const int64_t pad; // bytes after the m_v rows of each batch of a, so nb[2] of a is not a multiple of nb[1] const size_t offset_a; // byte offset of the a view when k_v > k or m_v > m const int scale_gran; // weight scale granularity for non-needs_scale types: 0=none, 1=per-tensor, 2=per-output-channel + const bool scale_view; std::string vars() override { - return VARS_TO_STR15(type_a, type_b, m, n, k, bs, nr, per, k_v, o, src_overlap, m_v, pad, offset_a, scale_gran); + return VARS_TO_STR16(type_a, type_b, m, n, k, bs, nr, per, k_v, o, src_overlap, m_v, pad, offset_a, scale_gran, scale_view); } double max_nmse_err() override { @@ -5433,8 +5434,8 @@ struct test_mul_mat : public test_case { std::array bs = {10, 10}, std::array nr = {2, 2}, std::array per = {0, 1, 2, 3}, - int64_t k_v = 0, uint32_t o = 1, bool src_overlap = false, int64_t m_v = 0, int64_t pad = 0, size_t offset_a = 0, int scale_gran = 0) - : type_a(type_a), type_b(type_b), m(m), n(n), k(k), bs(bs), nr(nr), per(per), k_v(k_v), o(o), src_overlap(src_overlap), m_v(m_v), pad(pad), offset_a(offset_a), scale_gran(scale_gran) {} + int64_t k_v = 0, uint32_t o = 1, bool src_overlap = false, int64_t m_v = 0, int64_t pad = 0, size_t offset_a = 0, int scale_gran = 0, bool scale_view = false) + : type_a(type_a), type_b(type_b), m(m), n(n), k(k), bs(bs), nr(nr), per(per), k_v(k_v), o(o), src_overlap(src_overlap), m_v(m_v), pad(pad), offset_a(offset_a), scale_gran(scale_gran), scale_view(scale_view) {} ggml_tensor * build_graph(ggml_context * ctx) override { // C^T = A * B^T: (k, m) * (k, n) => (m, n) @@ -5516,6 +5517,11 @@ struct test_mul_mat : public test_case { } else if (!s && scale_gran == 2) { s = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, m); } + if (s && scale_view) { + s = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 2, s->ne[0], s->ne[1]); + s = ggml_view_3d(ctx, s, 1, s->ne[1], s->ne[2], s->nb[1], s->nb[2], sizeof(float)); + s = ggml_permute(ctx, s, 2, 0, 1, 3); + } if (s) { ggml_set_name(s, "s"); } @@ -5672,9 +5678,10 @@ struct test_mul_mat_id : public test_case { const bool ids_offset; const bool skip; // set some ids to a negative value const int scale_gran; // weight scale granularity for non-needs_scale types: 0=none, 1=per-expert scalar, 2=per-channel-per-expert + const bool scale_view; std::string vars() override { - return VARS_TO_STR12(type_a, type_b, n_mats, n_used, b, m, n, k, m_v, ids_offset, skip, scale_gran); + return VARS_TO_STR13(type_a, type_b, n_mats, n_used, b, m, n, k, m_v, ids_offset, skip, scale_gran, scale_view); } double max_nmse_err() override { @@ -5707,9 +5714,9 @@ struct test_mul_mat_id : public test_case { test_mul_mat_id(ggml_type type_a = GGML_TYPE_F32, ggml_type type_b = GGML_TYPE_F32, int n_mats = 8, int n_used = 2, bool b = false, int64_t m = 32, int64_t n = 32, int64_t k = 32, - int64_t m_v = 0, bool ids_offset = false, bool skip = false, int scale_gran = 0) + int64_t m_v = 0, bool ids_offset = false, bool skip = false, int scale_gran = 0, bool scale_view = false) : type_a(type_a), type_b(type_b), n_mats(n_mats), n_used(n_used), b(b), - m(m), n(n), k(k), m_v(m_v), ids_offset(ids_offset), skip(skip), scale_gran(scale_gran) { + m(m), n(n), k(k), m_v(m_v), ids_offset(ids_offset), skip(skip), scale_gran(scale_gran), scale_view(scale_view) { GGML_ASSERT(n_used <= n_mats); GGML_ASSERT(m_v == 0 || m_v > m); } @@ -5738,6 +5745,11 @@ struct test_mul_mat_id : public test_case { } else if (!s && scale_gran == 2) { s = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, m, n_mats); } + if (s && scale_view) { + s = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 2, s->ne[0], s->ne[1]); + s = ggml_view_3d(ctx, s, 1, s->ne[1], s->ne[2], s->nb[1], s->nb[2], sizeof(float)); + s = ggml_permute(ctx, s, 2, 0, 1, 3); + } if (s) { ggml_set_name(s, "s"); } @@ -11159,12 +11171,18 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat_id_w4a4(GGML_TYPE_MXFP4, GGML_TYPE_F32, 8, 2, false, 32, 32, 256)); for (int scale_gran : {1, 2}) { - test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F32, GGML_TYPE_F32, 1, 1, false, 32, 4, 32, 1.0f, 0, false, scale_gran)); + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F32, GGML_TYPE_F32, 1, 1, false, 32, 4, 32, 0, false, false, scale_gran)); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 32, 4, 32, {1, 1}, {1, 1}, {0, 1, 2, 3}, 0, 1, false, 0, 0, 0, scale_gran, true)); + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F32, GGML_TYPE_F32, 4, 2, false, 32, 4, 32, 0, false, false, scale_gran, true)); for (int64_t n : {1, 8}) { test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_K, GGML_TYPE_F32, 32, n, 256, {1, 1}, {1, 1}, {0, 1, 2, 3}, 0, 1, false, 0, 0, 0, scale_gran)); } } + for (int64_t n : {1, 16}) { + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_NVFP4, GGML_TYPE_F32, 4, 2, false, 32, n, 256, 0, false, false, 0, true)); + } + // FP8 weight-scale coverage (Qwen3.5-35B-A3B-FP8 shapes): dense per-channel + MoE per-channel-per-expert for (int scale_gran : {1, 2}) { // dense attention/shexp: (k, m)