OperationTraits.Template.ixx File Reference

Unified compile-time dispatch template mapping (OperationType, DeviceType, TPrecision, TPolicy) to a concrete operation type. More...

import Dnn.TensorDataType;
import Compute.OperationType;
import Compute.DeviceType;

Namespaces
namespace	Mila
	Mila main API namespace.
namespace	Mila::Dnn
namespace	Mila::Dnn::Compute

Concepts
concept	Mila::Dnn::Compute::GqaOpConcept
	Contract for GroupedQueryAttentionOp: positional forward and backward.
concept	Mila::Dnn::Compute::LinearOpConcept
	Contract for LinearOp: typed forward matmul and backward weight/input gradients.
concept	Mila::Dnn::Compute::SamplingOpConcept
	Contract for SamplingOp: in-place token sampling from a logits tensor.
concept	Mila::Dnn::Compute::UnaryOpConcept
	Contract for policy-free unary ops (Softmax, RmsNorm, LayerNorm, Residual, ...).

Detailed Description

Unified compile-time dispatch template mapping (OperationType, DeviceType, TPrecision, TPolicy) to a concrete operation type.

Replaces all per-component *OpTypeMap templates with a single keyed dispatch table. Backend specializations live in OperationTraits:Cuda and OperationTraits:Cpu partition modules co-located with their concrete op implementations.

A missing specialization for any combination is a hard compile error – the correct diagnostic for an unsupported (OperationType, DeviceType, precision, policy) tuple.

Policy conventions: TPolicy = void policy-free ops (Softmax, RmsNorm, RoPE, Residual, ...) TPolicy = WeightQuantPolicy LinearOp (NoWeightQuant, PerChannelFp8<>, ...) TPolicy = KvCachePolicy GQA (NoKvCompression, PerChannelKvFp8<>, ...)

Components hold the concrete op type via a local alias:

using OpType = OperationTraits<OperationType::LinearOp, TDeviceType, TPrecision, TWeightQuant>::type;