submission 103911

+ """
+ CuTe DSL implementation of NVFP4 block-scaled GEMV.
+  
+ This is a simplified version that follows the same pattern as submission_cute.py
+ but with cleaner structure. The kernel processes all batches in a single launch.
+ """
+  
  import torch
  from task import input_t, output_t
   
  ⋯ 4 unchanged lines
   
  from cutlass import Float32
  from cutlass.cutlass_dsl import T, dsl_user_op
- from cutlass._mlir.dialects import nvvm, llvm
+ from cutlass._mlir.dialects import nvvm
   
+ # Kernel configuration parameters
+ 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
+ accum_dtype = cutlass.Float32
+ sf_vec_size = 16  # Scale factor block size (16 elements share one scale)
+  
+ # Thread block configuration
+ threads_per_m = 32
+ threads_per_k = 4
+ blk_k = 256  # K tile size
+  
+ # Tile sizes for the mainloop
+ mma_tiler_mnk = (threads_per_m, 1, blk_k)
+  
+  
+ def ceil_div(a, b):
+     return (a + b - 1) // b
+  
+  
  @dsl_user_op
  def atomic_add_fp32(a: float | Float32, gmem_ptr: cute.Pointer, *, loc=None, ip=None) -> None:
      nvvm.atomicrmw(
          res=T.f32(), op=nvvm.AtomicOpKind.FADD, ptr=gmem_ptr.llvm_ptr, a=Float32(a).ir_value()
      )
   
+  
  @dsl_user_op
  def elem_pointer(x: cute.Tensor, coord: cute.Coord, *, loc=None, ip=None) -> cute.Pointer:
      return x.iterator + cute.crd2idx(coord, x.layout, loc=loc, ip=ip)
   
+  
  @cute.jit
  def scalar_to_ssa(a: cute.Numeric, dtype) -> cute.TensorSSA:
-     """ Convert a scalar to a cute TensorSSA of shape (1,) and given dtype """
+     """Convert a scalar to a cute TensorSSA of shape (1,) and given dtype."""
      vec = cute.make_fragment(1, dtype)
      vec[0] = a
      return vec.load()
   
- # Kernel configuration parameters
   
- 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
- accum_dtype = cutlass.Float32
- sf_vec_size = 16  # Scale factor block size (16 elements share one scale)
- threads_per_cta = 128  # Number of threads per CUDA thread block
- threads_per_m = 32  # Number of threads per CUDA thread block
- threads_per_k  = 32
- mma_tiler_mnk = (threads_per_m, 1, 256)  # Tile sizes for M, N, K dimensions
-  
-  
- # Helper function for ceiling division
- def ceil_div(a, b):
-     return (a + b - 1) // b
-  
-  
- # The CuTe reference implementation for NVFP4 block-scaled GEMV
  @cute.kernel
- def kernel(
+ def gemv_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
+     """
+     Block-scaled GEMV kernel.
+     
+     Computes: C[m, 1, l] = sum_k(A[m, k, l] * SFA[m, k, l] * B[n, k, l] * SFB[n, k, l])
+     
+     Grid: (ceil(m/threads_per_m), 1, l)
+     Block: (threads_per_m, threads_per_k, 1)
+     """
      bidx, bidy, bidz = cute.arch.block_idx()
      tidx, tidy, _ = cute.arch.thread_idx()
-  
-     # Extract the local tile for input matrix A (shape: [block_M, block_K, rest_M, rest_K, rest_L])
+     
+     # Extract tiles for A and its scale factors
      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])
+     
+     # Extract tiles for B and its scale factors
      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])
+     
+     # Extract tiles for output C
      gC_mnl = cute.local_tile(
          mC_mnl, cute.slice_(mma_tiler_mnk, (None, None, 0)), (None, None, None)
      )
-  
-     # Select output element corresponding to this thread and block indices
+     
+     # Select output element for this thread
      tCgC = gC_mnl[tidx, None, bidx, bidy, bidz]
      tCgC = cute.make_tensor(tCgC.iterator, 1)
+     
+     # Initialize accumulator in FP32
      res = cute.zeros_like(tCgC, accum_dtype)
-  
-     # Shared Memory
+     
+     # Shared memory for reduction across K dimension
      allocator = cutlass.utils.SmemAllocator()
-     smem_layout = cute.make_layout(mma_tiler_mnk[0])
-     shared_res = allocator.allocate_tensor(element_type=cutlass.Float32, layout=smem_layout)
-  
+     smem_layout = cute.make_layout(threads_per_m)
+     shared_res = allocator.allocate_tensor(
+         element_type=cutlass.Float32, layout=smem_layout
+     )
+     
+     # Initialize shared memory
      if tidy == 0:
          shared_res[tidx] = 0.0
      cute.arch.sync_threads()
