submission 107128

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

cuda_source = """
__global__ void gemv_kernel(
    __nv_fp4x2_e2m1* A,
    __nv_fp4x2_e2m1* b,
    __nv_fp8_e4m3* sfa,
    __nv_fp8_e4m3* sfb,
    __half* c,
    int M,
    int K) {
  int tid = blockIdx.x * blockDim.x + threadIdx.x;
  int tidy = blockIdx.y * blockDim.y + threadIdx.y;
  if ((tid + tidy) >= M) return;

  float acc = 0.0f;
  for (int k=0; k<K; k++) {
    __nv_fp4x2_storage_t* a_storage = reinterpret_cast<__nv_fp4x2_storage_t*>(A);
    __nv_fp4x2_storage_t* b_storage = reinterpret_cast<__nv_fp4x2_storage_t*>(b);
    __half2 A_f2 = __nv_cvt_fp4x2_to_halfraw2(a_storage[tid * K + k], __NV_E2M1);
    __half2 b_f2 = __nv_cvt_fp4x2_to_halfraw2(b_storage[k], __NV_E2M1);

    __nv_fp8_storage_t* sfb_storage = reinterpret_cast<__nv_fp8_storage_t*>(sfb);
    __half sfb_val = __nv_cvt_fp8_to_halfraw(sfb_storage[(k*2) >> 4], __NV_E4M3);
    __half2 sfb_val_f2 = __half2half2(sfb_val);

    __nv_fp8_storage_t* sfa_storage = reinterpret_cast<__nv_fp8_storage_t*>(sfa);
    __half sfa_val = __nv_cvt_fp8_to_halfraw(sfa_storage[tid * ((K*2) >> 4) + ((k*2) >> 4)], __NV_E4M3);
    __half2 sfa_val_f2 = __half2half2(sfa_val);

    __half2 product = __hmul2_rn(__hmul2_rn(sfa_val_f2, A_f2), __hmul2_rn(sfb_val_f2, b_f2));
    acc += __half2float(__hadd(product.x, product.y));
  }
  c[tid] = __float2half(acc);
}

torch::Tensor gemv(torch::Tensor A, torch::Tensor b, torch::Tensor sfa, torch::Tensor sfb, int M, int K) {
  auto c = torch::empty({M, 1}, torch::dtype(torch::kHalf).device(torch::kCUDA));
  __half* pc = reinterpret_cast<__half*>(c.data_ptr());
  __nv_fp4x2_e2m1* pA = reinterpret_cast<__nv_fp4x2_e2m1*>(A.data_ptr());
  __nv_fp4x2_e2m1* pb = reinterpret_cast<__nv_fp4x2_e2m1*>(b.data_ptr());
  __nv_fp8_e4m3* psfa = reinterpret_cast<__nv_fp8_e4m3*>(sfa.data_ptr());
  __nv_fp8_e4m3* psfb = reinterpret_cast<__nv_fp8_e4m3*>(sfb.data_ptr());

  dim3 threadsPerBlock(128);
  dim3 blocks((M+threadsPerBlock.x-1)/threadsPerBlock.x);
  gemv_kernel<<<blocks, threadsPerBlock>>>(pA, pb, psfa, psfb, pc, M, K);

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

  return c;
}
"""

# NOTE ------------- KERNEL END ------------------- NOTE

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

torch::Tensor gemv(torch::Tensor A, torch::Tensor b, torch::Tensor sfa, torch::Tensor sfb, int M, int K);
"""

module = load_inline(
    name="my_inline",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["gemv"],
    verbose=False,
)

def custom_kernel(data: input_t) -> output_t:

    a_ref, b_ref, sfa_ref, sfb_ref, _, _, c_ref = data

    M, K, L = a_ref.shape 

    for l_idx in range(L):
        A = a_ref[:, :, l_idx]
        b = b_ref[0, :, l_idx]
        sfa = sfa_ref[:, :, l_idx]
        sfb = sfb_ref[0, :, l_idx]
        res = module.gemv(A, b, sfa, sfb, M, K)
        c_ref[:, 0, l_idx] = res[:, 0]
    return c_ref