Merge branch 'dev' of https://github.com/ostris/ai-toolkit into dev

manager/util.py

@@ -13,6 +13,11 @@ IS_WINDOWS = platform.system() == "Windows"
 IS_MAC = platform.system() == "Darwin"
 IS_LINUX = platform.system() == "Linux"
 
+USER_AGENT = (
+    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
+    "(KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36"
+)
+
 # When --json is used, human output goes to stderr so stdout stays machine-readable
 _json_mode = False
 
@@ -180,24 +185,38 @@ def clean_env(extra=None):
 
 
 def download(url, dest, label=None):
-    """Download url to dest (stdlib only), logging progress every ~10%."""
+    """Download url to dest (stdlib only), logging progress every ~10%.
+
+    Sends a browser User-Agent: some mirrors (ffmpeg.martin-riedl.de) return
+    403 for the default "Python-urllib/x.y" agent.
+    """
     import urllib.request
 
     label = label or os.path.basename(dest)
     info("Downloading %s ..." % label)
     last = [-1]
 
-    def hook(blocks, block_size, total):
+    def report(read, total):
         if total <= 0:
             return
-        pct = min(100, int(blocks * block_size * 100 / total))
+        pct = min(100, int(read * 100 / total))
         if pct >= last[0] + 10:
             last[0] = pct
             info("  %s: %d%%" % (label, pct))
 
+    req = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
     tmp = dest + ".part"
     try:
-        urllib.request.urlretrieve(url, tmp, reporthook=hook)
+        with urllib.request.urlopen(req) as resp, open(tmp, "wb") as out:
+            total = int(resp.headers.get("Content-Length") or 0)
+            read = 0
+            while True:
+                chunk = resp.read(1 << 16)
+                if not chunk:
+                    break
+                out.write(chunk)
+                read += len(chunk)
+                report(read, total)
     except Exception as e:  # noqa: BLE001 - surface any network failure clearly
         if os.path.exists(tmp):
             os.remove(tmp)

scripts/test_quantizations.py

@@ -17,15 +17,20 @@ Measures, per qtype:
     accumulated through the block stack
   - weight reconstruction error and one-time quantize (conversion) time
 
+Runs on CUDA or on Apple Silicon (MPS). See DEVICE NOTES below for what MPS
+can and cannot measure, and which qtypes it cannot run at all.
+
 Usage:
     python scripts/test_quantizations.py --gpu 1
     python scripts/test_quantizations.py --gpu 1 --qtypes bf16 convrot8
