CUDA execution context specialization. More...

Inheritance diagram for Mila::Dnn::Compute::CudaExecutionContext:

Public Member Functions
	CudaExecutionContext (const CudaExecutionContext &)=delete
	CudaExecutionContext (CudaExecutionContext &&)=delete
	CudaExecutionContext (DeviceId device_id)
	Constructs CUDA execution context for a specific device.
	~CudaExecutionContext ()
	Destructor with proper CUDA resource cleanup.
cublasLtHandle_t	getCublasLtHandle () const
	Gets a cuBLASLt handle for this context.
void *	getCublasLtWorkspace () const noexcept
	Gets the cuBLASLt workspace buffer.
size_t	getCublasLtWorkspaceSize () const noexcept
	Gets the cuBLASLt workspace buffer size in bytes.
curandGenerator_t	getCurandGenerator () const
	Gets the cuRAND generator bound to this context's stream.
DeviceId	getDeviceId () const noexcept override
	Gets the device identifier.
void *	getDeviceScratchBuffer (size_t required_bytes) const
	Gets the device scratch buffer shared by operations during forward passes.
void *	getLoadStagingBuffer (size_t required_bytes) const
	Gets or grows the device buffer a load fits full-precision weights through.
void *	getPinnedStagingBuffer (size_t required_bytes) const
	Gets or grows the pinned host staging buffer for Host->Device transfers.
cudaStream_t	getStream () const noexcept
	Gets the CUDA stream for asynchronous operations.
CudaExecutionContext &	operator= (const CudaExecutionContext &)=delete
CudaExecutionContext &	operator= (CudaExecutionContext &&)=delete
void	releaseLoadStaging () noexcept override
	Free the device buffer a load staged full-precision weights through.
void	reserveScratch (std::size_t bytes) override
	Allocate the forward scratch buffer at exactly the size a network reports.
void	synchronize () override
	Synchronizes the CUDA stream.
Public Member Functions inherited from Mila::Dnn::Compute::IExecutionContext
const ActivationObserver &	getActivationObserver () const noexcept
bool	hasActivationObserver () const noexcept
void	setActivationObserver (ActivationObserver observer)
	Install the observer components publish their activations to.

Static Public Member Functions
static constexpr size_t	getLoadStagingLimitBytes () noexcept
	Largest staging request one fit of full-precision weights may make, in bytes.

Detailed Description

CUDA execution context specialization.

Manages CUDA execution resources including streams and library handles. Each context owns an independent CUDA stream and owns the library handles used by callers bound to this context's stream.

Thread Safety:

Multiple contexts can safely exist on the same device.
Each context has its own stream for isolated execution.
Library handles are owned by the context and must not be used concurrently from multiple threads without external synchronization.

Constructor & Destructor Documentation

◆ CudaExecutionContext()

Mila::Dnn::Compute::CudaExecutionContext::CudaExecutionContext ( DeviceId device_id )

inlineexplicit

Constructs CUDA execution context for a specific device.

Creates an independent CUDA stream.

Parameters

device_id CUDA device identifier.

Exceptions

std::invalid_argument	If device_id type is not Cuda.
std::runtime_error	If CUDA stream creation fails.

Member Function Documentation

◆ getCublasLtHandle()

cublasLtHandle_t Mila::Dnn::Compute::CudaExecutionContext::getCublasLtHandle ( ) const

inlinenodiscard

Gets a cuBLASLt handle for this context.

Created lazily and owned by this context.

Returns: cublasLtHandle_t The cuBLASLt handle.

Exceptions

std::runtime_error If handle creation fails.

◆ getCublasLtWorkspace()

void * Mila::Dnn::Compute::CudaExecutionContext::getCublasLtWorkspace ( ) const

inlinenodiscardnoexcept

Gets the cuBLASLt workspace buffer.

Shared scratch memory for cuBLASLt algorithm execution. Allocated once at context creation, valid for the lifetime of this context. Size is kCublasLtWorkspaceSize bytes.

Returns: void* Pointer to device workspace buffer.

◆ getCublasLtWorkspaceSize()

size_t Mila::Dnn::Compute::CudaExecutionContext::getCublasLtWorkspaceSize ( ) const

inlinenodiscardnoexcept

Gets the cuBLASLt workspace buffer size in bytes.

Returns: size_t Workspace size in bytes.

◆ getCurandGenerator()

curandGenerator_t Mila::Dnn::Compute::CudaExecutionContext::getCurandGenerator ( ) const

inlinenodiscard

Gets the cuRAND generator bound to this context's stream.

Created lazily, seeded from Core::RandomGenerator, and bound to this context's stream. Reuse across calls avoids per-call creation overhead.

Returns: curandGenerator_t The cuRAND generator bound to this stream.

Exceptions

std::runtime_error If generator creation or seeding fails.

◆ getDeviceId()

DeviceId Mila::Dnn::Compute::CudaExecutionContext::getDeviceId ( ) const

inlinenodiscardoverridevirtualnoexcept

