nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/dtype/__init__.py

@@ -0,0 +1 @@
+"""Dtypes: the one table of widths (`dtypes.yml`) and its reader (`itemsize`)."""

nbxtensor/dtype/convert.py

@@ -0,0 +1,107 @@
+"""Host conversions between bfloat16, float16 and float32 — the one module that does them.
+
+numpy has no bfloat16. A bfloat16 array on the host is carried as its 16-bit patterns (`uint16`), or as the 2-byte
+void array a bfloat16 tensor's host copy comes back as (`<V2`, the array interface's typestr for a type numpy does
+not know). Every conversion between that carrier and real numbers lives here, once: rounding float32 to bfloat16,
+widening bfloat16 to float32, and narrowing bfloat16 to float16.
+
+Pure functions of numpy arrays: no device, no global state, no thread pool of their own (R-NBXT 11). A caller that
+splits a large array across workers calls them on each span with `out=`.
+"""
+from __future__ import annotations
+
+from typing import Optional
+
+import numpy as np
+
+__all__ = ["f32_to_bf16_bits", "bf16_bits_to_f32", "bf16_bits_to_f16", "bf16_bits", "bf16_abs_max", "saturate_scalar"]
+
+# The quiet NaN every NaN becomes, sign dropped: PyTorch's `round_to_nearest_even`
+# (torch/headeronly/util/BFloat16.h: `if (std::isnan(src)) return UINT16_C(0x7FC0);`).
+_QUIET_NAN = np.uint16(0x7FC0)
+# The elements one step of the rounding works on: its temporaries are one step's, never the operand's size (a 5.1 GB
+# operand took four full-size uint32 temporaries and the Mac's certifier to its 14 GB guard, 2026-09-28).
+_STEP = 1 << 18
+
+
+def bf16_bits(a: np.ndarray) -> np.ndarray:
+    """The `uint16` patterns of a bfloat16 host carrier: the array itself when it already is `uint16`, a view when
+    it is the 2-byte void carrier. Anything else is refused by name."""
+    a = np.asarray(a)
+    if a.dtype == np.uint16:
+        return a
+    if a.dtype.kind == "V" and a.dtype.itemsize == 2:
+        return np.ascontiguousarray(a).view(np.uint16)
+    raise TypeError(f"a bfloat16 host carrier is uint16 bits or a 2-byte void array, not {a.dtype}")
+
+
+def f32_to_bf16_bits(a: np.ndarray, out: Optional[np.ndarray] = None) -> np.ndarray:
+    """float32 values narrowed to bfloat16, returned as their `uint16` patterns: round to nearest, ties to even —
+    the rounding `torch.Tensor.to(torch.bfloat16)` applies and the device cast applies. `+0x7FFF` alone rounds an
+    exact half down; `((u >> 16) & 1)` adds the one that makes it ties-to-even. Values past the bfloat16 maximum
+    become infinities; every NaN becomes the quiet NaN 0x7FC0. Streamed in fixed steps: the peak is the output plus
+    one step."""
+    a = np.ascontiguousarray(a, dtype=np.float32)
+    flat = a.reshape(-1)
+    u = flat.view(np.uint32)
+    bits = np.empty(flat.shape, dtype=np.uint16) if out is None else out.reshape(-1)
+    for i in range(0, flat.size, _STEP):
+        uc = u[i:i + _STEP]
+        t = np.right_shift(uc, np.uint32(16))
+        np.bitwise_and(t, np.uint32(1), out=t)
+        np.add(t, np.uint32(0x7FFF), out=t)
+        np.add(t, uc, out=t)
+        np.right_shift(t, np.uint32(16), out=t)
+        bits[i:i + _STEP] = t
+        nan = np.isnan(flat[i:i + _STEP])
+        if nan.any():
+            bits[i:i + _STEP][nan] = _QUIET_NAN
+    return bits.reshape(a.shape) if out is None else out
+
+
+def bf16_bits_to_f32(bits: np.ndarray) -> np.ndarray:
+    """The exact float32 values of bfloat16 patterns (`uint16`, or the 2-byte void carrier): a bfloat16 is the top
+    sixteen bits of a float32, so the widening loses nothing."""
+    b = np.ascontiguousarray(bf16_bits(bits))
+    return (b.astype(np.uint32) << np.uint32(16)).view(np.float32).reshape(b.shape)
+
+
+def bf16_bits_to_f16(bits: np.ndarray, out: Optional[np.ndarray] = None) -> np.ndarray:
+    """bfloat16 patterns narrowed to float16, exactly as widening to float32 then numpy's float32 -> float16 cast
+    (round to nearest even; past the float16 range: infinity). `out`, when given, is a float16 array of the same
+    number of elements the result is written into (a caller splitting the work across workers passes spans)."""
+    src = np.ascontiguousarray(bf16_bits(bits)).reshape(-1)
+    dst = np.empty(src.shape, dtype=np.float16) if out is None else out.reshape(-1)
+    np.copyto(dst, (src.astype(np.uint32) << np.uint32(16)).view(np.float32), casting="unsafe")
+    return dst.reshape(np.shape(bits)) if out is None else out
+
+
+def bf16_abs_max(bits: np.ndarray) -> float:
+    """The largest magnitude among bfloat16 patterns, NaN ignored (an infinity is infinite): what a caller compares to
+    a narrower dtype's `finite_max` to know whether the values fit it. Streamed in fixed steps."""
+    flat = np.ascontiguousarray(bf16_bits(bits)).reshape(-1)
+    top = 0.0
+    for i in range(0, flat.size, _STEP):
+        v = np.abs((flat[i:i + _STEP].astype(np.uint32) << np.uint32(16)).view(np.float32))
+        v = v[~np.isnan(v)]
+        if v.size:
+            top = max(top, float(v.max()))
+    return top
+
+
+def saturate_scalar(value, dtype):
+    """A Python number clamped to the finite range of a half floating `dtype` (a name or an enum member; the range is
+    the dtype table's `max`), anything else unchanged: a bool, a non-number, a dtype wider than two bytes or not
+    floating, a NaN, an infinity the graph wrote on purpose. A mask sentinel written for one half format then holds in
+    the other (bfloat16's -3.39e38 becomes float16's -65504)."""
+    from nbxtensor.dtype.itemsize import canonical_name, finite_max, itemsize, kind
+    if isinstance(value, bool) or not isinstance(value, (int, float)):
+        return value
+    name = canonical_name(dtype)
+    if kind(name) != "float" or itemsize(name) != 2:
+        return value
+    v = float(value)
+    if v != v or v in (float("inf"), float("-inf")):
+        return value
+    top = finite_max(name)
+    return max(-top, min(top, v))