+    python scripts/test_quantizations.py --device mps
 """
 
 import argparse
 import math
 import os
 import sys
+import threading
 import time
 
 # set cuda bus ordering to be pcie
@@ -35,6 +40,124 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
 import torch  # noqa: E402
 
+# ---------------------------------------------------------------- DEVICE NOTES
+#
+# MPS differs from CUDA in three ways that matter to a benchmark, so the numbers
+# it prints are not interchangeable with CUDA's:
+#
+# 1. No peak-memory API. torch.mps has current_allocated_memory() but no
+#    max_memory_allocated()/reset_peak_memory_stats(), so the vram peaks are
+#    SAMPLED by a side thread polling the allocator (see _MpsPeakSampler)
+#    instead of read exactly. Steady-state values (resident weights) are exact.
+# 2. No fp8 dtype. torch.float8_e4m3fn is undefined on MPS, which rules out the
+#    qfloat8 qtype and convrot4 (its nvfp4 block scales are stored as e4m3).
+#    MPS_UNSUPPORTED lists them; they are dropped from a default run.
+# 3. No int8/fp4 tensor cores, so the convrot backends run their W8A16 fallback
+#    path rather than the W8A8/W4A4 fast path. Latency here says what Apple
+#    hardware does, not what the format is worth on a GPU that can run it.
+
+
+def pick_device(name: str, gpu: int) -> torch.device:
+    if name == "auto":
+        if torch.cuda.is_available():
+            name = "cuda"
+        elif torch.backends.mps.is_available():
+            name = "mps"
+        else:
+            raise SystemExit("no cuda or mps device available")
+    if name == "cuda":
+        if not torch.cuda.is_available():
+            raise SystemExit("--device cuda requested but cuda is not available")
+        device = torch.device(f"cuda:{gpu}")
+        torch.cuda.set_device(device)
+        return device
+    if name == "mps":
+        if not torch.backends.mps.is_available():
+            raise SystemExit("--device mps requested but mps is not available")
+        return torch.device("mps")
+    raise SystemExit(f"unsupported device {name!r}")
+
+
+def describe_device(device: torch.device) -> str:
+    if device.type == "cuda":
+        p = torch.cuda.get_device_properties(device)
+        return (f"{device} ({p.name}, sm_{p.major}{p.minor}, "
+                f"{p.total_memory / 1e9:.0f} GB)")
+    import platform
+    return (f"{device} (Apple {platform.machine()}, "
+            f"{torch.mps.recommended_max_memory() / 1e9:.0f} GB recommended max)")
+
+
+def sync(device: torch.device) -> None:
+    if device.type == "cuda":
+        torch.cuda.synchronize(device)
+    elif device.type == "mps":
+        torch.mps.synchronize()
+
+
+def empty_cache(device: torch.device) -> None:
+    if device.type == "cuda":
+        torch.cuda.empty_cache()
+    elif device.type == "mps":
+        torch.mps.empty_cache()
+
+
+def mem_allocated(device: torch.device) -> int:
+    if device.type == "cuda":
+        return torch.cuda.memory_allocated(device)
+    if device.type == "mps":
+        return torch.mps.current_allocated_memory()
+    return 0
+
+
+class _MpsPeakSampler:
+    """Approximate max_memory_allocated for MPS by polling the allocator.
+
+    Tensors are allocated on the calling thread as ops are enqueued, so a
+    fine-grained poll from a side thread does observe the transients; on the
+    probe cases it recovered exact expected sizes (a 512 MiB transient and a
+    24 MiB matmul output). It can still miss a transient shorter than the poll
+    interval, so treat MPS peaks as a lower bound, not a hard number.
+    """
+
+    def __init__(self, device, interval=0.0002):
+        self.device, self.interval = device, interval
+        self.peak = 0
+        self._stop = threading.Event()
+
+    def __enter__(self):
+        self.peak = mem_allocated(self.device)
+        self._thread = threading.Thread(target=self._poll, daemon=True)
+        self._thread.start()
+        return self
+
+    def _poll(self):
+        while not self._stop.is_set():
+            self.peak = max(self.peak, mem_allocated(self.device))
+            time.sleep(self.interval)
+
+    def __exit__(self, *exc):
+        sync(self.device)
+        self.peak = max(self.peak, mem_allocated(self.device))
+        self._stop.set()
+        self._thread.join()
+        return False
+
+
+def measure_peak(run, device: torch.device, base: int) -> int:
+    """Peak bytes allocated during run(), over base. run() is called once first
+    so lazy init and compilation are not charged to the steady-state peak."""
+    run()
+    sync(device)
+    if device.type == "cuda":
+        torch.cuda.reset_peak_memory_stats(device)
+        run()
+        sync(device)
+        return torch.cuda.max_memory_allocated(device) - base
+    with _MpsPeakSampler(device) as sampler:
+        run()
+    return sampler.peak - base
+
 # (tokens, in_features, out_features) — FLUX/Wan-class projections
 SPEED_SHAPES = [
     (4096, 3072, 3072),
@@ -54,6 +177,12 @@ QTYPES = [
     "convrotint2", "convrotbitnet", "convrotcomfyw4a4",
 ]
 
+# qtypes that cannot run on MPS at all (see DEVICE NOTES: no fp8 dtype).
+# quanto's float8 path is worse than a hard failure — it catches the dtype error,
+# prints "Failed to quantize", and leaves the layer in bf16, so it would benchmark
+# as bf16 under a quantized label. Dropped for the same reason.
+MPS_UNSUPPORTED = {"qfloat8", "float8", "convrot4"}
+
 STACK_KEY = f"{VRAM_BLOCKS}-block stack"
 
 
@@ -93,11 +222,11 @@ def fp_weight(module: torch.nn.Linear) -> torch.Tensor:
 def bench(fn, iters: int, device) -> float:
     for _ in range(max(3, iters // 5)):
         fn()
-    torch.cuda.synchronize(device)
+    sync(device)
     t0 = time.perf_counter()
     for _ in range(iters):
         fn()
-    torch.cuda.synchronize(device)
+    sync(device)
     return (time.perf_counter() - t0) / iters * 1000  # ms
 
 
@@ -195,20 +324,15 @@ def run_speed(qtype: str, device, iters: int, results: dict):
             )
         except Exception as e:
             print(f"  [{qtype}] compiled train failed for {m}x{k}->{n}: {e}")
-        torch.cuda.empty_cache()
+        empty_cache(device)
 
 
 def _stack_fwd_peak(blocks, x, device, base) -> int:
-    # warm up first so lazy-init allocations (and compilation) are not counted
-    # as steady-state peak
-    with torch.no_grad():
-        stack_forward(blocks, x)
-    torch.cuda.synchronize(device)
-    torch.cuda.reset_peak_memory_stats(device)
-    with torch.no_grad():
-        stack_forward(blocks, x)
-    torch.cuda.synchronize(device)
-    return torch.cuda.max_memory_allocated(device) - base
+    def fwd():
+        with torch.no_grad():
+            stack_forward(blocks, x)
+
+    return measure_peak(fwd, device, base)
 
 
 def _stack_train_peak(blocks, x, device, base, checkpoint) -> int:
@@ -217,24 +341,19 @@ def _stack_train_peak(blocks, x, device, base, checkpoint) -> int:
         xi = x.detach().requires_grad_(True)
         stack_forward(blocks, xi, checkpoint).float().pow(2).mean().backward()
 
-    train_step()
-    torch.cuda.synchronize(device)
-    torch.cuda.reset_peak_memory_stats(device)
-    train_step()
-    torch.cuda.synchronize(device)
-    return torch.cuda.max_memory_allocated(device) - base
+    return measure_peak(train_step, device, base)
 
 
 def run_vram(qtype: str, device, results: dict):
-    torch.cuda.empty_cache()
-    base = torch.cuda.memory_allocated(device)
+    empty_cache(device)
+    base = mem_allocated(device)
 
     blocks = make_stack(device)
     for b in blocks:
         for i in range(len(b)):
             b[i] = convert(b[i], qtype)
-    torch.cuda.empty_cache()
-    results[(qtype, "vram_weights")] = torch.cuda.memory_allocated(device) - base
+    empty_cache(device)
+    results[(qtype, "vram_weights")] = mem_allocated(device) - base
 
     x = torch.randn(VRAM_TOKENS, 3072, device=device, dtype=torch.bfloat16)
 
@@ -256,7 +375,7 @@ def run_vram(qtype: str, device, results: dict):
         print(f"  [{qtype}] compiled vram measurement failed: {e}")
 
     blocks = x = None  # release before the allocator accounting of the next run
-    torch.cuda.empty_cache()
+    empty_cache(device)
 
 
 def run_drift(qtype: str, device, results: dict):
@@ -270,7 +389,7 @@ def run_drift(qtype: str, device, results: dict):
             lin = convert(lin, qtype)
             y_q = lin(x).float()
         results[(qtype, "drift", (m, k, n))] = ((y_q - y_ref).norm() / y_ref.norm()).item()
-        torch.cuda.empty_cache()
+        empty_cache(device)
 
     blocks = make_stack(device)
     x = torch.randn(VRAM_TOKENS, 3072, device=device, dtype=torch.bfloat16)
@@ -282,24 +401,24 @@ def run_drift(qtype: str, device, results: dict):
         y_q = stack_forward(blocks, x).float()
     results[(qtype, "drift", STACK_KEY)] = ((y_q - y_ref).norm() / y_ref.norm()).item()
     blocks = x = None
-    torch.cuda.empty_cache()
+    empty_cache(device)
 
 
 def run_quality_and_quantize_time(qtype: str, device, results: dict):
     torch.manual_seed(0)
     lin = make_layer(3072, 3072, device)
     w0 = lin.weight.detach().float().clone()
-    torch.cuda.synchronize(device)
+    sync(device)
     t0 = time.perf_counter()
     lin = convert(lin, qtype)
-    torch.cuda.synchronize(device)
+    sync(device)
     results[(qtype, "quantize_ms")] = (time.perf_counter() - t0) * 1000
     if qtype == "bf16":
         results[(qtype, "weight_err")] = 0.0
     else:
         wq = fp_weight(lin)
         results[(qtype, "weight_err")] = ((wq - w0).norm() / w0.norm()).item()
-    torch.cuda.empty_cache()
+    empty_cache(device)
 
 
 def print_speed_table(title: str, kind: str, qts, results):
@@ -325,23 +444,37 @@ def print_speed_table(title: str, kind: str, qts, results):
 
 def main():
     ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument("--device", default="auto", choices=["auto", "cuda", "mps"],
+                    help="accelerator to run on (default: cuda if present, else mps)")
     ap.add_argument("--gpu", type=int, default=0, help="cuda device id to run on")
-    ap.add_argument("--qtypes", nargs="+", default=QTYPES, help=f"subset of {QTYPES}")
+    ap.add_argument("--qtypes", nargs="+", default=None, help=f"subset of {QTYPES}")
     ap.add_argument("--iters", type=int, default=50, help="timing iterations per case")
     args = ap.parse_args()
 
-    device = torch.device(f"cuda:{args.gpu}")
-    torch.cuda.set_device(device)
-    props = torch.cuda.get_device_properties(device)
-    print(f"device: cuda:{args.gpu} ({props.name}, sm_{props.major}{props.minor}, "
-          f"{props.total_memory / 1e9:.0f} GB)")
-    print(f"torch {torch.__version__}\n")
+    device = pick_device(args.device, args.gpu)
+    print(f"device: {describe_device(device)}")
+    print(f"torch {torch.__version__}")
+
+    explicit_qtypes = args.qtypes is not None
+    qtypes = args.qtypes if explicit_qtypes else list(QTYPES)
+    if device.type == "mps":
+        blocked = [qt for qt in qtypes if qt in MPS_UNSUPPORTED]
+        if blocked and not explicit_qtypes:
+            qtypes = [qt for qt in qtypes if qt not in MPS_UNSUPPORTED]
+            print(f"note: skipping {', '.join(blocked)} — no fp8 dtype on MPS")
+        elif blocked:
+            # asked for by name: try anyway, but say what is expected to happen
+            print(f"note: {', '.join(blocked)} need an fp8 dtype MPS does not have; "
+                  "expect them to fail or to silently stay in bf16")
+        print("note: MPS vram peaks are sampled, not exact; convrot backends run "
+              "their W8A16 fallback (no int8/fp4 tensor cores)")
+    print()
 
     # warm the toolkit import chain (module imports + custom-op registration) so it
     # isn't charged to the first qtype's quantize timing
-    if any(qt != "bf16" for qt in args.qtypes):
+    if any(qt != "bf16" for qt in qtypes):
         from toolkit.util.ostris_quant import get_ostris_quantizer
-        for qt in args.qtypes:
+        for qt in qtypes:
             if qt != "bf16":
                 get_ostris_quantizer(qt)
 
@@ -351,15 +484,36 @@ def main():
     torch._dynamo.config.cache_size_limit = 4096
 
     results = {}
-    for qt in args.qtypes:
+    ran = []
+    for qt in qtypes:
         print(f"benchmarking {qt} ...")
         torch._dynamo.reset()  # drop the previous qtype's compiled artifacts
-        run_quality_and_quantize_time(qt, device, results)
-        run_drift(qt, device, results)
-        run_speed(qt, device, args.iters, results)
-        run_vram(qt, device, results)
+        try:
+            run_quality_and_quantize_time(qt, device, results)
+            run_drift(qt, device, results)
+            run_speed(qt, device, args.iters, results)
+            run_vram(qt, device, results)
+        except Exception as e:
+            # one unsupported backend should not cost the whole run
+            print(f"  [{qt}] FAILED, dropped from the tables: {type(e).__name__}: {e}")
+            for key in list(results):
+                if key[0] == qt:
+                    del results[key]
+            empty_cache(device)
+            continue
+        if qt != "bf16" and results.get((qt, "weight_err")) == 0.0:
+            print(f"  [{qt}] quantization was a no-op (weights unchanged) — "
+                  "dropped so it is not reported as a quantized result")
+            for key in list(results):
+                if key[0] == qt:
+                    del results[key]
+            continue
+        ran.append(qt)
 
-    qts = args.qtypes
+    if not ran:
+        raise SystemExit("no qtype completed successfully")
+
+    qts = ran
     print_speed_table("layer latency, inference", "inf", qts, results)
     print_speed_table("layer latency, inference (compiled)", "inf_comp", qts, results)
     print_speed_table("layer latency, train fwd+bwd", "train", qts, results)

toolkit/config_modules.py

@@ -711,11 +711,15 @@ class ModelConfig:
         if self.layer_offloading and self.qtype_te == "qfloat8":
             self.qtype_te = "float8"
             
-        # Mac mps only works with torachao uint
+        # MPS has no fp8 dtype, so qfloat8 has to become an 8 bit integer format.
+        # convrot8, not torchao int8: measured on an M3 against bf16, convrot8
+        # trains at 0.79x and holds 1.04 GB of resident weight where torchao int8
+        # trains at 0.52x and holds 1.21 GB, and convrot8 quantizes in 19ms
+        # against 2.8s. See scripts/test_quantizations.py --device mps.
         if torch.backends.mps.is_available() and self.qtype == "qfloat8":
-            self.qtype = "int8"
+            self.qtype = "convrot8"
         if torch.backends.mps.is_available() and self.qtype_te == "qfloat8":
-            self.qtype_te = "int8"
+            self.qtype_te = "convrot8"
         
         # 0 is off and 1.0 is 100% of the layers
         self.layer_offloading_transformer_percent = kwargs.get("layer_offloading_transformer_percent", 1.0)

toolkit/util/convrot_quant.py

@@ -95,9 +95,17 @@ def regular_hadamard(rot_size: int, device, dtype=torch.bfloat16) -> torch.Tenso
     key = (rot_size, str(device), dtype)
 
     def build():
+        # fp32, not fp64: entries stay exactly +-1 through the krons and rot_size
+        # is a power of 4, so dividing by its (power-of-two) root is exact — this
+        # build is bit-identical to an fp64 one at every supported rot_size, and
+        # verified so. fp64 is not merely unnecessary here but harmful: MPS has no
+        # float64, and on a cache miss inside a torch.compile trace inductor lifts
+        # the matrix in as a graph constant and dies moving it to the device
+        # ("Cannot convert a MPS Tensor to float64").
         r4 = torch.tensor(
             [[1.0, 1, 1, -1], [1, 1, -1, 1], [1, -1, 1, 1], [-1, 1, 1, 1]],
-            dtype=torch.float64,
+            dtype=torch.float32,
+            device="cpu",
         )
         h = r4.clone()
         while h.shape[0] < rot_size:
@@ -1040,19 +1048,223 @@ def _int8_gemm_supported(device) -> bool:
     global _warned_no_int8
     if not supported and not _warned_no_int8:
         _warned_no_int8 = True
-        print_acc(
-            f"ConvRot: int8 matmul (torch._int_mm) is not usable on this device "
-            f"({device}). Inference falls back to dequantized bf16 matmuls: correct "
-            "output but NO speedup, and inference activations stay unquantized "
-            "(W8A16 numerics instead of W8A8). The training path is unaffected "
-            "(it always simulates W8A8 via fake-quant)."
-        )
+        if _mps_int8pack_device_ok(device):
+            print_acc(
+                f"ConvRot: int8 tensor-core matmul (torch._int_mm) is not implemented "
+                f"on this device ({device}). Inference uses torch._weight_int8pack_mm "
+                "instead — the int8 weight feeds the gemm directly, with no bf16 "
+                "dequant transient, which is the whole win on a bandwidth-bound Apple "
+                "GPU. Activations stay unquantized (W8A16 numerics instead of W8A8); "
+                "Apple GPUs have no int8 tensor cores, so quantizing them would only "
+                "add overhead. Layers whose in/out features are not both divisible by "
+                "32 still use a dequantized bf16 matmul. The training path is "
+                "unaffected (it always simulates W8A8 via fake-quant)."
+            )
+        else:
+            print_acc(
+                f"ConvRot: int8 matmul (torch._int_mm) is not usable on this device "
+                f"({device}). Inference falls back to dequantized bf16 matmuls: correct "
+                "output but NO speedup, and inference activations stay unquantized "
+                "(W8A16 numerics instead of W8A8). The training path is unaffected "
+                "(it always simulates W8A8 via fake-quant)."
+            )
     return supported
 
 
 _warned_no_int8 = False
 
 
+# ---------------- mps int8-weight gemm ----------------
+#
+# Apple GPUs have no int8 tensor cores and torch._int_mm has no MPS kernel, so
+# there is no W8A8 path here and no reason to want one: a quantized activation
+# would buy nothing back on hardware that multiplies in fp anyway. What DOES pay
+# on a unified-memory Apple GPU is not materializing the dequantized weight —
+# torch._weight_int8pack_mm (int8 (out, in) weight + per-row scales, exactly the
+# layout convrot8/convrotint already store) feeds the codes to the gemm directly.
+#
+# Measured on an M3, convrot8 3072x3072, vs the dequant+F.linear fallback it
+# replaces: 18x at M=1, 5.7x at M=64, 1.2x at M=1024, ~1.0x at M=4096. The
+# fallback's dequant is a fixed ~2.7ms that amortizes away as M grows, so the win
+# concentrates at decode/small-batch shapes and large batches are a wash. The
+# bitpacked convrotint*/bitnet backends see only 1.1-1.4x — their per-forward
+# weight unpack dominates on MPS (no triton fused gemv path there). Error vs the
+# unquantized bf16 reference is marginally LOWER than the fallback's. Holding no
+# dequantized weight transient is a win at every M on unified memory.
+
+_int8pack_mm_ok = None
+
+
+def _mps_int8pack_device_ok(device) -> bool:
+    """Whether torch._weight_int8pack_mm has a working kernel on this device."""
+    global _int8pack_mm_ok
+    if torch.device(device).type != "mps":
+        return False
+    if _int8pack_mm_ok is None:
+        try:
+            torch._weight_int8pack_mm(
+                torch.zeros(1, 32, dtype=torch.bfloat16, device=device),
+                torch.zeros(32, 32, dtype=torch.int8, device=device),
+                torch.ones(32, dtype=torch.bfloat16, device=device),
+            )
+            _int8pack_mm_ok = True
+        except Exception:
+            _int8pack_mm_ok = False
+    return _int8pack_mm_ok
+
+
+def _mps_int8pack_usable(device, in_f: int, out_f: int, dtype: torch.dtype) -> bool:
+    # the kernel asserts N % 32 == 0 && K % 32 == 0; can_quantize only guarantees
+    # in % 16 and out % 8, so the odd layer still needs the dequant fallback
+    return (
+        in_f % 32 == 0
+        and out_f % 32 == 0
+        and dtype in (torch.bfloat16, torch.float16, torch.float32)
+        and _mps_int8pack_device_ok(device)
+    )
+
+
+# registered as a custom op so torch.compile treats it as an opaque node with a
+# known output shape (see _nvfp4_act_quant_op). Calling the aten op directly
+# compiles fine on cuda but inductor's MPS lowering asserts on it as soon as its
+# input is an unrealized Pointwise — which is exactly what the rotate() feeding
+# it produces, so every compiled convrot8 inference on mac fell back to eager.
+@torch.library.custom_op("ostris::convrot_mps_int8pack_mm", mutates_args=())
+def _mps_int8pack_mm_op(
+    x2d: torch.Tensor, qdata: torch.Tensor, scales: torch.Tensor
+) -> torch.Tensor:
+    # scales is cast to x2d's dtype deliberately: the kernel reads the scale
+    # buffer AS the activation dtype with no check, so handing it the fp32 scales
+    # convrot stores returns silent garbage (nan / zeros) rather than an error —
+    # the one sharp edge here. It also rejects non-contiguous operands outright;
+    # both are contiguous today, so those calls are free no-ops that keep this
+    # correct if a caller's layout ever changes.
+    return torch._weight_int8pack_mm(
+        x2d.contiguous(), qdata.contiguous(), scales.to(x2d.dtype)
+    )
+
+
+@_mps_int8pack_mm_op.register_fake
+def _mps_int8pack_mm_fake(x2d, qdata, scales):
+    return torch.empty(
+        x2d.shape[0], qdata.shape[0], device=x2d.device, dtype=x2d.dtype
+    )
+
+
+def _mps_int8pack_linear(x2d, qdata, scales, bias) -> torch.Tensor:
+    """out = x2d @ dequant(qdata).T (+ bias) via the MPS int8-weight gemm."""
+    out = _mps_int8pack_mm_op(x2d, qdata, scales)
+    if bias is not None:
+        out = out + bias
+    return out
+
+
+# ---------------- mps: fused elementwise chains ----------------
+#
+# On MPS the gemm runs in bf16 whatever the weight storage is, so there is no
+# matmul-level win to chase: torch._weight_int8pack_mm measures within 2% of an
+# equal-shape bf16 F.linear for M >= 256 (it only pulls ahead at decode-size
+# batches — 2.1x at M=1, 1.2x at M=64, measured on an M3 at 3072x3072). What
+# separates convrot8 from plain bf16 at training shapes is therefore ENTIRELY the
+# elementwise work wrapped around the gemm: the activation fake-quant and the
+# weight dequant.
+#
+# Eager runs each of those as a chain of ~8 kernels over full-size fp32
+# temporaries, which on unified memory costs more than the quantization saves —
+# the activation fake-quant alone measures 9.1ms against a 25.8ms bf16 gemm (35%
+# overhead for one of two chains). Inductor's MPS backend fuses each chain into a
+# single Metal kernel and takes that 9.1ms to 1.0ms, which is what moves convrot8
+# training from 0.69x of bf16 to ~0.99x.
+#
+# The two chains want OPPOSITE compile modes, both for measured reasons — hence
+# the per-chain `dynamic` below rather than one global setting:
+#
+#   activations (_fake_quant_rows), dynamic=True: the batch dim genuinely varies,
+#     and a static compile lowers the fp32 divide to a reciprocal-multiply, which
+#     shifts ~0.002% of codes by 1. (Towards an fp64 reference, as it happens —
+#     but it is still a silent numerics change, and bit-identity with the cuda
+#     path is this backend's whole contract.) dynamic=True keeps the divide and
+#     is bit-identical, one compile covers every shape, and it costs ~0.4ms on a
+#     52ms step against the static build.
+#
+#   weights (_dequant_rows), dynamic=False: here dynamic shapes destroy the win
+#     outright — 2.5ms vs 0.38ms static, against 2.8ms eager, i.e. no better than
+#     not compiling at all (inductor cannot vectorize the broadcast multiply
+#     without a static inner width). All modes are bit-identical for this chain,
+#     so static is free of the numerics objection. Weight shapes are also fixed
+#     per layer, so this only recompiles once per distinct layer shape.
+#
+# The static build does hit dynamo's recompile limit (8) on a model with more
+# than 8 distinct linear shapes. That degrades gracefully and was measured, not
+# assumed: the first 8 shapes keep their compiled kernels and everything past
+# them runs the eager chain, which is exactly today's behaviour. Nothing silently
+# breaks, the tail just stops getting faster.
+
+_mps_fused_cache = {}
+_mps_fuse_disabled = False
+
+
+def _mps_fused(fn, dynamic):
+    """Run `fn` through inductor, falling back to eager for good on error.
+
+    Callers reach these only on mps — cuda has the triton kernels above and keeps
+    its existing code path untouched. Compilation is also skipped while tracing,
+    so an outer torch.compile fuses the chain into its own graph rather than
+    meeting an opaque nested compile.
+    """
+
+    def run(*args):
+        global _mps_fuse_disabled
+        if _mps_fuse_disabled or torch.compiler.is_compiling():
+            return fn(*args)
+        compiled = _mps_fused_cache.get(fn)
+        if compiled is None:
+            compiled = torch.compile(fn, dynamic=dynamic)
+            _mps_fused_cache[fn] = compiled
+        try:
+            return compiled(*args)
+        except Exception as e:
+            _mps_fuse_disabled = True
+            print_acc(
+                f"ConvRot: inductor fusion is unavailable on this mps build "
+                f"({type(e).__name__}: {e}) — falling back to the eager "
+                "elementwise path. Output is unchanged; expect roughly 0.7x bf16 "
+                "training speed instead of ~1.0x."
+            )
+            return fn(*args)
+
+    return run
+
+
+def _fake_quant_rows_impl(x: torch.Tensor, qmax: int) -> torch.Tensor:
+    """dequant(quant(x)) per row, returned in x's dtype — the value half of the
+    activation STE.
+
+    Bit-identical to quantize_int8_rows followed by the dequant it feeds, with
+    both full-size fp32 temporaries dropped: amax over the input equals amax over
+    its fp32 upcast (bf16 -> fp32 is exact and a max is a selection, not an
+    accumulation), and the int8 round trip is lossless on a value already rounded
+    and clamped into [-qmax, qmax].
+    """
+    scales = x.abs().amax(dim=1).float() / qmax
+    scales = torch.where(scales > 0, scales, torch.ones_like(scales))
+    s = scales.unsqueeze(1)
+    return (torch.round(x.float() / s).clamp_(-qmax, qmax) * s).to(x.dtype)
+
+
+_fake_quant_rows = _mps_fused(_fake_quant_rows_impl, dynamic=True)
+
+
+def _dequant_rows_impl(
+    qdata: torch.Tensor, scales: torch.Tensor, dtype: torch.dtype
+) -> torch.Tensor:
+    """dequant of per-row-scaled integer codes: the weight-side chain."""
+    return (qdata.float() * scales.unsqueeze(1)).to(dtype)
+
+
+_dequant_rows = _mps_fused(_dequant_rows_impl, dynamic=False)
+
+
 class ConvRotInt8Quantizer(OstrisQuantizer):
     """ConvRot W8A8 backend: shared regular-Hadamard rotation + per-token /
     per-output-channel symmetric int8 with torch._int_mm. One instance per qtype,
