#include <cuda_fp16.h>
#include <cuda_runtime.h>
#if CUDART_VERSION >= 11080
#include <cuda_fp8.h>
#define MAKEPAD_GGML_CUDA_FP8_AVAILABLE
#endif
#include <limits.h>
#include <stdint.h>

static constexpr int WARP_SIZE = 32;
static constexpr uint32_t QK8_1 = 32;
static constexpr uint32_t QK_NVFP4 = 64;
static constexpr uint32_t QK_NVFP4_SUB = 16;
static constexpr uint32_t QI8_0 = 8;
static constexpr uint32_t QI8_1 = 8;
static constexpr uint32_t QI_NVFP4 = 8;
static constexpr int CUDA_QUANTIZE_BLOCK_SIZE_MMQ = 128;
static constexpr int MMQ_TILE_NE_K = 32;
static constexpr int MMQ_TILE_Y_K = MMQ_TILE_NE_K + MMQ_TILE_NE_K / QI8_1;
static constexpr int MMQ_ITER_K = 256;
static constexpr int MMQ_NWARPS = 8;
static constexpr int MMQ_WARP_THREADS = WARP_SIZE * MMQ_NWARPS;
static constexpr int MMQ_Y = 128;
static constexpr int MMQ_MMA_TILE_X_K_NVFP4 = 2 * MMQ_TILE_NE_K + MMQ_TILE_NE_K / 2 + 4;
static constexpr int BLOCK_Q8_1_MMQ_INTS = 36;

typedef struct {
    uint16_t d;
    uint16_t s;
    int8_t qs[QK8_1];
} __align__(4) block_q8_1;

typedef struct {
    union {
        float d4[4];
    };
    int8_t qs[4 * QK8_1];
} __align__(4) block_q8_1_mmq;

typedef struct {
    uint8_t d[QK_NVFP4 / QK_NVFP4_SUB];
    uint8_t qs[QK_NVFP4 / 2];
} __align__(4) block_nvfp4;

static_assert(sizeof(block_q8_1) == 36, "wrong q8_1 block size");
static_assert(sizeof(block_q8_1_mmq) == 4 * sizeof(block_q8_1), "wrong q8_1 mmq block size");
static_assert(sizeof(block_nvfp4) == 36, "wrong nvfp4 block size");
static_assert(BLOCK_Q8_1_MMQ_INTS == static_cast<int>(sizeof(block_q8_1_mmq) / sizeof(int)), "wrong q8_1 mmq int size");

__device__ __constant__ int8_t KVALUES_MXFP4_X2[16] = {
    0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12,
};

template <int WIDTH>
static __device__ __forceinline__ float makepad_cuda_warp_reduce_sum(float x) {
#pragma unroll
    for (int offset = WIDTH / 2; offset > 0; offset >>= 1) {
        x += __shfl_xor_sync(0xffffffffu, x, offset, WIDTH);
    }
    return x;
}

template <int WIDTH>
static __device__ __forceinline__ float makepad_cuda_warp_reduce_max(float x) {
#pragma unroll
    for (int offset = WIDTH / 2; offset > 0; offset >>= 1) {
        x = fmaxf(x, __shfl_xor_sync(0xffffffffu, x, offset, WIDTH));
    }
    return x;
}

static __device__ __forceinline__ float makepad_cuda_ue4m3_to_fp32(uint8_t x) {
#if defined(MAKEPAD_GGML_CUDA_FP8_AVAILABLE)
    const uint32_t bits = x * (x != 0x7F && x != 0xFF);
    const __nv_fp8_e4m3 xf = *reinterpret_cast<const __nv_fp8_e4m3 *>(&bits);
    return static_cast<float>(xf) * 0.5f;
#else
    if (x == 0 || x == 0x7F || x == 0xFF) {
        return 0.0f;
    }
    const int exp = (x >> 3) & 0xF;
    const int man = x & 0x7;
    const float raw = exp == 0 ? ldexpf((float) man, -9) : ldexpf(1.0f + (float) man / 8.0f, exp - 7);
    return raw * 0.5f;
#endif
}

static __device__ __forceinline__ int makepad_cuda_dp4a(const int a, const int b, int c) {
#if __CUDA_ARCH__ >= 610
    return __dp4a(a, b, c);
#else
    const int8_t * a8 = reinterpret_cast<const int8_t *>(&a);
    const int8_t * b8 = reinterpret_cast<const int8_t *>(&b);
    return c + a8[0] * b8[0] + a8[1] * b8[1] + a8[2] * b8[2] + a8[3] * b8[3];
#endif
}

static __device__ __forceinline__ int2 makepad_cuda_get_int_from_table_16(
        const int q4,
        const int8_t * table) {
    const uint32_t * table32 = reinterpret_cast<const uint32_t *>(table);
    uint32_t tmp[2];
    const uint32_t low_high_selection_indices = 0x32103210u | ((static_cast<uint32_t>(q4) & 0x88888888u) >> 1);

#pragma unroll
    for (uint32_t i = 0; i < 2; ++i) {
        const uint32_t shift = 16u * i;
        const uint32_t low = __byte_perm(table32[0], table32[1], static_cast<uint32_t>(q4) >> shift);
        const uint32_t high = __byte_perm(table32[2], table32[3], static_cast<uint32_t>(q4) >> shift);
        tmp[i] = __byte_perm(low, high, low_high_selection_indices >> shift);
    }

    return make_int2(
        __byte_perm(tmp[0], tmp[1], 0x6420),
        __byte_perm(tmp[0], tmp[1], 0x7531));
}

template <int vdr>
static __device__ __forceinline__ float makepad_cuda_vec_dot_q8_0_16_q8_1_impl(
        const int * v,
        const int * u,
        const float * d8_0,
        const float d8_1) {
    float sumf = 0.0f;

#pragma unroll
    for (int i0 = 0; i0 < vdr; i0 += QI8_0 / 2) {
        int sumi = 0;

#pragma unroll
        for (int i = i0; i < i0 + QI8_0 / 2; ++i) {
            sumi = makepad_cuda_dp4a(v[i], u[i], sumi);
        }

        sumf += d8_0[i0 / (QI8_0 / 2)] * static_cast<float>(sumi);
    }

    return d8_1 * sumf;
}

static __host__ __device__ __forceinline__ int makepad_cuda_pad(int value, int align) {
    return ((value + align - 1) / align) * align;
}

template <int NBYTES>
static __device__ __forceinline__ void makepad_cuda_memcpy_1(
        void * __restrict__ dst,
        const void * __restrict__ src) {
#pragma unroll
    for (int i = 0; i < NBYTES / 4; ++i) {
        reinterpret_cast<int *>(dst)[i] = reinterpret_cast<const int *>(src)[i];
    }
}

template <int I, int J>
struct makepad_cuda_mma_tile_int;

template <>
struct makepad_cuda_mma_tile_int<8, 4> {
    static constexpr int I_VALUE = 8;
    static constexpr int J_VALUE = 4;
    static constexpr int ne = I_VALUE * J_VALUE / WARP_SIZE;
    int x[ne] = {0};

    static __device__ __forceinline__ int get_i(int) {
        return threadIdx.x / 4;
    }

    static __device__ __forceinline__ int get_j(int) {
        return threadIdx.x % 4;
    }
};

