For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).
Mojo function
register
def register(name: StringSlice[StaticConstantOrigin], type: StringSlice[StaticConstantOrigin] = StringSlice(""), api: StringSlice[StaticConstantOrigin] = StringSlice(""), arch: StringSlice[StaticConstantOrigin] = StringSlice(""), model: StringSlice[StaticConstantOrigin] = StringSlice(""))
Registers a struct as a kernel implementation for a MAX Graph op.
At compile time, the Graph Compiler selects the most specific registered kernel that matches the runtime device. The four device fields narrow the target from coarse to fine:
typeβ broad device category:"cpu","gpu", or a custom label.apiβ compute backend:"cuda","hip","metal", etc.archβ microarchitecture:"sm_100a","sm_90a","gfx942", etc.modelβ exact hardware model:"NVIDIA B200","AMD Instinct MI355X", etc.
An empty string for any field acts as a wildcard β it matches any device
value for that field.
Fields are ordered from least to most specific: type < api < arch <
model. Among all matching candidates, the most specific registration wins:
a registration that sets arch is more specific than one that sets only
type and api, because arch sits higher in the hierarchy.
If a user-defined kernel and a built-in kernel are equally specific for a device, the user-defined kernel takes precedence. Registering the same op with identical device fields more than once β whether within the same library or across multiple user libraries β is an error reported by the Graph Compiler.
Example β device field combinations and their reach:
# Matches all devices (backward-compatible default).
@compiler.register("mo.matmul")
# Matches all CPU devices.
@compiler.register("mo.matmul", type="cpu")
# Matches all GPU devices regardless of vendor or architecture.
@compiler.register("mo.matmul", type="gpu")
# Matches all CUDA GPUs (any architecture).
@compiler.register("mo.matmul", type="gpu", api="cuda")
# Matches only the NVIDIA SM100A architecture.
@compiler.register("mo.matmul", type="gpu", api="cuda", arch="sm_100a")
# Matches only the NVIDIA B200.
@compiler.register("mo.matmul", type="gpu", api="cuda", arch="sm_100a", model="NVIDIA B200")Example β selection with multiple registrations in scope:
@compiler.register("mo.matmul") # wildcard
@compiler.register("mo.matmul", type="cpu") # least specific
@compiler.register("mo.matmul", type="gpu") # least specific
@compiler.register("mo.matmul", type="gpu", api="cuda", arch="sm_100a") # more specific- SM100A CUDA GPU β all four match;
archregistration wins (most specific). - SM90A CUDA GPU β wildcard and
type="gpu"match;type="gpu"wins. - CPU β wildcard and
type="cpu"match;type="cpu"wins. - Some NPU β only the wildcard matches; wildcard wins.
Args:
- βname (
StringSlice[StaticConstantOrigin]): The MAX Graph op name to register (e.g."mo.matmul"). - βtype (
StringSlice[StaticConstantOrigin]): Broad device category β"cpu","gpu", or a custom accelerator label such as"npu-xxx". Corresponds to the label used inDeviceRef(e.g.DeviceRef.GPU()). Empty matches all. - βapi (
StringSlice[StaticConstantOrigin]): Programming API or compute backend for the device β for example"cuda"(NVIDIA CUDA),"hip"(AMD ROCm/HIP), or"metal"(Apple Metal). Empty matches all. - βarch (
StringSlice[StaticConstantOrigin]): Microarchitecture identifier β for example"sm_90a"(Hopper),"sm_100a"(Blackwell), or"gfx942"(AMD CDNA3). Empty matches all. - βmodel (
StringSlice[StaticConstantOrigin]): Specific device model β for example"NVIDIA B200"or"AMD Instinct MI355X". Empty matches all.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!