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 <ynankani@nvidia.com> * skip moe experts and allow others based on k geometry (allow only small idle tail) Signed-off-by: ynankani <ynankani@nvidia.com> * rename MMVQ DGX Spark params to GB10 and fix MSVC constexpr lambda capture Signed-off-by: ynankani <ynankani@nvidia.com> --------- Signed-off-by: ynankani <ynankani@nvidia.com>
This commit is contained in:
parent
71759b7c9f
commit
51319a2cd9
|
|
@ -4,6 +4,7 @@
|
|||
#include "vecdotq.cuh"
|
||||
|
||||
#include <cstdint>
|
||||
#include <type_traits>
|
||||
|
||||
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 <ggml_type type, int ncols_dst, bool has_fusion, bool small_k = false>
|
||||
__launch_bounds__(calc_nwarps(type, ncols_dst, get_device_table_id())*ggml_cuda_get_physical_warp_size(), 1)
|
||||
template <ggml_type type, int ncols_dst, bool has_fusion, bool small_k = false, bool halve_iters = false>
|
||||
__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<type>::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<ggml_type type>
|
||||
static std::pair<dim3, dim3> 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<dim3, dim3> calc_launch_params(
|
|||
return {block_nums, block_dims};
|
||||
}
|
||||
|
||||
template<ggml_type type, int c_ncols_dst, bool small_k = false>
|
||||
template<ggml_type type, int c_ncols_dst, bool small_k = false, bool halve_iters = false>
|
||||
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<type, c_ncols_dst, true, small_k>, launch_params,
|
||||
ggml_cuda_kernel_launch(mul_mat_vec_q<type, c_ncols_dst, true, small_k, halve_iters>, 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<type, c_ncols_dst, false, small_k>, launch_params,
|
||||
ggml_cuda_kernel_launch(mul_mat_vec_q<type, c_ncols_dst, false, small_k, halve_iters>, 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<type>::qk;
|
||||
constexpr int qi = ggml_cuda_type_traits<type>::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<type>::qk;
|
||||
constexpr int qi = ggml_cuda_type_traits<type>::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<ggml_type, 2> 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<type>(
|
||||
|
|
@ -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<dim3, dim3> dims = calc_launch_params<type>(c_ncols_dst, nrows_x, nchannels_dst,
|
||||
nsamples_dst, warp_size, table_id, true);
|
||||
mul_mat_vec_q_switch_fusion<type, c_ncols_dst, true>(
|
||||
constexpr bool c_halve_iters = decltype(halve_iters_tag)::value && c_promoted;
|
||||
|
||||
const std::pair<dim3, dim3> dims = calc_launch_params<type>(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<type, c_ncols_dst, c_small_k, c_halve_iters>(
|
||||
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<dim3, dim3> dims = calc_launch_params<type>(c_ncols_dst, nrows_x, nchannels_dst,
|
||||
nsamples_dst, warp_size, table_id);
|
||||
mul_mat_vec_q_switch_fusion<type, c_ncols_dst>(
|
||||
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: {
|
||||
|
|
|
|||
Loading…
Reference in New Issue