Module precision_schedule 

Per-layer + per-tensor precision overrides (C.3, 2026-04-25).

Reference: NVIDIA Transformer Engine 2.14 + EAQuant (arXiv:2506.13329)

• community 2025 mixed-precision recipes. In MoE models, the quantization-sensitivity hierarchy holds across re-tested benchmarks:

  1. Router (gate weights): hidden × num_experts, tiny in memory, but routing accuracy collapses fast under quant. Keep BF16 wherever feasible.
  2. LM head: hidden × vocab, large but determines output logit fidelity. BF16 closes the dominant chunk of perplexity gap.
  3. First 1-2 transformer blocks + last 2-3 blocks: the embedding-adjacent layers carry sink-token outliers and the output-adjacent layers shape final logits. Keep at FP8 (one tier above the bulk).
  4. Bulk MoE experts: NVFP4 / FP8 — the model has the most slack here.

Scope

This module ships:

• Role — semantic tag for each tensor the loader wants to classify (router, lm_head, attention, expert, etc.).
• Dtype — target precision values the schedule emits.
• PrecisionSchedule — the per-(layer, role) → dtype decision table, built from [precision] in MODEL.toml.

The loader consults schedule.dtype_for(layer_idx, role) at tensor-load time and chooses the matching path. When the schedule is in its default() state (no [precision] block in MODEL.toml), every lookup returns Dtype::Inherit — meaning “use whatever the existing per-checkpoint logic decides.” This keeps the pre-2026-04-25 behaviour bit-exact.

Structs

PrecisionSchedule

Per-(layer, role) precision schedule built from MODEL.toml’s [precision] block. Lookups are O(1) for the common tables; per-layer overrides hit a small sorted set.

Enums

Dtype

Target precision for a tensor. Inherit means “let the existing per-checkpoint logic decide” (preserves pre-C.3 behaviour); the other variants are hard requests.

Role

Semantic role of a tensor, used for precision lookups. The set is closed and minimal — adding a new role requires extending the Dtype::for_role match.