template <>
struct makepad_cuda_mma_tile_int<16, 4> {
    static constexpr int I_VALUE = 16;
    static constexpr int J_VALUE = 4;
    static constexpr int ne = I_VALUE * J_VALUE / WARP_SIZE;
    int x[ne] = {0};

    static __device__ __forceinline__ int get_i(int l) {
        return l * 8 + threadIdx.x / 4;
    }

    static __device__ __forceinline__ int get_j(int) {
        return threadIdx.x % 4;
    }
};

template <>
struct makepad_cuda_mma_tile_int<16, 8> {
    static constexpr int I_VALUE = 16;
    static constexpr int J_VALUE = 8;
    static constexpr int ne = I_VALUE * J_VALUE / WARP_SIZE;
    int x[ne] = {0};

    static __device__ __forceinline__ int get_i(int l) {
        return (l / 2) * 8 + threadIdx.x / 4;
    }

    static __device__ __forceinline__ int get_j(int l) {
        return (threadIdx.x % 4) * 2 + (l % 2);
    }
};

template <int I, int J>
static __device__ __forceinline__ void makepad_cuda_load_generic(
        makepad_cuda_mma_tile_int<I, J> & tile,
        const int * __restrict__ src,
        int stride) {
#pragma unroll
    for (int l = 0; l < tile.ne; ++l) {
        tile.x[l] = src[tile.get_i(l) * stride + tile.get_j(l)];
    }
}

static __device__ __forceinline__ void makepad_cuda_load_ldmatrix(
        makepad_cuda_mma_tile_int<16, 4> & tile,
        const int * __restrict__ src,
        int stride) {
#if __CUDA_ARCH__ >= 750
    const int * xs = src + (threadIdx.x % tile.I_VALUE) * stride;
    asm volatile("ldmatrix.sync.aligned.m8n8.x2.b16 {%0, %1}, [%2];"
        : "=r"(tile.x[0]), "=r"(tile.x[1])
        : "l"(xs));
#else
    makepad_cuda_load_generic(tile, src, stride);
#endif
}

static __device__ __forceinline__ void makepad_cuda_load_ldmatrix(
        makepad_cuda_mma_tile_int<16, 8> & tile,
        const int * __restrict__ src,
        int stride) {
#if __CUDA_ARCH__ >= 750
    const int * xs =
        src + (threadIdx.x % tile.I_VALUE) * stride + (threadIdx.x / tile.I_VALUE) * (tile.J_VALUE / 2);
    asm volatile("ldmatrix.sync.aligned.m8n8.x4.b16 {%0, %1, %2, %3}, [%4];"
        : "=r"(tile.x[0]), "=r"(tile.x[1]), "=r"(tile.x[2]), "=r"(tile.x[3])
        : "l"(xs));
#else
    makepad_cuda_load_generic(tile, src, stride);
#endif
}

static __device__ __forceinline__ void makepad_cuda_mma(
        makepad_cuda_mma_tile_int<16, 8> & d,
        const makepad_cuda_mma_tile_int<16, 4> & a,
        const makepad_cuda_mma_tile_int<8, 4> & b) {
#if __CUDA_ARCH__ >= 800
    asm("mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};"
        : "+r"(d.x[0]), "+r"(d.x[1]), "+r"(d.x[2]), "+r"(d.x[3])
        : "r"(a.x[0]), "r"(a.x[1]), "r"(b.x[0]));
#elif __CUDA_ARCH__ >= 750
    asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0, %1}, {%2}, {%3}, {%0, %1};"
        : "+r"(d.x[0]), "+r"(d.x[1])
        : "r"(a.x[0]), "r"(b.x[0]));
    asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0, %1}, {%2}, {%3}, {%0, %1};"
        : "+r"(d.x[2]), "+r"(d.x[3])
        : "r"(a.x[1]), "r"(b.x[0]));
#endif
}

static __global__ void makepad_cuda_quantize_q8_1_mmq_f32_kernel(
        const float * __restrict__ input_f32,
        block_q8_1_mmq * __restrict__ output_q8_1_mmq,
        uint32_t n_cols,
        uint32_t n_rows,
        uint32_t padded_rows) {
    const int64_t i0 = (static_cast<int64_t>(blockDim.x) * blockIdx.y + threadIdx.x) * 4;

    if (i0 >= n_cols || blockIdx.x >= padded_rows) {
        return;
    }

    const int64_t row = blockIdx.x;
    const int64_t ib = (i0 / (4 * QK8_1)) * padded_rows + row;
    const int64_t iqs = i0 % (4 * QK8_1);

    const float4 * input_f32x4 = reinterpret_cast<const float4 *>(input_f32);
    const float4 xi = row < n_rows
        ? input_f32x4[(row * n_cols + i0) / 4]
        : make_float4(0.0f, 0.0f, 0.0f, 0.0f);

    float amax = fabsf(xi.x);
    amax = fmaxf(amax, fabsf(xi.y));
    amax = fmaxf(amax, fabsf(xi.z));
    amax = fmaxf(amax, fabsf(xi.w));

#pragma unroll
    for (int offset = 4; offset > 0; offset >>= 1) {
        amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFFu, amax, offset, WARP_SIZE));
    }

    const float d_inv = amax == 0.0f ? 0.0f : 127.0f / amax;
    char4 q;
    q.x = static_cast<int8_t>(roundf(xi.x * d_inv));
    q.y = static_cast<int8_t>(roundf(xi.y * d_inv));
    q.z = static_cast<int8_t>(roundf(xi.z * d_inv));
    q.w = static_cast<int8_t>(roundf(xi.w * d_inv));

    char4 * output_qs4 = reinterpret_cast<char4 *>(output_q8_1_mmq[ib].qs);
    output_qs4[iqs / 4] = q;

    if (iqs % QK8_1 != 0) {
        return;
    }

    output_q8_1_mmq[ib].d4[iqs / QK8_1] = amax / 127.0f;
}