nbxtensor/dtype/dtypes.yml

@@ -0,0 +1,70 @@
+# The dtype table: the width of every dtype the engine knows, one entry each.
+#
+# Read through ONE reader and nowhere else: `nbxtensor/dtype/itemsize.py` (both engine branches, Prism, NBXTensor, the
+# wrappers, the launch keys, the certifier). A dtype absent from this file is refused BY NAME: no width is ever guessed
+# (the silent `.get(dtype, 4)` sized a complex64 at half its bytes until 2026-09-22, Kokoro-82M's vocoder).
+#
+# Per entry:
+#   bits           storage bits per element (bytes = bits / 8 for every dense dtype)
+#   typestr        the array-interface type string of the element (`__cuda_array_interface__`, numpy's
+#                  `__array_interface__`); bfloat16, which numpy cannot name, is the 2-byte void '<V2'
+#   numpy          the numpy dtype name, when numpy has one
+#   aliases        other spellings a caller may name it by (torch's `half`, `float`, `double`, `int`,
+#                  `long`, `byte`)
+#   kind           float, int, uint, bool or complex — the class the promotion rule reads
+#                  (nbxtensor.dtype.promotion)
+#   max            a float's largest finite value: (2 - 2^(1-p)) x 2^emax for an IEEE-style format of p
+#                  significant bits (binary64, binary32, binary16; bfloat16 = binary32's exponent with
+#                  8 significant bits; TF32 = 11 significant bits), and the OCP OFP8 v1.0 table for the
+#                  two float8 formats (E4M3: 448, no infinity; E5M2: 57344)
+#   safetensors    the dtype code a safetensors header writes for it (spec:
+#                  https://github.com/huggingface/safetensors — `Dtype`), when it has one
+#   autotune_keys  the spellings a launch key / certified directory carries for it
+#                  (`kernels/autotune_certified._DTYPES`)
+#   without_fp64   what a store with no float64 holds in its place (the device's or the engine's
+#                  fp64 ceiling: MPS refuses a float64 tensor and the vendors' own code narrows to
+#                  float32 there; the Triton kernels are fp32-max, `TritonSequence._parse_dtype`)
+#
+# A hi+lo split operand is TWO tensors of the low dtype: that is a representation, computed by
+# the readers' `representation_bytes(dtype, numel, split=True)`, never an entry here.
+
+dtypes:
+  # IEEE 754-2008 binary64.
+  float64: {bits: 64, kind: float, max: 1.7976931348623157e+308, safetensors: F64, autotune_keys: [fp64], without_fp64: float32, typestr: '<f8', numpy: float64, aliases: [double]}
+  # IEEE 754-2008 binary32.
+  float32: {bits: 32, kind: float, max: 3.4028234663852886e+38, safetensors: F32, autotune_keys: [fp32], typestr: '<f4', numpy: float32, aliases: [float]}
+  # NVIDIA TF32 (A100 architecture whitepaper): a tensor-core COMPUTE format of 19 significant
+  # bits, stored and moved in a full 32-bit fp32 word.
+  tfloat32: {bits: 32, kind: float, max: 3.4011621342146535e+38}
+  # Google Brain bfloat16: the top 16 bits of a binary32.
+  bfloat16: {bits: 16, kind: float, max: 3.3895313892515355e+38, safetensors: BF16, autotune_keys: [bf16], typestr: '<V2'}
+  # IEEE 754-2008 binary16.
+  float16: {bits: 16, kind: float, max: 65504.0, safetensors: F16, autotune_keys: [fp16], typestr: '<f2', numpy: float16, aliases: [half]}
+  # OCP 8-bit Floating Point Specification (OFP8) v1.0, E4M3 (finite, no inf) and E5M2;
+  # torch.float8_e4m3fn / torch.float8_e5m2. The fnuz variants appear nowhere in this code.
+  float8_e4m3fn: {bits: 8, kind: float, max: 448.0, safetensors: F8_E4M3}
+  float8_e5m2: {bits: 8, kind: float, max: 57344.0, safetensors: F8_E5M2}
+  # Two's-complement integers of the stated width (C99 <stdint.h>, numpy, torch).
+  int64: {bits: 64, kind: int, safetensors: I64, autotune_keys: [int64], typestr: '<i8', numpy: int64, aliases: [long]}
+  int32: {bits: 32, kind: int, safetensors: I32, autotune_keys: [int32], typestr: '<i4', numpy: int32, aliases: [int]}
+  int16: {bits: 16, kind: int, safetensors: I16, autotune_keys: [int16], typestr: '<i2', numpy: int16}
+  int8: {bits: 8, kind: int, safetensors: I8, autotune_keys: [int8], typestr: '<i1', numpy: int8}
+  # Unsigned integers of the stated width; uint16/32/64 are carried by launch keys (a bf16 bit
+  # carrier, sort and RNG kernels), not by any container.
+  uint64: {bits: 64, kind: uint, autotune_keys: [uint64]}
+  uint32: {bits: 32, kind: uint, autotune_keys: [uint32]}
+  uint16: {bits: 16, kind: uint, autotune_keys: [uint16]}
+  uint8: {bits: 8, kind: uint, safetensors: U8, autotune_keys: [uint8], typestr: '<u1', numpy: uint8, aliases: [byte]}
+  # One byte per element in torch and numpy (no bit packing); `int1` is Triton's spelling.
+  bool: {bits: 8, kind: bool, safetensors: BOOL, autotune_keys: [int1, bool], typestr: '|b1', numpy: bool}
+  # Two binary32 / two binary64, interleaved [real, imag] (numpy and torch layout).
+  complex64: {bits: 64, kind: complex, typestr: '<c8', numpy: complex64}
+  complex128: {bits: 128, kind: complex, without_fp64: complex64, typestr: '<c16', numpy: complex128}
+
+# Packed weight encodings: bits per element is a fraction of a byte and the group scales live
+# beside the packed values. Only the encodings this engine executes
+# (`core/runtime/weight_encoding.SUPPORTED_WEIGHT_ENCODINGS`).
+packed:
+  # Weight-only int4, groups of 128 along in_features, asymmetric: one float16 scale and one
+  # float16 minimum per group; 8 nibbles per int32 (core/runtime/weight_encoding.py).
+  int4-g128-asym: {bits: 4, group_size: 128, scale_dtype: float16, min_dtype: float16}

