Mila
Deep Neural Network Library
Loading...
Searching...
No Matches
LanguageModelConfig.ixx File Reference

CRTP base configuration for all deployable Mila language models. More...

#include <format>
#include <stdexcept>
#include <string>
#include <string_view>
import Dnn.TensorTypes;

Classes

struct  Mila::Dnn::LanguageModelConfig< TDerived >
 CRTP base configuration for all deployable Mila language models. More...

Namespaces

namespace  Mila
 Mila main API namespace.

Enumerations

enum class  Mila::Dnn::KvCacheCompression { None , FP8 }
 KV cache storage and compression strategy for GroupedQueryAttention. More...
enum class  Mila::Dnn::WeightQuantization { None , FP8 , FP4 , Plan }
 Weight storage and matmul strategy for Linear components. More...

Functions

bool Mila::Dnn::isDerivableFromReferenceWeights (WeightQuantization quantization)
 True when a load can derive this format from reference weights.
void Mila::Dnn::requireStoredQuantizationMatches (std::string_view caller, std::string_view weights_path, std::string_view stored_quantization, WeightQuantization requested_quantization, int fp4_group_size=128)
 Refuse a load whose stored weights were packed for a different policy.
std::string Mila::Dnn::weightQuantizationName (WeightQuantization quantization, int fp4_group_size=128)
 The scheme name recorded in a model's weights and in its manifest.

Detailed Description

CRTP base configuration for all deployable Mila language models.

LanguageModelConfig<TDerived> owns the deployment concerns that are universal across all language model architectures:

  1. context_length – maximum sequence length the model is built for. RoPE embeddings and KV cache buffers are sized to this.
  2. WeightQuantization – weight storage and matmul strategy for Linear components. Defaults to WeightQuantization::None (BF16 weights).
  3. KvCacheCompression – KV cache storage and compression strategy for GroupedQueryAttention components. Defaults to KvCacheCompression::None (no compression).

CRTP Pattern

All fluent setters return TDerived& so that chains work correctly across both base and derived methods without casting at the call site:

QwenModelConfig config = QwenModelConfig( context_length )
.withFP8Quantization() // returns QwenModelConfig&
.withThinkingMode(); // returns QwenModelConfig&

Relationship to ModelConfig

ModelConfig<TDevice, TPrecision> is the structural base for all Mila models. LanguageModelConfig is the deployment configuration counterpart for the language model branch of that hierarchy. Vision model configurations would derive from a sibling VisionModelConfig<TDerived>, not from this class.

Relationship to BuildContext

LanguageModelConfig is the public API surface for deployment configuration. BuildContext is the internal carrier through the component tree. load() projects LanguageModelConfig into BuildContext once – they are never the same object.

Quantization Presets vs Fine-Grained Control

Convenience preset methods express common deployment decisions in user vocabulary. Fine-grained setters are available for atypical configurations:

// Preset -- FP8 weights + FP8 KV cache
LlamaModelConfig config = LlamaModelConfig( context_length )
.withFP8Quantization();
// Fine-grained -- FP4 weights, no KV compression
LlamaModelConfig config = LlamaModelConfig( context_length )
.withWeightQuantization( WeightQuantization::FP4 )
.withKvCacheCompression( KvCacheCompression::None );

Enumeration Type Documentation

◆ KvCacheCompression

enum class Mila::Dnn::KvCacheCompression
exportstrong

KV cache storage and compression strategy for GroupedQueryAttention.

Maps to the TKvPolicy template parameter on GroupedQueryAttention and CudaGqaOp via the load() runtime->compile-time bridge. The mapping is:

None -> NoKvCompression (BF16 cache, no compression overhead) FP8 -> PerChannelKvFp8<> (FP8_E4M3 cache, per-head per-token float32 scales)

New compression algorithms (SlidingWindow, LowRank, TurboQuant) add a value here and a corresponding policy struct in KvCache.QuantPolicy – no other changes are required at this level.

Enumerator
None 

No compression – default; BF16 KV cache.

FP8 

FP8_E4M3 per-head per-token KV cache compression – Alpha.6 target.

◆ WeightQuantization

enum class Mila::Dnn::WeightQuantization
exportstrong

