From 282ed52a6c25f2a81da7c4f5eca336d692add89b Mon Sep 17 00:00:00 2001 From: Pooya Khosravi <16091823+pooyakhosravi@users.noreply.github.com> Date: Sat, 29 Aug 2026 22:56:17 -0700 Subject: [PATCH 1/4] cuda: optimize Bonsai low-bit kernels for GB10 Assisted-by: OpenAI Codex --- ggml/src/ggml-cuda/concat.cu | 52 +++++ ggml/src/ggml-cuda/gated_delta_net.cu | 145 ++++++++---- ggml/src/ggml-cuda/ggml-cuda.cu | 237 ++++++++++++++++++++ ggml/src/ggml-cuda/mmq-config-blackwell.cuh | 4 + ggml/src/ggml-cuda/mmq-vec-dot.cuh | 1 - ggml/src/ggml-cuda/mmq.cu | 101 ++++++++- ggml/src/ggml-cuda/mmq.cuh | 74 +++++- ggml/src/ggml-cuda/mmvq.cu | 111 ++++++--- ggml/src/ggml-cuda/norm.cu | 144 +++++++++++- ggml/src/ggml-cuda/norm.cuh | 6 + ggml/src/ggml-cuda/quantize.cu | 85 ++++++- ggml/src/ggml-cuda/quantize.cuh | 10 + ggml/src/ggml-cuda/scale.cu | 29 +++ tests/test-backend-ops.cpp | 1 + 14 files changed, 907 insertions(+), 93 deletions(-) diff --git a/ggml/src/ggml-cuda/concat.cu b/ggml/src/ggml-cuda/concat.cu index 6df89013ca79..f59708892021 100644 --- a/ggml/src/ggml-cuda/concat.cu +++ b/ggml/src/ggml-cuda/concat.cu @@ -139,6 +139,41 @@ static __global__ void __launch_bounds__(CUDA_CONCAT_BLOCK_SIZE) } } +// Transpose a dimension-0 view while prepending a small contiguous state. The +// generic non-contiguous kernel assigns a CTA to each output row, which makes +// every warp gather with a large stride for this layout. +static __global__ void concat_dim0_transpose_u32( + const uint32_t * src0, const char * src1, uint32_t * dst, + int ne00, int rows, int cols, uint64_t src1_col_stride, int dst_row_stride) { + __shared__ uint32_t tile[32][33]; + + const int input_x = (int) blockIdx.x*32 + threadIdx.x; // output row / input column + const int input_y0 = (int) blockIdx.y*32 + threadIdx.y; // output column / input row + +#pragma unroll + for (int j = 0; j < 32; j += 8) { + const int input_y = input_y0 + j; + if (input_x < rows && input_y < cols) { + tile[threadIdx.y + j][threadIdx.x] = + *(const uint32_t *)(src1 + (uint64_t) input_y*src1_col_stride + (uint64_t) input_x*sizeof(uint32_t)); + } + } + __syncthreads(); + + const int output_row0 = (int) blockIdx.x*32 + threadIdx.y; + const int output_col = (int) blockIdx.y*32 + threadIdx.x; +#pragma unroll + for (int j = 0; j < 32; j += 8) { + const int output_row = output_row0 + j; + if (output_row < rows && output_col < cols) { + dst[(int64_t) output_row*dst_row_stride + ne00 + output_col] = tile[threadIdx.x][threadIdx.y + j]; + } + if (blockIdx.y == 0 && threadIdx.x < ne00 && output_row < rows) { + dst[(int64_t) output_row*dst_row_stride + threadIdx.x] = src0[(int64_t) output_row*ne00 + threadIdx.x]; + } + } +} + template static void concat_cuda(const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, int dim, cudaStream_t stream) { if (dim != 3 && ggml_is_contiguous_to_3(src0) && ggml_is_contiguous_to_3(src1)) { @@ -164,6 +199,23 @@ static void concat_cuda(const ggml_tensor * src0, const ggml_tensor * src1, ggml GGML_ASSERT(!ggml_is_quantized(src0->type)); dim3 grid_dim(dst->ne[1], dst->ne[2], dst->ne[3]); + if constexpr (sizeof(T) == sizeof(uint32_t)) { + const bool transpose_dim0 = ggml_cuda_info().devices[ggml_cuda_get_device()].cc == GGML_CUDA_CC_DGX_SPARK && + dim == 0 && src0->ne[2] == 1 && src0->ne[3] == 1 && src1->ne[2] == 1 && src1->ne[3] == 1 && + dst->ne[2] == 1 && dst->ne[3] == 1 && src0->ne[0] <= 8 && + src0->nb[0] == sizeof(uint32_t) && src0->nb[1] == (uint64_t) src0->ne[0]*sizeof(uint32_t) && + src1->nb[1] == sizeof(uint32_t) && dst->nb[0] == sizeof(uint32_t) && + dst->nb[1] == (uint64_t) dst->ne[0]*sizeof(uint32_t) && + src0->ne[1] == src1->ne[1] && src0->ne[1] == dst->ne[1] && + src0->ne[0] + src1->ne[0] == dst->ne[0]; + if (transpose_dim0) { + const dim3 grid((dst->ne[1] + 31)/32, (src1->ne[0] + 31)/32, 1); + concat_dim0_transpose_u32<<>>( + (const uint32_t *) src0->data, (const char *) src1->data, (uint32_t *) dst->data, + src0->ne[0], dst->ne[1], src1->ne[0], src1->nb[0], dst->nb[1]/sizeof(uint32_t)); + return; + } + } auto launch_kernel = [&](auto dim) { concat_non_cont<<>>( (const char *) src0->data, (const char *) src1->data, (char *) dst->data, diff --git a/ggml/src/ggml-cuda/gated_delta_net.cu b/ggml/src/ggml-cuda/gated_delta_net.cu index 1b431a724d72..f52a2368283f 100644 --- a/ggml/src/ggml-cuda/gated_delta_net.cu +++ b/ggml/src/ggml-cuda/gated_delta_net.cu @@ -1,7 +1,14 @@ #include "gated_delta_net.cuh" #include "ggml-cuda/common.cuh" -template +static __global__ void gdn_precompute_exp(const float * g, float * g_exp, int64_t n) { + for (int64_t i = (int64_t) blockIdx.x*blockDim.x + threadIdx.x; i < n; + i += (int64_t) blockDim.x*gridDim.x) { + g_exp[i] = expf(g[i]); + } +} + +template __global__ void __launch_bounds__((ggml_cuda_get_physical_warp_size() < S_v ? ggml_cuda_get_physical_warp_size() : S_v) * 4, 2) gated_delta_net_cuda(const float * q, const float * k, @@ -30,9 +37,14 @@ gated_delta_net_cuda(const float * q, int K) { const uint32_t h_idx = blockIdx.x; const uint32_t sequence = blockIdx.y; - // each warp owns one column, using warp-level primitives to reduce across rows + // Each warp owns one or more columns, using warp-level primitives to reduce across rows. const int lane = threadIdx.x; - const int col = blockIdx.z * blockDim.y + threadIdx.y; +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK + constexpr int cols_per_warp = S_v == 128 && !KDA ? 4 : 1; +#else + constexpr int cols_per_warp = 1; +#endif + const int col = (blockIdx.z * blockDim.y + threadIdx.y) * cols_per_warp; const uint32_t iq1 = fastmodulo(h_idx, neqk1_magic); const uint32_t iq3 = fastdiv(sequence, rq3_magic); @@ -44,20 +56,23 @@ gated_delta_net_cuda(const float * q, const int64_t state_in_offset = sequence * H * S_v * S_v + h_idx * S_v * S_v; const int64_t state_out_offset = (sequence * H + h_idx) * S_v * S_v; state += state_out_offset; - curr_state += state_in_offset + col * S_v; + curr_state += state_in_offset; attn_data += (sequence * n_tokens * H + h_idx) * S_v; constexpr int warp_size = ggml_cuda_get_physical_warp_size() < S_v ? ggml_cuda_get_physical_warp_size() : S_v; static_assert(S_v % warp_size == 0, "S_v must be a multiple of warp_size"); constexpr int rows_per_lane = (S_v + warp_size - 1) / warp_size; - float s_shard[rows_per_lane]; + float s_shard[cols_per_warp][rows_per_lane]; // state is stored transposed: M[col][i] = S[i][col], row col is contiguous ggml_cuda_pdl_sync(); #pragma unroll - for (int r = 0; r < rows_per_lane; r++) { - const int i = r * warp_size + lane; - s_shard[r] = curr_state[i]; + for (int c = 0; c < cols_per_warp; ++c) { +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + const int i = r * warp_size + lane; + s_shard[c][r] = curr_state[(col + c) * S_v + i]; + } } for (int t = 0; t < n_tokens; t++) { @@ -82,32 +97,32 @@ gated_delta_net_cuda(const float * q, } if constexpr (!KDA) { - const float g_val = expf(*g_t); + const float g_val = G_PRECOMPUTED ? *g_t : expf(*g_t); - // kv[col] = (S^T @ k)[col] = sum_i S[i][col] * k[i] - float kv_shard = 0.0f; + // Each warp owns one or more columns and reuses the common q/k registers. #pragma unroll - for (int r = 0; r < rows_per_lane; r++) { - kv_shard += s_shard[r] * k_reg[r]; - } - float kv_col = warp_reduce_sum(kv_shard); + for (int c = 0; c < cols_per_warp; ++c) { + float kv_shard = 0.0f; +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + kv_shard += s_shard[c][r] * k_reg[r]; + } + float kv_col = warp_reduce_sum(kv_shard); - // delta[col] = (v[col] - g * kv[col]) * beta - float delta_col = (v_t[col] - g_val * kv_col) * beta_val; + float delta_col = (v_t[col + c] - g_val * kv_col) * beta_val; - // fused: S[i][col] = g * S[i][col] + k[i] * delta[col] - // attn[col] = (S^T @ q)[col] = sum_i S[i][col] * q[i] - float attn_partial = 0.0f; + float attn_partial = 0.0f; #pragma unroll - for (int r = 0; r < rows_per_lane; r++) { - s_shard[r] = g_val * s_shard[r] + k_reg[r] * delta_col; - attn_partial += s_shard[r] * q_reg[r]; - } + for (int r = 0; r < rows_per_lane; r++) { + s_shard[c][r] = g_val * s_shard[c][r] + k_reg[r] * delta_col; + attn_partial += s_shard[c][r] * q_reg[r]; + } - float attn_col = warp_reduce_sum(attn_partial); + float attn_col = warp_reduce_sum(attn_partial); - if (lane == 0) { - attn_data[col] = attn_col * scale; + if (lane == 0) { + attn_data[col + c] = attn_col * scale; + } } } else { // kv[col] = sum_i g[i] * S[i][col] * k[i] @@ -115,7 +130,7 @@ gated_delta_net_cuda(const float * q, #pragma unroll for (int r = 0; r < rows_per_lane; r++) { const int i = r * warp_size + lane; - kv_shard += expf(g_t[i]) * s_shard[r] * k_reg[r]; + kv_shard += expf(g_t[i]) * s_shard[0][r] * k_reg[r]; } float kv_col = warp_reduce_sum(kv_shard); @@ -129,8 +144,8 @@ gated_delta_net_cuda(const float * q, #pragma unroll for (int r = 0; r < rows_per_lane; r++) { const int i = r * warp_size + lane; - s_shard[r] = expf(g_t[i]) * s_shard[r] + k_reg[r] * delta_col; - attn_partial += s_shard[r] * q_reg[r]; + s_shard[0][r] = expf(g_t[i]) * s_shard[0][r] + k_reg[r] * delta_col; + attn_partial += s_shard[0][r] * q_reg[r]; } float attn_col = warp_reduce_sum(attn_partial); @@ -149,24 +164,31 @@ gated_delta_net_cuda(const float * q, if (target_slot >= 0 && target_slot < K) { float * curr_state = state + target_slot * state_slot_stride; #pragma unroll - for (int r = 0; r < rows_per_lane; r++) { - const int i = r * warp_size + lane; - curr_state[col * S_v + i] = s_shard[r]; + for (int c = 0; c < cols_per_warp; ++c) { +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + const int i = r * warp_size + lane; + curr_state[(col + c) * S_v + i] = s_shard[c][r]; + } } } } + } if constexpr (!keep_rs_t) { #pragma unroll - for (int r = 0; r < rows_per_lane; r++) { - const int i = r * warp_size + lane; - state[col * S_v + i] = s_shard[r]; + for (int c = 0; c < cols_per_warp; ++c) { +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + const int i = r * warp_size + lane; + state[(col + c) * S_v + i] = s_shard[c][r]; + } } } } -template +template static void launch_gated_delta_net( const float * q_d, const float * k_d, const float * v_d, const float * g_d, const float * b_d, const float * s_d, @@ -179,8 +201,10 @@ static void launch_gated_delta_net( float scale, int64_t state_slot_stride, int K, cudaStream_t stream) { //TODO: Add chunked kernel for even faster pre-fill const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; const int num_warps = 4; - dim3 grid_dims(H, n_seqs, (S_v + num_warps - 1) / num_warps); + const int cols_per_warp = cc == GGML_CUDA_CC_DGX_SPARK && S_v == 128 && !KDA ? 4 : 1; + dim3 grid_dims(H, n_seqs, (S_v + num_warps * cols_per_warp - 1) / (num_warps * cols_per_warp)); dim3 block_dims(warp_size <= S_v ? warp_size : S_v, num_warps, 1); const uint3 neqk1_magic = init_fastdiv_values(neqk1); @@ -189,26 +213,26 @@ static void launch_gated_delta_net( const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, stream); switch (S_v) { case 16: - ggml_cuda_kernel_launch(gated_delta_net_cuda<16, KDA, keep_rs_t>, launch_params, + ggml_cuda_kernel_launch(gated_delta_net_cuda<16, KDA, keep_rs_t, G_PRECOMPUTED>, launch_params, q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K); break; case 32: - ggml_cuda_kernel_launch(gated_delta_net_cuda<32, KDA, keep_rs_t>, launch_params, + ggml_cuda_kernel_launch(gated_delta_net_cuda<32, KDA, keep_rs_t, G_PRECOMPUTED>, launch_params, q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K); break; case 64: { - ggml_cuda_kernel_launch(gated_delta_net_cuda<64, KDA, keep_rs_t>, launch_params, + ggml_cuda_kernel_launch(gated_delta_net_cuda<64, KDA, keep_rs_t, G_PRECOMPUTED>, launch_params, q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K); break; } case 128: { - ggml_cuda_kernel_launch(gated_delta_net_cuda<128, KDA, keep_rs_t>, launch_params, + ggml_cuda_kernel_launch(gated_delta_net_cuda<128, KDA, keep_rs_t, G_PRECOMPUTED>, launch_params, q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K); @@ -294,6 +318,19 @@ static void ggml_cuda_op_gated_delta_net_impl( state_slot_stride = cache->slot_stride; } + ggml_cuda_pool_alloc g_exp_alloc(ctx.pool()); + bool g_precomputed = false; + if (!kda && S_v == 128 && n_tokens >= 32 && + ggml_cuda_info().devices[ggml_cuda_get_device()].cc == GGML_CUDA_CC_DGX_SPARK) { + const int64_t n_g = ggml_nelements(src_g); + g_exp_alloc.alloc(n_g*sizeof(float)); + const int block = 256; + const int grid = std::min((n_g + block - 1)/block, 4096); + gdn_precompute_exp<<>>(g_d, g_exp_alloc.ptr, n_g); + g_d = g_exp_alloc.ptr; + g_precomputed = true; + } + if (kda) { if (keep_rs) { launch_gated_delta_net(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, @@ -306,13 +343,25 @@ static void ggml_cuda_op_gated_delta_net_impl( } } else { if (keep_rs) { - launch_gated_delta_net(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, - S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, stream); + if (g_precomputed) { + launch_gated_delta_net(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, + S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, + sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, stream); + } else { + launch_gated_delta_net(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, + S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, + sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, stream); + } } else { - launch_gated_delta_net(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, - S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, stream); + if (g_precomputed) { + launch_gated_delta_net(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, + S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, + sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, stream); + } else { + launch_gated_delta_net(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, + S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, + sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, stream); + } } } } diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 9178d32fe717..c066e53a6bf8 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -3291,6 +3291,96 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph ggml_tensor * node = cgraph->nodes[i]; + // GB10 prefill: apply SWIGLU while quantizing the activation consumed by the + // following low-bit down projection. This avoids materializing and rereading + // the GLU output and removes one launch per transformer block. + if (node->op == GGML_OP_GLU && i + 1 < cgraph->n_nodes) { + ggml_tensor * mm = cgraph->nodes[i + 1]; + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + const ggml_op ops[] = { GGML_OP_GLU, GGML_OP_MUL_MAT }; + const int out_nodes[] = { i + 1 }; + if (cc == GGML_CUDA_CC_DGX_SPARK && mm->op == GGML_OP_MUL_MAT && mm->src[1] == node && + node->src[0] && node->src[1] && node->src[0]->type == GGML_TYPE_F32 && + node->src[1]->type == GGML_TYPE_F32 && ggml_get_glu_op(node) == GGML_GLU_OP_SWIGLU && + ggml_get_op_params_i32(node, 1) == 0 && ggml_are_same_shape(node->src[0], node->src[1]) && + ggml_is_contiguous(node->src[0]) && ggml_is_contiguous(node->src[1]) && + mm->src[1]->ne[1] >= 32 && + ggml_cuda_should_use_mmq(mm->src[0]->type, cc, mm->src[1]->ne[1], 0) && + ggml_can_fuse_subgraph(cgraph, i, 2, ops, out_nodes, 1) && + ggml_cuda_check_fusion_memory_ranges(cgraph, i, 2, out_nodes, 1)) { + ggml_cuda_mul_mat_q(*cuda_ctx, mm->src[0], node->src[1], nullptr, mm, node->src[0]); + return 1; + } + } + + // Post-attention normalization feeds adjacent gate and up projections. + // Quantize the normalized residual once and reuse it for both MMQs. + if (node->op == GGML_OP_ADD && i + 4 < cgraph->n_nodes) { + ggml_tensor * rms_norm = cgraph->nodes[i + 1]; + ggml_tensor * mul = cgraph->nodes[i + 2]; + ggml_tensor * mm_a = cgraph->nodes[i + 3]; + ggml_tensor * mm_b = cgraph->nodes[i + 4]; + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + const ggml_op ops[] = { GGML_OP_ADD, GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_MUL_MAT, GGML_OP_MUL_MAT }; + const int out_nodes[] = { i, i + 3, i + 4 }; + const char * dual_rms_q8 = getenv("GGML_CUDA_GB10_DUAL_RMS_Q8"); + if ((!dual_rms_q8 || std::atoi(dual_rms_q8) != 0) && + cc == GGML_CUDA_CC_DGX_SPARK && rms_norm->op == GGML_OP_RMS_NORM && + mul->op == GGML_OP_MUL && mm_a->op == GGML_OP_MUL_MAT && mm_b->op == GGML_OP_MUL_MAT && + mm_a->src[1] == mul && mm_b->src[1] == mul && + (mul->src[0] == rms_norm || mul->src[1] == rms_norm) && rms_norm->src[0] == node && + node->src[0] && node->src[1] && node->type == GGML_TYPE_F32 && + node->src[0]->type == GGML_TYPE_F32 && node->src[1]->type == GGML_TYPE_F32 && + rms_norm->type == GGML_TYPE_F32 && mul->type == GGML_TYPE_F32 && + ggml_are_same_shape(node->src[0], node->src[1]) && ggml_are_same_shape(node, rms_norm) && + ggml_is_contiguous(node->src[0]) && ggml_is_contiguous(node->src[1]) && + ggml_is_contiguous(node) && ggml_is_contiguous(rms_norm) && ggml_is_contiguous(mul) && + mul->ne[1] >= 32 && mm_a->src[0]->type == mm_b->src[0]->type && + (mm_a->src[0]->type == GGML_TYPE_Q1_0 || + mm_a->src[0]->type == GGML_TYPE_Q2_0 || + mm_a->src[0]->type == GGML_TYPE_PQ2_0) && + ggml_cuda_should_use_mmq(mm_a->src[0]->type, cc, mul->ne[1], 0) && + ggml_cuda_should_use_mmq(mm_b->src[0]->type, cc, mul->ne[1], 0) && + ggml_can_fuse_subgraph(cgraph, i, 5, ops, out_nodes, 3)) { + const ggml_tensor * weight = mul->src[0] == rms_norm ? mul->src[1] : mul->src[0]; + if (weight && weight->type == GGML_TYPE_F32 && ggml_is_contiguous(weight) && + weight->ne[0] == node->ne[0] && ggml_nrows(weight) == 1) { + ggml_cuda_pool_alloc row_scale(cuda_ctx->pool(), ggml_nrows(node)); + ggml_cuda_op_add_rms_norm_scale_fused(*cuda_ctx, node, rms_norm, row_scale.get()); + ggml_cuda_mul_mat_q_fused_two(*cuda_ctx, mm_a->src[0], mm_b->src[0], node, + mm_a, mm_b, weight, row_scale.get()); + return 4; + } + } + } + + // Preserve the residual sum for later graph consumers while normalizing it + // in the same launch. This removes a full read of the residual tensor. + if (node->op == GGML_OP_ADD && i + 2 < cgraph->n_nodes) { + ggml_tensor * rms_norm = cgraph->nodes[i + 1]; + ggml_tensor * mul = cgraph->nodes[i + 2]; + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + const ggml_op ops[] = { GGML_OP_ADD, GGML_OP_RMS_NORM, GGML_OP_MUL }; + const int out_nodes[] = { i, i + 2 }; + if (cc == GGML_CUDA_CC_DGX_SPARK && rms_norm->op == GGML_OP_RMS_NORM && + mul->op == GGML_OP_MUL && (mul->src[0] == rms_norm || mul->src[1] == rms_norm) && + rms_norm->src[0] == node && node->src[0] && node->src[1] && + node->type == GGML_TYPE_F32 && node->src[0]->type == GGML_TYPE_F32 && + node->src[1]->type == GGML_TYPE_F32 && rms_norm->type == GGML_TYPE_F32 && mul->type == GGML_TYPE_F32 && + ggml_are_same_shape(node->src[0], node->src[1]) && + ggml_are_same_shape(node, rms_norm) && ggml_is_contiguous(node->src[0]) && + ggml_is_contiguous(node->src[1]) && ggml_is_contiguous(node) && + ggml_is_contiguous(rms_norm) && ggml_is_contiguous(mul) && + ggml_can_fuse_subgraph(cgraph, i, 3, ops, out_nodes, 2)) { + const ggml_tensor * weight = mul->src[0] == rms_norm ? mul->src[1] : mul->src[0]; + if (weight && weight->type == GGML_TYPE_F32 && ggml_is_contiguous(weight) && + weight->ne[0] == node->ne[0] && ggml_nrows(weight) == 1) { + ggml_cuda_op_add_rms_norm_fused(*cuda_ctx, node, rms_norm, mul); + return 2; + } + } + } + // gated_delta_net -> cpy: scatter recurrent-state snapshots into the cache if (node->op == GGML_OP_GATED_DELTA_NET) { ggml_cuda_gated_delta_net_fused_cache fused_state_cpy; @@ -4053,6 +4143,35 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud ggml_cuda_concurrent_event * concurrent_event = nullptr; bool should_launch_concurrent_events = false; + struct gb10_shared_q8_entry { + const ggml_tensor * src1; + ggml_type type; + std::unique_ptr> data; + bool quantized = false; + int remaining = 0; + }; + std::vector gb10_shared_q8; + + struct gb10_virtual_rms_entry { + const ggml_tensor * normalized; + ggml_tensor * residual; + const ggml_tensor * weight; + ggml_type type; + std::unique_ptr> row_scale; + int remaining; + }; + std::vector gb10_virtual_rms; + + const auto gb10_shared_q8_consumer_count = [&](const ggml_tensor * src1, ggml_type type) { + int count = 0; + for (int j = 0; j < cgraph->n_nodes; ++j) { + const ggml_tensor * candidate = cgraph->nodes[j]; + count += candidate->op == GGML_OP_MUL_MAT && candidate->src[0] && + candidate->src[1] == src1 && candidate->src[0]->type == type; + } + return count; + }; + const auto try_launch_concurrent_event = [&](const ggml_tensor * node) { if (stream_ctx.concurrent_events.find(node) != stream_ctx.concurrent_events.end()) { concurrent_event = &stream_ctx.concurrent_events[node]; @@ -4073,6 +4192,10 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud }; while (!graph_evaluated_or_captured) { + for (auto & entry : gb10_shared_q8) { + entry.quantized = false; + entry.remaining = gb10_shared_q8_consumer_count(entry.src1, entry.type); + } // Only perform the graph execution if CUDA graphs are not enabled, or we are capturing the graph. // With the use of CUDA graphs, the execution will be performed by the graph launch. if (!use_cuda_graph || cuda_graph_update_required) { @@ -4182,6 +4305,61 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud continue; } + // The normalized pre-attention residual is consumed only by a + // group of low-bit projections. Preserve residual + one scale per + // row and let their shared Q8 quantizer apply the norm weight. + const char * virtual_rms_env = getenv("GGML_CUDA_GB10_VIRTUAL_RMS_Q8"); + const int virtual_rms_cc = ggml_cuda_info().devices[cuda_ctx->device].cc; + if ((!virtual_rms_env || std::atoi(virtual_rms_env) != 0) && + virtual_rms_cc == GGML_CUDA_CC_DGX_SPARK && !is_concurrent_event_active && + node->op == GGML_OP_ADD && i + 2 < cgraph->n_nodes) { + ggml_tensor * rms = cgraph->nodes[i + 1]; + ggml_tensor * mul = cgraph->nodes[i + 2]; + const ggml_op ops[] = { GGML_OP_ADD, GGML_OP_RMS_NORM, GGML_OP_MUL }; + const int out_nodes[] = { i, i + 2 }; + const ggml_tensor * weight = mul->src[0] == rms ? mul->src[1] : mul->src[0]; + int consumers = 0; + ggml_type consumer_type = GGML_TYPE_COUNT; + bool consumers_ok = true; + for (int j = i + 3; j < cgraph->n_nodes; ++j) { + ggml_tensor * consumer = cgraph->nodes[j]; + for (int s = 0; s < GGML_MAX_SRC; ++s) { + if (consumer->src[s] != mul) { + continue; + } + const bool is_mmq = s == 1 && consumer->op == GGML_OP_MUL_MAT && consumer->src[0] && + (consumer->src[0]->type == GGML_TYPE_Q1_0 || + consumer->src[0]->type == GGML_TYPE_Q2_0 || + consumer->src[0]->type == GGML_TYPE_PQ2_0) && + ggml_cuda_should_use_mmq(consumer->src[0]->type, virtual_rms_cc, mul->ne[1], 0); + if (!is_mmq || (consumers > 0 && consumer_type != consumer->src[0]->type)) { + consumers_ok = false; + } else { + consumer_type = consumer->src[0]->type; + ++consumers; + } + } + } + if (consumers_ok && consumers > 1 && rms->op == GGML_OP_RMS_NORM && + mul->op == GGML_OP_MUL && (mul->src[0] == rms || mul->src[1] == rms) && + rms->src[0] == node && node->src[0] && node->src[1] && weight && + node->type == GGML_TYPE_F32 && node->src[0]->type == GGML_TYPE_F32 && + node->src[1]->type == GGML_TYPE_F32 && rms->type == GGML_TYPE_F32 && + mul->type == GGML_TYPE_F32 && weight->type == GGML_TYPE_F32 && + ggml_are_same_shape(node->src[0], node->src[1]) && + ggml_are_same_shape(node, rms) && ggml_is_contiguous(node->src[0]) && + ggml_is_contiguous(node->src[1]) && ggml_is_contiguous(node) && + ggml_is_contiguous(mul) && ggml_is_contiguous(weight) && + weight->ne[0] == node->ne[0] && ggml_nrows(weight) == 1 && node->ne[1] >= 32 && + ggml_can_fuse_subgraph(cgraph, i, 3, ops, out_nodes, 2)) { + auto row_scale = std::make_unique>(cuda_ctx->pool(), ggml_nrows(node)); + ggml_cuda_op_add_rms_norm_scale_fused(*cuda_ctx, node, rms, row_scale->get()); + gb10_virtual_rms.push_back({ mul, node, weight, consumer_type, std::move(row_scale), consumers }); + i += 2; + continue; + } + } + int nodes_to_skip = ggml_cuda_try_fuse(cuda_ctx, cgraph, i); if (nodes_to_skip != 0) { @@ -4194,6 +4372,56 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud i += nodes_to_skip; continue; } + + // Several Qwen3.5 attention projections consume the exact same + // normalized activation. Quantize it once per graph execution and + // reuse the Q8 tile for the later MMQs on the same CUDA stream. + const char * shared_q8_env = getenv("GGML_CUDA_GB10_SHARED_Q8"); + const int cc = ggml_cuda_info().devices[cuda_ctx->device].cc; + if ((!shared_q8_env || std::atoi(shared_q8_env) != 0) && + cc == GGML_CUDA_CC_DGX_SPARK && !is_concurrent_event_active && + node->op == GGML_OP_MUL_MAT && node->src[0] && node->src[1] && + node->src[1]->type == GGML_TYPE_F32 && ggml_is_contiguous(node->src[1]) && + node->src[1]->ne[1] >= 32 && + (node->src[0]->type == GGML_TYPE_Q1_0 || + node->src[0]->type == GGML_TYPE_Q2_0 || + node->src[0]->type == GGML_TYPE_PQ2_0) && + ggml_cuda_should_use_mmq(node->src[0]->type, cc, node->src[1]->ne[1], 0) && + gb10_shared_q8_consumer_count(node->src[1], node->src[0]->type) > 1) { + auto virtual_it = std::find_if(gb10_virtual_rms.begin(), gb10_virtual_rms.end(), [&](const auto & entry) { + return entry.normalized == node->src[1] && entry.type == node->src[0]->type; + }); + const bool virtual_rms = virtual_it != gb10_virtual_rms.end() && virtual_it->row_scale; + auto it = std::find_if(gb10_shared_q8.begin(), gb10_shared_q8.end(), [&](const auto & entry) { + return entry.src1 == node->src[1] && entry.type == node->src[0]->type; + }); + if (it == gb10_shared_q8.end()) { + auto data = std::make_unique>(cuda_ctx->pool(), + ggml_cuda_mul_mat_q_q8_size(node->src[0], node->src[1])); + gb10_shared_q8.push_back({ node->src[1], node->src[0]->type, std::move(data), false, + gb10_shared_q8_consumer_count(node->src[1], node->src[0]->type) }); + it = std::prev(gb10_shared_q8.end()); + } else if (!it->data) { + it->data = std::make_unique>(cuda_ctx->pool(), + ggml_cuda_mul_mat_q_q8_size(node->src[0], node->src[1])); + } + ggml_cuda_mul_mat_q(*cuda_ctx, node->src[0], virtual_rms ? virtual_it->residual : node->src[1], + nullptr, node, nullptr, virtual_rms ? virtual_it->weight : nullptr, nullptr, + it->data->get(), !it->quantized, + virtual_rms ? virtual_it->row_scale->get() : nullptr); + it->quantized = true; + if (--it->remaining == 0) { + it->data.reset(); + } + if (virtual_rms && --virtual_it->remaining == 0) { + virtual_it->row_scale.reset(); + } + if (!is_concurrent_event_active) { + try_launch_concurrent_event(node); + } + continue; + } + #ifndef NDEBUG // On integrated GPUs (APUs, e.g. RDNA3.5) the scheduler may place a // node's output on the host-visible buffer, which the compute path @@ -4258,6 +4486,15 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud graph_evaluated_or_captured = true; #endif // USE_CUDA_GRAPH } + + // The VMM scratch pool is stack-like. Persistent shared-Q8 allocations are + // created in graph order, so release them explicitly in reverse order. + for (auto it = gb10_shared_q8.rbegin(); it != gb10_shared_q8.rend(); ++it) { + it->data.reset(); + } + for (auto it = gb10_virtual_rms.rbegin(); it != gb10_virtual_rms.rend(); ++it) { + it->row_scale.reset(); + } } #ifdef USE_CUDA_GRAPH diff --git a/ggml/src/ggml-cuda/mmq-config-blackwell.cuh b/ggml/src/ggml-cuda/mmq-config-blackwell.cuh index 9fbe32b6972b..7ead66f5d169 100644 --- a/ggml/src/ggml-cuda/mmq-config-blackwell.cuh +++ b/ggml/src/ggml-cuda/mmq-config-blackwell.cuh @@ -35,3 +35,7 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf return ggml_cuda_mmq_get_config_ampere(type, J, fallback); } + +static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_gb10(ggml_type type, int J, bool fallback) { + return ggml_cuda_mmq_get_config_blackwell(type, J, fallback); +} diff --git a/ggml/src/ggml-cuda/mmq-vec-dot.cuh b/ggml/src/ggml-cuda/mmq-vec-dot.cuh index d573433865f8..be091d270717 100644 --- a/ggml/src/ggml-cuda/mmq-vec-dot.cuh +++ b/ggml/src/ggml-cuda/mmq-vec-dot.cuh @@ -1248,4 +1248,3 @@ template static __device__ __forceinline_ } } } - diff --git a/ggml/src/ggml-cuda/mmq.cu b/ggml/src/ggml-cuda/mmq.cu index 3a08735ec9bf..148a6646d80f 100644 --- a/ggml/src/ggml-cuda/mmq.cu +++ b/ggml/src/ggml-cuda/mmq.cu @@ -86,10 +86,17 @@ static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, con } void ggml_cuda_mul_mat_q( - ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst) { + ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, + ggml_tensor * dst, const ggml_tensor * gate, + const ggml_tensor * norm_weight, const ggml_tensor * norm_scale, + void * external_q8, bool quantize_external, const float * external_norm_scale) { GGML_ASSERT( src1->type == GGML_TYPE_F32); GGML_ASSERT( dst->type == GGML_TYPE_F32); GGML_ASSERT(!ids || ids->type == GGML_TYPE_I32); // Optional, used for batched GGML_MUL_MAT_ID. + GGML_ASSERT(!gate || (!ids && gate->type == GGML_TYPE_F32)); + GGML_ASSERT((norm_weight == nullptr) == (norm_scale == nullptr && external_norm_scale == nullptr)); + GGML_ASSERT(!norm_weight || (!ids && !gate && norm_weight->type == GGML_TYPE_F32 && + ((norm_scale && norm_scale->type == GGML_TYPE_F32) || external_norm_scale))); GGML_TENSOR_BINARY_OP_LOCALS; @@ -138,13 +145,15 @@ void ggml_cuda_mul_mat_q( if (!ids) { const size_t nbytes_src1_q8_1 = ne13*ne12 * ne11*ne10_padded * y_block_size/y_values_per_block + ggml_cuda_mmq_get_J_max(src0->type, fallback, cc, ne11) * sizeof(block_q8_1_mmq); - ggml_cuda_pool_alloc src1_q8_1(ctx.pool(), nbytes_src1_q8_1); + GGML_ASSERT(!external_q8 || (!gate && !use_native_fp4)); + ggml_cuda_pool_alloc src1_q8_1(ctx.pool()); + char * src1_q8_ptr = external_q8 ? (char *) external_q8 : src1_q8_1.alloc(nbytes_src1_q8_1); ggml_cuda_pool_alloc src1_scale(ctx.pool()); if (src0->type == GGML_TYPE_NVFP4 && use_native_fp4) { src1_scale.alloc(ne13*ne12*ne11); } - { + if (!external_q8 || quantize_external) { const int64_t s11 = src1->nb[1] / ts_src1; const int64_t s12 = src1->nb[2] / ts_src1; const int64_t s13 = src1->nb[3] / ts_src1; @@ -152,12 +161,28 @@ void ggml_cuda_mul_mat_q( static constexpr size_t align_float8 = 32; const bool use_aligned_float8 = ggml_cuda_is_aligned(src1, align_float8); static_assert(sizeof(block_fp4_mmq) == 4 * sizeof(block_q8_1)); - quantize_mmq_fp4_cuda(src1_d, nullptr, src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10, s11, s12, s13, ne10_padded, + quantize_mmq_fp4_cuda(src1_d, nullptr, src1_q8_ptr, src1_scale.ptr, src0->type, use_aligned_float8, ne10, s11, s12, s13, ne10_padded, ne11, ne12, ne13, stream); } else { - quantize_mmq_q8_1_cuda(src1_d, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, - ne11, ne12, ne13, stream); + if (norm_weight) { + GGML_ASSERT(norm_weight->ne[0] == src1->ne[0] && ggml_nrows(norm_weight) == 1); + GGML_ASSERT(ggml_is_contiguous(norm_weight)); + GGML_ASSERT(external_norm_scale || (ggml_is_contiguous(norm_scale) && + ggml_nelements(norm_scale) >= ggml_nrows(src1))); + quantize_mmq_q8_1_rms_cuda(src1_d, (const float *) norm_weight->data, + external_norm_scale ? external_norm_scale : (const float *) norm_scale->data, + src1_q8_ptr, src0->type, + ne10, s11, s12, s13, ne10_padded, ne11, ne12, ne13, stream); + } else if (gate) { + GGML_ASSERT(ggml_are_same_shape(src1, gate)); + GGML_ASSERT(ggml_is_contiguous(gate)); + quantize_mmq_q8_1_swiglu_cuda(src1_d, (const float *) gate->data, src1_q8_ptr, src0->type, + ne10, s11, s12, s13, ne10_padded, ne11, ne12, ne13, stream); + } else { + quantize_mmq_q8_1_cuda(src1_d, nullptr, src1_q8_ptr, src0->type, ne10, s11, s12, s13, ne10_padded, + ne11, ne12, ne13, stream); + } } CUDA_CHECK(cudaGetLastError()); } @@ -169,7 +194,7 @@ void ggml_cuda_mul_mat_q( const int64_t s13 = ne12*s12; const mmq_args args = { - src0_d, src0->type, (const int *) src1_q8_1.ptr, nullptr, nullptr, dst_d, + src0_d, src0->type, (const int *) src1_q8_ptr, nullptr, nullptr, dst_d, src0->type == GGML_TYPE_NVFP4 && use_native_fp4 ? src1_scale.ptr : nullptr, ne00, ne01, ne1, s01, ne11, s1, ne02, ne12, s02, s12, s2, @@ -259,6 +284,68 @@ void ggml_cuda_mul_mat_q( ggml_cuda_mul_mat_q_switch_type(ctx, args, stream); } +size_t ggml_cuda_mul_mat_q_q8_size(const ggml_tensor * src0, const ggml_tensor * src1) { + GGML_ASSERT(src0 && src1 && src1->type == GGML_TYPE_F32); + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + const int64_t ne10_padded = GGML_PAD(src1->ne[0], MATRIX_ROW_PADDING); + const bool fallback = src0->ne[1] % 128 != 0; + return src1->ne[3]*src1->ne[2]*src1->ne[1]*ne10_padded*sizeof(block_q8_1_mmq)/QK8_1_MMQ + + ggml_cuda_mmq_get_J_max(src0->type, fallback, cc, src1->ne[1])*sizeof(block_q8_1_mmq); +} + +void ggml_cuda_mul_mat_q_fused_two( + ggml_backend_cuda_context & ctx, + const ggml_tensor * src0_a, const ggml_tensor * src0_b, const ggml_tensor * src1, + ggml_tensor * dst_a, ggml_tensor * dst_b, + const ggml_tensor * norm_weight, const float * norm_scale) { + GGML_ASSERT(src0_a && src0_b && src1 && dst_a && dst_b && norm_weight && norm_scale); + GGML_ASSERT(src0_a->type == src0_b->type && src1->type == GGML_TYPE_F32 && + dst_a->type == GGML_TYPE_F32 && dst_b->type == GGML_TYPE_F32); + GGML_ASSERT(src0_a->ne[0] == src1->ne[0] && src0_b->ne[0] == src1->ne[0]); + GGML_ASSERT(src0_a->ne[2] == src0_b->ne[2] && src0_a->ne[3] == src0_b->ne[3]); + GGML_ASSERT(ggml_is_contiguous(src1) && ggml_is_contiguous(norm_weight)); + GGML_ASSERT(norm_weight->type == GGML_TYPE_F32 && norm_weight->ne[0] == src1->ne[0] && + ggml_nrows(norm_weight) == 1); + + cudaStream_t stream = ctx.stream(); + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + const int64_t ne10 = src1->ne[0]; + const int64_t ne11 = src1->ne[1]; + const int64_t ne12 = src1->ne[2]; + const int64_t ne13 = src1->ne[3]; + const int64_t ne10_padded = GGML_PAD(ne10, MATRIX_ROW_PADDING); + const bool fallback = src0_a->ne[1] % 128 != 0; + GGML_ASSERT(fallback == (src0_b->ne[1] % 128 != 0)); + + const size_t nbytes_src1_q8_1 = ne13*ne12*ne11*ne10_padded*sizeof(block_q8_1_mmq)/QK8_1_MMQ + + ggml_cuda_mmq_get_J_max(src0_a->type, fallback, cc, ne11)*sizeof(block_q8_1_mmq); + ggml_cuda_pool_alloc src1_q8_1(ctx.pool(), nbytes_src1_q8_1); + const int64_t s11 = src1->nb[1] / sizeof(float); + const int64_t s12_src = src1->nb[2] / sizeof(float); + const int64_t s13_src = src1->nb[3] / sizeof(float); + quantize_mmq_q8_1_rms_cuda((const float *) src1->data, (const float *) norm_weight->data, + norm_scale, src1_q8_1.get(), src0_a->type, + ne10, s11, s12_src, s13_src, ne10_padded, ne11, ne12, ne13, stream); + CUDA_CHECK(cudaGetLastError()); + + const int64_t stride_q_channel = ne11*ne10_padded*sizeof(block_q8_1)/(QK8_1*sizeof(int)); + const int64_t stride_q_sample = ne12*stride_q_channel; + auto launch_one = [&](const ggml_tensor * src0, ggml_tensor * dst) { + GGML_ASSERT(src0->type == src0_a->type && src0->ne[0] == ne10); + const size_t ts0 = ggml_type_size(src0->type); + const mmq_args args = { + (const char *) src0->data, src0->type, (const int *) src1_q8_1.ptr, nullptr, nullptr, (float *) dst->data, + nullptr, + src0->ne[0], src0->ne[1], dst->ne[1], (int64_t)(src0->nb[1]/ts0), ne11, (int64_t)(dst->nb[1]/sizeof(float)), + src0->ne[2], ne12, (int64_t)(src0->nb[2]/ts0), stride_q_channel, (int64_t)(dst->nb[2]/sizeof(float)), + src0->ne[3], ne13, (int64_t)(src0->nb[3]/ts0), stride_q_sample, (int64_t)(dst->nb[3]/sizeof(float)), + dst->ne[1]}; + ggml_cuda_mul_mat_q_switch_type(ctx, args, stream); + }; + launch_one(src0_a, dst_a); + launch_one(src0_b, dst_b); +} + bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t n_experts) { #ifdef GGML_CUDA_FORCE_CUBLAS return false; diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh index 7cfcf01c485d..cfe7d4a49e04 100644 --- a/ggml/src/ggml-cuda/mmq.cuh +++ b/ggml/src/ggml-cuda/mmq.cuh @@ -1,6 +1,7 @@ #pragma once #include "common.cuh" +#include "cp-async.cuh" #include #include @@ -224,6 +225,7 @@ struct ggml_cuda_mmq_config { #include "mmq-config-rdna3-5.cuh" #include "mmq-config-rdna4.cuh" + #undef CASE static __host__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(const ggml_type type, const int J, const bool fallback, const int cc) { @@ -242,6 +244,9 @@ static __host__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(const ggml_type ty } return ggml_cuda_mmq_get_config_rdna2(type, J, fallback); } + if (cc == GGML_CUDA_CC_DGX_SPARK) { + return ggml_cuda_mmq_get_config_gb10(type, J, fallback); + } if (blackwell_mma_available(cc)) { return ggml_cuda_mmq_get_config_blackwell(type, J, fallback); } @@ -266,7 +271,11 @@ static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_t #endif // CDNA #else #ifdef BLACKWELL_MMA_AVAILABLE +#if __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK + return ggml_cuda_mmq_get_config_gb10(type, J, fallback); +#else return ggml_cuda_mmq_get_config_blackwell(type, J, fallback); +#endif #elif __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA return ggml_cuda_mmq_get_config_ampere(type, J, fallback); #else @@ -896,7 +905,15 @@ static __device__ __forceinline__ void mul_mat_q_process_tile( extern __shared__ int data_mul_mat_q[]; int * tile_y = data_mul_mat_q + J; - int * tile_x = tile_y + GGML_PAD(J*MMQ_TILE_Y_K, nwarps*warp_size); + constexpr int tile_y_stride = GGML_PAD(J*MMQ_TILE_Y_K, nwarps*warp_size); +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK + constexpr bool async_buffer_y = + type == GGML_TYPE_Q1_0 || type == GGML_TYPE_Q2_0 || type == GGML_TYPE_PQ2_0; +#else + constexpr bool async_buffer_y = false; +#endif + int * tile_y_next = tile_y + tile_y_stride; + int * tile_x = tile_y + (async_buffer_y ? 2 : 1)*tile_y_stride; #if defined(BLACKWELL_MMA_AVAILABLE) // FP4 tile stores 8 blocks @@ -913,8 +930,20 @@ static __device__ __forceinline__ void mul_mat_q_process_tile( constexpr int sz = sizeof(block_q8_1_mmq) / sizeof(int); for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) { - load_tiles(x, tile_x, offset_x + kb0, tile_x_max_i, stride_row_x); - { + if constexpr (async_buffer_y) { + const char * by0 = reinterpret_cast( + y + ncols_y * (kb0 * qk / ne_block) * sz); + char * tile_y_bytes = reinterpret_cast(tile_y); + const int tid = threadIdx.y*warp_size + threadIdx.x; +#pragma unroll + for (int byte0 = 16*tid; byte0 < J*MMQ_TILE_Y_K*int(sizeof(int)); byte0 += 16*nwarps*warp_size) { + cp_async_cg_16<256>( + ggml_cuda_cvta_generic_to_shared(tile_y_bytes + byte0), by0 + byte0); + } + load_tiles(x, tile_x, offset_x + kb0, tile_x_max_i, stride_row_x); + cp_async_wait_all(); + } else { + load_tiles(x, tile_x, offset_x + kb0, tile_x_max_i, stride_row_x); const int * by0 = y + ncols_y * (kb0 * qk / ne_block) * sz; #pragma unroll for (int l0 = 0; l0 < J * MMQ_TILE_Y_K; l0 += nwarps * warp_size) { @@ -926,6 +955,25 @@ static __device__ __forceinline__ void mul_mat_q_process_tile( __syncthreads(); + if constexpr (async_buffer_y) { + const char * by1 = reinterpret_cast( + y + ncols_y * ((kb0 * qk / ne_block) * sz + sz)); + char * tile_y_next_bytes = reinterpret_cast(tile_y_next); + const int tid = threadIdx.y*warp_size + threadIdx.x; +#pragma unroll + for (int byte0 = 16*tid; byte0 < J*MMQ_TILE_Y_K*int(sizeof(int)); byte0 += 16*nwarps*warp_size) { + cp_async_cg_16<256>( + ggml_cuda_cvta_generic_to_shared(tile_y_next_bytes + byte0), by1 + byte0); + } + + vec_dot(tile_x, tile_y, sum, 0); + cp_async_wait_all(); + __syncthreads(); + vec_dot(tile_x, tile_y_next, sum, MMQ_TILE_NE_K); + __syncthreads(); + continue; + } + vec_dot(tile_x, tile_y, sum, 0); __syncthreads(); @@ -1395,7 +1443,11 @@ static size_t mmq_get_nbytes_shared(const ggml_cuda_mmq_config & config, const i const size_t nbs_ids = config.J*sizeof(int); const size_t nbs_x = ggml_cuda_mmq_get_nbytes_shared_x(config, cc); const size_t nbs_y = config.J * (sizeof(block_q8_1_mmq)); - return nbs_ids + nbs_x + GGML_PAD(nbs_y, config.nthreads*sizeof(int)); + const size_t nbs_y_padded = GGML_PAD(nbs_y, config.nthreads*sizeof(int)); + const bool async_buffer_y = cc == GGML_CUDA_CC_DGX_SPARK && + (config.type == GGML_TYPE_Q1_0 || config.type == GGML_TYPE_Q2_0 || config.type == GGML_TYPE_PQ2_0); + const int y_buffers = async_buffer_y ? 2 : 1; + return nbs_ids + nbs_x + y_buffers*nbs_y_padded; } template @@ -1607,6 +1659,18 @@ extern DECL_MMQ_CASE(GGML_TYPE_NVFP4); // ------------------------------------------------------------------------------------------------------------------------- void ggml_cuda_mul_mat_q( - ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst); + ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, + ggml_tensor * dst, const ggml_tensor * gate = nullptr, + const ggml_tensor * norm_weight = nullptr, const ggml_tensor * norm_scale = nullptr, + void * external_q8 = nullptr, bool quantize_external = true, + const float * external_norm_scale = nullptr); + +size_t ggml_cuda_mul_mat_q_q8_size(const ggml_tensor * src0, const ggml_tensor * src1); + +void ggml_cuda_mul_mat_q_fused_two( + ggml_backend_cuda_context & ctx, + const ggml_tensor * src0_a, const ggml_tensor * src0_b, const ggml_tensor * src1, + ggml_tensor * dst_a, ggml_tensor * dst_b, + const ggml_tensor * norm_weight, const float * norm_scale); bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t n_experts); diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index 768b2394abc3..2f716efede72 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -501,9 +501,13 @@ static constexpr __host__ __device__ int calc_nwarps(ggml_type type, int ncols_d } if (table_id == MMVQ_PARAMETERS_GB10) { const int generic = calc_nwarps(type, ncols_dst, MMVQ_PARAMETERS_GENERIC); - // Only worth the wider block when it actually retires the K loop in half the trips (Observation) + // Only worth a wider block when it materially reduces K-loop trips (observation). if (ncols_dst == 1 && !small_k && halve_iters) { switch (type) { + case GGML_TYPE_Q1_0: + return generic; + case GGML_TYPE_PQ2_0: + return generic + generic / 2; case GGML_TYPE_Q4_0: case GGML_TYPE_Q4_1: case GGML_TYPE_Q5_0: @@ -543,7 +547,7 @@ static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int return 1; } -template +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( const void * vx_ptr, const void * vy_ptr, const int32_t * ids_ptr, const ggml_cuda_mm_fusion_args_device fusion, float * dst_ptr, @@ -586,7 +590,7 @@ static __global__ void mul_mat_vec_q( const uint32_t sample_x = fastdiv(sample_dst, sample_ratio); const uint32_t sample_y = sample_dst; - bool use_gate = false; + constexpr bool use_gate = has_gate; bool use_bias = false; bool use_gate_bias = false; bool use_scale = false; @@ -599,7 +603,6 @@ static __global__ void mul_mat_vec_q( ggml_glu_op active_glu; if constexpr (has_fusion) { - use_gate = fusion.gate != nullptr; use_bias = fusion.x_bias != nullptr; use_gate_bias = fusion.gate_bias != nullptr && use_gate; vgate = fusion.gate; @@ -657,22 +660,65 @@ static __global__ void mul_mat_vec_q( const block_q8_1 * y = ((const block_q8_1 *) vy) + sample_y*stride_sample_y + channel_y*stride_channel_y; const int kbx_offset = sample_x*stride_sample_x + channel_x*stride_channel_x + row0*stride_row_x; - for (int kbx = tid / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) { - const int kby = kbx * (qk/QK8_1); // y block index that aligns with kbx + if constexpr ((type == GGML_TYPE_Q1_0 || type == GGML_TYPE_Q2_0 || type == GGML_TYPE_PQ2_0) && + table_id == MMVQ_PARAMETERS_GB10) { + using block_t = std::conditional_t>; + // These packed AoS formats are scoreboard-latency bound on GB10. Prefetch one + // K iteration ahead; only the first lane that consumes a quant block issues it. +#pragma unroll (type == GGML_TYPE_Q2_0 ? 2 : 1) + for (int kbx = tid / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) { + const int kby = kbx * (qk/QK8_1); + const int kqs = vdr * (tid % (qi/vdr)); + const int kbx_prefetch = kbx + blocks_per_iter; + if (kbx_prefetch < blocks_per_row_x && tid % (qi/vdr) == 0) { +#pragma unroll + for (int i = 0; i < rows_per_cuda_block; ++i) { +#if defined(__CUDA_ARCH__) + const block_t * prefetch_ptr = (const block_t *) vx + + kbx_offset + i*stride_row_x + kbx_prefetch; + asm volatile("prefetch.global.L2 [%0];" :: "l"(__cvta_generic_to_global(prefetch_ptr))); + if constexpr (has_fusion && has_gate && type != GGML_TYPE_PQ2_0) { + const block_t * gate_prefetch_ptr = (const block_t *) vgate + + kbx_offset + i*stride_row_x + kbx_prefetch; + asm volatile("prefetch.global.L2 [%0];" :: "l"(__cvta_generic_to_global(gate_prefetch_ptr))); + } +#endif + } + } +#pragma unroll + for (int j = 0; j < ncols_dst; ++j) { +#pragma unroll + for (int i = 0; i < rows_per_cuda_block; ++i) { + tmp[j][i] += vec_dot_q_cuda( + vx, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + if constexpr (has_fusion) { + if constexpr (has_gate) { + tmp_gate[j][i] += vec_dot_q_cuda( + vgate, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + } + } + } + } + } + } else { + for (int kbx = tid / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) { + const int kby = kbx * (qk/QK8_1); // y block index that aligns with kbx - // x block quant index when casting the quants to int - const int kqs = vdr * (tid % (qi/vdr)); + // x block quant index when casting the quants to int + const int kqs = vdr * (tid % (qi/vdr)); #pragma unroll - for (int j = 0; j < ncols_dst; ++j) { + for (int j = 0; j < ncols_dst; ++j) { #pragma unroll - for (int i = 0; i < rows_per_cuda_block; ++i) { - tmp[j][i] += vec_dot_q_cuda( - vx, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); - if constexpr (has_fusion) { - if (use_gate) { - tmp_gate[j][i] += vec_dot_q_cuda( - vgate, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + for (int i = 0; i < rows_per_cuda_block; ++i) { + tmp[j][i] += vec_dot_q_cuda( + vx, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + if constexpr (has_fusion) { + if constexpr (has_gate) { + tmp_gate[j][i] += vec_dot_q_cuda( + vgate, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + } } } } @@ -689,7 +735,7 @@ static __global__ void mul_mat_vec_q( for (int i = 0; i < rows_per_cuda_block; ++i) { tmp_shared[threadIdx.y-1][j][i][threadIdx.x] = tmp[j][i]; if constexpr (has_fusion) { - if (use_gate) { + if constexpr (has_gate) { tmp_shared_gate[threadIdx.y-1][j][i][threadIdx.x] = tmp_gate[j][i]; } } @@ -712,14 +758,14 @@ static __global__ void mul_mat_vec_q( for (int l = 0; l < nwarps-1; ++l) { tmp[j][i] += tmp_shared[l][j][i][threadIdx.x]; if constexpr (has_fusion) { - if (use_gate) { + if constexpr (has_gate) { tmp_gate[j][i] += tmp_shared_gate[l][j][i][threadIdx.x]; } } } tmp[j][i] = warp_reduce_sum(tmp[j][i]); if constexpr (has_fusion) { - if (use_gate) { + if constexpr (has_gate) { tmp_gate[j][i] = warp_reduce_sum(tmp_gate[j][i]); } } @@ -731,7 +777,7 @@ static __global__ void mul_mat_vec_q( result *= x_scales; } result += x_biases[j]; - if (use_gate) { + if constexpr (has_gate) { float gate_value = tmp_gate[j][i]; if constexpr (type == GGML_TYPE_NVFP4) { gate_value *= gate_scales; @@ -863,10 +909,17 @@ static void mul_mat_vec_q_switch_fusion( if constexpr (c_ncols_dst == 1) { if (has_fusion) { const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, nbytes_shared, stream); - ggml_cuda_kernel_launch(mul_mat_vec_q, launch_params, - vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst, - channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, - sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride); + if (fusion.gate != nullptr) { + ggml_cuda_kernel_launch(mul_mat_vec_q, launch_params, + vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, + sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride); + } else { + ggml_cuda_kernel_launch(mul_mat_vec_q, launch_params, + vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, + sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride); + } return; } } @@ -874,7 +927,7 @@ static void mul_mat_vec_q_switch_fusion( GGML_ASSERT(!has_fusion && "fusion only supported for ncols_dst=1"); const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, nbytes_shared, stream); - ggml_cuda_kernel_launch(mul_mat_vec_q, launch_params, + ggml_cuda_kernel_launch(mul_mat_vec_q, launch_params, vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst, channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride); @@ -980,6 +1033,10 @@ static void mul_mat_vec_q_switch_ncols_dst( return false; } + if (type == GGML_TYPE_PQ2_0 && ncols_x == 6144 && nrows_x == 2048) { + return false; + } + // Expert rows are gathered per token, so a wider block adds reduction work without reuse. if (has_ids) { return false; @@ -1021,8 +1078,8 @@ static void mul_mat_vec_q_switch_ncols_dst( constexpr bool c_halve_iters = decltype(halve_iters_tag)::value && c_promoted; - const std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst, - nsamples_dst, warp_size, table_id, c_small_k, c_halve_iters); + std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst, + nsamples_dst, warp_size, table_id, c_small_k, c_halve_iters); mul_mat_vec_q_switch_fusion( vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio_fd, diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu index c3758cd50cfe..59eac4583b82 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu @@ -154,6 +154,59 @@ static __global__ void rms_norm_f32(const float * x, } } +template +static __global__ void add_rms_norm_f32( + const float * a, const float * b, const float * weight, float * sum, float * dst, + const int ncols, const float eps) { + const int64_t row = (int64_t) blockIdx.z*gridDim.y*gridDim.x + + (int64_t) blockIdx.y*gridDim.x + blockIdx.x; + const int tid = threadIdx.x; + a += row*ncols; + b += row*ncols; + sum += row*ncols; + dst += row*ncols; + + float tmp = 0.0f; + for (int col = tid; col < ncols; col += block_size) { + const float xi = __fadd_rn(a[col], b[col]); + sum[col] = xi; + tmp += xi*xi; + } + + extern __shared__ float s_sum[]; + tmp = block_reduce(tmp, s_sum); + const float scale = rsqrtf(tmp/ncols + eps); + + for (int col = tid; col < ncols; col += block_size) { + dst[col] = scale*sum[col]*weight[col]; + } +} + +template +static __global__ void add_rms_norm_scale_f32( + const float * a, const float * b, float * sum, float * row_scale, + const int ncols, const float eps) { + const int64_t row = (int64_t) blockIdx.z*gridDim.y*gridDim.x + + (int64_t) blockIdx.y*gridDim.x + blockIdx.x; + const int tid = threadIdx.x; + a += row*ncols; + b += row*ncols; + sum += row*ncols; + + float tmp = 0.0f; + for (int col = tid; col < ncols; col += block_size) { + const float xi = __fadd_rn(a[col], b[col]); + sum[col] = xi; + tmp += xi*xi; + } + + extern __shared__ float s_sum[]; + tmp = block_reduce(tmp, s_sum); + if (tid == 0) { + row_scale[row] = rsqrtf(tmp/ncols + eps); + } +} + template static __global__ void rms_norm_back_f32( const float * grad, const float * xf, float * dst, const int ncols, const float eps) { @@ -281,7 +334,18 @@ static void norm_f32_cuda( const float * x, float * dst, const int ncols, const int nrows, const int nchannels, const int nsamples, const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, cudaStream_t stream) { const dim3 blocks_num(nrows, nchannels, nsamples); - if (ncols < 1024) { + const char * rms128_env = getenv("GGML_CUDA_GB10_RMS128"); + const bool use_rms128 = ggml_cuda_info().devices[ggml_cuda_get_device()].cc == GGML_CUDA_CC_DGX_SPARK && + ncols <= 128 && (!rms128_env || std::atoi(rms128_env) != 0); + if (use_rms128) { + const dim3 block_dims(128, 1, 1); + const ggml_cuda_kernel_launch_params launch_params = {blocks_num, block_dims, 32 * sizeof(float), stream}; + ggml_cuda_kernel_launch(rms_norm_f32<128, false>, launch_params, + x, dst, ncols, stride_row, stride_channel, stride_sample, eps, + // underlying cudaLaunchKernelEx does not support default params + nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), + nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0)); + } else if (ncols < 1024) { const dim3 block_dims(WARP_SIZE, 1, 1); norm_f32<<>>(x, dst, ncols, stride_row, stride_channel, stride_sample, eps); } else { @@ -360,7 +424,18 @@ static void rms_norm_mul_f32_cuda(const float * x, const uint3 mul_nrows_packed = init_fastdiv_values(mul_nrows); const uint3 mul_nchannels_packed = init_fastdiv_values(mul_nchannels); const uint3 mul_nsamples_packed = init_fastdiv_values(mul_nsamples); - if (ncols < 1024) { + const char * rms128_env = getenv("GGML_CUDA_GB10_RMS128"); + const bool use_rms128 = ggml_cuda_info().devices[ggml_cuda_get_device()].cc == GGML_CUDA_CC_DGX_SPARK && + ncols <= 128 && (!rms128_env || std::atoi(rms128_env) != 0); + if (use_rms128) { + const dim3 block_dims(128, 1, 1); + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, 32 * sizeof(float), stream}; + ggml_cuda_kernel_launch(rms_norm_f32<128, true>, launch_params, + x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, + mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, + // underlying cudaLaunchKernelEx does not support default params + nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0)); + } else if (ncols < 1024) { const dim3 block_dims(256, 1, 1); const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream}; ggml_cuda_kernel_launch(rms_norm_f32<256, true>, launch_params, @@ -499,6 +574,71 @@ void ggml_cuda_op_rms_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { rms_norm_f32_cuda(src0_d, dst_d, ne00, ne01, ne02, ne03, s01, s02, s03, eps, stream); } +void ggml_cuda_op_add_rms_norm_fused( + ggml_backend_cuda_context & ctx, ggml_tensor * add, ggml_tensor * rms_norm, ggml_tensor * mul) { + const ggml_tensor * a = add->src[0]; + const ggml_tensor * b = add->src[1]; + GGML_ASSERT(a && b && rms_norm->src[0] == add && mul); + const ggml_tensor * weight = mul->src[0] == rms_norm ? mul->src[1] : mul->src[0]; + GGML_ASSERT(mul->src[0] == rms_norm || mul->src[1] == rms_norm); + GGML_ASSERT(a->type == GGML_TYPE_F32 && b->type == GGML_TYPE_F32 && + add->type == GGML_TYPE_F32 && rms_norm->type == GGML_TYPE_F32 && + weight->type == GGML_TYPE_F32 && mul->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_are_same_shape(a, b) && ggml_are_same_shape(a, add) && + ggml_are_same_shape(add, rms_norm) && ggml_are_same_shape(rms_norm, mul)); + GGML_ASSERT(ggml_is_contiguous(a) && ggml_is_contiguous(b) && + ggml_is_contiguous(add) && ggml_is_contiguous(rms_norm) && + ggml_is_contiguous(weight) && ggml_is_contiguous(mul)); + GGML_ASSERT(weight->ne[0] == add->ne[0] && ggml_nrows(weight) == 1); + + float eps; + memcpy(&eps, rms_norm->op_params, sizeof(float)); + GGML_ASSERT(eps >= 0.0f); + + const int ncols = add->ne[0]; + const dim3 grid(add->ne[1], add->ne[2], add->ne[3]); + cudaStream_t stream = ctx.stream(); + if (ncols < 1024) { + add_rms_norm_f32<256><<>>( + (const float *) a->data, (const float *) b->data, (const float *) weight->data, + (float *) add->data, (float *) mul->data, ncols, eps); + } else { + add_rms_norm_f32<1024><<>>( + (const float *) a->data, (const float *) b->data, (const float *) weight->data, + (float *) add->data, (float *) mul->data, ncols, eps); + } +} + +void ggml_cuda_op_add_rms_norm_scale_fused( + ggml_backend_cuda_context & ctx, ggml_tensor * add, ggml_tensor * rms_norm, float * row_scale) { + const ggml_tensor * a = add->src[0]; + const ggml_tensor * b = add->src[1]; + GGML_ASSERT(a && b && rms_norm->src[0] == add && row_scale); + GGML_ASSERT(a->type == GGML_TYPE_F32 && b->type == GGML_TYPE_F32 && + add->type == GGML_TYPE_F32 && rms_norm->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_are_same_shape(a, b) && ggml_are_same_shape(a, add) && + ggml_are_same_shape(add, rms_norm)); + GGML_ASSERT(ggml_is_contiguous(a) && ggml_is_contiguous(b) && + ggml_is_contiguous(add)); + + float eps; + memcpy(&eps, rms_norm->op_params, sizeof(float)); + GGML_ASSERT(eps >= 0.0f); + + const int ncols = add->ne[0]; + const dim3 grid(add->ne[1], add->ne[2], add->ne[3]); + cudaStream_t stream = ctx.stream(); + if (ncols < 1024) { + add_rms_norm_scale_f32<256><<>>( + (const float *) a->data, (const float *) b->data, (float *) add->data, + row_scale, ncols, eps); + } else { + add_rms_norm_scale_f32<1024><<>>( + (const float *) a->data, (const float *) b->data, (float *) add->data, + row_scale, ncols, eps); + } +} + void ggml_cuda_op_rms_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * mul_tensor) { const ggml_tensor * rms_norm_src = (ggml_tensor *) dst->src[0]; float eps = 0.0f; diff --git a/ggml/src/ggml-cuda/norm.cuh b/ggml/src/ggml-cuda/norm.cuh index a74f6376720a..1f7417f04c0c 100644 --- a/ggml/src/ggml-cuda/norm.cuh +++ b/ggml/src/ggml-cuda/norm.cuh @@ -6,6 +6,12 @@ void ggml_cuda_op_group_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) void ggml_cuda_op_rms_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst); +void ggml_cuda_op_add_rms_norm_fused( + ggml_backend_cuda_context & ctx, ggml_tensor * add, ggml_tensor * rms_norm, ggml_tensor * mul); + +void ggml_cuda_op_add_rms_norm_scale_fused( + ggml_backend_cuda_context & ctx, ggml_tensor * add, ggml_tensor * rms_norm, float * row_scale); + void ggml_cuda_op_rms_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * mul_tensor); void ggml_cuda_op_rms_norm_fused_add(ggml_backend_cuda_context & ctx, diff --git a/ggml/src/ggml-cuda/quantize.cu b/ggml/src/ggml-cuda/quantize.cu index bcc7723957b5..52a025373aae 100644 --- a/ggml/src/ggml-cuda/quantize.cu +++ b/ggml/src/ggml-cuda/quantize.cu @@ -1,4 +1,5 @@ #include "quantize.cuh" +#include "unary.cuh" #include #if defined(BLACKWELL_MMA_AVAILABLE) @@ -454,11 +455,14 @@ static __global__ void quantize_mmq_mxfp4(const float * __restrict__ x, } // scatter: grid over tokens, quantize once, write to all the token's compact rows -template +template static __global__ void quantize_mmq_q8_1( const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy, const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03, - const int64_t ne0, const int ne1, const int ne2, const int n_expert_used) { + const int64_t ne0, const int ne1, const int ne2, const int n_expert_used, + const float * __restrict__ gate = nullptr, + const float * __restrict__ norm_weight = nullptr, + const float * __restrict__ norm_scale = nullptr) { constexpr int vals_per_scale = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 64 : 32; constexpr int vals_per_sum = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 16 : 32; @@ -489,7 +493,26 @@ static __global__ void quantize_mmq_q8_1( const int64_t iqs = i0 % QK8_1_MMQ; // quant index in block // Load 4 floats per thread and calculate max. abs. value between them: - const float4 xi = i0 < ne00 ? x4[(base_idx + i00)/4] : make_float4(0.0f, 0.0f, 0.0f, 0.0f); + float4 xi = i0 < ne00 ? x4[(base_idx + i00)/4] : make_float4(0.0f, 0.0f, 0.0f, 0.0f); + if constexpr (swiglu) { + const float4 gi = ((const float4 *) gate)[(base_idx + i00)/4]; + xi.x *= ggml_cuda_op_silu_single(gi.x); + xi.y *= ggml_cuda_op_silu_single(gi.y); + xi.z *= ggml_cuda_op_silu_single(gi.z); + xi.w *= ggml_cuda_op_silu_single(gi.w); + } + if constexpr (rms_scale) { + const int64_t row = ((int64_t) blockIdx.z)*ne1 + blockIdx.x; + const float scale = norm_scale[row]; + const float4 wi = i0 < ne00 ? ((const float4 *) norm_weight)[i00/4] : make_float4(0, 0, 0, 0); + // Preserve the operation order of the unfused RMS_NORM -> MUL graph: + // (scale * x) * weight. Reassociating this as x * (scale * weight) + // can change the last bit and, in turn, a greedy decoding trajectory. + xi.x = (scale*xi.x)*wi.x; + xi.y = (scale*xi.y)*wi.y; + xi.z = (scale*xi.z)*wi.z; + xi.w = (scale*xi.w)*wi.w; + } float amax = fabsf(xi.x); amax = fmaxf(amax, fabsf(xi.y)); amax = fmaxf(amax, fabsf(xi.z)); @@ -555,6 +578,62 @@ static __global__ void quantize_mmq_q8_1( GGML_UNUSED(n_expert_used); } +void quantize_mmq_q8_1_swiglu_cuda( + const float * x, const float * gate, void * vy, const ggml_type type_src0, + const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03, + const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, cudaStream_t stream) { + GGML_ASSERT(ne00 % 4 == 0); + GGML_ASSERT(ne0 % QK8_1_MMQ == 0); + + const int64_t block_num_y = (ne0 + 4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ); + const dim3 num_blocks(ne1, block_num_y, ne2*ne3); + const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1); + switch (mmq_get_q8_1_ds_layout(type_src0)) { + case MMQ_Q8_1_DS_LAYOUT_D4: + quantize_mmq_q8_1<<>>( + x, nullptr, vy, ne00, s01, s02, s03, ne0, ne1, ne2, 0, gate); + break; + case MMQ_Q8_1_DS_LAYOUT_DS4: + quantize_mmq_q8_1<<>>( + x, nullptr, vy, ne00, s01, s02, s03, ne0, ne1, ne2, 0, gate); + break; + case MMQ_Q8_1_DS_LAYOUT_D2S6: + quantize_mmq_q8_1<<>>( + x, nullptr, vy, ne00, s01, s02, s03, ne0, ne1, ne2, 0, gate); + break; + default: + GGML_ABORT("fatal error"); + } +} + +void quantize_mmq_q8_1_rms_cuda( + const float * x, const float * weight, const float * scale, void * vy, const ggml_type type_src0, + const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03, + const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, cudaStream_t stream) { + GGML_ASSERT(ne00 % 4 == 0); + GGML_ASSERT(ne0 % QK8_1_MMQ == 0); + + const int64_t block_num_y = (ne0 + 4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ); + const dim3 num_blocks(ne1, block_num_y, ne2*ne3); + const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1); + switch (mmq_get_q8_1_ds_layout(type_src0)) { + case MMQ_Q8_1_DS_LAYOUT_D4: + quantize_mmq_q8_1<<>>( + x, nullptr, vy, ne00, s01, s02, s03, ne0, ne1, ne2, 0, nullptr, weight, scale); + break; + case MMQ_Q8_1_DS_LAYOUT_DS4: + quantize_mmq_q8_1<<>>( + x, nullptr, vy, ne00, s01, s02, s03, ne0, ne1, ne2, 0, nullptr, weight, scale); + break; + case MMQ_Q8_1_DS_LAYOUT_D2S6: + quantize_mmq_q8_1<<>>( + x, nullptr, vy, ne00, s01, s02, s03, ne0, ne1, ne2, 0, nullptr, weight, scale); + break; + default: + GGML_ABORT("fatal error"); + } +} + void quantize_row_q8_1_cuda( const float * x, const int32_t * ids, void * vy, const ggml_type type_src0, const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03, diff --git a/ggml/src/ggml-cuda/quantize.cuh b/ggml/src/ggml-cuda/quantize.cuh index 5f08dcbfe331..335ff587adb5 100644 --- a/ggml/src/ggml-cuda/quantize.cuh +++ b/ggml/src/ggml-cuda/quantize.cuh @@ -26,6 +26,16 @@ void quantize_mmq_q8_1_cuda( ggml_type type_src0, int64_t ne00, int64_t s01, int64_t s02, int64_t s03, int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream); +void quantize_mmq_q8_1_swiglu_cuda( + const float * x, const float * gate, void * vy, ggml_type type_src0, + int64_t ne00, int64_t s01, int64_t s02, int64_t s03, + int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream); + +void quantize_mmq_q8_1_rms_cuda( + const float * x, const float * weight, const float * scale, void * vy, ggml_type type_src0, + int64_t ne00, int64_t s01, int64_t s02, int64_t s03, + int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream); + void quantize_mmq_fp4_cuda(const float * x, const int32_t * ids, void * vy, diff --git a/ggml/src/ggml-cuda/scale.cu b/ggml/src/ggml-cuda/scale.cu index 7b2e59a4383f..c7c80335343e 100644 --- a/ggml/src/ggml-cuda/scale.cu +++ b/ggml/src/ggml-cuda/scale.cu @@ -13,7 +13,36 @@ static __global__ void scale_f32(const float * x, float * dst, const float scale } } +static __global__ void scale_f32_vec4( + const float4 * x, float4 * dst, const float scale, const float bias, const int64_t nelements4) { + ggml_cuda_pdl_lc(); + const int64_t tid = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + const int64_t stride = (int64_t) blockDim.x * gridDim.x; + + ggml_cuda_pdl_sync(); + for (int64_t i = tid; i < nelements4; i += stride) { + const float4 v = x[i]; + dst[i] = make_float4( + scale * v.x + bias, + scale * v.y + bias, + scale * v.z + bias, + scale * v.w + bias); + } +} + static void scale_f32_cuda(const float * x, float * dst, const float scale, const float bias, const int64_t nelements, cudaStream_t stream) { + const int device = ggml_cuda_get_device(); + const int cc = ggml_cuda_info().devices[device].cc; + if (cc == GGML_CUDA_CC_DGX_SPARK && nelements >= 1024 && nelements % 4 == 0 && + (uintptr_t(x) & 0x0F) == 0 && (uintptr_t(dst) & 0x0F) == 0) { + const int64_t nelements4 = nelements / 4; + const int64_t num_blocks = (nelements4 + CUDA_SCALE_BLOCK_SIZE - 1) / CUDA_SCALE_BLOCK_SIZE; + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params( + MIN(MAX_GRIDDIM_X, num_blocks), CUDA_SCALE_BLOCK_SIZE, 0, stream); + ggml_cuda_kernel_launch(scale_f32_vec4, launch_params, + (const float4 *) x, (float4 *) dst, scale, bias, nelements4); + return; + } const int64_t num_blocks = (nelements + CUDA_SCALE_BLOCK_SIZE - 1) / CUDA_SCALE_BLOCK_SIZE; const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(MIN(MAX_GRIDDIM_X, num_blocks), CUDA_SCALE_BLOCK_SIZE, 0, stream); ggml_cuda_kernel_launch(scale_f32, launch_params, x, dst, scale, bias, nelements); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index cca08f5dc5fd..3714b727775c 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10182,6 +10182,7 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 100, 1)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 200, 1)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 127, 2)); + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 128, 64, 1)); // S=128 long-token CUDA path test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 64, 1, 1, false, true)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 33, 1, 1, false, true)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 100, 1, 1, false, true)); From fff11bb57420d5cb352b5e649af88271004c744e Mon Sep 17 00:00:00 2001 From: Pooya Khosravi <16091823+pooyakhosravi@users.noreply.github.com> Date: Thu, 3 Sep 2026 12:50:41 -0700 Subject: [PATCH 2/4] cuda: address GB10 review feedback --- ggml/src/ggml-cuda/gated_delta_net.cu | 2 +- ggml/src/ggml-cuda/ggml-cuda.cu | 95 +++++++++++++++------------ ggml/src/ggml-cuda/mmq.cu | 29 +++++--- ggml/src/ggml-cuda/mmq.cuh | 2 + ggml/src/ggml-cuda/norm.cu | 36 +++++----- ggml/src/ggml-cuda/quantize.cu | 3 +- tests/test-backend-ops.cpp | 1 - 7 files changed, 100 insertions(+), 68 deletions(-) diff --git a/ggml/src/ggml-cuda/gated_delta_net.cu b/ggml/src/ggml-cuda/gated_delta_net.cu index f52a2368283f..8ac2afd39c4d 100644 --- a/ggml/src/ggml-cuda/gated_delta_net.cu +++ b/ggml/src/ggml-cuda/gated_delta_net.cu @@ -323,7 +323,7 @@ static void ggml_cuda_op_gated_delta_net_impl( if (!kda && S_v == 128 && n_tokens >= 32 && ggml_cuda_info().devices[ggml_cuda_get_device()].cc == GGML_CUDA_CC_DGX_SPARK) { const int64_t n_g = ggml_nelements(src_g); - g_exp_alloc.alloc(n_g*sizeof(float)); + g_exp_alloc.alloc(n_g); const int block = 256; const int grid = std::min((n_g + block - 1)/block, 4096); gdn_precompute_exp<<>>(g_d, g_exp_alloc.ptr, n_g); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index c066e53a6bf8..4504cc2a6e26 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -3288,6 +3288,10 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph if (disable_fusion) { return 0; } + static const bool dual_rms_q8_enabled = [] { + const char * env = getenv("GGML_CUDA_GB10_DUAL_RMS_Q8"); + return !env || std::atoi(env) != 0; + }(); ggml_tensor * node = cgraph->nodes[i]; @@ -3323,9 +3327,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; const ggml_op ops[] = { GGML_OP_ADD, GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_MUL_MAT, GGML_OP_MUL_MAT }; const int out_nodes[] = { i, i + 3, i + 4 }; - const char * dual_rms_q8 = getenv("GGML_CUDA_GB10_DUAL_RMS_Q8"); - if ((!dual_rms_q8 || std::atoi(dual_rms_q8) != 0) && - cc == GGML_CUDA_CC_DGX_SPARK && rms_norm->op == GGML_OP_RMS_NORM && + if (dual_rms_q8_enabled && cc == GGML_CUDA_CC_DGX_SPARK && rms_norm->op == GGML_OP_RMS_NORM && mul->op == GGML_OP_MUL && mm_a->op == GGML_OP_MUL_MAT && mm_b->op == GGML_OP_MUL_MAT && mm_a->src[1] == mul && mm_b->src[1] == mul && (mul->src[0] == rms_norm || mul->src[1] == rms_norm) && rms_norm->src[0] == node && @@ -3341,7 +3343,8 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph mm_a->src[0]->type == GGML_TYPE_PQ2_0) && ggml_cuda_should_use_mmq(mm_a->src[0]->type, cc, mul->ne[1], 0) && ggml_cuda_should_use_mmq(mm_b->src[0]->type, cc, mul->ne[1], 0) && - ggml_can_fuse_subgraph(cgraph, i, 5, ops, out_nodes, 3)) { + ggml_can_fuse_subgraph(cgraph, i, 5, ops, out_nodes, 3) && + ggml_cuda_check_fusion_memory_ranges(cgraph, i, 5, out_nodes, 3)) { const ggml_tensor * weight = mul->src[0] == rms_norm ? mul->src[1] : mul->src[0]; if (weight && weight->type == GGML_TYPE_F32 && ggml_is_contiguous(weight) && weight->ne[0] == node->ne[0] && ggml_nrows(weight) == 1) { @@ -3371,7 +3374,8 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph ggml_are_same_shape(node, rms_norm) && ggml_is_contiguous(node->src[0]) && ggml_is_contiguous(node->src[1]) && ggml_is_contiguous(node) && ggml_is_contiguous(rms_norm) && ggml_is_contiguous(mul) && - ggml_can_fuse_subgraph(cgraph, i, 3, ops, out_nodes, 2)) { + ggml_can_fuse_subgraph(cgraph, i, 3, ops, out_nodes, 2) && + ggml_cuda_check_fusion_memory_ranges(cgraph, i, 3, out_nodes, 2)) { const ggml_tensor * weight = mul->src[0] == rms_norm ? mul->src[1] : mul->src[0]; if (weight && weight->type == GGML_TYPE_F32 && ggml_is_contiguous(weight) && weight->ne[0] == node->ne[0] && ggml_nrows(weight) == 1) { @@ -4146,7 +4150,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud struct gb10_shared_q8_entry { const ggml_tensor * src1; ggml_type type; - std::unique_ptr> data; + ggml_cuda_pool_alloc * data; bool quantized = false; int remaining = 0; }; @@ -4157,21 +4161,35 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud ggml_tensor * residual; const ggml_tensor * weight; ggml_type type; - std::unique_ptr> row_scale; + ggml_cuda_pool_alloc * row_scale; int remaining; }; std::vector gb10_virtual_rms; + // Keep every persistent pool allocation in true allocation order. The VMM + // scratch pool is stack-like, so all releases happen globally in reverse. + std::vector>> gb10_pool_allocations; - const auto gb10_shared_q8_consumer_count = [&](const ggml_tensor * src1, ggml_type type) { - int count = 0; - for (int j = 0; j < cgraph->n_nodes; ++j) { - const ggml_tensor * candidate = cgraph->nodes[j]; - count += candidate->op == GGML_OP_MUL_MAT && candidate->src[0] && - candidate->src[1] == src1 && candidate->src[0]->type == type; + std::map> gb10_shared_q8_consumer_counts; + for (int j = 0; j < cgraph->n_nodes; ++j) { + const ggml_tensor * candidate = cgraph->nodes[j]; + if (candidate->op == GGML_OP_MUL_MAT && candidate->src[0] && candidate->src[1]) { + ++gb10_shared_q8_consumer_counts[candidate->src[1]][candidate->src[0]->type]; } - return count; + } + const auto gb10_shared_q8_consumer_count = [&](const ggml_tensor * src1, ggml_type type) { + const auto it = gb10_shared_q8_consumer_counts.find(src1); + return it == gb10_shared_q8_consumer_counts.end() ? 0 : it->second[type]; }; + static const bool virtual_rms_q8_enabled = [] { + const char * env = getenv("GGML_CUDA_GB10_VIRTUAL_RMS_Q8"); + return !env || std::atoi(env) != 0; + }(); + static const bool shared_q8_enabled = [] { + const char * env = getenv("GGML_CUDA_GB10_SHARED_Q8"); + return !env || std::atoi(env) != 0; + }(); + const auto try_launch_concurrent_event = [&](const ggml_tensor * node) { if (stream_ctx.concurrent_events.find(node) != stream_ctx.concurrent_events.end()) { concurrent_event = &stream_ctx.concurrent_events[node]; @@ -4308,10 +4326,8 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud // The normalized pre-attention residual is consumed only by a // group of low-bit projections. Preserve residual + one scale per // row and let their shared Q8 quantizer apply the norm weight. - const char * virtual_rms_env = getenv("GGML_CUDA_GB10_VIRTUAL_RMS_Q8"); const int virtual_rms_cc = ggml_cuda_info().devices[cuda_ctx->device].cc; - if ((!virtual_rms_env || std::atoi(virtual_rms_env) != 0) && - virtual_rms_cc == GGML_CUDA_CC_DGX_SPARK && !is_concurrent_event_active && + if (virtual_rms_q8_enabled && virtual_rms_cc == GGML_CUDA_CC_DGX_SPARK && !is_concurrent_event_active && node->op == GGML_OP_ADD && i + 2 < cgraph->n_nodes) { ggml_tensor * rms = cgraph->nodes[i + 1]; ggml_tensor * mul = cgraph->nodes[i + 2]; @@ -4351,10 +4367,15 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud ggml_is_contiguous(node->src[1]) && ggml_is_contiguous(node) && ggml_is_contiguous(mul) && ggml_is_contiguous(weight) && weight->ne[0] == node->ne[0] && ggml_nrows(weight) == 1 && node->ne[1] >= 32 && - ggml_can_fuse_subgraph(cgraph, i, 3, ops, out_nodes, 2)) { - auto row_scale = std::make_unique>(cuda_ctx->pool(), ggml_nrows(node)); - ggml_cuda_op_add_rms_norm_scale_fused(*cuda_ctx, node, rms, row_scale->get()); - gb10_virtual_rms.push_back({ mul, node, weight, consumer_type, std::move(row_scale), consumers }); + ggml_can_fuse_subgraph(cgraph, i, 3, ops, out_nodes, 2) && + ggml_cuda_check_fusion_memory_ranges(cgraph, i, 3, out_nodes, 2)) { + auto row_scale = std::make_unique>( + cuda_ctx->pool(), ggml_nrows(node)*sizeof(float)); + ggml_cuda_pool_alloc * row_scale_ptr = row_scale.get(); + gb10_pool_allocations.push_back(std::move(row_scale)); + ggml_cuda_op_add_rms_norm_scale_fused( + *cuda_ctx, node, rms, (float *) row_scale_ptr->get()); + gb10_virtual_rms.push_back({ mul, node, weight, consumer_type, row_scale_ptr, consumers }); i += 2; continue; } @@ -4376,10 +4397,8 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud // Several Qwen3.5 attention projections consume the exact same // normalized activation. Quantize it once per graph execution and // reuse the Q8 tile for the later MMQs on the same CUDA stream. - const char * shared_q8_env = getenv("GGML_CUDA_GB10_SHARED_Q8"); const int cc = ggml_cuda_info().devices[cuda_ctx->device].cc; - if ((!shared_q8_env || std::atoi(shared_q8_env) != 0) && - cc == GGML_CUDA_CC_DGX_SPARK && !is_concurrent_event_active && + if (shared_q8_enabled && cc == GGML_CUDA_CC_DGX_SPARK && !is_concurrent_event_active && node->op == GGML_OP_MUL_MAT && node->src[0] && node->src[1] && node->src[1]->type == GGML_TYPE_F32 && ggml_is_contiguous(node->src[1]) && node->src[1]->ne[1] >= 32 && @@ -4398,23 +4417,20 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud if (it == gb10_shared_q8.end()) { auto data = std::make_unique>(cuda_ctx->pool(), ggml_cuda_mul_mat_q_q8_size(node->src[0], node->src[1])); - gb10_shared_q8.push_back({ node->src[1], node->src[0]->type, std::move(data), false, + ggml_cuda_pool_alloc * data_ptr = data.get(); + gb10_pool_allocations.push_back(std::move(data)); + gb10_shared_q8.push_back({ node->src[1], node->src[0]->type, data_ptr, false, gb10_shared_q8_consumer_count(node->src[1], node->src[0]->type) }); it = std::prev(gb10_shared_q8.end()); - } else if (!it->data) { - it->data = std::make_unique>(cuda_ctx->pool(), - ggml_cuda_mul_mat_q_q8_size(node->src[0], node->src[1])); } ggml_cuda_mul_mat_q(*cuda_ctx, node->src[0], virtual_rms ? virtual_it->residual : node->src[1], nullptr, node, nullptr, virtual_rms ? virtual_it->weight : nullptr, nullptr, it->data->get(), !it->quantized, - virtual_rms ? virtual_it->row_scale->get() : nullptr); + virtual_rms ? (float *) virtual_it->row_scale->get() : nullptr); it->quantized = true; - if (--it->remaining == 0) { - it->data.reset(); - } - if (virtual_rms && --virtual_it->remaining == 0) { - virtual_it->row_scale.reset(); + --it->remaining; + if (virtual_rms) { + --virtual_it->remaining; } if (!is_concurrent_event_active) { try_launch_concurrent_event(node); @@ -4487,13 +4503,10 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud #endif // USE_CUDA_GRAPH } - // The VMM scratch pool is stack-like. Persistent shared-Q8 allocations are - // created in graph order, so release them explicitly in reverse order. - for (auto it = gb10_shared_q8.rbegin(); it != gb10_shared_q8.rend(); ++it) { - it->data.reset(); - } - for (auto it = gb10_virtual_rms.rbegin(); it != gb10_virtual_rms.rend(); ++it) { - it->row_scale.reset(); + // The VMM scratch pool is stack-like, so release all persistent allocations + // explicitly in reverse order across both shared-Q8 and row-scale buffers. + for (auto it = gb10_pool_allocations.rbegin(); it != gb10_pool_allocations.rend(); ++it) { + it->reset(); } } diff --git a/ggml/src/ggml-cuda/mmq.cu b/ggml/src/ggml-cuda/mmq.cu index 148a6646d80f..d68331d78829 100644 --- a/ggml/src/ggml-cuda/mmq.cu +++ b/ggml/src/ggml-cuda/mmq.cu @@ -85,6 +85,10 @@ static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, con } } +static size_t ggml_cuda_mmq_q8_buffer_size( + ggml_type type, bool fallback, int cc, + int64_t ne10_padded, int64_t ne11, int64_t ne12, int64_t ne13); + void ggml_cuda_mul_mat_q( ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, const ggml_tensor * gate, @@ -143,8 +147,10 @@ void ggml_cuda_mul_mat_q( const size_t y_values_per_block = use_native_fp4 ? QK_FP4_MMQ : QK8_1_MMQ; if (!ids) { - const size_t nbytes_src1_q8_1 = ne13*ne12 * ne11*ne10_padded * y_block_size/y_values_per_block + - ggml_cuda_mmq_get_J_max(src0->type, fallback, cc, ne11) * sizeof(block_q8_1_mmq); + const size_t nbytes_src1_q8_1 = use_native_fp4 ? + ne13*ne12*ne11*ne10_padded*y_block_size/y_values_per_block + + ggml_cuda_mmq_get_J_max(src0->type, fallback, cc, ne11)*sizeof(block_q8_1_mmq) : + ggml_cuda_mmq_q8_buffer_size(src0->type, fallback, cc, ne10_padded, ne11, ne12, ne13); GGML_ASSERT(!external_q8 || (!gate && !use_native_fp4)); ggml_cuda_pool_alloc src1_q8_1(ctx.pool()); char * src1_q8_ptr = external_q8 ? (char *) external_q8 : src1_q8_1.alloc(nbytes_src1_q8_1); @@ -284,13 +290,18 @@ void ggml_cuda_mul_mat_q( ggml_cuda_mul_mat_q_switch_type(ctx, args, stream); } +static size_t ggml_cuda_mmq_q8_buffer_size( + ggml_type type, bool fallback, int cc, + int64_t ne10_padded, int64_t ne11, int64_t ne12, int64_t ne13) { + return ne13*ne12*ne11*ne10_padded*sizeof(block_q8_1_mmq)/QK8_1_MMQ + + ggml_cuda_mmq_get_J_max(type, fallback, cc, ne11)*sizeof(block_q8_1_mmq); +} + size_t ggml_cuda_mul_mat_q_q8_size(const ggml_tensor * src0, const ggml_tensor * src1) { GGML_ASSERT(src0 && src1 && src1->type == GGML_TYPE_F32); const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; - const int64_t ne10_padded = GGML_PAD(src1->ne[0], MATRIX_ROW_PADDING); - const bool fallback = src0->ne[1] % 128 != 0; - return src1->ne[3]*src1->ne[2]*src1->ne[1]*ne10_padded*sizeof(block_q8_1_mmq)/QK8_1_MMQ + - ggml_cuda_mmq_get_J_max(src0->type, fallback, cc, src1->ne[1])*sizeof(block_q8_1_mmq); + return ggml_cuda_mmq_q8_buffer_size(src0->type, src0->ne[1] % 128 != 0, cc, + GGML_PAD(src1->ne[0], MATRIX_ROW_PADDING), src1->ne[1], src1->ne[2], src1->ne[3]); } void ggml_cuda_mul_mat_q_fused_two( @@ -317,8 +328,8 @@ void ggml_cuda_mul_mat_q_fused_two( const bool fallback = src0_a->ne[1] % 128 != 0; GGML_ASSERT(fallback == (src0_b->ne[1] % 128 != 0)); - const size_t nbytes_src1_q8_1 = ne13*ne12*ne11*ne10_padded*sizeof(block_q8_1_mmq)/QK8_1_MMQ + - ggml_cuda_mmq_get_J_max(src0_a->type, fallback, cc, ne11)*sizeof(block_q8_1_mmq); + const size_t nbytes_src1_q8_1 = ggml_cuda_mmq_q8_buffer_size( + src0_a->type, fallback, cc, ne10_padded, ne11, ne12, ne13); ggml_cuda_pool_alloc src1_q8_1(ctx.pool(), nbytes_src1_q8_1); const int64_t s11 = src1->nb[1] / sizeof(float); const int64_t s12_src = src1->nb[2] / sizeof(float); @@ -328,7 +339,7 @@ void ggml_cuda_mul_mat_q_fused_two( ne10, s11, s12_src, s13_src, ne10_padded, ne11, ne12, ne13, stream); CUDA_CHECK(cudaGetLastError()); - const int64_t stride_q_channel = ne11*ne10_padded*sizeof(block_q8_1)/(QK8_1*sizeof(int)); + const int64_t stride_q_channel = ne11*ne10_padded*sizeof(block_q8_1_mmq)/(QK8_1_MMQ*sizeof(int)); const int64_t stride_q_sample = ne12*stride_q_channel; auto launch_one = [&](const ggml_tensor * src0, ggml_tensor * dst) { GGML_ASSERT(src0->type == src0_a->type && src0->ne[0] == ne10); diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh index cfe7d4a49e04..339ec9bc3a51 100644 --- a/ggml/src/ggml-cuda/mmq.cuh +++ b/ggml/src/ggml-cuda/mmq.cuh @@ -969,6 +969,8 @@ static __device__ __forceinline__ void mul_mat_q_process_tile( vec_dot(tile_x, tile_y, sum, 0); cp_async_wait_all(); __syncthreads(); + // by1 is the second half of the same K tile; using a separate shared-memory + // buffer changes only its address, while MMQ_TILE_NE_K preserves its logical offset. vec_dot(tile_x, tile_y_next, sum, MMQ_TILE_NE_K); __syncthreads(); continue; diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu index 59eac4583b82..d8760a4ba158 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu @@ -334,18 +334,7 @@ static void norm_f32_cuda( const float * x, float * dst, const int ncols, const int nrows, const int nchannels, const int nsamples, const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, cudaStream_t stream) { const dim3 blocks_num(nrows, nchannels, nsamples); - const char * rms128_env = getenv("GGML_CUDA_GB10_RMS128"); - const bool use_rms128 = ggml_cuda_info().devices[ggml_cuda_get_device()].cc == GGML_CUDA_CC_DGX_SPARK && - ncols <= 128 && (!rms128_env || std::atoi(rms128_env) != 0); - if (use_rms128) { - const dim3 block_dims(128, 1, 1); - const ggml_cuda_kernel_launch_params launch_params = {blocks_num, block_dims, 32 * sizeof(float), stream}; - ggml_cuda_kernel_launch(rms_norm_f32<128, false>, launch_params, - x, dst, ncols, stride_row, stride_channel, stride_sample, eps, - // underlying cudaLaunchKernelEx does not support default params - nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), - nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0)); - } else if (ncols < 1024) { + if (ncols < 1024) { const dim3 block_dims(WARP_SIZE, 1, 1); norm_f32<<>>(x, dst, ncols, stride_row, stride_channel, stride_sample, eps); } else { @@ -369,7 +358,21 @@ static void rms_norm_f32_cuda( const float * x, float * dst, const int ncols, const int nrows, const int nchannels, const int nsamples, const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, cudaStream_t stream) { const dim3 blocks_num(nrows, nchannels, nsamples); - if (ncols < 1024) { + static const bool rms128_enabled = [] { + const char * env = getenv("GGML_CUDA_GB10_RMS128"); + return !env || std::atoi(env) != 0; + }(); + const bool use_rms128 = ggml_cuda_info().devices[ggml_cuda_get_device()].cc == GGML_CUDA_CC_DGX_SPARK && + ncols <= 128 && rms128_enabled; + if (use_rms128) { + const dim3 block_dims(128, 1, 1); + const ggml_cuda_kernel_launch_params launch_params = {blocks_num, block_dims, 32 * sizeof(float), stream}; + ggml_cuda_kernel_launch(rms_norm_f32<128, false>, launch_params, + x, dst, ncols, stride_row, stride_channel, stride_sample, eps, + // underlying cudaLaunchKernelEx does not support default params + nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), + nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0)); + } else if (ncols < 1024) { const dim3 block_dims(256, 1, 1); const ggml_cuda_kernel_launch_params launch_params = {blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream}; ggml_cuda_kernel_launch(rms_norm_f32<256, false>, launch_params, @@ -424,9 +427,12 @@ static void rms_norm_mul_f32_cuda(const float * x, const uint3 mul_nrows_packed = init_fastdiv_values(mul_nrows); const uint3 mul_nchannels_packed = init_fastdiv_values(mul_nchannels); const uint3 mul_nsamples_packed = init_fastdiv_values(mul_nsamples); - const char * rms128_env = getenv("GGML_CUDA_GB10_RMS128"); + static const bool rms128_enabled = [] { + const char * env = getenv("GGML_CUDA_GB10_RMS128"); + return !env || std::atoi(env) != 0; + }(); const bool use_rms128 = ggml_cuda_info().devices[ggml_cuda_get_device()].cc == GGML_CUDA_CC_DGX_SPARK && - ncols <= 128 && (!rms128_env || std::atoi(rms128_env) != 0); + ncols <= 128 && rms128_enabled; if (use_rms128) { const dim3 block_dims(128, 1, 1); const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, 32 * sizeof(float), stream}; diff --git a/ggml/src/ggml-cuda/quantize.cu b/ggml/src/ggml-cuda/quantize.cu index 52a025373aae..fbbc6314ab33 100644 --- a/ggml/src/ggml-cuda/quantize.cu +++ b/ggml/src/ggml-cuda/quantize.cu @@ -495,7 +495,8 @@ static __global__ void quantize_mmq_q8_1( // Load 4 floats per thread and calculate max. abs. value between them: float4 xi = i0 < ne00 ? x4[(base_idx + i00)/4] : make_float4(0.0f, 0.0f, 0.0f, 0.0f); if constexpr (swiglu) { - const float4 gi = ((const float4 *) gate)[(base_idx + i00)/4]; + const float4 gi = i0 < ne00 ? ((const float4 *) gate)[(base_idx + i00)/4] : + make_float4(0.0f, 0.0f, 0.0f, 0.0f); xi.x *= ggml_cuda_op_silu_single(gi.x); xi.y *= ggml_cuda_op_silu_single(gi.y); xi.z *= ggml_cuda_op_silu_single(gi.z); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 3714b727775c..cca08f5dc5fd 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10182,7 +10182,6 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 100, 1)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 200, 1)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 127, 2)); - test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 128, 64, 1)); // S=128 long-token CUDA path test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 64, 1, 1, false, true)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 33, 1, 1, false, true)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 100, 1, 1, false, true)); From 601219072dfd96aa97693949b825005d3d425c03 Mon Sep 17 00:00:00 2001 From: Pooya Khosravi <16091823+pooyakhosravi@users.noreply.github.com> Date: Thu, 3 Sep 2026 13:06:38 -0700 Subject: [PATCH 3/4] cuda: clarify async MMQ Y-tile addressing --- ggml/src/ggml-cuda/mmq.cuh | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh index 339ec9bc3a51..3e8d0dc4b713 100644 --- a/ggml/src/ggml-cuda/mmq.cuh +++ b/ggml/src/ggml-cuda/mmq.cuh @@ -969,8 +969,9 @@ static __device__ __forceinline__ void mul_mat_q_process_tile( vec_dot(tile_x, tile_y, sum, 0); cp_async_wait_all(); __syncthreads(); - // by1 is the second half of the same K tile; using a separate shared-memory - // buffer changes only its address, while MMQ_TILE_NE_K preserves its logical offset. + // vec_dot applies k00 only to the X-tile indices; Y is indexed locally from + // the supplied base. Both paths stage by1 at local Y offset 0, so passing + // tile_y_next changes only its storage address while k00 selects X's second half. vec_dot(tile_x, tile_y_next, sum, MMQ_TILE_NE_K); __syncthreads(); continue; From 037953ae6c8dd5e0a7d1d681a114e3d0edd05ebf Mon Sep 17 00:00:00 2001 From: Pooya Khosravi <16091823+pooyakhosravi@users.noreply.github.com> Date: Thu, 3 Sep 2026 14:38:54 -0700 Subject: [PATCH 4/4] cuda: prefetch Q1 MMQ tiles on GB10 --- ggml/src/ggml-cuda/mmq.cuh | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh index 3e8d0dc4b713..209248752f49 100644 --- a/ggml/src/ggml-cuda/mmq.cuh +++ b/ggml/src/ggml-cuda/mmq.cuh @@ -966,6 +966,25 @@ static __device__ __forceinline__ void mul_mat_q_process_tile( ggml_cuda_cvta_generic_to_shared(tile_y_next_bytes + byte0), by1 + byte0); } + if constexpr (type == GGML_TYPE_Q1_0 && J == 128) { + const int kb0_next = kb0 + blocks_per_iter; + if (kb0_next < kb0_stop) { + constexpr int hint_stride = 32; + constexpr int hints_per_row = + (blocks_per_iter*sizeof(block_q1_0) + hint_stride - 1) / hint_stride; +#pragma unroll + for (int linear = tid; linear < I*hints_per_row; linear += nwarps*warp_size) { + const int row = min(linear / hints_per_row, tile_x_max_i); + const int hint = linear % hints_per_row; + const char * next = reinterpret_cast(reinterpret_cast(x) + + offset_x + row*stride_row_x + kb0_next) + hint*hint_stride; +#if defined(__CUDA_ARCH__) + asm volatile("prefetch.global.L2 [%0];" :: "l"(__cvta_generic_to_global(next))); +#endif + } + } + } + vec_dot(tile_x, tile_y, sum, 0); cp_async_wait_all(); __syncthreads();