Skip to content

Getting Started

You have a model, and you notice that some part of it is taking a significant amount of time. Standard Kernels enables you to swap that part out for an optimized module without rewriting the surrounding model.

Example: speeding up a GEGLU block

Suppose that in your profiling, you find that this GEGLU block is a particular bottleneck in your workload.

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

This GEGLU block is a good candidate for fusion, as it often manifests as several separate kernel launches (the LayerNorms, the gate and up projections, the GELU, the elementwise multiply, the down projection, and the residual add are often each a kernel).

It would be great if we had a fused kernel on hand for this very GEGLU block! As it happens, Standard Kernels has exactly this ;)

Step 1: Check that your configuration is supported

Before you write any code, you'll want to confirm that Standard Kernels indeed has our desired GEGLU in the Kernel Coverage table.

So, going by example, Kernel Coverage tells us that yes, our pre-norm GEGLU feed-forward block is supported for bf16 on SM90:

Precision Hardware Shape Provenance Training Inference
bf16 SM90 (H100) (1024, 768), static In-house (SKC) ✓ ✓

Valid inputs outside optimized coverage

If you invoke our GEGLU on a shape/precision/hardware that we have not published a kernel variant for yet, Standard Kernels falls back to the equivalent PyTorch operations.

To request an optimized variant, run your workload and write a coverage report:

standardkernels.coverage.write_report("standardkernels-coverage.json")

Then submit an RFC and attach the report or copy and paste its contents.

If you do not want Standard Kernels to fall back to PyTorch, use strict mode:

with standardkernels.coverage.strict():
    output = block(x)

Strict mode is useful for benchmarks, CI coverage checks, and latency-sensitive deployments. In those cases, a mathematically correct PyTorch fallback could conceal a missing kernel and invalidate performance expectations. Ordinary training and inference should generally retain the fallback.

Step 2: Swapping in your optimized GEGLU module

Now you open the API reference for that particular GEGLU, and it tells you that:

  • the name of said op is standardkernels.ops.fused_pre_norm_geglu

  • the name of said module is standardkernels.nn.PreNormGEGLUFFN

We will go over how to use both ops and modules, with the former being the more common case.

Ops vs. modules

Our library provides two kinds of things:

  • Ops, under standardkernels.ops, are single fused kernels exposed as plain functions, for calling into a PyTorch module you're already writing yourself.
  • Modules, under standardkernels.nn, are complete nn.Modules we've rewritten to call our kernels internally, owning their own parameters like any other PyTorch module.

We recommend using modules only when there exists one that matches your model exactly. Otherwise, use ops to write a custom module of your own. If your module ends up solving a common enough problem that others would want it too, we highly recommend submitting it to our library!

Method 1: Writing your custom module with our optimized ops

First, if we don't have a prebuilt module that matches your model exactly, you can use our optimized ops to write a custom module of your own.

You likely already have your own feed-forward module, and it probably looks something like this:

class MyGEGLUBlock(torch.nn.Module):
    def __init__(self, d_model=768, d_ff=1448, eps=1e-5):
        super().__init__()
        self.my_norm = torch.nn.LayerNorm(d_model, eps=eps)
        self.my_gate = torch.nn.Linear(d_model, d_ff)
        self.my_up = torch.nn.Linear(d_model, d_ff)
        self.my_down = torch.nn.Linear(d_ff, d_model)
        self.eps = eps

    def forward(self, x):
        h = self.my_norm(x)
        return x + self.my_down(torch.nn.functional.gelu(self.my_gate(h)) * self.my_up(h))

Look there, in the forward() pass, there's our notoriously slow x + down(gelu(gate(norm(x))) * up(norm(x)))! Well, good news, Standard Kernels provides an optimized op, fused_pre_norm_geglu, that computes exactly that in a single fused kernel call.

Let's rewrite forward() to call our fused_pre_norm_geglu:

     def forward(self, x):
-        h = self.my_norm(x)
-        return x + self.my_down(torch.nn.functional.gelu(self.my_gate(h)) * self.my_up(h))
+        return standardkernels.ops.fused_pre_norm_geglu(
+            x,
+            self.my_norm.weight, self.my_norm.bias,
+            self.my_gate.weight, self.my_gate.bias,
+            self.my_up.weight, self.my_up.bias,
+            self.my_down.weight, self.my_down.bias,
+            self.eps,
+        )

With the forward pass implemented, now you can run inference on this module like any other Pytorch module.

block = MyGEGLUBlock().cuda()

# Inference
block.eval()
with torch.inference_mode():
    x = torch.randn(1024, 768, device="cuda", dtype=torch.bfloat16)
    output = block(x)

Method 2: Using our prebuilt module

Meanwhile, if your existing Pytorch module already matches the structure of one of our prebuilt modules exactly, you can skip writing your own module and use that prebuilt module instead.

