From 51319a2cd902ed0a67c96ac29db7219c48d17e35 Mon Sep 17 00:00:00 2001 From: ynankani Date: Tue, 18 Aug 2026 04:15:53 +0000 Subject: [PATCH] CUDA: MMVQ nwarps=8 for bs=1 for dense models on DGX Spark (llama/26843) * CUDA: MMVQ nwarps=8 for bs=1 for dense models on DGX Spark Signed-off-by: ynankani * skip moe experts and allow others based on k geometry (allow only small idle tail) Signed-off-by: ynankani * rename MMVQ DGX Spark params to GB10 and fix MSVC constexpr lambda capture Signed-off-by: ynankani --------- Signed-off-by: ynankani --- ggml/src/ggml-cuda/mmvq.cu | 122 +++++++++++++++++++++++++++---------- 1 file changed, 91 insertions(+), 31 deletions(-) diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index 0589e65bd..c99923804 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -4,6 +4,7 @@ #include "vecdotq.cuh" #include +#include typedef float (*vec_dot_q_cuda_t)(const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs); @@ -69,7 +70,8 @@ enum mmvq_parameter_table_id { MMVQ_PARAMETERS_GCN, MMVQ_PARAMETERS_RDNA2, MMVQ_PARAMETERS_RDNA3_0, - MMVQ_PARAMETERS_RDNA4 + MMVQ_PARAMETERS_RDNA4, + MMVQ_PARAMETERS_GB10 }; static constexpr __device__ mmvq_parameter_table_id get_device_table_id() { @@ -83,6 +85,8 @@ static constexpr __device__ mmvq_parameter_table_id get_device_table_id() { return MMVQ_PARAMETERS_GCN; #elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= GGML_CUDA_CC_TURING && __CUDA_ARCH__ < GGML_CUDA_CC_AMPERE return MMVQ_PARAMETERS_TURING; +#elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK + return MMVQ_PARAMETERS_GB10; #else return MMVQ_PARAMETERS_GENERIC; #endif @@ -104,6 +108,9 @@ static __host__ mmvq_parameter_table_id get_device_table_id(int cc) { if (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_TURING && ggml_cuda_highest_compiled_arch(cc) < GGML_CUDA_CC_AMPERE) { return MMVQ_PARAMETERS_TURING; } + if (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) == GGML_CUDA_CC_DGX_SPARK) { + return MMVQ_PARAMETERS_GB10; + } return MMVQ_PARAMETERS_GENERIC; } @@ -351,7 +358,7 @@ static constexpr __device__ int get_mmvq_mmid_max_batch_for_device() { #endif } -static constexpr __host__ __device__ int calc_nwarps(ggml_type type, int ncols_dst, mmvq_parameter_table_id table_id) { +static constexpr __host__ __device__ int calc_nwarps(ggml_type type, int ncols_dst, mmvq_parameter_table_id table_id, bool small_k = false, bool halve_iters = false) { if (table_id == MMVQ_PARAMETERS_GENERIC) { switch (ncols_dst) { case 1: @@ -454,11 +461,32 @@ static constexpr __host__ __device__ int calc_nwarps(ggml_type type, int ncols_d return 1; } } + 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) + if (ncols_dst == 1 && !small_k && halve_iters) { + switch (type) { + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ4_NL: + return 2 * generic; + default: + break; + } + } + return generic; + } return 1; } static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int table_id, bool small_k = false, int nwarps = 1) { - if (table_id == MMVQ_PARAMETERS_GENERIC || table_id == MMVQ_PARAMETERS_GCN || table_id == MMVQ_PARAMETERS_TURING) { + if (table_id == MMVQ_PARAMETERS_GENERIC || table_id == MMVQ_PARAMETERS_GCN || table_id == MMVQ_PARAMETERS_TURING || table_id == MMVQ_PARAMETERS_GB10) { switch (ncols_dst) { case 1: return small_k ? nwarps : 1; @@ -477,8 +505,8 @@ static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int return 1; } -template -__launch_bounds__(calc_nwarps(type, ncols_dst, get_device_table_id())*ggml_cuda_get_physical_warp_size(), 1) +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, const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t stride_row_x, const uint32_t stride_col_y, @@ -495,7 +523,7 @@ static __global__ void mul_mat_vec_q( constexpr int qi = ggml_cuda_type_traits::qi; constexpr int vdr = get_vdr_mmvq(type); constexpr mmvq_parameter_table_id table_id = get_device_table_id(); - constexpr int nwarps = calc_nwarps(type, ncols_dst, table_id); + constexpr int nwarps = calc_nwarps(type, ncols_dst, table_id, small_k, halve_iters); constexpr int rows_per_cuda_block = calc_rows_per_block(ncols_dst, table_id, small_k, nwarps); constexpr int warp_size = ggml_cuda_get_physical_warp_size(); @@ -773,8 +801,8 @@ static __global__ void mul_mat_vec_q_moe( template static std::pair calc_launch_params( const int ncols_dst, const int nrows_x, const int nchannels_dst, const int nsamples_or_ntokens, - const int warp_size, const mmvq_parameter_table_id table_id, const bool small_k = false) { - const int nwarps = calc_nwarps(type, ncols_dst, table_id); + const int warp_size, const mmvq_parameter_table_id table_id, const bool small_k = false, const bool halve_iters = false) { + const int nwarps = calc_nwarps(type, ncols_dst, table_id, small_k, halve_iters); const int rpb = calc_rows_per_block(ncols_dst, table_id, small_k, nwarps); const int64_t nblocks = (nrows_x + rpb - 1) / rpb; const dim3 block_nums(nblocks, nchannels_dst, nsamples_or_ntokens); @@ -782,7 +810,7 @@ static std::pair calc_launch_params( return {block_nums, block_dims}; } -template +template static void mul_mat_vec_q_switch_fusion( const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst, const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t stride_row_x, const uint32_t stride_col_y, @@ -797,7 +825,7 @@ 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, + 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); @@ -808,7 +836,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); @@ -860,16 +888,18 @@ static void mul_mat_vec_q_switch_ncols_dst( const bool has_ids = ids != nullptr; + // How the K loop divides up at the baseline block width, both decisions below use these. + constexpr int qk = ggml_cuda_type_traits::qk; + constexpr int qi = ggml_cuda_type_traits::qi; + constexpr int vdr = get_vdr_mmvq(type); + const int blocks_per_row_x = ncols_x / qk; + const int blocks_per_iter_1warp = vdr * warp_size / qi; + const auto should_use_small_k = [&](int c_ncols_dst) { // When K is small, increase rows_per_block to match nwarps so each warp has more work to do // Trigger when the full thread block covers all K blocks in a single loop iteration and few threads remain idle. - constexpr int qk = ggml_cuda_type_traits::qk; - constexpr int qi = ggml_cuda_type_traits::qi; - constexpr int vdr = get_vdr_mmvq(type); - const int blocks_per_row_x = ncols_x / qk; - const int blocks_per_iter_1warp = vdr * warp_size / qi; - const int nwarps = calc_nwarps(type, c_ncols_dst, table_id); - bool use = nwarps > 1 && blocks_per_row_x < nwarps * blocks_per_iter_1warp; + const int nwarps = calc_nwarps(type, c_ncols_dst, table_id); + bool use = nwarps > 1 && blocks_per_row_x < nwarps * blocks_per_iter_1warp; constexpr std::array iq_slow_turing = { GGML_TYPE_IQ3_XXS, @@ -902,6 +932,28 @@ static void mul_mat_vec_q_switch_ncols_dst( return use; }; + // Whether doubling nwarps pays off on the ncols_dst == 1 path, where K sets the K loop trip count. + const auto should_halve_iters = [&] { + if (table_id != MMVQ_PARAMETERS_GB10) { + return false; + } + + // Expert rows are gathered per token, so a wider block adds reduction work without reuse. + if (has_ids) { + return false; + } + + const int blocks_per_iter = calc_nwarps(type, 1, table_id) * blocks_per_iter_1warp; + const int iters = (blocks_per_row_x + blocks_per_iter - 1) / blocks_per_iter; + const int iters_wide = (blocks_per_row_x + blocks_per_iter * 2 - 1) / (blocks_per_iter * 2); + + // An odd trip count leaves half the wider block idle for its last iteration, that tail is + // only affordable once the loop is long enough to dilute it to an eighth of the work (observation). + const int idle = iters_wide * 2 - iters; + + return idle * 8 <= iters_wide * 2; + }; + if (has_ids && ncols_dst > 1) { // Multi-token MUL_MAT_ID path - dedicated MoE kernel mul_mat_vec_q_moe_launch( @@ -914,26 +966,34 @@ static void mul_mat_vec_q_switch_ncols_dst( switch (ncols_dst) { case 1: { - constexpr int c_ncols_dst = 1; + // static, else MSVC lambda capture breaks the constexpr uses below + static constexpr int c_ncols_dst = 1; - bool use_small_k = should_use_small_k(c_ncols_dst); + // Tag types keep the flags compile-time, so __launch_bounds__ matches what is launched. + const auto launch = [&](auto small_k_tag, auto halve_iters_tag) { + constexpr bool c_small_k = decltype(small_k_tag)::value; + // Types the table does not promote would compile a second, identical kernel. + constexpr bool c_promoted = + calc_nwarps(type, c_ncols_dst, MMVQ_PARAMETERS_GB10, false, true) != + calc_nwarps(type, c_ncols_dst, MMVQ_PARAMETERS_GB10, false, false); - if (use_small_k) { - std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst, - nsamples_dst, warp_size, table_id, true); - mul_mat_vec_q_switch_fusion( + 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); + 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, stride_sample_x, stride_sample_y, stride_sample_dst, dims.first, dims.second, 0, ids_stride, stream); + }; + + if (should_use_small_k(c_ncols_dst)) { + launch(std::true_type{}, std::false_type{}); + } else if (should_halve_iters()) { + launch(std::false_type{}, std::true_type{}); } else { - std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst, - nsamples_dst, warp_size, table_id); - 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, - stride_sample_x, stride_sample_y, stride_sample_dst, dims.first, dims.second, 0, ids_stride, - stream); + launch(std::false_type{}, std::false_type{}); } } break; case 2: {