submission 555000

  ⋯ 7 unchanged lines
  import helion.language as hl
   
   
- # Per-shape configs: map (B, T, H, K, V) to optimized helion.Config objects.
- # Autotuned config from B200
- _TUNED = helion.Config(block_sizes=[], indexing=['pointer', 'pointer', 'pointer', 'pointer', 'tensor_descriptor', 'pointer', 'tensor_descriptor'], l2_groupings=[1], load_eviction_policies=['first', '', '', 'first', 'first'], loop_orders=[[0, 1]], num_stages=1, num_warps=4, pid_type='flat', range_flattens=[None, None], range_multi_buffers=[None, None], range_num_stages=[0, 0], range_unroll_factors=[0, 0], range_warp_specializes=[None, None])
+ # Based on ramizik's config (range_num_stages=[0,3] is key) + our optimizations:
+ # tf32 dots, exp2, acc=state fusion, diff gating
+ _LOG2E = 1.4426950408889634
   
+ _TUNED = helion.Config(block_sizes=[], indexing=['tensor_descriptor', 'pointer', 'tensor_descriptor', 'tensor_descriptor', 'pointer', 'tensor_descriptor', 'pointer'], l2_groupings=[1], load_eviction_policies=['first', '', '', '', ''], loop_orders=[[1, 0]], num_stages=3, num_warps=4, pid_type='flat', range_flattens=[None, True], range_multi_buffers=[None, False], range_num_stages=[0, 3], range_unroll_factors=[0, 0], range_warp_specializes=[None, None])
+ # BH=1 shapes: higher warps for occupancy
+ _TINY = helion.Config(block_sizes=[], indexing=['tensor_descriptor', 'pointer', 'tensor_descriptor', 'tensor_descriptor', 'pointer', 'tensor_descriptor', 'pointer'], l2_groupings=[1], load_eviction_policies=['first', '', '', '', ''], loop_orders=[[1, 0]], num_stages=3, num_warps=16, pid_type='flat', range_flattens=[None, True], range_multi_buffers=[None, False], range_num_stages=[0, 3], range_unroll_factors=[0, 0], range_warp_specializes=[None, None])
+  
  SHAPE_CONFIGS: dict[tuple, helion.Config] = {
      # Test shapes
-     (1, 64, 2, 64, 64): _TUNED,
+     (1, 64, 2, 64, 64): _TINY,
      (2, 128, 4, 64, 64): _TUNED,
      (1, 256, 4, 64, 128): _TUNED,
      # Benchmark shapes
-     (1, 64, 1, 64, 64): _TUNED,
+     (1, 64, 1, 64, 64): _TINY,     # BH=1: higher warps
      (2, 512, 3, 64, 64): _TUNED,
      (2, 1024, 3, 64, 64): _TUNED,
-     (3, 1024, 4, 100, 100): _TUNED,
-     (4, 1024, 4, 128, 128): _TUNED,
-     (2, 1536, 4, 128, 128): _TUNED,
-     (4, 2048, 8, 64, 64): _TUNED,
  }
   
   
  def _make_kernel(config: helion.Config):
-     @helion.kernel(static_shapes=True, dot_precision="ieee", config=config)
+     @helion.kernel(static_shapes=True, dot_precision="tf32", config=config)
      def kernel(
          k: torch.Tensor,   # [B, T, H, K]
          w: torch.Tensor,   # [B, T, H, K]
  ⋯ 9 unchanged lines
          NT = (T + C - 1) // C
          h_out = torch.empty(B, NT, H, K, V, dtype=k.dtype, device=k.device)
          v_out = torch.empty_like(u)
-  
          BH = B * H
   
          for flat, tv in hl.tile([BH, V], block_size=[1, 8]):
  ⋯ 5 unchanged lines
                  chunk_idx = tc.begin // C
                  t_end = min(tc.begin + C, T) - 1
   
-                 # Store current state as h_out for this chunk
                  h_out[b_idx, chunk_idx, h_idx, :, tv] = state.to(k.dtype)
   
-                 # Compute v_new = u - w @ state
                  proj = hl.dot(
                      w[b_idx, tc, h_idx, :], state, out_dtype=torch.float32
                  )
                  diff = u[b_idx, tc, h_idx, tv].to(torch.float32) - proj
                  v_out[b_idx, tc, h_idx, tv] = diff.to(u.dtype)
   
-                 # Update state: state = state * exp(g_end) + k^T @ (v_new * gate)
                  g_end = g[b_idx, t_end, h_idx].to(torch.float32)
                  g_t = g[b_idx, tc, h_idx].to(torch.float32)
                  valid = tc.index < T
-                 alpha = torch.where(valid, torch.exp(g_end - g_t), 0.0)
-                 k_adj = k[b_idx, tc, h_idx, :] * alpha[:, None]
+                 alpha = torch.where(valid, torch.exp2((g_end - g_t) * _LOG2E), 0.0)
+                 diff_gated = diff * alpha[:, None]
   
-                 state = state * torch.exp(g_end)
-                 upd = hl.dot(k_adj.T, diff, out_dtype=torch.float32)
-                 state = state + upd
+                 state = state * torch.exp2(g_end * _LOG2E)
+                 state = hl.dot(
+                     k[b_idx, tc, h_idx, :].T, diff_gated, acc=state, out_dtype=torch.float32
+                 )
   
          return h_out, v_out