Mila::Dnn::GemmaBlockWorkspace< TDeviceType, TPrecision > Struct Template Reference

Transformer-owned shared activation workspace for GemmaBlock (pooling). More...

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

Public Attributes
std::shared_ptr< TensorType >	attn
std::shared_ptr< TensorType >	ffn_act
std::shared_ptr< TensorType >	ffn_down
std::shared_ptr< TensorType >	ffn_in
std::shared_ptr< TensorType >	ffn_normed
std::shared_ptr< TensorType >	gate_up
std::shared_ptr< TensorType >	k
std::shared_ptr< TensorType >	k_normed
std::shared_ptr< TensorType >	normed
std::shared_ptr< TensorType >	o
std::shared_ptr< TensorType >	o_normed
std::shared_ptr< TensorType >	q
std::shared_ptr< TensorType >	q_normed
std::shared_ptr< TensorType >	qkv
std::shared_ptr< TensorType >	res1
std::shared_ptr< TensorType >	stream
std::shared_ptr< TensorType >	v
std::shared_ptr< TensorType >	v_normed

Detailed Description

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

Transformer-owned shared activation workspace for GemmaBlock (pooling).

One slot per block-graph position, shared by every layer: the inference path is strictly sequential, so exactly one block is live at a time and 47/48 of per-layer retained activations are never read again. Slots are sized [B, chunk, max(local, global) width]; components view prefixes (the GQA workspace max-geometry convention). The single stream slot is alias-safe: a block's input is last read at res_1 (mid-block) and only overwritten by its own res_2 at block end.

---

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

• Mila/Src/Dnn/Components/Transformers/Gemma/Gemma.Block.ixx