GqaWorkspace.ixx File Reference

The owning GQA transient scratch one attention stack shares, and its factory. More...

#include <cstddef>
#include <memory>
#include <string>
import Compute.DeviceTypeTraits;
import Compute.DeviceId;
import Compute.GqaState;
import Dnn.TensorDataTypeTraits;
import Compute.DeviceType;
import Compute.DeviceAllocation;
import Dnn.TensorDataType;
import Dnn.TensorTypes;
import Dnn.Tensor;

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

Namespaces
namespace	Mila
	Mila main API namespace.

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.

Detailed Description

The owning GQA transient scratch one attention stack shares, and its factory.

GqaState is the non-owning view of it that the operation takes. Family-neutral, so a transformer and a layer-streamed harness allocate the identical set.

Function Documentation

gqaWorkspaceDeviceBytes()

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 )

export

Device bytes makeGqaWorkspace would allocate with the same arguments, without allocating.

Parameters

granularity	The allocation granularity of the device the workspace would live on.
B,num_heads,head_dim,T_ctx,prefill_chunk,score_width	As for makeGqaWorkspace.

makeGqaWorkspace()

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 )

export

Allocate the shared GQA transient.

Parameters

device	The device every workspace tensor is allocated on.
B	Batch size.
num_heads	Query heads.
head_dim	The widest head any layer of the stack uses; narrower layers view a prefix.
T_ctx	Context length; sizes the decode score buffers.
prefill_chunk	Tokens per prefill chunk; sizes the prefill buffers.
score_width	The caller's flash decision made concrete: the flash path reads no score buffer, while the cuBLASLt path needs the full context and would overflow a narrow one. Taken explicitly so a caller cannot allocate for one path and then run the other.
name_prefix	Prepended to each tensor's name.