cubie.memory.stream_groups

Stream group management for organizing instances and CUDA streams.

This module provides functionality to group instances together and assign them to shared CUDA streams for coordinated memory operations and kernel execution.

Classes

StreamGroups([groups, streams])

Container for organizing instances into groups with shared CUDA streams.

class cubie.memory.stream_groups.StreamGroups(groups: dict[str, list[int]] | None = NOTHING, streams: dict[str, FakeStream | int] = NOTHING)

Bases: object

Container for organizing instances into groups with shared CUDA streams.

Parameters:

• groups (dict of str to list of int, optional) – Dictionary mapping group names to lists of instance IDs.

• streams (dict of str to Stream or int, optional) – Dictionary mapping group names to CUDA streams.

groups

Dictionary mapping group names to lists of instance IDs.

Type:

dict of str to list of int

streams

Dictionary mapping group names to CUDA streams.

Type:

dict of str to Stream or int

Notes

Each group has an associated CUDA stream that all instances in the group share for coordinated operations. The ‘default’ group is created automatically.

add_instance(instance, group)

Add an instance to a stream group.

Parameters:

• instance (object or int) – The instance to add to the group, or its ID.

• group (str) – Name of the group to add the instance to.

Raises:

ValueError – If the instance is already in a stream group.

Notes

If the group doesn’t exist, it will be created with a new CUDA stream.

change_group(instance, new_group)

Move an instance to another stream group.

Parameters:

• instance (object or int) – The instance to move, or its ID.

• new_group (str) – Name of the group to move the instance to.

Notes

If the new group doesn’t exist, it will be created with a new CUDA stream.

get_group(instance)

Get the stream group associated with an instance.

Parameters:

instance (object or int) – The instance to find the group for, or its ID.

Returns:

Name of the group containing the instance.

Return type:

str

Raises:

ValueError – If the instance is not in any stream groups.

get_instances_in_group(group)

Get all instances in a stream group.

Parameters:

group (str) – Name of the group to get instances for.

Returns:

List of instance IDs in the group, or empty list if group doesn’t exist.

Return type:

list of int

get_stream(instance)

Get the CUDA stream associated with an instance.

Parameters:

instance (object or int) – The instance to get the stream for, or its ID.

Returns:

CUDA stream associated with the instance’s group.

Return type:

Stream or int

groups: dict[str, list[int]] | None

reinit_streams()

Reinitialize all streams after a context reset.

Notes

Called after CUDA context reset to create fresh streams for all groups.

streams: dict[str, FakeStream | int]