submission 380716

  """
- Outgoing TriMul (AlphaFold‑3) – Triton‑accelerated forward pass
+ Outgoing TriMul (AlphaFold‑3) – Triton forward pass
   
- Algorithm
- ---------
- 1️⃣ Row‑wise LayerNorm over the last axis (C).  The mean/variance are
-    computed in FP32, the normalised output is stored in FP16.
+ The implementation follows the reference PyTorch `TriMul` module but fuses the
+ following steps in custom Triton kernels:
   
- 2️⃣ Fused projection + gating + optional scalar mask.
-    For each flat row r = b·N² + i·N + k we compute
-        left_proj   = x_norm[r]·W_left_proj         (FP16)
-        left_gate   = sigmoid( x_norm[r]·W_left_gate )   (FP32)
-        right_proj  = x_norm[r]·W_right_proj
-        right_gate  = sigmoid( x_norm[r]·W_right_gate )
-        out_gate    = sigmoid( x_norm[r]·W_out_gate )
-    The scalar mask (if present) multiplies the projected tensors.
-    The results are written as
-        left [b, h, i, k] , right [b, h, j, k]  (FP16)
-        out_gate [b, i, j, h]                     (FP16)
+ 1️⃣ Row‑wise LayerNorm over the last dimension (fp16 output, fp32 accumulator).
+ 2️⃣ Fused linear projection, gating (sigmoid) and optional scalar mask.
+ 3️⃣ Batched GEMM (left @ rightᵀ) for the N×N pairwise multiplication.
+ 4️⃣ Fused hidden‑dim LayerNorm → out‑gate multiplication → final linear
+    projection (fp16→fp32 accumulate).
   
- 3️⃣ Batched GEMM (Tensor‑Core) :  left @ rightᵀ  → hidden[b, h, i, j]  (FP16).
-  
- 4️⃣ Fused hidden‑dim LayerNorm → element‑wise out‑gate → final linear.
-    The hidden tensor is normalised across the hidden dimension H
-    (per (b,i,j)).  The normalised tensor is multiplied by the
-    out_gate, then projected back to the original channel dimension C.
-    The final result is FP32.
-  
- All kernels use mixed‑precision (FP16 compute, FP32 accumulation) and
- are tuned for NVIDIA H100.  The only operation left to PyTorch is the
- batched GEMM, which already runs at peak Tensor‑Core efficiency.
+ All kernels use mixed‑precision (fp16 compute, fp32 accumulation) and are
+ tuned for NVIDIA H100 (8 warps for LN, 4 warps for the other kernels).  The
+ final tensor has shape `[B, N, N, dim]` and dtype `torch.float32`.
  """
   
  from typing import Tuple, Dict
+  
  import torch
  import triton
  import triton.language as tl
   
+  
  # ----------------------------------------------------------------------
- # 1) Row‑wise LayerNorm (FP16 output, FP32 accumulator)
+ # 1) Row‑wise LayerNorm (fp16 out, fp32 accumulator)
  # ----------------------------------------------------------------------
  @triton.jit
  def _row_ln_fp16_kernel(
      X_ptr, Y_ptr,          # (M, C) input / output
-     w_ptr, b_ptr,          # LN weight & bias (FP32)
-     M, C: tl.constexpr,    # rows, columns (C must be const)
+     w_ptr, b_ptr,          # LN weight & bias (fp32)
+     M, C: tl.constexpr,    # rows, columns (C must be constexpr)
      eps,
      BLOCK_M: tl.constexpr,
      BLOCK_C: tl.constexpr,
  ⋯ 3 unchanged lines
      rows = row_start + tl.arange(0, BLOCK_M)
      row_mask = rows < M
   
-     # --------------------------------------------------------------
-     # Compute mean & variance (FP32)
-     # --------------------------------------------------------------
+     # ---------- compute mean & variance (fp32) ----------
      sum_val = tl.zeros([BLOCK_M], dtype=tl.float32)
      sumsq_val = tl.zeros([BLOCK_M], dtype=tl.float32)
   
  ⋯ 4 unchanged lines
              X_ptr + rows[:, None] * C + cur_c[None, :],
              mask=row_mask[:, None] & col_mask[None, :],
              other=0.0,
-         ).to(tl.float32)
+         ).to(tl.float32)                     # (BLOCK_M, BLOCK_C)
          sum_val += tl.sum(x, axis=1)
          sumsq_val += tl.sum(x * x, axis=1)
   
  ⋯ 1 unchanged lines
      var = sumsq_val / C - mean * mean
      inv_std = 1.0 / tl.sqrt(var + eps)
   
