submission 553897

  ⋯ 8 unchanged lines
  # Autotune locally for each shape, then paste the best config here.
  SHAPE_CONFIGS: dict[tuple, helion.Config] = {
      # Test shapes
-     (1, 64, 2, 64, 64): helion.Config(block_sizes=[], num_warps=1, num_stages=1),  # TODO: use any config that passes correctness check
-     (2, 128, 4, 64, 64): helion.Config(block_sizes=[], num_warps=1, num_stages=1),  # TODO: use any config that passes correctness check
-     (1, 256, 4, 64, 128): helion.Config(block_sizes=[], num_warps=1, num_stages=1),  # TODO: use any config that passes correctness check
+     (1, 64, 2, 64, 64): helion.Config(block_sizes=[], num_stages=1, num_warps=8, pid_type='flat'),
+     (2, 128, 4, 64, 64): helion.Config(block_sizes=[], num_stages=1, num_warps=8, pid_type='flat'),
+     (1, 256, 4, 64, 128): helion.Config(block_sizes=[], num_stages=1, num_warps=8, pid_type='flat'),
      # Benchmark shapes
-     (1, 64, 1, 64, 64): helion.Config(block_sizes=[], num_warps=1, num_stages=1),  # TODO: replace with your autotuned config
-     (2, 512, 3, 64, 64): helion.Config(block_sizes=[], num_warps=1, num_stages=1),  # TODO: replace with your autotuned config
-     (2, 1024, 3, 64, 64): helion.Config(block_sizes=[], num_warps=1, num_stages=1),  # TODO: replace with your autotuned config
-     (3, 1024, 4, 100, 100): helion.Config(block_sizes=[], num_warps=1, num_stages=1),  # TODO: replace with your autotuned config
-     (4, 1024, 4, 128, 128): helion.Config(block_sizes=[], num_warps=1, num_stages=1),  # TODO: replace with your autotuned config
-     (2, 1536, 4, 128, 128): helion.Config(block_sizes=[], num_warps=1, num_stages=1),  # TODO: replace with your autotuned config
-     (4, 2048, 8, 64, 64): helion.Config(block_sizes=[], num_warps=1, num_stages=1),  # TODO: replace with your autotuned config
+     (1, 64, 1, 64, 64): helion.Config(block_sizes=[], num_stages=1, num_warps=8, pid_type='flat'),
+     (2, 512, 3, 64, 64): helion.Config(block_sizes=[], num_stages=1, num_warps=8, pid_type='flat'),
+     (2, 1024, 3, 64, 64): helion.Config(block_sizes=[], num_stages=1, num_warps=8, pid_type='flat'),
+     (3, 1024, 4, 100, 100): helion.Config(block_sizes=[], num_stages=1, num_warps=16, pid_type='flat'),
+     (4, 1024, 4, 128, 128): helion.Config(block_sizes=[], loop_orders=[[1, 0]], num_stages=1, num_warps=32, pid_type='flat'),
+     (2, 1536, 4, 128, 128): helion.Config(block_sizes=[], loop_orders=[[1, 0]], num_stages=1, num_warps=32, pid_type='flat'),
+     (4, 2048, 8, 64, 64): helion.Config(block_sizes=[], num_stages=1, num_warps=8, pid_type='flat'),
  }
   
   
  ⋯ 2 unchanged lines
  #     helion.Config(..., advanced_controls_file="/opt/booster_pack/recompute_w_u_fwd_0.acf")
   
   
- # NOTE: This is an intentionally inefficient baseline implementation.
  def _make_kernel(config: helion.Config):
      @helion.kernel(static_shapes=True, dot_precision="ieee", config=config)
      def kernel(
  ⋯ 4 unchanged lines
          g: torch.Tensor,     # [B, T, H]
      ) -> tuple[torch.Tensor, torch.Tensor]:
          B, T, H, K = k.shape
-         V = v.shape[-1]
+         V = hl.specialize(v.shape[-1])
          C = hl.specialize(A.shape[-1])
          K = hl.specialize(K)
