submission 107336

  ⋯ 31 unchanged lines
      return dst;
  }
   
- __global__ void gemv_v8(
+ __device__ __forceinline__ half2 process_16_elements(uint32_t a_lo, uint32_t a_hi, uint32_t b_lo, uint32_t b_hi, half2 scale_h2) {
+     uint32_t a0, a1, a2, a3, a4, a5, a6, a7;
+     uint32_t b0, b1, b2, b3, b4, b5, b6, b7;
+     cvt_f4x8_to_f16x8(a_lo, a0, a1, a2, a3);
+     cvt_f4x8_to_f16x8(a_hi, a4, a5, a6, a7);
+     cvt_f4x8_to_f16x8(b_lo, b0, b1, b2, b3);
+     cvt_f4x8_to_f16x8(b_hi, b4, b5, b6, b7);
+     
+     half2 ab0 = __hmul2(*reinterpret_cast<half2*>(&a0), *reinterpret_cast<half2*>(&b0));
+     half2 ab1 = __hmul2(*reinterpret_cast<half2*>(&a1), *reinterpret_cast<half2*>(&b1));
+     half2 ab2 = __hmul2(*reinterpret_cast<half2*>(&a2), *reinterpret_cast<half2*>(&b2));
+     half2 ab3 = __hmul2(*reinterpret_cast<half2*>(&a3), *reinterpret_cast<half2*>(&b3));
+     half2 ab4 = __hmul2(*reinterpret_cast<half2*>(&a4), *reinterpret_cast<half2*>(&b4));
+     half2 ab5 = __hmul2(*reinterpret_cast<half2*>(&a5), *reinterpret_cast<half2*>(&b5));
+     half2 ab6 = __hmul2(*reinterpret_cast<half2*>(&a6), *reinterpret_cast<half2*>(&b6));
+     half2 ab7 = __hmul2(*reinterpret_cast<half2*>(&a7), *reinterpret_cast<half2*>(&b7));
+     
+     half2 sum01 = __hadd2(ab0, ab1);
+     half2 sum23 = __hadd2(ab2, ab3);
+     half2 sum45 = __hadd2(ab4, ab5);
+     half2 sum67 = __hadd2(ab6, ab7);
+     half2 sum0123 = __hadd2(sum01, sum23);
+     half2 sum4567 = __hadd2(sum45, sum67);
+     half2 local_sum = __hadd2(sum0123, sum4567);
+     
+     return __hmul2(local_sum, scale_h2);
+ }
+  
+ __global__ void gemv_uint4(
      const uint8_t* __restrict__ A,
      const uint8_t* __restrict__ B,
      const uint8_t* __restrict__ SFA,
  ⋯ 1 unchanged lines
      half* __restrict__ C,
      int M, int K, int L,
      int64_t a_s0, int64_t a_s2, int64_t b_s2,
-     int64_t sfa_s0, int64_t sfa_s1, int64_t sfa_s2, int64_t sfa_s3, int64_t sfa_s4, int64_t sfa_s5,
-     int64_t sfb_s3, int64_t sfb_s4, int64_t sfb_s5,
+     int64_t sfa_s0, int64_t sfa_s1, int64_t sfa_s2,
+     int64_t sfb_s1, int64_t sfb_s2,
      int64_t c_s0, int64_t c_s2
  ) {
      constexpr int WARPS_PER_BLOCK = 4;
-     
      int warp_id = threadIdx.x / 32;
      int lane_id = threadIdx.x % 32;
      int row = blockIdx.x * WARPS_PER_BLOCK + warp_id;
  ⋯ 1 unchanged lines
      
      if (row >= M) return;
      
-     int64_t sfa_row_base = (row & 31) * sfa_s0 + ((row & 127) >> 5) * sfa_s1 + 
-                            (row / 128) * sfa_s2 + batch * sfa_s5;
-     int64_t sfb_batch_base = batch * sfb_s5;
+     int64_t sfa_base = row * sfa_s0 + batch * sfa_s2;
+     int64_t sfb_base = batch * sfb_s2;
      
      const uint8_t* A_row = A + row * a_s0 + batch * a_s2;
      const uint8_t* B_batch = B + batch * b_s2;
  ⋯ 1 unchanged lines
      half2 acc_h2 = __float2half2_rn(0.0f);
      int K_scales = K / 16;
      
-     for (int scale_base = 0; scale_base < K_scales; scale_base += 32) {
-         int k_scale = scale_base + lane_id;
-         if (k_scale >= K_scales) break;
+     // Each lane processes 2 scale groups (32 FP4 elements) per iteration
+     // 32 lanes * 2 groups = 64 scale groups per iteration
+     for (int scale_base = 0; scale_base < K_scales; scale_base += 64) {
+         int k_scale0 = scale_base + lane_id * 2;
+         int k_scale1 = k_scale0 + 1;
          
-         int sfa_idx = sfa_row_base + (k_scale & 3) * sfa_s3 + (k_scale >> 2) * sfa_s4;
-         int sfb_idx = sfb_batch_base + (k_scale & 3) * sfb_s3 + (k_scale >> 2) * sfb_s4;
+         if (k_scale1 >= K_scales) {
+             // Handle tail: only process first group if valid
+             if (k_scale0 < K_scales) {
+                 int sfa_idx = sfa_base + k_scale0 * sfa_s1;
+                 int sfb_idx = sfb_base + k_scale0 * sfb_s1;
+                 uint8_t sfa_val = SFA[sfa_idx];
+                 uint8_t sfb_val = SFB[sfb_idx];
+                 uint16_t sf_packed = (uint16_t(sfb_val) << 8) | uint16_t(sfa_val);
+                 uint32_t sf_f16x2 = cvt_f8x2_to_f16x2(sf_packed);
+                 half2 sf_h2 = *reinterpret_cast<half2*>(&sf_f16x2);
+                 half scale_h = __hmul(sf_h2.x, sf_h2.y);
+                 half2 scale_h2 = __half2half2(scale_h);
+                 
+                 int k_byte = k_scale0 * 8;
+                 uint2 a_vec = *reinterpret_cast<const uint2*>(A_row + k_byte);
+                 uint2 b_vec = *reinterpret_cast<const uint2*>(B_batch + k_byte);
+                 acc_h2 = __hadd2(acc_h2, process_16_elements(a_vec.x, a_vec.y, b_vec.x, b_vec.y, scale_h2));
+             }
+             break;
+         }
          
-         uint8_t sfa_val = SFA[sfa_idx];
-         uint8_t sfb_val = SFB[sfb_idx];
-         uint16_t sf_packed = (uint16_t(sfb_val) << 8) | uint16_t(sfa_val);
-         uint32_t sf_f16x2 = cvt_f8x2_to_f16x2(sf_packed);
-         half2 sf_h2 = *reinterpret_cast<half2*>(&sf_f16x2);
-         half scale_h = __hmul(sf_h2.x, sf_h2.y);
-         half2 scale_h2 = __half2half2(scale_h);
+         // Load 16 bytes = 32 FP4 elements (2 scale groups)
+         int k_byte = k_scale0 * 8;
+         uint4 a_vec = *reinterpret_cast<const uint4*>(A_row + k_byte);
+         uint4 b_vec = *reinterpret_cast<const uint4*>(B_batch + k_byte);
          
-         int k_byte_start = k_scale * 8;
-         uint2 a_vec = *reinterpret_cast<const uint2*>(A_row + k_byte_start);
-         uint2 b_vec = *reinterpret_cast<const uint2*>(B_batch + k_byte_start);
+         // Load 2 scale factors for A and B
+         int sfa_idx0 = sfa_base + k_scale0 * sfa_s1;
+         int sfa_idx1 = sfa_base + k_scale1 * sfa_s1;
+         int sfb_idx0 = sfb_base + k_scale0 * sfb_s1;
+         int sfb_idx1 = sfb_base + k_scale1 * sfb_s1;
          
-         uint32_t a_f16_0, a_f16_1, a_f16_2, a_f16_3, a_f16_4, a_f16_5, a_f16_6, a_f16_7;
-         cvt_f4x8_to_f16x8(a_vec.x, a_f16_0, a_f16_1, a_f16_2, a_f16_3);
-         cvt_f4x8_to_f16x8(a_vec.y, a_f16_4, a_f16_5, a_f16_6, a_f16_7);
+         uint8_t sfa0 = SFA[sfa_idx0], sfa1 = SFA[sfa_idx1];
+         uint8_t sfb0 = SFB[sfb_idx0], sfb1 = SFB[sfb_idx1];
          
-         uint32_t b_f16_0, b_f16_1, b_f16_2, b_f16_3, b_f16_4, b_f16_5, b_f16_6, b_f16_7;
-         cvt_f4x8_to_f16x8(b_vec.x, b_f16_0, b_f16_1, b_f16_2, b_f16_3);
-         cvt_f4x8_to_f16x8(b_vec.y, b_f16_4, b_f16_5, b_f16_6, b_f16_7);
+         // Convert scales to fp16
+         uint16_t sf_packed0 = (uint16_t(sfb0) << 8) | uint16_t(sfa0);
+         uint16_t sf_packed1 = (uint16_t(sfb1) << 8) | uint16_t(sfa1);
+         uint32_t sf0_f16x2 = cvt_f8x2_to_f16x2(sf_packed0);
+         uint32_t sf1_f16x2 = cvt_f8x2_to_f16x2(sf_packed1);
+         half2 sf0_h2 = *reinterpret_cast<half2*>(&sf0_f16x2);
+         half2 sf1_h2 = *reinterpret_cast<half2*>(&sf1_f16x2);
+         half scale0 = __hmul(sf0_h2.x, sf0_h2.y);
+         half scale1 = __hmul(sf1_h2.x, sf1_h2.y);
+         half2 scale0_h2 = __half2half2(scale0);
+         half2 scale1_h2 = __half2half2(scale1);
          
-         half2 ab0 = __hmul2(*reinterpret_cast<half2*>(&a_f16_0), *reinterpret_cast<half2*>(&b_f16_0));
-         half2 ab1 = __hmul2(*reinterpret_cast<half2*>(&a_f16_1), *reinterpret_cast<half2*>(&b_f16_1));
-         half2 ab2 = __hmul2(*reinterpret_cast<half2*>(&a_f16_2), *reinterpret_cast<half2*>(&b_f16_2));
-         half2 ab3 = __hmul2(*reinterpret_cast<half2*>(&a_f16_3), *reinterpret_cast<half2*>(&b_f16_3));
-         half2 ab4 = __hmul2(*reinterpret_cast<half2*>(&a_f16_4), *reinterpret_cast<half2*>(&b_f16_4));
-         half2 ab5 = __hmul2(*reinterpret_cast<half2*>(&a_f16_5), *reinterpret_cast<half2*>(&b_f16_5));
-         half2 ab6 = __hmul2(*reinterpret_cast<half2*>(&a_f16_6), *reinterpret_cast<half2*>(&b_f16_6));
-         half2 ab7 = __hmul2(*reinterpret_cast<half2*>(&a_f16_7), *reinterpret_cast<half2*>(&b_f16_7));
-         
-         half2 sum01 = __hadd2(ab0, ab1);
-         half2 sum23 = __hadd2(ab2, ab3);
-         half2 sum45 = __hadd2(ab4, ab5);
-         half2 sum67 = __hadd2(ab6, ab7);
-         half2 sum0123 = __hadd2(sum01, sum23);
-         half2 sum4567 = __hadd2(sum45, sum67);
-         half2 local_sum = __hadd2(sum0123, sum4567);
-         
-         acc_h2 = __hfma2(local_sum, scale_h2, acc_h2);
+         // Process first 16 elements (first scale group)
+         acc_h2 = __hadd2(acc_h2, process_16_elements(a_vec.x, a_vec.y, b_vec.x, b_vec.y, scale0_h2));
+         // Process second 16 elements (second scale group)
+         acc_h2 = __hadd2(acc_h2, process_16_elements(a_vec.z, a_vec.w, b_vec.z, b_vec.w, scale1_h2));
      }
      
      float acc = __half2float(acc_h2.x) + __half2float(acc_h2.y);
  ⋯ 4 unchanged lines
      }
  }
   
- torch::Tensor gemv_cuda(torch::Tensor a, torch::Tensor b, torch::Tensor sfa, torch::Tensor sfb, torch::Tensor c) {
+ torch::Tensor gemv_cuda(torch::Tensor a, torch::Tensor b,
+                         torch::Tensor sfa, torch::Tensor sfb,
+                         torch::Tensor c) {
      int M = a.size(0), K = a.size(1) * 2, L = a.size(2);
      constexpr int WARPS_PER_BLOCK = 4;
      dim3 grid((M + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK, L);
      dim3 block(32 * WARPS_PER_BLOCK);
      
-     gemv_v8<<<grid, block>>>(
+     gemv_uint4<<<grid, block>>>(
          reinterpret_cast<const uint8_t*>(a.data_ptr()),
          reinterpret_cast<const uint8_t*>(b.data_ptr()),
          reinterpret_cast<const uint8_t*>(sfa.data_ptr()),
          reinterpret_cast<const uint8_t*>(sfb.data_ptr()),
          reinterpret_cast<half*>(c.data_ptr()),
-         M, K, L, a.stride(0), a.stride(2), b.stride(2),
-         sfa.stride(0), sfa.stride(1), sfa.stride(2), sfa.stride(3), sfa.stride(4), sfa.stride(5),
-         sfb.stride(3), sfb.stride(4), sfb.stride(5), c.stride(0), c.stride(2)
+         M, K, L,
+         a.stride(0), a.stride(2), b.stride(2),
+         sfa.stride(0), sfa.stride(1), sfa.stride(2),
+         sfb.stride(1), sfb.stride(2),
+         c.stride(0), c.stride(2)
      );
      return c;
  }
  ⋯ 2 unchanged lines
  cpp_source = "torch::Tensor gemv_cuda(torch::Tensor a, torch::Tensor b, torch::Tensor sfa, torch::Tensor sfb, torch::Tensor c);"
   
  module = load_inline(
-     name='gemv_v8',
+     name='gemv_uint4',
      cpp_sources=cpp_source,
      cuda_sources=cuda_source,
      functions=['gemv_cuda'],
  ⋯ 3 unchanged lines
   
  def custom_kernel(data: input_t) -> output_t:
      a, b, sfa, sfb, sfa_perm, sfb_perm, c = data
-     return module.gemv_cuda(a, b, sfa_perm, sfb_perm, c)
   No newline at end of file
+     return module.gemv_cuda(a, b, sfa, sfb, c)
   No newline at end of file