-     # --------------------------------------------------------------
-     # Normalise + affine (FP16)
-     # --------------------------------------------------------------
+     # ---------- normalize + affine (fp16) ----------
      for c in range(0, C, BLOCK_C):
          cur_c = c + tl.arange(0, BLOCK_C)
          col_mask = cur_c < C
  ⋯ 15 unchanged lines
              mask=row_mask[:, None] & col_mask[None, :],
          )
   
+  
  def _row_layernorm_fp16(
      x: torch.Tensor,
      weight: torch.Tensor,
  ⋯ 26 unchanged lines
   
   
  # ----------------------------------------------------------------------
- # 2) Fused projection + gating + optional mask
+ # 2) Fused projection + gating + optional scalar mask
  # ----------------------------------------------------------------------
  @triton.jit
  def _proj_gate_mask_kernel(
-     x_ptr,                         # (M, C) FP16
-     mask_ptr,                      # (M,) FP16   (if MASKED==1)
-     left_proj_w_ptr,               # (C, H) FP16
-     left_gate_w_ptr,               # (C, H) FP16
-     right_proj_w_ptr,              # (C, H) FP16
-     right_gate_w_ptr,              # (C, H) FP16
-     out_gate_w_ptr,                # (C, H) FP16
-     left_ptr,                      # (B, H, N, N) FP16
-     right_ptr,                     # (B, H, N, N) FP16
-     out_gate_ptr,                  # (B, N, N, H) FP16
+     x_ptr,                         # (M, C) fp16
+     mask_ptr,                      # (M,) fp16  (if MASKED==1)
+     left_proj_w_ptr,               # (C, H) fp16
+     left_gate_w_ptr,               # (C, H) fp16
+     right_proj_w_ptr,              # (C, H) fp16
+     right_gate_w_ptr,              # (C, H) fp16
+     out_gate_w_ptr,                # (C, H) fp16
+     left_ptr,                      # (B, H, N, N) fp16
+     right_ptr,                     # (B, H, N, N) fp16
+     out_gate_ptr,                  # (B, N, N, H) fp16
      M, N, C: tl.constexpr, H: tl.constexpr,
      BLOCK_M: tl.constexpr,
      BLOCK_H: tl.constexpr,
      BLOCK_K: tl.constexpr,
      MASKED: tl.constexpr,
  ):
-     pid_m = tl.program_id(0)   # row block
-     pid_h = tl.program_id(1)   # hidden block
+     pid_m = tl.program_id(0)  # row block
+     pid_h = tl.program_id(1)  # hidden block
   
      row_start = pid_m * BLOCK_M
      hid_start = pid_h * BLOCK_H
  ⋯ 10 unchanged lines
      else:
          mask_val = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
   
-     # ---- accumulators (FP32) ----
+     # ---- accumulators (fp32) ----
      acc_lp = tl.zeros((BLOCK_M, BLOCK_H), dtype=tl.float32)  # left proj
      acc_lg = tl.zeros((BLOCK_M, BLOCK_H), dtype=tl.float32)  # left gate
      acc_rp = tl.zeros((BLOCK_M, BLOCK_H), dtype=tl.float32)  # right proj
  ⋯ 4 unchanged lines
          cur_k = k + tl.arange(0, BLOCK_K)
          k_mask = cur_k < C
   
-         # input tile (FP16 → FP32)
+         # input tile (fp16 → fp32)
          a = tl.load(
              x_ptr + rows[:, None] * C + cur_k[None, :],
              mask=row_mask[:, None] & k_mask[None, :],
              other=0.0,
-         )  # (BLOCK_M, BLOCK_K) FP16
+         )  # (BLOCK_M, BLOCK_K) fp16
   
