API reference
The API reference is generated from public Python docstrings by mkdocstrings.
Available computations
| Computation | Op | Module |
|---|---|---|
| Pre-norm GEGLU | standardkernels.ops.fused_pre_norm_geglu |
standardkernels.nn.PreNormGEGLUFFN |
standardkernels.ops
Stateless functional operators backed by Standard Kernels.
fused_pre_norm_geglu(x, norm_weight, norm_bias, gate_weight, gate_bias, up_weight, up_bias, down_weight, down_bias, eps)
This PyTorch op computes x + down(gelu(gate(norm(x))) * up(norm(x))).
To see what shapes or precisions are supported, see GEGLU kernel coverage and submit an RFC if needed.
The optimized CUDA variant is used when available. Otherwise-valid calls
use equivalent PyTorch math and record a local coverage miss. Use
:func:standardkernels.coverage.strict in tests or benchmarks when falling
back should be an error.
standardkernels.nn
Parameter-owning PyTorch modules backed by Standard Kernels.
PreNormGEGLUFFN
Bases: Module
This module computes x + down(gelu(gate(norm(x))) * up(norm(x))).
To see what shapes or precisions are supported, see GEGLU kernel coverage and submit an RFC if needed.
Optimized variants are selected automatically. Valid inputs outside current optimized coverage use PyTorch math and remain trainable.
from_module(module)
classmethod
Copy a structurally equivalent module without sharing parameters.
Construct the optimizer only after installing the returned module. An existing optimizer continues to reference the source parameters.
prepare_cuda_graph(x)
Pin this module's prepared binding for one CUDA Graph.
Execution features
standardkernels.cuda_graphs
CUDA Graph preparation and resource-lifetime helpers.
CudaGraphBindingHandle
Pin one private GEGLU binding for exactly one CUDA Graph capture.
capture(graph)
Activate this binding around one torch.cuda.graph block.
replay()
Serialize one replay on the stream used to prepare this handle.
release()
Release pinned resources after the graph's final replay.
prepare_fused_pre_norm_geglu(x, norm_weight, norm_bias, gate_weight, gate_bias, up_weight, up_bias, down_weight, down_bias, eps)
Prepare one private optimized binding for a single CUDA Graph capture.
make_graphed_callables(callables, sample_args, **kwargs)
Wrap supported Standard Kernels modules with PyTorch CUDA Graph callables.
Diagnostics
standardkernels.coverage
Local observability for valid calls that miss optimized kernel coverage.
KernelCoverageWarning
Bases: UserWarning
A valid operation used PyTorch because no optimized variant matched.
UnsupportedKernelVariantError
Bases: RuntimeError
Strict mode rejected an otherwise-valid optimized coverage miss.
write_report(path)
Write coverage misses as formatted JSON without sending it anywhere.
clear()
Atomically clear counters and warn-once state.
Callbacks already in flight are allowed to finish.
add_callback(callback)
Subscribe to each coverage miss; duplicate registration is ignored.
remove_callback(callback)
Remove a callback if it is currently registered.
strict(enabled=True)
Temporarily turn optimized coverage misses into errors.
The setting is nestable and local to the current thread or async context.