Repository: ggerganov/llama.cpp
File: ggml/src/ggml-cuda/mmq.cuh (CUDA)

Code before:
#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
}

template <int mmq_x, int mmq_y>
static __device__ __forceinline__ void vec_dot_mxfp4_mxfp4_mma(const int * __restrict__ x,
                                                               const int * __restrict__ y,
                                                               float * __restrict__ sum,
                                                               const int k00) {
    typedef tile<16, 8, int>   tile_A;
    typedef tile<8, 8, int>    tile_B;
    typedef tile<16, 8, float> tile_C;  // Output is float for native scaled MMA

    constexpr int granularity   = mmq_get_granularity_device(mmq_x);
    constexpr int rows_per_warp = 2 * granularity;
    constexpr int ntx           = rows_per_warp / tile_C::I;  // Number of x minitiles per warp.

    y += (threadIdx.y % ntx) * (tile_C::J * MMQ_TILE_Y_FP4_K);

    // Match layout from load_tiles_mxfp4_fp4
    const int *      x_qs = (const int *) x;
    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 * QI_MXFP4)];
    uint32_t scaleA[ntx][MMQ_TILE_NE_K / (2 * QI_MXFP4)];

    // Block scale
    // Each thread has to point to a 4 byte scale value
    // https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-block-scaling

    const int i0 = (threadIdx.y / ntx) * rows_per_warp;

#pragma unroll
    for (int n = 0; n < ntx; ++n) {
#pragma unroll
        for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 2 * QI_MXFP4) {
            const int k0 = k00 + k01;

            load_ldmatrix(A[n][k01 / (2 * QI_MXFP4)], x_qs + (i0 + n * tile_A::I) * MMQ_MMA_TILE_X_K_FP4 + k0,
                          MMQ_MMA_TILE_X_K_FP4);

            // based on block-scaling document, 2 threads in each quad need to supply to the scale value
            const int tidx         = threadIdx.x / 4 + (threadIdx.x % 2) * 8;
            scaleA[n][k01 / (2 * QI_MXFP4)] =
                *(x_sc + (i0 + n * tile_A::I + tidx) * MMQ_MMA_TILE_X_K_FP4 + k0 / (2 * QI_MXFP4));
        }
    }

#pragma unroll
    for (int j0 = 0; j0 < mmq_x; j0 += ntx * tile_C::J) {

Review comment (suggestion):
```suggestion
    // tile_A has a length of 64 logical values vs. 32 values in block_mxfp4
    tile_A   A[ntx][MMQ_TILE_NE_K / (2 * QI_MXFP4)];
```

Code after:
#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
}

template <int mmq_x, int mmq_y>
static __device__ __forceinline__ void vec_dot_mxfp4_mxfp4_mma(const int * __restrict__ x,
                                                               const int * __restrict__ y,
                                                               float * __restrict__ sum,
                                                               const int k00) {
    typedef tile<16, 8, int>   tile_A;
    typedef tile<8, 8, int>    tile_B;
    typedef tile<16, 8, float> tile_C;  // Output is float for native scaled MMA

    constexpr int granularity   = mmq_get_granularity_device(mmq_x);
    constexpr int rows_per_warp = 2 * granularity;
    constexpr int ntx           = rows_per_warp / tile_C::I;  // Number of x minitiles per warp.

    y += (threadIdx.y % ntx) * (tile_C::J * MMQ_TILE_Y_FP4_K);

    // Match layout from load_tiles_mxfp4_fp4
    const int *      x_qs = (const int *) x;
    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 has a length of 64 logical values vs. 32 values in block_mxfp4
    tile_A   A[ntx][MMQ_TILE_NE_K / (2 * QI_MXFP4)];
    uint32_t scaleA[ntx][MMQ_TILE_NE_K / (2 * QI_MXFP4)];

    // Block scale
    // Each thread has to point to a 4 byte scale value
    // https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-block-scaling

    const int i0 = (threadIdx.y / ntx) * rows_per_warp;

#pragma unroll
    for (int n = 0; n < ntx; ++n) {
#pragma unroll
        for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 2 * QI_MXFP4) {
            const int k0 = k00 + k01;

            load_ldmatrix(A[n][k01 / (2 * QI_MXFP4)], x_qs + (i0 + n * tile_A::I) * MMQ_MMA_TILE_X_K_FP4 + k0,
                          MMQ_MMA_TILE_X_K_FP4);

            // based on block-scaling document, 2 threads in each quad need to supply to the scale value
            const int tidx         = threadIdx.x / 4 + (threadIdx.x % 2) * 8;
            scaleA[n][k01 / (2 * QI_MXFP4)] =
                *(x_sc + (i0 + n * tile_A::I + tidx) * MMQ_MMA_TILE_X_K_FP4 + k0 / (2 * QI_MXFP4));
        }
    }

#pragma unroll