-         # weight tiles (row‑major (C, H))
+         # weight tiles (C,H) column‑major
          w_lp = tl.load(left_proj_w_ptr + cur_k[:, None] * H + hids[None, :],
                         mask=k_mask[:, None] & hid_mask[None, :],
                         other=0.0)
  ⋯ 10 unchanged lines
                         mask=k_mask[:, None] & hid_mask[None, :],
                         other=0.0)
   
-         # FP16·FP16 → FP32 dot products
+         # fp16·fp16 → fp32 dot products
          acc_lp += tl.dot(a, w_lp)
          acc_lg += tl.dot(a, w_lg)
          acc_rp += tl.dot(a, w_rp)
          acc_rg += tl.dot(a, w_rg)
          acc_og += tl.dot(a, w_og)
   
-     # ---- sigmoid (FP32) ----
+     # ---- sigmoid (fp32) ----
      left_gate  = 1.0 / (1.0 + tl.exp(-acc_lg))
      right_gate = 1.0 / (1.0 + tl.exp(-acc_rg))
      out_gate   = 1.0 / (1.0 + tl.exp(-acc_og))
   
-     # ---- apply scalar mask and per‑row gates ----
+     # ---- apply mask and per‑row gates ----
      left_out  = acc_lp * left_gate * mask_val[:, None]
      right_out = acc_rp * right_gate * mask_val[:, None]
   
  ⋯ 32 unchanged lines
  # ----------------------------------------------------------------------
  @triton.jit
  def _ln_gate_out_linear_fused_kernel(
-     hidden_ptr,           # (B*H*N*N,) FP16 flattened
-     out_gate_ptr,         # (B*N*N*H,) FP16 flattened
-     ln_w_ptr, ln_b_ptr,  # (H,) FP32
-     w_out_ptr,            # (H, D) FP16
-     out_ptr,              # (B, N, N, D) FP32
+     hidden_ptr,           # (B*H*N*N,) fp16 flattened
+     out_gate_ptr,         # (B*N*N*H,) fp16 flattened
+     ln_w_ptr, ln_b_ptr,  # (H,) fp32
+     w_out_ptr,            # (H, D) fp16
+     out_ptr,              # (B, N, N, D) fp32
      B, N, H, D: tl.constexpr,
      eps: tl.constexpr,
      BLOCK_M: tl.constexpr,
  ⋯ 7 unchanged lines
   
      N_sq = N * N
      b_idx = rows // N_sq
-     rem   = rows - b_idx * N_sq
+     rem = rows - b_idx * N_sq
      i_idx = rem // N
      j_idx = rem - i_idx * N
   
-     # --------------------------------------------------------------
-     # Load hidden tensor tile (BLOCK_M, BLOCK_H)
-     # --------------------------------------------------------------
+     # hidden tile (BLOCK_M, BLOCK_H)
      hids = tl.arange(0, BLOCK_H)
      hid_mask = hids < H
   
  ⋯ 2 unchanged lines
          hidden_ptr + hidden_off,
          mask=row_mask[:, None] & hid_mask[None, :],
          other=0.0,
-     )  # FP16
+     )  # fp16
+  
      hidden_fp32 = hidden_tile.to(tl.float32)
   
-     # --------------------------------------------------------------
-     # Mean / variance across H (FP32)
-     # --------------------------------------------------------------
-     sum_val = tl.sum(hidden_fp32, axis=1)
-     sumsq_val = tl.sum(hidden_fp32 * hidden_fp32, axis=1)
+     # ---- mean / variance across H (fp32) ----
+     sum_val = tl.sum(hidden_fp32, axis=1)                     # (BLOCK_M,)
+     sumsq_val = tl.sum(hidden_fp32 * hidden_fp32, axis=1)    # (BLOCK_M,)
      mean = sum_val / H
      var = sumsq_val / H - mean * mean
