diff --git a/scripts/test_quantizations.py b/scripts/test_quantizations.py index 445edaa..a997cd1 100644 --- a/scripts/test_quantizations.py +++ b/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) diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index a0ef971..0726f52 100644 --- a/toolkit/config_modules.py +++ b/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) diff --git a/toolkit/util/convrot_quant.py b/toolkit/util/convrot_quant.py index 9293507..8b7234c 100644 --- a/toolkit/util/convrot_quant.py +++ b/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)