submission 71049

  ⋯ 2 unchanged lines
   
  torch._dynamo.config.cache_size_limit = 32
   
- # Helper function to convert scale factor tensor to blocked format
- @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()
+ # Convert all scale factors to blocked formats
   
- @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])
+ @torch.compile(dynamic=False, fullgraph=True)
+ def to_blocked_3d(input_matrix):
+     # input_matrix is rows x cols x l
+     rows, cols, l = input_matrix.shape
   
-     # (m, k) @ (n, k).T -> (m, n)
-     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,
-     )[:, 0]
+     data = input_matrix.permute(2, 0, 1)
   
+     return data.view(l, rows // 128, 128, cols // 4, 4) \
+         .transpose(2, 3) \
+         .reshape(l, -1, 4, 32, 4) \
+         .transpose(2, 3) \
+         .flatten(1)
   
- @torch.compile(mode="max-autotune-no-cudagraphs")
- def batched_gemv_impl(a_ref, b_ref, c_ref, sfa_ref, sfb_ref):
+ @torch.compile(dynamic=False, mode="max-autotune-no-cudagraphs")
+ def batched_gemv_impl(a_ref, b_ref, c_ref, sfa, sfb):
      _, _, l = b_ref.shape
   
+     sfa_blocked = to_blocked_3d(sfa)
+     sfb_blocked = to_blocked_3d(sfb)
+  
      for l_idx in range(l):
-         c_ref[:, 0, l_idx] = _inner(a_ref, b_ref, sfa_ref, sfb_ref, l_idx)
+         c_ref[:, 0, l_idx] = torch._scaled_mm(
+             a_ref[..., l_idx],
+             b_ref[..., l_idx].t(),
+             sfa_blocked[l_idx, ...],
+             sfb_blocked[l_idx, ...],
+             bias=None,
+             out_dtype=torch.float16,
+         )[:, 0]
   
      return c_ref
   
  ⋯ 3 unchanged lines
      """
      PyTorch reference implementation of NVFP4 block-scaled GEMV.
      """
-     a_ref, b_ref, sfa_ref_cpu, sfb_ref_cpu, _, _, c_ref = data
+     a_ref, b_ref, sfa, sfb, _, _, c_ref = data
   
-     sfa_ref_gpu = sfa_ref_cpu.cuda()
-     sfb_ref_gpu = sfb_ref_cpu.cuda()
-  
+     # a_ref is [m, k//2, l]
+     # b_ref is [n, k//2, l], n=1 padded to n=128
+     # c_ref is [m, 1, l]
      return batched_gemv_impl(
          a_ref,
          b_ref,
          c_ref,
-         sfa_ref_gpu,
-         sfb_ref_gpu,
+         sfa,
+         sfb,
      )