-     # Get the number of k tiles (depth dimension) for the reduction loop
+     
+     # Get K tile count for reduction loop
      k_tile_cnt = gA_mkl.layout[3].shape
+     
+     # Main reduction loop over K tiles
+     # Each thread in tidy processes a subset of K tiles
      for k_tile in range(tidy, k_tile_cnt, threads_per_k, unroll_full=True):
+         # Load A tile and scale factors
          tAgA = gA_mkl[tidx, None, bidx, k_tile, bidz]
-         tBgB = gB_nkl[0, None, bidy, k_tile, bidz]
          tAgSFA = gSFA_mkl[tidx, None, bidx, k_tile, bidz]
+         
+         # Load B tile and scale factors (B is broadcast across M)
+         tBgB = gB_nkl[0, None, bidy, k_tile, bidz]
          tBgSFB = gSFB_nkl[0, None, bidy, k_tile, bidz]
-  
+         
+         # Create register tensors
          tArA = cute.make_rmem_tensor_like(tAgA, c_dtype)
          tBrB = cute.make_rmem_tensor_like(tBgB, c_dtype)
-         tABrAB = cute.make_rmem_tensor_like(tAgA, c_dtype)
          tArSFA = cute.make_rmem_tensor_like(tAgSFA, accum_dtype)
          tBrSFB = cute.make_rmem_tensor_like(tBgSFB, accum_dtype)
-         tSFrSF = cute.make_rmem_tensor_like(tAgSFA, accum_dtype)
-  
-         # 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()
-  
-         # Convert loaded values to float32 for computation (FFMA)
-         a_val = a_val_nvfp4.to(c_dtype)
-         b_val = b_val_nvfp4.to(c_dtype)
-         sfa_val = sfa_val_fp8.to(accum_dtype)
-         sfb_val = sfb_val_fp8.to(accum_dtype)
-  
-         # Store the converted values to RMEM CuTe tensors
+         
+         # Load from global memory and convert types
+         a_val = tAgA.load().to(c_dtype)
+         b_val = tBgB.load().to(c_dtype)
+         sfa_val = tAgSFA.load().to(accum_dtype)
+         sfb_val = tBgSFB.load().to(accum_dtype)
+         
+         # Store to register tensors
          tArA.store(a_val)
          tBrB.store(b_val)
          tArSFA.store(sfa_val)
          tBrSFB.store(sfb_val)
-  
-         tABrAB.store(tArA.load() * tBrB.load())
-         tSFrSF.store(tArSFA.load() * tBrSFB.load())
-  
-         # Iterate over SF vector tiles and compute the scale&matmul accumulation
-         for i in cutlass.range_constexpr(mma_tiler_mnk[2]):
+         
+         # Compute block-scaled dot product for this K tile
+         for i in cutlass.range_constexpr(blk_k):
              res += tArA[i] * tArSFA[i] * tBrB[i] * tBrSFB[i]
      
+     # Reduce across K dimension using atomic add to shared memory
      atomic_add_fp32(res[0], elem_pointer(shared_res, tidx))
      cute.arch.sync_threads()
+     
+     # Final store to global memory (only thread 0 in K dimension)
      if tidy == 0:
          out = scalar_to_ssa(shared_res[tidx], cutlass.Float32)
-         # Store the final float16 result back to global memory
          tCgC.store(out.to(cutlass.Float16))
+     
      return
   
+  
  @cute.jit
- def my_kernel(
+ def gemv_launcher(
      a_ptr: cute.Pointer,
      b_ptr: cute.Pointer,
      sfa_ptr: cute.Pointer,
  ⋯ 1 unchanged lines
      c_ptr: cute.Pointer,
      problem_size: tuple,
  ):
-     """
-     Host-side JIT function to prepare tensors and launch GPU kernel.
-     """
+     """Host-side JIT function to prepare tensors and launch kernel."""
      m, _, k, l = problem_size
-     # Create CuTe Tensor via pointer and problem size.
+     
+     # Create A tensor: [m, k, l] K-major
      a_tensor = cute.make_tensor(
          a_ptr,
          cute.make_layout(
  ⋯ 1 unchanged lines
              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
+     
+     # Create B tensor: [n_padded, k, l] K-major
+     n_padded = 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)),
+             (n_padded, cute.assume(k, 32), l),
+             stride=(cute.assume(k, 32), 1, cute.assume(n_padded * k, 32)),
          ),
      )
+     
+     # Create C tensor: [m, 1, l]
      c_tensor = cute.make_tensor(
-         c_ptr, cute.make_layout((cute.assume(m, 32), 1, l), stride=(1, 1, m))
+         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))
+     
+     # Create scale factor tensors with MMA layout
      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], threads_per_m),
          1,
          c_tensor.shape[2],
      )
-  
-     # Launch the CUDA kernel
-     kernel(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor).launch(
+     
+     # Launch kernel
+     gemv_kernel(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor).launch(
          grid=grid,
          block=[threads_per_m, threads_per_k, 1],
          cluster=(1, 1, 1),
      )
+     
      return
   
   
  ⋯ 1 unchanged lines
  _compiled_kernel_cache = None
   
   
- # 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():
-     """
-     Compile the kernel once and cache it.
-     This should be called before any timing measurements.
-  
-     Returns:
-         The compiled kernel function
-     """
+     """Compile the kernel once and cache it."""
      global _compiled_kernel_cache
-  
+     
      if _compiled_kernel_cache is not None:
          return _compiled_kernel_cache
-  
-     # Create CuTe pointers for A/B/C/SFA/SFB via torch tensor data pointer
+     
+     # Create placeholder pointers for compilation
      a_ptr = make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16)
      b_ptr = make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16)
      c_ptr = make_ptr(c_dtype, 0, cute.AddressSpace.gmem, assumed_align=16)
      sfa_ptr = make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32)
      sfb_ptr = make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32)
-  
-     # Compile the kernel
+     
      try:
          _compiled_kernel_cache = cute.compile(
-             my_kernel, a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (0, 0, 0, 0)
+             gemv_launcher, a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (0, 0, 0, 0)
          )
      except Exception as e:
-         msg = f"cute.compile(my_kernel, ...) failed with error: {e}"
-         raise RuntimeError(msg)
+         raise RuntimeError(f"Kernel compilation failed: {e}")
+     
      return _compiled_kernel_cache
   
   
  def custom_kernel(data: input_t) -> output_t:
      """
      Execute the block-scaled GEMV kernel.
-  
-     This is the main entry point called by the evaluation framework.
-     It converts PyTorch tensors to CuTe tensors, launches the kernel,
-     and returns the result.
-  
+     
+     This implementation processes all batches in a single kernel launch.
+     
      Args:
-         data: Tuple of (a, b, sfa_cpu, sfb_cpu, sfa_permuted, sfb_permuted, c) PyTorch tensors
-             a: [m, k, l] - Input matrix in float4e2m1fn
-             b: [1, k, l] - Input vector in float4e2m1fn
-             sfa_cpu: [m, k, l] - Scale factors in float8_e4m3fn (not used, kept for compatibility)
-             sfb_cpu: [1, k, l] - Scale factors in float8_e4m3fn (not used, kept for compatibility)
-             sfa_permuted: [32, 4, rest_m, 4, rest_k, l] - Scale factors in float8_e4m3fn
-             sfb_permuted: [32, 4, rest_n, 4, rest_k, l] - Scale factors in float8_e4m3fn
+         data: Tuple of (a, b, sfa_ref, sfb_ref, sfa_permuted, sfb_permuted, c) tensors
+             a: [m, k/2, l] - Input matrix in float4e2m1fn_x2
+             b: [n_pad, k/2, l] - Input vector (padded to 128) in float4e2m1fn_x2
+             sfa_ref: [m, sf_k, l] - Scale factors for A (not used)
+             sfb_ref: [n_pad, sf_k, l] - Scale factors for B (not used)
+             sfa_permuted: [32, 4, rest_m, 4, rest_k, l] - Scale factors for A (MMA layout)
+             sfb_permuted: [32, 4, rest_n, 4, rest_k, l] - Scale factors for B (MMA layout)
              c: [m, 1, l] - Output vector in float16
-  
+     
      Returns:
          Output tensor c with computed GEMV results
      """
      a, b, _, _, sfa_permuted, sfb_permuted, c = data
-  
-     # Ensure kernel is compiled (will use cached version if available)
-     # To avoid the compilation overhead, we compile the kernel once and cache it.
+     
+     # Compile kernel (uses cache if available)
      compiled_func = compile_kernel()
-  
-     # Get dimensions from MxKxL layout
-     m, k, l = a.shape
-     # Torch use e2m1_x2 data type, thus k is halved
-     k = k * 2
-     # GEMV N dimension is always 1
-     n = 1
-  
-     # Create CuTe pointers for A/B/C/SFA/SFB via torch tensor data pointer
+     
+     # Get dimensions
+     m, k_packed, l = a.shape
+     k = k_packed * 2  # FP4 packed: 2 elements per byte
+     n = 1  # GEMV
+     
+     # Create CuTe pointers from PyTorch tensors
      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_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
-     )
-     sfb_ptr = make_ptr(
-         sf_dtype, sfb_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
-     )
-  
-     # Execute the compiled kernel
+     sfa_ptr = make_ptr(sf_dtype, sfa_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32)
+     sfb_ptr = make_ptr(sf_dtype, sfb_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32)
+     
+     # Execute kernel - processes all batches in a single launch
      compiled_func(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l))
-  
+     
      return c