submission 75989

  ⋯ 6 unchanged lines
  import cutlass.utils.blockscaled_layout as blockscaled_utils
   
  # Kernel configuration parameters
- m_dim = 128
- mma_tiler_mnk = (m_dim, 1, 256)  # Tile sizes for M, N, K dimensions (default)
  ab_dtype = cutlass.Float4E2M1FN  # FP4 data type for A and B
  sf_dtype = cutlass.Float8E4M3FN  # FP8 data type for scale factors
  c_dtype = cutlass.Float16  # FP16 output type
  sf_vec_size = 16  # Scale factor block size (16 elements share one scale)
- threads_per_cta = m_dim  # Number of threads per CUDA thread block
+ threads_per_cta = 128  # Number of threads per CUDA thread block
   
   
  # Helper function for ceiling division
  def ceil_div(a, b):
      return (a + b - 1) // b
   
- # The optimized CuTe implementation for NVFP4 block-scaled GEMV
- @cute.kernel
- def kernel(
-     mA_mkl: cute.Tensor,
-     mB_nkl: cute.Tensor,
-     mSFA_mkl: cute.Tensor,
-     mSFB_nkl: cute.Tensor,
-     mC_mnl: cute.Tensor,
- ):
-     # Get CUDA block and thread indices
-     bidx, bidy, bidz = cute.arch.block_idx()
-     tidx, _, _ = cute.arch.thread_idx()
   