-     inv_std = 1.0 / tl.sqrt(var + eps)                     # (BLOCK_M,)
+     inv_std = 1.0 / tl.sqrt(var + eps)                       # (BLOCK_M,)
   
-     # --------------------------------------------------------------
-     # Layer‑norm (FP32) + affine
-     # --------------------------------------------------------------
-     w_ln = tl.load(ln_w_ptr + hids, mask=hid_mask, other=0.0)   # (H,)
-     b_ln = tl.load(ln_b_ptr + hids, mask=hid_mask, other=0.0)   # (H,)
+     # ---- layer‑norm (fp32) ----
+     w_ln = tl.load(ln_w_ptr + hids, mask=hid_mask, other=0.0)  # (H,)
+     b_ln = tl.load(ln_b_ptr + hids, mask=hid_mask, other=0.0)  # (H,)
      hidden_norm = (hidden_fp32 - mean[:, None]) * inv_std[:, None]
-     hidden_norm = hidden_norm * w_ln[None, :] + b_ln[None, :]
+     hidden_norm = hidden_norm * w_ln[None, :] + b_ln[None, :]   # (BLOCK_M, BLOCK_H)
   
-     # --------------------------------------------------------------
-     # Apply out‑gate (FP32)
-     # --------------------------------------------------------------
+     # ---- out‑gate (fp32) ----
      out_gate_off = rows[:, None] * H + hids[None, :]
      out_gate_tile = tl.load(
          out_gate_ptr + out_gate_off,
          mask=row_mask[:, None] & hid_mask[None, :],
          other=0.0,
-     ).to(tl.float32)
+     ).to(tl.float32)                                          # (BLOCK_M, BLOCK_H)
   
-     gated = hidden_norm * out_gate_tile
+     gated = hidden_norm * out_gate_tile                        # (BLOCK_M, BLOCK_H)
   
-     # --------------------------------------------------------------
-     # Final linear projection (H → D)
-     # --------------------------------------------------------------
+     # Convert to fp16 for the final matrix‑multiply (TensorCore friendly)
      gated_fp16 = gated.to(tl.float16)
+  
+     # ---- final linear projection (fp32) ----
      for d0 in range(0, D, BLOCK_D):
          cols = d0 + tl.arange(0, BLOCK_D)
          col_mask = cols < D
  ⋯ 1 unchanged lines
              w_out_ptr + hids[:, None] * D + cols[None, :],
              mask=hid_mask[:, None] & col_mask[None, :],
              other=0.0,
-         )  # (H, BLOCK_D) FP16
+         )  # (BLOCK_H, BLOCK_D) fp16
   
