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

DeviceStream

struct DeviceStream

Represents a CUDA/HIP stream for asynchronous GPU operations.

A DeviceStream provides a queue for GPU operations that can execute concurrently with operations in other streams. Operations within a single stream execute in the order they are issued, but operations in different streams may execute in any relative order or concurrently.

This abstraction allows for better utilization of GPU resources by enabling overlapping of computation and data transfers.

Example:

from max.gpu.host import DeviceContext

var ctx = DeviceContext(0)  # Select first GPU
var stream = ctx.create_stream()

# Launch operations on the stream
# ...

# Wait for all operations in the stream to complete
stream.synchronize()

Implemented traits​

AnyType, Copyable, Deinitable, ImplicitlyCopyable, Movable, _FunctionEnqueuer

Methods​

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 stream.

Forwards directly to AsyncRT_DeviceStream_enqueueFunctionDirect, scheduling the kernel on the underlying CUDA/HIP stream. 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.

synchronize​

def synchronize(self)

Blocks the calling CPU thread until all operations in this stream complete.

This function waits until all previously issued commands in this stream have completed execution. It provides a synchronization point between host and device code.

Example:

from max.gpu.host import DeviceContext

var ctx = DeviceContext()
var stream = ctx.create_stream()

# Launch kernel or memory operations on the stream
# ...

# Wait for completion
stream.synchronize()

# Now it's safe to use results on the host

Raises:

If synchronization fails.

enqueue_wait_for​

def enqueue_wait_for(self, event: DeviceEvent)

Makes this stream wait for the specified event.

This function inserts a wait operation into this stream that will block all subsequent operations in the stream until the specified event has been recorded and completed.

Args:

Raises:

If the wait operation fails.

record_event​

def record_event(self, event: DeviceEvent)

Records an event in this stream.

This function records the given event at the current point in this stream. All operations in the stream that were enqueued before this call will complete before the event is triggered.

Example:

from max.gpu.host import DeviceContext

var ctx = DeviceContext()

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

# Create event on the context
var event = ctx.create_event()

# Wait for the event on the new stream
new_stream.enqueue_wait_for(event)

# Stream 2 can continue
default_stream.record_event(event)

Args:

Raises:

If event recording fails.

enqueue_host_func​

def enqueue_host_func(self, func: def(Pointer[NoneType, MutAnyOrigin]) thin -> None, user_data: Pointer[NoneType, MutAnyOrigin])

Enqueues a host callback to run on this stream.

This corresponds to CUDA's cuLaunchHostFunc. The callback func runs on a driver thread once all preceding work on this stream has completed, and receives user_data as its only argument. Per the CUDA contract, the callback must not call any device APIs.

Currently only implemented for CUDA streams; other backends raise.

Args:

  • ​func (def(Pointer[NoneType, MutAnyOrigin]) thin -> None): A thin C-compatible function pointer that accepts a single void* argument.
  • ​user_data (Pointer[NoneType, MutAnyOrigin]): An opaque pointer passed through to func when it runs.

Raises:

If the underlying device does not support host callbacks, or if the driver rejects the enqueue.

wait_for_host_value​

def wait_for_host_value(self, flag: CompletionFlag, value: UInt64)

Stalls this stream until flag's 64-bit value equals value.

Corresponds to CUDA's cuStreamWaitValue64. The stream blocks at this node until the 64-bit slot owned by flag (allocated in device-mapped pinned host memory by its owning C++ DeviceContext) holds value. A CPU thread, or the AsyncRT worker dispatched by enqueue_host_func, calling the C++ producer-side signal(value) lets the GPU stream synchronize on CPU-produced data without a second stream or a blocking host-function callback on the consumer's critical path.

Captures cleanly into a CUDA graph as a wait-value (batch-mem-op) node, so this operation can be placed between graph-captured kernels to gate a downstream consumer on CPU-produced data.

Currently only implemented for CUDA streams; other backends raise.

Args:

  • ​flag (CompletionFlag): A non-owning handle to a M::Driver::CompletionFlag allocated by the same device's C++ context.
  • ​value (UInt64): The 64-bit value to wait for (equality).

Raises:

If the underlying device does not support stream memory ops, or if the driver rejects the enqueue.

enqueue_function​

def enqueue_function[*Ts: DevicePassable](self, f: DeviceFunction[target=f.target, compile_options=f.compile_options, link_options=f.link_options, _ptxas_info_verbose=f._ptxas_info_verbose], *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)))

Enqueues a checked compiled function for execution on this stream.

Parameters:

Args:

Raises:

If the operation fails.

Models

Cookbook

Puzzles

Blog

Community

Contact