diff --git a/ggml/src/ggml-cuda/mmq.cu b/ggml/src/ggml-cuda/mmq.cu index bf9f5d5264..34e5cbf8cb 100644 --- a/ggml/src/ggml-cuda/mmq.cu +++ b/ggml/src/ggml-cuda/mmq.cu @@ -122,11 +122,12 @@ void ggml_cuda_mul_mat_q( const bool fallback = ne01 % 128 != 0; - // TODO: tighter pool buffer size vs q8 path const bool use_native_fp4 = blackwell_mma_available(cc) && (src0->type == GGML_TYPE_MXFP4 || src0->type == GGML_TYPE_NVFP4); + const size_t y_block_size = use_native_fp4 ? sizeof(block_fp4_mmq) : sizeof(block_q8_1_mmq); + 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 * sizeof(block_q8_1)/QK8_1 + + 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); @@ -148,7 +149,7 @@ void ggml_cuda_mul_mat_q( // Stride depends on quantization format const int64_t s12 = use_native_fp4 ? - ne11 * ne10_padded * sizeof(block_fp4_mmq) / (QK_K * sizeof(int)) : // block_fp4_mmq holds 256 values + ne11 * ne10_padded * sizeof(block_fp4_mmq) / (QK_FP4_MMQ * sizeof(int)) : ne11 * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int)); const int64_t s13 = ne12*s12; @@ -184,7 +185,7 @@ void ggml_cuda_mul_mat_q( CUDA_CHECK(cudaGetLastError()); } - const size_t nbytes_src1_q8_1 = ne12*n_expert_used*ne10_padded * sizeof(block_q8_1)/QK8_1 + + const size_t nbytes_src1_q8_1 = ne12*n_expert_used*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); @@ -207,8 +208,8 @@ void ggml_cuda_mul_mat_q( CUDA_CHECK(cudaGetLastError()); } - static_assert(QK_K == 8 * QK_MXFP4, "QK_K needs to be 8 * QK_MXFP4"); - const int64_t s12 = use_native_fp4 ? ne11 * ne10_padded * sizeof(block_fp4_mmq) / (QK_K * sizeof(int)) : + static_assert(QK_FP4_MMQ == 8 * QK_MXFP4, "QK_FP4_MMQ needs to be 8 * QK_MXFP4"); + const int64_t s12 = use_native_fp4 ? ne11 * ne10_padded * sizeof(block_fp4_mmq) / (QK_FP4_MMQ * sizeof(int)) : ne11 * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int)); const int64_t s13 = ne12*s12; diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh index 607e433bf1..ad3d6a5a68 100644 --- a/ggml/src/ggml-cuda/mmq.cuh +++ b/ggml/src/ggml-cuda/mmq.cuh @@ -21,6 +21,9 @@ enum mmq_q8_1_ds_layout { MMQ_Q8_1_DS_LAYOUT_D2S6, }; +static constexpr int QK8_1_MMQ = 4*QK8_1; +static constexpr int QK_FP4_MMQ = 2*QK8_1_MMQ; + struct block_q8_1_mmq { // The y float data is converted to a data layout that can simply be copied to shared memory as a contiguous block. // The y float data is first grouped as blocks of 128 values. @@ -39,7 +42,7 @@ struct block_q8_1_mmq { half d2s6[8]; // 1 16 bit scale per 64 values + 1 16 bit partial sum per 16 values for the first 96 values, // stored as d0,d1,s1,s2,s3,s4,s5 }; - int8_t qs[4*QK8_1]; // 128 values quantized to 8 bit each + int8_t qs[QK8_1_MMQ]; }; // this struct is used for fp4 data types (currently only used for Blackwell) @@ -47,10 +50,10 @@ struct block_q8_1_mmq { // nvfp4 has block size 16, each int32 of d4 contains 4 ue4m3 scales struct block_fp4_mmq { uint32_t d4[4]; - int8_t qs[4 * 32]; // 256 FP4 values packed as 4-bit pairs (2 per byte) + int8_t qs[QK_FP4_MMQ / 2]; }; -static_assert(sizeof(block_q8_1_mmq) == 4*QK8_1 + 4*sizeof(half2), "Unexpected block_q8_1_mmq size"); +static_assert(sizeof(block_q8_1_mmq) == QK8_1_MMQ + 4*sizeof(half2), "Unexpected block_q8_1_mmq size"); static_assert(sizeof(block_q8_1_mmq) == 4*sizeof(block_q8_1), "Unexpected block_q8_1_mmq size"); static_assert(sizeof(block_fp4_mmq) == sizeof(block_q8_1_mmq), "Unexpected block_fp4_mmq size"); @@ -833,9 +836,9 @@ static __device__ __forceinline__ void mul_mat_q_process_tile( #if defined(BLACKWELL_MMA_AVAILABLE) // FP4 tile stores 8 blocks - constexpr int ne_block = (type == GGML_TYPE_MXFP4 || type == GGML_TYPE_NVFP4) ? QK_K : 4 * QK8_1; + constexpr int ne_block = (type == GGML_TYPE_MXFP4 || type == GGML_TYPE_NVFP4) ? QK_FP4_MMQ : QK8_1_MMQ; #else - constexpr int ne_block = 4 * QK8_1; + constexpr int ne_block = QK8_1_MMQ; #endif // defined(BLACKWELL_MMA_AVAILABLE) constexpr int ITER_K = ggml_cuda_mmq_get_K_vram(type, J, fallback); diff --git a/ggml/src/ggml-cuda/quantize.cu b/ggml/src/ggml-cuda/quantize.cu index 39a500a170..a7d450737f 100644 --- a/ggml/src/ggml-cuda/quantize.cu +++ b/ggml/src/ggml-cuda/quantize.cu @@ -90,8 +90,8 @@ static __global__ void quantize_mmq_nvfp4( const int64_t i2 = blockIdx.z % ne2; const int64_t i3 = blockIdx.z / ne2; const int64_t i01 = ids ? ids[i1] : i1; - const int64_t k_block = i0_base / QK_K; - const int64_t blocks_per_col = (ne0 + QK_K - 1) / QK_K; + const int64_t k_block = i0_base / QK_FP4_MMQ; + const int64_t blocks_per_col = (ne0 + QK_FP4_MMQ - 1) / QK_FP4_MMQ; if (k_block >= blocks_per_col) { return; } @@ -100,7 +100,7 @@ static __global__ void quantize_mmq_nvfp4( block_fp4_mmq * y = (block_fp4_mmq *) vy; block_fp4_mmq * yb = y + ib; - const int sub = (i0_base % QK_K) / QK_NVFP4_SUB; + const int sub = (i0_base % QK_FP4_MMQ) / QK_NVFP4_SUB; float vals_raw[QK_NVFP4_SUB]; float amax_raw = 0.0f; @@ -207,7 +207,7 @@ static __global__ void quantize_mmq_mxfp4(const float * __restrict__ x, block_fp4_mmq * y = (block_fp4_mmq *) vy; - const int64_t block_fp4_mmq_size = 8 * QK_MXFP4; // 256 values + const int64_t block_fp4_mmq_size = QK_FP4_MMQ; const int64_t ib0 = blockIdx.z * ((int64_t) ne1 * (ne0 / block_fp4_mmq_size)); const int64_t ib = ib0 + (warp_start_offset / block_fp4_mmq_size) * ne1 + blockIdx.x; const int64_t quad_idx_in_block = (warp_start_offset % block_fp4_mmq_size) / vals_per_warp; @@ -303,8 +303,8 @@ static __global__ void quantize_mmq_q8_1( block_q8_1_mmq * y = (block_q8_1_mmq *) vy; const int64_t ib0 = blockIdx.z*((int64_t)gridDim.x*gridDim.y*blockDim.x/QK8_1); // first block of channel - const int64_t ib = ib0 + (i0 / (4*QK8_1))*ne1 + blockIdx.x; // block index in channel - const int64_t iqs = i0 % (4*QK8_1); // quant index in block + const int64_t ib = ib0 + (i0 / QK8_1_MMQ)*ne1 + blockIdx.x; // block index in channel + 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[(i03*s03 + i02*s02 + i01*s01 + i00)/4] : make_float4(0.0f, 0.0f, 0.0f, 0.0f); @@ -394,7 +394,7 @@ void quantize_mmq_q8_1_cuda( 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 % (4*QK8_1) == 0); + GGML_ASSERT(ne0 % QK8_1_MMQ == 0); // ne1 tends to assume the highest values, therefore use it as the "x" dimension of the CUDA grid: const int64_t block_num_y = (ne0 + 4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ);