submission 69558

  import torch
  from task import input_t, output_t
   
+ torch._dynamo.config.cache_size_limit = 32
   
- # Helper function for ceiling division
- def ceil_div(a, b):
-     return (a + b - 1) // b
-  
-  
  # Helper function to convert scale factor tensor to blocked format
- @torch.compile()
+ @torch.compile(dynamic=False, mode="reduce-overhead", fullgraph=True)
  def to_blocked(input_matrix):
      rows, cols = input_matrix.shape
+     blocks = input_matrix.view(rows // 128, 128, cols // 4, 4).permute(0, 2, 1, 3)
+     #rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
+     #return rearranged.flatten()
+     return blocks.reshape(-1, 4, 32, 4).transpose(1, 2).flatten()
   
-     # Please ensure rows and cols are multiples of 128 and 4 respectively
-     n_row_blocks = ceil_div(rows, 128)
-     n_col_blocks = ceil_div(cols, 4)
-  
-     padded = input_matrix
-     blocks = padded.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3)
-     rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
-  
-     return rearranged.flatten()
-  
-  
+ @torch.compile(dynamic=False, mode="reduce-overhead", fullgraph=True)
  def _inner(a_ref, b_ref, sfa_ref, sfb_ref, l_idx):
-     scale_a = to_blocked(sfa_ref[:, :, l_idx])
-     scale_b = to_blocked(sfb_ref[:, :, l_idx])
+     scale_a = to_blocked(sfa_ref[..., l_idx])
+     scale_b = to_blocked(sfb_ref[..., l_idx])
+  
      # (m, k) @ (n, k).T -> (m, n)
-     # (m, k) @ (n, k).T -> (m, n)
-     res = torch._scaled_mm(
-         a_ref[:, :, l_idx],
-         b_ref[:, :, l_idx].transpose(0, 1),
+     return torch._scaled_mm(
+         a_ref[..., l_idx],
+         b_ref[..., l_idx].transpose(0, 1),
          scale_a,
          scale_b,
          bias=None,
          out_dtype=torch.float16,
-     )
-     return res[:, 0]
-     
+     )[:, 0]
+  
+  
+ @torch.compile(mode="max-autotune-no-cudagraphs")
+ def batched_gemv_impl(a_ref, b_ref, c_ref, sfa_ref, sfb_ref):
+     _, _, l = b_ref.shape
+  
+     for l_idx in range(l):
+         c_ref[:, 0, l_idx] = _inner(a_ref, b_ref, sfa_ref, sfb_ref, l_idx)
+  
+     return c_ref
+  
  def custom_kernel(
      data: input_t,
  ) -> output_t:
  ⋯ 2 unchanged lines
      """
      a_ref, b_ref, sfa_ref_cpu, sfb_ref_cpu, _, _, c_ref = data
   
-     # Get dimensions from MxNxL layout
-     _, _, l = c_ref.shape
-  
      sfa_ref_gpu = sfa_ref_cpu.cuda()
      sfb_ref_gpu = sfb_ref_cpu.cuda()
   
-     # Call torch._scaled_mm to compute the GEMV result
-     for l_idx in range(l):
-         c_ref[:, 0, l_idx] = _inner(a_ref, b_ref, sfa_ref_gpu, sfb_ref_gpu, l_idx)
+     return batched_gemv_impl(
+         a_ref,
+         b_ref,
+         c_ref,
+         sfa_ref_gpu,
+         sfb_ref_gpu,
+     )
   
-     return c_ref