submission 380656

  """
- Outgoing TriMul (AlphaFold‑3) – Triton accelerated forward pass.
+ Outgoing TriMul (AlphaFold‑3) – Triton‑accelerated forward pass
   
- The implementation follows the reference ``TriMul`` module but fuses the
- expensive kernels:
+ Algorithm
+ ---------
+ 1️⃣ Row‑wise LayerNorm over the last axis (C).  The mean/variance are
+    computed in FP32, the normalised output is stored in FP16.
   
- 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 + 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)
   
- The output tensor has shape ``[B, N, N, dim]`` and dtype ``float32``.
+ 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.
  """
   
  from typing import Tuple, Dict
  ⋯ 1 unchanged lines
  import triton
  import triton.language as tl
   
-  
  # ----------------------------------------------------------------------
  # 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 is compile‑time constant)
+     w_ptr, b_ptr,          # LN weight & bias (FP32)
+     M, C: tl.constexpr,    # rows, columns (C must be const)
      eps,
      BLOCK_M: tl.constexpr,
      BLOCK_C: tl.constexpr,
  ⋯ 3 unchanged lines
      rows = row_start + tl.arange(0, BLOCK_M)
      row_mask = rows < M
   
-     # ---------- mean / var (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)                     # (BLOCK_M, BLOCK_C)
-  
+         ).to(tl.float32)
          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)
   
-     # ---------- normalize + affine (fp16) ----------
+     # --------------------------------------------------------------
+     # Normalise + 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,
  ⋯ 30 unchanged lines
  # ----------------------------------------------------------------------
  @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,
  ⋯ 12 unchanged lines
      row_mask = rows < M
      hid_mask = hids < H
   
-     # ---------------- mask (scalar per row) ----------------
+     # ---- scalar mask per row (if any) ----
      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
  ⋯ 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 (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,
-         )
+         # weight tiles (row‑major (C, H))
+         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
+         # 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 gates -------------------------
+     # ---- 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 mask and per‑row gates ----------
+     # ---- apply scalar mask and per‑row gates ----
      left_out  = acc_lp * left_gate * mask_val[:, None]
      right_out = acc_rp * right_gate * mask_val[:, None]
   
-     # ---------------- store left/right (B,H,N,N) -------------
+     # ---- map flat row index (b,i,k) → coordinates ----
      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) → flat index:
-     off = ((b_idx[:, None] * H + hids[None, :]) * N_sq) + i_idx[:, None] * N + k_idx[:, None]
+     # 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]
   
      tl.store(
-         left_ptr + off,
+         left_ptr + left_offset,
          left_out.to(tl.float16),
          mask=row_mask[:, None] & hid_mask[None, :],
      )
      tl.store(
-         right_ptr + off,
+         right_ptr + left_offset,
          right_out.to(tl.float16),
          mask=row_mask[:, None] & hid_mask[None, :],
      )
   
-     # ---------------- store out_gate (B,N,N,H) ---------------
-     out_off = rows[:, None] * H + hids[None, :]
+     # out_gate layout: (B, N, N, H)
+     out_gate_offset = rows[:, None] * H + hids[None, :]
      tl.store(
-         out_gate_ptr + out_off,
+         out_gate_ptr + out_gate_offset,
          out_gate.to(tl.float16),
          mask=row_mask[:, None] & hid_mask[None, :],
      )
  ⋯ 4 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 slice (BLOCK_M, BLOCK_H) ------------
+     # --------------------------------------------------------------
+     # Load hidden tensor 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 / 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 / variance across H (FP32)
+     # --------------------------------------------------------------
+     sum_val = tl.sum(hidden_fp32, axis=1)
+     sumsq_val = tl.sum(hidden_fp32 * hidden_fp32, axis=1)
      mean = sum_val / H
      var = sumsq_val / H - mean * mean
-     inv_std = 1.0 / tl.sqrt(var + eps)
+     inv_std = 1.0 / tl.sqrt(var + eps)                     # (BLOCK_M,)
   
