IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo struct

DeviceContext

struct DeviceContext

Represents a single stream of execution on a particular accelerator (GPU).

A DeviceContext serves as the low-level interface to the accelerator inside a MAX custom operation and provides methods for allocating buffers on the device, copying data between host and device, and for compiling and running functions (also known as kernels) on the device.

The device context can be used as a context manager. For example:

from max.gpu.host import DeviceContext
from std.gpu import thread_idx

def kernel():
    print("hello from thread:", thread_idx.x, thread_idx.y, thread_idx.z)

with DeviceContext() as ctx:
    ctx.enqueue_function[kernel](grid_dim=1, block_dim=(2, 2, 2))
    ctx.synchronize()

A custom operation receives a DeviceContext directly:

from max.gpu.host import DeviceContext
from extensibility import register

@register("custom_op")
struct CustomOp:
    @staticmethod
    def execute(ctx: DeviceContext) raises:
        ctx.enqueue_function[kernel, kernel](grid_dim=1, block_dim=(2, 2, 2))
        ctx.synchronize()

Implemented traits​

AnyType, Copyable, Deinitable, ImplicitlyCopyable, Movable, RegisterPassable, _FunctionEnqueuer

comptime members​

default_device_info​

comptime default_device_info = GPUInfo.from_name[_accelerator_arch()]()

GPUInfo object for the default accelerator.

Methods​

__init__​

def __init__(out self, device_id: Int = Int(0), *, var api: String = DeviceContext.default_device_info.api)

Constructs a DeviceContext for the specified device.

This initializer creates a new device context for the specified accelerator device. The device context provides an interface for interacting with the GPU, including memory allocation, data transfer, and kernel execution.

Example:

from max.gpu.host import DeviceContext

# Create a context for the default GPU
var ctx = DeviceContext()

# Create a context for a specific GPU (device 1)
var ctx2 = DeviceContext(1)

Args:

  • ​device_id (Int): ID of the accelerator device. If not specified, uses the default accelerator (device 0).
  • ​api (String): Requested device API (for example, "cuda" or "hip"). Defaults to the device API specified by current target accelerator.

Raises:

If device initialization fails or the specified device is not available.

def __init__(*, copy: Self) -> Self

Creates a copy of an existing device context by incrementing its reference count.

This copy constructor creates a new reference to the same underlying device context by incrementing the reference count of the native context object. Both the original and the copy will refer to the same device context.

Args:

  • ​copy (Self): The device context to copy.

__deinit__​

def __deinit__(deinit self)

Releases resources associated with this device context.

This destructor decrements the reference count of the native device context. When the reference count reaches zero, the underlying resources are released, including any cached memory buffers and compiled device functions.

enqueue​