-     # Extract the local tile for input matrix A (shape: [block_M, block_K, rest_M, rest_K, rest_L])
-     gA_mkl = cute.local_tile(
-         mA_mkl, cute.slice_(mma_tiler_mnk, (None, 0, None)), (None, None, None)
-     )
-     
-     # Extract the local tile for scale factor tensor for A with optimized tiling
-     # OPTIMIZATION: Use reduced K dimension for scale factors (64 // 16 = 4)
-     sf_tiler_mnk = (mma_tiler_mnk[0], mma_tiler_mnk[1], mma_tiler_mnk[2] // sf_vec_size)
-     gSFA_mkl = cute.local_tile(
-         mSFA_mkl, cute.slice_(sf_tiler_mnk, (None, 0, None)), (None, None, None)
-     )
-     
-     # Extract the local tile for input matrix B (shape: [block_N, block_K, rest_N, rest_K, rest_L])
-     gB_nkl = cute.local_tile(
-         mB_nkl, cute.slice_(mma_tiler_mnk, (0, None, None)), (None, None, None)
-     )
-     
-     # Extract the local tile for scale factor tensor for B with optimized tiling
-     gSFB_nkl = cute.local_tile(
-         mSFB_nkl, cute.slice_(sf_tiler_mnk, (0, None, None)), (None, None, None)
-     )
-     
-     # Extract the local tile for output matrix C (shape: [block_M, block_N, rest_M, rest_N, rest_L])
-     gC_mnl = cute.local_tile(
-         mC_mnl, cute.slice_(mma_tiler_mnk, (None, None, 0)), (None, None, None)
-     )
+ # Function to create kernel with specific tile size
+ def create_kernel_functions(mma_tiler_mnk):
+     """
+     Create kernel and my_kernel functions with a specific mma_tiler_mnk.
+     This is needed because CuTe kernels need compile-time constants.
+     """
+     # The CuTe reference implementation for NVFP4 block-scaled GEMV
+     @cute.kernel
+     def kernel(
+         mA_mkl: cute.Tensor,
+         mB_nkl: cute.Tensor,
+         mSFA_mkl: cute.Tensor,
+         mSFB_nkl: cute.Tensor,
+         mC_mnl: cute.Tensor,
+     ):
+         # Get CUDA block and thread indices
+         bidx, bidy, bidz = cute.arch.block_idx()
+         tidx, _, _ = cute.arch.thread_idx()
   
-     # Select output element corresponding to this thread and block indices
-     tCgC = gC_mnl[tidx, None, bidx, bidy, bidz]
-     tCgC = cute.make_tensor(tCgC.iterator, 1)
-     res = cute.zeros_like(tCgC, cutlass.Float32)
+         # Extract the local tile for input matrix A (shape: [block_M, block_K, rest_M, rest_K, rest_L])
+         gA_mkl = cute.local_tile(
+             mA_mkl, cute.slice_(mma_tiler_mnk, (None, 0, None)), (None, None, None)
+         )
+         # Extract the local tile for scale factor tensor for A (same shape as gA_mkl)
+         # Here, block_M = (32, 4); block_K = (16, 4)
+         gSFA_mkl = cute.local_tile(
+             mSFA_mkl, cute.slice_(mma_tiler_mnk, (None, 0, None)), (None, None, None)
+         )
+         # Extract the local tile for input matrix B (shape: [block_N, block_K, rest_N, rest_K, rest_L])
+         gB_nkl = cute.local_tile(
+             mB_nkl, cute.slice_(mma_tiler_mnk, (0, None, None)), (None, None, None)
+         )
+         # Extract the local tile for scale factor tensor for B (same shape as gB_nkl)
+         gSFB_nkl = cute.local_tile(
+             mSFB_nkl, cute.slice_(mma_tiler_mnk, (0, None, None)), (None, None, None)
+         )
+         # Extract the local tile for output matrix C (shape: [block_M, block_N, rest_M, rest_N, rest_L])
+         gC_mnl = cute.local_tile(
+             mC_mnl, cute.slice_(mma_tiler_mnk, (None, None, 0)), (None, None, None)
+         )
   
-     # Get the number of k tiles (depth dimension) for the reduction loop
-     k_tile_cnt = gA_mkl.layout[3].shape
-     
-     for k_tile in range(k_tile_cnt):
-         # Load data tile (128 elements)
-         tAgA = gA_mkl[tidx, None, bidx, k_tile, bidz]
-         tBgB = gB_nkl[0, None, bidy, k_tile, bidz]
-         
-         # Load scale factor tile (4 scale factors for the 128 elements)
-         # OPTIMIZATION: Use k_tile directly since both tensors have same number of tiles
-         tAgSFA = gSFA_mkl[tidx, None, bidx, k_tile, bidz]
-         tBgSFB = gSFB_nkl[0, None, bidy, k_tile, bidz]
+         # Select output element corresponding to this thread and block indices
+         tCgC = gC_mnl[tidx, None, bidx, bidy, bidz]
+         tCgC = cute.make_tensor(tCgC.iterator, 1)
+         res = cute.zeros_like(tCgC, cutlass.Float32)
   
-         # Create register tensors - keep structure similar to original
-         tArA = cute.make_rmem_tensor_like(tAgA, cutlass.Float16)
-         tBrB = cute.make_rmem_tensor_like(tBgB, cutlass.Float16)
-         tArSFA = cute.make_rmem_tensor_like(tAgSFA, cutlass.Float32)
-         tBrSFB = cute.make_rmem_tensor_like(tBgSFB, cutlass.Float32)
+         # Get the number of k tiles (depth dimension) for the reduction loop
+         k_tile_cnt = gA_mkl.layout[3].shape
+         for k_tile in range(k_tile_cnt):
+             tAgA = gA_mkl[tidx, None, bidx, k_tile, bidz]
+             tBgB = gB_nkl[0, None, bidy, k_tile, bidz]
+             tAgSFA = gSFA_mkl[tidx, (0, None, None), bidx, k_tile, bidz]
+             tBgSFB = gSFB_nkl[0, (0, None, None), bidy, k_tile, bidz]
   
-         # Load values from global memory
-         a_val_nvfp4 = tAgA.load()
-         b_val_nvfp4 = tBgB.load()
-         sfa_val_fp8 = tAgSFA.load()
-         sfb_val_fp8 = tBgSFB.load()
+             tArA = cute.make_rmem_tensor_like(tAgA, cutlass.Float16)
+             tBrB = cute.make_rmem_tensor_like(tBgB, cutlass.Float16)
+             tArSFA = cute.make_rmem_tensor_like(tAgSFA, cutlass.Float32)
+             tBrSFB = cute.make_rmem_tensor_like(tBgSFB, cutlass.Float32)
   
-         # OPTIMIZATION: Fuse conversion with store operations
-         tArA.store(a_val_nvfp4.to(cutlass.Float16))
-         tBrB.store(b_val_nvfp4.to(cutlass.Float16))
-         tArSFA.store(sfa_val_fp8.to(cutlass.Float32))
-         tBrSFB.store(sfb_val_fp8.to(cutlass.Float32))
+             # Load NVFP4 or FP8 values from global memory
+             a_val_nvfp4 = tAgA.load()
+             b_val_nvfp4 = tBgB.load()
+             sfa_val_fp8 = tAgSFA.load()
+             sfb_val_fp8 = tBgSFB.load()
   
-         # OPTIMIZED ACCUMULATION: Fuse operations in inner loop
-         num_sf_blocks = mma_tiler_mnk[2] // sf_vec_size
-         for sf_block in cutlass.range_constexpr(num_sf_blocks):
-             scale_prod = tArSFA[sf_block] * tBrSFB[sf_block]
-             base = sf_block * sf_vec_size
+             # Store the converted values to RMEM CuTe tensors
+             tArA.store(a_val_nvfp4.to(cutlass.Float16))
+             tBrB.store(b_val_nvfp4.to(cutlass.Float16))
+             tArSFA.store(sfa_val_fp8.to(cutlass.Float32))
+             tBrSFB.store(sfb_val_fp8.to(cutlass.Float32))
   
-             # OPTIMIZATION: Use fused multiply-add pattern
-             for offset in cutlass.range_constexpr(sf_vec_size):
-                 element_idx = base + offset
-                 # Fuse: res += scale_prod * (a * b)
-                 res += scale_prod * (tArA[element_idx] * tBrB[element_idx])
+             # Iterate over SF vector tiles and compute the scale&matmul accumulation
+             for sf_block in cutlass.range_constexpr(mma_tiler_mnk[2] // sf_vec_size):
+                 tmp = cute.zeros_like(tCgC, cutlass.Float32)
+                 # base = sf_block * sf_vec_size
   
-     # Store the final float16 result back to global memory
-     tCgC.store(res.to(cutlass.Float16))
-     return
+                 for offset in cutlass.range_constexpr(sf_vec_size):
+                     tmp += tArA[sf_block * sf_vec_size + offset] * tBrB[sf_block * sf_vec_size + offset]
+                 res += tArSFA[sf_block] * tBrSFB[sf_block] * tmp
   
+         # Store the final float16 result back to global memory
+         tCgC.store(res.to(cutlass.Float16))
+         return
   
- @cute.jit
- def my_kernel(
-     a_ptr: cute.Pointer,
-     b_ptr: cute.Pointer,
-     sfa_ptr: cute.Pointer,
-     sfb_ptr: cute.Pointer,
-     c_ptr: cute.Pointer,
-     problem_size: tuple,
- ):
-     """
-     Host-side JIT function to prepare tensors and launch GPU kernel.
-     """
-     m, _, k, l = problem_size
-     
-     # Create CuTe Tensor via pointer and problem size.
-     a_tensor = cute.make_tensor(
-         a_ptr,
-         cute.make_layout(
-             (m, cute.assume(k, 32), l),
-             stride=(cute.assume(k, 32), 1, cute.assume(m * k, 32)),
-         ),
-     )
-     
-     # We use n=128 to create the torch tensor to do fp4 computation via torch._scaled_mm
-     # then copy torch tensor to cute tensor for cute customize kernel computation
-     # therefore we need to ensure b_tensor has the right stride with this 128 padded size on n.
-     n_padded_128 = 128
-     b_tensor = cute.make_tensor(
-         b_ptr,
-         cute.make_layout(
-             (n_padded_128, cute.assume(k, 32), l),
-             stride=(cute.assume(k, 32), 1, cute.assume(n_padded_128 * k, 32)),
-         ),
-     )
-     
-     c_tensor = cute.make_tensor(
-         c_ptr, cute.make_layout((cute.assume(m, 32), 1, l), stride=(1, 1, m))
-     )
-     
-     k_sf = k // sf_vec_size
-     
-     # K-major order for scale factors: K//16 dimension is NOT contiguous!
-     # For [M, K//16, L] in K-major
-     sfa_tensor = cute.make_tensor(
-         sfa_ptr, 
-         cute.make_layout(
-             (m, k_sf, l),
-             stride=(k_sf, 1, m * k_sf)
+     @cute.jit
+     def my_kernel(
+         a_ptr: cute.Pointer,
+         b_ptr: cute.Pointer,
+         sfa_ptr: cute.Pointer,
+         sfb_ptr: cute.Pointer,
+         c_ptr: cute.Pointer,
+         problem_size: tuple,
+     ):
+         """
+         Host-side JIT function to prepare tensors and launch GPU kernel.
+         """
+         m, _, k, l = problem_size
+         # Create CuTe Tensor via pointer and problem size.
+         a_tensor = cute.make_tensor(
+             a_ptr,
+             cute.make_layout(
+                 (m, cute.assume(k, 32), l),
+                 stride=(cute.assume(k, 32), 1, cute.assume(m * k, 32)),
+             ),
          )
-     )
-  
-     # For [1, K//16, L] in K-major
-     sfb_tensor = cute.make_tensor(
-         sfb_ptr, 
-         cute.make_layout(
-             (n_padded_128, k_sf, l),
-             stride=(k_sf, 1, n_padded_128 * k_sf)
+         # We use n=128 to create the torch tensor to do fp4 computation via torch._scaled_mm
+         # then copy torch tensor to cute tensor for cute customize kernel computation
+         # therefore we need to ensure b_tensor has the right stride with this 128 padded size on n.
+         n_padded_128 = 128
+         b_tensor = cute.make_tensor(
+             b_ptr,
+             cute.make_layout(
+                 (n_padded_128, cute.assume(k, 32), l),
+                 stride=(cute.assume(k, 32), 1, cute.assume(n_padded_128 * k, 32)),
+             ),
          )
-     )
+         c_tensor = cute.make_tensor(
+             c_ptr, cute.make_layout((cute.assume(m, 32), 1, l), stride=(1, 1, m))
+         )
+         # Convert scale factor tensors to MMA layout
+         # The layout matches Tensor Core requirements: (((32, 4), REST_M), ((SF_K, 4), REST_K), (1, REST_L))
+         sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(a_tensor.shape, sf_vec_size)
+         sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout)
   
+         sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(b_tensor.shape, sf_vec_size)
+         sfb_tensor = cute.make_tensor(sfb_ptr, sfb_layout)
   
-     # Compute grid dimensions
-     # Grid is (M_blocks, 1, L) where:
-     # - M_blocks = ceil(M / 128) to cover all output rows
-     # - L = batch size
-     grid = (
-         cute.ceil_div(c_tensor.shape[0], m_dim),
-         1,
-         c_tensor.shape[2],
-     )
+         # Compute grid dimensions
+         # Grid is (M_blocks, 1, L) where:
+         # - M_blocks = ceil(M / 128) to cover all output rows
+         # - L = batch size
+         grid = (
+             cute.ceil_div(c_tensor.shape[0], 128),
+             1,
+             c_tensor.shape[2],
+         )
   
-     # Launch the CUDA kernel
-     kernel(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor).launch(
-         grid=grid,
-         block=[threads_per_cta, 1, 1],
-         cluster=(1, 1, 1),
-     )
-     return
+         # Launch the CUDA kernel
+         kernel(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor).launch(
+             grid=grid,
+             block=[threads_per_cta, 1, 1],
+             cluster=(1, 1, 1),
+         )
+         return
   
+     return my_kernel
   
- # Global cache for compiled kernels keyed by MMA tiler configuration
+  
+ # Global cache for compiled kernels by tile size
  _compiled_kernel_cache = {}
   
   
- def compile_kernel(tile_config):
+ # This function is used to compile the kernel once and cache it and then allow users to
+ # run the kernel multiple times to get more accurate timing results.
+ def compile_kernel(mma_tiler_mnk):
      """
-     Compile the kernel once and cache it.
+     Compile the kernel with a specific tile size and cache it.
      This should be called before any timing measurements.
   
+     Args:
+         mma_tiler_mnk: Tuple of (M, N, K) tile sizes
+  
      Returns:
          The compiled kernel function
      """
-     global _compiled_kernel_cache, mma_tiler_mnk
+     global _compiled_kernel_cache
   
-     if tile_config in _compiled_kernel_cache:
-         # Ensure global tiler matches the cached configuration before launching
-         mma_tiler_mnk = tile_config
-         return _compiled_kernel_cache[tile_config]
+     # Check if already cached
+     if mma_tiler_mnk in _compiled_kernel_cache:
+         return _compiled_kernel_cache[mma_tiler_mnk]
   
-     # Update global tiler configuration for compilation
-     mma_tiler_mnk = tile_config
+     # Create kernel functions with the specific tile size
+     my_kernel = create_kernel_functions(mma_tiler_mnk)
   
      # Create CuTe pointers for A/B/C/SFA/SFB via torch tensor data pointer
      a_ptr = make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16)
  ⋯ 3 unchanged lines
      sfb_ptr = make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32)
   
      # Compile the kernel
-     compiled = cute.compile(
+     compiled_kernel = cute.compile(
          my_kernel, a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (0, 0, 0, 0)
      )
   
-     _compiled_kernel_cache[tile_config] = compiled
+     # Cache the compiled kernel
+     _compiled_kernel_cache[mma_tiler_mnk] = compiled_kernel
   
-     return compiled
+     return compiled_kernel
   
   
  def custom_kernel(data: input_t) -> output_t:
  ⋯ 17 unchanged lines
      Returns:
          Output tensor c with computed GEMV results
      """
-     a, b, sfa_cpu, sfb_cpu, _, _, c = data
+     a, b, _, _, sfa_permuted, sfb_permuted, c = data
   
      # Get dimensions from MxKxL layout
      m, k, l = a.shape
  ⋯ 2 unchanged lines
      # GEMV N dimension is always 1
      n = 1
   
-     # Select MMA tiler configuration based on K dimension
-     tile_k = 256 if k < 512 else 512
-     tile_config = (m_dim, 1, tile_k)
+     # Determine tile size based on k
+     if k < 512:
+         mma_tiler_mnk = (128, 1, 256)
+     else:
+         mma_tiler_mnk = (128, 1, 512)
   
-     # Ensure kernel is compiled (will use cached version if available) for the chosen tiler.
-     compiled_func = compile_kernel(tile_config)
+     # Ensure kernel is compiled (will use cached version if available)
+     # To avoid the compilation overhead, we compile the kernel once and cache it.
+     compiled_func = compile_kernel(mma_tiler_mnk)
   
      # Create CuTe pointers for A/B/C/SFA/SFB via torch tensor data pointer
      a_ptr = make_ptr(ab_dtype, a.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
      b_ptr = make_ptr(ab_dtype, b.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
      c_ptr = make_ptr(c_dtype, c.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
      sfa_ptr = make_ptr(
-         sf_dtype, sfa_cpu.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
+         sf_dtype, sfa_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
      )
      sfb_ptr = make_ptr(
-         sf_dtype, sfb_cpu.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
+         sf_dtype, sfb_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
      )
   
      # Execute the compiled kernel