Getting Started: GEGLU

GEGLU is a gated feed-forward activation used in Transformer-style networks. The current Standard Kernels module fuses a complete pre-norm feed-forward block:

x + down(gelu(gate(layer_norm(x))) * up(layer_norm(x)))

Run the fused module

import torch
from standardkernels.nn import PreNormGEGLUFFN

block = PreNormGEGLUFFN().cuda().eval()
x = torch.randn(1024, 768, device="cuda", dtype=torch.bfloat16)

with torch.inference_mode():
    output = block(x)

print(output.shape)

Current contract

Property	Current proof-of-concept contract
Input shape	(1024, 768)
Input dtype	torch.bfloat16
LayerNorm parameters	torch.float32
Device	CUDA SM90 with 132 SMs
Execution	Inference / no gradient tracking
Live instances	One process-global binding
Dynamic shapes	Not yet supported
Backpropagation	Not yet supported

The module owns ordinary PyTorch parameters, but its fused extension has no backward implementation. Do not use it in training or gradient-enabled forward.