diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 3cb4e07b48ff..4aff1f232333 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4363,6 +4363,17 @@ static void ggml_cuda_norm_fwht_q8_register(ggml_backend_cuda_context & ctx, con } } +struct ggml_cuda_tail_verify_entry { + float * scratch; int tag; + // q8_1 check of the attn_norm output (GGML_CUDA_TAIL_FUSION=3): the fused rows and a scratch for the reference quantize + const void * q8_fused = nullptr; void * q8_ref = nullptr; const ggml_tensor * norm_out = nullptr; + int64_t ne0_pad = 0; int64_t nrows = 0; +}; +static std::map & ggml_cuda_tail_verify_pending() { + static std::map m; + return m; +} + static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i) { static bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION")); @@ -4376,6 +4387,133 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph ggml_tensor * node = cgraph->nodes[i]; + // Gemma 4 MoE layer tail: [RMS_NORM, MUL] [RMS_NORM, MUL, ADD] [RMS_NORM, MUL, ADD] MUL(layer scalar) -> one launch. + // With GGML_CUDA_TAIL_FUSION=2 the next layer's attn_norm [RMS_NORM, MUL] joins the same launch; =3 also writes its q8_1 rows for the mat-vecs. + // GGML_CUDA_TAIL_FUSION=1/2 enables; GGML_CUDA_TAIL_VERIFY=1 additionally checks it bit-for-bit against the unfused nodes. + if (node->op == GGML_OP_RMS_NORM && i + 8 < cgraph->n_nodes) { + static const int tail_level = getenv("GGML_CUDA_TAIL_FUSION") != nullptr ? atoi(getenv("GGML_CUDA_TAIL_FUSION")) : 0; + static const bool tail_verify = getenv("GGML_CUDA_TAIL_VERIFY") != nullptr && atoi(getenv("GGML_CUDA_TAIL_VERIFY")) != 0; + if (tail_level > 0) { + ggml_tensor * const * n = cgraph->nodes + i; + static const ggml_op ops[] = { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD, + GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD, GGML_OP_MUL, GGML_OP_RMS_NORM, GGML_OP_MUL }; + bool ok = true; + for (int q = 0; q < 9 && ok; ++q) ok = n[q]->op == ops[q] && n[q]->type == GGML_TYPE_F32; + auto has = [](const ggml_tensor * t, const ggml_tensor * s) { return t->src[0] == s || t->src[1] == s; }; + auto other = [](const ggml_tensor * t, const ggml_tensor * s) { return t->src[0] == s ? t->src[1] : t->src[0]; }; + if (ok) { + ok = has(n[1], n[0]) && has(n[3], n[2]) && has(n[4], n[1]) && has(n[4], n[3]) && n[5]->src[0] == n[4] && + has(n[6], n[5]) && has(n[7], n[6]) && has(n[8], n[7]); + } + if (ok) { + const int64_t nc = n[0]->src[0]->ne[0]; + const ggml_tensor * ws[3] = { other(n[1], n[0]), other(n[3], n[2]), other(n[6], n[5]) }; + const ggml_tensor * rs = other(n[7], n[6]); + const ggml_tensor * ls = other(n[8], n[7]); + float e0, e2, e5; + memcpy(&e0, n[0]->op_params, 4); memcpy(&e2, n[2]->op_params, 4); memcpy(&e5, n[5]->op_params, 4); + ok = nc <= 3*1024 && nc % 1 == 0 && e0 == e2 && e0 == e5 && + n[0]->src[0]->ne[2] == 1 && n[0]->src[0]->ne[3] == 1 && n[2]->src[0]->ne[0] == nc && + n[2]->src[0]->ne[1] == n[0]->src[0]->ne[1] && n[2]->src[0]->ne[2] == 1 && n[2]->src[0]->ne[3] == 1 && + n[0]->src[0]->type == GGML_TYPE_F32 && n[2]->src[0]->type == GGML_TYPE_F32 && + n[0]->src[0]->nb[0] == 4 && n[2]->src[0]->nb[0] == 4 && + ggml_are_same_shape(n[4], n[0]) && ggml_are_same_shape(n[8], n[0]) && ggml_are_same_shape(rs, n[0]) && + rs->type == GGML_TYPE_F32 && rs->nb[0] == 4 && ggml_is_contiguous(n[8]) && + ls->type == GGML_TYPE_F32 && ggml_nelements(ls) == 1; + for (int q = 0; q < 3 && ok; ++q) { + ok = ws[q]->type == GGML_TYPE_F32 && ggml_is_contiguous(ws[q]) && ws[q]->ne[0] == nc && ggml_nrows(ws[q]) == 1; + } + // the next layer's attn_norm: RMS_NORM(l_out) * weight, not followed by a sign MUL (that is the + // norm+Hadamard fusion, which keeps its own path) + int n_fused = 9; + if (ok && tail_level >= 2 && i + 10 < cgraph->n_nodes && n[9]->op == ops[9] && n[10]->op == ops[10] && + n[9]->src[0] == n[8] && has(n[10], n[9]) && n[9]->type == GGML_TYPE_F32 && n[10]->type == GGML_TYPE_F32 && + strncmp(n[10]->name, "attn_norm", 9) == 0 && ggml_are_same_shape(n[10], n[8]) && ggml_is_contiguous(n[10])) { + const ggml_tensor * w4 = other(n[10], n[9]); + float e9; + memcpy(&e9, n[9]->op_params, 4); + bool ext = e9 == e0 && w4->type == GGML_TYPE_F32 && ggml_is_contiguous(w4) && w4->ne[0] == nc && ggml_nrows(w4) == 1; + for (int k = i + 11; ext && k < cgraph->n_nodes && k <= i + 26; ++k) { + ext = !(cgraph->nodes[k]->op == GGML_OP_MUL && has(cgraph->nodes[k], n[10])); + } + if (ext) n_fused = 11; + } + int outs[] = { i + 8, i + 10 }; + if (n_fused == 11 && !(ggml_can_fuse_subgraph(cgraph, i, 11, ops, outs, 2) && + ggml_cuda_check_fusion_memory_ranges(cgraph, i, 11, outs, 2, false, true))) { + n_fused = 9; + } + if (n_fused == 9) { + ok = ggml_can_fuse_subgraph(cgraph, i, 9, ops, outs, 1) && ggml_cuda_check_fusion_memory_ranges(cgraph, i, 9, outs, 1, false, true); + } + // level 3: also write the q8_1 rows of the normalized activation that the mat-vecs of the next + // attention block would quantize (all consumers are plain MUL_MATs with the AOS layout) + bool with_q8 = false; + int64_t ne0_pad = 0; + if (ok && n_fused == 11 && tail_level >= 3) { + const int64_t nrows = n[0]->src[0]->ne[1]; + ne0_pad = GGML_PAD(nc, MATRIX_ROW_PADDING); + int uses_found = 0; + with_q8 = nc % QK8_1 == 0 && ne0_pad <= 3*1024 && nrows <= 8 && n[10]->ne[2] == 1 && n[10]->ne[3] == 1; + for (int k = i + 11; with_q8 && k < cgraph->n_nodes && k <= i + 40; ++k) { + const ggml_tensor * t = cgraph->nodes[k]; + for (int q = 0; with_q8 && q < GGML_MAX_SRC; ++q) { + if (t->src[q] != n[10]) continue; + ++uses_found; + with_q8 = t->op == GGML_OP_MUL_MAT && q == 1 && t->src[0]->ne[2] == 1 && t->src[0]->ne[3] == 1 && + ggml_cuda_q8_1_layout_host(t->src[0]->type, (int) nrows, false) == GGML_CUDA_Q8_1_AOS; + } + } + with_q8 = with_q8 && uses_found > 0 && uses_found == ggml_node_get_use_count(cgraph, i + 10); + } + if (ok) { + if (tail_verify) { + static float * scratch = nullptr; static int tag = 0; + const size_t cap = (size_t) 32*1024*1024; + if (!scratch) { CUDA_CHECK(cudaMalloc(&scratch, 4*cap)); } + const size_t need = ggml_nelements(n[8])*sizeof(float); + if (need <= cap && tag < 1000000) { + float * scratch2 = (float *) ((char *) scratch + cap); + char * q8_fused = (char *) scratch + 2*cap; + char * q8_ref = (char *) scratch + 3*cap; + ggml_cuda_op_tail_fused(*cuda_ctx, n, n_fused, scratch, scratch2, with_q8 ? (block_q8_1 *) q8_fused : nullptr, (int) ne0_pad); + ggml_cuda_tail_verify_pending()[n[8]] = { scratch, tag }; + if (n_fused == 11) { + ggml_cuda_tail_verify_entry e = { scratch2, 1000000 + tag }; + if (with_q8) { + e.q8_fused = q8_fused; e.q8_ref = q8_ref; e.norm_out = n[10]; e.ne0_pad = ne0_pad; e.nrows = n[0]->src[0]->ne[1]; + } + ggml_cuda_tail_verify_pending()[n[10]] = e; + } + static int n_with_norm = 0, n_with_q8 = 0; + n_with_norm += n_fused == 11; + n_with_q8 += with_q8; + ++tag; + if ((tag & (tag - 1)) == 0) fprintf(stderr, "TAILVERIFY groups verified: %d, of which with the next attn_norm: %d, with q8_1: %d\n", tag, n_with_norm, n_with_q8); + } + return 0; + } + if (with_q8) { + const size_t q8_bytes = (size_t) (n[0]->src[0]->ne[1] * ne0_pad) * sizeof(block_q8_1) / QK8_1; + auto blk = std::make_unique>(cuda_ctx->pool(), q8_bytes); + char * q8 = blk->get(); + cuda_ctx->fwht_q8().held.push_back(std::move(blk)); + ggml_cuda_op_tail_fused(*cuda_ctx, n, n_fused, nullptr, nullptr, (block_q8_1 *) q8, (int) ne0_pad); + ggml_cuda_fwht_q8 e; + e.layout = GGML_CUDA_Q8_1_AOS; + e.ne0 = ne0_pad; + e.ncols = n[0]->src[0]->ne[1]; + e.data = q8; + cuda_ctx->fwht_q8().set(n[10], e); + } else { + ggml_cuda_op_tail_fused(*cuda_ctx, n, n_fused, nullptr, nullptr, nullptr, 0); + } + return n_fused - 1; + } + } + } + } + if (node->op == GGML_OP_CPY) { static const bool batched_cpy = [] { const char * env = getenv("GGML_CUDA_BATCHED_CPY"); @@ -5790,6 +5928,22 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud } GGML_ASSERT(ok); + if (!ggml_cuda_tail_verify_pending().empty()) { + auto it = ggml_cuda_tail_verify_pending().find(node); + if (it != ggml_cuda_tail_verify_pending().end()) { + ggml_cuda_op_tail_verify_cmp(*cuda_ctx, node, it->second.scratch, it->second.tag); + if (it->second.q8_fused) { + const ggml_cuda_tail_verify_entry & ve = it->second; + quantize_row_q8_1_cuda((const float *) ve.norm_out->data, nullptr, ve.q8_ref, GGML_CUDA_Q8_1_AOS, ve.norm_out->ne[0], + ve.norm_out->nb[1] / sizeof(float), 0, 0, ve.ne0_pad, ve.nrows, 1, 1, cuda_ctx->stream()); + ggml_cuda_op_tail_verify_cmp_q8(*cuda_ctx, ve.q8_ref, ve.q8_fused, (size_t) (ve.nrows * ve.ne0_pad) * sizeof(block_q8_1) / QK8_1, ve.tag); + } + static int nlaunch = 0; + ++nlaunch; + if ((nlaunch & (nlaunch - 1)) == 0) fprintf(stderr, "TAILVERIFY compare launches captured/issued: %d\n", nlaunch); + } + } + if (!is_concurrent_event_active) { try_launch_concurrent_event(node); } diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index 4a1712d17193..712c4030c851 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -582,6 +582,12 @@ static constexpr __host__ __device__ int calc_rows_per_block(ggml_type type, int // rows of the weight matrix per CTA for a tall one-column mat-vec (see should_use_tall below) #define MMVQ_TALL_ROWS_PER_BLOCK 8 +// Types besides PQ2_0 that also take the tall geometry (one instantiation more each): the common quantized types. +static constexpr bool mmvq_tall_ext_type(ggml_type type) { + return type == GGML_TYPE_Q8_0 || type == GGML_TYPE_Q4_0 || type == GGML_TYPE_Q2_0 || + type == GGML_TYPE_Q4_K || type == GGML_TYPE_Q6_K; +} + template __launch_bounds__(calc_nwarps(type, ncols_dst, get_device_table_id(), small_k, halve_iters)*ggml_cuda_get_physical_warp_size(), 1) static __global__ void mul_mat_vec_q( @@ -1343,8 +1349,11 @@ static void mul_mat_vec_q_switch_ncols_dst( // (Bonsai 2 9B LM head, 248320 x 4096: 249 us, about 1.1 TB/s on an H200). Eight rows per CTA keep the same K // partition and reduction order for every row (each row's sum is bit-identical) and cut that to 121 us; the // 4096-12288 row trunk projections gain 28-42% as well. Matrices under 4096 rows have too few CTAs for it. + // The same geometry helps any type whose one-column default gives one 128-thread CTA per row (tall output heads). + // GGML_CUDA_MMVQ_TALL_EXT=0 keeps it PQ2_0-only. const auto should_use_tall = [&] { - return type == GGML_TYPE_PQ2_0 && table_id == MMVQ_PARAMETERS_GENERIC && GGML_CUDA_CC_IS_NVIDIA(cc) && + static const bool ext = !getenv("GGML_CUDA_MMVQ_TALL_EXT") || atoi(getenv("GGML_CUDA_MMVQ_TALL_EXT")) != 0; + return (type == GGML_TYPE_PQ2_0 || (ext && mmvq_tall_ext_type(type))) && table_id == MMVQ_PARAMETERS_GENERIC && GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_HOPPER && cc < GGML_CUDA_CC_BLACKWELL && !has_ids && !y_soa && nrows_x >= 4096 && nrows_x % MMVQ_TALL_ROWS_PER_BLOCK == 0; // whole CTAs: rows past the end would be read, not only skipped }; @@ -1388,7 +1397,7 @@ static void mul_mat_vec_q_switch_ncols_dst( launch(std::true_type{}, std::false_type{}, std::false_type{}); } else if (should_halve_iters()) { launch(std::false_type{}, std::true_type{}, std::false_type{}); - } else if constexpr (type == GGML_TYPE_PQ2_0) { + } else if constexpr (type == GGML_TYPE_PQ2_0 || mmvq_tall_ext_type(type)) { if (should_use_tall()) { launch(std::false_type{}, std::false_type{}, std::true_type{}); } else { diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu index 136a4ace466e..29e48203cbae 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu @@ -1016,3 +1016,181 @@ void ggml_cuda_op_l2_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { l2_norm_f32_cuda(src0_d, dst_d, ne00, ne01, ne02, ne03, s01, s02, s03, eps, stream); } + +// Gemma 4 MoE layer tail in one launch (one block per row, ncols <= 3*1024): +// mlp = rms_norm(x1)*w1 ; comb = rms_norm(x2)*w2 + mlp ; res = rms_norm(comb)*w3 + resid ; dst = res * layer_scalar +// Same per-thread partial sums, same warp/block reduction (float2 reduce is component-wise identical) and same +// expression forms (fused multiply-add where the unfused rms_norm+mul+add kernel contracts) as the 4 unfused kernels. +template +static __global__ void __launch_bounds__(block_size) tail_fused_f32( + const float * x1, const float * w1, const float * x2, const float * w2, const float * w3, + const float * resid, const float * lscale, float * dst, const float * w4, float * dst4, block_q8_1 * q8, const int ne0_pad, + const int ncols, const int64_t x1_stride, const int64_t x2_stride, const int64_t resid_stride, const float eps) { + ggml_cuda_pdl_lc(); + const int row = blockIdx.x; + const int tid = threadIdx.x; + __shared__ float2 s_sum2[32]; + __shared__ float s_sum1[32]; + __shared__ float s_sum4[32]; + + ggml_cuda_pdl_sync(); + + const float * x1r = x1 + row*x1_stride; + const float * x2r = x2 + row*x2_stride; + const float * rr = resid + row*resid_stride; + float * dr = dst + (int64_t) row*ncols; + + float a[3], b[3]; + float t1 = 0.0f, t2 = 0.0f; +#pragma unroll + for (int k = 0; k < 3; ++k) { + const int col = tid + k*block_size; + if (col < ncols) { + a[k] = x1r[col]; + b[k] = x2r[col]; + t1 += a[k] * a[k]; + t2 += b[k] * b[k]; + } + } + const float2 r12 = block_reduce(make_float2(t1, t2), s_sum2); + const float scale1 = rsqrtf(r12.x / ncols + eps); + const float scale2 = rsqrtf(r12.y / ncols + eps); + + float c[3]; + float o[3]; + [[maybe_unused]] float y4[3]; + float t3 = 0.0f; +#pragma unroll + for (int k = 0; k < 3; ++k) { + const int col = tid + k*block_size; + if (col < ncols) { + const float mlp = __fmul_rn(__fmul_rn(scale1, a[k]), w1[col]); + c[k] = __fmaf_rn(__fmul_rn(scale2, b[k]), w2[col], mlp); + t3 += c[k] * c[k]; + } + } + const float r3 = block_reduce(t3, s_sum1); + const float scale3 = rsqrtf(r3 / ncols + eps); + const float ls = lscale[0]; + float t4 = 0.0f; +#pragma unroll + for (int k = 0; k < 3; ++k) { + const int col = tid + k*block_size; + if (col < ncols) { + const float res = __fmaf_rn(__fmul_rn(scale3, c[k]), w3[col], rr[col]); + o[k] = __fmul_rn(res, ls); + dr[col] = o[k]; + if constexpr (next_norm) { + t4 += o[k] * o[k]; + } + } + } + if constexpr (next_norm) { + // the next layer's attn_norm: rms_norm(l_out) * w4 + const float r4 = block_reduce(t4, s_sum4); + const float scale4 = rsqrtf(r4 / ncols + eps); + float * d4 = dst4 + (int64_t) row*ncols; +#pragma unroll + for (int k = 0; k < 3; ++k) { + const int col = tid + k*block_size; + if (col < ncols) { + y4[k] = __fmul_rn(__fmul_rn(scale4, o[k]), w4[col]); + d4[col] = y4[k]; + } + } + if constexpr (quant) { + // the q8_1 activation quantize of the normalized row (ggml_cuda_mul_mat_vec_q's quantize_q8_1, + // same expressions), padded to ne0_pad with zero blocks; one warp holds one 32-block (ncols % 32 == 0) +#pragma unroll + for (int k = 0; k < 3; ++k) { + const int col = tid + k*block_size; + if (col < ne0_pad) { + const float xi = col < ncols ? y4[k] : 0.0f; + float amax = fabsf(xi); + float sum = xi; + amax = warp_reduce_max(amax); + sum = warp_reduce_sum(sum); + const float d = amax / 127.0f; + const int8_t q = amax == 0.0f ? 0 : roundf(xi / d); + const int64_t i_cont = (int64_t) row*ne0_pad + col; + const int64_t ib = i_cont / QK8_1; + const int iqs = i_cont % QK8_1; + q8[ib].qs[iqs] = q; + if (iqs == 0) { + q8[ib].ds = make_half2(d, sum); + } + } + } + } + } +} + +static __global__ void tail_verify_cmp_f32(const float * a, const float * b, const int64_t n, const int tag) { + for (int64_t i = blockIdx.x*(int64_t) blockDim.x + threadIdx.x; i < n; i += (int64_t) gridDim.x*blockDim.x) { + if (__float_as_uint(a[i]) != __float_as_uint(b[i])) { + printf("TAILVERIFY MISMATCH tag=%d i=%lld ref=%.9g fused=%.9g\n", tag, (long long) i, a[i], b[i]); + } + } +} + +// nodes: rms1, mul1, rms2, mul2, add2, rms3, mul3, add3, lscale_mul[, rms4, mul4 (next layer attn_norm) when n_fused == 11]. +// verify_scratch != nullptr: write there (verify_scratch2 for the attn_norm output) instead of the graph tensors. +void ggml_cuda_op_tail_fused(ggml_backend_cuda_context & ctx, ggml_tensor * const * nodes, int n_fused, + float * verify_scratch, float * verify_scratch2, block_q8_1 * q8, int ne0_pad) { + const ggml_tensor * rms1 = nodes[0]; const ggml_tensor * mul1 = nodes[1]; + const ggml_tensor * rms2 = nodes[2]; const ggml_tensor * mul2 = nodes[3]; + const ggml_tensor * rms3 = nodes[5]; const ggml_tensor * mul3 = nodes[6]; const ggml_tensor * add3 = nodes[7]; + ggml_tensor * out = nodes[8]; + const float * w1 = (const float *) (mul1->src[0] == rms1 ? mul1->src[1] : mul1->src[0])->data; + const float * w2 = (const float *) (mul2->src[0] == rms2 ? mul2->src[1] : mul2->src[0])->data; + const float * w3 = (const float *) (mul3->src[0] == rms3 ? mul3->src[1] : mul3->src[0])->data; + const ggml_tensor * resid_t = add3->src[0] == mul3 ? add3->src[1] : add3->src[0]; + const ggml_tensor * ls_t = out->src[0] == nodes[7] ? out->src[1] : out->src[0]; + const float * w4 = nullptr; + float * dst4 = nullptr; + if (n_fused == 11) { + const ggml_tensor * mul4 = nodes[10]; + w4 = (const float *) (mul4->src[0] == nodes[9] ? mul4->src[1] : mul4->src[0])->data; + dst4 = verify_scratch2 ? verify_scratch2 : (float *) mul4->data; + } + float eps; + memcpy(&eps, rms1->op_params, sizeof(float)); + const int ncols = rms1->src[0]->ne[0]; + const int nrows = rms1->src[0]->ne[1]; + const dim3 blocks_num(nrows, 1, 1); + const dim3 block_dims(1024, 1, 1); + const ggml_cuda_kernel_launch_params launch_params{ blocks_num, block_dims, 0, ctx.stream() }; + float * dst = verify_scratch ? verify_scratch : (float *) out->data; + const int64_t s1 = rms1->src[0]->nb[1] / sizeof(float); + const int64_t s2 = rms2->src[0]->nb[1] / sizeof(float); + const int64_t sr = resid_t->nb[1] / sizeof(float); + if (n_fused == 11 && q8) { + ggml_cuda_kernel_launch(tail_fused_f32<1024, true, true>, launch_params, + (const float *) rms1->src[0]->data, w1, (const float *) rms2->src[0]->data, w2, w3, + (const float *) resid_t->data, (const float *) ls_t->data, dst, w4, dst4, q8, ne0_pad, ncols, s1, s2, sr, eps); + } else if (n_fused == 11) { + ggml_cuda_kernel_launch(tail_fused_f32<1024, true, false>, launch_params, + (const float *) rms1->src[0]->data, w1, (const float *) rms2->src[0]->data, w2, w3, + (const float *) resid_t->data, (const float *) ls_t->data, dst, w4, dst4, q8, ne0_pad, ncols, s1, s2, sr, eps); + } else { + ggml_cuda_kernel_launch(tail_fused_f32<1024, false, false>, launch_params, + (const float *) rms1->src[0]->data, w1, (const float *) rms2->src[0]->data, w2, w3, + (const float *) resid_t->data, (const float *) ls_t->data, dst, w4, dst4, q8, ne0_pad, ncols, s1, s2, sr, eps); + } +} + +void ggml_cuda_op_tail_verify_cmp(ggml_backend_cuda_context & ctx, const ggml_tensor * out, const float * scratch, int tag) { + tail_verify_cmp_f32<<<8, 256, 0, ctx.stream()>>>((const float *) out->data, scratch, ggml_nelements(out), tag); +} + +static __global__ void tail_verify_cmp_words(const int * a, const int * b, const int64_t n, const int tag) { + for (int64_t i = blockIdx.x*(int64_t) blockDim.x + threadIdx.x; i < n; i += (int64_t) gridDim.x*blockDim.x) { + if (a[i] != b[i]) { + printf("TAILVERIFY MISMATCH q8 tag=%d word=%lld ref=%08x fused=%08x\n", tag, (long long) i, (unsigned) a[i], (unsigned) b[i]); + } + } +} + +void ggml_cuda_op_tail_verify_cmp_q8(ggml_backend_cuda_context & ctx, const void * ref, const void * fused, size_t nbytes, int tag) { + tail_verify_cmp_words<<<8, 256, 0, ctx.stream()>>>((const int *) ref, (const int *) fused, (int64_t) (nbytes/4), tag); +} diff --git a/ggml/src/ggml-cuda/norm.cuh b/ggml/src/ggml-cuda/norm.cuh index 1ca5713aa371..4486ff970f02 100644 --- a/ggml/src/ggml-cuda/norm.cuh +++ b/ggml/src/ggml-cuda/norm.cuh @@ -29,3 +29,9 @@ void ggml_cuda_op_rms_norm_fused_add(ggml_backend_cuda_context & ctx, void ggml_cuda_op_rms_norm_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_l2_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +// Gemma 4 MoE layer tail (3 rms_norm+mul(+add) groups and the layer-scalar mul, optionally the next attn_norm) in one launch; nodes are the 9 or 11 graph nodes +void ggml_cuda_op_tail_fused(ggml_backend_cuda_context & ctx, ggml_tensor * const * nodes, int n_fused, + float * verify_scratch, float * verify_scratch2, block_q8_1 * q8, int ne0_pad); +void ggml_cuda_op_tail_verify_cmp(ggml_backend_cuda_context & ctx, const ggml_tensor * out, const float * scratch, int tag); +void ggml_cuda_op_tail_verify_cmp_q8(ggml_backend_cuda_context & ctx, const void * ref, const void * fused, size_t nbytes, int tag); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 3bf8245e406c..8f6516be0227 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1287,6 +1287,7 @@ void llama_model_base::load_hparams(llama_model_loader & ml) { case LLM_ARCH_QWEN35MOE: case LLM_ARCH_QWEN3NEXT: case LLM_ARCH_DSPARK: + case LLM_ARCH_GEMMA4: break; default: throw std::runtime_error(format( diff --git a/src/models/gemma4.cpp b/src/models/gemma4.cpp index e44f423bdbc5..3130b8284b5e 100644 --- a/src/models/gemma4.cpp +++ b/src/models/gemma4.cpp @@ -336,6 +336,7 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para // residual connection cur = ggml_add(ctx0, cur, attn_out); + cb(cur, "ffn_res", il); // per-layer embedding if (inp_per_layer) {