nbxtensor/dtype/itemsize.py

@@ -0,0 +1,166 @@
+"""The dtype table and its reader — the ONE place a dtype's width is written (`dtypes.yml`, beside this module).
+
+Every width a dtype has is read from `dtypes.yml` through this module, and through no other reader: NeuroBrix's two
+engine branches read it here.
+
+Torch-free and numpy-free, at import and at execution.
+
+A dtype is named as the code passes it today: a name ("float16"), a torch spelling ("torch.float16", as a traced
+graph writes some — read as text, torch is never imported), or an enum member read by its NAME (NBXDtype is an IntEnum,
+and `str()` of one is its bare number on Python 3.11+ — never parsed). A name the table does not carry is refused by
+name, never given a width.
+"""
+from __future__ import annotations
+
+import copy
+import functools
+from pathlib import Path
+from enum import Enum
+from typing import Any, Dict, Tuple
+
+_SOURCE = "nbxtensor/dtype/dtypes.yml"
+_TABLE = Path(__file__).resolve().parent / "dtypes.yml"
+_TORCH_PREFIX = "torch."
+_KINDS = ("float", "int", "uint", "bool", "complex")
+
+
+@functools.lru_cache(maxsize=1)
+def _raw_table() -> Dict[str, Any]:
+    import yaml
+    if not _TABLE.is_file():
+        raise FileNotFoundError(f"ZERO FALLBACK: the dtype table is not found: {_TABLE}")
+    table = yaml.safe_load(_TABLE.read_text()) or {}
+    for section in ("dtypes", "packed"):
+        if not isinstance(table.get(section), dict):
+            raise KeyError(f"ZERO FALLBACK: {_TABLE} carries no {section!r} section")
+    return table
+
+
+def dtype_table() -> Dict[str, Any]:
+    """The whole table as written (`dtypes`, `packed`, and every other section), a copy: the reader a client keeps of
+    its own is fed from here, never from a second file."""
+    return copy.deepcopy(_raw_table())
+
+
+@functools.lru_cache(maxsize=1)
+def _tables() -> Tuple[Dict[str, Dict[str, Any]], Dict[str, Dict[str, Any]], Dict[str, str], Dict[str, str]]:
+    table = _raw_table()
+    dense, packed = dict(table["dtypes"]), dict(table["packed"])
+    for name, entry in dense.items():
+        bits = entry.get("bits")
+        if not isinstance(bits, int) or bits <= 0 or bits % 8:
+            raise ValueError(f"ZERO FALLBACK: {_SOURCE} gives dtype {name!r} {bits!r} bits — a dense "
+                             f"dtype is a whole number of bytes")
+        if entry.get("kind") not in _KINDS:
+            raise ValueError(f"ZERO FALLBACK: {_SOURCE} gives dtype {name!r} kind {entry.get('kind')!r} — one of "
+                             f"{', '.join(_KINDS)}")
+    for name, entry in packed.items():
+        for key in ("bits", "group_size", "scale_dtype", "min_dtype"):
+            if key not in entry:
+                raise KeyError(f"ZERO FALLBACK: {_SOURCE} packed encoding {name!r} states no {key}")
+    safetensors = {e["safetensors"]: n for n, e in dense.items() if "safetensors" in e}
+    keys = {k: n for n, e in dense.items() for k in (e.get("autotune_keys") or ())}
+    return dense, packed, safetensors, keys
+
+
+def canonical_name(dtype: Any) -> str:
+    """The table's name for `dtype` (not checked against the table: `itemsize` refuses)."""
+    if isinstance(dtype, Enum):
+        name = dtype.name
+        if name.endswith("_"):                 # `bool_`: Python's keyword escape, not a dtype name
+            name = name[:-1]
+    elif isinstance(dtype, str):
+        name = dtype
+    elif type(dtype).__module__ == "torch":    # a torch.dtype: its str() is its torch spelling
+        name = str(dtype)
+    else:
+        raise TypeError(f"ZERO FALLBACK: {dtype!r} ({type(dtype).__name__}) is not a dtype name, "
+                        f"a torch dtype or a dtype enum member ({_SOURCE})")
+    return name[len(_TORCH_PREFIX):] if name.startswith(_TORCH_PREFIX) else name
+
+
+def _dense_entry(dtype: Any) -> Tuple[str, Dict[str, Any]]:
+    dense, packed, _, _ = _tables()
+    name = canonical_name(dtype)
+    if name in dense:
+        return name, dense[name]
+    if name in packed:
+        raise ValueError(f"ZERO FALLBACK: {name!r} is a packed encoding of {packed[name]['bits']} bits "
+                         f"per element ({_SOURCE}) and has no whole itemsize — its bytes are "
+                         f"representation_bytes({name!r}, numel)")
+    raise ValueError(f"ZERO FALLBACK: unknown dtype {dtype!r}: {_SOURCE} has no entry {name!r} "
+                     f"(known: {', '.join(sorted(dense))}). Add it there with its source; no width "
+                     f"is ever guessed.")
+
+
+def known_dtypes() -> Tuple[str, ...]:
+    """Every dense dtype the table carries, sorted."""
+    return tuple(sorted(_tables()[0]))
+
+
+def itemsize(dtype: Any) -> int:
+    """Bytes per element of `dtype`, refused by name when the table does not carry it."""
+    return _dense_entry(dtype)[1]["bits"] // 8
+
+
+def kind(dtype: Any) -> str:
+    """The class of `dtype`: float, int, uint, bool or complex (the table's `kind`)."""
+    return _dense_entry(dtype)[1]["kind"]
+
+
+def finite_max(dtype: Any) -> float:
+    """The largest finite value of a floating `dtype` (the table's `max`); any other kind is refused by name."""
+    name, entry = _dense_entry(dtype)
+    if entry.get("kind") != "float":
+        raise ValueError(f"ZERO FALLBACK: {name!r} is not a floating dtype: it has no finite maximum")
+    if "max" not in entry:
+        raise KeyError(f"ZERO FALLBACK: {_SOURCE} states no max for {name!r}")
+    return float(entry["max"])
+
+
+def itemsize_table() -> Dict[str, int]:
+    """{name: bytes per element} over every dense dtype (a copy: the caller may not edit the table)."""
+    return {name: entry["bits"] // 8 for name, entry in _tables()[0].items()}
+
+
+def from_safetensors(code: str) -> str:
+    """The dtype name of a safetensors header code ("BF16" -> "bfloat16"), refused by name."""
+    names = _tables()[2]
+    if code not in names:
+        raise ValueError(f"ZERO FALLBACK: safetensors dtype code {code!r} has no entry in {_SOURCE} "
+                         f"(known: {', '.join(sorted(names))})")
+    return names[code]
+
+
+def from_autotune_key(key: str) -> str:
+    """The dtype name of a launch-key spelling ("fp16" -> "float16", "int1" -> "bool")."""
+    names = _tables()[3]
+    if key not in names:
+        raise ValueError(f"ZERO FALLBACK: launch-key dtype {key!r} has no entry in {_SOURCE} "
+                         f"(known: {', '.join(sorted(names))})")
+    return names[key]
+
+
+def storage_dtype(dtype: Any, stores_fp64: bool) -> str:
+    """The dtype a tensor of `dtype` is STORED in by a store that does (`stores_fp64`) or does not
+    hold 64-bit floats: float64 -> float32 and complex128 -> complex64 without, unchanged with."""
+    name, entry = _dense_entry(dtype)
+    if stores_fp64:
+        return name
+    return entry.get("without_fp64", name)
+
+
+def representation_bytes(dtype: Any, numel: int, split: bool = False) -> int:
+    """Bytes `numel` elements of `dtype` occupy as represented: a dense dtype at its itemsize, a
+    hi+lo SPLIT operand as two tensors of `dtype` (its low dtype), a packed encoding as its packed
+    values plus one scale and one minimum per group."""
+    numel = int(numel)
+    packed = _tables()[1]
+    name = canonical_name(dtype)
+    if name in packed:
+        if split:
+            raise ValueError(f"ZERO FALLBACK: a packed encoding ({name!r}) is not split hi+lo")
+        e = packed[name]
+        groups = -(-numel // int(e["group_size"]))
+        return -(-numel * int(e["bits"]) // 8) + groups * (itemsize(e["scale_dtype"]) + itemsize(e["min_dtype"]))
+    return numel * itemsize(dtype) * (2 if split else 1)

nbxtensor/dtype/nbx_dtype.py

@@ -0,0 +1,72 @@
+"""NBXDtype and its tables: the dtype enum NBXTensor carries (its values frozen: new members are appended, never
+renumbered), and the tables derived from the ONE dtype table (`dtypes.yml`): byte sizes, the spellings it parses (a
+table name, a launch-key spelling, an alias such as torch's `half`; a `torch.` prefix read as text), the numpy names
+and the array-interface type strings, the floating and complex kinds."""
+
+from __future__ import annotations
+
+from enum import IntEnum
+
+# ============================================================================
+# DTYPE SYSTEM
+# ============================================================================
+
+class NBXDtype(IntEnum):
+    float16 = 0
+    bfloat16 = 1
+    float32 = 2
+    float64 = 3
+    int8 = 4
+    int16 = 5
+    int32 = 6
+    int64 = 7
+    uint8 = 8
+    bool_ = 9
+    complex64 = 10
+    complex128 = 11
+
+
+def _table_name(m: "NBXDtype") -> str:
+    """A member's name in the dtype table (`bool_` is Python's keyword escape of `bool`)."""
+    return "bool" if m is NBXDtype.bool_ else m.name
+
+
+def _derive():
+    """Every table of the enum, read from the one dtype table (`dtypes.yml`): widths, the spellings a caller may name
+    a dtype by, the numpy names, the array-interface type strings, the floating and complex kinds."""
+    from .itemsize import dtype_table
+    entries = dtype_table()["dtypes"]
+    sizes, names, numpy_names, typestr, floating, complex_ = {}, {}, {}, {}, set(), set()
+    for m in NBXDtype:
+        e = entries[_table_name(m)]
+        sizes[m] = int(e["bits"]) // 8
+        for spelling in [_table_name(m), *(e.get("autotune_keys") or ()), *(e.get("aliases") or ())]:
+            names[spelling] = m
+        if "numpy" in e:
+            numpy_names[e["numpy"]] = m
+        typestr[m] = e["typestr"]
+        if e["kind"] == "float":
+            floating.add(m)
+        elif e["kind"] == "complex":
+            complex_.add(m)
+    return sizes, names, numpy_names, typestr, frozenset(floating), frozenset(complex_)
+
+
+# Read once from the table: no width, spelling or type string is written here (R-NBXT 2, one fact once).
+_DTYPE_SIZES, _DTYPE_FROM_STR, _NUMPY_DTYPE_MAP, _DTYPE_TYPESTR, _FLOATING_DTYPES, _COMPLEX_DTYPES = _derive()
+
+
+def dtype_size(dtype: NBXDtype) -> int:
+    return _DTYPE_SIZES[dtype]
+
+
+def parse_dtype(d) -> NBXDtype:
+    """Parse dtype from string or NBXDtype."""
+    if isinstance(d, NBXDtype):
+        return d
+    if isinstance(d, str):
+        s = d.replace('torch.', '')
+        return _DTYPE_FROM_STR[s]
+    # Unknown type — try string conversion
+    s = str(d).replace('torch.', '')
+    return _DTYPE_FROM_STR[s]

nbxtensor/dtype/operands.py

@@ -0,0 +1,158 @@
+"""The operand decisions of the kernels, as dtype names (DESIGN rule 12: the mechanism in the library): the dtype a
+contraction's partials are summed in, how a matrix unit carries an operand, how an attention's operands are aligned
+when they disagree, the scale of a split representation. Pure functions of dtype names and of a device's matrix unit
+(`ops._configs.matrix_unit`); the application decides which device and which policy, never these rules.
+"""
+from __future__ import annotations
+
+from typing import Optional
+
+from nbxtensor.dtype import itemsize as _itemsize
+
+__all__ = ["contraction_accumulator_dtype", "contraction_accumulator_bytes", "matrix_unit_representation",
+           "matrix_unit_flash_representation", "ATTENTION_OPERAND_ALIGNMENTS", "attention_operand_alignment",
+           "attention_operand_dtypes", "matrix_unit_split_scale", "constant_load_dtype"]
+
+
+def contraction_accumulator_dtype(dtype: str) -> str:
+    """The dtype (a name) the partial results of a contraction split over the axis it reduces
+    are summed in, when a stretch runs in slices of that axis (`core/strategies/chunked_piece`):
+    the split contraction is the SAME contraction, so its partials are summed in the accumulator
+    its kernel keeps — float32 for a half dtype (every mm/bmm/conv kernel accumulates fp32, and
+    `sum` is an AMP_FP32 op), the dtype itself for float32 and float64 — and the sum is stored
+    once, in the dtype the contraction's own output was given. A partial summed at the pass
+    dtype rounds once per slice where the whole op rounds once. Any other dtype is refused by
+    name: a contraction over a token axis yields a floating value.
+
+    The one definition both engines read."""
+    if dtype in ("float16", "bfloat16"):
+        return "float32"
+    if dtype in ("float32", "float64"):
+        return dtype
+    raise ValueError(f"contraction_accumulator_dtype: {dtype!r} is not a floating dtype a "
+                     f"contraction's partials are summed in")
+
+
+def contraction_accumulator_bytes(store_bytes: int) -> int:
+    """`contraction_accumulator_dtype` read on a WIDTH (bytes per element), for Prism, which prices
+    a plan in widths: a 2-byte store (fp16 or bf16 alike) accumulates in 4 bytes; 4 and 8 are their
+    own accumulator. The same rule (a unit test holds the two equal over every floating dtype)."""
+    if store_bytes == 2:
+        return 4
+    if store_bytes in (4, 8):
+        return store_bytes
+    raise ValueError(f"contraction_accumulator_bytes: a {store_bytes}-byte store is no floating "
+                     f"dtype a contraction's partials are summed in")
+
+
+def matrix_unit_representation(dtype: str, unit: dict) -> Optional[str]:
+    """How the hardware profile's matrix unit (`matrix_unit`, nbxtensor.ops._configs.matrix_unit) carries one GEMM
+    operand whose dtype in memory is `dtype` (a name): "native" when it IS the unit's operand dtype (one MMA term
+    per operand), "split" when the profile declares `matrix_unit.fp32_split` with this dtype among its `operands`
+    (the value, widened to fp32, is carried as hi + lo of the operand dtype under a power-of-two scale, see
+    `matrix_unit_split_scale`; hi*hi + hi*lo + lo*hi with fp32 accumulation outside the unit, the vendor's fp32
+    arithmetic kept), else None (the GEMM keeps its tl.dot kernel). The ONE decision: the wrappers launch with it
+    and the derived census keys with it (nbxtensor.ops.launch_keys.matrix_unit_operands)."""
+    if not unit or not unit.get("mm"):
+        return None
+    if dtype == unit["operand_dtype"]:
+        return "native"
+    split = unit.get("fp32_split") or {}
+    if split.get("mm") and dtype in (split.get("operands") or ()):
+        return "split"
+    return None
+
+
+def matrix_unit_flash_representation(q: str, k: str, v: str, unit: dict) -> Optional[str]:
+    """How the profile's matrix unit carries an attention's three operands (names, already aligned by
+    `attention_operand_dtypes`) in its flash kernel (nbxtensor/ops/flash_attention_m8n8k4.py): "native" when all three
+    ARE the unit's operand dtype and the profile lists `matrix_unit.flash` rows; "split" when all three are among
+    `matrix_unit.fp32_split.operands` and the profile lists `matrix_unit.fp32_split.flash` rows — Q, K, V and the
+    softmax's P carried as hi + lo as `matrix_unit_representation` carries a GEMM operand, the vendor's fp32
+    arithmetic kept; else None (the attention keeps its tl.dot kernel). The ONE decision the wrapper launches with
+    and the derived census routes with (nbxtensor.ops.launch_keys.unit_flash)."""
+    if not unit or not (q == k == v):
+        return None
+    if q == unit["operand_dtype"]:
+        return "native" if unit.get("flash") else None
+    split = unit.get("fp32_split") or {}
+    if split.get("flash") and q in (split.get("operands") or ()):
+        return "split"
+    return None
+
+
+#: The ways a hardware profile may align attention operands that disagree in dtype
+#: (`precision.attention_operands`): every operand takes the narrowest of the three, or the widest.
+ATTENTION_OPERAND_ALIGNMENTS = ("narrowest", "widest")
+#: The operand dtypes the alignment chooses among; their widths are the dtype table's.
+_ATTENTION_OPERAND_DTYPES = ("float16", "bfloat16", "float32", "float64")
+
+
+def attention_operand_alignment(profile: dict) -> str:
+    """The alignment the hardware profile declares for attention operands that disagree in dtype
+    (`precision.attention_operands`, one of ATTENTION_OPERAND_ALIGNMENTS). A profile that declares
+    none, or another word, is refused by name: the choice moves a whole attention between the
+    matrix unit and the scalar units, so it is never guessed.
+
+    The one definition both engines read."""
+    v = (profile.get("precision") or {}).get("attention_operands")
+    if v not in ATTENTION_OPERAND_ALIGNMENTS:
+        raise ValueError(f"hardware profile {profile.get('architecture')!r}: "
+                         f"precision.attention_operands is {v!r}, not one of "
+                         f"{ATTENTION_OPERAND_ALIGNMENTS}")
+    return v
+
+
+def attention_operand_dtypes(q: str, k: str, v: str, alignment: str,
+                             q_round: Optional[str] = None) -> tuple:
+    """The (q, k, v, q_round) dtypes (names) an attention computes with — the ONE decision both
+    engines' attention reads (R30): the Triton wrapper (`launch_keys.sdpa_operand_dtypes`), the
+    derived census and Prism's width estimate through it, and the compiled engine through its
+    twin. The kernels need matching operands: three that already match (Q judged at its KV-cache
+    rounding `q_round` when the cache asks for one) are returned unchanged; three that disagree
+    all take the narrowest or the widest of them, as the profile's `alignment` declares
+    (`attention_operand_alignment`), the first in (q, k, v) order on a tie, and q_round is
+    dropped. A non-floating operand is refused by name.
+
+    Measured on V100 (2026-10-08, nbx/campaigns/2026_10_08_allegro_prof/probe): an fp16 attention
+    whose V alone came out of an fp32-stored GEMM ran 22.0 s/op on the scalar flash when aligned
+    to the widest, 1.16 s/op on the matrix-unit flash when aligned to the narrowest.
+
+    The one definition both engines read."""
+    qj = q_round if q_round is not None else q
+    for d in (qj, k, v):
+        if d not in _ATTENTION_OPERAND_DTYPES:
+            raise ValueError(f"attention_operand_dtypes: {d!r} is not a floating dtype an "
+                             f"attention computes in")
+    if qj == k == v:
+        return q, k, v, q_round
+    if alignment not in ATTENTION_OPERAND_ALIGNMENTS:
+        raise ValueError(f"attention_operand_dtypes: alignment {alignment!r} is not one of "
+                         f"{ATTENTION_OPERAND_ALIGNMENTS}")
+    pick = min if alignment == "narrowest" else max
+    t = pick((qj, k, v), key=_itemsize.itemsize)
+    return t, t, t, None
+
+
+def matrix_unit_split_scale(unit: dict) -> tuple:
+    """(hi_exp, lo_shift) of the split representation, derived from the unit's operand dtype (never written in a
+    kernel): each operand tile is scaled by the power of two that brings its largest magnitude into
+    [2^hi_exp, 2^(hi_exp+1)) — the top binade below the operand dtype's overflow, so no hi overflows and every
+    value down to 2^-(hi_exp+1+|min exponent|) of the tile's largest stays normal — and lo, the residual
+    x - hi, is stored times 2^lo_shift (the operand dtype's significand width) so it does not underflow where hi
+    does not (Ootomo & Yokota 2022, arXiv 2203.03341, eq. 19-24)."""
+    import numpy as np
+    fi = np.finfo(np.dtype(unit["operand_dtype"]))
+    return int(fi.maxexp) - 2, int(fi.nmant) + 1
+
+
+def constant_load_dtype(traced: str, compute: str, stores_fp64: bool) -> str:
+    """The dtype (a name) an embedded graph constant is bound in by the Triton engines
+    (`GraphExecutor._load_constant_triton`): a bfloat16 constant is decoded to the compute dtype
+    when that is half (fp16 bits from bf16 on hardware computing fp16, the bf16 bits kept under
+    bf16), to float32 otherwise; float64 / complex128 are held as the branch holds them
+    (`dtype.itemsize.storage_dtype`); every other dtype is kept as traced. The loader and Prism's width
+    pass both ask it."""
+    if traced == "bfloat16":
+        return compute if compute in ("float16", "bfloat16") else "float32"
+    return _itemsize.storage_dtype(traced, stores_fp64)

nbxtensor/dtype/outputs.py

@@ -0,0 +1,64 @@
+"""The output dtype of an op (DESIGN rule 12: the mechanism in the library): read from the op's traced record, or
+decided by the cast-back rule of an op computed in float32 inside a half forward. Pure functions of dtype names and of
+a plain dict; the one definition every engine imports (an engine holding torch dtypes, NBXDtype or strings reads the
+same answer: only the last dotted segment of a dtype's name is read).
+"""
+from __future__ import annotations
+
+from typing import Any, Dict, FrozenSet, Optional
+
+__all__ = ["COMPLEX_DTYPE_NAMES", "traced_output_dtype_name", "traced_output_is_complex", "amp_fp32_output_dtype"]
+
+COMPLEX_DTYPE_NAMES: FrozenSet[str] = frozenset({"complex32", "complex64", "complex128", "chalf", "cfloat", "cdouble"})
+_GRAPH_FLOAT_NAMES: FrozenSet[str] = frozenset({"float16", "bfloat16", "float32", "float64"})
+_HALF_GRAPH_NAMES: FrozenSet[str] = frozenset({"float16", "bfloat16"})
+
+
+def _name(dtype) -> str:
+    return str(dtype).rsplit(".", 1)[-1]
+
+
+def traced_output_dtype_name(record: Optional[Dict[str, Any]]) -> Optional[str]:
+    """The plain name of the op's first traced output dtype (`output_dtypes[0]` at the record's TOP level), or None
+    when the record states none."""
+    if not isinstance(record, dict):
+        return None
+    dts = record.get("output_dtypes") or ()
+    return _name(dts[0]) if dts else None
+
+
+def traced_output_is_complex(record: Optional[Dict[str, Any]]) -> bool:
+    """True when the op's traced record types an output complex.
+
+    The traced OUTPUT dtype is the only signal: an op's name cannot answer (a plain `mul` by the scalar `1j` produces
+    complex), nor can its inputs (that scalar is no input tensor). It matters because a half input at a
+    complex-producing op yields complex32, which `exp`, `angle` and the inverse real FFT have no kernel for. Only
+    `output_dtypes` at the record's TOP level is read: a nested `attributes` dict reads as "not complex"."""
+    if not isinstance(record, dict):
+        return False
+    return any(_name(d) in COMPLEX_DTYPE_NAMES for d in (record.get("output_dtypes") or ()))
+
+
+def amp_fp32_output_dtype(compute_dtype: str, graph_dtype: Optional[str], safe: bool,
+                          narrowed: bool, traced: Optional[str] = None) -> str:
+    """The OUTPUT dtype (a name) of an op computed in float32 inside under the half compute dtype `compute_dtype`.
+
+    * bfloat16 on a bfloat16 graph (`graph_dtype`, the dtype the model was traced in): bfloat16 — a float32 result
+      always fits bfloat16's range, and the model's own forward returns its input's dtype there.
+    * bfloat16 on a graph of another floating dtype (a float32 model computed in bfloat16): float32, the dtype the
+      model's forward returned. A graph dtype that is no floating dtype is refused by name.
+    * float16: float16 only under the precision contract — `safe` (narrows every such op) or `narrowed` (this op is
+      in the narrowed set); float32 otherwise. An op `traced` in float32 inside a half graph is the model's own
+      float32 island: only `narrowed` brings it to float16, never `safe`."""
+    if compute_dtype == "bfloat16":
+        if graph_dtype not in _GRAPH_FLOAT_NAMES:
+            raise ValueError(
+                f"amp_fp32_output_dtype: under bfloat16 compute the output follows the graph dtype, and "
+                f"{graph_dtype!r} is none of {sorted(_GRAPH_FLOAT_NAMES)} — the traced graph must state it")
+        return "bfloat16" if graph_dtype == "bfloat16" else "float32"
+    if compute_dtype == "float16":
+        if traced == "float32" and graph_dtype in _HALF_GRAPH_NAMES:
+            return "float16" if narrowed else "float32"
+        return "float16" if (safe or narrowed) else "float32"
+    raise ValueError(f"amp_fp32_output_dtype: compute dtype {compute_dtype!r} is not a half dtype — the cast-back "
+                     f"rule exists for bfloat16 and float16 only")