For this example, we conveniently have the prebuilt module PreNormGEGLUFFN waiting for us.

We can import it like so:

from standardkernels.nn import PreNormGEGLUFFN

and run inference as per usual: 

block = PreNormGEGLUFFN().cuda()

# Inference
block.eval()
with torch.inference_mode():
    x = torch.randn(1024, 768, device="cuda", dtype=torch.bfloat16)
    output = block(x)

Additional workflows after Step 2

Once you know which op or module you want to use, the following workflows are independent. You can follow either one, both, or neither: training does not require CUDA Graphs, and CUDA Graph inference does not require completing the training section.

Training

There are two ways to swap in the module for training, depending on whether you want to swap it into a model that already has trained weights, or train a new model from scratch.

Case 1: "I'm training a new model from scratch"

Then you can construct the module directly, with randomly-initialized weights:

model = MyModel()
model.ffn = PreNormGEGLUFFN()
model = model.cuda().train()

Case 2: "I want to swap in the module on a model that already has trained weights"

You should use from_module to inherit the checkpoint's trained weights. Suppose the pre-trained model is model, and the module that we're swapping out is model.ffn.

model.ffn = PreNormGEGLUFFN.from_module(model.ffn).train()

This line effectively means, "Please create a module PreNormGEGLUFFN from this pre-existing module model.ffn, and swap the newly created module in place of model.ffn," which basically copies the trained weights to your new module and moves the new module to whichever device that the pre-trained model's parameters are already on.

Word on optimizers

Do NOT construct the optimizer first and then replace the module:

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
model.ffn = PreNormGEGLUFFN.from_module(model.ffn).train()

That won't work, since the optimizer still points to the old module's parameters instead of the newly copied parameters.

Instead, create the optimizer only after swapping PreNormGEGLUFFN into the model:

model.ffn = PreNormGEGLUFFN.from_module(model.ffn).train()
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

If you are resuming an optimizer from a checkpoint, construct it against the final model first and then call optimizer.load_state_dict(...). The parameter groups and their order must match the optimizer that wrote the checkpoint.

Keeping the original model for comparison

In our example, we swap our optimized module directly into the original model.

If you want to benchmark or compare the original and optimized models, you may construct a second model from the same checkpoint before swapping its FFN:

reference_model = MyModel()
reference_model.load_state_dict(checkpoint)

optimized_model = MyModel()
optimized_model.load_state_dict(checkpoint)
optimized_model.ffn = PreNormGEGLUFFN.from_module(
    optimized_model.ffn
).train()

If you are comparing training, you should give each model its own optimizer, since the two models and their parameters are independent.

Sample training loop

Once the final model and optimizer have been constructed, training works like an ordinary PyTorch training loop:

model.train()
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

for x, target in dataloader:
    x = x.to(device="cuda", dtype=torch.bfloat16)
    target = target.to(device="cuda")

    optimizer.zero_grad(set_to_none=True)

    output = model(x)
    loss = torch.nn.functional.mse_loss(
        output.float(),
        target.float(),
    )

    loss.backward()
    optimizer.step()

Where FP32 and BF16 are used

PreNormGEGLUFFN creates FP32 parameters by default. These are the master parameters updated by the optimizer, and their gradients are FP32 as well.

For the optimized variant, the input activation is BF16. LayerNorm reads that activation and its FP32 scale and bias, performs its reductions in FP32, and writes a BF16 normalized activation. The projection parameters may remain FP32 master parameters: the kernel maintains refreshed BF16 copies for projection compute, while its matrix multiplications accumulate in FP32. The residual output is BF16.

If you use CUDA BF16 autocast, an FP32 input activation is selectively converted to BF16 for this op while the FP32 LayerNorm and master parameters stay FP32:

with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
    output = model(x)  # x may start as FP32
    loss = loss_fn(output.float(), target.float())

Outside autocast, an FP32 input is still valid, but it uses the PyTorch fallback because the current optimized variant expects BF16.

CUDA Graph inference

Independently of training, the optimized GEGLU forward can be captured for inference. Use make_graphed_callables to warm up and capture the module once:

from standardkernels.cuda_graphs import make_graphed_callables

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

with torch.inference_mode():
    graphed_block = make_graphed_callables(
        block,
        (sample_x,),
    )

Then you can call the returned module (graphed_block) normally. It copies each new input into the graph's input buffers, replays the graph, and returns its static output:

with torch.inference_mode():
    output = graphed_block(next_x)

Every input must have the same shape, dtype, device, and layout as sample_x. Parameters are immutable for the lifetime of graphed_block. The returned output uses storage that a later call may overwrite, so clone it when it must remain unchanged across replays:

with torch.inference_mode():
    saved_output = graphed_block(next_x).clone()