TokenEmbedding.ixx File Reference

Device-templated TokenEmbedding component. More...

#include <memory>
#include <vector>
#include <string>
#include <sstream>
#include <iostream>
#include <stdexcept>
#include <optional>
#include <cmath>
#include <algorithm>
#include <numeric>
#include <format>
#include <type_traits>
import Logging.Logger;
import Dnn.TensorHelpers;
import Serialization.Mode;
import Serialization.ModelArchive;
import Compute.Observation;
import Compute.CpuMemoryResource;
import Compute.MemoryResource;
import Compute.ExecutionContextFactory;
import Compute.DeviceTypeTraits;
import Compute.DeviceType;
import Serialization.SafeTensors;
import Compute.ExecutionContext;
import Dnn.TensorDataType;
import Dnn.ComponentType;
import Dnn.TensorOps;
import Compute.Device;
import Dnn.Tensor;
import Compute.DeviceAllocation;
import Serialization.Tensor;
import Dnn.Component;
import Dnn.TensorTypes;
import Serialization.Metadata;
import Compute.OperationTraits;
import Dnn.ITensor;
import Dnn.TensorDataTypeTraits;
import Dnn.Quantization.Weight.Policies;
import Dnn.Components.TokenEmbeddingConfig;
import Compute.DeviceId;

Classes
class	Mila::Dnn::TokenEmbedding< TDeviceType, TIndex, TPrecision, TTableQuantization, TTableResidency >
	Pure token embedding component (device-templated). More...

Namespaces
namespace	Mila
	Mila main API namespace.

Enumerations
enum class	Mila::Dnn::EmbeddingTableResidency { Device , Host }
	Where the embedding table's bytes live. More...

Detailed Description

Device-templated TokenEmbedding component.

Pure vocabulary lookup: maps token indices [B, T] to dense vectors [B, T, C]. Owns the wte parameter and its gradient. No positional encoding – that is handled downstream by a dedicated encoding component (RoPE, ALiBi, or Learned).

Derived from Lpe with all wpe / IPositionalDecode / decode() concerns removed.

Enumeration Type Documentation

EmbeddingTableResidency

enum class Mila::Dnn::EmbeddingTableResidency

exportstrong

Where the embedding table's bytes live.

Device is what every other parameter in the tree does and stays the default. Host places the table in pinned host memory, which is device-addressable under unified virtual addressing, so the gather kernel reads rows across PCIe and is unchanged – the table simply costs no VRAM. It is worth doing only for a table large enough to matter and only because a lookup is a gather and never a multiply: the traffic is one row per token, not one pass per matrix.

See Specifications/Qwen3.8.md section 5, which spends this on 1.271 B parameters (2.37 GiB) to bring a 27B model inside a 12 GiB card.