submission 407546

import torch
import triton
import triton.language as tl

@triton.jit
def _fused_preprocess_kernel(
    X_ptr, M_ptr, L_ptr, R_ptr, OG_ptr,
    W_norm_ptr, B_norm_ptr,
    W_L_ptr, W_R_ptr, W_LG_ptr, W_RG_ptr, W_OG_ptr,
    stride_xb, stride_xi, stride_xj, stride_xc,
    stride_mb, stride_mi, stride_mj,
    B, N, C, D: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_C: tl.constexpr
):
    pid_b = tl.program_id(0)
    pid_i = tl.program_id(1)
    pid_j_start = tl.program_id(2) * BLOCK_SIZE_N

    offsets_j = pid_j_start + tl.arange(0, BLOCK_SIZE_N)
    mask_j = offsets_j < N

    # 1. Compute LayerNorm for the block [BLOCK_SIZE_N, C]
    acc_sum = tl.zeros([BLOCK_SIZE_N], dtype=tl.float32)
    acc_sum_sq = tl.zeros([BLOCK_SIZE_N], dtype=tl.float32)
    
    for c_offset in range(0, C, BLOCK_SIZE_C):
        cols = c_offset + tl.arange(0, BLOCK_SIZE_C)
        c_mask = cols < C
        x_ptr = X_ptr + pid_b * stride_xb + pid_i * stride_xi + offsets_j[:, None] * stride_xj + cols[None, :]
        x_chunk = tl.load(x_ptr, mask=(mask_j[:, None] & c_mask[None, :]), other=0.0).to(tl.float32)
        acc_sum += tl.sum(x_chunk, axis=1)
        acc_sum_sq += tl.sum(x_chunk * x_chunk, axis=1)

    mean = acc_sum / C
    var = (acc_sum_sq / C) - (mean * mean)
    rstd = 1.0 / tl.sqrt(var + 1e-5)

    # 2. Compute Projections
    acc_l = tl.zeros([BLOCK_SIZE_N, D], dtype=tl.float32)
    acc_lg = tl.zeros([BLOCK_SIZE_N, D], dtype=tl.float32)
    acc_r = tl.zeros([BLOCK_SIZE_N, D], dtype=tl.float32)
    acc_rg = tl.zeros([BLOCK_SIZE_N, D], dtype=tl.float32)
    acc_og = tl.zeros([BLOCK_SIZE_N, D], dtype=tl.float32)

    off_d = tl.arange(0, D)

    for c_offset in range(0, C, BLOCK_SIZE_C):
        cols = c_offset + tl.arange(0, BLOCK_SIZE_C)
        c_mask = cols < C
        
        x_ptr = X_ptr + pid_b * stride_xb + pid_i * stride_xi + offsets_j[:, None] * stride_xj + cols[None, :]
        x_chunk = tl.load(x_ptr, mask=(mask_j[:, None] & c_mask[None, :]), other=0.0).to(tl.float32)
        w_n = tl.load(W_norm_ptr + cols, mask=c_mask, other=0.0)
        b_n = tl.load(B_norm_ptr + cols, mask=c_mask, other=0.0)
        x_n = (x_chunk - mean[:, None]) * rstd[:, None] * w_n[None, :] + b_n[None, :]
        x_n = x_n.to(tl.float16)

        # Projection weights [C, D]
        w_l = tl.load(W_L_ptr + cols[:, None] * D + off_d[None, :], mask=c_mask[:, None], other=0.0).to(tl.float16)
        acc_l += tl.dot(x_n, w_l)
        
        w_lg = tl.load(W_LG_ptr + cols[:, None] * D + off_d[None, :], mask=c_mask[:, None], other=0.0).to(tl.float16)
        acc_lg += tl.dot(x_n, w_lg)
        
        w_r = tl.load(W_R_ptr + cols[:, None] * D + off_d[None, :], mask=c_mask[:, None], other=0.0).to(tl.float16)
        acc_r += tl.dot(x_n, w_r)
        
        w_rg = tl.load(W_RG_ptr + cols[:, None] * D + off_d[None, :], mask=c_mask[:, None], other=0.0).to(tl.float16)
        acc_rg += tl.dot(x_n, w_rg)
        
        w_og = tl.load(W_OG_ptr + cols[:, None] * D + off_d[None, :], mask=c_mask[:, None], other=0.0).to(tl.float16)
        acc_og += tl.dot(x_n, w_og)

    # 3. Apply Gating and Masking
    m_ptr = M_ptr + pid_b * stride_mb + pid_i * stride_mi + offsets_j
    mask_val = tl.load(m_ptr, mask=mask_j, other=0.0).to(tl.float32)

    l_final = acc_l * tl.sigmoid(acc_lg) * mask_val[:, None]
    r_final = acc_r * tl.sigmoid(acc_rg) * mask_val[:, None]
    og_final = tl.sigmoid(acc_og)

    # 4. Store results
    idx_nn = pid_i * N + offsets_j
    # L and R stored in [B, D, N, N] for BMM efficiency
    off_l_r = pid_b * D * N * N + off_d[None, :] * N * N + idx_nn[:, None]
    tl.store(L_ptr + off_l_r, l_final.to(tl.float16), mask=(mask_j[:, None] & (off_d[None, :] < D)))
    tl.store(R_ptr + off_l_r, r_final.to(tl.float16), mask=(mask_j[:, None] & (off_d[None, :] < D)))
    
    # OG stored in [B, N, N, D]
    off_og = pid_b * N * N * D + idx_nn[:, None] * D + off_d[None, :]
    tl.store(OG_ptr + off_og, og_final.to(tl.float16), mask=(mask_j[:, None] & (off_d[None, :] < D)))

