Mila::Dnn::BuildContext Class Reference

Build-time context for Component::build(). More...

Public Member Functions
	BuildContext ()
	Default constructor – sentinel value for pre-build state.
	BuildContext (shape_t input_shape, RuntimeMode runtime_mode, std::optional< bool > initialize_parameters=std::nullopt)
	Construct from all six concerns explicitly.
int64_t	getPrefillSize () const noexcept
	Number of tokens processed per prefill pass.
RuntimeMode	getRuntimeMode () const noexcept
	The runtime mode governing output buffer allocation.
bool	hasInstalledOutput () const noexcept
	True if the caller will install this component's output buffer.
const shape_t &	inputShape () const noexcept
	The full input shape this component receives.
bool	isInferenceMode () const noexcept
	True if output buffers should be allocated at T=1.
bool	isTrainingMode () const noexcept
	True if output buffers should be allocated at full input shape sequence length with gradient buffers.
bool	shouldInitializeParameters () const noexcept
	True if components should initialize parameter values after allocation.
BuildContext	withInstalledOutput (bool installed) const
	Declare that the caller will install this component's output buffer.
BuildContext	withPrefillSize (int64_t prefill_size) const
	Return a copy of this context with a different prefill size.
BuildContext	withShape (shape_t new_shape) const
	Return a copy of this context with a different input shape.

Detailed Description

Build-time context for Component::build().

Carries six orthogonal concerns down the Component hierarchy:

1. Input shape – the full input shape the component receives. Used for parameter sizing, output buffer allocation, and build-time validation against component config.

2. RuntimeMode – allocation policy governing output buffer sizing and gradient buffer allocation.

   Inference – T=1 decode path output buffers. Training – full sequence output buffers, gradient buffers allocated.

3. Parameter initialization – whether components should initialize parameter tensors after allocation. Set to false when building for a pretrained weight load to avoid computing initializers (Xavier, normal, zeros) that are immediately overwritten by loadParameter(). When not specified, the default is derived from RuntimeMode: Training initializes (train from scratch), Inference skips (weights are loaded). An inference-mode build therefore cannot silently run then discard parameter initialization by omitting the flag.

Caller responsibility

The Network or Transformer constructing BuildContext is responsible for providing the correct full input shape for each child component. Each component validates this shape against its own config in onBuilding().

Threading

Not synchronized. Used only during the single-threaded build phase.

Constructor & Destructor Documentation

BuildContext() [1/2]

Mila::Dnn::BuildContext::BuildContext ( )

inline

Default constructor – sentinel value for pre-build state.

Produces a minimal valid BuildContext with parameter initialization enabled, Auto precision policy, and no quantization. Never read before build() is called – Component::ensureBuilt() guards all access paths.

BuildContext() [2/2]

Mila::Dnn::BuildContext::BuildContext ( shape_t input_shape, RuntimeMode runtime_mode, std::optional< bool > initialize_parameters = std::nullopt )

inlineexplicit

Construct from all six concerns explicitly.

precision_policy and quantization are extracted from ModelConfig by fromPretrained() and passed here as raw values, keeping BuildContext free of any model-layer dependency.

Parameters

input_shape	Complete input shape this component receives. Must have at least one dimension.
runtime_mode	Allocation policy: Inference or Training.
initialize_parameters	When false, components allocate parameter tensors but skip value initialization. When omitted (nullopt), the default is derived from runtime_mode – Training initializes, Inference skips – so a load path cannot regress by forgetting the flag.

Exceptions

std::invalid_argument

if input_shape is empty.

Member Function Documentation

getPrefillSize()

int64_t Mila::Dnn::BuildContext::getPrefillSize ( ) const

inlinenoexcept

Number of tokens processed per prefill pass.

The tuned prefill chunk size, computed once at network build time and threaded down to every component that sizes prefill buffers or attention scratch. Zero on training-mode contexts (no chunking).

withInstalledOutput()

BuildContext Mila::Dnn::BuildContext::withInstalledOutput ( bool installed ) const

inlinenodiscard

Declare that the caller will install this component's output buffer.

A composite that pools activations installs a shared slot into each child before calling build(), so the child skips self-allocating its output. That decision is invisible to getRequiredMemory(), which runs before any installation has happened and would otherwise count a buffer the build never allocates – once in the child and again in the pooling parent.

Only prediction reads this; onBuilding() continues to use the component's own installed flag, which by then is accurate.

withPrefillSize()

BuildContext Mila::Dnn::BuildContext::withPrefillSize ( int64_t prefill_size ) const

inlinenodiscard

Return a copy of this context with a different prefill size.

All other fields are preserved. Used by the network to stamp the tuned prefill chunk size onto the contexts it builds its child components with.

Parameters

prefill_size

Tokens per prefill pass.

Returns

New BuildContext with prefill_size and all other fields unchanged.

withShape()

BuildContext Mila::Dnn::BuildContext::withShape ( shape_t new_shape ) const

inlinenodiscard

Return a copy of this context with a different input shape.

All other fields are preserved – RuntimeMode, prefill_size, and initialize_parameters.

Parameters

new_shape

Replacement input shape. Must be non-empty.

Returns

New BuildContext with new_shape and all other fields unchanged.

---

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

• Mila/Src/Dnn/Core/Component.BuildContext.ixx