submission 113511

#!POPCORN leaderboard nvfp4_gemv

import torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t


# CUDA SOURCE CODE

cuda_source = """
#include <cuda_fp4.h>
#include <cuda_fp8.h>
#include <cuda_fp16.h>
#include <cuda_pipeline.h>
#include <cuda/ptx>
#include<cuda_awbarrier_primitives.h>

namespace ptx = cuda::ptx;


#define FULL_MASK 0xffffffff

__inline__ __device__ void multiply_and_accumulate(
    int4 a_packed,
    int4 b_packed,
    __nv_fp8x2_storage_t sfa_fp8x2,
    __nv_fp8x2_storage_t sfb_fp8x2,
    int* result_0,
    int* result_1,
    int* result_2,
    int* result_3
) {
    asm volatile( \\
        "{\\n" \\
        // declare registers for A / B tensors
        ".reg .b8 byte0_0, byte0_1, byte0_2, byte0_3;\\n" \\
        ".reg .b8 byte0_4, byte0_5, byte0_6, byte0_7;\\n" \\
        ".reg .b8 byte1_0, byte1_1, byte1_2, byte1_3;\\n" \\
        ".reg .b8 byte1_4, byte1_5, byte1_6, byte1_7;\\n" \\
        ".reg .b8 byte2_0, byte2_1, byte2_2, byte2_3;\\n" \\
        ".reg .b8 byte2_4, byte2_5, byte2_6, byte2_7;\\n" \\
        ".reg .b8 byte3_0, byte3_1, byte3_2, byte3_3;\\n" \\
        ".reg .b8 byte3_4, byte3_5, byte3_6, byte3_7;\\n" \\

        // declare registers for accumulators
        ".reg .f16x2 accum_0_0, accum_0_1, accum_0_2, accum_0_3;\\n" \\
        ".reg .f16x2 accum_1_0, accum_1_1, accum_1_2, accum_1_3;\\n" \\
        ".reg .f16x2 accum_2_0, accum_2_1, accum_2_2, accum_2_3;\\n" \\
        ".reg .f16x2 accum_3_0, accum_3_1, accum_3_2, accum_3_3;\\n" \\

        // declare registers for scaling factors
        ".reg .f16x2 sfa_f16x2;\\n" \\
        ".reg .f16x2 sfb_f16x2;\\n" \\
        ".reg .f16x2 sf_f16x2;\\n" \\
        
        // declare registers for conversion
        ".reg .f16x2 cvt_0_0, cvt_0_1, cvt_0_2, cvt_0_3;\\n" \\
        ".reg .f16x2 cvt_0_4, cvt_0_5, cvt_0_6, cvt_0_7;\\n" \\
        ".reg .f16x2 cvt_1_0, cvt_1_1, cvt_1_2, cvt_1_3;\\n" \\
        ".reg .f16x2 cvt_1_4, cvt_1_5, cvt_1_6, cvt_1_7;\\n" \\
        ".reg .f16x2 cvt_2_0, cvt_2_1, cvt_2_2, cvt_2_3;\\n" \\
        ".reg .f16x2 cvt_2_4, cvt_2_5, cvt_2_6, cvt_2_7;\\n" \\
        ".reg .f16x2 cvt_3_0, cvt_3_1, cvt_3_2, cvt_3_3;\\n" \\
        ".reg .f16x2 cvt_3_4, cvt_3_5, cvt_3_6, cvt_3_7;\\n" \\
        ".reg .f16 result_f16, lane0, lane1;\\n" \\
        ".reg .f16x2 mul_f16x2_0, mul_f16x2_1;\\n" \\

        // convert scaling factors from fp8 to f16x2
        "cvt.rn.f16x2.e4m3x2 sfa_f16x2, %4;\\n" \\
        "cvt.rn.f16x2.e4m3x2 sfb_f16x2, %5;\\n" \\
        
        // clear accumulators
        "mov.b32 accum_0_0, 0;\\n" \\
        "mov.b32 accum_0_1, 0;\\n" \\
        "mov.b32 accum_0_2, 0;\\n" \\
        "mov.b32 accum_0_3, 0;\\n" \\
        "mov.b32 accum_1_0, 0;\\n" \\
        "mov.b32 accum_1_1, 0;\\n" \\
        "mov.b32 accum_1_2, 0;\\n" \\
        "mov.b32 accum_1_3, 0;\\n" \\
        "mov.b32 accum_2_0, 0;\\n" \\
        "mov.b32 accum_2_1, 0;\\n" \\
        "mov.b32 accum_2_2, 0;\\n" \\
        "mov.b32 accum_2_3, 0;\\n" \\
        "mov.b32 accum_3_0, 0;\\n" \\
        "mov.b32 accum_3_1, 0;\\n" \\
        "mov.b32 accum_3_2, 0;\\n" \\
        "mov.b32 accum_3_3, 0;\\n" \\
        
        // multiply, unpacking and permuting scale factors
        "mul.rn.f16x2 sf_f16x2, sfa_f16x2, sfb_f16x2;\\n" \\
        "mov.b32 {lane0, lane1}, sf_f16x2;\\n" \\
        "mov.b32 mul_f16x2_0, {lane0, lane0};\\n" \\
        "mov.b32 mul_f16x2_1, {lane1, lane1};\\n" \\

        // unpacking A and B tensors
        "mov.b32 {byte0_0, byte0_1, byte0_2, byte0_3}, %6;\\n" \\
        "mov.b32 {byte0_4, byte0_5, byte0_6, byte0_7}, %7;\\n" \\
        "mov.b32 {byte1_0, byte1_1, byte1_2, byte1_3}, %8;\\n" \\
        "mov.b32 {byte1_4, byte1_5, byte1_6, byte1_7}, %9;\\n" \\
        "mov.b32 {byte2_0, byte2_1, byte2_2, byte2_3}, %10;\\n" \\
        "mov.b32 {byte2_4, byte2_5, byte2_6, byte2_7}, %11;\\n" \\
        "mov.b32 {byte3_0, byte3_1, byte3_2, byte3_3}, %12;\\n" \\
        "mov.b32 {byte3_4, byte3_5, byte3_6, byte3_7}, %13;\\n" \\

        // convert A and B tensors from fp4 to f16x2

        // A[0 - 7] and B[0 - 7]
        "cvt.rn.f16x2.e2m1x2 cvt_0_0, byte0_0;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_0_1, byte0_1;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_0_2, byte0_2;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_0_3, byte0_3;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_0_4, byte0_4;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_0_5, byte0_5;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_0_6, byte0_6;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_0_7, byte0_7;\\n" \\

        // A[8 - 15] and B[8 - 15]
        "cvt.rn.f16x2.e2m1x2 cvt_1_0, byte1_0;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_1_1, byte1_1;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_1_2, byte1_2;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_1_3, byte1_3;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_1_4, byte1_4;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_1_5, byte1_5;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_1_6, byte1_6;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_1_7, byte1_7;\\n" \\

        // A[16 - 23] and B[16 - 23]
        "cvt.rn.f16x2.e2m1x2 cvt_2_0, byte2_0;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_2_1, byte2_1;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_2_2, byte2_2;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_2_3, byte2_3;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_2_4, byte2_4;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_2_5, byte2_5;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_2_6, byte2_6;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_2_7, byte2_7;\\n" \\

        // A[24 - 31] and B[24 - 31]
        "cvt.rn.f16x2.e2m1x2 cvt_3_0, byte3_0;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_3_1, byte3_1;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_3_2, byte3_2;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_3_3, byte3_3;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_3_4, byte3_4;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_3_5, byte3_5;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_3_6, byte3_6;\\n" \\
        "cvt.rn.f16x2.e2m1x2 cvt_3_7, byte3_7;\\n" \\

        // fma for A[0 - 7] and B[0 - 7]
        "fma.rn.f16x2 accum_0_0, cvt_0_0, cvt_0_4, accum_0_0;\\n" \\
        "fma.rn.f16x2 accum_0_1, cvt_0_1, cvt_0_5, accum_0_1;\\n" \\
        "fma.rn.f16x2 accum_0_2, cvt_0_2, cvt_0_6, accum_0_2;\\n" \\
        "fma.rn.f16x2 accum_0_3, cvt_0_3, cvt_0_7, accum_0_3;\\n" \\

        // fma for A[8 - 15] and B[8 - 15]
        "fma.rn.f16x2 accum_1_0, cvt_1_0, cvt_1_4, accum_1_0;\\n" \\
        "fma.rn.f16x2 accum_1_1, cvt_1_1, cvt_1_5, accum_1_1;\\n" \\
        "fma.rn.f16x2 accum_1_2, cvt_1_2, cvt_1_6, accum_1_2;\\n" \\
        "fma.rn.f16x2 accum_1_3, cvt_1_3, cvt_1_7, accum_1_3;\\n" \\

        // fma for A[16 - 23] and B[16 - 23]
        "fma.rn.f16x2 accum_2_0, cvt_2_0, cvt_2_4, accum_2_0;\\n" \\
        "fma.rn.f16x2 accum_2_1, cvt_2_1, cvt_2_5, accum_2_1;\\n" \\
        "fma.rn.f16x2 accum_2_2, cvt_2_2, cvt_2_6, accum_2_2;\\n" \\
        "fma.rn.f16x2 accum_2_3, cvt_2_3, cvt_2_7, accum_2_3;\\n" \\

        // fma for A[24 - 31] and B[24 - 31]
        "fma.rn.f16x2 accum_3_0, cvt_3_0, cvt_3_4, accum_3_0;\\n" \\
        "fma.rn.f16x2 accum_3_1, cvt_3_1, cvt_3_5, accum_3_1;\\n" \\
        "fma.rn.f16x2 accum_3_2, cvt_3_2, cvt_3_6, accum_3_2;\\n" \\
        "fma.rn.f16x2 accum_3_3, cvt_3_3, cvt_3_7, accum_3_3;\\n" \\

        // tree reduction for accumulators
        "add.rn.f16x2 accum_0_0, accum_0_0, accum_0_1;\\n" \\
        "add.rn.f16x2 accum_0_2, accum_0_2, accum_0_3;\\n" \\
        "add.rn.f16x2 accum_1_0, accum_1_0, accum_1_1;\\n" \\
        "add.rn.f16x2 accum_1_2, accum_1_2, accum_1_3;\\n" \\
        "add.rn.f16x2 accum_2_0, accum_2_0, accum_2_1;\\n" \\
        "add.rn.f16x2 accum_2_2, accum_2_2, accum_2_3;\\n" \\
        "add.rn.f16x2 accum_3_0, accum_3_0, accum_3_1;\\n" \\
        "add.rn.f16x2 accum_3_2, accum_3_2, accum_3_3;\\n" \\

        "fma.rn.f16x2 %0, accum_0_0, mul_f16x2_0, %0;\\n" \\
        "fma.rn.f16x2 %1, accum_0_2, mul_f16x2_0, %1;\\n" \\
        "fma.rn.f16x2 %2, accum_1_0, mul_f16x2_0, %2;\\n" \\
        "fma.rn.f16x2 %3, accum_1_2, mul_f16x2_0, %3;\\n" \\
        

        "fma.rn.f16x2 %0, accum_2_0, mul_f16x2_1, %0;\\n" \\
        "fma.rn.f16x2 %1, accum_2_2, mul_f16x2_1, %1;\\n" \\
        "fma.rn.f16x2 %2, accum_3_0, mul_f16x2_1, %2;\\n" \\
        "fma.rn.f16x2 %3, accum_3_2, mul_f16x2_1, %3;\\n" \\

        "}\\n"
        : "+r"(*result_0), "+r"(*result_1), "+r"(*result_2), "+r"(*result_3)    // 0, 1, 2, 3
        : "h"(sfa_fp8x2), "h"(sfb_fp8x2),                   // 4, 5
            "r"(a_packed.x), "r"(b_packed.x),               // 6, 7
            "r"(a_packed.y), "r"(b_packed.y),               // 8, 9
            "r"(a_packed.z), "r"(b_packed.z),               // 10, 11
            "r"(a_packed.w), "r"(b_packed.w)                // 12, 13
    );
}


__global__ void gemv_kernel_4096_7168(
    const __nv_fp4x2_storage_t* __restrict__ a,
    const __nv_fp4x2_storage_t* __restrict__ b,
    const __nv_fp8_e4m3* __restrict__ sfa,
    const __nv_fp8_e4m3* __restrict__ sfb,
    __half* __restrict__ c
) {
    const int M = 4096;
    const int K = 7168;

    extern __shared__ unsigned char shared_storage[];
    auto* b_shared = reinterpret_cast<__nv_fp4x2_storage_t*>(shared_storage);
    auto* sfb_shared = reinterpret_cast<__nv_fp8_e4m3*>(b_shared + (K / 2));
    __shared__ __half c_shared[32];

    b += blockIdx.y * (K / 2) * 128;
    sfb += blockIdx.y * (K / 16) * 128;

    for (int i = threadIdx.y * 32 + threadIdx.x; i < K / 32; i += blockDim.y * blockDim.x) {
        reinterpret_cast<int4*>(b_shared)[i] = reinterpret_cast<const int4*>(b)[i];
    }
    for (int i = threadIdx.y * 32 + threadIdx.x; i < K / 256; i += blockDim.y * blockDim.x) {
        reinterpret_cast<int4*>(sfb_shared)[i] = reinterpret_cast<const int4*>(sfb)[i];
    }
    __syncthreads();

    // Each warp computes one result and saves it to shared memory
    int result_0 = 0;
    int result_1 = 0;
    int result_2 = 0;
    int result_3 = 0;
    int offset = blockIdx.y * (K * M / 2) + (blockIdx.x * 32 + threadIdx.y) * (K / 2);
    a += offset;
    sfa += offset / 8;
    
    for (int i = threadIdx.x; i < K / 32; i += 32) {
        int4 a_packed = reinterpret_cast<const int4*>(a)[i];
        int4 b_packed = reinterpret_cast<int4*>(b_shared)[i];
        
        __nv_fp8x2_storage_t sfa_fp8x2 = reinterpret_cast<const __nv_fp8x2_storage_t*>(sfa)[i];
        __nv_fp8x2_storage_t sfb_fp8x2 = reinterpret_cast<__nv_fp8x2_storage_t*>(sfb_shared)[i];

        multiply_and_accumulate(a_packed, b_packed, sfa_fp8x2, sfb_fp8x2, &result_0, &result_1, &result_2, &result_3);
    }


    // Reduce the result and store it in shared memory
    __half2 reduction_result_0 = __hadd2(reinterpret_cast<const __half2&>(result_0),
            reinterpret_cast<const __half2&>(result_1));
    __half2 reduction_result_1 = __hadd2(reinterpret_cast<const __half2&>(result_2),
            reinterpret_cast<const __half2&>(result_3));
    reduction_result_0 = __hadd2(reduction_result_0, reduction_result_1);
    float final_result_f = __half22float2(reduction_result_0).x + __half22float2(reduction_result_0).y;
    for (int offset = 16; offset > 0; offset /= 2) {
        final_result_f += __shfl_down_sync(FULL_MASK, final_result_f, offset);
    }
    if (threadIdx.x == 0) {
        int c_offset = blockIdx.y * M + blockIdx.x * 32 + threadIdx.y;
        c[c_offset] = __float2half_rn(final_result_f);
    }
}


__global__ void gemv_kernel_7168_2048(
    const __nv_fp4x2_storage_t* __restrict__ a,
    const __nv_fp4x2_storage_t* __restrict__ b,
    const __nv_fp8_e4m3* __restrict__ sfa,
    const __nv_fp8_e4m3* __restrict__ sfb,
    __half* __restrict__ c
) {
    const int M = 7168;
    const int K = 2048;

    extern __shared__ unsigned char shared_storage[];
    auto* b_shared = reinterpret_cast<__nv_fp4x2_storage_t*>(shared_storage);
    auto* sfb_shared = reinterpret_cast<__nv_fp8_e4m3*>(b_shared + (K / 2));
    __shared__ __half c_shared[32];

    b += blockIdx.y * (K / 2) * 128;
    sfb += blockIdx.y * (K / 16) * 128;

    for (int i = threadIdx.y * 32 + threadIdx.x; i < K / 32; i += blockDim.y * blockDim.x) {
        reinterpret_cast<int4*>(b_shared)[i] = reinterpret_cast<const int4*>(b)[i];
    }
    for (int i = threadIdx.y * 32 + threadIdx.x; i < K / 256; i += blockDim.y * blockDim.x) {
        reinterpret_cast<int4*>(sfb_shared)[i] = reinterpret_cast<const int4*>(sfb)[i];
    }
    __syncthreads();

    // Each warp computes one result and saves it to shared memory
    int result_0 = 0;
    int result_1 = 0;
    int result_2 = 0;
    int result_3 = 0;
    int offset = blockIdx.y * (K * M / 2) + (blockIdx.x * 32 + threadIdx.y) * (K / 2);
    a += offset;
    sfa += offset / 8;
    
    for (int i = threadIdx.x; i < K / 32; i += 32) {
        int4 a_packed = reinterpret_cast<const int4*>(a)[i];
        int4 b_packed = reinterpret_cast<int4*>(b_shared)[i];
        
        __nv_fp8x2_storage_t sfa_fp8x2 = reinterpret_cast<const __nv_fp8x2_storage_t*>(sfa)[i];
        __nv_fp8x2_storage_t sfb_fp8x2 = reinterpret_cast<__nv_fp8x2_storage_t*>(sfb_shared)[i];

        multiply_and_accumulate(a_packed, b_packed, sfa_fp8x2, sfb_fp8x2, &result_0, &result_1, &result_2, &result_3);
    }


    // Reduce the result and store it in shared memory
    __half2 reduction_result_0 = __hadd2(reinterpret_cast<const __half2&>(result_0),
            reinterpret_cast<const __half2&>(result_1));
    __half2 reduction_result_1 = __hadd2(reinterpret_cast<const __half2&>(result_2),
            reinterpret_cast<const __half2&>(result_3));
    reduction_result_0 = __hadd2(reduction_result_0, reduction_result_1);
    float final_result_f = __half22float2(reduction_result_0).x + __half22float2(reduction_result_0).y;
    for (int offset = 16; offset > 0; offset /= 2) {
        final_result_f += __shfl_down_sync(FULL_MASK, final_result_f, offset);
    }
    if (threadIdx.x == 0) {
        int c_offset = blockIdx.y * M + blockIdx.x * 32 + threadIdx.y;
        c[c_offset] = __float2half_rn(final_result_f);
    }
}



__global__ void
__launch_bounds__(832)
gemv_kernel_7168_16384(
    const int4* __restrict__ a,
    const int4* __restrict__ b,
    const int* __restrict__ sfa,
    const int* __restrict__ sfb,
    __half* __restrict__ c
) {
    const int M = 7168;
    const int K = 16384;
    const int Q_SIZE = 2;
    const int active_warps = (blockIdx.x < 32) ? 26 : 25;

    __shared__ int4 a_shared[Q_SIZE + 1][25][2][32];
    __shared__ int sfa_shared[Q_SIZE + 1][25][32];

    // We will load all b and sfb, because we can, it simplifies the logic
    __shared__ int4 b_shared[8][2][32];
    __shared__ int sfb_shared[8][32];

    __shared__ __mbarrier_t bar[8];
    if (threadIdx.y == 0 && threadIdx.x == 0) {
        #pragma unroll
        for (int i = 0; i < 8; i++) {
            __mbarrier_init(&bar[i], 32);
        }
    }
    __syncthreads();

    if (threadIdx.y == 0) {
        // ========== WARP 0: Load b and sfb for all columns ==========
        #pragma unroll
        for (int col_idx = 0; col_idx < 8; col_idx++) {
            __pipeline_memcpy_async(&b_shared[col_idx][0][threadIdx.x], &b[col_idx * 64 + threadIdx.x], sizeof(int4));
            __pipeline_memcpy_async(&b_shared[col_idx][1][threadIdx.x], &b[col_idx * 64 + 32 + threadIdx.x], sizeof(int4));
            __pipeline_memcpy_async(&sfb_shared[col_idx][threadIdx.x], &sfb[col_idx * 32 + threadIdx.x], sizeof(int));
            __pipeline_arrive_on(&bar[col_idx]);
            __mbarrier_arrive(&bar[col_idx]);
        }
    } else if (threadIdx.y < active_warps) {
        // ========== COMPUTE WARPS: Load a/sfa and compute ==========
        int offset = (blockIdx.x * 24 + min(blockIdx.x, 32) + threadIdx.y - 1) * 2 * (K / 2);
        a += offset / 16;
        sfa += offset / 32;

        int result[2][4] = {0};

        // Prologue: prefetch col 0 (both rows)
        __pipeline_memcpy_async(&a_shared[0][threadIdx.y - 1][0][threadIdx.x], &a[0 * (K / 32) + 0 * 64 + threadIdx.x], sizeof(int4));
        __pipeline_memcpy_async(&a_shared[0][threadIdx.y - 1][1][threadIdx.x], &a[0 * (K / 32) + 0 * 64 + 32 + threadIdx.x], sizeof(int4));
        __pipeline_memcpy_async(&sfa_shared[0][threadIdx.y - 1][threadIdx.x], &sfa[0 * (K / 64) + 0 * 32 + threadIdx.x], sizeof(int));
        __pipeline_commit();
        __pipeline_memcpy_async(&a_shared[1][threadIdx.y - 1][0][threadIdx.x], &a[1 * (K / 32) + 0 * 64 + threadIdx.x], sizeof(int4));
        __pipeline_memcpy_async(&a_shared[1][threadIdx.y - 1][1][threadIdx.x], &a[1 * (K / 32) + 0 * 64 + 32 + threadIdx.x], sizeof(int4));
        __pipeline_memcpy_async(&sfa_shared[1][threadIdx.y - 1][threadIdx.x], &sfa[1 * (K / 64) + 0 * 32 + threadIdx.x], sizeof(int));
        __pipeline_commit();

        // Main loop: process columns 0-6, prefetch next column
        #pragma unroll
        for (int col_idx = 0; col_idx < 7; col_idx++) {
            int next_col = col_idx + 1;

            // Prefetch row 0 for next column
            __pipeline_memcpy_async(&a_shared[(col_idx * 2 + 2) % (Q_SIZE + 1)][threadIdx.y - 1][0][threadIdx.x], &a[0 * (K / 32) + next_col * 64 + threadIdx.x], sizeof(int4));
            __pipeline_memcpy_async(&a_shared[(col_idx * 2 + 2) % (Q_SIZE + 1)][threadIdx.y - 1][1][threadIdx.x], &a[0 * (K / 32) + next_col * 64 + 32 + threadIdx.x], sizeof(int4));
            __pipeline_memcpy_async(&sfa_shared[(col_idx * 2 + 2) % (Q_SIZE + 1)][threadIdx.y - 1][threadIdx.x], &sfa[0 * (K / 64) + next_col * 32 + threadIdx.x], sizeof(int));
            __pipeline_commit();

            // Wait for b/sfb data, load once for this column
            while (!ptx::mbarrier_try_wait_parity(&bar[col_idx], 0)) {}
            int4 b_packed_0 = b_shared[col_idx][0][threadIdx.x];
            int4 b_packed_1 = b_shared[col_idx][1][threadIdx.x];
            __nv_fp8x2_storage_t sfb_fp8x2_0 = reinterpret_cast<__nv_fp8x2_storage_t*>(sfb_shared[col_idx])[threadIdx.x];
            __nv_fp8x2_storage_t sfb_fp8x2_1 = reinterpret_cast<__nv_fp8x2_storage_t*>(sfb_shared[col_idx])[threadIdx.x + 32];

            // Load and compute row 0
            __pipeline_wait_prior(Q_SIZE);
            int4 a_packed_r0_0 = a_shared[(col_idx * 2) % (Q_SIZE + 1)][threadIdx.y - 1][0][threadIdx.x];
            int4 a_packed_r0_1 = a_shared[(col_idx * 2) % (Q_SIZE + 1)][threadIdx.y - 1][1][threadIdx.x];
            __nv_fp8x2_storage_t sfa_fp8x2_r0_0 = reinterpret_cast<__nv_fp8x2_storage_t*>(sfa_shared[(col_idx * 2) % (Q_SIZE + 1)][threadIdx.y - 1])[threadIdx.x];
            __nv_fp8x2_storage_t sfa_fp8x2_r0_1 = reinterpret_cast<__nv_fp8x2_storage_t*>(sfa_shared[(col_idx * 2) % (Q_SIZE + 1)][threadIdx.y - 1])[threadIdx.x + 32];
            multiply_and_accumulate(a_packed_r0_0, b_packed_0, sfa_fp8x2_r0_0, sfb_fp8x2_0, &result[0][0], &result[0][1], &result[0][2], &result[0][3]);
            multiply_and_accumulate(a_packed_r0_1, b_packed_1, sfa_fp8x2_r0_1, sfb_fp8x2_1, &result[0][0], &result[0][1], &result[0][2], &result[0][3]);

            // Prefetch row 1 for next column
            __pipeline_memcpy_async(&a_shared[(col_idx * 2 + 3) % (Q_SIZE + 1)][threadIdx.y - 1][0][threadIdx.x], &a[1 * (K / 32) + next_col * 64 + threadIdx.x], sizeof(int4));
            __pipeline_memcpy_async(&a_shared[(col_idx * 2 + 3) % (Q_SIZE + 1)][threadIdx.y - 1][1][threadIdx.x], &a[1 * (K / 32) + next_col * 64 + 32 + threadIdx.x], sizeof(int4));
            __pipeline_memcpy_async(&sfa_shared[(col_idx * 2 + 3) % (Q_SIZE + 1)][threadIdx.y - 1][threadIdx.x], &sfa[1 * (K / 64) + next_col * 32 + threadIdx.x], sizeof(int));
            __pipeline_commit();

            // Load and compute row 1
            __pipeline_wait_prior(Q_SIZE);
            int4 a_packed_r1_0 = a_shared[(col_idx * 2 + 1) % (Q_SIZE + 1)][threadIdx.y - 1][0][threadIdx.x];
            int4 a_packed_r1_1 = a_shared[(col_idx * 2 + 1) % (Q_SIZE + 1)][threadIdx.y - 1][1][threadIdx.x];
            __nv_fp8x2_storage_t sfa_fp8x2_r1_0 = reinterpret_cast<__nv_fp8x2_storage_t*>(sfa_shared[(col_idx * 2 + 1) % (Q_SIZE + 1)][threadIdx.y - 1])[threadIdx.x];
            __nv_fp8x2_storage_t sfa_fp8x2_r1_1 = reinterpret_cast<__nv_fp8x2_storage_t*>(sfa_shared[(col_idx * 2 + 1) % (Q_SIZE + 1)][threadIdx.y - 1])[threadIdx.x + 32];
            multiply_and_accumulate(a_packed_r1_0, b_packed_0, sfa_fp8x2_r1_0, sfb_fp8x2_0, &result[1][0], &result[1][1], &result[1][2], &result[1][3]);
            multiply_and_accumulate(a_packed_r1_1, b_packed_1, sfa_fp8x2_r1_1, sfb_fp8x2_1, &result[1][0], &result[1][1], &result[1][2], &result[1][3]);
        }

        // Epilogue: process last column (col_idx = 7)
        {
            const int col_idx = 7;

            // Wait for b/sfb data, load once
            while (!ptx::mbarrier_try_wait_parity(&bar[col_idx], 0)) {}
            int4 b_packed_0 = b_shared[col_idx][0][threadIdx.x];
            int4 b_packed_1 = b_shared[col_idx][1][threadIdx.x];
            __nv_fp8x2_storage_t sfb_fp8x2_0 = reinterpret_cast<__nv_fp8x2_storage_t*>(sfb_shared[col_idx])[threadIdx.x];
            __nv_fp8x2_storage_t sfb_fp8x2_1 = reinterpret_cast<__nv_fp8x2_storage_t*>(sfb_shared[col_idx])[threadIdx.x + 32];

            // Load and compute row 0
            __pipeline_wait_prior(1);
            int4 a_packed_r0_0 = a_shared[(col_idx * 2) % (Q_SIZE + 1)][threadIdx.y - 1][0][threadIdx.x];
            int4 a_packed_r0_1 = a_shared[(col_idx * 2) % (Q_SIZE + 1)][threadIdx.y - 1][1][threadIdx.x];
            __nv_fp8x2_storage_t sfa_fp8x2_r0_0 = reinterpret_cast<__nv_fp8x2_storage_t*>(sfa_shared[(col_idx * 2) % (Q_SIZE + 1)][threadIdx.y - 1])[threadIdx.x];
            __nv_fp8x2_storage_t sfa_fp8x2_r0_1 = reinterpret_cast<__nv_fp8x2_storage_t*>(sfa_shared[(col_idx * 2) % (Q_SIZE + 1)][threadIdx.y - 1])[threadIdx.x + 32];
            multiply_and_accumulate(a_packed_r0_0, b_packed_0, sfa_fp8x2_r0_0, sfb_fp8x2_0, &result[0][0], &result[0][1], &result[0][2], &result[0][3]);
            multiply_and_accumulate(a_packed_r0_1, b_packed_1, sfa_fp8x2_r0_1, sfb_fp8x2_1, &result[0][0], &result[0][1], &result[0][2], &result[0][3]);

            // Load and compute row 1
            __pipeline_wait_prior(0);
            int4 a_packed_r1_0 = a_shared[(col_idx * 2 + 1) % (Q_SIZE + 1)][threadIdx.y - 1][0][threadIdx.x];
            int4 a_packed_r1_1 = a_shared[(col_idx * 2 + 1) % (Q_SIZE + 1)][threadIdx.y - 1][1][threadIdx.x];
            __nv_fp8x2_storage_t sfa_fp8x2_r1_0 = reinterpret_cast<__nv_fp8x2_storage_t*>(sfa_shared[(col_idx * 2 + 1) % (Q_SIZE + 1)][threadIdx.y - 1])[threadIdx.x];
            __nv_fp8x2_storage_t sfa_fp8x2_r1_1 = reinterpret_cast<__nv_fp8x2_storage_t*>(sfa_shared[(col_idx * 2 + 1) % (Q_SIZE + 1)][threadIdx.y - 1])[threadIdx.x + 32];
            multiply_and_accumulate(a_packed_r1_0, b_packed_0, sfa_fp8x2_r1_0, sfb_fp8x2_0, &result[1][0], &result[1][1], &result[1][2], &result[1][3]);
            multiply_and_accumulate(a_packed_r1_1, b_packed_1, sfa_fp8x2_r1_1, sfb_fp8x2_1, &result[1][0], &result[1][1], &result[1][2], &result[1][3]);
        }

        // Reduction and store
        float final_result_f[2];
        #pragma unroll
        for (int i = 0; i < 2; i++) {
            __half2 reduction_result_0 = __hadd2(reinterpret_cast<const __half2&>(result[i][0]),
                    reinterpret_cast<const __half2&>(result[i][1]));
            __half2 reduction_result_1 = __hadd2(reinterpret_cast<const __half2&>(result[i][2]),
                    reinterpret_cast<const __half2&>(result[i][3]));
            reduction_result_0 = __hadd2(reduction_result_0, reduction_result_1);
            final_result_f[i] = __half22float2(reduction_result_0).x + __half22float2(reduction_result_0).y;
            for (int offset = 16; offset > 0; offset /= 2) {
                final_result_f[i] += __shfl_down_sync(FULL_MASK, final_result_f[i], offset);
            }
        }
        if (threadIdx.x == 0) {
            __half final_result[2];
            final_result[0] = __float2half_rn(final_result_f[0]);
            final_result[1] = __float2half_rn(final_result_f[1]);
            int c_offset = (blockIdx.x * 24 + min((int)blockIdx.x, 32) + threadIdx.y - 1);
            reinterpret_cast<int*>(c)[c_offset] = reinterpret_cast<int&>(final_result);
        }
    }
}



__global__ void gemv_kernel(
    const __nv_fp4x2_storage_t* __restrict__ a,
    const __nv_fp4x2_storage_t* __restrict__ b,
    const __nv_fp8_e4m3* __restrict__ sfa,
    const __nv_fp8_e4m3* __restrict__ sfb,
    __half* __restrict__ c,
    int M,
    int K
) {
    extern __shared__ unsigned char shared_storage[];
    auto* b_shared = reinterpret_cast<__nv_fp4x2_storage_t*>(shared_storage);
    auto* sfb_shared = reinterpret_cast<__nv_fp8_e4m3*>(b_shared + (K / 2));
    __shared__ __half c_shared[32];

    b += blockIdx.y * (K / 2) * 128;
    sfb += blockIdx.y * (K / 16) * 128;

    for (int i = threadIdx.y * 32 + threadIdx.x; i < K / 32; i += blockDim.y * blockDim.x) {
        reinterpret_cast<int4*>(b_shared)[i] = reinterpret_cast<const int4*>(b)[i];
    }
    for (int i = threadIdx.y * 32 + threadIdx.x; i < K / 256; i += blockDim.y * blockDim.x) {
        reinterpret_cast<int4*>(sfb_shared)[i] = reinterpret_cast<const int4*>(sfb)[i];
    }
    __syncthreads();

    // Each warp computes one result and saves it to shared memory
    int result_0 = 0;
    int result_1 = 0;
    int result_2 = 0;
    int result_3 = 0;
    int offset = blockIdx.y * (K * M / 2) + (blockIdx.x * 32 + threadIdx.y) * (K / 2);
    a += offset;
    sfa += offset / 8;
    
    for (int i = threadIdx.x; i < K / 32; i += 32) {
        int4 a_packed = reinterpret_cast<const int4*>(a)[i];
        int4 b_packed = reinterpret_cast<int4*>(b_shared)[i];
        
        __nv_fp8x2_storage_t sfa_fp8x2 = reinterpret_cast<const __nv_fp8x2_storage_t*>(sfa)[i];
        __nv_fp8x2_storage_t sfb_fp8x2 = reinterpret_cast<__nv_fp8x2_storage_t*>(sfb_shared)[i];

        multiply_and_accumulate(a_packed, b_packed, sfa_fp8x2, sfb_fp8x2, &result_0, &result_1, &result_2, &result_3);
    }


    // Reduce the result and store it in shared memory
    __half2 reduction_result_0 = __hadd2(reinterpret_cast<const __half2&>(result_0),
            reinterpret_cast<const __half2&>(result_1));
    __half2 reduction_result_1 = __hadd2(reinterpret_cast<const __half2&>(result_2),
            reinterpret_cast<const __half2&>(result_3));
    reduction_result_0 = __hadd2(reduction_result_0, reduction_result_1);
    float final_result_f = __half22float2(reduction_result_0).x + __half22float2(reduction_result_0).y;
    for (int offset = 16; offset > 0; offset /= 2) {
        final_result_f += __shfl_down_sync(FULL_MASK, final_result_f, offset);
    }
    if (threadIdx.x == 0) {
        c_shared[threadIdx.y] = __float2half_rn(final_result_f);
    }
    __syncthreads();
    
    // Write the result to global memory
    if (threadIdx.y == 0) {
        int c_offset = blockIdx.y * M + blockIdx.x * 32 + threadIdx.x;
        c[c_offset] = c_shared[threadIdx.x];
    }
}



torch::Tensor gemv_cuda(torch::Tensor a, torch::Tensor b, torch::Tensor sfa, torch::Tensor sfb, torch::Tensor c) {
    const int64_t M = a.size(0);
    const int64_t K = a.size(1) * 2;
    const int64_t L = a.size(2);


    dim3 block_dim(32, 32, 1);
    dim3 grid_dim(M / 32, L, 1);
    const auto* a_ptr = reinterpret_cast<const __nv_fp4x2_storage_t*>(a.data_ptr());
    const auto* b_ptr = reinterpret_cast<const __nv_fp4x2_storage_t*>(b.data_ptr());
    const auto* sfa_ptr = reinterpret_cast<const __nv_fp8_e4m3*>(sfa.data_ptr());
    const auto* sfb_ptr = reinterpret_cast<const __nv_fp8_e4m3*>(sfb.data_ptr());
    auto* c_ptr = reinterpret_cast<__half*>(c.data_ptr<c10::Half>());

    size_t shared_mem_bytes =
        (static_cast<size_t>(K) / 2) * sizeof(__nv_fp4x2_storage_t) +
        (static_cast<size_t>(K) / 16) * sizeof(__nv_fp8_e4m3);
    
    if (M == 4096 && K == 7168) {
        gemv_kernel_4096_7168<<<grid_dim, block_dim, shared_mem_bytes>>>(
            a_ptr,
            b_ptr,
            sfa_ptr,
            sfb_ptr,
            c_ptr
        );
    } else if (M == 7168 && K == 2048) {
        gemv_kernel_7168_2048<<<grid_dim, block_dim, shared_mem_bytes>>>(
            a_ptr,
            b_ptr,
            sfa_ptr,
            sfb_ptr,
            c_ptr
        );
    } else if (M == 7168 && K == 16384) {
        grid_dim = dim3(148, 1, 1);
        block_dim = dim3(32, 26, 1);
        gemv_kernel_7168_16384<<<grid_dim, block_dim>>>(
            reinterpret_cast<const int4*>(a.data_ptr()),
            reinterpret_cast<const int4*>(b.data_ptr()),
            reinterpret_cast<const int*>(sfa.data_ptr()),
            reinterpret_cast<const int*>(sfb.data_ptr()),
            c_ptr
        );
    } else {
        gemv_kernel<<<grid_dim, block_dim, shared_mem_bytes>>>(
            a_ptr,
            b_ptr,
            sfa_ptr,
            sfb_ptr,
            c_ptr,
            static_cast<int>(M),
            static_cast<int>(K)
        );
    }
    return c;
}
"""


cpp_source = """
#include <torch/extension.h>

torch::Tensor gemv_cuda(torch::Tensor a, torch::Tensor b, torch::Tensor sfa, torch::Tensor sfb, torch::Tensor c);
"""

gemv_module = load_inline(
    name='gemv_cuda',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['gemv_cuda'],
    verbose=True,
    extra_cuda_cflags=['-arch=compute_100a', '-code=sm_100a', '-O3'],
)




def custom_kernel(
    data: input_t,
) -> output_t:
    """
    PyTorch reference implementation of NVFP4 block-scaled GEMV.
    """

    a, b, sfa, sfb, _, _, c = data

    return gemv_module.gemv_cuda(a, b, sfa, sfb, c)