1{
2  "op": "replace",
3  "filename": "nanoGPT/custom_pretrain.py",
4  "start_line": 33,
5  "end_line": 48,
6  "content": "import triton\nimport triton.language as tl\n\n@triton.jit\ndef _sym_cross_glu_fwd_kernel(\n    h_ptr, out_ptr, half_n, n_elements,\n    BLOCK_SIZE: tl.constexpr,\n):\n    \"\"\"Fused SymCrossGLU forward: out[i] = silu(h1[i]) * h2[i] for first half,\n       silu(h2[i]) * h1[i] for second half. Each half gates the other.\"\"\"\n    pid = tl.program_id(0)\n    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)\n    mask = offsets < n_elements\n\n    # Determine row and column\n    row = offsets // (half_n * 2)\n    col = offsets % (half_n * 2)\n\n    # For each position, determine if it's in the first or second half\n    is_first_half = col < half_n\n    # Compute partner offset (the corresponding element in the other half)\n    partner_offsets = tl.where(is_first_half, offsets + half_n, offsets - half_n)\n\n    # Load both this element and its partner\n    x = tl.load(h_ptr + offsets, mask=mask).to(tl.float32)\n    partner = tl.load(h_ptr + partner_offsets, mask=mask).to(tl.float32)\n\n    # silu(x) = x * sigmoid(x)\n    sigmoid_x = tl.sigmoid(x)\n    silu_x = x * sigmoid_x\n\n    # Output: silu(self) * partner\n    result = silu_x * partner\n    tl.store(out_ptr + offsets, result.to(tl.bfloat16), mask=mask)\n\n@triton.jit\ndef _sym_cross_glu_bwd_kernel(\n    h_ptr, grad_out_ptr, grad_h_ptr, half_n, n_elements,\n    BLOCK_SIZE: tl.constexpr,\n):\n    \"\"\"Backward for SymCrossGLU.\n    For first half output: o1 = silu(h1) * h2\n      d_loss/d_h1 = d_loss/d_o1 * h2 * silu'(h1) + d_loss/d_o2 * silu(h2)\n      (h1 appears in o1 via silu(h1) and in o2 as the multiplier)\n    Similarly for h2.\n    silu'(x) = sigmoid(x) + x * sigmoid(x) * (1 - sigmoid(x))\n             = sigmoid(x) * (1 + x * (1 - sigmoid(x)))\n    \"\"\"\n    pid = tl.program_id(0)\n    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)\n    mask = offsets < n_elements\n\n    col = offsets % (half_n * 2)\n    is_first_half = col < half_n\n    partner_offsets = tl.where(is_first_half, offsets + half_n, offsets - half_n)\n\n    # Load h values\n    x = tl.load(h_ptr + offsets, mask=mask).to(tl.float32)\n    partner = tl.load(h_ptr + partner_offsets, mask=mask).to(tl.float32)\n\n    # Load grad_output for both this position and partner position\n    grad_self = tl.load(grad_out_ptr + offsets, mask=mask).to(tl.float32)\n    grad_partner = tl.load(grad_out_ptr + partner_offsets, mask=mask).to(tl.float32)\n\n    # Compute sigmoid and silu for self\n    sigmoid_x = tl.sigmoid(x)\n    silu_x = x * sigmoid_x\n    # silu derivative\n    dsilu_x = sigmoid_x * (1.0 + x * (1.0 - sigmoid_x))\n\n    # Compute sigmoid and silu for partner\n    sigmoid_p = tl.sigmoid(partner)\n    silu_p = partner * sigmoid_p\n\n    # Gradient contribution:\n    # From output at self position: o_self = silu(x) * partner\n    #   d_loss/d_x from this = grad_self * partner * dsilu(x)\n    # From output at partner position: o_partner = silu(partner) * x\n    #   d_loss/d_x from this = grad_partner * silu(partner)\n    grad_x = grad_self * partner * dsilu_x + grad_partner * silu_p\n\n    tl.store(grad_h_ptr + offsets, grad_x.to(tl.bfloat16), mask=mask)\n\nclass _SymCrossGLU_MLP(torch.autograd.Function):\n    @staticmethod\n    def forward(ctx, x, w_fc, w_proj):\n        h = x @ w_fc.t()  # (B*T, 4*n_embd)\n        M, N = h.shape\n        half_n = N // 2\n        act = torch.empty_like(h)\n        n = h.numel()\n        BLOCK = 1024\n        grid = ((n + BLOCK - 1) // BLOCK,)\n        _sym_cross_glu_fwd_kernel[grid](h, act, half_n, n, BLOCK_SIZE=BLOCK)\n        out = act @ w_proj.t()\n        # Memory-efficient: only save h (recompute act in backward if needed)\n        ctx.save_for_backward(x, w_fc, w_proj, h)\n        ctx.half_n = half_n\n        return out\n\n    @staticmethod\n    def backward(ctx, grad_output):\n        x, w_fc, w_proj, h = ctx.saved_tensors\n        half_n = ctx.half_n\n        dtype = grad_output.dtype\n\n        # Recompute activation for w_proj gradient\n        n = h.numel()\n        BLOCK = 1024\n        grid = ((n + BLOCK - 1) // BLOCK,)\n        act = torch.empty_like(h)\n        _sym_cross_glu_fwd_kernel[grid](h, act, half_n, n, BLOCK_SIZE=BLOCK)\n\n        # Gradient through second linear\n        d_act = grad_output @ w_proj.to(dtype)\n        grad_w_proj = grad_output.reshape(-1, grad_output.shape[-1]).t() @ act.to(dtype).reshape(-1, act.shape[-1])\n\n        # Gradient through SymCrossGLU activation using fused Triton kernel\n        d_h = torch.empty_like(h)\n        _sym_cross_glu_bwd_kernel[grid](h, d_act, d_h, half_n, n, BLOCK_SIZE=BLOCK)\n\n        # Gradient through first linear\n        grad_x = d_h.to(dtype) @ w_fc.to(dtype)\n        grad_w_fc = d_h.reshape(-1, d_h.shape[-1]).to(dtype).t() @ x.to(dtype).reshape(-1, x.shape[-1])\n\n        return grad_x, grad_w_fc, grad_w_proj\n\ndef fused_mlp_forward(x, w_fc, w_proj):\n    \"\"\"MLP forward with Symmetric Cross-Gated Linear Unit (SymCrossGLU).\n\n    Novel activation: splits hidden dim in half, each half gates the other via SiLU.\n    This introduces GLU-style gating without changing architecture dimensions.\n    Implemented as fused Triton kernels with memory-efficient backward (recomputation).\n\n    Args:\n        x: input tensor (B*T, n_embd)\n        w_fc: first linear weight (4*n_embd, n_embd)\n        w_proj: second linear weight (n_embd, 4*n_embd)\n    Returns:\n        output tensor (B*T, n_embd)\n    \"\"\"\n    return _SymCrossGLU_MLP.apply(x, w_fc, w_proj)\n"
7}

