# Examples of custom CPU / GPU operations in Mojo

**URL:** <https://forum.modular.com/t/examples-of-custom-cpu-gpu-operations-in-mojo/348>\
**Category:** MAX\
**Tags:** discussion, 24\_6\
**Created:** [December 18, 2024, 6:25pm UTC](https://forum.modular.com/t/examples-of-custom-cpu-gpu-operations-in-mojo/348 "2024-12-18T18:25:41Z")\
**Posts on this page:** 20\
**Page:** 1

<div class="post-metadata">

**Author:** ![BradLarson](https://sea1.discourse-cdn.com/flex001/user_avatar/forum.modular.com/bradlarson/32/62_2.png) [@BradLarson](https://forum.modular.com/u/BradLarson)\
**Post date:** [December 18, 2024, 6:25pm UTC](https://forum.modular.com/t/examples-of-custom-cpu-gpu-operations-in-mojo/348/1 "2024-12-18T18:25:41Z")

</div>

We have [a ton of stuff](https://docs.modular.com/max/changelog) launching with MAX in 24.6, including a preview of GPU support. All of those new additions have great new [docs](https://docs.modular.com) and tutorials (for example, [how to serve the Llama 3.1 model on GPUs at scale](https://docs.modular.com/max/tutorials/max-serve-local-to-cloud/)).

One undocumented and experimental capability that we’re sharing with the community is very early support for writing custom GPU operations in Mojo. We currently have two examples of this in the nightly branch of the MAX GitHub repository under [examples/custom\_ops](https://github.com/modularml/max/tree/nightly/examples/custom_ops): a very basic “hello world” sample that adds 1 to each element of a tensor, and a kernel that calculates the number of iterations to escape in the Mandelbrot set. Both use a Mojo API for defining custom operations, then they show how to [construct simple computational graphs in Python](https://docs.modular.com/max/tutorials/get-started-with-max-graph-in-python) and use these operations within them.

I’ll caution that these examples are very much subject to change or breakage as we work towards our next stable release, and have little to no documentation, but we wanted to provide an early preview of this capability to the MAX community. As people have already started to try out these GPU programming examples, I felt it would be useful to start a thread to discuss them and to gather issues and requests.

---

<div class="post-metadata">

**Author:** ![BradLarson](https://sea1.discourse-cdn.com/flex001/user_avatar/forum.modular.com/bradlarson/32/62_2.png) [@BradLarson](https://forum.modular.com/u/BradLarson)\
**Post date:** [December 19, 2024, 4:44pm UTC](https://forum.modular.com/t/examples-of-custom-cpu-gpu-operations-in-mojo/348/2 "2024-12-19T16:44:16Z")

</div>

One issue that was surfaced in a Discord discussion, and that I wanted to document here for reference: to use custom Mojo operations using our in-development extensibility API, you need to make sure the environment variable `MODULAR_ONLY_USE_NEW_EXTENSIBILITY_API` is set to `true`. This is done for you as part of the Magic invocation:

```auto
[tasks]
addition = { cmd="mojo package kernels/ -o kernels.mojopkg && python addition.py", env={ MODULAR_ONLY_USE_NEW_EXTENSIBILITY_API="true" }}

```

The approximate sequence of commands that would be used to build and run one of these examples if you were doing this at the shell level would be:

```auto
export MODULAR_ONLY_USE_NEW_EXTENSIBILITY_API=true
mojo package kernels/ -o kernels.mojopkg
python addition.py

```

This activates the new extensibility API, builds the Mojo custom ops into a package, and then uses the ops within a new graph on the GPU or CPU.

---

<div class="post-metadata">

**Author:** ![BradLarson](https://sea1.discourse-cdn.com/flex001/user_avatar/forum.modular.com/bradlarson/32/62_2.png) [@BradLarson](https://forum.modular.com/u/BradLarson)\
**Post date:** [January 11, 2025, 3:47pm UTC](https://forum.modular.com/t/examples-of-custom-cpu-gpu-operations-in-mojo/348/3 "2025-01-11T15:47:59Z")

</div>

As a quick update on these examples, the latest nightly contains a series of improvements:

- The addition example now shows a simpler way to create a single-operation graph in our Python API.
- The Mandelbrot set example now calculates the initial value for the complex C inside the operation, avoiding a wasteful calculation and copy from the host. It shows how you can define an operation with no tensor inputs, but that produces a tensor output.
- The latest nightly version of MAX no longer needs the environment variable `MODULAR_ONLY_USE_NEW_EXTENSIBILITY_API` to be set to true for these custom operations to work.
- The type for an accelerator `Device` in the Python Driver API in the latest nightlies has changed from `CUDA` to `Accelerator` to reflect the generalized nature of accelerator support in MAX.
- We’re in the process of gradually adding API documentation for the Mojo custom operation types and functions used for GPU programming. These will appear first in the nightly documentation, such as the new entry for [ManagedTensorSlice](https://docs.modular.com/nightly/max/api/mojo/tensor_utils/managed_tensor_slice/ManagedTensorSlice). Keep an eye on these nightly API docs for more.

---

<div class="post-metadata">

**Author:** ![BradLarson](https://sea1.discourse-cdn.com/flex001/user_avatar/forum.modular.com/bradlarson/32/62_2.png) [@BradLarson](https://forum.modular.com/u/BradLarson)\
**Post date:** [January 25, 2025, 3:41pm UTC](https://forum.modular.com/t/examples-of-custom-cpu-gpu-operations-in-mojo/348/4 "2025-01-25T15:41:00Z")

</div>

Note: we’re working on exposing a nightly changelog for MAX in the same way as has been done for Mojo. When that is live, you’ll be able to read these feature updates within the MAX changelog. Stay tuned!

The latest MAX nightly adds a lot of new content for GPU programming in MAX:

- As a source-breaking change, `ManagedTensorSlice` and `foreach` have moved from the `tensor_utils` module to `max.tensor`.
- Two new [custom operation programming examples](https://github.com/modular/max/tree/nightly/examples/custom_ops) have been added that show more of the power that Mojo provides for programming GPUs:
  - The `vector_addition` example demonstrates how to write device-specific codepaths, as well as how to manually dispatch Mojo functions on the GPU within a custom operation. This mode of programming may be much more familiar to those used to CUDA C programming. Note that the `foreach` abstraction performs elementwise calculations far more efficiently than the manual functions here, due to its hardware optimization. This is merely an instructive example.
  - The `top_k` example shows a practical use case for a custom operation that is used today within large language model graphs: a [top-K token sampler](https://dataforest.ai/glossary/top-k-sampling). The Mojo code contains a much more complex calculation, as well as how to construct a custom shape function for the operation. The Python-side code also hosts a showcase for how such an operation is used in practice.

- The `synchronous` parameter has been removed from the interfaces in the custom operation examples. It was only useful in a few cases, and we’re evaluating removing it overall. That simplifies the overall operation interface.
- Once the nightly docs update (the team is working hard on deploying these right now), initial API docs will appear for the `gpu` and `layout` modules and their dependencies.

---

<div class="post-metadata">

**Author:** ![BradLarson](https://sea1.discourse-cdn.com/flex001/user_avatar/forum.modular.com/bradlarson/32/62_2.png) [@BradLarson](https://forum.modular.com/u/BradLarson)\
**Post date:** [February 6, 2025, 3:54pm UTC](https://forum.modular.com/t/examples-of-custom-cpu-gpu-operations-in-mojo/348/5 "2025-02-06T15:54:59Z")

</div>

A couple more recent additions to the custom operations programming examples:

- We added a GPU-specific path in [the top-K token sampler custom operation example](https://github.com/modular/max/blob/main/examples/custom_ops/kernels/top_k.mojo) which demonstrates a few more GPU programming techniques in Mojo.
- Custom operations can now be compile-time parameterized on user-supplied custom Int or StringLiteral parameters that you pass in via the Graph API. The [parametric\_addition](https://github.com/modular/max/blob/main/examples/custom_ops/parametric_addition.py) example shows how this works, and I had a little fun with it when implementing [image blend modes](https://github.com/BradLarson/max-cv/blob/main/operations/blend.mojo#L65) parameterized by the name of the blend function.

---

<div class="post-metadata">

**Author:** ![owenhilyard](https://sea1.discourse-cdn.com/flex001/user_avatar/forum.modular.com/owenhilyard/32/41_2.png) [@owenhilyard](https://forum.modular.com/u/owenhilyard)\
**Post date:** [February 6, 2025, 4:46pm UTC](https://forum.modular.com/t/examples-of-custom-cpu-gpu-operations-in-mojo/348/6 "2025-02-06T16:46:29Z")

</div>

In the top-K example, `var top_k_sram = external_memory[...]()` is pulling data from the array with size given by `shared_mem_bytes`. Is there a reason to do it this way instead of allocating the buffer beforehand and passing it as an argument?

---

<div class="post-metadata">

**Author:** ![MMabrouk](https://avatars.discourse-cdn.com/v4/letter/m/3ec8ea/32.png) [@MMabrouk](https://forum.modular.com/u/MMabrouk)\
**Post date:** [February 6, 2025, 5:14pm UTC](https://forum.modular.com/t/examples-of-custom-cpu-gpu-operations-in-mojo/348/7 "2025-02-06T17:14:42Z")

</div>

the work on custom-ops is really interesting. AFAIK you still need the MAX engine to compile/run/partition the computational graph. Any hope that we can run this graph from mojo soon?

---

<div class="post-metadata">

**Author:** ![jack.clayton](https://sea1.discourse-cdn.com/flex001/user_avatar/forum.modular.com/jack.clayton/32/57_2.png) [@jack.clayton](https://forum.modular.com/u/jack.clayton)\
**Post date:** [February 6, 2025, 5:36pm UTC](https://forum.modular.com/t/examples-of-custom-cpu-gpu-operations-in-mojo/348/9 "2025-02-06T17:36:38Z")

</div>

Hi @owenhilyard `shared_mem_bytes` shares memory between every thread in a block, we don’t need to access that memory from the CPU host. It’d be wasted resources to allocate it on host especially if it’s a very large batch size in this example, as each batch gets it’s own block, so for a 1M size batch we’d need to allocate `100,000,000 * K * sizeof[T]`.

---

<div class="post-metadata">

**Author:** ![BradLarson](https://sea1.discourse-cdn.com/flex001/user_avatar/forum.modular.com/bradlarson/32/62_2.png) [@BradLarson](https://forum.modular.com/u/BradLarson)\
**Post date:** [February 6, 2025, 5:40pm UTC](https://forum.modular.com/t/examples-of-custom-cpu-gpu-operations-in-mojo/348/10 "2025-02-06T17:40:28Z")

</div>

MAX has a lot of interface levels from low to high: from Driver to Graph to Serve. Mojo is part of MAX as a language to define high-performance computation. So far, we’ve shown how to use Mojo to program GPUs within the Graph portion of the API, and you can see the use of the Driver to allocate memory on host, manage interactions with accelerators, and transfer memory between host and device.

There are both Python and Mojo interfaces to the Driver and Graph APIs, although the Python APIs have gotten the most attention recently for various reasons: progressive introduction into Python codebases, Mojo has changed a lot since these were first designed, etc. The Python interfaces provide zero-cost direct interoperability with PyTorch tensors and NumPy arrays, for example, which is huge for usability. When staging a computational graph, overhead from the host language tends to be minimal, so Python works well for graph computation.

We believe that moving computation into a graph is the best path for large-scale computation, like is done in ML models. However, we also recognize that there’s some overhead in doing so for setting up and compiling a graph. Without spoiling anything, we’re working on something in MAX that I think you’ll like and will satisfy your core question here.

---

<div class="post-metadata">

**Author:** ![owenhilyard](https://sea1.discourse-cdn.com/flex001/user_avatar/forum.modular.com/owenhilyard/32/41_2.png) [@owenhilyard](https://forum.modular.com/u/owenhilyard)\
**Post date:** [February 6, 2025, 5:47pm UTC](https://forum.modular.com/t/examples-of-custom-cpu-gpu-operations-in-mojo/348/11 "2025-02-06T17:47:14Z")

</div>

I was imagining using cudaMalloc (or a mojo binding) to allocate the memory on the GPU and get a device memory pointer back, then you submit those 8 bytes as an argument to the kernel. This would let MAX reuse that buffer if multiple calls are made.

---

<div class="post-metadata">

**Author:** ![jack.clayton](https://sea1.discourse-cdn.com/flex001/user_avatar/forum.modular.com/jack.clayton/32/57_2.png) [@jack.clayton](https://forum.modular.com/u/jack.clayton)\
**Post date:** [February 6, 2025, 8:07pm UTC](https://forum.modular.com/t/examples-of-custom-cpu-gpu-operations-in-mojo/348/12 "2025-02-06T20:07:07Z")

</div>

cudaMalloc allocates to global memory, and is accessible from any thread in any block, it’s not as fast as shared memory which is only shared by a single block

---

<div class="post-metadata">

**Author:** ![owenhilyard](https://sea1.discourse-cdn.com/flex001/user_avatar/forum.modular.com/owenhilyard/32/41_2.png) [@owenhilyard](https://forum.modular.com/u/owenhilyard)\
**Post date:** [February 6, 2025, 8:08pm UTC](https://forum.modular.com/t/examples-of-custom-cpu-gpu-operations-in-mojo/348/13 "2025-02-06T20:08:13Z")

</div>

Aha, that makes sense, thank you.

---

<div class="post-metadata">

**Author:** ![arthur](https://sea1.discourse-cdn.com/flex001/user_avatar/forum.modular.com/arthur/32/139_2.png) [@arthur](https://forum.modular.com/u/arthur)\
**Post date:** [February 10, 2025, 6:23pm UTC](https://forum.modular.com/t/examples-of-custom-cpu-gpu-operations-in-mojo/348/14 "2025-02-10T18:23:41Z")

</div>

Just adding to Jack’s answer. In the custom ops examples, the input tensors are allocated in the Python code and moved to device (GPU) memory. So when you see the custom op accessing `ManagedTensorSlice` instances, it’s accessing a pointer to GPU global memory.

There are also low-level Mojo APIs for allocating GPU memory and copying to/from memory, including methods in [`gpu.memory`](https://docs.modular.com/mojo/stdlib/gpu/memory/) to allocate and copy to/from shared memory. There are also methods in `gpu.host.DeviceContext` for allocating global memory buffers, but you’re less likely to need those in a custom op.

The driver API also provides methods for moving data to and from the GPU. We’re hoping to add API docs for the Mojo driver API soon.

---

<div class="post-metadata">

**Author:** ![owenhilyard](https://sea1.discourse-cdn.com/flex001/user_avatar/forum.modular.com/owenhilyard/32/41_2.png) [@owenhilyard](https://forum.modular.com/u/owenhilyard)\
**Post date:** [February 10, 2025, 7:04pm UTC](https://forum.modular.com/t/examples-of-custom-cpu-gpu-operations-in-mojo/348/15 "2025-02-10T19:04:39Z")

</div>

Thanks for the extra context! I spend a fair amount of energy trying to design programs to do all of their allocation up front, so I’ll need to spend some time poking around there.

I’m also hoping to see some unified memory options in there at some point since DIGITS is rumored to have MI300A-style unified memory.

---

<div class="post-metadata">

**Author:** ![TilliFe](https://sea1.discourse-cdn.com/flex001/user_avatar/forum.modular.com/tillife/32/419_2.png) [@TilliFe](https://forum.modular.com/u/TilliFe)\
**Post date:** [February 12, 2025, 11:22am UTC](https://forum.modular.com/t/examples-of-custom-cpu-gpu-operations-in-mojo/348/16 "2025-02-12T11:22:09Z")

</div>

Re. GPU Sessions in Mojo API:

Is it already possible in the Mojo API of MAX (nightly) to create Graphs that execute on a GPU?

For example, I would like to have the following code translated from Python to Mojo:

```Mojo
def main():
    rows = 2
    columns = 3
    dtype = DType.float32

    graph = Graph(
        name="graph-with-single-sin-op",
        forward=lambda x: ops.sin(x),
        input_types=[TensorType(dtype, shape=[rows, columns])]
    )

    # Create session on GPU if available, otherwise use CPU.
    device = CPU() if accelerator_count() == 0 else Accelerator()
    session = InferenceSession(devices=[device],)
    model = session.load(graph)

    # creat inputs and execute model
    x_values = np.random.uniform(size=(rows, columns)).astype(np.float32)
    x = Tensor.from_numpy(x_values).to(device)
    result = model.execute(x)[0]
    assert isinstance(result, Tensor)
    result = result.to(CPU())

```

To clarify, I don’t target custom ops for now, just the existing max.ops. If this is already possible, how would that look like?

---

<div class="post-metadata">

**Author:** ![owenhilyard](https://sea1.discourse-cdn.com/flex001/user_avatar/forum.modular.com/owenhilyard/32/41_2.png) [@owenhilyard](https://forum.modular.com/u/owenhilyard)\
**Post date:** [February 12, 2025, 7:37pm UTC](https://forum.modular.com/t/examples-of-custom-cpu-gpu-operations-in-mojo/348/17 "2025-02-12T19:37:42Z")

</div>

I think a direct translation looks something like this:

```mojo
from max.graph import Graph
from max.graph import TensorType, Type
from sys.info import has_nvidia_gpu_accelerator
from max.driver import cpu_device, DeviceTensor
from max.driver._cuda import cuda_device
from max.engine import InferenceSession, SessionOptions
from tensor import TensorSpec
from max.graph.ops import sin
import random

fn build_graph(tensor_type: TensorType) raises -> Graph:
    var graph = Graph(
        name="graph-with-single-sin-op",
        in_types=List(Type(tensor_type)),
        out_types=List(Type(tensor_type)),
    )

    graph.output(sin(graph[0])) 

    return graph

fn main() raises:
    alias rows = 2
    alias columns = 3
    alias dtype = DType.float32
    alias tensor_type = TensorType(dtype, rows, columns)

    var graph = build_graph(tensor_type)
    var device = cuda_device() if has_nvidia_gpu_accelerator() else cpu_device() 
    var session = InferenceSession(SessionOptions(device))
    var model = session.load(graph)

    var input_tensor = DeviceTensor(TensorSpec(dtype, rows, columns), cpu_device(), None).to_tensor[dtype, 2]()
    
    var num_bytes = input_tensor.spec().num_elements()
    var ptr = input_tensor.unsafe_ptr()
    
    random.rand(ptr, num_bytes)

    var device_tensor = input_tensor.to_device_tensor().move_to(device)

    var result = model.execute(device_tensor)
    var output_tensor = result[0].take().to_device_tensor().move_to(cpu_device())

    var output = output_tensor.to_tensor[dtype, 2]()

    print(output)

```

No custom ops required since `max.graph.ops.sin` exists.

---

<div class="post-metadata">

**Author:** ![BradLarson](https://sea1.discourse-cdn.com/flex001/user_avatar/forum.modular.com/bradlarson/32/62_2.png) [@BradLarson](https://forum.modular.com/u/BradLarson)\
**Post date:** [February 12, 2025, 7:41pm UTC](https://forum.modular.com/t/examples-of-custom-cpu-gpu-operations-in-mojo/348/18 "2025-02-12T19:41:04Z")

</div>

The Mojo Graph and Driver APIs do support running graphs (with custom ops) on a GPU. I will warn that the Mojo Graph and Driver interfaces have diverged from the Python APIs, and may not have all the features that are present in the Python APIs.

With the latest MAX nightly, the following code should place a graph on an accelerator. For inputs and outputs, Driver Tensors are allocated (note: not a `max.tensor` Tensor) and placed on the accelerator, with the input being initialized on host and moved to accelerator and output being retrieved from the accelerator and moved to the host.

```auto
from max.driver import accelerator_device, cpu_device, Tensor, Device
from max.engine import InferenceSession, SessionOptions
from max.graph import Graph, TensorType, ops

def main():
    alias VECTOR_WIDTH = 6

    host_device = cpu_device()
    gpu_device = accelerator_device()

    graph = Graph(TensorType(DType.float32, VECTOR_WIDTH))
    result = ops.sin(graph[0])
    graph.output(result)

    options = SessionOptions(gpu_device)
    session = InferenceSession(options)
    model = session.load(graph)

    input_tensor = Tensor[DType.float32, 1]((VECTOR_WIDTH), host_device)
    for i in range(VECTOR_WIDTH):
        input_tensor[i] = 1.25

    results = model.execute(input_tensor^.move_to(gpu_device))
    output = results[0].take().to_device_tensor().move_to(host_device).to_tensor[DType.float32, 1]()
    print(output)

```

Owen and I collided on this, but one difference in our two examples is that we recently generalized the Mojo Driver interface to not be hardcoded to CUDA-compatible devices to start to align with the Python Driver API.

---

<div class="post-metadata">

**Author:** ![TilliFe](https://sea1.discourse-cdn.com/flex001/user_avatar/forum.modular.com/tillife/32/419_2.png) [@TilliFe](https://forum.modular.com/u/TilliFe)\
**Post date:** [February 13, 2025, 8:36am UTC](https://forum.modular.com/t/examples-of-custom-cpu-gpu-operations-in-mojo/348/19 "2025-02-13T08:36:20Z")

</div>

Works perfectly! Thank you so much.

---

<div class="post-metadata">

**Author:** ![TilliFe](https://sea1.discourse-cdn.com/flex001/user_avatar/forum.modular.com/tillife/32/419_2.png) [@TilliFe](https://forum.modular.com/u/TilliFe)\
**Post date:** [February 13, 2025, 3:36pm UTC](https://forum.modular.com/t/examples-of-custom-cpu-gpu-operations-in-mojo/348/20 "2025-02-13T15:36:44Z")

</div>

The Mojo API Reference about the MAX Driver Module used to be present on the Modular Webpage, but at the moment it seems to be removed both from the stable and nightly API. I’d be super glad to have it back! 🙏

---

<div class="post-metadata">

**Author:** ![joshpeterson](https://sea1.discourse-cdn.com/flex001/user_avatar/forum.modular.com/joshpeterson/32/14_2.png) [@joshpeterson](https://forum.modular.com/u/joshpeterson)\
**Post date:** [February 13, 2025, 3:49pm UTC](https://forum.modular.com/t/examples-of-custom-cpu-gpu-operations-in-mojo/348/21 "2025-02-13T15:49:31Z")

</div>

Yeah, this was removed a few months ago to avoid confusion with the Python Driver API as that was brought up. We’re looking to publish the Mojo Driver API docs soon though. They will likely be in nightly docs first.

[Next page](https://forum.modular.com/t/examples-of-custom-cpu-gpu-operations-in-mojo/348.md?page=2)