template <bool need_check>
static __device__ __forceinline__ void makepad_cuda_load_tiles_nvfp4(
        const block_nvfp4 * __restrict__ weights_nvfp4,
        int * __restrict__ x_tile,
        const int kb0,
        const int i_max,
        const int stride) {
    int * x_qs = x_tile;
#if __CUDA_ARCH__ >= 750
    float * x_df = reinterpret_cast<float *>(x_qs + MMQ_TILE_NE_K * 2);
#else
    float * x_df = reinterpret_cast<float *>(x_qs + MMQ_Y * MMQ_TILE_NE_K * 2 + MMQ_Y);
#endif

    constexpr int threads_per_row = MMQ_ITER_K / QK_NVFP4;
    constexpr int rows_per_warp = WARP_SIZE / threads_per_row;
    const int kbx = threadIdx.x % threads_per_row;
    const int row_in_warp = threadIdx.x / threads_per_row;

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

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

        const block_nvfp4 * block = weights_nvfp4 + kb0 + i * stride + kbx;
        const uint32_t * src_qs = reinterpret_cast<const uint32_t *>(block->qs);
        const int kqs = 16 * kbx;
        const int ksc = 4 * kbx;

#pragma unroll
        for (int sub = 0; sub < static_cast<int>(QK_NVFP4 / QK_NVFP4_SUB); ++sub) {
            const int2 q0 = makepad_cuda_get_int_from_table_16(
                static_cast<int>(src_qs[2 * sub + 0]),
                KVALUES_MXFP4_X2);
            const int2 q1 = makepad_cuda_get_int_from_table_16(
                static_cast<int>(src_qs[2 * sub + 1]),
                KVALUES_MXFP4_X2);

#if __CUDA_ARCH__ >= 750
            x_qs[i * MMQ_MMA_TILE_X_K_NVFP4 + kqs + 4 * sub + 0] = q0.x;
            x_qs[i * MMQ_MMA_TILE_X_K_NVFP4 + kqs + 4 * sub + 1] = q1.x;
            x_qs[i * MMQ_MMA_TILE_X_K_NVFP4 + kqs + 4 * sub + 2] = q0.y;
            x_qs[i * MMQ_MMA_TILE_X_K_NVFP4 + kqs + 4 * sub + 3] = q1.y;
            x_df[i * MMQ_MMA_TILE_X_K_NVFP4 + ksc + sub] =
                makepad_cuda_ue4m3_to_fp32(block->d[sub]);
#else
            x_qs[i * (2 * MMQ_TILE_NE_K + 1) + kqs + 4 * sub + 0] = q0.x;
            x_qs[i * (2 * MMQ_TILE_NE_K + 1) + kqs + 4 * sub + 1] = q1.x;
            x_qs[i * (2 * MMQ_TILE_NE_K + 1) + kqs + 4 * sub + 2] = q0.y;
            x_qs[i * (2 * MMQ_TILE_NE_K + 1) + kqs + 4 * sub + 3] = q1.y;
            x_df[i * (2 * MMQ_TILE_NE_K * 2 / QI_NVFP4) + i / (QK_NVFP4_SUB / QI_NVFP4) + ksc + sub] =
                makepad_cuda_ue4m3_to_fp32(block->d[sub]);
#endif
        }
    }
}

template <int mmq_x>
static __device__ __forceinline__ void makepad_cuda_vec_dot_q8_0_16_q8_1_dp4a(
        const int * __restrict__ x,
        const int * __restrict__ y,
        float * __restrict__ sum,
        const int k00) {
    const int * x_qs = x;
    const float * x_df = reinterpret_cast<const float *>(x_qs + MMQ_Y * MMQ_TILE_NE_K * 2 + MMQ_Y);
    const int * y_qs = y + 4;
    const float * y_df = reinterpret_cast<const float *>(y);

    for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) {
        const int k0 = k00 + k01;

#pragma unroll
        for (int j0 = 0; j0 < mmq_x; j0 += MMQ_NWARPS) {
            const int j = j0 + threadIdx.y;

#pragma unroll
            for (int i0 = 0; i0 < MMQ_Y; i0 += WARP_SIZE) {
                const int i = i0 + threadIdx.x;

                sum[j0 / MMQ_NWARPS * (MMQ_Y / WARP_SIZE) + i0 / WARP_SIZE] +=
                    makepad_cuda_vec_dot_q8_0_16_q8_1_impl<QI8_0>(
                        &x_qs[i * (2 * MMQ_TILE_NE_K + 1) + k0],
                        &y_qs[j * MMQ_TILE_Y_K + k01],
                        &x_df[i * (2 * MMQ_TILE_NE_K * 2 / QI8_0) + i / (QI8_0 / 4) + k0 / (QI8_0 / 2)],
                        y_df[j * MMQ_TILE_Y_K + k01 / QI8_1]);
            }
        }
    }
}

template <int mmq_x>
static __device__ __forceinline__ void makepad_cuda_vec_dot_q8_0_16_q8_1_mma(
        const int * __restrict__ x,
        const int * __restrict__ y,
        float * __restrict__ sum,
        const int k00) {
#if __CUDA_ARCH__ >= 750
    using tile_a = makepad_cuda_mma_tile_int<16, 4>;
    using tile_a8 = makepad_cuda_mma_tile_int<16, 8>;
    using tile_b = makepad_cuda_mma_tile_int<8, 4>;
    using tile_c = makepad_cuda_mma_tile_int<16, 8>;

    constexpr int granularity = mmq_x >= 48 ? 16 : 8;
    constexpr int rows_per_warp = 2 * granularity;
    constexpr int ntx = rows_per_warp / tile_c::I_VALUE;

    y += (threadIdx.y % ntx) * (tile_c::J_VALUE * MMQ_TILE_Y_K);

    const int * x_qs = x;
    const float * x_df = reinterpret_cast<const float *>(x_qs) + 2 * MMQ_TILE_NE_K;
    const int * y_qs = y + 4;
    const float * y_df = reinterpret_cast<const float *>(y);
    const int i0 = (threadIdx.y / ntx) * (ntx * tile_a::I_VALUE);

    tile_a a[ntx][8];
    float d_a[ntx][tile_c::ne / 2][8];

#pragma unroll
    for (int n = 0; n < ntx; ++n) {
#pragma unroll
        for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 8) {
            const int k0 = k00 + k01;
            makepad_cuda_load_ldmatrix(
                reinterpret_cast<tile_a8 *>(a[n])[k01 / 8],
                x_qs + (i0 + n * tile_a::I_VALUE) * MMQ_MMA_TILE_X_K_NVFP4 + k0,
                MMQ_MMA_TILE_X_K_NVFP4);
        }

#pragma unroll
        for (int l = 0; l < tile_c::ne / 2; ++l) {
            const int i = i0 + n * tile_c::I_VALUE + tile_c::get_i(2 * l);

#pragma unroll
            for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) {
                const int k0 = k00 + k01;
                d_a[n][l][k01 / 4] = x_df[i * MMQ_MMA_TILE_X_K_NVFP4 + k0 / 4];
            }
        }
    }

#pragma unroll
    for (int j0 = 0; j0 < mmq_x; j0 += ntx * tile_c::J_VALUE) {
#pragma unroll
        for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 8) {
            tile_b b[2];
            float d_b[tile_c::ne / 2];

            makepad_cuda_load_generic(
                b[0],
                y_qs + j0 * MMQ_TILE_Y_K + k01,
                MMQ_TILE_Y_K);
            makepad_cuda_load_generic(
                b[1],
                y_qs + j0 * MMQ_TILE_Y_K + tile_b::J_VALUE + k01,
                MMQ_TILE_Y_K);

#pragma unroll
            for (int l = 0; l < tile_c::ne / 2; ++l) {
                const int j = j0 + tile_c::get_j(l);
                d_b[l] = y_df[j * MMQ_TILE_Y_K + k01 / QI8_1];
            }

#pragma unroll
            for (int n = 0; n < ntx; ++n) {
                tile_c c[2];
                makepad_cuda_mma(c[0], a[n][k01 / 4 + 0], b[0]);
                makepad_cuda_mma(c[1], a[n][k01 / 4 + 1], b[1]);

#pragma unroll
                for (int l = 0; l < tile_c::ne; ++l) {
                    sum[(j0 / tile_c::J_VALUE + n) * tile_c::ne + l] +=
                        d_b[l % 2] *
                        (c[0].x[l] * d_a[n][l / 2][k01 / 4 + 0] +
                         c[1].x[l] * d_a[n][l / 2][k01 / 4 + 1]);
                }
            }
        }
    }
#else
    (void) x;
    (void) y;
    (void) sum;
    (void) k00;
#endif
}

