diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 18d1be1c9..28fd850d8 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1869,8 +1869,7 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { if (cc <= GGML_CUDA_CC_PASCAL) { return false; } - //we only support fusion for ncols_dst = 1 - if (tensor->op == GGML_OP_MUL_MAT && dst->ne[1] != 1) { + if (tensor->op == GGML_OP_MUL_MAT && !ggml_cuda_mmvq_fusion_supported(src0->type, cc, dst->ne[1])) { return false; } diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index 8fb790d73..a93c22dcd 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -2,6 +2,7 @@ #include "quantize.cuh" #include "unary.cuh" #include "vecdotq.cuh" +#include "sm75-iq4-reuse.cuh" #include #include @@ -719,6 +720,30 @@ static __global__ void mul_mat_vec_q( // x block quant index when casting the quants to int const int kqs = vdr * (tid % (qi/vdr)); +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == GGML_CUDA_CC_TURING && GGML_SM75_IQ4_REUSE + if constexpr (type == GGML_TYPE_IQ4_XS && ncols_dst >= 2 && ncols_dst <= 4) { +#pragma unroll + for (int i = 0; i < rows_per_cuda_block; ++i) { + const int bx = kbx_offset + i * stride_row_x + kbx; + const sm75_iq4_decoded w = sm75_iq4_decode(vx, bx, kqs); +#pragma unroll + for (int j = 0; j < ncols_dst; ++j) { + tmp[j][i] += sm75_iq4_dot(w, &y[j * stride_col_y + kby + kqs / 4]); + } + if constexpr (has_fusion) { + if (use_gate) { + const sm75_iq4_decoded gate = sm75_iq4_decode(vgate, bx, kqs); +#pragma unroll + for (int j = 0; j < ncols_dst; ++j) { + tmp_gate[j][i] += sm75_iq4_dot(gate, &y[j * stride_col_y + kby + kqs / 4]); + } + } + } + } + continue; + } +#endif + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK // start the next iterations' weight loads early if constexpr (mmvq_should_prefetch(type)) { @@ -1020,7 +1045,7 @@ static void mul_mat_vec_q_switch_fusion( const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr || fusion.x_scale != nullptr || fusion.gate_scale != nullptr; - if constexpr (c_ncols_dst == 1) { + if constexpr (ggml_cuda_mmvq_fusion_supported(type, GGML_CUDA_CC_TURING, c_ncols_dst)) { 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, @@ -1031,7 +1056,7 @@ static void mul_mat_vec_q_switch_fusion( } } - GGML_ASSERT(!has_fusion && "fusion only supported for ncols_dst=1"); + GGML_ASSERT(!has_fusion && "unsupported MMVQ fusion shape"); 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, @@ -1450,7 +1475,7 @@ void ggml_cuda_mul_mat_vec_q( if (fusion) { const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; GGML_ASSERT( !ids || dst->ne[2] <= get_mmvq_mmid_max_batch(src0->type, cc)); - GGML_ASSERT( ids || dst->ne[1] == 1); + GGML_ASSERT( ids || ggml_cuda_mmvq_fusion_supported(src0->type, cc, dst->ne[1])); // Scale fusion is only allowed for NVFP4 currently as the cost of checking this at run-time in the prologue is // non-negligible for some models such as gpt-oss-20b GGML_ASSERT((fusion->x_scale == nullptr && fusion->gate_scale == nullptr) || src0->type == GGML_TYPE_NVFP4); diff --git a/ggml/src/ggml-cuda/mmvq.cuh b/ggml/src/ggml-cuda/mmvq.cuh index 5605bf7a4..ec6b840ee 100644 --- a/ggml/src/ggml-cuda/mmvq.cuh +++ b/ggml/src/ggml-cuda/mmvq.cuh @@ -1,9 +1,15 @@ +#pragma once + #include "common.cuh" #define MMVQ_MAX_BATCH_SIZE 8 // Max. batch size for which to use MMVQ kernels. bool ggml_cuda_should_use_mmvq(enum ggml_type type, int cc, int64_t ne11); +static constexpr bool ggml_cuda_mmvq_fusion_supported(ggml_type type, int cc, int64_t ncols) { + return ncols == 1 || (cc == GGML_CUDA_CC_TURING && type == GGML_TYPE_IQ4_XS && ncols >= 2 && ncols <= 4); +} + // Returns the maximum batch size for which MMVQ should be used for MUL_MAT_ID, // based on the quantization type and GPU architecture (compute capability). int get_mmvq_mmid_max_batch(ggml_type type, int cc); diff --git a/ggml/src/ggml-cuda/sm75-iq4-reuse.cuh b/ggml/src/ggml-cuda/sm75-iq4-reuse.cuh new file mode 100644 index 000000000..cfb682bfe --- /dev/null +++ b/ggml/src/ggml-cuda/sm75-iq4-reuse.cuh @@ -0,0 +1,51 @@ +#pragma once + +#include + +#ifndef GGML_SM75_IQ4_REUSE +#define GGML_SM75_IQ4_REUSE 1 +#endif + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == GGML_CUDA_CC_TURING && GGML_SM75_IQ4_REUSE + +static_assert(QI4_XS == 32 && VDR_IQ4_XS_Q8_1_MMVQ == 4, "IQ4_XS reuse expects four words per lane"); +static_assert(sizeof(block_iq4_xs) % 8 == 0, "IQ4_XS block stride must preserve 8-byte alignment"); +static_assert(offsetof(block_iq4_xs, qs) % 8 == 0, "IQ4_XS quants must have 8-byte alignment"); + +struct sm75_iq4_decoded { + int2 values[4]; + int scale; + float d; +}; + +static __device__ __forceinline__ sm75_iq4_decoded sm75_iq4_decode(const void * vx, int kbx, int iqs) { + const block_iq4_xs * b = (const block_iq4_xs *) vx + kbx; + const int2 * packed = (const int2 *) (b->qs + sizeof(int) * iqs); + const int2 p0 = packed[0]; + const int2 p1 = packed[1]; + const int words[4] = {p0.x, p0.y, p1.x, p1.y}; + const uint2 header = *(const uint2 *) b; + sm75_iq4_decoded result; +#pragma unroll + for (int j = 0; j < 4; ++j) { + result.values[j] = get_int_from_table_16(words[j], kvalues_iq4nl); + } + const int ls = int(((header.y >> iqs) & 15) | (((header.x >> (16 + iqs / 2)) & 3) << 4)); + result.scale = ls - 32; + result.d = __half2float(__ushort_as_half((uint16_t) header.x)); + return result; +} + +static __device__ __forceinline__ float sm75_iq4_dot(const sm75_iq4_decoded & w, const block_q8_1 * y) { + int sumi = 0; +#pragma unroll + for (int j = 0; j < 4; ++j) { + sumi = ggml_cuda_dp4a(w.values[j].x, get_int_b4(y->qs, j), sumi); + sumi = ggml_cuda_dp4a(w.values[j].y, get_int_b4(y->qs, j + 4), sumi); + } + sumi *= w.scale; + const float d = w.d * __low2float(y->ds); + return d * sumi; +} + +#endif