Gemma.Block.Workspace.ixx File Reference

Transformer-owned shared activation slots for GemmaBlock (pooling), and their factory. More...

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

Classes
struct	Mila::Dnn::GemmaBlockWorkspace< TDeviceType, TPrecision >
	Transformer-owned shared activation workspace for GemmaBlock (pooling). More...
struct	Mila::Dnn::GemmaBlockWorkspaceWidths
	Max-geometry slot widths, shared by the allocation and the transformer's footprint. More...

Namespaces
namespace	Mila
	Mila main API namespace.

Functions
GemmaBlockWorkspaceWidths	Mila::Dnn::gemmaBlockWorkspaceWidths (const GemmaConfig &config)
template<DeviceType TDeviceType, TensorDataType TPrecision> requires PrecisionSupportedOnDevice<TPrecision, TDeviceType>
GemmaBlockWorkspace< TDeviceType, TPrecision >	Mila::Dnn::makeGemmaBlockWorkspace (const GemmaConfig &config, DeviceId device, dim_t B, dim_t prefill_chunk, const std::string &name_prefix)
	Allocate the shared block workspace at the widths gemmaBlockWorkspaceWidths() reports.

Detailed Description

Transformer-owned shared activation slots for GemmaBlock (pooling), and their factory.

One slot set serves every layer; the transformer owns it and accounts for it.

Function Documentation

makeGemmaBlockWorkspace()

template<DeviceType TDeviceType, TensorDataType TPrecision>
requires PrecisionSupportedOnDevice<TPrecision, TDeviceType>

GemmaBlockWorkspace< TDeviceType, TPrecision > Mila::Dnn::makeGemmaBlockWorkspace ( const GemmaConfig & config, DeviceId device, dim_t B, dim_t prefill_chunk, const std::string & name_prefix )

export

Allocate the shared block workspace at the widths gemmaBlockWorkspaceWidths() reports.

Lives beside the struct because the transformer is not the only caller: a layer-streamed parity harness holds one block at a time and must measure the geometry the model builds, not a copy of it.