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

Code before:
    constexpr int warp_size = ggml_cuda_get_physical_warp_size();

#if defined(BLACKWELL_MMA_AVAILABLE)
    int *      x_qs = (int *) x_tile;
    uint32_t * x_sc = (uint32_t *) (x_qs + 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

#pragma unroll
    for (int i0 = 0; i0 < mmq_y; i0 += rows_per_warp * nwarps) {
        int i = i0 + threadIdx.y * rows_per_warp + row_in_warp;

        if (need_check) {
            i = min(i, i_max);
        }

        const block_mxfp4 * bxi = (const block_mxfp4 *) x + kbx0 + i * stride + kbx;

        int aux_q4[4];
        memcpy(aux_q4, bxi->qs, 16);

        // Compress: extract low nibbles from each byte and pack into 16 bits
        // Input byte layout: [hi3|lo3][hi2|lo2][hi1|lo1][hi0|lo0]
        // Output: [lo3|lo2|lo1|lo0] as 16 bits
        const auto compress = [](const int x) -> int {
            const int m = x & 0x0F0F0F0F;  // isolate low nibbles: 0x0lo30lo20lo10lo0
            // Pack nibbles: shift and combine
            const int t1 = (m | (m >> 4)) & 0x00FF00FF;  // 0x00_lo3lo2_00_lo1lo0
            return (t1 | (t1 >> 8)) & 0x0000FFFF;         // 0x0000_lo3lo2lo1lo0
        };

        const int k0 = kbx * 4;

        x_qs[i * MMQ_MMA_TILE_X_K_FP4 + k0 + 0] = compress(aux_q4[1]) << 16 | compress(aux_q4[0]);
        x_qs[i * MMQ_MMA_TILE_X_K_FP4 + k0 + 1] = compress(aux_q4[3]) << 16 | compress(aux_q4[2]);
        x_qs[i * MMQ_MMA_TILE_X_K_FP4 + k0 + 2] = compress(aux_q4[1] >> 4) << 16 | compress(aux_q4[0] >> 4);
        x_qs[i * MMQ_MMA_TILE_X_K_FP4 + k0 + 3] = compress(aux_q4[3] >> 4) << 16 | compress(aux_q4[2] >> 4);

        // Load E8M0 scales: pack 2 consecutive scales into one uint32
        if (kbx % 2 == 0) {
            uint32_t e = bxi->e;
            e |= ((bxi + 1)->e << 8);
            x_sc[i * MMQ_MMA_TILE_X_K_FP4 + kbx / 2] = e;
        }
    }

Review comment (suggestion):
Not to say that this should be addressed in this PR, but I think that the `block_mxfp4` datatype has poor alignment. Ideally one would have values that are aligned to at least 4 bytes. It my make sense to convert weights from a format of 32 values + 1 scale to one with 128 values + 4 scales. Or to come up with a better way to load weights with poor alignment (something like load them to shared memory in an aligned way first).

Code after:
template <int mmq_y, bool need_check>
static __device__ __forceinline__ void load_tiles_mxfp4_fp4(const char * __restrict__ x,
                                                            int * __restrict__ x_tile,
                                                            const int kbx0,
                                                            const int i_max,
                                                            const int stride) {
    constexpr int nwarps = mmq_get_nwarps_device();
    constexpr int warp_size = ggml_cuda_get_physical_warp_size();

    int *      x_qs = (int *) x_tile;
    uint32_t * x_sc = (uint32_t *) (x_qs + 2 * MMQ_TILE_NE_K);

    const int txi = threadIdx.x;

    constexpr int iter_k = get_iter_k(GGML_TYPE_MXFP4);

    constexpr int threads_per_row = iter_k / QK_MXFP4;  // each thread processes 1 block
    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) {
        int i = i0 + threadIdx.y * rows_per_warp + row_in_warp;

        if constexpr (need_check) {
            i = min(i, i_max);
        }

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

        // Load E8M0 scales: pack 2 consecutive scales into one uint32
        if (kbx % 2 == 0) {
            uint32_t e = bxi->e;
            e |= ((bxi + 1)->e << 8);
            x_sc[i * MMQ_MMA_TILE_X_K_FP4 + kbx / 2] = e;
        }
    }
}

template <int mmq_x, int mmq_y>
static __device__ __forceinline__ void vec_dot_q8_0_q8_1_dp4a(
    const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) {
    constexpr int nwarps = mmq_get_nwarps_device();
    constexpr int warp_size = ggml_cuda_get_physical_warp_size();

    constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q8_0, mmq_y);