-         out = tl.dot(gated_fp16, w_out)                # (BLOCK_M, BLOCK_D) FP32
+         out = tl.dot(gated_fp16, w_out)                      # (BLOCK_M, BLOCK_D) fp32
+  
          tl.store(
              out_ptr + rows[:, None] * D + cols[None, :],
              out,
  ⋯ 2 unchanged lines
   
   
  # ----------------------------------------------------------------------
- # Entry point
+ # 4) Entry point
  # ----------------------------------------------------------------------
  def custom_kernel(
      data: Tuple[torch.Tensor, torch.Tensor, Dict[str, torch.Tensor], Dict]
  ⋯ 38 unchanged lines
      )  # (B, N, N, C) fp16
   
      # --------------------------------------------------------------
-     # 2) Prepare projection / gate weights (C, H) fp16, row‑major
+     # 2) Prepare projection / gate weights (C, H) in fp16, column‑major
      # --------------------------------------------------------------
      left_proj_w_T  = weights["left_proj.weight"].t().contiguous().to(torch.float16)
      right_proj_w_T = weights["right_proj.weight"].t().contiguous().to(torch.float16)
  ⋯ 2 unchanged lines
      out_gate_w_T   = weights["out_gate.weight"].t().contiguous().to(torch.float16)
   
      # --------------------------------------------------------------
-     # 3) Mask handling – optional scalar mask per (b,i,k)
+     # 3) Mask handling (flattened) – optional
      # --------------------------------------------------------------
      if not nomask and mask is not None:
          mask_flat = mask.reshape(M).to(torch.float16).contiguous()
  ⋯ 3 unchanged lines
          MASKED = 0
   
      # --------------------------------------------------------------
-     # 4) Allocate buffers for fused projection / gating kernel
+     # 4) Allocate buffers for the fused projection / gating kernel
      # --------------------------------------------------------------
      left = torch.empty((B, hidden_dim, N, N), dtype=torch.float16, device=device)
      right = torch.empty_like(left)
      out_gate = torch.empty((B, N, N, hidden_dim), dtype=torch.float16, device=device)
   
      # --------------------------------------------------------------
-     # 5) Fused projection + gating (+ optional mask)
+     # 5) Fused projection + gating + (optional) mask
      # --------------------------------------------------------------
-     BLOCK_M = 64      # number of flat rows per program (covers B·N²)
-     BLOCK_H = 64      # hidden‑dim block size
-     BLOCK_K = 32      # inner‑channel tile size
+     BLOCK_M = 64      # rows per program (B·N·N)
+     BLOCK_H = 64      # hidden‑dim block
+     BLOCK_K = 32      # input‑channel tile size
   
      grid_proj = (triton.cdiv(M, BLOCK_M), triton.cdiv(hidden_dim, BLOCK_H))
      _proj_gate_mask_kernel[grid_proj](
  ⋯ 21 unchanged lines
      # --------------------------------------------------------------
      # 6) Pairwise multiplication (batched GEMM)
      # --------------------------------------------------------------
-     # left/right shapes: (B, H, N, N) → view as (B*H, N, N)
-     # torch.bmm will use Tensor‑Core fp16 matmul on H100
-     left_mat = left.view(B * hidden_dim, N, N)                     # (B*H, N, N)
-     right_mat = right.view(B * hidden_dim, N, N).transpose(1, 2)   # (B*H, N, N)
-     hidden_fp16 = torch.bmm(left_mat, right_mat)                   # (B*H, N, N) fp16
-     hidden = hidden_fp16.view(B, hidden_dim, N, N)                 # (B, H, N, N) fp16
+     left_mat = left.view(B * hidden_dim, N, N)                 # (B*H, N, N)
+     right_mat = right.view(B * hidden_dim, N, N).transpose(1, 2)  # (B*H, N, N)
+     hidden_fp16 = torch.bmm(left_mat, right_mat)                # (B*H, N, N) fp16
+     hidden = hidden_fp16.view(B, hidden_dim, N, N)               # (B, H, N, N) fp16
   
      # --------------------------------------------------------------
      # 7) Fused hidden‑dim LN → out‑gate → final linear
      # --------------------------------------------------------------
      to_out_norm_w = weights["to_out_norm.weight"]   # (H,) fp32
      to_out_norm_b = weights["to_out_norm.bias"]    # (H,) fp32
-     to_out_w_T = weights["to_out.weight"].t().contiguous().to(torch.float16)  # (H, C)
+     to_out_w_T = weights["to_out.weight"].t().contiguous().to(torch.float16)   # (H, C)
   
      out = torch.empty((B, N, N, dim), dtype=torch.float32, device=device)
   
      BLOCK_M_OUT = 64
-     BLOCK_D_OUT = 64                     # D‑tile for final projection
+     BLOCK_H_OUT = 128   # covers the whole hidden dimension (H≤128)
+     BLOCK_D_OUT = 64
   
      grid_out = (triton.cdiv(B * N * N, BLOCK_M_OUT),)
+  
      _ln_gate_out_linear_fused_kernel[grid_out](
          hidden.view(-1),                     # flat fp16 hidden
          out_gate.view(-1),                   # flat fp16 out‑gate
  ⋯ 7 unchanged lines
          dim,
          eps,
          BLOCK_M=BLOCK_M_OUT,
-         BLOCK_H=hidden_dim,          # H fits into one block for our configs
+         BLOCK_H=BLOCK_H_OUT,
          BLOCK_D=BLOCK_D_OUT,
          num_warps=4,
      )