Skip to content

Registering a kernel shape

GEGLU dispatch uses exact variants. Non-matching calls use the PyTorch fallback and record a coverage miss.

Register the variant

Put the implementation under src/standardkernels/_kernels/, then add one entry to src/standardkernels/variants.toml:

[[variants]]
op = "standardkernels.fused_pre_norm_geglu"
name = "GEGLU"
precision = "bf16"
hardware = "SM90a (H100)"
provenance = "In-house (SKC)"
train = true
infer = true
backward_kernel = false

[variants.target]
runtime = "cuda"
architecture = "sm90a"

[variants.implementation]
entrypoint = "standardkernels._kernels.cuda.geglu_bf16_sm90_512"
language = "cuda-cpp"

[variants.constraints]
input_shape = [512, 768]
d_ff = 1448
input_dtype = "bfloat16"
norm_dtype = "float32"
eps = 1e-5

The TOML entry tells Standard Kernels when your kernel can be used. target identifies where it runs (hardware, runtime environment), while the constraints table describes the inputs that it supports, and implementation.entrypoint points to its Python kernel file.

The kernel file must also provide a get_kernel() function. Standard Kernels uses this function to load and run the implementation.

Test the shape

Follow Using sk-check on the target accelerator with strict coverage. The candidate must import the public API so fallback cannot pass unnoticed.