Gets the device identifier.

Returns: DeviceId The CUDA device identifier (type + index).

Implements Mila::Dnn::Compute::IExecutionContext.

◆ getDeviceScratchBuffer()

void * Mila::Dnn::Compute::CudaExecutionContext::getDeviceScratchBuffer ( size_t required_bytes ) const

inlinenodiscard

Gets the device scratch buffer shared by operations during forward passes.

Used by operations that need a temporary device buffer during forward passes (e.g. FP8->BF16 weight dequantization before a cuBLASLt GEMM). Because all operations on this context share a single stream, the buffer can be safely reused across sequential ops – each op finishes before the next one writes to the buffer.

Once a network has reserved it, a request larger than the reservation throws: the reservation is the footprint's prediction, and exceeding it is a defect in that prediction. Before any reservation the buffer grows on demand. Freed in releaseResources().

Parameters

required_bytes Minimum number of bytes required.

Returns: void* Device buffer of at least required_bytes.

Exceptions

std::logic_error	If the request exceeds a reservation.
std::runtime_error	If allocation fails.

◆ getLoadStagingBuffer()

void * Mila::Dnn::Compute::CudaExecutionContext::getLoadStagingBuffer ( size_t required_bytes ) const

inlinenodiscard

Gets or grows the device buffer a load fits full-precision weights through.

Separate from the forward scratch: it lives only while a load runs, and the load frees it through releaseLoadStaging() when it returns. Grown on demand in between; the returned pointer is invalidated by a later call with a larger request.

Parameters

required_bytes Minimum number of bytes required.

Returns: void* Device buffer of at least required_bytes.

Exceptions

std::runtime_error If allocation fails.

◆ getLoadStagingLimitBytes()

constexpr size_t Mila::Dnn::Compute::CudaExecutionContext::getLoadStagingLimitBytes ( )

inlinestaticnodiscardconstexprnoexcept

Largest staging request one fit of full-precision weights may make, in bytes.

A tensor larger than this is fitted in row blocks that fit, which costs more copies and no more memory.

◆ getPinnedStagingBuffer()

void * Mila::Dnn::Compute::CudaExecutionContext::getPinnedStagingBuffer ( size_t required_bytes ) const

inlinenodiscard

Gets or grows the pinned host staging buffer for Host->Device transfers.

Page-locked via cudaHostAlloc so cudaMemcpyAsync from it to device memory is a direct DMA transfer. Grown on demand, never shrunk. Freed in releaseResources().

Parameters

required_bytes Minimum number of bytes required.

Returns: void* Pinned host buffer of at least required_bytes.

Exceptions

std::runtime_error If allocation fails.

◆ getStream()

cudaStream_t Mila::Dnn::Compute::CudaExecutionContext::getStream ( ) const

inlinenodiscardnoexcept

Gets the CUDA stream for asynchronous operations.

Returns: cudaStream_t The CUDA stream owned by this context.

◆ releaseLoadStaging()

void Mila::Dnn::Compute::CudaExecutionContext::releaseLoadStaging ( )

inlineoverridevirtualnoexcept

Free the device buffer a load staged full-precision weights through.

A load that fits full-precision weights on the device stages each tensor through one buffer, which is needed only until the last tensor has landed. The model load calls this when it returns, so the buffer is not held for the model's lifetime. Safe to call when nothing was staged; a no-op for contexts that stage nothing.

Reimplemented from Mila::Dnn::Compute::IExecutionContext.

◆ reserveScratch()

void Mila::Dnn::Compute::CudaExecutionContext::reserveScratch ( std::size_t bytes )

inlineoverridevirtual

Allocate the forward scratch buffer at exactly the size a network reports.

Parameters

bytes The largest scratch request any operation on this context makes.

Exceptions

std::runtime_error If allocation fails.

Reimplemented from Mila::Dnn::Compute::IExecutionContext.

◆ synchronize()

void Mila::Dnn::Compute::CudaExecutionContext::synchronize ( )

inlineoverridevirtual

Synchronizes the CUDA stream.

Blocks until all operations submitted to this context's stream complete.

Exceptions

std::runtime_error If stream synchronization fails.

Implements Mila::Dnn::Compute::IExecutionContext.

The documentation for this class was generated from the following file:

Mila/Src/Dnn/Compute/Devices/Cuda/CudaExecutionContext.ixx

Public Member Functions

Static Public Member Functions

Detailed Description

Constructor & Destructor Documentation

◆ CudaExecutionContext()

Member Function Documentation

◆ getCublasLtHandle()

◆ getCublasLtWorkspace()

◆ getCublasLtWorkspaceSize()

◆ getCurandGenerator()

◆ getDeviceId()

◆ getDeviceScratchBuffer()

◆ getLoadStagingBuffer()

◆ getLoadStagingLimitBytes()

◆ getPinnedStagingBuffer()

◆ getStream()

◆ releaseLoadStaging()

◆ reserveScratch()

◆ synchronize()