Compute.GqaWorkspace Module Reference

Classes
struct	Mila::Dnn::Compute::GqaWorkspace< TDeviceType, TPrecision >
	The GQA transient the attention layers of one stack share, owned as one unit. More...

Functions
template<TensorDataType TPrecision>
std::size_t	Mila::Dnn::Compute::gqaWorkspaceDeviceBytes (std::size_t granularity, dim_t B, dim_t num_heads, dim_t head_dim, dim_t T_ctx, dim_t prefill_chunk, dim_t score_width)
	Device bytes makeGqaWorkspace would allocate with the same arguments, without allocating.
template<DeviceType TDeviceType, TensorDataType TPrecision> requires PrecisionSupportedOnDevice<TPrecision, TDeviceType>
GqaWorkspace< TDeviceType, TPrecision >	Mila::Dnn::Compute::makeGqaWorkspace (DeviceId device, dim_t B, dim_t num_heads, dim_t head_dim, dim_t T_ctx, dim_t prefill_chunk, dim_t score_width, const std::string &name_prefix)
	Allocate the shared GQA transient.

Files
file	Mila/Src/Dnn/Compute/Operations/GqaWorkspace.ixx
	The owning GQA transient scratch one attention stack shares, and its factory.