template <int mmq_x, bool need_check>
static __device__ __forceinline__ void makepad_cuda_mmq_write_back(
        const float * __restrict__ sum,
        float * __restrict__ dst,
        const int stride,
        const int i_max,
        const int j_max) {
#pragma unroll
    for (int j0 = 0; j0 < mmq_x; j0 += MMQ_NWARPS) {
        const int j = j0 + threadIdx.y;

        if (j > j_max) {
            return;
        }

#pragma unroll
        for (int i0 = 0; i0 < MMQ_Y; i0 += WARP_SIZE) {
            const int i = i0 + threadIdx.x;

            if constexpr (need_check) {
                if (i > i_max) {
                    continue;
                }
            }

            dst[j * stride + i] = sum[(j0 / MMQ_NWARPS) * (MMQ_Y / WARP_SIZE) + i0 / WARP_SIZE];
        }
    }
}

template <int mmq_x, bool need_check>
static __device__ __forceinline__ void makepad_cuda_mmq_write_back_mma(
        const float * __restrict__ sum,
        float * __restrict__ dst,
        const int stride,
        const int i_max,
        const int j_max) {
#if __CUDA_ARCH__ >= 750
    using tile_c = makepad_cuda_mma_tile_int<16, 8>;

    constexpr int granularity = mmq_x >= 48 ? 16 : 8;
    constexpr int rows_per_warp = 2 * granularity;
    constexpr int ntx = rows_per_warp / tile_c::I_VALUE;
    const int i0 = (threadIdx.y / ntx) * (ntx * tile_c::I_VALUE);

#pragma unroll
    for (int j0 = 0; j0 < mmq_x; j0 += ntx * tile_c::J_VALUE) {
#pragma unroll
        for (int n = 0; n < ntx; ++n) {
#pragma unroll
            for (int l = 0; l < tile_c::ne; ++l) {
                const int j = j0 + (threadIdx.y % ntx) * tile_c::J_VALUE + tile_c::get_j(l);
                if (j > j_max) {
                    continue;
                }
                const int i = i0 + n * tile_c::I_VALUE + tile_c::get_i(l);
                if constexpr (need_check) {
                    if (i > i_max) {
                        continue;
                    }
                }
                dst[j * stride + i] = sum[(j0 / tile_c::J_VALUE + n) * tile_c::ne + l];
            }
        }
    }
#else
    makepad_cuda_mmq_write_back<mmq_x, need_check>(sum, dst, stride, i_max, j_max);
#endif
}

template <int mmq_x, bool need_check>
__launch_bounds__(MMQ_WARP_THREADS, 1)
static __global__ void makepad_cuda_nvfp4_q8_1_mmq_matmul_kernel(
        const block_nvfp4 * __restrict__ weights_nvfp4,
        const int * __restrict__ input_q8_1_mmq,
        float * __restrict__ output_f32,
        const int ncols_x,
        const int out_rows,
        const int input_rows) {
    const int it = blockIdx.x;
    const int jt = blockIdx.y;

    const int i_max = out_rows - it * MMQ_Y - 1;
    const int j_max = input_rows - jt * mmq_x - 1;
    if (i_max < 0 || j_max < 0) {
        return;
    }

    const int tile_y_ints = makepad_cuda_pad(mmq_x * MMQ_TILE_Y_K, MMQ_WARP_THREADS);
    extern __shared__ int shared_data[];
    int * tile_y = shared_data + mmq_x;
    int * tile_x = tile_y + tile_y_ints;

    constexpr int qk = QK_NVFP4;
    constexpr int ne_block = 4 * QK8_1;
    constexpr int blocks_per_iter = MMQ_ITER_K / qk;
    constexpr int sum_elems = mmq_x * MMQ_Y / (MMQ_NWARPS * WARP_SIZE);

    float sum[sum_elems] = {0.0f};

    const int stride_row_x = ncols_x / qk;
    const int offset_x = it * MMQ_Y * stride_row_x;
    const int offset_y = jt * mmq_x * BLOCK_Q8_1_MMQ_INTS;
    float * dst = output_f32 + jt * mmq_x * out_rows + it * MMQ_Y;

    for (int kb0 = 0; kb0 < stride_row_x; kb0 += blocks_per_iter) {
        makepad_cuda_load_tiles_nvfp4<need_check>(
            weights_nvfp4,
            tile_x,
            offset_x + kb0,
            i_max,
            stride_row_x);

        const int * by0 =
            input_q8_1_mmq + input_rows * (kb0 * qk / ne_block) * BLOCK_Q8_1_MMQ_INTS + offset_y;

#pragma unroll
        for (int l0 = 0; l0 < mmq_x * MMQ_TILE_Y_K; l0 += MMQ_WARP_THREADS) {
            const int l = l0 + threadIdx.y * WARP_SIZE + threadIdx.x;
            if (l < mmq_x * MMQ_TILE_Y_K) {
                const int col = l / MMQ_TILE_Y_K;
                tile_y[l] = col <= j_max ? by0[l] : 0;
            }
        }

        __syncthreads();
#if __CUDA_ARCH__ >= 750
        makepad_cuda_vec_dot_q8_0_16_q8_1_mma<mmq_x>(tile_x, tile_y, sum, 0);
#else
        makepad_cuda_vec_dot_q8_0_16_q8_1_dp4a<mmq_x>(tile_x, tile_y, sum, 0);
#endif
        __syncthreads();

        const int * by1 = by0 + input_rows * BLOCK_Q8_1_MMQ_INTS;

#pragma unroll
        for (int l0 = 0; l0 < mmq_x * MMQ_TILE_Y_K; l0 += MMQ_WARP_THREADS) {
            const int l = l0 + threadIdx.y * WARP_SIZE + threadIdx.x;
            if (l < mmq_x * MMQ_TILE_Y_K) {
                const int col = l / MMQ_TILE_Y_K;
                tile_y[l] = col <= j_max ? by1[l] : 0;
            }
        }

        __syncthreads();
#if __CUDA_ARCH__ >= 750
        makepad_cuda_vec_dot_q8_0_16_q8_1_mma<mmq_x>(tile_x, tile_y, sum, MMQ_TILE_NE_K);
#else
        makepad_cuda_vec_dot_q8_0_16_q8_1_dp4a<mmq_x>(tile_x, tile_y, sum, MMQ_TILE_NE_K);
#endif
        __syncthreads();
    }

#if __CUDA_ARCH__ >= 750
    makepad_cuda_mmq_write_back_mma<mmq_x, need_check>(sum, dst, out_rows, i_max, j_max);
#else
    makepad_cuda_mmq_write_back<mmq_x, need_check>(sum, dst, out_rows, i_max, j_max);
#endif
}