@@ -1124,7 +1336,10 @@ class ConvRotInt8Quantizer(OstrisQuantizer):
         return torch.round(w_rot.float() / s).clamp_(-127, 127) * s
 
     def _dequantize_rotated(self, module, dtype: torch.dtype) -> torch.Tensor:
-        w = self._qdata(module).float() * self._scales(module).unsqueeze(1)
+        qdata, scales = self._qdata(module), self._scales(module)
+        if qdata.device.type == "mps":
+            return _dequant_rows(qdata, scales, dtype)
+        w = qdata.float() * scales.unsqueeze(1)
         return w.to(dtype)
 
     def dequantize(self, module) -> torch.Tensor:
@@ -1176,8 +1391,13 @@ class ConvRotInt8Quantizer(OstrisQuantizer):
             # no int8 hardware: straight-through fake-quant + bf16 matmul
             x2d = rotate(x, rot).reshape(-1, in_f)
             with torch.no_grad():
-                aq, a_s = quantize_int8_rows(x2d.detach(), self.act_qmax)
-                x_dq = (aq.float() * a_s.unsqueeze(1)).to(x.dtype)
+                if x.device.type == "mps":
+                    # same arithmetic, as one fused Metal kernel instead of a
+                    # chain over full-size fp32 temporaries (see _mps_fused)
+                    x_dq = _fake_quant_rows(x2d.detach(), self.act_qmax)
+                else:
+                    aq, a_s = quantize_int8_rows(x2d.detach(), self.act_qmax)
+                    x_dq = (aq.float() * a_s.unsqueeze(1)).to(x.dtype)
                 w = self._dequantize_rotated(module, x.dtype)
             x_ste = x2d + (x_dq - x2d).detach()
             out = F.linear(x_ste, w, module.bias)
