submission 104178

  ⋯ 61 unchanged lines
   
   
  #define FULL_MASK 0xffffffff
- #define ROWS_PER_BLOCK 32
   
- __device__ void mul_fp4x8_to_half2(
-     int a_packed,
-     int b_packed,
-     __half2  out_pair[4])   // 4 half2 → 8 results
- {
-     #pragma unroll
-     for (int pair = 0; pair < 4; ++pair) {
-         unsigned shift = 8 * pair;
   
-         __nv_fp4x2_storage_t a_pair =
-             static_cast<__nv_fp4x2_storage_t>((a_packed >> shift) & 0xFFu);
-         __nv_fp4x2_storage_t b_pair =
-             static_cast<__nv_fp4x2_storage_t>((b_packed >> shift) & 0xFFu);
+ __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;
   
-         __half2_raw a_raw = __nv_cvt_fp4x2_to_halfraw2(a_pair, __NV_E2M1);
-         __half2_raw b_raw = __nv_cvt_fp4x2_to_halfraw2(b_pair, __NV_E2M1);
+     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];
   
-         // __half2 has a constructor from __half2_raw in recent CUDA versions. :contentReference[oaicite:5]{index=5}
-         __half2 a_h2(a_raw);
-         __half2 b_h2(b_raw);
+     b += blockIdx.y * (K / 2) * 128;
+     sfb += blockIdx.y * (K / 16) * 128;
   
-         out_pair[3 - pair] = __hmul2(a_h2, b_h2);
+     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();
   
- __device__ void mul_fp8x8_to_half2(
-     int2       a_packed,
-     int2       b_packed,
-     __half2    out_pair[4])   // 4 half2 → 8 results
- {
-     #pragma unroll
-     for (int pair = 0; pair < 4; ++pair) {
-         // Select which 32-bit word (x or y) and which 16-bit half inside it.
-         int word_a = (pair < 2) ? a_packed.x : a_packed.y;
-         int word_b = (pair < 2) ? b_packed.x : b_packed.y;
+     // 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];
   
-         unsigned shift = (pair & 1) * 16u;   // 0 or 16 bits
+         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" \\
   
-         __nv_fp8x2_storage_t a_pair =
-             static_cast<__nv_fp8x2_storage_t>((static_cast<unsigned>(word_a) >> shift) & 0xFFFFu);
-         __nv_fp8x2_storage_t b_pair =
-             static_cast<__nv_fp8x2_storage_t>((static_cast<unsigned>(word_b) >> shift) & 0xFFFFu);
+             // 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" \\
   
-         // Convert fp8x2(e4m3) → half2_raw
-         __half2_raw a_raw = __nv_cvt_fp8x2_to_halfraw2(a_pair, __NV_E4M3);
-         __half2_raw b_raw = __nv_cvt_fp8x2_to_halfraw2(b_pair, __NV_E4M3);
+             // 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" \\
   
-         // __half2 has a constructor from __half2_raw in recent CUDA versions.
-         __half2 a_h2(a_raw);
-         __half2 b_h2(b_raw);
+             // 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" \\
   
-         out_pair[3 - pair] = __hmul2(a_h2, b_h2);
+             // 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
+         );
      }
+  
+  
+     // 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];
+     }
  }
   
   