Weight storage and matmul strategy for Linear components.

Maps to the TWeightQuant template parameter on Linear and CudaLinearOp via the load() runtime->compile-time bridge. The mapping is:

None -> NoWeightQuant (BF16 weights, standard cuBLASLt plan) FP8 -> PerChannelFp8<> (FP8_E4M3 weights, per-channel float32 scales) FP4 -> PerGroupFp4<128> (FP4_E2M1 weights, per-group float32 scales), or <64> for a model whose projection widths are not multiples of 128

This enum is Mila API vocabulary. Callers set it via fluent methods on the concrete model config – they do not interact with the policy structs directly.

Enumerator
None 

BF16 weights – default; no quantization overhead.

FP8 

FP8_E4M3 per-channel weight quantization – Alpha.5 target.

FP4 

Per-group FP4 weight quantization – future target.

Plan 

The family's own designed per-role allocation, rather than one uniform format.

The three values above name a STORAGE FORMAT that applies to every Linear alike. This one does not name a format at all: it says "build this model the way its designers allocated its bits", and which formats that means is the family's to define – Qwen 3.8 spends 2.5 bits on the feed-forward gate/up pair and 4.125 on full attention (Specifications/Qwen3.8.md section 5).

A family with no plan must REFUSE this value rather than fall back to a uniform policy, which is why dispatchWeightQuantization handles it explicitly.

Function Documentation

◆ isDerivableFromReferenceWeights()

bool Mila::Dnn::isDerivableFromReferenceWeights ( WeightQuantization quantization)
inlineexport

True when a load can derive this format from reference weights.

The distinction the artifact check turns on. FP4 and FP8 are computed from the weights at load time – absmax scales and a format-defined level table – so a BF16 artifact is a legitimate source for them, and every family already relies on that: Qwen's own packed artifact carries codebook tensors only and quantizes its attention and head projections on load (Qwen3.8.md section 8).

A plan's codebooks are the opposite case. They are FITTED offline against calibration data, so nothing in a BF16 tensor recovers them and the artifact must carry them.

Refusing both alike would be the safe-looking answer and the wrong one: it would make a uniform FP4 build of any family unreachable without a repack that adds nothing, which is exactly what the Phase 5 FP4 oracle needs to load.

◆ requireStoredQuantizationMatches()

void Mila::Dnn::requireStoredQuantizationMatches ( std::string_view caller,
std::string_view weights_path,
std::string_view stored_quantization,
WeightQuantization requested_quantization,
int fp4_group_size = 128 )
inlineexport

Refuse a load whose stored weights were packed for a different policy.

Nothing downstream can tell the two apart: packed codes reinterpreted as BF16, or a BF16 blob decoded through a codebook, produce a model that loads and runs and is wrong. The storage dtype cannot stand in for this check – FP4 at group 128 and at group 64 are both U8. An empty stored name means reference weights; the reader normalizes the writer's "none" to empty, so the two spellings compare as one.

The one asymmetry is deliberate: reference weights ARE a valid source for a derivable format, because the load computes those scales from the weights. See isDerivableFromReferenceWeights for why a codebook is the opposite case.

It lives here rather than in a family because the rule is the writer's and the reader's alike, and a family that omits it does not fail – it runs wrong.

Parameters
callerQualified name of the calling factory, for the message.
weights_pathPath to the weights being loaded.
stored_quantizationScheme the weights declare; empty for reference weights.
requested_quantizationPolicy this build was compiled for.
fp4_group_sizeFP4 group this build was compiled for; FP4 at another group is refused.

◆ weightQuantizationName()

std::string Mila::Dnn::weightQuantizationName ( WeightQuantization quantization,
int fp4_group_size = 128 )
inlineexport

The scheme name recorded in a model's weights and in its manifest.

It lives beside the enum because it is written by the model that saves the weights and read by the tool that packages them, and those two must agree exactly: the load side refuses weights whose scheme disagrees with the build's compile-time policy, since the bytes are packed differently per scheme and reinterpreting them produces a model that runs and is wrong. It was previously spelled out in both places.

Parameters
quantizationThe scheme to name.
fp4_group_sizeThe FP4 group the model's geometry requires; ignored for other schemes.