-         V = hl.specialize(V)
   
          w_out = torch.empty_like(k)
          u_out = torch.empty_like(v)
  ⋯ 3 unchanged lines
              b_idx = flat_bh.begin // H
              h_idx = flat_bh.begin % H
   
-             w_acc1 = hl.zeros([rt, K], dtype=torch.float32)
-             u_acc1 = hl.zeros([rt, V], dtype=torch.float32)
-             w_acc2 = hl.zeros([rt, K], dtype=torch.float32)
-             u_acc2 = hl.zeros([rt, V], dtype=torch.float32)
+             A_chunk = A[b_idx, rt, h_idx, :].to(torch.float32)
+             k_chunk = k[b_idx, rt, h_idx, :].to(torch.float32)
+             v_chunk = v[b_idx, rt, h_idx, :].to(torch.float32)
+             beta_chunk = beta[b_idx, rt, h_idx].to(torch.float32)
+             g_chunk = g[b_idx, rt, h_idx].to(torch.float32)
   
-             for ci in range(C):
-                 t_ci = rt.begin + ci
-                 a_col = A[b_idx, rt, h_idx, ci].to(torch.float32)
-                 coeff_ci = beta[b_idx, t_ci, h_idx].to(torch.float32)
-                 decay_ci = torch.exp(g[b_idx, t_ci, h_idx].to(torch.float32))
+             v_scaled = v_chunk * beta_chunk[:, None]
+             k_scaled = k_chunk * (beta_chunk * torch.exp(g_chunk))[:, None]
   
-                 k_ci = k[b_idx, t_ci, h_idx, :].to(torch.float32)
-                 v_ci = v[b_idx, t_ci, h_idx, :].to(torch.float32)
+             u = hl.dot(A_chunk, v_scaled, out_dtype=torch.float32)
+             w = hl.dot(A_chunk, k_scaled, out_dtype=torch.float32)
   
-                 w_acc1 = w_acc1 + a_col[:, None] * (k_ci * coeff_ci * decay_ci)[None, :]
-                 u_acc1 = u_acc1 + a_col[:, None] * (v_ci * coeff_ci)[None, :]
+             w_out[b_idx, rt, h_idx, :] = w.to(k.dtype)
+             u_out[b_idx, rt, h_idx, :] = u.to(v.dtype)
   
-             for ci in range(C - 1, -1, -1):
-                 t_ci = rt.begin + ci
-                 a_col = A[b_idx, rt, h_idx, ci].to(torch.float32)
-                 coeff_ci = beta[b_idx, t_ci, h_idx].to(torch.float32)
-                 decay_ci = torch.exp(g[b_idx, t_ci, h_idx].to(torch.float32))
-  
-                 k_ci = k[b_idx, t_ci, h_idx, :].to(torch.float32)
-                 v_ci = v[b_idx, t_ci, h_idx, :].to(torch.float32)
-  
-                 w_acc2 = w_acc2 + a_col[:, None] * (k_ci * coeff_ci * decay_ci)[None, :]
-                 u_acc2 = u_acc2 + a_col[:, None] * (v_ci * coeff_ci)[None, :]
-  
-             w_out[b_idx, rt, h_idx, :] = ((w_acc1 + w_acc2) * 0.5).to(k.dtype)
-             u_out[b_idx, rt, h_idx, :] = ((u_acc1 + u_acc2) * 0.5).to(v.dtype)
-  
          return w_out, u_out
   
      return kernel
   
   
- _KERNELS = {shape: _make_kernel(cfg) for shape, cfg in SHAPE_CONFIGS.items()}
+ _KERNEL_CACHE: dict[tuple, object] = {}
   
   
  def custom_kernel(data: input_t) -> output_t:
      k, v, beta, A, g = data
      B, T, H, K = k.shape
      V = v.shape[-1]
-     kernel = _KERNELS[(B, T, H, K, V)]
+     key = (B, T, H, K, V)
+     if key not in _KERNEL_CACHE:
+         _KERNEL_CACHE[key] = _make_kernel(SHAPE_CONFIGS[key])
+     kernel = _KERNEL_CACHE[key]
      return kernel(k, v, beta, A, g)