template <int mmq_x, bool need_check, bool check_y, bool fixup>
static __device__ __forceinline__ void makepad_cuda_nvfp4_q8_1_mmq_process_tile(
        const block_nvfp4 * __restrict__ weights_nvfp4,
        const int * __restrict__ input_q8_1_mmq,
        float * __restrict__ output_f32,
        float * __restrict__ tmp_fixup,
        const int ncols_x,
        const int out_rows,
        const int input_rows,
        const int it,
        const int jt,
        const int kb0_start,
        const int kb0_stop) {
    const int i_max = out_rows - it * MMQ_Y - 1;
    const int j_max = input_rows - jt * mmq_x - 1;
    if (i_max < 0 || j_max < 0) {
        return;
    }

    const int tile_y_ints = makepad_cuda_pad(mmq_x * MMQ_TILE_Y_K, MMQ_WARP_THREADS);
    extern __shared__ int shared_data[];
    int * tile_y = shared_data + mmq_x;
    int * tile_x = tile_y + tile_y_ints;

    constexpr int qk = QK_NVFP4;
    constexpr int ne_block = 4 * QK8_1;
    constexpr int blocks_per_iter = MMQ_ITER_K / qk;
    constexpr int sum_elems = mmq_x * MMQ_Y / (MMQ_NWARPS * WARP_SIZE);

    float sum[sum_elems] = {0.0f};

    const int stride_row_x = ncols_x / qk;
    const int offset_x = it * MMQ_Y * stride_row_x;
    const int offset_y = jt * mmq_x * BLOCK_Q8_1_MMQ_INTS;

    for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) {
        makepad_cuda_load_tiles_nvfp4<need_check>(
            weights_nvfp4,
            tile_x,
            offset_x + kb0,
            i_max,
            stride_row_x);

        const int * by0 =
            input_q8_1_mmq + input_rows * (kb0 * qk / ne_block) * BLOCK_Q8_1_MMQ_INTS + offset_y;

#pragma unroll
        for (int l0 = 0; l0 < mmq_x * MMQ_TILE_Y_K; l0 += MMQ_WARP_THREADS) {
            const int l = l0 + threadIdx.y * WARP_SIZE + threadIdx.x;
            if (l < mmq_x * MMQ_TILE_Y_K) {
                if constexpr (check_y) {
                    const int col = l / MMQ_TILE_Y_K;
                    tile_y[l] = col <= j_max ? by0[l] : 0;
                } else {
                    tile_y[l] = by0[l];
                }
            }
        }

        __syncthreads();
#if __CUDA_ARCH__ >= 750
        makepad_cuda_vec_dot_q8_0_16_q8_1_mma<mmq_x>(tile_x, tile_y, sum, 0);
#else
        makepad_cuda_vec_dot_q8_0_16_q8_1_dp4a<mmq_x>(tile_x, tile_y, sum, 0);
#endif
        __syncthreads();

        const int * by1 = by0 + input_rows * BLOCK_Q8_1_MMQ_INTS;

#pragma unroll
        for (int l0 = 0; l0 < mmq_x * MMQ_TILE_Y_K; l0 += MMQ_WARP_THREADS) {
            const int l = l0 + threadIdx.y * WARP_SIZE + threadIdx.x;
            if (l < mmq_x * MMQ_TILE_Y_K) {
                if constexpr (check_y) {
                    const int col = l / MMQ_TILE_Y_K;
                    tile_y[l] = col <= j_max ? by1[l] : 0;
                } else {
                    tile_y[l] = by1[l];
                }
            }
        }

        __syncthreads();
#if __CUDA_ARCH__ >= 750
        makepad_cuda_vec_dot_q8_0_16_q8_1_mma<mmq_x>(tile_x, tile_y, sum, MMQ_TILE_NE_K);
#else
        makepad_cuda_vec_dot_q8_0_16_q8_1_dp4a<mmq_x>(tile_x, tile_y, sum, MMQ_TILE_NE_K);
#endif
        __syncthreads();
    }

    if constexpr (fixup) {
#if __CUDA_ARCH__ >= 750
        makepad_cuda_mmq_write_back_mma<mmq_x, false>(
            sum,
            tmp_fixup + blockIdx.x * (mmq_x * MMQ_Y),
            MMQ_Y,
            MMQ_Y - 1,
            mmq_x - 1);
#else
        makepad_cuda_mmq_write_back<mmq_x, false>(
            sum,
            tmp_fixup + blockIdx.x * (mmq_x * MMQ_Y),
            MMQ_Y,
            MMQ_Y - 1,
            mmq_x - 1);
#endif
    } else {
        float * dst = output_f32 + jt * mmq_x * out_rows + it * MMQ_Y;
#if __CUDA_ARCH__ >= 750
        makepad_cuda_mmq_write_back_mma<mmq_x, need_check>(sum, dst, out_rows, i_max, j_max);
#else
        makepad_cuda_mmq_write_back<mmq_x, need_check>(sum, dst, out_rows, i_max, j_max);
#endif
    }
}

template <int mmq_x, bool need_check, bool check_y>
__launch_bounds__(MMQ_WARP_THREADS, 1)
static __global__ void makepad_cuda_nvfp4_q8_1_mmq_stream_k_kernel(
        const block_nvfp4 * __restrict__ weights_nvfp4,
        const int * __restrict__ input_q8_1_mmq,
        float * __restrict__ output_f32,
        float * __restrict__ tmp_fixup,
        const int ncols_x,
        const int out_rows,
        const int input_rows) {
    constexpr int qk = QK_NVFP4;
    constexpr int blocks_per_iter = MMQ_ITER_K / qk;

    const int blocks_per_row = ncols_x / qk;
    const int ntx = (input_rows + mmq_x - 1) / mmq_x;
    const int nty = (out_rows + MMQ_Y - 1) / MMQ_Y;
    const int64_t work = static_cast<int64_t>(ntx) * nty * blocks_per_row;

    int64_t kbc = static_cast<int64_t>(blockIdx.x) * work / gridDim.x;
    int64_t kbc_stop = static_cast<int64_t>(blockIdx.x + 1) * work / gridDim.x;

    kbc -= (kbc % blocks_per_row) % blocks_per_iter;
    kbc_stop -= (kbc_stop % blocks_per_row) % blocks_per_iter;

    int kb0_start = static_cast<int>(kbc % blocks_per_row);
    int kb0_stop = min(blocks_per_row, kb0_start + static_cast<int>(kbc_stop - kbc));

    while (kbc < kbc_stop && kb0_stop == blocks_per_row) {
        int64_t tmp = kbc;
        const int it = static_cast<int>(tmp / (static_cast<int64_t>(ntx) * blocks_per_row));
        tmp -= static_cast<int64_t>(it) * ntx * blocks_per_row;
        const int jt = static_cast<int>(tmp / blocks_per_row);

        makepad_cuda_nvfp4_q8_1_mmq_process_tile<mmq_x, need_check, check_y, false>(
            weights_nvfp4,
            input_q8_1_mmq,
            output_f32,
            tmp_fixup,
            ncols_x,
            out_rows,
            input_rows,
            it,
            jt,
            kb0_start,
            kb0_stop);

        kbc += blocks_per_row;
        kbc -= kbc % blocks_per_row;

        kb0_start = 0;
        kb0_stop = min(blocks_per_row, static_cast<int>(kbc_stop - kbc));
    }

    if (kbc >= kbc_stop) {
        return;
    }

    int64_t tmp = kbc;
    const int it = static_cast<int>(tmp / (static_cast<int64_t>(ntx) * blocks_per_row));
    tmp -= static_cast<int64_t>(it) * ntx * blocks_per_row;
    const int jt = static_cast<int>(tmp / blocks_per_row);

    makepad_cuda_nvfp4_q8_1_mmq_process_tile<mmq_x, need_check, check_y, true>(
        weights_nvfp4,
        input_q8_1_mmq,
        output_f32,
        tmp_fixup,
        ncols_x,
        out_rows,
        input_rows,
        it,
        jt,
        kb0_start,
        kb0_stop);
}