1{
2  "content": "def fused_mlp_forward(x, w_fc, w_proj):\n    \"\"\"MLP forward with Symmetric Cross-Gated Linear Unit (SymCrossGLU).\n\n    Novel activation: splits hidden dim in half, each half gates the other via SiLU.\n    This introduces GLU-style gating without changing architecture dimensions.\n    output = cat[silu(h1)*h2, silu(h2)*h1] where h = x @ w_fc.t(), h1,h2 = split(h).\n\n    Motivation: GLU variants (SwiGLU, GeGLU) are state-of-the-art activations in modern\n    LLMs (LLaMA, PaLM), but require architectural changes (2/3 expansion). SymCrossGLU\n    achieves gating within the standard 4x expansion by using symmetric cross-gating,\n    where each half of the hidden state provides a multiplicative gate for the other.\n\n    Args:\n        x: input tensor (B*T, n_embd)\n        w_fc: first linear weight (4*n_embd, n_embd)\n        w_proj: second linear weight (n_embd, 4*n_embd)\n    Returns:\n        output tensor (B*T, n_embd)\n    \"\"\"\n    h = x @ w_fc.t()  # (B*T, 4*n_embd)\n    h1, h2 = h.chunk(2, dim=-1)  # each (B*T, 2*n_embd)\n    # Symmetric cross-gating: each half gates the other via SiLU\n    gated = torch.cat([F.silu(h1) * h2, F.silu(h2) * h1], dim=-1)  # (B*T, 4*n_embd)\n    return gated @ w_proj.t()\n",
3  "end_line": 174,
4  "filename": "nanoGPT/custom_pretrain.py",
5  "op": "replace",
6  "start_line": 33
7}

1{
2  "n": 3
3}