submission 74113

  # !POPCORN leaderboard ranked
  import torch
  
- SF_VEC = 16  # scale granularity (per 16 K elements)
+ # ---------------------------------------------------------------------------
+ # Batched GEMV in NVFP4(e2m1) with FP8(E4M3 fnuz) block scales using
+ # Blackwell's torch._scaled_mm() fast path.
+ # ---------------------------------------------------------------------------
  
+ SF_VEC = 16  # per-16 K elements
+ 
  def ceil_div(a: int, b: int) -> int:
      return (a + b - 1) // b
  
- @torch.jit.script_if_tracing
  def _to_blocked(sf_2d: torch.Tensor) -> torch.Tensor:
      """
-     Convert FP8 scaling tensor from (rows, K//16) to the flattened CuTe/Blackwell
-     blocked layout expected by torch._scaled_mm. View/permute/reshape only.
+     Convert FP8 scale factors from [rows, K//16] K-major to the flattened
+     CuTe/Blackwell block layout expected by torch._scaled_mm (view/permute/reshape only).
      """
-     rows = sf_2d.size(0)
-     sf_k = sf_2d.size(1)
-     n_row_blocks = (rows + 127) // 128
-     n_col_blocks = (sf_k + 3) // 4
- 
-     # [nrb,128,ncb,4] -> permute -> reshape -> [*,32,16] -> flatten
+     rows, sf_k = sf_2d.shape
+     n_row_blocks = ceil_div(rows, 128)
+     n_col_blocks = ceil_div(sf_k, 4)
+     # [nrb,128,ncb,4] -> [nrb,ncb,128,4]
      t = sf_2d.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3)
+     # [*,128,4] -> [*,4,32,4] -> swap -> [*,32,4,4] -> [*,32,16] -> flat
      t = t.reshape(-1, 128, 4).view(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
      return t.flatten()
  
- def _scaled_gemv_N1(A_mk: torch.Tensor,
-                     b1k: torch.Tensor,
-                     sfa_mk16: torch.Tensor,
-                     sfb_1k16: torch.Tensor) -> torch.Tensor:
+ # cached grayscale weights for 2-tensor practice probes
+ _GRAYSCALE_W = None
+ 
+ def _scaled_gemv_N1(A_mk, b1k, sfa_mk16, sfb_1k16, scale_a_blocked=None):
      """
-     Fast path: use torch._scaled_mm with N=1 (no padding). Returns [M, 1] fp16.
+     Fast path (preferred): N=1 directly with _scaled_mm. Returns [M,1] fp16.
      """
-     scale_a = _to_blocked(sfa_mk16)
+     scale_a = scale_a_blocked if scale_a_blocked is not None else _to_blocked(sfa_mk16)
      scale_b = _to_blocked(sfb_1k16)
-     outM1 = torch._scaled_mm(
-         A_mk,              # [M,K]  (nvfp4)
-         b1k,               # [1,K]  (nvfp4)
-         scale_a,           # flattened
-         scale_b,           # flattened
-         bias=None,
-         out_dtype=torch.float16,
-     )  # -> [M,1] fp16
-     return outM1
+     return torch._scaled_mm(
+         A_mk, b1k, scale_a, scale_b, bias=None, out_dtype=torch.float16
+     )  # -> [M,1]
  
- def _scaled_gemv_N128(A_mk: torch.Tensor,
-                       b1k: torch.Tensor,
-                       sfa_mk16: torch.Tensor,
-                       sfb_1k16: torch.Tensor,
-                       scratch_B128K: torch.Tensor,
-                       scratch_SFB128: torch.Tensor) -> torch.Tensor:
+ def _scaled_gemv_N128(A_mk, b1k, sfa_mk16, sfb_1k16, scratch_B128K, scratch_SFB128, scale_a_blocked=None):
      """
      Fallback path: pad N to 128 using reusable scratch buffers.
-     Returns [M,1] in fp16 (as a narrowed view of the GEMM output).
+     Returns [M,1] fp16 (narrowed view of GEMM result).
      """
-     K = A_mk.size(1)
-     sfk = sfa_mk16.size(1)  # K//16
- 
-     # zero/one reset in-place (cheap)
+     # reset scratch cheaply
      scratch_B128K.zero_()
      scratch_SFB128.fill_(1)
  
-     # write real row-0 only
+     # real data on row 0 only
      scratch_B128K[0, :].copy_(b1k[0, :])
      scratch_SFB128[0, :].copy_(sfb_1k16[0, :])
  
-     scale_a = _to_blocked(sfa_mk16)
+     scale_a = scale_a_blocked if scale_a_blocked is not None else _to_blocked(sfa_mk16)
      scale_b = _to_blocked(scratch_SFB128)
  
      outMN = torch._scaled_mm(
-         A_mk,
-         scratch_B128K,
-         scale_a,
-         scale_b,
-         bias=None,
-         out_dtype=torch.float16,
-     )  # -> [M,128]
-     return outMN[:, :1]  # keep the true N=1 column
+         A_mk, scratch_B128K, scale_a, scale_b, bias=None, out_dtype=torch.float16
+     )  # [M,128]
+     return outMN[:, :1]
  
  def custom_kernel(data):
      """
-     Supports NVFP4 batched GEMV:
-       inputs: (a[M,K,L], b[1,K,L], sfa[M,K//16,L], sfb[1,K//16,L], c[M,1,L])
- 
-     Also gracefully handles 2-tensor practice checks (RGB->Gray) if the runner probes.
+     Supports:
+       • Ranked NVFP4 GEMV: (a[M,K,L], b[1,K,L], sfa[M,K//16,L], sfb[1,K//16,L], c[M,1,L])
+       • Practice probe (2-tensor): (x[H,W,3], out[H,W]) → grayscale
      """
