submission 371064

  """
- Outgoing TriMul (AlphaFold‑3) – Triton‑accelerated forward pass
+ Outgoing TriMul (AlphaFold‑3) – Triton accelerated forward pass.
   
- The implementation follows the reference PyTorch module but fuses the
- major compute stages into three Triton kernels:
+ The implementation follows the reference ``TriMul`` module but fuses the
+ expensive kernels:
   
- 1️⃣ Row‑wise LayerNorm over the last dimension (float32 accumulator,
-    fp16 output).
+ 1️⃣  Row‑wise LayerNorm over the last dimension (FP16 output, FP32 reduction).
+ 2️⃣  Fused projection, gating and optional scalar mask:
+     * left_proj, right_proj  = x_norm @ W_proj
+     * left_gate, right_gate, out_gate = sigmoid(x_norm @ W_gate)
+     * left  = left_proj  * left_gate  * mask
+     * right = right_proj * right_gate * mask
+ 3️⃣  Pairwise multiplication across the sequence dimension (batched GEMM on
+     fp16 tensors).
+ 4️⃣  Fused hidden‑dim LayerNorm → out‑gate multiplication → final linear
+     projection (all in one kernel, FP16 matmul with FP32 accumulation).
   
- 2️⃣ Fused projection + sigmoid gating + optional scalar mask.
-    For each flat row (b,i,k) we compute
-        left_proj = x_norm·W_left_proj
-        left_gate = sigmoid(x_norm·W_left_gate)
-        right_proj = x_norm·W_right_proj
-        right_gate = sigmoid(x_norm·W_right_gate)
-        out_gate   = sigmoid(x_norm·W_out_gate)
-    and store `left`, `right` (shape B×H×N×N) and `out_gate`
-    (shape B×N×N×H).
-  
- 3️⃣ Batched GEMM `left @ rightᵀ` (using torch.bmm on fp16 tensors), then
-    a fused hidden‑dim LayerNorm → element‑wise out‑gate → final linear
-    projection (H → C) producing the float32 output.
-  
- All kernels use mixed‑precision (fp16 compute, fp32 accumulation) and
- avoid host‑side synchronisation, yielding a fast forward pass on an
- H100 GPU.
+ The output tensor has shape ``[B, N, N, dim]`` and dtype ``float32``.
  """
   
  from typing import Tuple, Dict
  ⋯ 1 unchanged lines
  import triton
  import triton.language as tl
   
- # --------------------------------------------------------------
- # 1) Row‑wise LayerNorm (fp16 output, fp32 accumulator)
- # --------------------------------------------------------------
+  
+ # ----------------------------------------------------------------------
+ # 1) Row‑wise LayerNorm (FP16 output, 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 constexpr)
+     X_ptr, Y_ptr,          # (M, C) input / output
+     w_ptr, b_ptr,          # LN weight & bias (fp32)
+     M, C: tl.constexpr,    # rows, columns (C is compile‑time constant)
      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) ---
+     # ---------- mean / var (fp32) ----------
      sum_val = tl.zeros([BLOCK_M], dtype=tl.float32)
      sumsq_val = tl.zeros([BLOCK_M], dtype=tl.float32)
+  
      for c in range(0, C, BLOCK_C):
          cur_c = c + tl.arange(0, BLOCK_C)
          col_mask = cur_c < C
  ⋯ 1 unchanged lines
              X_ptr + rows[:, None] * C + cur_c[None, :],
              mask=row_mask[:, None] & col_mask[None, :],
              other=0.0,
-         ).to(tl.float32)                 # (BLOCK_M, BLOCK_C)
+         ).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,
  ⋯ 25 unchanged lines
      return y_flat.view(B, N, N, C)
   
   
- # --------------------------------------------------------------
- # 2) Fused projection + gating (+ optional mask)
- # --------------------------------------------------------------
+ # ----------------------------------------------------------------------
+ # 2) Fused projection + gating + optional mask
+ # ----------------------------------------------------------------------
  @triton.jit
  def _proj_gate_mask_kernel(
      x_ptr,                         # (M, C) fp16
-     mask_ptr,                      # (M,) fp16 (if MASKED==1)
+     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
  ⋯ 20 unchanged lines
      row_mask = rows < M
      hid_mask = hids < H
   
-     # ---- scalar mask per row (if any) ----
+     # ---------------- mask (scalar per row) ----------------
      if MASKED:
          mask_val = tl.load(mask_ptr + rows, mask=row_mask, other=0.0).to(tl.float32)  # (BLOCK_M,)
      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
  ⋯ 9 unchanged lines
              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 (C, H) row‑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)          # (BLOCK_K, BLOCK_H)
