submission 93858

#!POPCORN leaderboard nvfp4_gemv

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

# Kernel configuration parameters
sf_vec_size = 16

gemv_cuda_source = r"""
#include<cuda_fp4.h>
#include<cuda_fp16.h>

#define M_BLOCK 8
#define FP4X2_PER_16B 16
#define FP8X2_PER_16B 8
#define K_BLOCK 32 * FP4X2_PER_16B
#define ceilDiv(x, y) (((x) + (y) - 1) / (y))

__device__ __forceinline__ __half2 fp4x2_e2m1_to_half2(__nv_fp4x2_e2m1 v) {
    __nv_fp4x2_storage_t raw = v.__x;  // packed 2×fp4
    __half2_raw hraw = __nv_cvt_fp4x2_to_halfraw2(raw, __NV_E2M1);
    return *reinterpret_cast<__half2*>(&hraw);
}

__device__ __forceinline__ __half2 fp8x2_e4m3_to_half2(__nv_fp8x2_e4m3 v) {
    __nv_fp8x2_storage_t raw = v.__x;
    __half2_raw hraw = __nv_cvt_fp8x2_to_halfraw2(raw, __NV_E4M3);
    return *reinterpret_cast<__half2*>(&hraw);
}

__device__ __forceinline__ __half fp8_e4m3_to_half(__nv_fp8_e4m3 v) {
    __nv_fp8_storage_t raw = v.__x;
    __half_raw hraw = __nv_cvt_fp8_to_halfraw(raw, __NV_E4M3);
    return *reinterpret_cast<__half*>(&hraw);
}

__global__ void debug_print(
        const __nv_fp8x2_e4m3* SFA
        ) {
  __nv_fp8x2_e4m3 sfa_reg_fp8x2;
  float2 x;
  __half2 xh;
  for (int i=0; i<16; i++) {
    sfa_reg_fp8x2 = *(SFA + i);
    xh = fp8x2_e4m3_to_half2(sfa_reg_fp8x2);
    x = __half22float2(xh);
    printf("sfa[%d] %f, sfa[%d] %f \n", i, x.x, i+1, x.y);
  }
}

__global__ void debug_print_scalar(
        const __nv_fp8_e4m3* SFA
        ) {
  __nv_fp8_e4m3 sfa_reg_fp8;
  float x1, x2;
  for (int i=0; i<16; i++) {
    sfa_reg_fp8 = *(SFA + 2*i);
    x1 = __half2float(fp8_e4m3_to_half(sfa_reg_fp8));
    sfa_reg_fp8 = *(SFA + 2*i+1);
    x2 = __half2float(fp8_e4m3_to_half(sfa_reg_fp8));
    printf("sfa[%d] %f, sfa[%d] %f \n", i, x1, i+1, x2);;
  }
}

__global__ void gemv_kernel(
		const __nv_fp4x2_e2m1* A, 
		const __nv_fp4x2_e2m1* B, 
        const __nv_fp8x2_e4m3* SFA,
        const __nv_fp8x2_e4m3* SFB,
		half* C,
		int M,
		int N,
		int K,
		int L
		) {
  // warp layout
  // M/K
  // warp_0
  // warp_1
  // ...
  // warp_BM-1
  // block is 1D
  int threadID = threadIdx.x;
  int warpID = threadID / 32;
  int rowID = warpID;
  int laneID = threadID % 32;

  int blockRowIdx = blockIdx.x * M_BLOCK;
  int threadRowIdx = blockRowIdx + rowID;
  int batchBlockIdx = blockIdx.z;
  //int batchOffset = M * K * batchBlockIdx / 2;
  //int rowOffset =  K * threadRowIdx / 2;

  int batchOffset = M * K * batchBlockIdx;
  int bBatchOffset = N * K * batchBlockIdx;
  int rowOffset =  K * threadRowIdx;
  int cOffset = (M * batchBlockIdx + blockRowIdx);

  // scale factor offsets
  // Have K//16 fp8 values per row 
  // We are interpreting the pointer as fp8x2 so we have K//32 values per row
  //int sfaRowOffset = K * threadRowIdx / 32;
  //int sfaBatchOffset = M * K * batchBlockIdx / 32;
  //int sfbBatchOffset = K * batchBlockIdx / 32;

  int sfaBatchOffset = M * K * batchBlockIdx / 16;
  int sfbBatchOffset = N * K * batchBlockIdx / 16;
  int sfaRowOffset = K * threadRowIdx / 16;
  const unsigned FULL_MASK = 0xffffffff;

  __nv_fp4x2_e2m1 a_reg_fp4x2[16];
  __nv_fp4x2_e2m1 b_reg_fp4x2[16];
  uint4 * a_reg_ptr = reinterpret_cast<uint4 *>(&a_reg_fp4x2[0]);
  uint4 * b_reg_ptr = reinterpret_cast<uint4 *>(&b_reg_fp4x2[0]);


  __nv_fp8x2_e4m3 sfa_reg_fp8x2;
  __nv_fp8x2_e4m3 sfb_reg_fp8x2;
  //__nv_fp8x2_e4m3 *sfa_reg_ptr = reinterpret_cast<__nv_fp8x2_e4m3 *>(&sfa_reg_fp8x2);
  //__nv_fp8x2_e4m3 *sfb_reg_ptr = reinterpret_cast<__nv_fp8x2_e4m3 *>(&sfb_reg_fp8x2);
  __half2 sfa_reg_half2;
  __half2 sfb_reg_half2;
  //__half2 a_reg_half2[16];
  //__half2 b_reg_half2[16];
  float2 a_reg_float2[16];
  float2 b_reg_float2[16];
  //__half final_accum = __float2half(0.0f);
  //__half2 accum[8];
  float accum[8];
  float final_accum = 0.0f;

  __shared__ uint4 b_shared1[32];
  __shared__ uint4 b_shared2[32];
  __shared__ __nv_fp8x2_e4m3 sfb_shared1[32];
  __shared__ __nv_fp8x2_e4m3 sfb_shared2[32];

  uint4* b_bufs[2] = {b_shared1, b_shared2};
  __nv_fp8x2_e4m3* sfb_bufs[2] = {sfb_shared1, sfb_shared2};
  uint ctr = 0;

  for (int k_tile=0; k_tile<K; k_tile+=K_BLOCK) {
    if ((k_tile + laneID * FP4X2_PER_16B) < K) {
      if (warpID==0) {
        const uint4 *gB_ptr = reinterpret_cast<const uint4 *>(B + bBatchOffset + FP4X2_PER_16B*laneID + k_tile);
        const __nv_fp8x2_e4m3 *gSFB_ptr = (SFB + sfbBatchOffset + laneID + k_tile/16);
        //b_shared[laneID] = *gB_ptr;
        //sfb_shared[laneID] = *gSFB_ptr;
        b_bufs[ctr%2][laneID] = *gB_ptr;
        sfb_bufs[ctr%2][laneID] = *gSFB_ptr;
      }
    }
    __syncthreads();

    if ((k_tile + laneID * FP4X2_PER_16B) < K) {
    // read 16B from global to reg
    const uint4 *gA_ptr = reinterpret_cast<const uint4 *>(A + batchOffset + rowOffset + FP4X2_PER_16B*laneID + k_tile);
    const __nv_fp8x2_e4m3 *gSFA_ptr = (SFA + sfaBatchOffset + sfaRowOffset + laneID + k_tile/16);
    *a_reg_ptr = *gA_ptr;
    sfa_reg_fp8x2 = *gSFA_ptr;

    // TODO: look at coalescing these loads
    //const uint4 *gB_ptr = reinterpret_cast<const uint4 *>(B + bBatchOffset + FP4X2_PER_16B*laneID + k_tile);
    //const __nv_fp8x2_e4m3 *gSFB_ptr = (SFB + sfbBatchOffset + laneID + k_tile/16);
    //*b_reg_ptr = *gB_ptr;
    //sfb_reg_fp8x2 = *gSFB_ptr;
    //*b_reg_ptr = b_shared[laneID];
    //sfb_reg_fp8x2 = sfb_shared[laneID];
    *b_reg_ptr = b_bufs[ctr%2][laneID];
    sfb_reg_fp8x2 = sfb_bufs[ctr%2][laneID];

    // a reg is 16B so contains 32 fp4 vals
    // convert fp4x2 to __half2
    sfa_reg_half2 = fp8x2_e4m3_to_half2(sfa_reg_fp8x2);
    sfb_reg_half2 = fp8x2_e4m3_to_half2(sfb_reg_fp8x2);

    //__half2 sfa_0_half2 = __half2half2(__low2half(sfa_reg_half2));
    //__half2 sfa_1_half2 = __half2half2(__high2half(sfa_reg_half2));
    //__half2 sfb_0_half2 = __half2half2(__low2half(sfb_reg_half2));
    //__half2 sfb_1_half2 = __half2half2(__high2half(sfb_reg_half2));

    //float sfa_val_low = __half2float(__low2half(sfa_reg_half2));
    //float sfa_val_high = __half2float(__high2half(sfa_reg_half2));
    //float sfb_val_low = __half2float(__low2half(sfb_reg_half2));
    //float sfb_val_high = __half2float(__high2half(sfb_reg_half2));

    float sfa_vals[2] = {__half2float(__low2half(sfa_reg_half2)), __half2float(__high2half(sfa_reg_half2))};
    float sfb_vals[2] = {__half2float(__low2half(sfb_reg_half2)), __half2float(__high2half(sfb_reg_half2))};
    #pragma unroll
    for (int j=0; j<16; j++) {
      a_reg_float2[j] = __half22float2(fp4x2_e2m1_to_half2(a_reg_fp4x2[j]));
      b_reg_float2[j] = __half22float2(fp4x2_e2m1_to_half2(b_reg_fp4x2[j]));
      a_reg_float2[j].x = a_reg_float2[j].x * sfa_vals[j/8];
      a_reg_float2[j].y = a_reg_float2[j].y * sfa_vals[j/8];
      b_reg_float2[j].x = b_reg_float2[j].x * sfb_vals[j/8];
      b_reg_float2[j].y = b_reg_float2[j].y * sfb_vals[j/8];

      //float sfa_val = (j < 8 ? sfa_val_low : sfa_val_high);
      //float sfb_val = (j < 8 ? sfb_val_low : sfb_val_high);
      //a_reg_half2[j] = (fp4x2_e2m1_to_half2(a_reg_fp4x2[j]));
      //b_reg_half2[j] = (fp4x2_e2m1_to_half2(b_reg_fp4x2[j]));

      //__half2 sfa_val = (j < 8 ? sfa_0_half2 : sfa_1_half2);
      //__half2 sfb_val = (j < 8 ? sfb_1_half2 : sfb_1_half2);
      /*
      if ((k_tile==0) && (threadID==15) && (blockIdx.x==0) && (blockIdx.z==0)) {
        printf("j: %d, sfa: %f, %f \n", j, sfa_val_low, sfa_val_high);
      }
      */
      //a_reg_half2[j] = __hmul2(a_reg_half2[j], sfa_val);
      //b_reg_half2[j] = __hmul2(b_reg_half2[j], sfb_val);
    }
    // at end of this accum[0] contains the sum of this threads 32 vals 
    #pragma unroll
    for (int j=0; j<8; ++j) {
    /*
      accum[j] = __hadd2(
          (__hmul2(a_reg_half2[2*j  ], b_reg_half2[2*j  ])), 
          (__hmul2(a_reg_half2[2*j+1], b_reg_half2[2*j+1]))
          );
          */
      float2 prod1 = make_float2(a_reg_float2[2*j].x * b_reg_float2[2*j].x, a_reg_float2[2*j].y * b_reg_float2[2*j].y);
      float2 prod2 = make_float2(a_reg_float2[2*j + 1].x * b_reg_float2[2*j + 1].x, a_reg_float2[2*j + 1].y * b_reg_float2[2*j + 1].y);
      //float2 prod1 = __half22float2(__hmul2(a_reg_half2[2*j  ], b_reg_half2[2*j  ]));
      //float2 prod2 = __half22float2(__hmul2(a_reg_half2[2*j+1], b_reg_half2[2*j+1]));
      //float2 a1 = __half22float2(a_reg_half2[2*j]);
      //float2 a2 = __half22float2(a_reg_half2[2*j+1]);
      //float2 b1 = __half22float2(b_reg_half2[2*j]);
      //float2 b2 = __half22float2(b_reg_half2[2*j+1]);
      //float2 prod1 = make_float2(a1.x * b1.x, a1.y * b1.y);
      //float2 prod2 = make_float2(a2.x * b2.x, a2.y * b2.y);
      //accum[2*j]   =  prod1.x + prod2.x;
      //accum[2*j+1] = prod1.y + prod2.y;
      accum[j]   =  prod1.x + prod2.x;
      accum[j] += prod1.y + prod2.y;
    }
    /*
    #pragma unroll
    for (int j=0; j<8; ++j) {
      accum[j] = accum[2*j] + accum[2*j+1];
    }
    */
    accum[0] = accum[0] + accum[1];
    accum[2] = accum[2] + accum[3];
    accum[4] = accum[4] + accum[5];
    accum[6] = accum[6] + accum[7];
    accum[0] = accum[0] + accum[2];
    accum[4] = accum[4] + accum[6];
    accum[0] = accum[0] + accum[4];
    final_accum += accum[0];

    //accum[0] = __hadd2(accum[0], accum[1]);
    //accum[2] = __hadd2(accum[2], accum[3]);
    //accum[4] = __hadd2(accum[4], accum[5]);
    //accum[6] = __hadd2(accum[6], accum[7]);
    //accum[0] = __hadd2(accum[0], accum[2]);
    //accum[4] = __hadd2(accum[4], accum[6]);
    //accum[0] = __hadd2(accum[0], accum[4]);
    //final_accum = __hadd(final_accum, __hadd(__low2half(accum[0]),  __high2half(accum[0])));
    }
    //__syncthreads();

  }
  // at this point each thread contains the sum of it's strided values in the row
  // need to use a warp reduction on each warp to compute final row sum
	// Tree reduction: fold upper half onto lower half
  for (int offset = 16; offset > 0; offset >>= 1) {
    final_accum += __shfl_down_sync(FULL_MASK, final_accum, offset);
  }
  // now for all threads with laneID = 0 want to write the FP16 val to C
  __shared__ half shared_C[M_BLOCK];
  if (laneID == 0) {
    shared_C[warpID] = __float2half(final_accum);
  }
  __syncthreads();
  // each thread can write 8 FP16 values to global memory in one go
  // we have BLOCK_M FP16 values to write so need BLOCK_M // 8 threads to participate
  if (threadID < M_BLOCK/8) {
    *reinterpret_cast<uint4 *>(C + cOffset + 8*threadID) = *reinterpret_cast<uint4 *>(shared_C + 8*threadID);
  }
  /*
  if (laneID == 0) {
    C[cOffset + warpID] = __float2half(final_accum);
  }
  */
}

torch::Tensor gemv_cuda(torch::Tensor A, torch::Tensor B, torch::Tensor SFA, torch::Tensor SFB, torch::Tensor C) {
    TORCH_CHECK(A.device().is_cuda(), "Tensor A must be a CUDA tensor");
    TORCH_CHECK(B.device().is_cuda(), "Tensor B must be a CUDA tensor");
    TORCH_CHECK(SFA.device().is_cuda(), "Tensor SFA must be a CUDA tensor");
    TORCH_CHECK(SFB.device().is_cuda(), "Tensor SFB must be a CUDA tensor");
    TORCH_CHECK(C.device().is_cuda(), "Tensor C must be a CUDA tensor");
    
    torch::IntArrayRef a_sizes = A.sizes();
    torch::IntArrayRef b_sizes = B.sizes();
    int M = a_sizes[0];
    //int K = a_sizes[1] * 2;
    int K = a_sizes[1];
    int L = a_sizes[2];
    int N = b_sizes[0];

    const int threads = M_BLOCK * 32; 
    //printf("M_BLOCK : %d \n", M_BLOCK);
    const dim3 grid(ceilDiv(M, M_BLOCK), 1, L);
    //printf("Problem size M: %d, K: %d, N: %d, L: %d \n", M, K, N, L);
    //printf("Threads per block: %d, Block dims (%d, 1, %d)\n", threads, grid.x, grid.z);

    auto A_ptr = reinterpret_cast<__nv_fp4x2_e2m1*>(A.data_ptr());
    auto B_ptr = reinterpret_cast<__nv_fp4x2_e2m1*>(B.data_ptr());
    auto SFA_ptr = reinterpret_cast<__nv_fp8x2_e4m3*>(SFA.data_ptr());
    auto SFB_ptr = reinterpret_cast<__nv_fp8x2_e4m3*>(SFB.data_ptr());
    auto C_ptr = reinterpret_cast<__half*>(C.data_ptr());
    
    //auto SFA_ptr1 = reinterpret_cast<__nv_fp8_e4m3*>(SFA.data_ptr());
    //debug_print_scalar<<<1, 1>>>(SFA_ptr1);
    gemv_kernel<<<grid, threads>>>(
            A_ptr,
            B_ptr,
            SFA_ptr,
            SFB_ptr,
            C_ptr,
            M, N, K, L
        );

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        throw std::runtime_error(cudaGetErrorString(err));
    }
    return C;
}
"""

