mlx-dfloat 0.1.0__py3-none-any.whl

mlx_dfloat/__init__.py

@@ -0,0 +1,25 @@
+"""Run DFloat11 losslessly compressed BF16 models on Apple Silicon with MLX."""
+
+from mlx_dfloat._version import __version__
+from mlx_dfloat.errors import (
+    DFloatAccessError,
+    DFloatBackendError,
+    DFloatDependencyError,
+    DFloatError,
+    DFloatFormatError,
+    DFloatIntegrationError,
+    DFloatResourceError,
+    DFloatUnsupportedError,
+)
+
+__all__ = [
+    "DFloatAccessError",
+    "DFloatBackendError",
+    "DFloatDependencyError",
+    "DFloatError",
+    "DFloatFormatError",
+    "DFloatIntegrationError",
+    "DFloatResourceError",
+    "DFloatUnsupportedError",
+    "__version__",
+]

mlx_dfloat/_memory_caps.py

@@ -0,0 +1,74 @@
+"""Hardware-aware MLX memory caps (kernel-watchdog panic guard).
+
+Derives wired + memory caps from the device's reported working-set size and
+clamps strictly below it. Returns (0, 0) as a no-op signal on devices/CI images
+that report no working-set size. Mirrors the mlx-taef and mlx-quant-fidelity pattern.
+"""
+
+import mlx.core as mx
+
+DESIRED_WIRED_GB = 20
+DESIRED_MEMORY_GB = 22
+HEADROOM_GB = 2
+
+
+def _clamp_caps_gb(max_recommended_gb: int) -> tuple[int, int]:
+    """Clamp the desired caps to fit a device with `max_recommended_gb` working set."""
+    if max_recommended_gb <= 0:
+        return (0, 0)
+    wired_gb = min(DESIRED_WIRED_GB, max(1, max_recommended_gb - HEADROOM_GB))
+    memory_gb = min(DESIRED_MEMORY_GB, max(wired_gb + 1, max_recommended_gb))
+    return (wired_gb, memory_gb)
+
+
+def compute_safe_caps_gb() -> tuple[int, int]:
+    """Return (wired_gb, memory_gb) that fit the current device, or (0, 0)."""
+    try:
+        info = mx.device_info()
+        max_gb = int(info.get("max_recommended_working_set_size", 0)) // (1024**3)
+    except Exception:
+        return (0, 0)
+    return _clamp_caps_gb(max_gb)
+
+
+def device_string() -> str | None:
+    """Human-readable chip + unified-memory size (e.g. 'Apple M1 Max, 32 GB'), or None.
+
+    Report provenance only — never used for gating. Returns None when MLX does not
+    report a device name (CI containers, future backends).
+    """
+    try:
+        info = mx.device_info()
+        name = info.get("device_name")
+        if not isinstance(name, str) or not name:
+            return None
+        mem = info.get("memory_size")
+        if isinstance(mem, int) and mem > 0:
+            return f"{name}, {round(mem / 1024**3)} GB"
+        return name
+    except Exception:
+        return None
+
+
+def install_memory_caps() -> tuple[int, int]:
+    """Apply wired + memory caps for the current device. Idempotent; never raises.
+
+    Returns the (wired_gb, memory_gb) actually installed; each is 0 when that cap could not be
+    applied, and both are 0 on a device with no reported working-set size. The two caps are
+    applied independently, so a failed wired cap still leaves the memory cap in place.
+    """
+    wired_gb, memory_gb = compute_safe_caps_gb()
+    if wired_gb == 0:
+        return (0, 0)
+    try:
+        mx.set_wired_limit(wired_gb * 1024**3)
+    except Exception:
+        wired_gb = 0
+    try:
+        mx.set_memory_limit(memory_gb * 1024**3)
+    except Exception:
+        memory_gb = 0
+    return (wired_gb, memory_gb)
+
+
+__all__ = ["compute_safe_caps_gb", "device_string", "install_memory_caps"]

mlx_dfloat/_metal_decode.py

