submission 408335

  ⋯ 3 unchanged lines
   
   
  @triton.jit
- def layernorm_kernel_v11(
-     X, LN_W, LN_B, Out,
-     stride_xb, stride_xn1, stride_xn2, stride_xc,
-     B, N, C, 
-     BLOCK_N: tl.constexpr, BLOCK_C: tl.constexpr
- ):
+ def layernorm_kernel(X, LN_W, LN_B, Out,
+                      stride_xb, stride_xn1, stride_xn2, stride_xc,
+                      B, N, C, 
+                      BLOCK_N: tl.constexpr, BLOCK_C: tl.constexpr):
      pid_b, pid_n1 = tl.program_id(0), tl.program_id(1)
      pid_n2_start = tl.program_id(2) * BLOCK_N
      offs_n2 = pid_n2_start + tl.arange(0, BLOCK_N)
  ⋯ 26 unchanged lines
   
   
  @triton.jit
- def projection_kernel_v7(
+ def projection_kernel_optimized(
      X_norm, Mask, W_concat,
-     L_out, R_out, OG_out,
+     L_out, R_out,
      stride_xb, stride_xn1, stride_xn2, stride_xc,
      stride_mb, stride_mn1, stride_mn2,
      stride_ob, stride_od, stride_on1, stride_on2,
  ⋯ 9 unchanged lines
      acc_r = tl.zeros([BLOCK_N, D], dtype=tl.float32)
      acc_lg = tl.zeros([BLOCK_N, D], dtype=tl.float32)
      acc_rg = tl.zeros([BLOCK_N, D], dtype=tl.float32)
-     acc_og = tl.zeros([BLOCK_N, D], dtype=tl.float32)
   
      offs_d = tl.arange(0, D)
-  
-     # Use block pointers for X_norm to improve memory access efficiency
+     
+     # Pre-load weights into registers/shared memory if possible, but D=128 is large.
+     # We use block pointers for X_norm to improve coalescing.
      x_block_ptr = tl.make_block_ptr(
          base=X_norm + pid_b * stride_xb + pid_n1 * stride_xn1,
          shape=(N, C),
  ⋯ 5 unchanged lines
   
      for c_start in range(0, C, BLOCK_C):
          x = tl.load(x_block_ptr, boundary_check=(0, 1)).to(tl.float16)
-         
          offs_c = c_start + tl.arange(0, BLOCK_C)
          mask_c = offs_c < C
   
-         # Weights are [5*D, C]. We load chunks of [D, BLOCK_C] and transpose for dot.
-         # This is more efficient than standard pointer indexing.
+         # Load weights for all 4 projections
          w_l = tl.load(W_concat + (0*D + offs_d)[:, None] * C + offs_c[None, :], mask=mask_c[None, :], other=0.0).to(tl.float16)
          w_r = tl.load(W_concat + (1*D + offs_d)[:, None] * C + offs_c[None, :], mask=mask_c[None, :], other=0.0).to(tl.float16)
          w_lg = tl.load(W_concat + (2*D + offs_d)[:, None] * C + offs_c[None, :], mask=mask_c[None, :], other=0.0).to(tl.float16)
          w_rg = tl.load(W_concat + (3*D + offs_d)[:, None] * C + offs_c[None, :], mask=mask_c[None, :], other=0.0).to(tl.float16)
-         w_og = tl.load(W_concat + (4*D + offs_d)[:, None] * C + offs_c[None, :], mask=mask_c[None, :], other=0.0).to(tl.float16)
   
          acc_l += tl.dot(x, tl.trans(w_l))
          acc_r += tl.dot(x, tl.trans(w_r))
          acc_lg += tl.dot(x, tl.trans(w_lg))
          acc_rg += tl.dot(x, tl.trans(w_rg))
-         acc_og += tl.dot(x, tl.trans(w_og))
-         
          x_block_ptr = tl.advance(x_block_ptr, [0, BLOCK_C])
   
      m_ptr = Mask + pid_b * stride_mb + pid_n1 * stride_mn1 + offs_n2 * stride_mn2
  ⋯ 1 unchanged lines
   
      l_final = acc_l * (mask_val[:, None] * tl.sigmoid(acc_lg))
      r_final = acc_r * (mask_val[:, None] * tl.sigmoid(acc_rg))
-     og_final = tl.sigmoid(acc_og)
   
+     # Store in [B, D, N, N] format for optimal matmul
      out_off = pid_b * stride_ob + offs_d[:, None] * stride_od + pid_n1 * stride_on1 + offs_n2[None, :] * stride_on2
      tl.store(L_out + out_off, tl.trans(l_final).to(tl.float16), mask=mask_n2[None, :])
      tl.store(R_out + out_off, tl.trans(r_final).to(tl.float16), mask=mask_n2[None, :])
-     tl.store(OG_out + out_off, tl.trans(og_final).to(tl.float16), mask=mask_n2[None, :])
   
   
  @triton.jit
- def post_process_kernel_v62(
-     matmul_out, OG_p, LN_W, LN_B, TO_OUT_W_T, Out,
+ def post_process_optimized_kernel(
+     matmul_out, X_norm, OG_W, LN_W, LN_B, TO_OUT_W_T, Out,
      stride_mb, stride_md, stride_mn, stride_mm,
+     stride_xb, stride_xn1, stride_xn2, stride_xc,
      stride_ob, stride_on1, stride_on2, stride_oc,
      B, N, D: tl.constexpr, C: tl.constexpr,
      BLOCK_N: tl.constexpr, BLOCK_C: tl.constexpr
  ⋯ 4 unchanged lines
      mask_n2 = offs_n2 < N
      offs_d = tl.arange(0, D)
   
-     m_off = pid_b * stride_mb + offs_d[None, :] * stride_md + pid_n1 * stride_mn + offs_n2[:, None] * stride_mm
-     x = tl.load(matmul_out + m_off, mask=mask_n2[:, None], other=0.0).to(tl.float32)
-     og = tl.load(OG_p + m_off, mask=mask_n2[:, None], other=0.0).to(tl.float32)
+     # 1. Load Matmul Output [BLOCK_N, D]
+     m_block_ptr = tl.make_block_ptr(
+         base=matmul_out + pid_b * stride_mb + pid_n1 * stride_mn,
+         shape=(D, N),
+         strides=(stride_md, stride_mm),
+         offsets=(0, pid_n2_start),
+         block_shape=(D, BLOCK_N),
+         order=(0, 1)
+     )
+     x = tl.trans(tl.load(m_block_ptr, boundary_check=(1,))).to(tl.float32)
   
+     # 2. Recompute Out-Gate [BLOCK_N, D]
+     acc_og = tl.zeros([BLOCK_N, D], dtype=tl.float32)
+     for c_start in range(0, C, BLOCK_C):
+         offs_c = c_start + tl.arange(0, BLOCK_C)
+         mask_c = offs_c < C
+         xn = tl.load(X_norm + pid_b * stride_xb + pid_n1 * stride_xn1 + offs_n2[:, None] * stride_xn2 + offs_c[None, :], mask=(mask_n2[:, None] & mask_c[None, :]), other=0.0).to(tl.float16)
+         w_og = tl.load(OG_W + offs_d[:, None] * C + offs_c[None, :], mask=mask_c[None, :], other=0.0).to(tl.float16)
+         acc_og += tl.dot(xn, tl.trans(w_og))
+     og = tl.sigmoid(acc_og)
+  
+     # 3. LayerNorm on Matmul Output
      mean = tl.sum(x, axis=1) / D
      diff = x - mean[:, None]
      var = tl.sum(diff * diff, axis=1) / D
      rstd = 1.0 / tl.sqrt(var + 1e-5)
-     
      ln_w = tl.load(LN_W + offs_d)
      ln_b = tl.load(LN_B + offs_d)
-     x_norm = ((diff * rstd[:, None]) * ln_w[None, :] + ln_b[None, :]) * og
-     x_norm_f16 = x_norm.to(tl.float16)
+     x_normed = ((diff * rstd[:, None]) * ln_w[None, :] + ln_b[None, :]) * og
+     x_normed_f16 = x_normed.to(tl.float16)
   
-     # Using pre-transposed weights for the final projection
+     # 4. Final Projection to C - use larger BLOCK_C if possible
      for c_start in range(0, C, BLOCK_C):
          offs_c = c_start + tl.arange(0, BLOCK_C)
          mask_c = offs_c < C
-         # TO_OUT_W_T is [D, C]
          w = tl.load(TO_OUT_W_T + offs_d[:, None] * C + offs_c[None, :], mask=mask_c[None, :], other=0.0).to(tl.float16)
-         out_chunk = tl.dot(x_norm_f16, w)
+         out_chunk = tl.dot(x_normed_f16, w)
          tl.store(Out + pid_b * stride_ob + pid_n1 * stride_on1 + offs_n2[:, None] * stride_on2 + offs_c[None, :] * stride_oc, out_chunk.to(tl.float32), mask=(mask_n2[:, None] & mask_c[None, :]))
   
   
  ⋯ 3 unchanged lines
      D = config["hidden_dim"]
      device = x.device
   
+     # 1. LayerNorm
      x_norm = torch.empty_like(x, dtype=torch.float16)
-     layernorm_kernel_v11[(B, N, (N + 31) // 32)](
+     layernorm_kernel[(B, N, (N + 31) // 32)](
          x, weights["norm.weight"], weights["norm.bias"], x_norm,
          x.stride(0), x.stride(1), x.stride(2), x.stride(3),
          B, N, C, BLOCK_N=32, BLOCK_C=128, num_warps=4
      )
   
+     # 2. Projections
      W_concat = torch.cat([
          weights["left_proj.weight"], weights["right_proj.weight"], 
-         weights["left_gate.weight"], weights["right_gate.weight"], weights["out_gate.weight"]
+         weights["left_gate.weight"], weights["right_gate.weight"]
      ], dim=0).to(device=device, dtype=torch.float16)
   
      L_p = torch.empty((B, D, N, N), device=device, dtype=torch.float16)
      R_p = torch.empty((B, D, N, N), device=device, dtype=torch.float16)
-     OG_p = torch.empty((B, D, N, N), device=device, dtype=torch.float16)
      
      BLOCK_N_PROJ = 64
-     # Use num_stages=2 to avoid OutOfResources error
-     projection_kernel_v7[(B, N, (N + BLOCK_N_PROJ - 1) // BLOCK_N_PROJ)](
+     projection_kernel_optimized[(B, N, (N + BLOCK_N_PROJ - 1) // BLOCK_N_PROJ)](
          x_norm, mask, W_concat,
-         L_out=L_p, R_out=R_p, OG_out=OG_p,
+         L_out=L_p, R_out=R_p,
          stride_xb=x_norm.stride(0), stride_xn1=x_norm.stride(1), stride_xn2=x_norm.stride(2), stride_xc=x_norm.stride(3),
          stride_mb=mask.stride(0), stride_mn1=mask.stride(1), stride_mn2=mask.stride(2),
          stride_ob=L_p.stride(0), stride_od=L_p.stride(1), stride_on1=L_p.stride(2), stride_on2=L_p.stride(3),
-         B=B, N=N, C=C, D=D, BLOCK_N=BLOCK_N_PROJ, BLOCK_C=64, num_warps=8, num_stages=2
+         B=B, N=N, C=C, D=D, BLOCK_N=BLOCK_N_PROJ, BLOCK_C=64, num_warps=8, num_stages=3
      )
   
+     # 3. Matmul
      matmul_out = torch.matmul(L_p, R_p.transpose(-1, -2))
-  
-     # Pre-transpose final weight for better loading in kernel
+     
+     # 4. Post-process
      to_out_w_t = weights["to_out.weight"].t().contiguous().to(device=device, dtype=torch.float16)
-  
+     og_w = weights["out_gate.weight"].to(device=device, dtype=torch.float16)
      out = torch.empty((B, N, N, C), device=device, dtype=torch.float32)
+     
      BLOCK_N_POST = 64
-     post_process_kernel_v62[(B, N, (N + BLOCK_N_POST - 1) // BLOCK_N_POST)](
-         matmul_out, OG_p, weights["to_out_norm.weight"], weights["to_out_norm.bias"], to_out_w_t, out,
+     post_process_optimized_kernel[(B, N, (N + BLOCK_N_POST - 1) // BLOCK_N_POST)](
+         matmul_out, x_norm, og_w, weights["to_out_norm.weight"], weights["to_out_norm.bias"], to_out_w_t, out,
          matmul_out.stride(0), matmul_out.stride(1), matmul_out.stride(2), matmul_out.stride(3),
+         x_norm.stride(0), x_norm.stride(1), x_norm.stride(2), x_norm.stride(3),
          out.stride(0), out.stride(1), out.stride(2), out.stride(3),
-         B, N, D, C, BLOCK_N=BLOCK_N_POST, BLOCK_C=64, num_warps=8
+         B, N, D, C, BLOCK_N=BLOCK_N_POST, BLOCK_C=64, num_warps=8, num_stages=2
      )
   
      return out