use iter_k as 512, cleanup

Author: Aman Gupta

ggml/src/ggml-cuda/common.cuh

@@ -719,11 +719,6 @@ __device__ __forceinline__ uint8_t ggml_cuda_float_to_fp4_e2m1(float x, float e)
     // Positive LUT
     static constexpr float pos_lut[8] = { 0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f };
 
-    // Saturate to max representable magnitude
-    if (ax > pos_lut[7]) {
-        ax = pos_lut[7];
-    }
-
     int   best_i   = 0;
     float best_err = fabsf(ax - pos_lut[0]);
 

ggml/src/ggml-cuda/mmq.cu

@@ -140,7 +140,7 @@ void ggml_cuda_mul_mat_q(
         // Stride depends on quantization format
         const int64_t s12 = use_native_mxfp4 ?
                                 ne11 * ne10_padded * sizeof(block_fp4_mmq) /
-                                    (4 * QK_MXFP4 * sizeof(int))  // block_fp4_mmq holds 128 values
+                                    (8 * QK_MXFP4 * sizeof(int))  // block_fp4_mmq holds 256 values (8 blocks of 32)
                                 :
                                 ne11 * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int));
         const int64_t s13 = ne12*s12;
@@ -200,9 +200,8 @@ void ggml_cuda_mul_mat_q(
         CUDA_CHECK(cudaGetLastError());
     }
 
-    const int64_t s12 = use_native_mxfp4 ?
-                            ne11 * ne10_padded * sizeof(block_fp4_mmq) / (4 * QK_MXFP4 * sizeof(int)) :
-                            ne11 * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int));
+    const int64_t s12 = use_native_mxfp4 ? ne11 * ne10_padded * sizeof(block_fp4_mmq) / (8 * QK_MXFP4 * sizeof(int)) :
+                                           ne11 * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int));
     const int64_t s13 = ne12*s12;
 
     // Note that ne02 is used instead of ne12 because the number of y channels determines the z dimension of the CUDA grid.

ggml/src/ggml-cuda/mmq.cuh

@@ -11,6 +11,7 @@ using namespace ggml_cuda_mma;
 
 #define MMQ_DP4A_MAX_BATCH_SIZE 64 // Max. batch size to use for dp4a MMQ kernels when FP16 tensor cores are available.
 #define MMQ_ITER_K 256
+#define MMQ_ITER_K_MXFP4_FP4    512
 #define MMQ_NWARPS 8
 
 typedef void (*load_tiles_mmq_t)(const char * __restrict__ x, int * x_tile, const int kbx0, const int i_max, const int stride);
@@ -46,13 +47,13 @@ struct block_q8_1_mmq {
 };
 
 struct block_fp4_mmq {
-    uint32_t d4[2];       // 1 8 bit (e8m0) scale per 32 values, packed LSB as d0-d1 in d4[0] and d4[1]
-    int8_t   qs[2 * 32];  // 128 values to 4 bit each (4 blocks)
+    uint32_t d4[4];       // 8 E8M0 scales (1 per 32 values), 2 packed per uint32: d4[0]={s0,s1}, d4[1]={s2,s3}, etc.
+    int8_t   qs[4 * 32];  // 256 FP4 values packed as 4-bit pairs (2 per byte), 8 blocks of 32 values
 };
 
 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) == 4*sizeof(block_q8_1),      "Unexpected block_q8_1_mmq size");