template <int mmq_x, bool need_check>
__launch_bounds__(MMQ_WARP_THREADS, 1)
static __global__ void makepad_cuda_nvfp4_q8_1_mmq_stream_k_fixup_kernel(
        float * __restrict__ output_f32,
        const float * __restrict__ tmp_fixup,
        const int ncols_x,
        const int out_rows,
        const int input_rows) {
    constexpr int qk = QK_NVFP4;
    constexpr int blocks_per_iter = MMQ_ITER_K / qk;
    constexpr int sum_elems = mmq_x * MMQ_Y / (MMQ_NWARPS * WARP_SIZE);

    const int blocks_per_row = ncols_x / qk;
    const int ntx = (input_rows + mmq_x - 1) / mmq_x;
    const int nty = (out_rows + MMQ_Y - 1) / MMQ_Y;
    const int64_t work = static_cast<int64_t>(ntx) * nty * blocks_per_row;

    const int bidx0 = blockIdx.x;
    int64_t kbc0 = static_cast<int64_t>(bidx0) * work / gridDim.x;
    int64_t kbc0_stop = static_cast<int64_t>(bidx0 + 1) * work / gridDim.x;

    kbc0 -= (kbc0 % blocks_per_row) % blocks_per_iter;
    kbc0_stop -= (kbc0_stop % blocks_per_row) % blocks_per_iter;

    const bool did_not_have_any_data = kbc0 == kbc0_stop;
    const bool wrote_beginning_of_tile = kbc0 % blocks_per_row == 0;
    const bool did_not_write_last =
        kbc0 / blocks_per_row == kbc0_stop / blocks_per_row && kbc0_stop % blocks_per_row != 0;
    if (did_not_have_any_data || wrote_beginning_of_tile || did_not_write_last) {
        return;
    }

    bool any_fixup = false;
    float sum[sum_elems] = {0.0f};

    int64_t bidx = static_cast<int64_t>(bidx0) - 1;
    int64_t kbc_stop = kbc0;
    while (bidx >= 0) {
        int64_t kbc = bidx * work / gridDim.x;
        kbc -= (kbc % blocks_per_row) % blocks_per_iter;

        if (kbc == kbc_stop) {
            --bidx;
            kbc_stop = kbc;
            continue;
        }

        any_fixup = true;

#pragma unroll
        for (int j0 = 0; j0 < mmq_x; j0 += MMQ_NWARPS) {
            const int j = j0 + threadIdx.y;

#pragma unroll
            for (int i0 = 0; i0 < MMQ_Y; i0 += WARP_SIZE) {
                const int i = i0 + threadIdx.x;
                sum[(j0 / MMQ_NWARPS) * (MMQ_Y / WARP_SIZE) + i0 / WARP_SIZE] +=
                    tmp_fixup[bidx * (mmq_x * MMQ_Y) + j * MMQ_Y + i];
            }
        }

        if (kbc % blocks_per_row == 0 || kbc / blocks_per_row < kbc0 / blocks_per_row) {
            break;
        }
        --bidx;
        kbc_stop = kbc;
    }

    if (!any_fixup) {
        return;
    }

    int64_t tmp = kbc0;
    const int it = static_cast<int>(tmp / (static_cast<int64_t>(ntx) * blocks_per_row));
    tmp -= static_cast<int64_t>(it) * ntx * blocks_per_row;
    const int jt = static_cast<int>(tmp / blocks_per_row);

    float * dst = output_f32 + jt * mmq_x * out_rows + it * MMQ_Y;
    const int i_max = out_rows - it * MMQ_Y - 1;
    const int j_max = input_rows - jt * mmq_x - 1;

#pragma unroll
    for (int j0 = 0; j0 < mmq_x; j0 += MMQ_NWARPS) {
        const int j = j0 + threadIdx.y;

        if (j > j_max) {
            return;
        }

#pragma unroll
        for (int i0 = 0; i0 < MMQ_Y; i0 += WARP_SIZE) {
            const int i = i0 + threadIdx.x;

            if constexpr (need_check) {
                if (i > i_max) {
                    continue;
                }
            }

            dst[j * out_rows + i] += sum[(j0 / MMQ_NWARPS) * (MMQ_Y / WARP_SIZE) + i0 / WARP_SIZE];
        }
    }
}

static inline int makepad_cuda_nvfp4_mmq_shared_bytes(int mmq_x) {
    const int tile_y_ints = makepad_cuda_pad(mmq_x * MMQ_TILE_Y_K, MMQ_WARP_THREADS);
    const int tile_x_ints = MMQ_Y * MMQ_MMA_TILE_X_K_NVFP4;
    return (mmq_x + tile_y_ints + tile_x_ints) * static_cast<int>(sizeof(int));
}

static inline int makepad_cuda_nvfp4_mmq_max_shared_bytes() {
    int device = 0;
    if (cudaGetDevice(&device) != cudaSuccess) {
        return 48 * 1024;
    }

    int max_shared_optin = 0;
    if (cudaDeviceGetAttribute(&max_shared_optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, device) == cudaSuccess &&
        max_shared_optin > 0) {
        return max_shared_optin;
    }

    int max_shared = 0;
    if (cudaDeviceGetAttribute(&max_shared, cudaDevAttrMaxSharedMemoryPerBlock, device) == cudaSuccess &&
        max_shared > 0) {
        return max_shared;
    }

    return 48 * 1024;
}

static inline int makepad_cuda_select_mmq_x(
        uint32_t input_rows,
        int max_shared_bytes) {
    int mmq_x_best = 0;
    int ntiles_x_best = INT_MAX;

    for (int mmq_x = 8; mmq_x <= 128 && ntiles_x_best > 1; mmq_x += 8) {
        const int granularity = mmq_x >= 48 ? 16 : 8;
        if ((mmq_x % granularity) != 0) {
            continue;
        }

        if (makepad_cuda_nvfp4_mmq_shared_bytes(mmq_x) > max_shared_bytes) {
            continue;
        }

        const int ntiles_x = (static_cast<int>(input_rows) + mmq_x - 1) / mmq_x;
        if (ntiles_x < ntiles_x_best) {
            mmq_x_best = mmq_x;
            ntiles_x_best = ntiles_x;
        }
    }

    return mmq_x_best;
}

static inline cudaError_t makepad_cuda_mmq_sm_count(int * nsm) {
    int device = 0;
    cudaError_t status = cudaGetDevice(&device);
    if (status != cudaSuccess) {
        return status;
    }
    status = cudaDeviceGetAttribute(nsm, cudaDevAttrMultiProcessorCount, device);
    if (status != cudaSuccess) {
        return status;
    }
    return cudaSuccess;
}