def enqueue[args_origin: MutOrigin, //](self, func_handle: Optional[Pointer[_DeviceFunctionCpp, MutUntrackedOrigin]], grid_dim: Dim, block_dim: Dim, shared_mem_bytes: Int, attributes: Pointer[LaunchAttribute], num_attributes: Int, args: Pointer[Pointer[NoneType, args_origin]], arg_count: UInt32, arg_sizes: Optional[Pointer[UInt64, origin]]) -> Optional[CStringSlice[ImmUntrackedOrigin]]

Enqueues a kernel launch on this context's default stream.

Forwards directly to AsyncRT_DeviceContext_enqueueFunctionDirect. See _FunctionEnqueuer.enqueue for the full contract.

Args:

Returns:

Optional[CStringSlice[ImmUntrackedOrigin]]: A C-string carrying an error message on failure, or an empty string on success.

__enter__​

def __enter__(var self) -> Self

Enables the use of DeviceContext in a 'with' statement context manager.

This method allows DeviceContext to be used with Python-style context managers, which ensures proper resource management and cleanup when the context exits.

Example:

from max.gpu.host import DeviceContext

# Using DeviceContext as a context manager
with DeviceContext() as ctx:
    # Perform GPU operations
    # Resources are automatically released when exiting the block
    pass

Returns:

Self: The DeviceContext instance to be used within the context manager block.

name​

def name(self) -> String

Returns the device name, an ASCII string identifying this device, defined by the native device API.

This method queries the underlying GPU device for its name, which typically includes the model and other identifying information. This can be useful for logging, debugging, or making runtime decisions based on the specific GPU hardware.

Example:

from max.gpu.host import DeviceContext

var ctx = DeviceContext()
print("Running on device:", ctx.name())

Returns:

String: A string containing the device name.

api​

def api(self) -> String

Returns the name of the API used to program the device.

This method queries the underlying device context to determine which GPU programming API is being used for the current device. This information is useful for writing code that can adapt to different GPU architectures and programming models.

Possible values are:

  • "cpu": Generic host device (CPU).
  • "cuda": NVIDIA GPUs.
  • "hip": AMD GPUs.

Example:

from max.gpu.host import DeviceContext

var ctx = DeviceContext()
var api_name = ctx.api()
print("Using device API:", api_name)

# Conditionally execute code based on the API
if api_name == "cuda":
    print("Running on NVIDIA GPU")
elif api_name == "hip":
    print("Running on AMD GPU")

Returns:

String: A string identifying the device API.

enqueue_create_buffer​

def enqueue_create_buffer[dtype: DType](self, size: Int) -> DeviceBuffer[dtype]

Enqueues a buffer creation using the DeviceBuffer constructor.

For GPU devices, the space is allocated in the device's global memory.

Parameters:

  • ​dtype (DType): The data type to be stored in the allocated memory.

Args:

  • ​size (Int): The number of elements of type to allocate memory for.

Returns:

DeviceBuffer[dtype]: The allocated buffer.

Raises:

If the operation fails.

create_buffer_sync​

def create_buffer_sync[dtype: DType](self, size: Int) -> DeviceBuffer[dtype]

Creates a buffer synchronously using the DeviceBuffer constructor.

Parameters:

  • ​dtype (DType): The data type to be stored in the allocated memory.

Args:

  • ​size (Int): The number of elements of type to allocate memory for.

Returns:

DeviceBuffer[dtype]: The allocated buffer.

Raises:

If the operation fails.

enqueue_create_host_buffer​

def enqueue_create_host_buffer[dtype: DType](self, size: Int) -> HostBuffer[dtype]

Enqueues the creation of a HostBuffer.

This function allocates memory on the host that is accessible by the device. The memory is page-locked (pinned) for efficient data transfer between host and device.

Pinned memory is guaranteed to remain resident in the host's RAM, not be paged/swapped out to disk. Memory allocated normally (for example, using alloc()) is pageable—individual pages of memory can be moved to secondary storage (disk/SSD) when main memory fills up.

Using pinned memory allows devices to make fast transfers between host memory and device memory, because they can use direct memory access (DMA) to transfer data without relying on the CPU.

Allocating too much pinned memory can cause performance issues, since it reduces the amount of memory available for other processes.

Example:

from max.gpu.host import DeviceContext

with DeviceContext() as ctx:
    # Allocate host memory accessible by the device
    var host_buffer = ctx.enqueue_create_host_buffer[DType.float32](1024)

    # Use the host buffer for device operations
    # ...

Parameters:

  • ​dtype (DType): The data type to be stored in the allocated memory.

Args:

  • ​size (Int): The number of elements of type to allocate memory for.

Returns:

HostBuffer[dtype]: A HostBuffer object that wraps the allocated host memory.

Raises:

If memory allocation fails or if the device context is invalid.

compile_function​

def compile_function[declared_arg_types: TypeList[declared_arg_types.values], //, func: def(*args: *declared_arg_types) thin -> None, *, compile_options: StringSlice[ImmStaticOrigin] = CompilationTarget.default_compile_options(), link_options: StringSlice[ImmStaticOrigin] = StringSlice(""), dump_asm: Variant[Bool, Path, StringSlice[ImmStaticOrigin], def() capturing thin -> Path] = False, dump_llvm: Variant[Bool, Path, StringSlice[ImmStaticOrigin], def() capturing thin -> Path] = False, _dump_sass: Variant[Bool, Path, StringSlice[ImmStaticOrigin], def() capturing thin -> Path] = False, _ptxas_info_verbose: Bool = False](self, *, func_attribute: OptionalReg[FuncAttribute] = None, out result: DeviceFunction[func, declared_arg_types, target=Self.default_device_info.target(), compile_options=compile_options, link_options=link_options, _ptxas_info_verbose=_ptxas_info_verbose])

Compiles the provided function for execution on this device.

Parameters:

Args:

  • ​func_attribute (OptionalReg[FuncAttribute]): An attribute to use when compiling the code (such as maximum shared memory size).

Returns:

DeviceFunction[func, declared_arg_types, target=Self.default_device_info.target(), compile_options=compile_options, link_options=link_options, _ptxas_info_verbose=_ptxas_info_verbose]: The compiled function via the result output parameter.

Raises:

If the operation fails.

def compile_function[declared_arg_types: TypeList[declared_arg_types.values], //, func: def(*args: *declared_arg_types) capturing thin -> None, *, compile_options: StringSlice[ImmStaticOrigin] = CompilationTarget.default_compile_options(), link_options: StringSlice[ImmStaticOrigin] = StringSlice(""), dump_asm: Variant[Bool, Path, StringSlice[ImmStaticOrigin], def() capturing thin -> Path] = False, dump_llvm: Variant[Bool, Path, StringSlice[ImmStaticOrigin], def() capturing thin -> Path] = False, _dump_sass: Variant[Bool, Path, StringSlice[ImmStaticOrigin], def() capturing thin -> Path] = False, _ptxas_info_verbose: Bool = False](self, *, func_attribute: OptionalReg[FuncAttribute] = None, out result: DeviceFunction[func, declared_arg_types, target=Self.default_device_info.target(), compile_options=compile_options, link_options=link_options, _ptxas_info_verbose=_ptxas_info_verbose])

Compiles the provided function for execution on this device.

Parameters:

Args:

  • ​func_attribute (OptionalReg[FuncAttribute]): An attribute to use when compiling the code (such as maximum shared memory size).

Returns:

DeviceFunction[func, declared_arg_types, target=Self.default_device_info.target(), compile_options=compile_options, link_options=link_options, _ptxas_info_verbose=_ptxas_info_verbose]: The compiled function via the result output parameter.

Raises:

If the operation fails.

load_function​

def load_function[func_type: TrivialRegisterPassable, //, func: func_type](self, *, var function_name: String, var asm: String, func_attribute: OptionalReg[FuncAttribute] = None, out result: DeviceExternalFunction)

Loads a pre-compiled device function from assembly code.

This method loads an external GPU function from provided assembly code (PTX/SASS) rather than compiling it from Mojo source. This is useful for integrating with existing CUDA/HIP code or for using specialized assembly optimizations.

Example:

from max.gpu.host import DeviceContext
from max.gpu.host.device_context import DeviceExternalFunction

def func_signature(
    # Arguments being passed to the assembly code
    # e.g. two pointers and a length
    input: Pointer[Float32],
    output: Pointer[Float32],
    len: Int,
):
    # No body because that is passed as assembly code below.
    pass

var ctx = DeviceContext()
var ptx_code = "..."  # PTX assembly code
var ext_func = ctx.load_function[func_signature](
    function_name="my_kernel",
    asm=ptx_code,
)

Parameters:

  • ​func_type (TrivialRegisterPassable): The dtype of the function to load.
  • ​func (func_type): The function reference.

Args:

  • ​function_name (String): The name of the function in the assembly code.
  • ​asm (String): The assembly code (PTX/SASS) containing the function.
  • ​func_attribute (OptionalReg[FuncAttribute]): Optional attribute to apply to the function (such as maximum shared memory size).

Returns:

DeviceExternalFunction: The loaded function is stored in the result parameter.

Raises:

If loading the function fails or the assembly code is invalid.

enqueue_function​

def enqueue_function[*Ts: AnyType](self, f: DeviceExternalFunction, *args: *Ts.values, *, grid_dim: Dim, block_dim: Dim, cluster_dim: OptionalReg[Dim] = None, shared_mem_bytes: OptionalReg[Int] = None, var attributes: List[LaunchAttribute] = List(__list_literal__=NoneType(None)), var constant_memory: List[ConstantMemoryMapping] = List(__list_literal__=NoneType(None)), location: OptionalReg[SourceLocation] = None)

Enqueues an external device function for execution on this device.

This overload accepts a DeviceExternalFunction that was loaded from assembly code (PTX/SASS). External functions are pre-compiled GPU kernels that can be integrated with Mojo code.

Example:

from max.gpu.host import DeviceContext

def vec_add_sig(
    in0: Pointer[Float32],
    in1: Pointer[Float32],
    out: Pointer[Float32],
    len: Int,
):
    pass

with DeviceContext() as ctx:
    var func = ctx.load_function[vec_add_sig](
        function_name="vectorAdd",
        asm=ptx_code,
    )
    ctx.enqueue_function(
        func,
        in0_buf,
        in1_buf,
        out_buf,
        1024,
        grid_dim=Dim(32),
        block_dim=Dim(32),
    )
    ctx.synchronize()

Parameters:

  • ​*Ts (AnyType): Argument types to pass to the external function.

Args:

  • ​f (DeviceExternalFunction): The external device function to execute.
  • ​*args (*Ts.values): Arguments to pass to the function.
  • ​grid_dim (Dim): Dimensions of the compute grid, made up of thread blocks.
  • ​block_dim (Dim): Dimensions of each thread block in the grid.
  • ​cluster_dim (OptionalReg[Dim]): Dimensions of clusters (if the thread blocks are grouped into clusters).
  • ​shared_mem_bytes (OptionalReg[Int]): Amount of shared memory per thread block.
  • ​attributes (List[LaunchAttribute]): Launch attributes.
  • ​constant_memory (List[ConstantMemoryMapping]): Constant memory mapping.
  • ​location (OptionalReg[SourceLocation]): Source location for the function call.

Raises:

If the operation fails.

def enqueue_function[FuncType: def() -> None, //, dump_asm: Variant[Bool, Path, StringSlice[ImmStaticOrigin], def() capturing thin -> Path] = False, dump_llvm: Variant[Bool, Path, StringSlice[ImmStaticOrigin], def() capturing thin -> Path] = False, _dump_sass: Variant[Bool, Path, StringSlice[ImmStaticOrigin], def() capturing thin -> Path] = False, _ptxas_info_verbose: Bool = False](self, func: FuncType, grid_dim: Dim, block_dim: Dim, cluster_dim: OptionalReg[Dim] = None, shared_mem_bytes: OptionalReg[Int] = None, var attributes: List[LaunchAttribute] = List(__list_literal__=NoneType(None)), var constant_memory: List[ConstantMemoryMapping] = List(__list_literal__=NoneType(None)), func_attribute: OptionalReg[FuncAttribute] = None, location: OptionalReg[SourceLocation] = None)

Compiles and enqueues a capturing kernel for execution on this device with type checking.

This overload is for kernels that capture variables from their enclosing scope. The capturing annotation on the signature function indicates that the kernel can access variables from the surrounding context. Like the non-capturing overload, both func and signature_func should typically be the same kernel function.

Most parameters are inferred automatically. This overload is selected when your kernel captures variables from its surrounding scope:

from std.gpu import DeviceContext, global_idx
from layout import TileTensor, row_major


def main() raises:
    with DeviceContext() as ctx:
        var scale_factor: Float32 = 2.0

        var data_buffer = ctx.enqueue_create_buffer[DType.float32](100)
        var data = TileTensor(data_buffer, row_major[100]())
        with data_buffer.map_to_host() as h:
            for i in range(data.num_elements()):
                h[i] = Float32(i)

        # This kernel captures 'scale_factor' from the enclosing scope
        def scale_kernel() {var}:
            var i = global_idx.x
            if i >= 100:
                return
            data[i] = data[i] * scale_factor

        ctx.enqueue_function(scale_kernel, grid_dim=1, block_dim=256)
        ctx.synchronize()
        with data_buffer.map_to_host() as h:
            for i in range(data.num_elements()):
                print(h[i])

Parameters:

Args:

Raises:

If the operation fails.

def enqueue_function[declared_arg_types: TypeList[declared_arg_types.values], //, func: def(*args: *declared_arg_types) capturing thin -> None, *actual_arg_types: DevicePassable, *, link_options: StringSlice[ImmStaticOrigin] = StringSlice(""), dump_asm: Variant[Bool, Path, StringSlice[ImmStaticOrigin], def() capturing thin -> Path] = False, dump_llvm: Variant[Bool, Path, StringSlice[ImmStaticOrigin], def() capturing thin -> Path] = False, _dump_sass: Variant[Bool, Path, StringSlice[ImmStaticOrigin], def() capturing thin -> Path] = False, _ptxas_info_verbose: Bool = False](self, *args: *actual_arg_types.values, *, grid_dim: Dim, block_dim: Dim, cluster_dim: OptionalReg[Dim] = None, shared_mem_bytes: OptionalReg[Int] = None, var attributes: List[LaunchAttribute] = List(__list_literal__=NoneType(None)), var constant_memory: List[ConstantMemoryMapping] = List(__list_literal__=NoneType(None)), func_attribute: OptionalReg[FuncAttribute] = None, location: OptionalReg[SourceLocation] = None)

Compiles and enqueues a kernel for execution on this device. This overload takes in a function that's capturing.

You can pass the function directly to enqueue_function without compiling it first:

from max.gpu.host import DeviceContext

def kernel():
    print("hello from the GPU")

with DeviceContext() as ctx:
    ctx.enqueue_function[kernel](grid_dim=1, block_dim=1)
    ctx.synchronize()

If you are reusing the same function and parameters multiple times, this incurs 50-500 nanoseconds of overhead per enqueue, so you can compile it first to remove the overhead:

from max.gpu.host import DeviceContext

def kernel():
    print("hello from the GPU")

with DeviceContext() as ctx:
    var compiled_func = ctx.compile_function[kernel]()
    ctx.enqueue_function(compiled_func, grid_dim=1, block_dim=1)
    ctx.enqueue_function(compiled_func, grid_dim=1, block_dim=1)
    ctx.synchronize()

Parameters:

Args:

Raises:

If the operation fails.

enqueue_cpu_function​

def enqueue_cpu_function[func: def() capturing thin -> None](self)

Enqueues a function for execution on CPU.

Parameters:

  • ​func (def() capturing thin -> None): The function to execute.

Raises:

If the operation fails. If self is not a CPU DeviceContext.

def enqueue_cpu_function[FuncType: def() -> None](self, func: FuncType)

Enqueues a function for execution on CPU.

Parameters:

  • ​FuncType (def() -> None): The function type.

Args:

  • ​func (FuncType): The function to execute.

Raises:

If the operation fails. If self is not a CPU DeviceContext.

enqueue_cpu_range​

def enqueue_cpu_range[FuncType: def(Int) -> None](self, func: FuncType, count: Int)

Enqueues a function to be executed in parallel over a 1D range.

The function is called as func(i) for each i in range(count).

Instances of the function are executed in parallel, but it is not guaranteed that all instances will execute simultaneously.

Parameters:

  • ​FuncType (def(Int) -> None): The type of function to execute.

Args:

  • ​func (FuncType): The function closure to execute.
  • ​count (Int): The number of parallel instances of the function to enqueue.

Raises:

If the operation fails. If self is not a CPU DeviceContext.

execution_time​

def execution_time[func: def(Self) raises capturing thin -> None](self, num_iters: Int) -> Int

Measures the execution time of a function that takes a DeviceContext parameter.

This method times the execution of a provided function that requires the DeviceContext as a parameter. It runs the function for the specified number of iterations and returns the total elapsed time in nanoseconds.

Example:

from max.gpu.host import DeviceContext

def gpu_operation(ctx: DeviceContext) raises -> None:
    # Perform some GPU operation using ctx
    pass

with DeviceContext() as ctx:
    # Measure execution time of a function that uses the context
    var time_ns = ctx.execution_time[gpu_operation](10)
    print("Execution time for 10 iterations:", time_ns, "ns")

Parameters:

  • ​func (def(Self) raises capturing thin -> None): A function that takes a DeviceContext parameter to execute and time.

Args:

  • ​num_iters (Int): The number of iterations to run the function.

Returns:

Int: The total elapsed time in nanoseconds for all iterations.

Raises:

If the timer operations fail or if the function raises an exception.

def execution_time[FuncType: def(DeviceContext) raises -> None](self, func: FuncType, num_iters: Int) -> Int

Measures the execution time of a function that takes a DeviceContext parameter.

Parameters:

  • ​FuncType (def(DeviceContext) raises -> None): The body function type.

Args:

  • ​func (FuncType): The closure carrying the captured state of the body function.
  • ​num_iters (Int): The number of iterations to run the function.

Returns:

Int: The total elapsed time in nanoseconds for all iterations.

Raises:

If the timer operations fail.

def execution_time[func: def() raises capturing thin -> None](self, num_iters: Int) -> Int

Measures the execution time of a function over multiple iterations.

This method times the execution of a provided function that doesn't require the DeviceContext as a parameter. It runs the function for the specified number of iterations and returns the total elapsed time in nanoseconds.

Example:

from max.gpu.host import DeviceContext

def some_gpu_operation() raises -> None:
    # Perform some GPU operation
    pass

with DeviceContext() as ctx:
    # Measure execution time of a function
    var time_ns = ctx.execution_time[some_gpu_operation](10)
    print("Execution time:", time_ns, "ns")

Parameters:

  • ​func (def() raises capturing thin -> None): A function with no parameters to execute and time.

Args:

  • ​num_iters (Int): The number of iterations to run the function.

Returns:

Int: The total elapsed time in nanoseconds for all iterations.

Raises:

If the timer operations fail or if the function raises an exception.

def execution_time[FuncType: def() raises -> None](self, func: FuncType, num_iters: Int) -> Int

Measures the execution time of a function over multiple iterations.

Parameters:

  • ​FuncType (def() raises -> None): The body function type.

Args:

  • ​func (FuncType): The closure carrying the captured state of the body function.
  • ​num_iters (Int): The number of iterations to run the function.

Returns:

Int: The total elapsed time in nanoseconds for all iterations.

Raises:

If the timer operations fail.

push_context​

def push_context(self) -> _DeviceContextScope

Returns a context manager that ensures this device's driver context is active.

This method returns a context manager that pushes this device's driver context as the current context on entry and restores the previous context on exit. This is useful for operations that require a specific GPU context to be active, such as cuDNN operations on multi-GPU systems.

Example:

from max.gpu.host import DeviceContext

var ctx = DeviceContext(device_id=1)
# Ensure GPU 1's context is active for these operations.
with ctx.push_context():
    # All GPU operations here will use GPU 1's context.
    ...  # call external stateful APIs, such as cudnn.
# Previous context is automatically restored

Returns:

_DeviceContextScope: A context manager that manages the driver context stack.

Raises:

If there's an error switching contexts.

set_as_current​

def set_as_current(self)

For use with libraries that require a specific GPU context to be active. Sets the current device to the one associated with this DeviceContext.

Example:

from max.gpu.host import DeviceContext
var ctx = DeviceContext(device_id=1)
ctx.set_as_current()

Raises:

If there's an error setting the current device.

execution_time_iter​

def execution_time_iter[func: def(Self, Int) raises capturing thin -> None](self, num_iters: Int) -> Int

Measures the execution time of a function that takes iteration index as input.

This method times the execution of a provided function that requires both the DeviceContext and the current iteration index as parameters. It runs the function for the specified number of iterations, passing the iteration index to each call, and returns the total elapsed time in nanoseconds.

Example:

from max.gpu.host import DeviceContext

def benchmark_kernel(ctx: DeviceContext, i: Int) raises -> None:
    # Perform GPU operations using ctx, potentially varying by iteration
    pass

with DeviceContext() as ctx:
    # Measure execution time with iteration awareness
    var time_ns = ctx.execution_time_iter[benchmark_kernel](10)
    print("Total execution time:", time_ns, "ns")

Parameters:

  • ​func (def(Self, Int) raises capturing thin -> None): A function that takes the DeviceContext and an iteration index.

Args:

  • ​num_iters (Int): The number of iterations to run the function.

Returns:

Int: The total elapsed time in nanoseconds for all iterations.

Raises:

If the timer operations fail or if the function raises an exception.

def execution_time_iter[FuncType: def(DeviceContext, Int) raises -> None](self, func: FuncType, num_iters: Int) -> Int

Measures the execution time of a function that takes iteration index as input.

Parameters:

  • ​FuncType (def(DeviceContext, Int) raises -> None): The body function type.

Args:

  • ​func (FuncType): The closure carrying the captured state of the body function.
  • ​num_iters (Int): The number of iterations to run the function.

Returns:

Int: The total elapsed time in nanoseconds for all iterations.

Raises:

If the timer operations fail.

enqueue_copy​

def enqueue_copy[dtype: DType](self, dst_buf: DeviceBuffer[dtype], src_ptr: Pointer[Scalar[dtype]])

Enqueues an async copy from the host to the provided device buffer. The number of bytes copied is determined by the size of the device buffer.

Parameters:

  • ​dtype (DType): Type of the data being copied.

Args:

Raises:

If the operation fails.

def enqueue_copy[dtype: DType](self, dst_buf: HostBuffer[dtype], src_ptr: Pointer[Scalar[dtype]])

Enqueues an async copy from the host to the provided device buffer. The number of bytes copied is determined by the size of the device buffer.

Parameters:

  • ​dtype (DType): Type of the data being copied.

Args:

Raises:

If the operation fails.

def enqueue_copy[dtype: DType](self, dst_ptr: Pointer[Scalar[dtype]], src_buf: DeviceBuffer[dtype])

Enqueues an async copy from the device to the host. The number of bytes copied is determined by the size of the device buffer.

Parameters:

  • ​dtype (DType): Type of the data being copied.

Args:

Raises:

If the operation fails.

def enqueue_copy[dtype: DType](self, dst_ptr: Pointer[Scalar[dtype]], src_buf: HostBuffer[dtype])

Enqueues an async copy from the device to the host. The number of bytes copied is determined by the size of the device buffer.

Parameters:

  • ​dtype (DType): Type of the data being copied.

Args:

Raises:

If the operation fails.

def enqueue_copy[dtype: DType](self, dst_ptr: Pointer[Scalar[dtype]], src_ptr: Pointer[Scalar[dtype]], size: Int)

Enqueues an async copy of size elements from a device pointer to another device pointer.

Parameters:

  • ​dtype (DType): Type of the data being copied.

Args:

Raises:

If the operation fails.

def enqueue_copy[dtype: DType](self, dst_buf: DeviceBuffer[dtype], src: Span[Scalar[dtype]])

Enqueues an async copy from a host Span to a device buffer.

The number of bytes copied is determined by the size of the device buffer. The span must contain at least as many elements as the destination buffer; this invariant is checked via debug_assert.

Parameters:

  • ​dtype (DType): Type of the data being copied.

Args:

Raises:

If the operation fails.

def enqueue_copy[dtype: DType](self, dst_buf: HostBuffer[dtype], src: Span[Scalar[dtype]])

Enqueues an async copy from a host Span to a host buffer.

The number of bytes copied is determined by the size of the destination buffer. The span must contain at least as many elements as the destination buffer; this invariant is checked via debug_assert.

Parameters:

  • ​dtype (DType): Type of the data being copied.

Args:

Raises:

If the operation fails.

def enqueue_copy[dtype: DType](self, dst: Span[Scalar[dtype]], src_buf: DeviceBuffer[dtype])

Enqueues an async copy from a device buffer to a host Span.

The number of bytes copied is determined by the size of the device buffer. The span must contain at least as many elements as the source buffer; this invariant is checked via debug_assert (debug builds only).

Parameters:

  • ​dtype (DType): Type of the data being copied.

Args:

Raises:

If the operation fails.

def enqueue_copy[dtype: DType](self, dst: Span[Scalar[dtype]], src_buf: HostBuffer[dtype])

Enqueues an async copy from a host buffer to a host Span.

The number of bytes copied is determined by the size of the source buffer. The span must contain at least as many elements as the source buffer; this invariant is checked via debug_assert (debug builds only).

Parameters:

  • ​dtype (DType): Type of the data being copied.

Args:

Raises:

If the operation fails.

def enqueue_copy[dtype: DType](self, dst_buf: DeviceBuffer[dtype], src_buf: DeviceBuffer[dtype])

Enqueues an async copy from one device buffer to another. The amount of data transferred is determined by the size of the destination buffer.

Parameters:

  • ​dtype (DType): Type of the data being copied.

Args:

Raises:

If the operation fails.

def enqueue_copy[dtype: DType](self, dst_buf: DeviceBuffer[dtype], src_buf: HostBuffer[dtype])

Enqueues an async copy from one device buffer to another. The amount of data transferred is determined by the size of the destination buffer.

Parameters:

  • ​dtype (DType): Type of the data being copied.

Args:

Raises:

If the operation fails.

def enqueue_copy[dtype: DType](self, dst_buf: HostBuffer[dtype], src_buf: DeviceBuffer[dtype])

Enqueues an async copy from one device buffer to another. The amount of data transferred is determined by the size of the destination buffer.

Parameters:

  • ​dtype (DType): Type of the data being copied.

Args:

Raises:

If the operation fails.

def enqueue_copy[dtype: DType](self, dst_buf: HostBuffer[dtype], src_buf: HostBuffer[dtype])

Enqueues an async copy from one device buffer to another. The amount of data transferred is determined by the size of the destination buffer.

Parameters:

  • ​dtype (DType): Type of the data being copied.

Args:

Raises:

If the operation fails.

enqueue_copy_no_cross_stream_sync​

def enqueue_copy_no_cross_stream_sync[dtype: DType](self, dst_buf: DeviceBuffer[dtype], src_buf: DeviceBuffer[dtype])

Enqueues a device-to-device copy without cross-stream synchronization.

This behaves like enqueue_copy for two device buffers, except that when the source and destination are on different streams the driver does not insert the events that normally synchronize them. The caller is responsible for ensuring the source data is ready before the copy and that the source buffer is not reused until the copy completes. This is used by the graph compiler, which emits explicit synchronization ops around the copy.

Parameters:

  • ​dtype (DType): Type of the data being copied.

Args:

Raises:

If the operation fails.

enqueue_memset​

def enqueue_memset[dtype: DType](self, dst: DeviceBuffer[dtype], val: Scalar[dtype])

Enqueues an async memset operation, setting all of the elements in the destination device buffer to the specified value.

Parameters:

  • ​dtype (DType): Type of the data stored in the buffer.

Args:

Raises:

If the operation fails.

def enqueue_memset[dtype: DType](self, dst: HostBuffer[dtype], val: Scalar[dtype])

Enqueues an async memset operation, setting all of the elements in the destination host buffer to the specified value.

Parameters:

  • ​dtype (DType): Type of the data stored in the buffer.

Args:

Raises:

If the operation fails.

create_event​

def create_event[*, blocking_sync: Bool = False, disable_timing: Bool = True, interprocess: Bool = False](self) -> DeviceEvent

Creates a new event for synchronization between streams.

Provides the best performance by default, disabling timing and blocking sync. DeviceContext.execution_time() provides the functionality required for timing kernels by passing it a closure, and is functionally equivalent to recording start and end events, then calculating the elapsed time.

Example:

from max.gpu.host import DeviceContext

var ctx = DeviceContext()

var default_stream = ctx.stream()
var new_stream = ctx.create_stream()

# Create an event
var event = ctx.create_event()

# Wait for the event in new_stream
new_stream.enqueue_wait_for(event)

# new_stream can continue
default_stream.record_event(event)
default_stream.synchronize()

Parameters:

  • ​blocking_sync (Bool): Enable event.synchronize() to block until the event has been recorded. Incurs overhead compared to stream.enqueue_wait_for(event) (default: False).
  • ​disable_timing (Bool): Remove timing overhead (default: True).
  • ​interprocess (Bool): Enable interprocess synchronization, currently unimplemented. (default: False).

Returns:

DeviceEvent: A DeviceEvent that can be used for synchronization.

Raises:

If event creation fails.

stream_priority_range​

def stream_priority_range(self) -> StreamPriorityRange

Returns the range of stream priorities supported by this device context.

Returns:

StreamPriorityRange: A StreamPriorityRange object containing the minimum and maximum stream priorities.

Raises:

If the operation fails.

create_stream​

def create_stream(self, *, priority: Int = Int(0)) -> DeviceStream

Creates a new stream associated with the given device context.

To create a stream with the highest priority, use:

from max.gpu.host import DeviceContext
var ctx = DeviceContext()
var priority = ctx.stream_priority_range().greatest
var stream = ctx.create_stream(priority=priority)

Args:

  • ​priority (Int): The priority of the stream (default: 0).

Returns:

DeviceStream: The newly created device stream with the specified priority.

Raises:

If stream creation fails.

create_external_stream​

def create_external_stream(self, external_stream: Optional[Pointer[NoneType, origin]]) -> DeviceStream

Creates a non-owning stream wrapper around an externally managed GPU stream.

The returned DeviceStream does not take ownership of the underlying stream. The caller is responsible for ensuring the external stream remains valid for the lifetime of the returned wrapper.

Args:

Returns:

DeviceStream: A DeviceStream wrapping the external stream without taking ownership of it.

Raises:

If wrapping the external stream fails.

synchronize​

def synchronize(self)

Blocks until all asynchronous calls on the stream associated with this device context have completed.

Raises:

If the operation fails. This should never be necessary when writing a custom operation.

enqueue_wait_for​

def enqueue_wait_for(self, other: Self)

Enqueues a wait operation for another device context to complete its work.

This method creates a dependency between two device contexts, ensuring that operations in the current context will not begin execution until all previously enqueued operations in the other context have completed. This is useful for synchronizing work across multiple devices or streams.

Example:

from max.gpu.host import DeviceContext

# Create two device contexts
var ctx1 = DeviceContext(0)  # First GPU
var ctx2 = DeviceContext(1)  # Second GPU

# Enqueue operations on ctx1
# ...

# Make ctx2 wait for ctx1 to complete before proceeding
ctx2.enqueue_wait_for(ctx1)

# Enqueue operations on ctx2 that depend on ctx1's completion
# ...

Args:

  • ​other (Self): The device context whose operations must complete before operations in this context can proceed.

Raises:

If there's an error enqueuing the wait operation or if the operation is not supported by the underlying device API.

num_streams​

def num_streams(self) -> Int

Returns the number of streams available on this device context.

Returns:

Int: The number of streams available on this device context.

select_stream​

def select_stream(self, stream_id: Int) -> Self

Returns a view of this device context bound to the given stream.

The returned context shares this context's full stream set, driver context, and device memory pool; only the current-stream selector differs, so work enqueued on it runs on stream stream_id. Stream 0 is the base/default stream. Backends without a multi-stream model return a view equivalent to this context.

Args:

  • ​stream_id (Int): Index of the stream the returned view submits to.

Returns:

Self: A device context view bound to stream stream_id.

Raises:

If the stream cannot be selected or created.

get_api_version​

def get_api_version(self) -> Int

Returns the API version associated with this device.

This method retrieves the version number of the GPU driver currently installed on the system for the device associated with this context. The version is returned as an integer that can be used to check compatibility with specific features or to troubleshoot driver-related issues.

Example:

from max.gpu.host import DeviceContext

with DeviceContext() as ctx:
    # Get the API version
    var api_version = ctx.get_api_version()
    print("GPU API version:", api_version)

Returns:

Int: An integer representing the driver version.

Raises:

If the driver version cannot be retrieved or if the device context is invalid.

get_attribute​

def get_attribute(self, attr: DeviceAttribute) -> Int

Returns the specified attribute for this device.

Use the aliases defined by DeviceAttribute to specify attributes. For example:

from max.gpu.host import DeviceAttribute, DeviceContext

def main() raises:
    var ctx = DeviceContext()
    var attr = DeviceAttribute.MAX_BLOCKS_PER_MULTIPROCESSOR
    var max_blocks = ctx.get_attribute(attr)
    print(max_blocks)

Args:

Returns:

Int: The value for attr on this device.

Raises:

If the operation fails.

is_compatible​

def is_compatible(self) -> Bool

Returns True if this device is compatible with MAX.

This method checks whether the current device is compatible with the Modular Accelerated Execution (MAX) runtime. It's useful for validating that the device can execute the compiled code before attempting operations.

Example:

from max.gpu.host import DeviceContext

var ctx = DeviceContext()
print("Device is compatible with MAX:", ctx.is_compatible())

Returns:

Bool: True if the device is compatible with MAX, False otherwise.

run_healthcheck​

def run_healthcheck(self)

Runs lightweight GPU health validation.

Checks for hardware throttling, uncorrectable ECC errors, and stuck VRAM. Raises an error if the GPU is unhealthy. The healthcheck runs automatically during device initialization; this method allows re-running it explicitly.

Disable with MODULAR_DEVICE_CONTEXT_DISABLE_HEALTHCHECK=true.

Raises:

Error: If the GPU is in an unhealthy state.

id​

def id(self) -> Int64

Returns the ID associated with this device.

This method retrieves the unique identifier for the current device. Device IDs are used to distinguish between multiple devices in a system and are often needed for multi-GPU programming.

Example:

from max.gpu.host import DeviceContext

var ctx = DeviceContext()
try:
    var device_id = ctx.id()
    print("Using device with ID:", device_id)
except:
    print("Failed to get device ID")

Returns:

Int64: The unique device ID as an Int64.

Raises:

If there's an error retrieving the device ID.

get_memory_info​

def get_memory_info(self) -> Tuple[UInt, UInt]

Returns the free and total memory size for this device.

This method queries the current state of device memory, providing information about how much memory is available and the total memory capacity of the device. This is useful for memory management and determining if there's enough space for planned operations.

Example:

from max.gpu.host import DeviceContext

var ctx = DeviceContext()
try:
    (free, total) = ctx.get_memory_info()
    print("Free memory:", free / (1024*1024), "MB")
    print("Total memory:", total / (1024*1024), "MB")
except:
    print("Failed to get memory information")

Returns:

Tuple[UInt, UInt]: A tuple of (free memory, total memory) in bytes.

Raises:

If there's an error retrieving the memory information.

max_single_alloc_size​

def max_single_alloc_size(self) -> c_size_t

Returns the largest single contiguous allocation, in bytes.

On Metal this is maxBufferLength; on other backends it is the device's total memory.

Returns:

c_size_t: The maximum size, in bytes, of a single contiguous allocation supported by this device.

Raises:

If the underlying device query fails.

can_access​

def can_access(self, peer: Self) -> Bool

Returns True if this device can access the identified peer device.

This method checks whether the current device can directly access memory on the specified peer device. Peer-to-peer access allows for direct memory transfers between devices without going through host memory, which can significantly improve performance in multi-GPU scenarios.

Example:

from max.gpu.host import DeviceContext
var ctx1 = DeviceContext(0)  # First GPU
var ctx2 = DeviceContext(1)  # Second GPU

try:
    if ctx1.can_access(ctx2):
        print("Direct peer access is possible")
        ctx1.enable_peer_access(ctx2)
    else:
        print("Direct peer access is not supported")
except:
    print("Failed to check peer access capability")

Args:

  • ​peer (Self): The peer device to check for accessibility.

Returns:

Bool: True if the current device can access the peer device, False otherwise.

Raises:

If there's an error checking peer access capability.

enable_peer_access​

def enable_peer_access(self, peer: Self)

Enables direct memory access to the peer device.

This method establishes peer-to-peer access from the current device to the specified peer device. Once enabled, the current device can directly read from and write to memory allocated on the peer device without going through host memory, which can significantly improve performance for multi-GPU operations.

Notes:

  • It's recommended to call can_access() first to check if peer access is possible.
  • Peer access is not always symmetric; you may need to enable access in both directions.

Example:

from max.gpu.host import DeviceContext

var ctx1 = DeviceContext(0)  # First GPU
var ctx2 = DeviceContext(1)  # Second GPU

try:
    if ctx1.can_access(ctx2):
        ctx1.enable_peer_access(ctx2)
        print("Peer access enabled from device 0 to device 1")

        # For bidirectional access
        if ctx2.can_access(ctx1):
            ctx2.enable_peer_access(ctx1)
            print("Peer access enabled from device 1 to device 0")
    else:
        print("Peer access not supported between these devices")
except:
    print("Failed to enable peer access")

Args:

  • ​peer (Self): The peer device to enable access to.

Raises:

If there's an error enabling peer access or if peer access is not supported between the devices.

supports_multicast​

def supports_multicast(self) -> Bool

Returns True if this device supports multicast memory mappings.

Returns:

Bool: True if the current device supports multicast memory, False otherwise.

Raises:

If there's an error checking peer access capability.

number_of_devices​

static def number_of_devices(*, var api: String = DeviceContext.default_device_info.api) -> Int

Returns the number of devices available that support the specified API.

This function queries the system for available devices that support the requested API (such as CUDA or HIP). It's useful for determining how many accelerators are available before allocating resources or distributing work.

Example:

from max.gpu.host import DeviceContext

# Get number of CUDA devices
var num_cuda_devices = DeviceContext.number_of_devices(api="cuda")

# Get number of devices for the default API
var num_devices = DeviceContext.number_of_devices()

Args:

  • ​api (String): Requested device API (for example, "cuda" or "hip"). Defaults to the device API specified by current target accelerator.

Returns:

Int: The number of available devices supporting the specified API.

enable_all_peer_access​

static def enable_all_peer_access()

Enable peer-to-peer memory access between all available accelerators.

This function detects all available accelerators in the system and enables peer-to-peer (P2P) memory access between every pair of devices.

When peer access is enabled, kernels running on one device can directly access memory allocated on another device without going through host memory. This is crucial for efficient multi-GPU operations like allreduce.

The function is a no-op when:

  • No accelerators are available
  • Only one accelerator is available
  • Peer access is already enabled between devices

Example:

from max.gpu.host import DeviceContext

# Enable P2P access between all GPUs
DeviceContext.enable_all_peer_access()

# Now GPUs can directly access each other's memory

Raises:

If peer access cannot be enabled between any pair of devices. This can happen if the hardware doesn't support P2P access or if there's a configuration issue.

all_peer_access_enabled​

static def all_peer_access_enabled() -> Bool

Check whether peer-to-peer memory access is enabled between all GPU pairs.

This function queries whether P2P access has been successfully enabled between all pairs of GPUs in the system. It returns True only if every GPU can directly access every other GPU's memory.

Example:

from max.gpu.host import DeviceContext

# P2P access is automatically enabled when devices are constructed.
# Check if it was successful for all pairs.
if DeviceContext.all_peer_access_enabled():
    print("P2P access enabled between all GPUs")
else:
    print("P2P access not available for all GPU pairs")

Returns:

Bool: True if P2P access is enabled between all GPU pairs, False otherwise. Returns False if there are fewer than 2 GPUs or if P2P is not supported between any pair.

Raises:

If there's an error querying the P2P access status.

Models

Cookbook

Puzzles

Blog

Community

Contact