gemv_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=gemv_cpp_source,
    cuda_sources=gemv_cuda_source,
    functions=['gemv_cuda'],
    verbose=True,
)


def gemv_cuda(A, B, SFA, SFB, C):
    if not A.is_cuda or not B.is_cuda or not SFA.is_cuda or not SFB.is_cuda or not C.is_cuda:
        raise RuntimeError("Both tensors must be on GPU")
    return gemv_module.gemv_cuda(A, B, SFA, SFB, C)


# Helper function for ceiling division
def ceil_div(a, b):
    return (a + b - 1) // b


def custom_kernel(
    data: input_t,
) -> output_t:
    """
    PyTorch reference implementation of NVFP4 block-scaled GEMV.
    """
    a_ref, b_ref, sfa, sfb, _, _, c_ref = data
    m, k, l = a_ref.shape
    n, k, l = b_ref.shape
    """
    print(f"K is {k}, n is {n}")
    print(f"A shape {a_ref.shape}")
    print(f"A shape {a_ref.stride()}")
    print(f"SFA shape {sfa.shape}")
    print(f"SFA shape {sfa.stride()}")
    print(f"B shape {b_ref.shape}")
    print(f"B shape {b_ref.stride()}")
    print(f"SFB shape {sfb.shape}")
    print(f"SFB shape {sfb.stride()}")
    print(f"C shape {c_ref.shape}")
    print(f"C shape {c_ref.stride()}")
    """

    # Get dimensions from MxNxL layout
    _, _, l = c_ref.shape
    #print(sfa.shape, sfa.stride())
    #print(f"SFA[0,0:32,0]: {sfa[0,:32,0].reshape(-1,2)}")
    gemv_cuda(a_ref, b_ref, sfa, sfb, c_ref)
    #torch.cuda.synchronize()
    #print(c_ref)
    return c_ref