""" Triton-fused two-layer MLPs in the style of `VariableUDLP`: out = (act(x @ W1.T))^2 @ W2.T (ReLU^2, LeakyReLU(0.5)^2) out = (silu(a) * b) @ W2.T (SwiGLU) The fused ReLU^2 and LeakyReLU^2 kernels are ports of: - modded-nanogpt linear_relu_square (Keller Jordan et al.) - parameter-golf linear_leaky_relu_square (samacqua / openai/parameter-golf) Both use the Hopper TensorDescriptor (TMA) API, so they need: - triton >= 3.2 - GPU compute capability >= 9.0 (H100 / GH200) On older GPUs / when triton is missing, the wrapper modules transparently fall back to a pure-torch path that is `torch.compile`-friendly. The autograd.Function wrappers are accepted natively by `torch.compile` (>=2.0): the compiler treats `.apply(...)` as an opaque op and traces around it. API: each module mirrors `VariableUDLP(d_model, d_output, M, dropout)`. """ from __future__ import annotations import torch import torch.nn as nn import torch.nn.functional as F from Vathos._basics import Layer from Vathos.functions import flag try: import triton import triton.language as tl from triton.tools.tensor_descriptor import TensorDescriptor _HAS_TRITON = True except Exception: _HAS_TRITON = False def _hopper_available() -> bool: if not (_HAS_TRITON and torch.cuda.is_available()): return False major, _ = torch.cuda.get_device_capability() return major >= 9 # ============================================================================= # Triton kernels # ============================================================================= if _HAS_TRITON: @triton.jit def _linear_relu_square_kernel( a_desc, b_desc, c_desc, aux_desc, M, N, K, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, NUM_SMS: tl.constexpr, FORWARD: tl.constexpr, ): """Port of modded-nanogpt's linear_relu_square_kernel. FORWARD: computes c = (x @ W1.T) and aux = relu(c)^2 (c is the pre-activation) BACKWARD: computes c = grad_pre = grad_post * 2 * relu(pre), reading pre from aux (caller passes aux=pre, grad_post is the input matmul result) """ dtype = tl.bfloat16 start_pid = tl.program_id(axis=0) num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) num_tiles = num_pid_m * num_pid_n k_tiles = tl.cdiv(K, BLOCK_SIZE_K) for tile_id in tl.range(start_pid, num_tiles, NUM_SMS, flatten=True): pid_m = tile_id // num_pid_n pid_n = tile_id % num_pid_n offs_am = pid_m * BLOCK_SIZE_M offs_bn = pid_n * BLOCK_SIZE_N accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) for ki in range(k_tiles): offs_k = ki * BLOCK_SIZE_K a = a_desc.load([offs_am, offs_k]) b = b_desc.load([offs_bn, offs_k]) accumulator = tl.dot(a, b.T, accumulator) acc = tl.reshape(accumulator, (BLOCK_SIZE_M, 2, BLOCK_SIZE_N // 2)) acc = tl.permute(acc, (0, 2, 1)) acc0, acc1 = tl.split(acc) c0 = acc0.to(dtype) if not FORWARD: pre0 = aux_desc.load([offs_am, offs_bn]) c0 = 2 * c0 * tl.where(pre0 > 0, pre0, 0) c_desc.store([offs_am, offs_bn], c0) if FORWARD: p0 = tl.maximum(c0, 0) aux_desc.store([offs_am, offs_bn], p0 * p0) c1 = acc1.to(dtype) if not FORWARD: pre1 = aux_desc.load([offs_am, offs_bn + BLOCK_SIZE_N // 2]) c1 = 2 * c1 * tl.where(pre1 > 0, pre1, 0) c_desc.store([offs_am, offs_bn + BLOCK_SIZE_N // 2], c1) if FORWARD: p1 = tl.maximum(c1, 0) aux_desc.store([offs_am, offs_bn + BLOCK_SIZE_N // 2], p1 * p1) @triton.jit def _linear_leaky_relu_square_kernel( a_desc, b_desc, c_desc, aux_desc, M, N, K, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, NUM_SMS: tl.constexpr, FORWARD: tl.constexpr, ): """Port of parameter-golf's linear_leaky_relu_square_kernel. Uses LeakyReLU with negative slope 0.5, then squares. Forward: aux = lrelu(c)^2, Backward: c = grad_post * 2 * (lrelu'(pre) * lrelu(pre)) = grad_post * (where(pre>0, 2*pre, 0.5*pre)) """ dtype = tl.bfloat16 start_pid = tl.program_id(axis=0) num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) num_tiles = num_pid_m * num_pid_n k_tiles = tl.cdiv(K, BLOCK_SIZE_K) for tile_id in tl.range(start_pid, num_tiles, NUM_SMS, flatten=True): pid_m = tile_id // num_pid_n pid_n = tile_id % num_pid_n offs_am = pid_m * BLOCK_SIZE_M offs_bn = pid_n * BLOCK_SIZE_N accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) for ki in range(k_tiles): offs_k = ki * BLOCK_SIZE_K a = a_desc.load([offs_am, offs_k]) b = b_desc.load([offs_bn, offs_k]) accumulator = tl.dot(a, b.T, accumulator) acc = tl.reshape(accumulator, (BLOCK_SIZE_M, 2, BLOCK_SIZE_N // 2)) acc = tl.permute(acc, (0, 2, 1)) acc0, acc1 = tl.split(acc) c0 = acc0.to(dtype) c1 = acc1.to(dtype) if not FORWARD: pre0 = aux_desc.load([offs_am, offs_bn]) pre1 = aux_desc.load([offs_am, offs_bn + BLOCK_SIZE_N // 2]) c0 = c0 * tl.where(pre0 > 0, 2.0 * pre0, 0.5 * pre0) c1 = c1 * tl.where(pre1 > 0, 2.0 * pre1, 0.5 * pre1) c_desc.store([offs_am, offs_bn], c0) c_desc.store([offs_am, offs_bn + BLOCK_SIZE_N // 2], c1) if FORWARD: aux0 = tl.where(c0 > 0, c0, 0.5 * c0) aux1 = tl.where(c1 > 0, c1, 0.5 * c1) aux_desc.store([offs_am, offs_bn], aux0 * aux0) aux_desc.store([offs_am, offs_bn + BLOCK_SIZE_N // 2], aux1 * aux1) # ----------------------------------------------------------------------------- # T4-friendly LeakyReLU(0.5)^2 kernel (no TMA, fp16, smaller tiles). # Designed for sm_75 (Turing, T4): pointer-based loads, BLOCK_M=BLOCK_N=64, # BLOCK_K=32, fp16 inputs/outputs, fp32 accumulator. Same fwd/bwd convention # as the Hopper kernels: forward stores c=pre and aux=post; backward reads # pre from aux and writes grad_pre into c. Output dtype follows the inputs # (cast to fp32 fails tl.dot on Turing — keep them fp16). # ----------------------------------------------------------------------------- if _HAS_TRITON: @triton.jit def _linear_lrelu2_kernel_t4( a_ptr, b_ptr, c_ptr, aux_ptr, M, N, K, stride_am, stride_ak, stride_bn, stride_bk, stride_cm, stride_cn, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, GROUP_M: tl.constexpr, FORWARD: tl.constexpr, ): pid = tl.program_id(axis=0) num_pid_m = tl.cdiv(M, BLOCK_M) num_pid_n = tl.cdiv(N, BLOCK_N) num_pid_in_group = GROUP_M * num_pid_n group_id = pid // num_pid_in_group first_pid_m = group_id * GROUP_M group_size_m = tl.minimum(num_pid_m - first_pid_m, GROUP_M) pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) pid_n = (pid % num_pid_in_group) // group_size_m offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) offs_k = tl.arange(0, BLOCK_K) a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak b_ptrs = b_ptr + offs_n[:, None] * stride_bn + offs_k[None, :] * stride_bk acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) for k in range(0, tl.cdiv(K, BLOCK_K)): k_remaining = K - k * BLOCK_K a_mask = (offs_m[:, None] < M) & (offs_k[None, :] < k_remaining) b_mask = (offs_n[:, None] < N) & (offs_k[None, :] < k_remaining) a = tl.load(a_ptrs, mask=a_mask, other=0.0) b = tl.load(b_ptrs, mask=b_mask, other=0.0) acc += tl.dot(a, b.T) a_ptrs += BLOCK_K * stride_ak b_ptrs += BLOCK_K * stride_bk out_dtype = c_ptr.dtype.element_ty c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn aux_ptrs = aux_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) if FORWARD: # acc here is pre = x @ W1.T. Save pre into c, post = lrelu²(pre) into aux. tl.store(c_ptrs, acc.to(out_dtype), mask=mask) scale = tl.where(acc > 0, 1.0, 0.25) # act² = z² * (1 if z>0 else 0.25) post = acc * acc * scale tl.store(aux_ptrs, post.to(out_dtype), mask=mask) else: # acc here is grad_post @ W2 (b=W2.T). Multiply elementwise by # d(lrelu²)/dpre = where(pre>0, 2*pre, 0.5*pre) using saved pre in aux. pre = tl.load(aux_ptrs, mask=mask, other=0.0).to(tl.float32) deriv = tl.where(pre > 0, 2.0 * pre, 0.5 * pre) tl.store(c_ptrs, (acc * deriv).to(out_dtype), mask=mask) def _launch_lrelu2_t4(a: torch.Tensor, b: torch.Tensor, aux: torch.Tensor | None): """Pointer-based kernel launcher for T4. Returns (pre, post) in fwd, grad_pre in bwd.""" M, K = a.shape N, K2 = b.shape assert K == K2 c = torch.empty((M, N), device=a.device, dtype=a.dtype) forward = aux is None if aux is None: aux = torch.empty((M, N), device=a.device, dtype=a.dtype) BLOCK_M, BLOCK_N, BLOCK_K = 64, 64, 32 grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N),) _linear_lrelu2_kernel_t4[grid]( a, b, c, aux, M, N, K, a.stride(0), a.stride(1), b.stride(0), b.stride(1), c.stride(0), c.stride(1), BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, GROUP_M=8, FORWARD=forward, num_warps=4, num_stages=2, ) return (c, aux) if forward else c class _FusedLinearLReLU2_T4(torch.autograd.Function): @staticmethod def forward(ctx, x, W1, W2): x_flat = x.reshape(-1, x.shape[-1]) pre, post = _launch_lrelu2_t4(x_flat, W1, None) out = post @ W2.T ctx.save_for_backward(x, W1, W2, pre, post) return out.view(*x.shape[:-1], W2.shape[0]) @staticmethod def backward(ctx, grad_out): x, W1, W2, pre, post = ctx.saved_tensors x_flat = x.reshape(-1, x.shape[-1]) go = grad_out.reshape(-1, W2.shape[0]) dW2 = go.T @ post dpre = _launch_lrelu2_t4(go, W2.T.contiguous(), pre) dW1 = dpre.T @ x_flat dx = dpre @ W1 return dx.view_as(x), dW1, dW2 def _can_use_triton_t4(x) -> bool: """T4 (CC 7.5) requires fp16 inputs.""" if not (_HAS_TRITON and torch.cuda.is_available()): return False major, minor = torch.cuda.get_device_capability() if (major, minor) != (7, 5): return False return x.dtype == torch.float16 # ============================================================================= # Host-side launchers (return both pre-activation and post-activation in fwd) # ============================================================================= def _launch_fused_act_square(kernel_fn, a: torch.Tensor, b: torch.Tensor, aux: torch.Tensor | None): """Shared launcher for both relu^2 and leaky_relu^2 kernels. Forward (aux=None): returns (pre, post) = (x@W.T, act(pre)^2) Backward (aux=pre): returns dpre = grad_post * d(act(pre)^2)/dpre """ M, K = a.shape N, K2 = b.shape assert K == K2 c = torch.empty((M, N), device=a.device, dtype=a.dtype) forward = aux is None if aux is None: aux = torch.empty((M, N), device=a.device, dtype=a.dtype) BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 128, 256, 64 num_stages = 4 if forward else 3 a_desc = TensorDescriptor.from_tensor(a, [BLOCK_SIZE_M, BLOCK_SIZE_K]) b_desc = TensorDescriptor.from_tensor(b, [BLOCK_SIZE_N, BLOCK_SIZE_K]) c_desc = TensorDescriptor.from_tensor(c, [BLOCK_SIZE_M, BLOCK_SIZE_N // 2]) aux_desc = TensorDescriptor.from_tensor(aux, [BLOCK_SIZE_M, BLOCK_SIZE_N // 2]) num_sms = torch.cuda.get_device_properties(a.device).multi_processor_count grid = lambda _meta: ( min(num_sms, triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(N, BLOCK_SIZE_N)), ) kernel_fn[grid]( a_desc, b_desc, c_desc, aux_desc, M, N, K, BLOCK_SIZE_M=BLOCK_SIZE_M, BLOCK_SIZE_N=BLOCK_SIZE_N, BLOCK_SIZE_K=BLOCK_SIZE_K, NUM_SMS=num_sms, FORWARD=forward, num_stages=num_stages, num_warps=8, ) return (c, aux) if forward else c # ============================================================================= # Autograd Functions — composes the fused kernel with the second matmul. # Compatible with torch.compile (autograd.Function is opaque to dynamo). # ============================================================================= class _FusedLinearReLU2(torch.autograd.Function): @staticmethod def forward(ctx, x, W1, W2): x_flat = x.reshape(-1, x.shape[-1]) pre, post = _launch_fused_act_square(_linear_relu_square_kernel, x_flat, W1, None) out = post @ W2.T ctx.save_for_backward(x, W1, W2, pre, post) return out.view(*x.shape[:-1], W2.shape[0]) @staticmethod def backward(ctx, grad_out): x, W1, W2, pre, post = ctx.saved_tensors x_flat = x.reshape(-1, x.shape[-1]) go = grad_out.reshape(-1, W2.shape[0]) dW2 = go.T @ post dpre = _launch_fused_act_square(_linear_relu_square_kernel, go, W2.T.contiguous(), pre) dW1 = dpre.T @ x_flat dx = dpre @ W1 return dx.view_as(x), dW1, dW2 class _FusedLinearLReLU2(torch.autograd.Function): @staticmethod def forward(ctx, x, W1, W2): x_flat = x.reshape(-1, x.shape[-1]) pre, post = _launch_fused_act_square(_linear_leaky_relu_square_kernel, x_flat, W1, None) out = post @ W2.T ctx.save_for_backward(x, W1, W2, pre, post) return out.view(*x.shape[:-1], W2.shape[0]) @staticmethod def backward(ctx, grad_out): x, W1, W2, pre, post = ctx.saved_tensors x_flat = x.reshape(-1, x.shape[-1]) go = grad_out.reshape(-1, W2.shape[0]) dW2 = go.T @ post dpre = _launch_fused_act_square(_linear_leaky_relu_square_kernel, go, W2.T.contiguous(), pre) dW1 = dpre.T @ x_flat dx = dpre @ W1 return dx.view_as(x), dW1, dW2 # ============================================================================= # Pure-torch fallbacks (torch.compile-friendly) # ============================================================================= def _eager_relu2_mlp(x, W1, W2): return F.relu(F.linear(x, W1)).square() @ W2.T def _eager_lrelu2_mlp(x, W1, W2): return F.leaky_relu(F.linear(x, W1), 0.5).square() @ W2.T def _can_use_triton_fused(x, M, N) -> bool: """The TMA kernels require shapes divisible by their tile sizes and a Hopper GPU.""" if not _hopper_available(): return False if x.dtype not in (torch.bfloat16, torch.float16): return False return (M % 128 == 0) and (N % 256 == 0) # ============================================================================= # VariableUDLP-compatible nn.Modules # ============================================================================= class _BaseTritonUDLP(Layer): """Shared structure: expand (d_model→M, no bias) -> fused act² -> contract (M→d_output, no bias). Matches `VariableUDLP(d_model, d_output, M, dropout)` so that it can be passed as the `UDLP=` argument of `ModdedFormer` without further changes. The `activation` parameter accepted by `VariableUDLP` is ignored here — the activation is fixed by the chosen kernel (subclass). """ _autograd_fn = None _eager_fn = None _name = "TritonUDLP" def __init__(self, d_model: int, d_output: int, M: int, activation=None, dropout: float = 0.0): super().__init__() self.d_model = d_model self.d_output = d_output self.M = M self.expand = nn.Linear(d_model, M, bias=False) self.contract = nn.Linear(M, d_output, bias=False) self.dropout = nn.Dropout(dropout) if activation is not None: flag(f"{self._name}: 'activation' arg ignored — kernel hardcodes the activation.", 3) def _init_weights(self): torch.nn.init.zeros_(self.contract.weight) def forward(self, x: torch.Tensor) -> torch.Tensor: W1 = self.expand.weight # [M, d_model] W2 = self.contract.weight # [d_output, M] if _can_use_triton_fused(x, self.M, self.d_output): out = self._autograd_fn.apply(x, W1, W2) else: out = self._eager_fn(x, W1, W2) return self.dropout(out) class TritonReLU2UDLP(_BaseTritonUDLP): """Fused ReLU² UDLP. Drop-in for VariableUDLP with activation=ReLU2.""" _autograd_fn = _FusedLinearReLU2 if _HAS_TRITON else None _eager_fn = staticmethod(_eager_relu2_mlp) _name = "TritonReLU2UDLP" class TritonLReLU2UDLP(_BaseTritonUDLP): """Fused LeakyReLU(0.5)² UDLP. Drop-in for VariableUDLP with activation=LeakyReLU2.""" _autograd_fn = _FusedLinearLReLU2 if _HAS_TRITON else None _eager_fn = staticmethod(_eager_lrelu2_mlp) _name = "TritonLReLU2UDLP" class TritonLReLU2UDLP_T4(_BaseTritonUDLP): """LeakyReLU(0.5)² UDLP tuned for T4 / Turing (sm_75). Uses a non-TMA fp16 kernel with 64×64×32 tiles. Falls back to the eager torch path on non-T4 GPUs or when triton is missing. NOTE: weights are kept fp32 by Linear; cast inputs to fp16 before forward (e.g. `model = model.half()` or `with torch.autocast('cuda', dtype=torch.float16)`). """ _autograd_fn = _FusedLinearLReLU2_T4 if _HAS_TRITON else None _eager_fn = staticmethod(_eager_lrelu2_mlp) _name = "TritonLReLU2UDLP_T4" def forward(self, x: torch.Tensor) -> torch.Tensor: W1 = self.expand.weight W2 = self.contract.weight if _can_use_triton_t4(x): out = self._autograd_fn.apply(x, W1, W2) else: out = self._eager_fn(x, W1, W2) return self.dropout(out) # ----------------------------------------------------------------------------- # SwiGLU — torch-only path (compile-friendly). # # SwiGLU is gated, so the expand projects to 2*M and the activation consumes # both halves: silu(a) * b. There is no clean port of this in modded-nanogpt # or parameter-golf; Liger-Kernel implements a full fused fwd+bwd for it but # that's a separate dependency. The torch path below compiles cleanly and is # usually within ~10% of a custom kernel for the typical M sizes we use. # ----------------------------------------------------------------------------- class TritonSwiGLUUDLP(Layer): """SwiGLU UDLP, drop-in for VariableUDLP with activation=SwiGLU. Compile-friendly torch implementation — no custom kernel. `M` is the post-gate inner dimension (i.e. expand projects to 2*M). """ def __init__(self, d_model: int, d_output: int, M: int, activation=None, dropout: float = 0.0): super().__init__() self.d_model = d_model self.d_output = d_output self.M = M self.expand = nn.Linear(d_model, 2 * M, bias=False) self.contract = nn.Linear(M, d_output, bias=False) self.dropout = nn.Dropout(dropout) if activation is not None: flag("TritonSwiGLUUDLP: 'activation' arg ignored — gate is fixed to SiLU.", 3) def _init_weights(self): torch.nn.init.zeros_(self.contract.weight) def forward(self, x: torch.Tensor) -> torch.Tensor: h = self.expand(x) a, b = h.chunk(2, dim=-1) return self.dropout(self.contract(F.silu(a) * b)) __all__ = [ "TritonReLU2UDLP", "TritonLReLU2UDLP", "TritonLReLU2UDLP_T4", "TritonSwiGLUUDLP", ]