Llama.Block.ixx File Reference

LLaMA transformer block – module partition of LlamaTransformer. More...

#include <memory>
#include <vector>
#include <string>
#include <format>
#include <iostream>
#include <sstream>
#include <stdexcept>
#include <optional>
import Serialization.Mode;
import Serialization.ModelArchive;
import Dnn.Components.Linear;
import Dnn.Components.Residual;
import Dnn.Components.Gqa;
import Dnn.Quantization.Weight.Policies;
import Dnn.Components.RmsNorm;
import Compute.GqaState;
import Compute.ExecutionContextFactory;
import Compute.IExecutionContext;
import Compute.ExecutionContext;
import Compute.DeviceTypeTraits;
import Compute.DeviceType;
import Dnn.Components.LlamaTransformer:Config;
import Dnn.CompositeComponent;
import Dnn.Tensor;
import Compute.Device;
import Dnn.Components.Swiglu;
import Dnn.TensorDataType;
import Dnn.TensorDataTypeTraits;
import Dnn.Quantization.KvCache.Policy;
import Dnn.TensorTypes;
import Dnn.TensorOps;
import Dnn.Component;
import Dnn.ITensor;
import Dnn.ComponentType;
import Dnn.Components.Rope;
import Compute.DeviceId;

Classes
struct	Mila::Dnn::LlamaBlock< TDeviceType, TPrecision, TWeightQuant, TKvPolicy >::BlockBuildContexts
	The per-child build contexts this block implies. More...
class	Mila::Dnn::LlamaBlock< TDeviceType, TPrecision, TWeightQuant, TKvPolicy >

Namespaces
namespace	Mila
	Mila main API namespace.

Detailed Description

LLaMA transformer block – module partition of LlamaTransformer.

Implements the correct Llama 3.x attention sub-graph using fused projections, zero-copy tensor views, and in-place RoPE rotation:

input [B, T, model_dim] +- RMSNorm (ln_1) +- fc_qkv_proj [model_dim -> (n_heads + 2*n_kv_heads) * head_dim] 1 GEMM +- view Q [B, T, n_heads * head_dim] –+ +- view K [B, T, n_kv_heads * head_dim] –+ RoPE in-place +- view V [B, T, n_kv_heads * head_dim] | (V untouched) +- GroupedQueryAttention (packed QKV) +- fc_out_proj [model_dim -> model_dim] +- Residual (input + out_proj) res_1 +- RMSNorm (ln_2) +- fc_gate_up [model_dim -> 2*hidden_dim] 1 GEMM +- SwiGLU -> [B, T, hidden_dim] +- fc_down [hidden_dim -> model_dim] +- Residual (res1 + ffn) res_2

Key design points:

• Single GEMM for QKV (GQA-correct output dim).
• RoPE applied in-place via zero-copy views – no concat/split ops needed.
• Single GEMM for gate+up (SwiGLU kernel expects [gate | up] layout).
• FFN composed directly from Linear + SwiGLU primitives; no MLP composite.
• Component names match convert_llama32.py tensor name mapping.

Weight loader note: HuggingFace stores fc_gate and fc_up as separate tensors. The weight loader must concatenate them along dim 0 into a single [2*hidden_dim, model_dim] matrix when loading into fc_gate_up.