-         w_lg = tl.load(left_gate_w_ptr + cur_k[:, None] * H + hids[None, :],
-                        mask=k_mask[:, None] & hid_mask[None, :],
-                        other=0.0)
-         w_rp = tl.load(right_proj_w_ptr + cur_k[:, None] * H + hids[None, :],
-                        mask=k_mask[:, None] & hid_mask[None, :],
-                        other=0.0)
-         w_rg = tl.load(right_gate_w_ptr + cur_k[:, None] * H + hids[None, :],
-                        mask=k_mask[:, None] & hid_mask[None, :],
-                        other=0.0)
-         w_og = tl.load(out_gate_w_ptr + cur_k[:, None] * H + hids[None, :],
-                        mask=k_mask[:, None] & hid_mask[None, :],
-                        other=0.0)
+         # 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,
+         )
+         w_lg = tl.load(
+             left_gate_w_ptr + cur_k[:, None] * H + hids[None, :],
+             mask=k_mask[:, None] & hid_mask[None, :],
+             other=0.0,
+         )
+         w_rp = tl.load(
+             right_proj_w_ptr + cur_k[:, None] * H + hids[None, :],
+             mask=k_mask[:, None] & hid_mask[None, :],
+             other=0.0,
+         )
+         w_rg = tl.load(
+             right_gate_w_ptr + cur_k[:, None] * H + hids[None, :],
+             mask=k_mask[:, None] & hid_mask[None, :],
+             other=0.0,
+         )
+         w_og = tl.load(
+             out_gate_w_ptr + cur_k[:, None] * H + hids[None, :],
+             mask=k_mask[:, None] & hid_mask[None, :],
+             other=0.0,
+         )
   
          # fp16·fp16 → fp32 dot products
          acc_lp += tl.dot(a, w_lp)
  ⋯ 2 unchanged lines
          acc_rg += tl.dot(a, w_rg)
          acc_og += tl.dot(a, w_og)
   
-     # ---- sigmoid (fp32) ----
+     # ---------------- sigmoid gates -------------------------
      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]
   
-     # ---- map flat row index (b,i,k) → coordinates ----
+     # ---------------- store left/right (B,H,N,N) -------------
      N_sq = N * N
      b_idx = rows // N_sq
      rem   = rows - b_idx * N_sq
      i_idx = rem // N
      k_idx = rem - i_idx * N
   
-     # layout for left/right: (B, H, N, N)
-     left_offset = ((b_idx[:, None] * H + hids[None, :]) * N_sq) + i_idx[:, None] * N + k_idx[:, None]
+     # layout for left/right: (B, H, N, N) → flat index:
+     off = ((b_idx[:, None] * H + hids[None, :]) * N_sq) + i_idx[:, None] * N + k_idx[:, None]
   
      tl.store(
-         left_ptr + left_offset,
+         left_ptr + off,
          left_out.to(tl.float16),
          mask=row_mask[:, None] & hid_mask[None, :],
      )
      tl.store(
-         right_ptr + left_offset,
+         right_ptr + off,
          right_out.to(tl.float16),
          mask=row_mask[:, None] & hid_mask[None, :],
      )
   
-     # out_gate layout: (B, N, N, H)
-     out_gate_offset = rows[:, None] * H + hids[None, :]
+     # ---------------- store out_gate (B,N,N,H) ---------------
+     out_off = rows[:, None] * H + hids[None, :]
      tl.store(
-         out_gate_ptr + out_gate_offset,
+         out_gate_ptr + out_off,
          out_gate.to(tl.float16),
          mask=row_mask[:, None] & hid_mask[None, :],
      )
   
   
- # --------------------------------------------------------------
- # 3) Fused hidden‑dim LN → out‑gate → final linear
- # --------------------------------------------------------------
+ # ----------------------------------------------------------------------
+ # 3) Fused hidden‑dim LayerNorm → out‑gate → final linear
+ # ----------------------------------------------------------------------
  @triton.jit
  def _ln_gate_out_linear_fused_kernel(
      hidden_ptr,           # (B*H*N*N,) fp16 flattened
  ⋯ 14 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 tile (BLOCK_M, BLOCK_H)
-     # --------------------------------------------------------------
+     # ----- load hidden slice (BLOCK_M, BLOCK_H) ------------
      hids = tl.arange(0, BLOCK_H)
      hid_mask = hids < H
   
  ⋯ 5 unchanged lines
      )  # 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 / var 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)
   
-     # --------------------------------------------------------------
-     # 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,)
+     # ----- LayerNorm (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)
-     # --------------------------------------------------------------
+     # ----- final linear projection (fp16 matmul, fp32 acc) -----
      gated_fp16 = gated.to(tl.float16)
