Mila::Dnn::Compute::GqaWorkspace< TDeviceType, TPrecision > Struct Template Reference

The GQA transient the attention layers of one stack share, owned as one unit. More...

Public Types
using	MR = typename DeviceTypeTraits<TDeviceType>::memory_resource
using	TensorType = Tensor<TPrecision, MR>

Public Member Functions
std::size_t	deviceStorageBytes () const
GqaState	state () const

Public Attributes
std::unique_ptr< TensorType >	att
std::unique_ptr< TensorType >	att_decode
std::unique_ptr< TensorType >	preatt
std::unique_ptr< TensorType >	preatt_decode
std::unique_ptr< TensorType >	q_permute
std::unique_ptr< TensorType >	v_out
std::unique_ptr< TensorType >	v_out_decode

Detailed Description

template<DeviceType TDeviceType, TensorDataType TPrecision>
requires PrecisionSupportedOnDevice<TPrecision, TDeviceType>
struct Mila::Dnn::Compute::GqaWorkspace< TDeviceType, TPrecision >

The GQA transient the attention layers of one stack share, owned as one unit.

Owned together rather than as seven loose tensors so a caller cannot allocate half of it: GqaState is a struct of raw pointers, and a null one is a crash at the kernel rather than an error at the call.

---

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

• Mila/Src/Dnn/Compute/Operations/GqaWorkspace.ixx