+  
  __global__ void gemv_kernel(
      const __nv_fp4x2_storage_t* __restrict__ a,
      const __nv_fp4x2_storage_t* __restrict__ b,
  ⋯ 11 unchanged lines
      b += blockIdx.y * (K / 2) * 128;
      sfb += blockIdx.y * (K / 16) * 128;
   
-     for (int i = threadIdx.y * 32 + threadIdx.x; i < K / 8; i += blockDim.y * blockDim.x) {
-         reinterpret_cast<int*>(b_shared)[i] = reinterpret_cast<const int*>(b)[i];
+     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 / 64; i += blockDim.y * blockDim.x) {
-         reinterpret_cast<int*>(sfb_shared)[i] = reinterpret_cast<const int*>(sfb)[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
-     __half2 result_0 = __float2half2_rn(0.0f);
-     __half2 result_1 = __float2half2_rn(0.0f);
-     __half2 result_2 = __float2half2_rn(0.0f);
-     __half2 result_3 = __float2half2_rn(0.0f);
+     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;
-     uchar4 a_packed;
-     uchar4 b_packed;
-     __half2 a_h2_0, a_h2_1, a_h2_2, a_h2_3, b_h2_0, b_h2_1, b_h2_2, b_h2_3;
-     __half2 prod_h2_0, prod_h2_1, prod_h2_2, prod_h2_3;
      
      for (int i = threadIdx.x; i < K / 32; i += 32) {
-         int4 a_packed_i4 = reinterpret_cast<const int4*>(a)[i];
-         int4 b_packed_i4 = reinterpret_cast<int4*>(b_shared)[i];
-         const uchar4* a_packed_u4 = reinterpret_cast<const uchar4*>(&a_packed_i4);
-         const uchar4* b_packed_u4 = reinterpret_cast<const uchar4*>(&b_packed_i4);
+         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];
-         __half2 sfa_h2 = __half2(__nv_cvt_fp8x2_to_halfraw2(sfa_fp8x2, __NV_E4M3));
-         __half2 sfb_h2 = __half2(__nv_cvt_fp8x2_to_halfraw2(sfb_fp8x2, __NV_E4M3));
-         __half2 sf_h2 = __hmul2(sfa_h2, sfb_h2);
-         __half2 sf_low_h2 = __low2half2(sf_h2);
-         __half2 sf_high_h2 = __high2half2(sf_h2);
   
-         a_packed = a_packed_u4[0];
-         b_packed = b_packed_u4[0];
-         a_h2_0 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(a_packed.x), __NV_E2M1);
-         a_h2_1 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(a_packed.y), __NV_E2M1);
-         a_h2_2 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(a_packed.z), __NV_E2M1);
-         a_h2_3 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(a_packed.w), __NV_E2M1);
-         b_h2_0 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(b_packed.x), __NV_E2M1);
-         b_h2_1 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(b_packed.y), __NV_E2M1);
-         b_h2_2 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(b_packed.z), __NV_E2M1);
-         b_h2_3 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(b_packed.w), __NV_E2M1);
-         prod_h2_0 = __hmul2(a_h2_0, b_h2_0);
-         prod_h2_1 = __hmul2(a_h2_1, b_h2_1);
-         prod_h2_2 = __hmul2(a_h2_2, b_h2_2);
-         prod_h2_3 = __hmul2(a_h2_3, b_h2_3);
-         result_0 = __hadd2(result_0, __hmul2(prod_h2_0, sf_low_h2));
-         result_1 = __hadd2(result_1, __hmul2(prod_h2_1, sf_low_h2));
-         result_2 = __hadd2(result_2, __hmul2(prod_h2_2, sf_low_h2));
-         result_3 = __hadd2(result_3, __hmul2(prod_h2_3, sf_low_h2));
+         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" \\
   
-         a_packed = a_packed_u4[1];
-         b_packed = b_packed_u4[1];
-         a_h2_0 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(a_packed.x), __NV_E2M1);
-         a_h2_1 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(a_packed.y), __NV_E2M1);
-         a_h2_2 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(a_packed.z), __NV_E2M1);
-         a_h2_3 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(a_packed.w), __NV_E2M1);
-         b_h2_0 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(b_packed.x), __NV_E2M1);
-         b_h2_1 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(b_packed.y), __NV_E2M1);
-         b_h2_2 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(b_packed.z), __NV_E2M1);
-         b_h2_3 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(b_packed.w), __NV_E2M1);
-         prod_h2_0 = __hmul2(a_h2_0, b_h2_0);
-         prod_h2_1 = __hmul2(a_h2_1, b_h2_1);
-         prod_h2_2 = __hmul2(a_h2_2, b_h2_2);
-         prod_h2_3 = __hmul2(a_h2_3, b_h2_3);
-         result_0 = __hadd2(result_0, __hmul2(prod_h2_0, sf_low_h2));
-         result_1 = __hadd2(result_1, __hmul2(prod_h2_1, sf_low_h2));
-         result_2 = __hadd2(result_2, __hmul2(prod_h2_2, sf_low_h2));
-         result_3 = __hadd2(result_3, __hmul2(prod_h2_3, sf_low_h2));
+             // 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" \\
   