-     # ----- LayerNorm (fp32) -----
+     # --------------------------------------------------------------
+     # 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,)
      hidden_norm = (hidden_fp32 - mean[:, None]) * inv_std[:, None]
-     hidden_norm = hidden_norm * w_ln[None, :] + b_ln[None, :]   # (BLOCK_M, BLOCK_H)
+     hidden_norm = hidden_norm * w_ln[None, :] + b_ln[None, :]
   
-     # ----- out‑gate (fp32) -----
+     # --------------------------------------------------------------
+     # Apply 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)                                            # (BLOCK_M, BLOCK_H)
+     ).to(tl.float32)
   
-     gated = hidden_norm * out_gate_tile                         # (BLOCK_M, BLOCK_H)
+     gated = hidden_norm * out_gate_tile
   
-     # ----- final linear projection (fp16 matmul, fp32 acc) -----
+     # --------------------------------------------------------------
+     # Final linear projection (H → D)
+     # --------------------------------------------------------------
      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
-  
          w_out = tl.load(
              w_out_ptr + hids[:, None] * D + cols[None, :],
              mask=hid_mask[:, None] & col_mask[None, :],
              other=0.0,
-         )  # (BLOCK_H, BLOCK_D) fp16
+         )  # (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
   
   
  # ----------------------------------------------------------------------
- # 4) Entrypoint
+ # Entry point
  # ----------------------------------------------------------------------
  def custom_kernel(
      data: Tuple[torch.Tensor, torch.Tensor, Dict[str, torch.Tensor], Dict]
  ⋯ 1 unchanged lines
      """
      Forward pass of the outgoing TriMul operator (no gradients).
   
-     Arguments
-     ---------
-     data : (input, mask, weights, config)
-         - input  : Tensor[B, N, N, C] (float32)
-         - mask   : Tensor[B, N, N] (bool/float) or None
+     Parameters
+     ----------
+     data : tuple
+         (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``
   
  ⋯ 1 unchanged lines
      -------
      Tensor[B, N, N, C] (float32)
      """
+     # --------------------------------------------------------------
+     # unpack arguments
+     # --------------------------------------------------------------
      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 output)
+     # 1) Row‑wise LayerNorm (fp16)
      # --------------------------------------------------------------
      x_norm = _row_layernorm_fp16(
          inp,
  ⋯ 3 unchanged lines
      )  # (B, N, N, C) fp16
   
      # --------------------------------------------------------------
-     # 2) Prepare projection / gate weights  (C, H) fp16, column‑major
+     # 2) Prepare projection / gate weights (C, H) fp16, row‑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)
+     # 3) Mask handling – optional scalar mask per (b,i,k)
      # --------------------------------------------------------------
      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
+     # 4) Allocate buffers for 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
-     BLOCK_H = 64
-     BLOCK_K = 32
-     grid_proj = (triton.cdiv(M, BLOCK_M), triton.cdiv(hidden_dim, BLOCK_H))
+     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
   
+     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) – left @ rightᵀ
+     # 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)ᵀ
+     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 LayerNorm → out‑gate → final linear
+     # 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
  ⋯ 2 unchanged lines
      out = torch.empty((B, N, N, dim), dtype=torch.float32, device=device)
   
      BLOCK_M_OUT = 64
-     BLOCK_H_OUT = hidden_dim      # cover the whole hidden dim in one kernel launch
-     BLOCK_D_OUT = 64
+     BLOCK_D_OUT = 64                     # D‑tile for final projection
   
      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,
  ⋯ 4 unchanged lines
          dim,
          eps,
          BLOCK_M=BLOCK_M_OUT,
-         BLOCK_H=BLOCK_H_OUT,
+         BLOCK_H=hidden_dim,          # H fits into one block for our configs
          BLOCK_D=BLOCK_D_OUT,
          num_warps=4,
      )