@@ -1214,6 +1434,16 @@ class ConvRotInt8Quantizer(OstrisQuantizer):
             )
             return out.reshape(*x.shape[:-1], out_f)
 
+        if _mps_int8pack_usable(x.device, in_f, out_f, x.dtype):
+            # W8A16 on the int8 weight codes, no dequantized weight transient
+            out = _mps_int8pack_linear(
+                rotate(x, rot).reshape(-1, in_f),
+                self._qdata(module),
+                self._scales(module),
+                module.bias,
+            )
+            return out.reshape(*x.shape[:-1], out_f)
+
         w = self._dequantize_rotated(module, x.dtype)
         out = F.linear(rotate(x, rot).reshape(-1, in_f), w, module.bias)
         return out.reshape(*x.shape[:-1], out_f)

ui/src/app/api/cpu/route.ts

@@ -1,9 +1,9 @@
 import { NextResponse } from 'next/server';
 import si from 'systeminformation';
-import { createRequire } from 'module';
 import os from 'os';
 import { CpuInfo } from '@/types';
 import { cached } from '@/server/apiCache';
+import { loadMacstats } from '@/server/macstats';
 
 const isMac = os.platform() === 'darwin';
 
@@ -13,8 +13,8 @@ async function getCpuInfo(): Promise<CpuInfo> {
 
   if (isMac) {
     try {
-      const nativeRequire = createRequire(import.meta.url);
-      const ms = nativeRequire('macstats') as any;
+      const ms = loadMacstats();
+      if (!ms) throw new Error('macstats unavailable');
       const ramData = ms.getRAMUsageSync();
       const cpuData = ms.getCpuDataSync();
 

ui/src/app/api/gpu/route.ts

@@ -1,9 +1,9 @@
 import { NextResponse } from 'next/server';
 import { exec } from 'child_process';
 import { promisify } from 'util';
-import { createRequire } from 'module';
 import os from 'os';
 import { cached } from '@/server/apiCache';
+import { loadMacstats } from '@/server/macstats';
 
 const execAsync = promisify(exec);
 
@@ -47,11 +47,8 @@ async function getMacGpuInfo(): Promise<MacGpuResult | null> {
     let memUsed = 0;
     let memTotal = memoryTotal;
 
-    try {
-      // Use createRequire to hide from webpack static analysis so it doesn't fail on non-mac platforms
-      const nativeRequire = createRequire(import.meta.url);
-      const ms = nativeRequire('macstats') as any;
-
+    const ms = loadMacstats();
+    if (ms) {
       try {
         const gpuData = ms.getGpuDataSync();
         temperature = gpuData.temperature || 0;
@@ -84,8 +81,6 @@ async function getMacGpuInfo(): Promise<MacGpuResult | null> {
       } catch {
         // ignore
       }
-    } catch (error) {
-      console.warn('macstats not available:', error);
     }
 
     return { name: gpuName, memUsed, memTotal, gpuLoad, temperature, fanSpeed, powerDraw };

ui/src/server/macstats.ts

@@ -0,0 +1,38 @@
+import * as nodeModule from 'module';
+import os from 'os';
+import path from 'path';
+
+/**
+ * macstats is a native, macOS-only optional dependency.
+ *
+ * It must not be bundled. Webpack special-cases `createRequire(...)`: with a literal base it
+ * resolves and inlines the module (a hard MODULE_NOT_FOUND throw when it can't), and with a
+ * computed base it replaces the whole call with `void 0` — both leave macstats unreachable at
+ * runtime even when it is installed, regardless of serverExternalPackages/externals. So look
+ * createRequire up dynamically, where the bundler can't recognize it, and base the require on
+ * process.cwd() (the ui/ directory the Next server runs in) so it resolves against
+ * ui/node_modules at runtime.
+ */
+
+// undefined = not tried yet, null = unavailable on this machine
+let cachedModule: any | null | undefined;
+
+export function loadMacstats(): any | null {
+  if (cachedModule !== undefined) return cachedModule;
+
+  if (os.platform() !== 'darwin') {
+    cachedModule = null;
+    return cachedModule;
+  }
+
+  try {
+    const createRequire = (nodeModule as any)['create' + 'Require'] as typeof nodeModule.createRequire;
+    const nativeRequire = createRequire(path.join(process.cwd(), 'package.json'));
+    cachedModule = nativeRequire('macstats');
+  } catch (error) {
+    console.warn('macstats not available:', error);
+    cachedModule = null;
+  }
+
+  return cachedModule;
+}