WaveCut/orbitquant-packed-matmul

1from __future__ import annotations2 3import argparse4import json5import os6import platform7import statistics8import time9 10import torch11from orbitquant_packed_matmul import matmul_packed_weight12 13 14def _pack(values: torch.Tensor, bits: int) -> torch.Tensor:15    flat = values.detach().to(device="cpu", dtype=torch.uint8).flatten()16    packed = torch.zeros((flat.numel() * bits + 7) // 8, dtype=torch.uint8)17    for value_index, value in enumerate(flat.tolist()):18        bit_start = value_index * bits19        byte_index = bit_start // 820        shift = bit_start % 821        packed[byte_index] |= (value << shift) & 0xFF22        if shift + bits > 8:23            packed[byte_index + 1] |= value >> (8 - shift)24    return packed25 26 27def _synchronize(device: str) -> None:28    if device == "cuda":29        torch.cuda.synchronize()30    elif device == "mps":31        torch.mps.synchronize()32 33 34def _time_call(device: str, fn) -> float:35    _synchronize(device)36    start = time.perf_counter_ns()37    fn()38    _synchronize(device)39    return (time.perf_counter_ns() - start) / 1_000_000_00040 41 42def _time_distribution(device: str, iters: int, fn) -> dict[str, float]:43    samples = []44    for _ in range(iters):45        samples.append(_time_call(device, fn))46    samples.sort()47    return {48        "mean": statistics.fmean(samples),49        "median": statistics.median(samples),50        "p95": samples[min(len(samples) - 1, int(len(samples) * 0.95))],51    }52 53 54def _parse_rows(raw: str) -> list[int]:55    rows_values = []56    for chunk in raw.split(","):57        chunk = chunk.strip()58        if not chunk:59            continue60        rows = int(chunk)61        if rows <= 0:62            raise argparse.ArgumentTypeError("--rows values must be positive")63        rows_values.append(rows)64    if not rows_values:65        raise argparse.ArgumentTypeError("--rows must list at least one row count")66    return rows_values67 68 69_DTYPES = {70    "float32": torch.float32,71    "float16": torch.float16,72    "bfloat16": torch.bfloat16,73}74 75 76def _benchmark_rows(args, rows: int, dtype: torch.dtype, weights) -> dict:77    packed, row_norms, centroids, indices_device = weights78    x = torch.randn(rows, args.in_features, device=args.device, dtype=dtype)79    bias = (80        torch.randn(args.out_features, device=args.device, dtype=dtype)81        if args.with_bias82        else None83    )84 85    def materialize_reference_weight() -> torch.Tensor:86        return (row_norms[:, None] * centroids[indices_device]).to(dtype)87 88    reference_weight = materialize_reference_weight()89    packed_weight_indices_bytes = packed.numel() * packed.element_size()90    row_norms_bytes = row_norms.numel() * row_norms.element_size()91    centroid_bytes = centroids.numel() * centroids.element_size()92    packed_weight_path_bytes = packed_weight_indices_bytes + row_norms_bytes + centroid_bytes93    materialized_weight_bytes = reference_weight.numel() * reference_weight.element_size()94 95    def packed_call() -> torch.Tensor:96        return matmul_packed_weight(97            x,98            packed,99            row_norms,100            centroids,101            bits=args.bits,102            out_features=args.out_features,103            in_features=args.in_features,104            bias=bias,105        )106 107    def predequantized_linear_call() -> torch.Tensor:108        return torch.nn.functional.linear(x, reference_weight, bias)109 110    def dequantize_then_linear_call() -> torch.Tensor:111        return torch.nn.functional.linear(x, materialize_reference_weight(), bias)112 113    packed_first_call_seconds = _time_call(args.device, packed_call)114    predequantized_first_call_seconds = _time_call(args.device, predequantized_linear_call)115    dequantize_then_first_call_seconds = _time_call(args.device, dequantize_then_linear_call)116 117    for _ in range(args.warmup):118        packed_call()119        predequantized_linear_call()120        dequantize_then_linear_call()121    packed_distribution = _time_distribution(args.device, args.iters, packed_call)122    predequantized_distribution = _time_distribution(123        args.device,124        args.iters,125        predequantized_linear_call,126    )127    dequantize_then_distribution = _time_distribution(128        args.device,129        args.iters,130        dequantize_then_linear_call,131    )132    # Headline numbers are hot-loop medians; the mean is retained alongside the133    # median/p95 so noisy first-iteration outliers cannot skew comparisons.134    packed_seconds = packed_distribution["median"]135    predequantized_linear_seconds = predequantized_distribution["median"]136    dequantize_then_linear_seconds = dequantize_then_distribution["median"]137 138    packed_output = packed_call()139    reference_output = predequantized_linear_call()140    _synchronize(args.device)141    error = packed_output.float() - reference_output.float()142    max_abs_error = error.abs().max().item()143    rmse = error.square().mean().sqrt().item()144    reference_rms = reference_output.float().square().mean().sqrt().item()145    relative_rmse = rmse / max(reference_rms, 1e-12)146 147    return {148        "device": args.device,149        "device_name": (150            torch.cuda.get_device_name(0)151            if args.device == "cuda"152            else "mps"153            if args.device == "mps"154            else f"{platform.processor() or platform.machine()} "155            f"({torch.backends.cpu.get_cpu_capability()})"156        ),157        "dtype": str(dtype).replace("torch.", ""),158        "bits": args.bits,159        "rows": rows,160        "in_features": args.in_features,161        "out_features": args.out_features,162        "iters": args.iters,163        "warmup": args.warmup,164        "threads": (165            os.environ.get("ORBITQUANT_CPU_THREADS", "runtime default")166            if args.device == "cpu"167            else None168        ),169        "torch_threads": torch.get_num_threads() if args.device == "cpu" else None,170        "with_bias": args.with_bias,171        "packed_seconds_per_iter": packed_seconds,172        "packed_first_call_seconds": packed_first_call_seconds,173        "packed_hot_mean_seconds": packed_distribution["mean"],174        "packed_hot_median_seconds": packed_distribution["median"],175        "packed_hot_p95_seconds": packed_distribution["p95"],176        "predequantized_f_linear_seconds_per_iter": predequantized_linear_seconds,177        "predequantized_first_call_seconds": predequantized_first_call_seconds,178        "predequantized_hot_mean_seconds": predequantized_distribution["mean"],179        "predequantized_hot_median_seconds": predequantized_distribution["median"],180        "predequantized_hot_p95_seconds": predequantized_distribution["p95"],181        "dequantize_then_f_linear_seconds_per_iter": dequantize_then_linear_seconds,182        "dequantize_then_first_call_seconds": dequantize_then_first_call_seconds,183        "dequantize_then_hot_mean_seconds": dequantize_then_distribution["mean"],184        "dequantize_then_hot_median_seconds": dequantize_then_distribution["median"],185        "dequantize_then_hot_p95_seconds": dequantize_then_distribution["p95"],186        "packed_weight_indices_bytes": packed_weight_indices_bytes,187        "row_norms_bytes": row_norms_bytes,188        "centroid_bytes": centroid_bytes,189        "packed_weight_path_bytes": packed_weight_path_bytes,190        "materialized_weight_bytes": materialized_weight_bytes,191        "packed_weight_path_vs_materialized_weight_ratio": packed_weight_path_bytes192        / materialized_weight_bytes193        if materialized_weight_bytes > 0194        else None,195        "packed_vs_predequantized_f_linear_speedup": predequantized_linear_seconds196        / packed_seconds197        if packed_seconds > 0198        else None,199        "packed_vs_dequantize_then_f_linear_speedup": dequantize_then_linear_seconds200        / packed_seconds201        if packed_seconds > 0202        else None,203        "reference_seconds_per_iter": predequantized_linear_seconds,204        "packed_vs_reference_speedup": predequantized_linear_seconds / packed_seconds205        if packed_seconds > 0206        else None,207        "max_abs_error": max_abs_error,208        "rmse": rmse,209        "relative_rmse": relative_rmse,210        "timing_headline": "hot-loop median seconds per iteration",211        "reference": (212            "predequantized PyTorch F.linear over a materialized dequantized "213            "weight matrix"214        ),215        "dequantize_reference": (216            "materialize the dequantized weight matrix, then call PyTorch F.linear"217        ),218    }219 220 221def main() -> None:222    parser = argparse.ArgumentParser()223    parser.add_argument("--device", choices=["cpu", "cuda", "mps"], default="cuda")224    parser.add_argument("--bits", type=int, default=4)225    parser.add_argument(226        "--rows",227        type=_parse_rows,228        default=[1, 8, 512, 4096],229        help="comma-separated row counts to sweep (default covers decode-bound "230        "small batches and GEMM-bound large batches)",231    )232    parser.add_argument(233        "--dtype",234        choices=["auto", *sorted(_DTYPES)],235        default="auto",236        help="activation dtype; auto picks float16 on mps and bfloat16 elsewhere",237    )238    parser.add_argument("--in-features", type=int, default=3072)239    parser.add_argument("--out-features", type=int, default=3072)240    parser.add_argument("--iters", type=int, default=20)241    parser.add_argument("--warmup", type=int, default=3)242    parser.add_argument("--seed", type=int, default=0)243    parser.add_argument("--threads", type=int, default=0)244    parser.add_argument("--with-bias", action="store_true")245    args = parser.parse_args()246 247    if args.threads < 0:248        parser.error("--threads must be non-negative")249    if args.iters <= 0 or args.warmup < 0:250        parser.error("--iters must be positive and --warmup must be non-negative")251    if args.device == "cpu" and args.threads > 0:252        os.environ["ORBITQUANT_CPU_THREADS"] = str(args.threads)253        torch.set_num_threads(args.threads)254 255    torch.manual_seed(args.seed)256    if args.dtype == "auto":257        dtype = torch.float16 if args.device == "mps" else torch.bfloat16258    else:259        dtype = _DTYPES[args.dtype]260    indices = torch.randint(261        0,262        2**args.bits,263        (args.out_features, args.in_features),264        dtype=torch.uint8,265    )266    packed = _pack(indices, args.bits).to(args.device)267    row_norms = torch.ones(args.out_features, device=args.device, dtype=torch.bfloat16)268    centroids = torch.linspace(-1.0, 1.0, 2**args.bits, device=args.device)269    indices_device = indices.long().to(args.device)270    weights = (packed, row_norms, centroids, indices_device)271 272    payloads = [_benchmark_rows(args, rows, dtype, weights) for rows in args.rows]273    if len(payloads) == 1:274        print(json.dumps(payloads[0], indent=2, sort_keys=True))275    else:276        print(json.dumps(payloads, indent=2, sort_keys=True))277 278 279if __name__ == "__main__":280    main()281