template <int mmq_x>
static cudaError_t makepad_cuda_launch_nvfp4_q8_1_mmq_matmul(
        const uint8_t * input_q8_1_mmq_bytes,
        const uint8_t * packed_weights_nvfp4_bytes,
        float * output_f32,
        float * tmp_fixup_f32,
        uint32_t tmp_fixup_f32_len,
        uint32_t ncols_x,
        uint32_t out_rows,
        uint32_t input_rows,
        cudaStream_t stream) {
    const dim3 block(WARP_SIZE, MMQ_NWARPS, 1);
    const int shared_bytes = makepad_cuda_nvfp4_mmq_shared_bytes(mmq_x);
    int nsm = 0;
    cudaError_t status = makepad_cuda_mmq_sm_count(&nsm);
    if (status != cudaSuccess) {
        return status;
    }
    if (nsm <= 0) {
        return cudaErrorInvalidDevice;
    }

    const uint32_t required_fixup_f32 = static_cast<uint32_t>(nsm) * mmq_x * MMQ_Y;
    if (tmp_fixup_f32 == nullptr || tmp_fixup_f32_len < required_fixup_f32) {
        return cudaErrorInvalidValue;
    }

    status = cudaFuncSetAttribute(
        makepad_cuda_nvfp4_q8_1_mmq_matmul_kernel<mmq_x, false>,
        cudaFuncAttributeMaxDynamicSharedMemorySize,
        shared_bytes);
    if (status != cudaSuccess) {
        return status;
    }

    status = cudaFuncSetAttribute(
        makepad_cuda_nvfp4_q8_1_mmq_matmul_kernel<mmq_x, true>,
        cudaFuncAttributeMaxDynamicSharedMemorySize,
        shared_bytes);
    if (status != cudaSuccess) {
        return status;
    }

    status = cudaFuncSetAttribute(
        makepad_cuda_nvfp4_q8_1_mmq_stream_k_kernel<mmq_x, false, false>,
        cudaFuncAttributeMaxDynamicSharedMemorySize,
        shared_bytes);
    if (status != cudaSuccess) {
        return status;
    }

    status = cudaFuncSetAttribute(
        makepad_cuda_nvfp4_q8_1_mmq_stream_k_kernel<mmq_x, false, true>,
        cudaFuncAttributeMaxDynamicSharedMemorySize,
        shared_bytes);
    if (status != cudaSuccess) {
        return status;
    }

    status = cudaFuncSetAttribute(
        makepad_cuda_nvfp4_q8_1_mmq_stream_k_kernel<mmq_x, true, false>,
        cudaFuncAttributeMaxDynamicSharedMemorySize,
        shared_bytes);
    if (status != cudaSuccess) {
        return status;
    }

    status = cudaFuncSetAttribute(
        makepad_cuda_nvfp4_q8_1_mmq_stream_k_kernel<mmq_x, true, true>,
        cudaFuncAttributeMaxDynamicSharedMemorySize,
        shared_bytes);
    if (status != cudaSuccess) {
        return status;
    }

    const dim3 grid(nsm, 1, 1);
    const int nty = (out_rows + MMQ_Y - 1) / MMQ_Y;
    const int ntx = (input_rows + mmq_x - 1) / mmq_x;
    const bool fixup_needed = (ntx * nty) % nsm != 0;
    const bool check_y = (input_rows % mmq_x) != 0;

    if ((out_rows % MMQ_Y) == 0) {
        if (check_y) {
            makepad_cuda_nvfp4_q8_1_mmq_stream_k_kernel<mmq_x, false, true><<<grid, block, shared_bytes, stream>>>(
                reinterpret_cast<const block_nvfp4 *>(packed_weights_nvfp4_bytes),
                reinterpret_cast<const int *>(input_q8_1_mmq_bytes),
                output_f32,
                tmp_fixup_f32,
                static_cast<int>(ncols_x),
                static_cast<int>(out_rows),
                static_cast<int>(input_rows));
        } else {
            makepad_cuda_nvfp4_q8_1_mmq_stream_k_kernel<mmq_x, false, false><<<grid, block, shared_bytes, stream>>>(
                reinterpret_cast<const block_nvfp4 *>(packed_weights_nvfp4_bytes),
                reinterpret_cast<const int *>(input_q8_1_mmq_bytes),
                output_f32,
                tmp_fixup_f32,
                static_cast<int>(ncols_x),
                static_cast<int>(out_rows),
                static_cast<int>(input_rows));
        }
        status = cudaGetLastError();
        if (status != cudaSuccess || !fixup_needed) {
            return status;
        }
        makepad_cuda_nvfp4_q8_1_mmq_stream_k_fixup_kernel<mmq_x, false><<<grid, block, 0, stream>>>(
            output_f32,
            tmp_fixup_f32,
            static_cast<int>(ncols_x),
            static_cast<int>(out_rows),
            static_cast<int>(input_rows));
    } else {
        if (check_y) {
            makepad_cuda_nvfp4_q8_1_mmq_stream_k_kernel<mmq_x, true, true><<<grid, block, shared_bytes, stream>>>(
                reinterpret_cast<const block_nvfp4 *>(packed_weights_nvfp4_bytes),
                reinterpret_cast<const int *>(input_q8_1_mmq_bytes),
                output_f32,
                tmp_fixup_f32,
                static_cast<int>(ncols_x),
                static_cast<int>(out_rows),
                static_cast<int>(input_rows));
        } else {
            makepad_cuda_nvfp4_q8_1_mmq_stream_k_kernel<mmq_x, true, false><<<grid, block, shared_bytes, stream>>>(
                reinterpret_cast<const block_nvfp4 *>(packed_weights_nvfp4_bytes),
                reinterpret_cast<const int *>(input_q8_1_mmq_bytes),
                output_f32,
                tmp_fixup_f32,
                static_cast<int>(ncols_x),
                static_cast<int>(out_rows),
                static_cast<int>(input_rows));
        }
        status = cudaGetLastError();
        if (status != cudaSuccess || !fixup_needed) {
            return status;
        }
        makepad_cuda_nvfp4_q8_1_mmq_stream_k_fixup_kernel<mmq_x, true><<<grid, block, 0, stream>>>(
            output_f32,
            tmp_fixup_f32,
            static_cast<int>(ncols_x),
            static_cast<int>(out_rows),
            static_cast<int>(input_rows));
    }

    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_nvfp4_q8_1_mmq_fixup_f32_len(
        uint32_t * len_out) {
    if (len_out == nullptr) {
        return cudaErrorInvalidValue;
    }
    int nsm = 0;
    cudaError_t status = makepad_cuda_mmq_sm_count(&nsm);
    if (status != cudaSuccess) {
        return status;
    }
    if (nsm <= 0) {
        return cudaErrorInvalidDevice;
    }
    *len_out = static_cast<uint32_t>(nsm) * 128u * static_cast<uint32_t>(MMQ_Y);
    return cudaSuccess;
}

extern "C" cudaError_t makepad_cuda_quantize_q8_1_mmq_f32(
        const float * input_f32,
        uint8_t * output_q8_1_mmq_bytes,
        uint32_t n_cols,
        uint32_t n_rows,
        cudaStream_t stream);

extern "C" cudaError_t makepad_cuda_quantize_q8_1_mmq_f32_padded(
        const float * input_f32,
        uint8_t * output_q8_1_mmq_bytes,
        uint32_t n_cols,
        uint32_t n_rows,
        uint32_t padded_rows,
        cudaStream_t stream) {
    if (n_cols == 0 || n_rows == 0 || padded_rows < n_rows || (n_cols % (4 * QK8_1)) != 0) {
        return cudaErrorInvalidValue;
    }

    const uint32_t block_num_y =
        (n_cols + 4 * CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / (4 * CUDA_QUANTIZE_BLOCK_SIZE_MMQ);
    const dim3 grid(padded_rows, block_num_y, 1);
    const dim3 block(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1);
    makepad_cuda_quantize_q8_1_mmq_f32_kernel<<<grid, block, 0, stream>>>(
        input_f32,
        reinterpret_cast<block_q8_1_mmq *>(output_q8_1_mmq_bytes),
        n_cols,
        n_rows,
        padded_rows);
    return cudaGetLastError();
}

extern "C" cudaError_t makepad_cuda_quantize_q8_1_mmq_f32(
        const float * input_f32,
        uint8_t * output_q8_1_mmq_bytes,
        uint32_t n_cols,
        uint32_t n_rows,
        cudaStream_t stream) {
    return makepad_cuda_quantize_q8_1_mmq_f32_padded(
        input_f32,
        output_q8_1_mmq_bytes,
        n_cols,
        n_rows,
        n_rows,
        stream);
}

extern "C" cudaError_t makepad_cuda_nvfp4_q8_1_mmq_matmul(
        const uint8_t * input_q8_1_mmq_bytes,
        const uint8_t * packed_weights_nvfp4_bytes,
        float * output_f32,
        float * tmp_fixup_f32,
        uint32_t tmp_fixup_f32_len,
        uint32_t n_cols,
        uint32_t out_rows,
        uint32_t input_rows,
        cudaStream_t stream) {
    if (n_cols == 0 || out_rows == 0 || input_rows == 0 || (n_cols % (4 * QK8_1)) != 0) {
        return cudaErrorInvalidValue;
    }

    const int max_shared_bytes = makepad_cuda_nvfp4_mmq_max_shared_bytes();
    switch (makepad_cuda_select_mmq_x(input_rows, max_shared_bytes)) {
        case 8:
            return makepad_cuda_launch_nvfp4_q8_1_mmq_matmul<8>(
                input_q8_1_mmq_bytes,
                packed_weights_nvfp4_bytes,
                output_f32,
                tmp_fixup_f32,
                tmp_fixup_f32_len,
                n_cols,
                out_rows,
                input_rows,
                stream);
        case 16:
            return makepad_cuda_launch_nvfp4_q8_1_mmq_matmul<16>(
                input_q8_1_mmq_bytes,
                packed_weights_nvfp4_bytes,
                output_f32,
                tmp_fixup_f32,
                tmp_fixup_f32_len,
                n_cols,
                out_rows,
                input_rows,
                stream);
        case 24:
            return makepad_cuda_launch_nvfp4_q8_1_mmq_matmul<24>(
                input_q8_1_mmq_bytes,
                packed_weights_nvfp4_bytes,
                output_f32,
                tmp_fixup_f32,
                tmp_fixup_f32_len,
                n_cols,
                out_rows,
                input_rows,
                stream);
        case 32:
            return makepad_cuda_launch_nvfp4_q8_1_mmq_matmul<32>(
                input_q8_1_mmq_bytes,
                packed_weights_nvfp4_bytes,
                output_f32,
                tmp_fixup_f32,
                tmp_fixup_f32_len,
                n_cols,
                out_rows,
                input_rows,
                stream);
        case 40:
            return makepad_cuda_launch_nvfp4_q8_1_mmq_matmul<40>(
                input_q8_1_mmq_bytes,
                packed_weights_nvfp4_bytes,
                output_f32,
                tmp_fixup_f32,
                tmp_fixup_f32_len,
                n_cols,
                out_rows,
                input_rows,
                stream);
        case 48:
            return makepad_cuda_launch_nvfp4_q8_1_mmq_matmul<48>(
                input_q8_1_mmq_bytes,
                packed_weights_nvfp4_bytes,
                output_f32,
                tmp_fixup_f32,
                tmp_fixup_f32_len,
                n_cols,
                out_rows,
                input_rows,
                stream);
        case 56:
            return makepad_cuda_launch_nvfp4_q8_1_mmq_matmul<56>(
                input_q8_1_mmq_bytes,
                packed_weights_nvfp4_bytes,
                output_f32,
                tmp_fixup_f32,
                tmp_fixup_f32_len,
                n_cols,
                out_rows,
                input_rows,
                stream);
        case 64:
            return makepad_cuda_launch_nvfp4_q8_1_mmq_matmul<64>(
                input_q8_1_mmq_bytes,
                packed_weights_nvfp4_bytes,
                output_f32,
                tmp_fixup_f32,
                tmp_fixup_f32_len,
                n_cols,
                out_rows,
                input_rows,
                stream);
        case 72:
            return makepad_cuda_launch_nvfp4_q8_1_mmq_matmul<72>(
                input_q8_1_mmq_bytes,
                packed_weights_nvfp4_bytes,
                output_f32,
                tmp_fixup_f32,
                tmp_fixup_f32_len,
                n_cols,
                out_rows,
                input_rows,
                stream);
        case 80:
            return makepad_cuda_launch_nvfp4_q8_1_mmq_matmul<80>(
                input_q8_1_mmq_bytes,
                packed_weights_nvfp4_bytes,
                output_f32,
                tmp_fixup_f32,
                tmp_fixup_f32_len,
                n_cols,
                out_rows,
                input_rows,
                stream);
        case 88:
            return makepad_cuda_launch_nvfp4_q8_1_mmq_matmul<88>(
                input_q8_1_mmq_bytes,
                packed_weights_nvfp4_bytes,
                output_f32,
                tmp_fixup_f32,
                tmp_fixup_f32_len,
                n_cols,
                out_rows,
                input_rows,
                stream);
        case 96:
            return makepad_cuda_launch_nvfp4_q8_1_mmq_matmul<96>(
                input_q8_1_mmq_bytes,
                packed_weights_nvfp4_bytes,
                output_f32,
                tmp_fixup_f32,
                tmp_fixup_f32_len,
                n_cols,
                out_rows,
                input_rows,
                stream);
        case 104:
            return makepad_cuda_launch_nvfp4_q8_1_mmq_matmul<104>(
                input_q8_1_mmq_bytes,
                packed_weights_nvfp4_bytes,
                output_f32,
                tmp_fixup_f32,
                tmp_fixup_f32_len,
                n_cols,
                out_rows,
                input_rows,
                stream);
        case 112:
            return makepad_cuda_launch_nvfp4_q8_1_mmq_matmul<112>(
                input_q8_1_mmq_bytes,
                packed_weights_nvfp4_bytes,
                output_f32,
                tmp_fixup_f32,
                tmp_fixup_f32_len,
                n_cols,
                out_rows,
                input_rows,
                stream);
        case 120:
            return makepad_cuda_launch_nvfp4_q8_1_mmq_matmul<120>(
                input_q8_1_mmq_bytes,
                packed_weights_nvfp4_bytes,
                output_f32,
                tmp_fixup_f32,
                tmp_fixup_f32_len,
                n_cols,
                out_rows,
                input_rows,
                stream);
        case 128:
            return makepad_cuda_launch_nvfp4_q8_1_mmq_matmul<128>(
                input_q8_1_mmq_bytes,
                packed_weights_nvfp4_bytes,
                output_f32,
                tmp_fixup_f32,
                tmp_fixup_f32_len,
                n_cols,
                out_rows,
                input_rows,
                stream);
        default:
            return cudaErrorInvalidValue;
    }
}