-static_assert(sizeof(block_fp4_mmq)  == 72, "Unexpected block_fp4_mmq size");
+static_assert(sizeof(block_fp4_mmq)  == sizeof(block_q8_1_mmq),    "Unexpected block_fp4_mmq size");
 
 static mmq_q8_1_ds_layout mmq_get_q8_1_ds_layout(const ggml_type type_x) {
     switch (type_x) {
@@ -136,6 +137,14 @@ static int get_mmq_y_host(const int cc) {
         ((GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 128 : 64);
 }
 
+static constexpr __device__ int get_iter_k([[maybe_unused]] const ggml_type type) {
+#if defined(BLACKWELL_MMA_AVAILABLE)
+    return type == GGML_TYPE_MXFP4 ? MMQ_ITER_K_MXFP4_FP4 : MMQ_ITER_K;
+#else
+    return MMQ_ITER_K;
+#endif // defined(BLACKWELL_MMA_AVAILABLE)
+}
+
 static constexpr __device__ int get_mmq_y_device() {
 #if defined(GGML_USE_HIP)
 #if defined(RDNA1)
@@ -198,7 +207,7 @@ static constexpr __host__ __device__ tile_x_sizes mmq_get_dp4a_tile_x_sizes(ggml
 }
 
 #define MMQ_MMA_TILE_X_K_Q8_0 (2*MMQ_TILE_NE_K + 2*MMQ_TILE_NE_K/QI8_0                   + 4)
-#define MMQ_MMA_TILE_X_K_FP4  (MMQ_TILE_NE_K + MMQ_TILE_NE_K / QI8_0)
+#define MMQ_MMA_TILE_X_K_FP4  (2*MMQ_TILE_NE_K + 2*MMQ_TILE_NE_K/QI8_0                   + 4)
 #define MMQ_MMA_TILE_X_K_Q8_1 (2*MMQ_TILE_NE_K + 2*MMQ_TILE_NE_K/QI8_0                   + 4)
 #define MMQ_MMA_TILE_X_K_Q2_K (2*MMQ_TILE_NE_K + MMQ_TILE_NE_K                           + 4)
 #define MMQ_MMA_TILE_X_K_Q3_K (2*MMQ_TILE_NE_K + MMQ_TILE_NE_K/2                         + 4)
@@ -209,7 +218,7 @@ static_assert(MMQ_MMA_TILE_X_K_Q8_1 % 8 == 4, "Wrong padding.");
 static_assert(MMQ_MMA_TILE_X_K_Q2_K % 8 == 4, "Wrong padding.");
 static_assert(MMQ_MMA_TILE_X_K_Q3_K % 8 == 4, "Wrong padding.");
 static_assert(MMQ_MMA_TILE_X_K_Q6_K % 8 == 4, "Wrong padding.");
-static_assert(MMQ_MMA_TILE_X_K_FP4 % 8 == 4, "Wrong padding.");
+static_assert(MMQ_MMA_TILE_X_K_FP4  % 8 == 4, "Wrong padding.");
 
 static constexpr __host__ __device__ int mmq_get_mma_tile_x_k(ggml_type type) {
     switch (type) {
@@ -218,11 +227,8 @@ static constexpr __host__ __device__ int mmq_get_mma_tile_x_k(ggml_type type) {
         case GGML_TYPE_Q5_0:    return MMQ_MMA_TILE_X_K_Q8_0;
         case GGML_TYPE_Q5_1:    return MMQ_MMA_TILE_X_K_Q8_1;
         case GGML_TYPE_Q8_0:    return MMQ_MMA_TILE_X_K_Q8_0;
-#ifdef BLACKWELL_MMA_AVAILABLE
-        case GGML_TYPE_MXFP4:   return MMQ_MMA_TILE_X_K_FP4;
-#else
+        // tile sizes are the same for Q8_1 and FP4 for blackwell
         case GGML_TYPE_MXFP4:   return MMQ_MMA_TILE_X_K_Q8_1;
-#endif
         case GGML_TYPE_Q2_K:    return MMQ_MMA_TILE_X_K_Q2_K;
         case GGML_TYPE_Q3_K:    return MMQ_MMA_TILE_X_K_Q3_K;
         case GGML_TYPE_Q4_K:    return MMQ_MMA_TILE_X_K_Q8_1;
@@ -242,7 +248,8 @@ static constexpr __host__ __device__ int mmq_get_mma_tile_x_k(ggml_type type) {
 
 // block_q8_1_mmq has (128 8-bit ints == 32 32-bit ints + 4 32-bit scales)
 #define MMQ_TILE_Y_K (MMQ_TILE_NE_K + MMQ_TILE_NE_K/QI8_1)
-#define MMQ_TILE_Y_FP4_K MMQ_TILE_Y_K / 2
+//#define MMQ_TILE_Y_FP4_K MMQ_TILE_Y_K / 2
+#define MMQ_TILE_Y_FP4_K MMQ_TILE_Y_K
 
 static int mmq_get_granularity_host(const int mmq_x, const int cc) {
     if (amd_mfma_available(cc) || amd_wmma_available(cc)) {
@@ -785,15 +792,14 @@ static __device__ __forceinline__ void load_tiles_mxfp4_fp4(const char * __restr
     constexpr int warp_size = ggml_cuda_get_physical_warp_size();
 
     int *      x_qs = (int *) x_tile;
-    uint32_t * x_sc = (uint32_t *) (x_qs + MMQ_TILE_NE_K);
+    uint32_t * x_sc = (uint32_t *) (x_qs + 2 * MMQ_TILE_NE_K);
 
     const int txi = threadIdx.x;
 
-    // Use all 32 threads: 8 threads per row, process 4 rows per warp per iteration
-    constexpr int threads_per_row = 8;  // 8 blocks per row
-    constexpr int rows_per_warp = warp_size / threads_per_row;  // 4 rows per warp
-    const int kbx = txi % threads_per_row;      // block id 0-7
-    const int row_in_warp = txi / threads_per_row;  // which of the 4 rows this thread handles
+    constexpr int threads_per_row = 16;  // 16 blocks per row for 512 values (ITER_K=512)
+    constexpr int rows_per_warp   = warp_size / threads_per_row;
+    const int     kbx             = txi % threads_per_row;
+    const int     row_in_warp     = txi / threads_per_row;
 
 #pragma unroll
     for (int i0 = 0; i0 < mmq_y; i0 += rows_per_warp * nwarps) {
@@ -805,6 +811,7 @@ static __device__ __forceinline__ void load_tiles_mxfp4_fp4(const char * __restr
 
         const block_mxfp4 * bxi = (const block_mxfp4 *) x + kbx0 + i * stride + kbx;
 
+        // quantize_mxfp4_mmq permutes nibbles to match the quantized format
         const int k0 = kbx * 4;
         memcpy(x_qs + i * MMQ_MMA_TILE_X_K_FP4 + k0, bxi->qs, 16);
 
@@ -1003,12 +1010,12 @@ static __device__ __forceinline__ void vec_dot_mxfp4_mxfp4_mma(const int * __res
 
     // Match layout from load_tiles_mxfp4_fp4
     const int *      x_qs = (const int *) x;
-    const uint32_t * x_sc = (const uint32_t *) (x_qs + MMQ_TILE_NE_K);  // E8M0 scales at same offset as load
-    const int *      y_qs = (const int *) y + 2;
-    const uint32_t * y_sc = (const uint32_t *) y;                       // E8M0 scales for Y
+    const uint32_t * x_sc = (const uint32_t *) (x_qs + 2 * MMQ_TILE_NE_K);
+    const int *      y_qs = (const int *) y + 4;
+    const uint32_t * y_sc = (const uint32_t *) y;
 
-    tile_A   A[ntx][MMQ_TILE_NE_K / (2 * QI8_0)];       // 2 x 4 A tiles. Per warp there will be 1 scale per tile
-    uint32_t scaleA[ntx][MMQ_TILE_NE_K / (2 * QI8_0)];  // per tile you would only have 1 scale per thread
+    tile_A   A[ntx][MMQ_TILE_NE_K / QI8_0];
+    uint32_t scaleA[ntx][MMQ_TILE_NE_K / QI8_0];
 
     // Block scale
     // Each thread has to point to a 4 byte scale value
@@ -1019,8 +1026,8 @@ static __device__ __forceinline__ void vec_dot_mxfp4_mxfp4_mma(const int * __res
 #pragma unroll
     for (int n = 0; n < ntx; ++n) {
 #pragma unroll
-        for (int k01 = 0; k01 < MMQ_TILE_NE_K / 2; k01 += QI8_0) {
-            const int k0 = k00 / 2 + k01;
+        for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) {
+            const int k0 = k00 + k01;
 
             load_ldmatrix(A[n][k01 / QI8_0], x_qs + (i0 + n * tile_A::I) * MMQ_MMA_TILE_X_K_FP4 + k0,
                           MMQ_MMA_TILE_X_K_FP4);
@@ -1034,7 +1041,7 @@ static __device__ __forceinline__ void vec_dot_mxfp4_mxfp4_mma(const int * __res
 #pragma unroll
     for (int j0 = 0; j0 < mmq_x; j0 += ntx * tile_C::J) {
 #pragma unroll
-        for (int k01 = 0; k01 < MMQ_TILE_NE_K / 2; k01 += QI8_0) {
+        for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) {
             tile_B   B;
             uint32_t scaleB;  // 2xN scales
 
@@ -3367,34 +3374,24 @@ static __device__ __forceinline__ void mul_mat_q_process_tile(
     constexpr mmq_write_back_t write_back = mmq_write_back_dp4a<mmq_x, mmq_y, need_check>;
 #endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
 
-#if defined(BLACKWELL_MMA_AVAILABLE)
-    constexpr bool use_native_mxfp4 = (type == GGML_TYPE_MXFP4);
-#else
-    constexpr bool use_native_mxfp4 = false;
-#endif // defined(BLACKWELL_MMA_AVAILBLE)
-
-
-    constexpr int blocks_per_iter = MMQ_ITER_K / qk;
+    constexpr int ITER_K          = get_iter_k(type);
+    constexpr int blocks_per_iter = ITER_K / qk;
 
     float sum[mmq_x*mmq_y / (nwarps*warp_size)] = {0.0f};
 
-    constexpr size_t sz       = use_native_mxfp4 ? sizeof(block_fp4_mmq) : sizeof(block_q8_1_mmq);
-    constexpr size_t y_stride = use_native_mxfp4 ? MMQ_TILE_Y_FP4_K : MMQ_TILE_Y_K;
+    constexpr int sz = sizeof(block_q8_1_mmq) / sizeof(int);
 
-    constexpr int y_block_stride = use_native_mxfp4 ? (sz / sizeof(int))  // 18 ints per block_fp4_mmq (covers 128 values = 4 qk-blocks)
-                                  :
-                                  (qk * sz / (4 * QK8_1 * sizeof(int)));  // original formula for Q8_1
+    // blocks_per_mmq: number of qk-blocks per Y-block structure
+    // MXFP4: block_fp4_mmq holds 8 qk-blocks (256 values)
+    // Others: block_q8_1_mmq holds 4 qk-blocks (128 values)
+    constexpr int blocks_per_mmq = (type == GGML_TYPE_MXFP4) ? 8 : 4;
 
     for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) {
         load_tiles(x, tile_x, offset_x + kb0, tile_x_max_i, stride_row_x);
         {
-            const int * by0 =
-                use_native_mxfp4 ?
-                    y + ncols_y * ((kb0 / 4) * y_block_stride)  // kb0/4 for MXFP4 since 4 qk-blocks per block_fp4_mmq
-                    :
-                    y + ncols_y * (kb0 * y_block_stride);       // original for Q8_1
+            const int * by0 = y + ncols_y * (kb0 / blocks_per_mmq) * sz;
 #pragma unroll
-            for (int l0 = 0; l0 < mmq_x * y_stride; l0 += nwarps * warp_size) {
+            for (int l0 = 0; l0 < mmq_x * MMQ_TILE_Y_K; l0 += nwarps * warp_size) {
                 int l = l0 + threadIdx.y*warp_size + threadIdx.x;
 
                 tile_y[l] = by0[l];
@@ -3408,14 +3405,9 @@ static __device__ __forceinline__ void mul_mat_q_process_tile(
         __syncthreads();
 
         {
-            const int * by0 =
-                use_native_mxfp4 ?
-                    y + ncols_y * ((kb0 / 4) * y_block_stride + y_block_stride)  // advance by one block_fp4_mmq
-                    :
-                    y + ncols_y * (kb0 * y_block_stride +
-                                   (int) (sz / sizeof(int)));  // original for Q8_1 (advance by one block)
-#pragma unroll
-            for (int l0 = 0; l0 < mmq_x * y_stride; l0 += nwarps * warp_size) {
+            const int * by0 = y + ncols_y * ((kb0 / blocks_per_mmq) * sz + sz);
+#pragma unroll
+            for (int l0 = 0; l0 < mmq_x * MMQ_TILE_Y_K; l0 += nwarps * warp_size) {
                 int l = l0 + threadIdx.y*warp_size + threadIdx.x;
 
                 tile_y[l] = by0[l];
@@ -3547,8 +3539,10 @@ static __global__ void mul_mat_q(
     }
 #endif // (defined(GGML_USE_HIP) && !defined(CDNA3)) || __CUDA_ARCH__ < GGML_CUDA_CC_VOLTA
 
+    constexpr int ITER_K = get_iter_k(type);
+
     const     int64_t blocks_per_ne00 = ncols_x / qk;
-    constexpr int     blocks_per_iter = MMQ_ITER_K / qk;
+    constexpr int     blocks_per_iter = ITER_K / qk;
 
     // kbc == k block continuous, current index in continuous ijk space.
     int64_t kbc      = (int64_t) blockIdx.x     *nsamples_y*nchannels_y*ntx*nty*blocks_per_ne00 / gridDim.x;
@@ -3609,8 +3603,7 @@ static __global__ void mul_mat_q(
             __syncthreads();
         }
 
-        constexpr size_t sz = type == GGML_TYPE_MXFP4 ? sizeof(block_fp4_mmq) : sizeof(block_q8_1_mmq);
-        offset_y += (col_low + jt * mmq_x) * (sz / sizeof(int));
+        offset_y += (col_low + jt * mmq_x) * (sizeof(block_q8_1_mmq) / sizeof(int));
         offset_dst += it*mmq_y;
 
         const int tile_x_max_i = nrows_x  - it*mmq_y - 1;
@@ -3677,8 +3670,7 @@ static __global__ void mul_mat_q(
         __syncthreads();
     }
 
-    constexpr size_t sz = type == GGML_TYPE_MXFP4 ? sizeof(block_fp4_mmq) : sizeof(block_q8_1_mmq);
-    offset_y += (col_low + jt * mmq_x) * (sz / sizeof(int));
+    offset_y += (col_low + jt * mmq_x) * (sizeof(block_q8_1_mmq) / sizeof(int));
     offset_dst += it*mmq_y;
 
     const int tile_x_max_i = nrows_x  - it*mmq_y - 1;
@@ -3701,7 +3693,9 @@ static __global__ void mul_mat_q_stream_k_fixup(
         const int ncols_max) {
     constexpr int     mmq_y           = get_mmq_y_device();
     constexpr int     qk              = ggml_cuda_type_traits<type>::qk;
-    constexpr int     blocks_per_iter = MMQ_ITER_K / qk;
+    constexpr int     ITER_K          = get_iter_k(type);
+
+    constexpr int     blocks_per_iter = ITER_K / qk;
     const     int64_t blocks_per_ne00 = ncols_x / qk;
 
     constexpr int nwarps = mmq_get_nwarps_device();