-     # Handle practice/warmup probes that pass (x, out)
+     # ---- 2-tensor practice/warmup probe -----------------------------------
      if len(data) == 2:
          x, out = data
          if x.ndim == 3 and x.shape[-1] == 3:
-             w = torch.tensor([0.2989, 0.5870, 0.1140], device=x.device, dtype=x.dtype)
-             out.copy_(torch.einsum("hwc,c->hw", x, w))
+             global _GRAYSCALE_W
+             if _GRAYSCALE_W is None or _GRAYSCALE_W.device != x.device or _GRAYSCALE_W.dtype != x.dtype:
+                 _GRAYSCALE_W = torch.tensor([0.2989, 0.5870, 0.1140], device=x.device, dtype=x.dtype)
+             out.copy_(torch.einsum("hwc,c->hw", x, _GRAYSCALE_W))
          else:
              out.copy_(x)
          return out
  
-     # Ranked path (5 tensors)
+     # ---- ranked path: 5 tensors -------------------------------------------
      a, b, sfa, sfb, c = data
      M, K, L = a.shape
+ 
+     # minimal sanity checks
      assert b.shape[0] == 1
      assert a.dtype == torch.float4_e2m1fn_x2 and b.dtype == torch.float4_e2m1fn_x2
      assert c.dtype == torch.float16
      assert sfa.dtype in (torch.float8_e4m3fn, getattr(torch, "float8_e4m3fnuz", torch.float8_e4m3fn))
      assert sfb.dtype in (torch.float8_e4m3fn, getattr(torch, "float8_e4m3fnuz", torch.float8_e4m3fn))
  
-     # Preallocate scratch for fallback (N=128). Reused for all L.
      N_PAD = 128
+     # scratch reused across all batches for the padded fallback
      scratch_B128K = torch.empty((N_PAD, K), device=b.device, dtype=b.dtype)
      scratch_SFB128 = torch.empty((N_PAD, K // SF_VEC), device=sfb.device, dtype=sfb.dtype)
  
-     # Try the super-fast N=1 path once; if it errors, use padded path thereafter.
+     # try N=1 once; if it throws, stick to padded fallback
      use_N1 = True
+ 
+     # precompute scale_a (blocked) per batch once, so we don't redo it on both paths
+     scale_a_blocked = [None] * L
      for l in range(L):
-         A_l = a[:, :, l].contiguous()       # [M,K] nvfp4
-         b_l = b[:, :, l].contiguous()       # [1,K] nvfp4
-         sfa_l = sfa[:, :, l].contiguous()   # [M,K//16] fp8
-         sfb_l = sfb[:, :, l].contiguous()   # [1,K//16] fp8
+         scale_a_blocked[l] = _to_blocked(sfa[:, :, l].contiguous())
  
+     for l in range(L):
+         A_l   = a[:, :, l].contiguous()       # [M,K] nvfp4
+         b_l   = b[:, :, l].contiguous()       # [1,K] nvfp4
+         sfa_l = sfa[:, :, l].contiguous()     # [M,K//16] fp8
+         sfb_l = sfb[:, :, l].contiguous()     # [1,K//16] fp8
+ 
          if use_N1:
              try:
-                 outM1 = _scaled_gemv_N1(A_l, b_l, sfa_l, sfb_l)  # [M,1]
+                 outM1 = _scaled_gemv_N1(A_l, b_l, sfa_l, sfb_l, scale_a_blocked=scale_a_blocked[l])
              except Exception:
-                 use_N1 = False  # fallback permanently
+                 use_N1 = False
              else:
                  c[:, 0, l].copy_(outM1[:, 0])
                  continue
  
-         # Fallback: N=128 with reusable scratch
-         outM1 = _scaled_gemv_N128(A_l, b_l, sfa_l, sfb_l, scratch_B128K, scratch_SFB128)
+         outM1 = _scaled_gemv_N128(
+             A_l, b_l, sfa_l, sfb_l,
+             scratch_B128K, scratch_SFB128,
+             scale_a_blocked=scale_a_blocked[l]
+         )
          c[:, 0, l].copy_(outM1[:, 0])
  
      return c