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 trait

DeviceGraphInput

A device graph input that contributes to the graph's cache key.

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, Movable

Required methods

write_graph_key

def write_graph_key(self, mut writer: T)

Writes this input's contribution to a graph's cache key.

Write whatever distinguishes graphs that must not be shared: two inputs that write the same bytes are treated as interchangeable. Contributions are separated from each other by DeviceGraphCache.make_key(), so there is no need to delimit or tag the output for that purpose.

Args:

• ​writer (T): The writer accumulating the cache key.

allocate_stable

def allocate_stable(self, mut builder: DeviceGraphBuilder) -> Self

Returns a graph-owned value of this type for the graph to record.

Allocate through DeviceGraphBuilder.create_input_buffer(), which reserves the address for the graph's lifetime and registers it, and describe the result with the same shape and dtype as self.

Read only self's shape, unless the input's data is host-resident, in which case also copy self's contents into the stable allocation before returning. Recorded ops read host values at enqueue time, so the stable location must already hold the caller's bytes while the build closure records.

A mutable input the graph writes in place (a KV cache, for example) must instead call DeviceGraphBuilder.register_in_place_input() and return a value aliasing self: a private stable copy would silently discard the graph's writes. Such an implementation must include the input's address in write_graph_key, so a moved buffer misses the cache and forces a rebuild instead of replaying stale addresses.

Args:

• ​builder (DeviceGraphBuilder): The builder for the graph under construction.

Returns:

_Self: A value of the same type backed by the graph's memory pool, or — for a mutable in-place input — a value aliasing self.

Raises:

If allocating from the graph's memory pool fails.