CUDA: tighter MMQ src1 buffer size for native fp4 (#25613)

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<char> 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<char> 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;
 

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);

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);