def custom_kernel(data):
    x, mask, weights, config = data
    B, N, _, C = x.shape
    D = config["hidden_dim"]
    device = x.device

    # Cast weights to FP16 for the fused kernel to avoid dtype mismatch in tl.dot
    # and for the final linear layer.
    w_fp16 = {k: v.to(torch.float16) for k, v in weights.items()}

    L = torch.empty((B, D, N, N), device=device, dtype=torch.float16)
    R = torch.empty((B, D, N, N), device=device, dtype=torch.float16)
    OG = torch.empty((B, N, N, D), device=device, dtype=torch.float16)

    BLOCK_SIZE_N = 32
    BLOCK_SIZE_C = 128
    
    grid = (B, N, (N + BLOCK_SIZE_N - 1) // BLOCK_SIZE_N)
    _fused_preprocess_kernel[grid](
        x, mask, L, R, OG,
        w_fp16["norm.weight"], w_fp16["norm.bias"],
        w_fp16["left_proj.weight"].t().contiguous(), w_fp16["right_proj.weight"].t().contiguous(),
        w_fp16["left_gate.weight"].t().contiguous(), w_fp16["right_gate.weight"].t().contiguous(),
        w_fp16["out_gate.weight"].t().contiguous(),
        x.stride(0), x.stride(1), x.stride(2), x.stride(3),
        mask.stride(0), mask.stride(1), mask.stride(2),
        B, N, C, D,
        BLOCK_SIZE_N=BLOCK_SIZE_N,
        BLOCK_SIZE_C=BLOCK_SIZE_C,
        num_stages=2
    )

    # Contraction: einsum('... i k d, ... j k d -> ... i j d')
    # L: [B, D, N, K], R: [B, D, N, K]. We want [B, D, N, N] where out[b, d, i, j] = sum_k L[b, d, i, k] * R[b, d, j, k]
    # This is L @ R.T
    bmm_out = torch.bmm(L.view(B * D, N, N), R.view(B * D, N, N).transpose(-1, -2))
    bmm_out = bmm_out.view(B, D, N, N).permute(0, 2, 3, 1) # [B, N, N, D]

    # Final layers: LayerNorm -> Gating -> Linear
    # Use FP32 for LayerNorm stability if needed, but here we use FP16 for speed
    out = torch.nn.functional.layer_norm(bmm_out, (D,), w_fp16["to_out_norm.weight"], w_fp16["to_out_norm.bias"])
    out = out * OG
    
    # Final projection to C
    return (out @ w_fp16["to_out.weight"].t()).to(torch.float32)