+  
      for d0 in range(0, D, BLOCK_D):
          cols = d0 + tl.arange(0, BLOCK_D)
          col_mask = cols < D
-         # weight shape (H, D) row‑major
+  
          w_out = tl.load(
              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,
  ⋯ 1 unchanged lines
          )
   
   
- # --------------------------------------------------------------
- # Entry point
- # --------------------------------------------------------------
+ # ----------------------------------------------------------------------
+ # 4) Entrypoint
+ # ----------------------------------------------------------------------
  def custom_kernel(
      data: Tuple[torch.Tensor, torch.Tensor, Dict[str, torch.Tensor], Dict]
  ) -> torch.Tensor:
      """
      Forward pass of the outgoing TriMul operator (no gradients).
   
-     Parameters
-     ----------
-     data : tuple
-         (input, mask, weights, config)
-         - input : Tensor[B, N, N, C]   (float32)
-         - mask  : Tensor[B, N, N]      (bool/float) or None
+     Arguments
+     ---------
+     data : (input, mask, weights, config)
+         - input  : Tensor[B, N, N, C] (float32)
+         - mask   : Tensor[B, N, N] (bool/float) or None
          - weights: dict of module parameters (float32)
          - config : dict with ``dim`` (C) and ``hidden_dim`` (H) and optional ``nomask``
   
  ⋯ 2 unchanged lines
      Tensor[B, N, N, C] (float32)
      """
      inp, mask, weights, cfg = data
-     dim = cfg["dim"]                     # C
-     hidden_dim = cfg["hidden_dim"]       # H
+     dim = cfg["dim"]                # C
+     hidden_dim = cfg["hidden_dim"]  # H
      nomask = cfg.get("nomask", True)
      eps = 1e-5
   
      device = inp.device
      B, N, _, _ = inp.shape
-     M = B * N * N                       # total rows for row‑wise ops
+     M = B * N * N                     # total rows for row‑wise ops
   
      # --------------------------------------------------------------
-     # 1) Row‑wise LayerNorm (fp16)
+     # 1) Row‑wise LayerNorm (fp16 output)
      # --------------------------------------------------------------
      x_norm = _row_layernorm_fp16(
          inp,
  ⋯ 3 unchanged lines
      )  # (B, N, N, C) fp16
   
      # --------------------------------------------------------------
-     # 2) Prepare projection / gate weight tensors (C, H) fp16, row‑major
+     # 2) Prepare projection / gate weights  (C, H) 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 (optional)
      # --------------------------------------------------------------
-     if (not nomask) and (mask is not None):
+     if not nomask and mask is not None:
          mask_flat = mask.reshape(M).to(torch.float16).contiguous()
          MASKED = 1
      else:
  ⋯ 1 unchanged lines
          MASKED = 0
   
      # --------------------------------------------------------------
-     # 4) Allocate buffers for fused projection / gating kernel
+     # 4) Allocate buffers for fused projection + gating
      # --------------------------------------------------------------
      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      # rows per program (covers B·N²)
-     BLOCK_H = 64      # hidden‑dim block
-     BLOCK_K = 32      # input channel tile size
-  
+     BLOCK_M = 64
+     BLOCK_H = 64
+     BLOCK_K = 32
      grid_proj = (triton.cdiv(M, BLOCK_M), triton.cdiv(hidden_dim, BLOCK_H))
+  
      _proj_gate_mask_kernel[grid_proj](
          x_norm,
          mask_flat,
  ⋯ 17 unchanged lines
      )
   
      # --------------------------------------------------------------
-     # 6) Pairwise multiplication (batched GEMM)
+     # 6) Pairwise multiplication (batched GEMM) – left @ rightᵀ
      # --------------------------------------------------------------
-     # left, right shape: (B, H, N, N) → view as (B*H, N, N)
      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)
+     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
+     hidden = hidden_fp16.view(B, hidden_dim, N, N)                # (B, H, N, N) fp16
   
      # --------------------------------------------------------------
-     # 7) Fused hidden‑dim LN → out‑gate → final linear
+     # 7) Fused hidden‑dim LayerNorm → 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
  ⋯ 2 unchanged lines
      out = torch.empty((B, N, N, dim), dtype=torch.float32, device=device)
   
      BLOCK_M_OUT = 64
-     BLOCK_H_OUT = hidden_dim          # fits in a single block for the tested configs
+     BLOCK_H_OUT = hidden_dim      # cover the whole hidden dim in one kernel launch
      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
+         hidden.view(-1),                 # flat fp16 hidden
+         out_gate.view(-1),               # flat fp16 out‑gate
          to_out_norm_w,
          to_out_norm_b,
          to_out_w_T,