-         a_packed = a_packed_u4[2];
-         b_packed = b_packed_u4[2];
-         a_h2_0 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(a_packed.x), __NV_E2M1);
-         a_h2_1 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(a_packed.y), __NV_E2M1);
-         a_h2_2 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(a_packed.z), __NV_E2M1);
-         a_h2_3 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(a_packed.w), __NV_E2M1);
-         b_h2_0 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(b_packed.x), __NV_E2M1);
-         b_h2_1 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(b_packed.y), __NV_E2M1);
-         b_h2_2 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(b_packed.z), __NV_E2M1);
-         b_h2_3 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(b_packed.w), __NV_E2M1);
-         prod_h2_0 = __hmul2(a_h2_0, b_h2_0);
-         prod_h2_1 = __hmul2(a_h2_1, b_h2_1);
-         prod_h2_2 = __hmul2(a_h2_2, b_h2_2);
-         prod_h2_3 = __hmul2(a_h2_3, b_h2_3);
-         result_0 = __hadd2(result_0, __hmul2(prod_h2_0, sf_high_h2));
-         result_1 = __hadd2(result_1, __hmul2(prod_h2_1, sf_high_h2));
-         result_2 = __hadd2(result_2, __hmul2(prod_h2_2, sf_high_h2));
-         result_3 = __hadd2(result_3, __hmul2(prod_h2_3, sf_high_h2));
+             // 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" \\
   
-         a_packed = a_packed_u4[3];
-         b_packed = b_packed_u4[3];
-         a_h2_0 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(a_packed.x), __NV_E2M1);
-         a_h2_1 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(a_packed.y), __NV_E2M1);
-         a_h2_2 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(a_packed.z), __NV_E2M1);
-         a_h2_3 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(a_packed.w), __NV_E2M1);
-         b_h2_0 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(b_packed.x), __NV_E2M1);
-         b_h2_1 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(b_packed.y), __NV_E2M1);
-         b_h2_2 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(b_packed.z), __NV_E2M1);
-         b_h2_3 = __nv_cvt_fp4x2_to_halfraw2(reinterpret_cast<__nv_fp4x2_storage_t>(b_packed.w), __NV_E2M1);
-         prod_h2_0 = __hmul2(a_h2_0, b_h2_0);
-         prod_h2_1 = __hmul2(a_h2_1, b_h2_1);
-         prod_h2_2 = __hmul2(a_h2_2, b_h2_2);
-         prod_h2_3 = __hmul2(a_h2_3, b_h2_3);
-         result_0 = __hadd2(result_0, __hmul2(prod_h2_0, sf_high_h2));
-         result_1 = __hadd2(result_1, __hmul2(prod_h2_1, sf_high_h2));
-         result_2 = __hadd2(result_2, __hmul2(prod_h2_2, sf_high_h2));
-         result_3 = __hadd2(result_3, __hmul2(prod_h2_3, sf_high_h2));
+             // 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
+         );
      }
   
   
      // Reduce the result and store it in shared memory
-     result_0 = __hadd2(result_0, result_1);
-     result_2 = __hadd2(result_2, result_3);
-     result_0 = __hadd2(result_0, result_2);
-     float final_result_f = __half22float2(result_0).x + __half22float2(result_0).y;
+     __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);
      }
  ⋯ 16 unchanged lines
      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());
  ⋯ 5 unchanged lines
      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);
-  
-     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)
-     );
+     
+     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 {
+         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;
  }
  """
  ⋯ 11 unchanged lines
      cuda_sources=cuda_source,
      functions=['gemv_cuda'],
      verbose=True,
-     extra_cuda_cflags=['-arch=sm_100a', '-O3'],
+     extra_cuda_cflags=['-arch=compute_100a', '-code=sm_100a', '-O3'],
  )