cubie.memory

GPU memory management subsystem for cubie.

This module provides GPU memory management capabilities including: - CuPy memory pool integration for efficient GPU memory allocation - Stream group management for asynchronous CUDA operations - Array request/response system for structured memory allocation - Manual or automatic allocation of VRAM to different processes - Automatic chunking for large allocations that exceed available memory

The main components are:

• MemoryManager: Singleton interface for managing all memory operations

• ArrayRequest: Specification for array allocation requests

• ArrayResponse: Results of array allocation operations

• StreamGroups: Management of CUDA stream groups for coordination

• current_cupy_stream(): Context manager for CuPy stream integration

• CuPyAsyncNumbaManager: Async CuPy memory pool integration

• CuPySyncNumbaManager: Sync CuPy memory pool integration

The default memory manager instance is available as default_memmgr.

class cubie.memory.CuPyAsyncNumbaManager(context)

Bases: CuPyNumbaManager

Numba EMM plugin using CuPy MemoryAsyncPool for allocation and freeing.

Parameters:

context (numba.cuda.cudadrv.driver.Context) – CUDA context for memory management.

Notes

Uses CuPy’s asynchronous memory pool which provides stream-ordered memory operations.

initialize()

Initialize the async memory pool.

memalloc(nbytes)

Allocate memory from the async CuPy pool.

Parameters:

nbytes (int) – Number of bytes to allocate.

Returns:

Numba memory pointer wrapping the CuPy allocation.

Return type:

MemoryPointer

class cubie.memory.CuPySyncNumbaManager(context)

Bases: CuPyNumbaManager

Numba EMM plugin using CuPy MemoryPool for allocation and freeing.

Parameters:

context (numba.cuda.cudadrv.driver.Context) – CUDA context for memory management.

Notes

Uses CuPy’s synchronous memory pool which provides standard memory operations.

initialize()

Initialize the sync memory pool.

memalloc(nbytes)

Allocate memory from the sync CuPy pool.

Parameters:

nbytes (int) – Number of bytes to allocate.

Returns:

Numba memory pointer wrapping the CuPy allocation.

Return type:

MemoryPointer

class cubie.memory.current_cupy_stream(nb_stream)

Bases: object

Context manager to override CuPy’s current stream with a Numba stream.

Parameters:

nb_stream (numba.cuda.cudadrv.driver.Stream) – Numba stream to use as CuPy’s current stream.

nb_stream

The Numba stream being used.

Type:

numba.cuda.cudadrv.driver.Stream

cupy_ext_stream

CuPy external stream wrapper around the Numba stream.

Type:

cupy.cuda.ExternalStream or None

Notes

This context manager only has effect when the current CUDA memory manager is a CuPy-based manager. Otherwise, it acts as a no-op context manager.

Modules

array_requests	Array request and response data structures for memory management.
cupy_emm	CuPy async/sync memory pool External Memory Manager plugin for Numba.
mem_manager	Comprehensive GPU memory management system for cubie.
stream_groups	Stream group management for organizing instances and CUDA streams.