@@ -0,0 +1,420 @@
+"""The Metal DF11 decode kernel: a per-thread, bounded, status-reporting mirror of the reference algorithm."""
+
+from dataclasses import dataclass
+from typing import Any
+
+import mlx.core as mx
+import numpy as np
+import numpy.typing as npt
+
+from mlx_dfloat.decode import DecodeResult
+from mlx_dfloat.errors import DFloatBackendError
+from mlx_dfloat.format import GroupArrays, MxGroup
+
+THREADS = 512
+CAP = 16192
+_SCRATCH_BYTES = 4 * (16 + 16 + 1 + 16)
+
+
+def _metal_static_bytes(declared: int) -> int:
+    """The static threadgroup memory Metal reserves for ``declared`` bytes of arrays.
+
+    Metal reports ``staticThreadgroupMemoryLength`` at 16-byte granularity (measured with
+    standalone compiles: 196 -> 208, 192 -> 192; mlx 0.32.2, macOS 27.0, M1 Max).
+    """
+    return -(-declared // 16) * 16
+
+
+THREADGROUP_BYTES_STAGED = _metal_static_bytes(
+    2 * CAP + _SCRATCH_BYTES
+)  # 32,580 -> 32,592 < 32,768
+THREADGROUP_BYTES_DIRECT = _metal_static_bytes(2 + _SCRATCH_BYTES)  # 198 -> 208
+
+_KERNEL: Any | None = None
+# (force_direct, poison_buf, unguarded_gap_read) -> warmed
+_PIPELINES: dict[tuple[bool, bool, bool], bool] = {}
+
+_SOURCE = r"""
+    const uint t = thread_position_in_threadgroup.x;
+    const uint b = threadgroup_position_in_grid.x;
+    const uint g = b * 512u + t;
+    const uint n_bytes = (uint)encoded_shape[0];
+    const uint n = (uint)sm_shape[0];
+    const uint n_luts = (uint)luts_shape[0];
+    const uint n_launch = (uint)positions_shape[0] - 1u;
+    const uint blk_start = positions[b];
+    const uint blk_end = positions[b + 1u];
+    const uint interval = blk_end - blk_start;
+    const bool staged = (!FORCE_DIRECT) && (interval <= (uint)CAP);
+    const uint last_real = (n_bytes - 1u) / 8u;
+    const bool real = (8u * g) < n_bytes;
+
+    threadgroup ushort buf[FORCE_DIRECT ? 1 : CAP];
+    threadgroup uint simd_totals[16];
+    threadgroup uint simd_pre[16];
+    threadgroup uint block_total[1];
+    threadgroup uint simd_bad[16];
+    if (POISON_BUF && staged) { for (uint i = t; i < (uint)CAP; i += 512u) buf[i] = (ushort)0xDEAD; }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    const uint gbit = 5u * g;
+    const uint gbyte = gbit >> 3;
+    const uint gsh = gbit & 7u;
+    const uint g0 = (uint)gaps[gbyte];
+    const uint g1 = (UNGUARDED_GAP_READ || gsh > 3u) ? (uint)gaps[gbyte + 1u] : 0u;
+    const uint gap = (((g0 << 8) | g1) >> (11u - gsh)) & 31u;
+
+    ulong w = 0;
+    for (uint k = 0; k < 8u; ++k) {
+        const uint i = 8u * g + k;
+        w = (w << 8) | (ulong)((i < n_bytes) ? encoded[i] : (uchar)0);
+    }
+    uint la = 0;
+    for (uint k = 8u; k < 12u; ++k) {
+        const uint i = 8u * g + k;
+        la = (la << 8) | (uint)((i < n_bytes) ? encoded[i] : (uchar)0);
+    }
+
+    uint count = 0;
+    uint p = gap;
+    bool bad = false;
+    if (real) {
+        for (uint it = 0; it < 64u && p < 64u; ++it) {
+            const uint word = (p <= 32u) ? (uint)(w >> (32u - p)) : ((uint)(w << (p - 32u)) | (la >> (64u - p)));
+            uint sym = (uint)luts[word >> 24];
+            for (uint level = 1u; level <= 3u && sym >= 240u; ++level) {
+                const uint row = 256u - sym;
+                if (row < 1u || row > n_luts - 2u) { sym = 240u; break; }
+                sym = (uint)luts[row * 256u + ((word >> (24u - 8u * level)) & 0xFFu)];
+            }
+            const uint len = (sym < 240u) ? (uint)luts[(n_luts - 1u) * 256u + sym] : 0u;
+            if (len == 0u) { bad = true; break; }
+            count += 1u;
+            p += len;
+        }
+    }
+    const uint sg = t / 32u;
+    const uint lane = t % 32u;
+    const uint local = simd_prefix_exclusive_sum(count);
+
+    // Chain check (reference `broken`): a clean thread before the last real one must end exactly where the
+    // next thread's gap says its first code starts, unless its last code runs past the stream end.
+    bool broken = false;
+    if (real && !bad && g < last_real) {
+        const uint nbit = 5u * (g + 1u);
+        const uint nb0 = (uint)gaps[nbit >> 3];
+        const uint nsh = nbit & 7u;
+        const uint nb1 = (nsh > 3u) ? (uint)gaps[(nbit >> 3) + 1u] : 0u;
+        const uint next_gap = (((nb0 << 8) | nb1) >> (11u - nsh)) & 31u;
+        const ulong end_bit = 64ul * (ulong)g + (ulong)p;
+        broken = (end_bit < 8ul * (ulong)n_bytes) && (p != 64u + next_gap);
+    }
+    const bool relevant_bad = bad && (g < last_real);
+    const uint vote = (simd_any(relevant_bad) ? 1u : 0u) | (simd_any(broken) ? 4u : 0u);   // uniform flow
+
+    if (lane == 31u) { simd_totals[sg] = local + count; simd_bad[sg] = vote; }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    if (sg == 0u) {
+        const uint v = (lane < 16u) ? simd_totals[lane] : 0u;
+        const uint pre = simd_prefix_exclusive_sum(v);
+        if (lane < 16u) simd_pre[lane] = pre;
+        if (lane == 15u) block_total[0] = pre + v;
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    const uint first = blk_start + simd_pre[sg] + local;
+    const uint total = block_total[0];
+
+    if (real && count > 0u) {
+        uint idx = first;
+        const uint stop = min(min(first + count, n), blk_end);
+        uint q = gap;
+        for (uint it = 0; it < 64u && idx < stop; ++it) {
+            const uint word = (q <= 32u) ? (uint)(w >> (32u - q)) : ((uint)(w << (q - 32u)) | (la >> (64u - q)));
+            uint sym = (uint)luts[word >> 24];
+            for (uint level = 1u; level <= 3u && sym >= 240u; ++level) {
+                const uint row = 256u - sym;
+                if (row < 1u || row > n_luts - 2u) { sym = 240u; break; }
+                sym = (uint)luts[row * 256u + ((word >> (24u - 8u * level)) & 0xFFu)];
+            }
+            const uint len = (sym < 240u) ? (uint)luts[(n_luts - 1u) * 256u + sym] : 0u;
+            if (len == 0u) break;
+            const uint smv = (uint)sm[idx];
+            const ushort v = (ushort)(((smv & 0x80u) << 8) | (sym << 7) | (smv & 0x7Fu));
+            if (staged) buf[idx - blk_start] = v; else out[idx] = v;
+            idx += 1u;
+            q += len;
+        }
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    if (staged) {
+        const uint limit = min(interval, n - blk_start);
+        for (uint i = t; i < limit; i += 512u) out[blk_start + i] = buf[i];
+    }
+    if (t == 511u) {
+        uint flags = 0u;
+        for (uint s = 0; s < 16u; ++s) flags |= simd_bad[s];
+        uint word = flags;                                          // 1 invalid code, 4 broken chain
+        const bool last_block = (b + 1u == n_launch);
+        if ((!last_block && total != interval) || (last_block && total < interval)) word |= 2u;
+        if (!staged) word |= 8u;                                    // informational: direct path taken
+        status[b] = word;
+    }
+"""
+
+
+def _build_kernel() -> Any:
+    global _KERNEL
+    if _KERNEL is None:
+        _KERNEL = mx.fast.metal_kernel(
+            name="df11_decode",
+            input_names=["encoded", "sm", "luts", "gaps", "positions"],
+            output_names=["out", "status"],
+            source=_SOURCE,
+        )
+    return _KERNEL
+
+
+def _dispatch(
+    group: MxGroup,
+    *,
+    force_direct: bool,
+    poison_buf: bool,
+    init_value: int | None,
+    unguarded_gap_read: bool = False,
+) -> tuple[mx.array, mx.array]:
+    out, status = _build_kernel()(
+        inputs=[
+            group.encoded_exponent,
+            group.sign_mantissa,
+            group.luts,
+            group.gaps,
+            group.positions,
+        ],
+        template=[
+            ("CAP", CAP),
+            ("FORCE_DIRECT", force_direct),
+            ("POISON_BUF", poison_buf),
+            ("UNGUARDED_GAP_READ", unguarded_gap_read),
+        ],
+        grid=(THREADS * group.n_launch, 1, 1),
+        threadgroup=(THREADS, 1, 1),
+        output_shapes=[(group.n_elements,), (group.n_launch,)],
+        output_dtypes=[mx.uint16, mx.uint32],
+        init_value=init_value,
+    )
+    return out, status
+
+
+# The warm-up groups. MLX binds a read-only input of fewer than 8 entries in the `constant` address
+# space and names (and compiles) one pipeline per input-binding signature, so readiness must warm every
+# signature a real group can present, not one of them:
+#
+# 1. `warmup-7-block`: seven full 4096-byte blocks, `positions` has eight entries and every input binds
+#    `device`, the signature of any group with at least seven launched blocks. Each thread's 8 bytes
+#    FF FF FF FF FF FF FF FE decode under the H1 tables (0xxxxxxx -> exponent 127, length 1;
+#    10xxxxxx -> 126, length 2; 11xxxxxx -> 128, length 3) to 21 codes `111` (exponent 128) then one
+#    `0` (exponent 127): 22 elements per thread, 11,264 per block, under CAP, so the staged
+#    instantiation stages every block. An all-zero stream would decode to 64 elements per thread,
+#    32,768 per block, over CAP, and never exercise staging.
+# 2. `warmup-1-block`: one thread of that same stream (8 bytes, 22 elements), so `positions` has two
+#    entries and binds `constant` while every other input still binds `device`: the signature of any
+#    group with fewer than seven launched blocks (positions.size <= 7).
+# 3. `warmup-tiny`: the H1 hand fixture (one byte 0x9B, three elements), so `encoded`, `sm` and
+#    `positions` all bind `constant`: the signature of a very small group.
+_WARMUP_BLOCKS = 7
+_WARMUP_CODES_PER_THREAD = 22
+_WARMUP_ELEMENTS_PER_BLOCK = THREADS * _WARMUP_CODES_PER_THREAD
+_WARMUP_THREAD_BYTES = (0xFF,) * 7 + (0xFE,)
+_GAPS_BYTES_PER_BLOCK = 320  # 512 threads x 5 bits
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class _Warmup:
+    """One warm-up group: the signature it exercises, its host arrays and its hand-derived BF16 bits."""
+
+    name: str
+    arrays: GroupArrays
+    expected: npt.NDArray[np.uint16]
+
+
+def _h1_luts() -> npt.NDArray[np.uint8]:
+    """H1: 127 `0`, 126 `10`, 128 `110` (EOF `111` inherits 128, as upstream); row 1 holds the lengths."""
+    row0 = np.zeros(256, np.uint8)
+    row0[0:128], row0[128:192], row0[192:256] = 127, 126, 128
+    lens = np.zeros(256, np.uint8)
+    lens[126], lens[127], lens[128] = 2, 1, 3
+    return np.stack([row0, lens])
+
+
+def _thread_stream_arrays(n_threads: int, n_blocks: int, name: str) -> _Warmup:
+    """`n_threads` copies of the FF..FE thread stream over `n_blocks` blocks, zero gaps, sm[i] = i & 0xFF."""
+    n_elements = n_threads * _WARMUP_CODES_PER_THREAD
+    i = np.arange(n_elements, dtype=np.uint32)
+    exponent = np.where(i % _WARMUP_CODES_PER_THREAD == _WARMUP_CODES_PER_THREAD - 1, 127, 128)
+    sm = i & 0xFF
+    expected = (((sm & 0x80) << 8) | (exponent.astype(np.uint32) << 7) | (sm & 0x7F)).astype(
+        np.uint16
+    )
+    per_block = _WARMUP_ELEMENTS_PER_BLOCK if n_blocks > 1 else n_elements
+    arrays = GroupArrays(
+        encoded_exponent=np.tile(np.array(_WARMUP_THREAD_BYTES, np.uint8), n_threads),
+        sign_mantissa=sm.astype(np.uint8),
+        luts=_h1_luts(),
+        gaps=np.zeros(_GAPS_BYTES_PER_BLOCK * n_blocks, np.uint8),
+        output_positions=(np.arange(n_blocks + 1) * per_block).astype(np.uint32),
+        split_positions=np.zeros(0, np.int64),
+    )
+    return _Warmup(name=name, arrays=arrays, expected=expected)
+
+
+def _warmups() -> tuple[_Warmup, ...]:
+    """The three warm-up groups, one per input-binding signature (see the note above)."""
+    tiny = GroupArrays(
+        encoded_exponent=np.array([0x9B], np.uint8),  # 10 | 0 | 110 -> exponents 126, 127, 128
+        sign_mantissa=np.array([0x00, 0x80, 0x7F], np.uint8),
+        luts=_h1_luts(),
+        gaps=np.zeros(_GAPS_BYTES_PER_BLOCK, np.uint8),
+        output_positions=np.array([0, 3], np.uint32),
+        split_positions=np.zeros(0, np.int64),
+    )
+    return (
+        _thread_stream_arrays(_WARMUP_BLOCKS * THREADS, _WARMUP_BLOCKS, "warmup-7-block"),
+        _thread_stream_arrays(1, 1, "warmup-1-block"),
+        _Warmup(
+            name="warmup-tiny",
+            arrays=tiny,
+            expected=np.array([0x3F00, 0xBF80, 0x407F], np.uint16),
+        ),
+    )
+
+
+def _warmup_arrays() -> GroupArrays:
+    """The seven-block warm-up group as host arrays."""
+    return _warmups()[0].arrays
+
+
+def _warmup_expected() -> npt.NDArray[np.uint16]:
+    """The seven-block warm-up group's BF16 bits, derived by hand (none equals 0)."""
+    return _warmups()[0].expected
+
+
+def _warmup_group() -> MxGroup:
+    """The seven-block warm-up group as evaluated MLX arrays (also the register-pressure probe's launch group)."""
+    return _warmup_arrays().to_mx(name="warmup-7-block")
+
+
+def _warmup_groups() -> list[tuple[MxGroup, npt.NDArray[np.uint16]]]:
+    """Every warm-up group as evaluated MLX arrays with its expected bits, in signature order."""
+    return [(w.arrays.to_mx(name=w.name), w.expected) for w in _warmups()]
+
+
+def ensure_pipeline(
+    *, force_direct: bool, poison_buf: bool = False, unguarded_gap_read: bool = False
+) -> None:
+    """Compile this instantiation's pipelines for every input-binding signature and prove each decodes bit-exactly.
+
+    Each distinct template tuple is its own JIT compile, so each is warmed once and the result
+    is cached in `_PIPELINES`. MLX compiles one pipeline per input-binding signature (a read-only
+    input of fewer than 8 entries binds `constant`, the rest `device`), so the warm-up dispatches
+    three groups: seven blocks with every input in `device` memory, one block whose `positions`
+    alone binds `constant`, and an all-tiny group whose `encoded`, `sign_mantissa` and `positions`
+    bind `constant`. Those are the signatures real groups present, so the pipelines compiled here
+    are the ones the caller's decodes reuse, whatever the group's size. Readiness means every one
+    of them compiled, dispatched 512-thread groups, and wrote every element of its group
+    bit-exactly, with no error bit on any block and every block on the expected write path
+    (staged unless `force_direct`; every warm-up block fits under `CAP`). Each output is prefilled
+    with 0, which no expected word equals, so an element the kernel fails to write cannot pass on
+    recycled memory. `unguarded_gap_read` is a test-only mutant (see `decode`).
+
+    Raises:
+        DFloatBackendError: Metal is unavailable, a pipeline fails to compile or dispatch here
+            (a pipeline ceiling below 512 threads, a driver refusal), a warm-up group decodes to
+            the wrong bits or reports an error bit on any block, or any block runs on the wrong
+            write path.
+    """
+    key = (force_direct, poison_buf, unguarded_gap_read)
+    if _PIPELINES.get(key):
+        return
+    if not mx.metal.is_available():
+        raise DFloatBackendError("Metal is not available on this machine")
+    for group, expected in _warmup_groups():
+        try:
+            out, status = _dispatch(
+                group,
+                force_direct=force_direct,
+                poison_buf=poison_buf,
+                init_value=0,
+                unguarded_gap_read=unguarded_gap_read,
+            )
+            mx.eval(out, status)
+        except Exception as exc:  # compile error, pipeline ceiling below 512, driver refusal
+            raise DFloatBackendError(
+                f"the Metal decode kernel cannot run here ({group.name}): {exc}"
+            ) from exc
+        words = np.array(status)
+        if not np.array_equal(np.array(out), expected) or np.any(words & 7):
+            raise DFloatBackendError(
+                f"the Metal decode kernel produced wrong bits on the warm-up group {group.name}"
+            )
+        if np.any((words & 8) != (8 if force_direct else 0)):
+            want = "direct" if force_direct else "staged"
+            raise DFloatBackendError(
+                f"the Metal decode kernel took the wrong write path on the warm-up group "
+                f"{group.name} (want {want})"
+            )
+    _PIPELINES[key] = True
+
+
+def metal_ready() -> bool:
+    """Whether both production instantiations (direct and staged) compile here and decode the warm-up groups.
+
+    True means each instantiation compiled a pipeline for every input-binding signature a real
+    group can present and decoded each warm-up group bit-exactly on its expected write path (see
+    `ensure_pipeline`), so no later decode meets a first-time compile.
+    """
+    try:
+        ensure_pipeline(force_direct=True)
+        ensure_pipeline(force_direct=False)
+    except DFloatBackendError:
+        return False
+    return True
+
+
+def decode(
+    group: MxGroup,
+    *,
+    force_direct: bool = False,
+    _init_value: int | None = None,
+    _poison_buf: bool = False,
+    _unguarded_gap_read: bool = False,
+) -> DecodeResult:
+    """Lazy Metal decode of one group.
+
+    `_init_value` and `_poison_buf` are test-only poisons for the write-once tests: they prefill
+    the output and the staging buffer so an unwritten element cannot pass as a correct one.
+    `_unguarded_gap_read` is a test-only mutant for the shader-validation test: it drops the
+    guard on the second `gaps` byte, so the last thread of a full group reads one byte past
+    `gaps` (the value is shifted out, so the bits stay correct).
+
+    Raises:
+        DFloatBackendError: The kernel instantiation cannot be warmed up here (see
+            `ensure_pipeline`).
+    """
+    ensure_pipeline(
+        force_direct=force_direct, poison_buf=_poison_buf, unguarded_gap_read=_unguarded_gap_read
+    )
+    out, status = _dispatch(
+        group,
+        force_direct=force_direct,
+        poison_buf=_poison_buf,
+        init_value=_init_value,
+        unguarded_gap_read=_unguarded_gap_read,
+    )
+    direct = group.n_launch if force_direct else int(np.count_nonzero(group.intervals > CAP))
+    return DecodeResult(
+        bits=out,
+        status=status,
+        backend="metal",
+        direct_blocks=direct,
+        threadgroup_bytes=THREADGROUP_BYTES_DIRECT if force_direct else THREADGROUP_BYTES_STAGED,
+    )