HDCharles · September 11, 2026 18:49
diff --git a/README.md b/README.md
diff --git a/bench_llama3.py b/bench_llama3.py
 import argparse
 import json
 import os
 import time

 import torch
 from compressed_tensors.quantization import QuantizationArgs, QuantizationScheme
 from phase_memory import install_phase_tracker
 from torch.utils.data import DataLoader
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from weight_hashes import write_quantized_weight_hashes

 from llmcompressor import oneshot
 from llmcompressor.modifiers.gptq import GPTQModifier

 DEFAULT_MODEL_PATH = (
    "/home/HDCharles/hf_hub/models--meta-llama--Meta-Llama-3-8B-Instruct/"
    "snapshots/8afb486c1db24fe5011ec46dfbe5b5dccdb575c2"
 )
 NUM_SAMPLES = 8
 SEQ_LEN = 256


 def make_calibration_loader(vocab_size: int) -> DataLoader:
    generator = torch.Generator().manual_seed(0)
    samples = [
        {
            "input_ids": torch.randint(
                0, vocab_size, (SEQ_LEN,), generator=generator, dtype=torch.long
            ),
            "attention_mask": torch.ones(SEQ_LEN, dtype=torch.long),
        }
        for _ in range(NUM_SAMPLES)
    ]
    return DataLoader(samples, batch_size=1)


 def cuda_stats():
    torch.cuda.synchronize()
    free, total = torch.cuda.mem_get_info()
    return {
        "allocated_mib": round(torch.cuda.memory_allocated() / 2**20, 1),
        "reserved_mib": round(torch.cuda.memory_reserved() / 2**20, 1),
        "free_mib": round(free / 2**20, 1),
        "total_mib": round(total / 2**20, 1),
    }


 def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--label", required=True)
    parser.add_argument("--model-path", default=DEFAULT_MODEL_PATH)
    parser.add_argument("--batching", action="store_true")
    parser.add_argument("--backend", choices=("eager", "triton"), default="triton")
    parser.add_argument("--linear-sequential-targets", action="store_true")
    parser.add_argument("--debug-phases", action="store_true")
    parser.add_argument("--weight-hashes-output")
    parser.add_argument("--single-layer", action="store_true")
    args = parser.parse_args()

    if args.backend == "eager":
        os.environ["LLMCOMPRESSOR_DISABLE_GPTQ_TRITON"] = "1"
    else:
        os.environ.pop("LLMCOMPRESSOR_DISABLE_GPTQ_TRITON", None)

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required")

    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
    model = AutoModelForCausalLM.from_pretrained(
        args.model_path, dtype=torch.bfloat16
    ).cuda()
    if args.single_layer:
        model.model.layers = torch.nn.ModuleList([model.model.layers[0]])
        model.config.num_hidden_layers = 1
    loader = make_calibration_loader(model.config.vocab_size)
    load_stats = cuda_stats()
    tracker = install_phase_tracker() if args.debug_phases else None
    if tracker is None:
        torch.cuda.reset_peak_memory_stats()

    scheme = QuantizationScheme(
        targets=["Linear"],
        weights=QuantizationArgs(
            num_bits=4,
            type="int",
            symmetric=True,
            strategy="group",
            group_size=128,
        ),
    )
    fields = getattr(GPTQModifier, "model_fields", {})
    modifier_kwargs = {
        "targets": "Linear",
        "ignore": ["lm_head"],
        "config_groups": {"group_0": scheme},
    }
    if "batched_quantization" in fields:
        modifier_kwargs["batched_quantization"] = args.batching
        modifier_kwargs["batch_memory_fraction"] = 0.75

    start = time.perf_counter()
    oneshot(
        model=model,
        tokenizer=tokenizer,
        dataset=loader,
        recipe=GPTQModifier(**modifier_kwargs),
        num_calibration_samples=NUM_SAMPLES,
        max_seq_length=SEQ_LEN,
        pipeline="sequential",
        sequential_targets=["Linear"] if args.linear_sequential_targets else None,
        save_compressed=False,
        output_dir=None,
    )
    torch.cuda.synchronize()
    elapsed = time.perf_counter() - start
    if tracker is not None:
        tracker.stop()
        peak_stats = tracker.result()
    else:
        peak_stats = {
            "peak_allocated_mib": round(torch.cuda.max_memory_allocated() / 2**20, 1),
            "peak_reserved_mib": round(torch.cuda.max_memory_reserved() / 2**20, 1),
        }
    if args.weight_hashes_output:
        write_quantized_weight_hashes(model, args.weight_hashes_output)
    print(
        json.dumps(
            {
                "label": args.label,
                "model": "Meta-Llama-3-8B-Instruct",
                "scheme": "W4A16-G128",
                "calibration": f"{NUM_SAMPLES}x{SEQ_LEN}",
                "batching_requested": args.batching,
                "backend": args.backend,
                "sequential_targets": (
                    "Linear" if args.linear_sequential_targets else "automatic"
                ),
                "batching_supported": "batched_quantization" in fields,
                "elapsed_seconds": round(elapsed, 2),
                "load": load_stats,
                "after": cuda_stats(),
                **peak_stats,
            }
        ),
        flush=True,
    )


 if __name__ == "__main__":
    main()
diff --git a/bench_qwen3_full.py b/bench_qwen3_full.py
 """Run the Qwen3-30B-A3B GPTQ memory benchmark without layer truncation."""

 import sys

 from bench_qwen3_moe import main

 if __name__ == "__main__":
    sys.argv.append("--full-model")
    main()
diff --git a/bench_qwen3_moe.py b/bench_qwen3_moe.py
 import argparse
 import json
 import os
 import threading
 import time

 import torch
 from phase_memory import (
    install_phase_tracker,
    start_allocation_history,
    stop_allocation_history,
 )
 from torch.utils.data import DataLoader
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
 from weight_hashes import write_quantized_weight_hashes

 from llmcompressor import oneshot
 from llmcompressor.modifiers.gptq import GPTQModifier
 from llmcompressor.utils import load_context

 DEFAULT_MODEL_PATH = (
    "/home/HDCharles/hf_hub/models--Qwen--Qwen3-30B-A3B/"
    "snapshots/ad44e777bcd18fa416d9da3bd8f70d33ebb85d39"
 )
 NUM_SAMPLES = 8
 SEQ_LEN = 256


 class StopAfterFirstGPTQ(RuntimeError):
    pass


 def make_calibration_loader(vocab_size: int) -> DataLoader:
    generator = torch.Generator().manual_seed(0)
    samples = [
        {
            "input_ids": torch.randint(
                0, vocab_size, (SEQ_LEN,), generator=generator, dtype=torch.long
            ),
            "attention_mask": torch.ones(SEQ_LEN, dtype=torch.long),
        }
        for _ in range(NUM_SAMPLES)
    ]
    return DataLoader(samples, batch_size=1)


 def cuda_stats():
    torch.cuda.synchronize()
    free, total = torch.cuda.mem_get_info()
    return {
        "allocated_mib": round(torch.cuda.memory_allocated() / 2**20, 1),
        "reserved_mib": round(torch.cuda.memory_reserved() / 2**20, 1),
        "free_mib": round(free / 2**20, 1),
        "total_mib": round(total / 2**20, 1),
    }


 def rss_mib():
    with open("/proc/self/statm") as statm:
        resident_pages = int(statm.read().split()[1])
    return resident_pages * os.sysconf("SC_PAGE_SIZE") / 2**20


 def sample_rss(stop: threading.Event, peak: list[float]):
    while not stop.wait(0.01):
        peak[0] = max(peak[0], rss_mib())


 def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--label", required=True)
    parser.add_argument("--model-path", default=DEFAULT_MODEL_PATH)
    parser.add_argument("--targets", default="Linear")
    parser.add_argument("--batching", action="store_true")
    parser.add_argument("--backend", choices=("eager", "triton"), default="triton")
    parser.add_argument("--linear-sequential-targets", action="store_true")
    parser.add_argument("--sequential-targets", nargs="+")
    parser.add_argument("--targets-per-subgraph", type=int, default=1)
    parser.add_argument("--debug-phases", action="store_true")
    parser.add_argument("--debug-allocation-history", action="store_true")
    parser.add_argument("--time-gptq", action="store_true")
    parser.add_argument("--weight-hashes-output")
    parser.add_argument("--offload-hessians", action="store_true")
    parser.add_argument("--stop-after-first-gptq", action="store_true")
    parser.add_argument("--full-model", action="store_true", help=argparse.SUPPRESS)
    args = parser.parse_args()

    if args.backend == "eager":
        os.environ["LLMCOMPRESSOR_DISABLE_GPTQ_TRITON"] = "1"
    else:
        os.environ.pop("LLMCOMPRESSOR_DISABLE_GPTQ_TRITON", None)

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required")

    config = AutoConfig.from_pretrained(args.model_path)
    if not args.full_model:
        config.num_hidden_layers = 1
    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
    load_kwargs = {}
    if args.full_model:
        load_kwargs = {
            "device_map": "auto_offload",
            "offload_folder": f"/tmp/qwen3-full-offload-{args.label}",
        }
    with load_context():
        model = AutoModelForCausalLM.from_pretrained(
            args.model_path,
            config=config,
            dtype=torch.bfloat16,
            low_cpu_mem_usage=True,
            **load_kwargs,
        )
    if not args.full_model:
        model = model.cuda()
    if not args.full_model:
        assert len(model.model.layers) == 1

    loader = make_calibration_loader(model.config.vocab_size)
    load_cuda = cuda_stats()
    load_rss = rss_mib()
    tracker = install_phase_tracker() if args.debug_phases else None
    if tracker is None:
        torch.cuda.reset_peak_memory_stats()
    history_baseline = (
        start_allocation_history() if args.debug_allocation_history else None
    )
    peak_rss = [load_rss]
    stop_sampling = threading.Event()
    sampler = threading.Thread(
        target=sample_rss, args=(stop_sampling, peak_rss), daemon=True
    )

    fields = getattr(GPTQModifier, "model_fields", {})
    modifier_kwargs = {
        "targets": args.targets,
        "scheme": "W4A16",
        "ignore": [
            "lm_head",
            "re:.*mlp.gate$",
            "re:.*mlp.shared_expert_gate$",
        ],
    }
    if "batched_quantization" in fields:
        modifier_kwargs["batched_quantization"] = args.batching
        modifier_kwargs["batch_memory_fraction"] = 0.75
    if args.offload_hessians:
        if "offload_hessians" not in fields:
            raise ValueError(
                "This GPTQ implementation does not support Hessian offload"
            )
        modifier_kwargs["offload_hessians"] = True

    stopped_after_first_gptq = False
    if args.stop_after_first_gptq:
        original_compress_module_list = GPTQModifier.compress_module_list

        def compress_module_list_once(modifier, module_list, *call_args, **call_kwargs):
            compressed_before = modifier._num_compressed_modules
            original_compress_module_list(
                modifier, module_list, *call_args, **call_kwargs
            )
            if modifier._num_compressed_modules > compressed_before:
                raise StopAfterFirstGPTQ

        GPTQModifier.compress_module_list = compress_module_list_once

    sequential_targets = args.sequential_targets
    if sequential_targets is None and args.linear_sequential_targets:
        sequential_targets = ["Linear"]

    sampler.start()
    start = time.perf_counter()
    try:
        try:
            oneshot(
                model=model,
                tokenizer=tokenizer,
                dataset=loader,
                recipe=GPTQModifier(**modifier_kwargs),
                num_calibration_samples=NUM_SAMPLES,
                max_seq_length=SEQ_LEN,
                pipeline="sequential",
                sequential_targets=sequential_targets,
                sequential_targets_per_subgraph=args.targets_per_subgraph,
                save_compressed=False,
                output_dir=None,
                moe_calibrate_all_experts=True,
            )
        except StopAfterFirstGPTQ:
            stopped_after_first_gptq = True
        torch.cuda.synchronize()
    finally:
        stop_sampling.set()
        sampler.join()
    elapsed = time.perf_counter() - start
    allocation_peak = (
        stop_allocation_history(history_baseline)
        if history_baseline is not None
        else None
    )
    if tracker is not None:
        tracker.stop()
        peak_stats = tracker.result()
    else:
        peak_stats = {
            "peak_allocated_mib": round(torch.cuda.max_memory_allocated() / 2**20, 1),
            "peak_reserved_mib": round(torch.cuda.max_memory_reserved() / 2**20, 1),
        }
    if args.weight_hashes_output:
        write_quantized_weight_hashes(model, args.weight_hashes_output)
    profile_sections = None
    if args.time_gptq:
        from llmcompressor.modifiers.gptq.gptq_quantize import get_gptq_timings

        profile_sections = get_gptq_timings()

    print(
        json.dumps(
            {
                "label": args.label,
                "model": (
                    "Qwen3-30B-A3B" if args.full_model else "Qwen3-30B-A3B-one-layer"
                ),
                "experts": model.config.num_experts,
                "scheme": "W4A16-G128",
                "targets": args.targets,
                "calibration": f"{NUM_SAMPLES}x{SEQ_LEN}",
                "batching_requested": args.batching,
                "backend": args.backend,
                "offload_hessians": args.offload_hessians,
                "stopped_after_first_gptq": stopped_after_first_gptq,
                "sequential_targets": sequential_targets or "automatic",
                "targets_per_subgraph": args.targets_per_subgraph,
                "batching_supported": "batched_quantization" in fields,
                "elapsed_seconds": round(elapsed, 2),
                "load_cuda": load_cuda,
                "after_cuda": cuda_stats(),
                "load_rss_mib": round(load_rss, 1),
                "peak_rss_mib": round(peak_rss[0], 1),
                "after_rss_mib": round(rss_mib(), 1),
                "allocation_peak": allocation_peak,
                "profile_sections": profile_sections,
                **peak_stats,
            }
        ),
        flush=True,
    )


 if __name__ == "__main__":
    main()
diff --git a/compare_gptq_debug.py b/compare_gptq_debug.py
 import argparse

 import torch


 def _compare(left, right, path=""):
    if isinstance(left, torch.Tensor) and isinstance(right, torch.Tensor):
        equal = left.dtype == right.dtype and torch.equal(left, right)
        if equal:
            return None
        result = {
            "path": path,
            "left_shape": tuple(left.shape),
            "right_shape": tuple(right.shape),
            "left_dtype": str(left.dtype),
            "right_dtype": str(right.dtype),
        }
        if left.shape == right.shape and left.is_floating_point():
            delta = (left.float() - right.float()).abs()
            result["mismatches"] = int((left != right).sum())
            result["max_abs_diff"] = float(delta.max())
        return result
    if isinstance(left, dict) and isinstance(right, dict):
        if left.keys() != right.keys():
            return {
                "path": path,
                "left_keys": sorted(left),
                "right_keys": sorted(right),
            }
        for key in left:
            mismatch = _compare(left[key], right[key], f"{path}.{key}".strip("."))
            if mismatch is not None:
                return mismatch
        return None
    if isinstance(left, (tuple, list)) and isinstance(right, (tuple, list)):
        if len(left) != len(right):
            return {"path": path, "left_len": len(left), "right_len": len(right)}
        for index, (left_item, right_item) in enumerate(zip(left, right)):
            mismatch = _compare(left_item, right_item, f"{path}[{index}]")
            if mismatch is not None:
                return mismatch
        return None
    if left != right:
        return {"path": path, "left": left, "right": right}
    return None


 def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("left")
    parser.add_argument("right")
    args = parser.parse_args()
    left = torch.load(args.left, map_location="cpu", weights_only=False)
    right = torch.load(args.right, map_location="cpu", weights_only=False)

    for section in ("before", "after", "layer_outputs"):
        mismatch = _compare(left[section], right[section], section)
        print(f"{section}: {mismatch or 'equal'}")


 if __name__ == "__main__":
    main()
diff --git a/compare_weight_hashes.py b/compare_weight_hashes.py
 import argparse
 import json


 def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("reference")
    parser.add_argument("candidate")
    args = parser.parse_args()

    with open(args.reference) as reference_file:
        reference = json.load(reference_file)
    with open(args.candidate) as candidate_file:
        candidate = json.load(candidate_file)

    reference_names = set(reference)
    candidate_names = set(candidate)
    missing = sorted(reference_names - candidate_names)
    unexpected = sorted(candidate_names - reference_names)
    mismatched = sorted(
        name
        for name in reference_names & candidate_names
        if reference[name] != candidate[name]
    )
    max_reported = 20
    print(
        json.dumps(
            {
                "reference_weights": len(reference),
                "candidate_weights": len(candidate),
                "bitwise_equal": not (missing or unexpected or mismatched),
                "mismatch_count": len(mismatched),
                "missing": missing,
                "unexpected": unexpected,
                "mismatched": mismatched[:max_reported],
                "mismatches_omitted": max(0, len(mismatched) - max_reported),
            }
        )
    )


 if __name__ == "__main__":
    main()
diff --git a/debug_gptq_parity.py b/debug_gptq_parity.py
 import argparse
 from pathlib import Path

 import torch
 from compressed_tensors.quantization import QuantizationArgs, QuantizationScheme
 from torch.utils.data import DataLoader
 from transformers import AutoModelForCausalLM, AutoTokenizer

 from llmcompressor import oneshot
 from llmcompressor.modifiers.gptq import GPTQModifier

 DEFAULT_MODEL_PATH = (
    "/home/HDCharles/hf_hub/models--meta-llama--Meta-Llama-3-8B-Instruct/"
    "snapshots/8afb486c1db24fe5011ec46dfbe5b5dccdb575c2"
 )
 NUM_SAMPLES = 8
 SEQ_LEN = 256


 def _cpu_copy(value):
    if isinstance(value, torch.Tensor):
        return value.detach().cpu().clone()
    if isinstance(value, (tuple, list)):
        return type(value)(_cpu_copy(item) for item in value)
    if isinstance(value, dict):
        return {key: _cpu_copy(item) for key, item in value.items()}
    return value


 def _module_snapshot(modifier, module):
    return {
        "weight": _cpu_copy(module.weight),
        "hessian": _cpu_copy(modifier._hessians[module]),
        "num_samples": _cpu_copy(modifier._num_samples[module]),
    }


 def _compressed_snapshot(module):
    result = {"weight": _cpu_copy(module.weight)}
    for name in (
        "weight_scale",
        "weight_zero_point",
        "weight_global_scale",
        "weight_g_idx",
    ):
        if hasattr(module, name):
            result[name] = _cpu_copy(getattr(module, name))
    return result


 def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--model-path", default=DEFAULT_MODEL_PATH)
    parser.add_argument("--batching", action="store_true")
    parser.add_argument("--num-layers", type=int, default=2)
    parser.add_argument("--module-suffix", default=".self_attn.q_proj")
    args = parser.parse_args()

    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
    model = AutoModelForCausalLM.from_pretrained(
        args.model_path, dtype=torch.bfloat16
    ).cuda()
    model.model.layers = torch.nn.ModuleList(
        list(model.model.layers[: args.num_layers])
    )
    model.config.num_hidden_layers = args.num_layers

    generator = torch.Generator().manual_seed(0)
    samples = [
        {
            "input_ids": torch.randint(
                0,
                model.config.vocab_size,
                (SEQ_LEN,),
                generator=generator,
                dtype=torch.long,
            ),
            "attention_mask": torch.ones(SEQ_LEN, dtype=torch.long),
        }
        for _ in range(NUM_SAMPLES)
    ]
    loader = DataLoader(samples, batch_size=1)

    snapshots = {"before": {}, "after": {}, "layer_outputs": {}}
    handles = []
    for layer_index, layer in enumerate(model.model.layers):
        outputs = snapshots["layer_outputs"].setdefault(str(layer_index), [])

        def capture_output(_module, _inputs, output, outputs=outputs):
            outputs.append(_cpu_copy(output))

        handles.append(layer.register_forward_hook(capture_output))

    original_compress = GPTQModifier.compress_module_list

    def traced_compress(self, module_list, *compress_args, **compress_kwargs):
        selected = [
            module
            for module in module_list
            if self._module_names[module].startswith(
                tuple(f"model.layers.{index}." for index in range(args.num_layers))
            )
            and self._module_names[module].endswith(args.module_suffix)
        ]
        for module in selected:
            name = self._module_names[module]
            snapshots["before"][name] = _module_snapshot(self, module)
        result = original_compress(self, module_list, *compress_args, **compress_kwargs)
        for module in selected:
            name = self._module_names[module]
            snapshots["after"][name] = _compressed_snapshot(module)
        return result

    GPTQModifier.compress_module_list = traced_compress
    scheme = QuantizationScheme(
        targets=["Linear"],
        weights=QuantizationArgs(
            num_bits=4,
            type="int",
            symmetric=True,
            strategy="group",
            group_size=128,
        ),
    )
    modifier_kwargs = {
        "targets": "Linear",
        "ignore": ["lm_head"],
        "config_groups": {"group_0": scheme},
    }
    fields = getattr(GPTQModifier, "model_fields", {})
    if "batched_quantization" in fields:
        modifier_kwargs["batched_quantization"] = args.batching

    try:
        oneshot(
            model=model,
            tokenizer=tokenizer,
            dataset=loader,
            recipe=GPTQModifier(**modifier_kwargs),
            num_calibration_samples=NUM_SAMPLES,
            max_seq_length=SEQ_LEN,
            pipeline="sequential",
            save_compressed=False,
            output_dir=None,
        )
    finally:
        GPTQModifier.compress_module_list = original_compress
        for handle in handles:
            handle.remove()

    args.output.parent.mkdir(parents=True, exist_ok=True)
    torch.save(snapshots, args.output)


 if __name__ == "__main__":
    main()
diff --git a/kernel_microbenchmark.py b/kernel_microbenchmark.py
 """Microbenchmark fused Triton GPTQ and eager GPTQ on representative shapes."""

 import argparse
 import copy
 import os
 import time

 import torch
 from compressed_tensors.quantization import QuantizationArgs, QuantizationScheme

 from llmcompressor.modifiers.gptq.gptq_quantize import (
    quantize_weight,
 )
 from llmcompressor.modifiers.quantization.calibration import (
    initialize_observer,
    observe,
 )

 DEV = "cuda"
 SCHEMES = {
    "NVFP4-g16": QuantizationArgs(
        num_bits=4, type="float", symmetric=True, strategy="tensor_group",
        group_size=16,
    ),
    "INT4-g128": QuantizationArgs(
        num_bits=4, symmetric=True, strategy="group", group_size=128
    ),
 }
 # (in_features, out_features, meaningful batch sizes). Batch keys require
 # exactly matching weight shapes: dense-model projections only form pairs,
 # whereas all same-shaped MoE experts can be batched together.
 SHAPES = {
    "qwen_attn_qo": (2048, 2048, (1, 2)),
    "qwen_attn_kv": (2048, 256, (1, 2)),
    "qwen_expert_gate_up": (2048, 768, (1, 8, 32, 64, 128)),
    "qwen_expert_down": (768, 2048, (1, 8, 32, 64, 128)),
    "llama_attn_qo": (4096, 4096, (1, 2)),
    "llama_attn_kv": (4096, 1024, (1, 2)),
    "llama_mlp_gate_up": (4096, 14336, (1, 2)),
    "llama_mlp_down": (14336, 4096, (1,)),
 }


 def make_mod(in_f, out_f, qa, seed):
    torch.manual_seed(seed)
    m = torch.nn.Linear(in_f, out_f, bias=False).to(DEV)
    m.quantization_scheme = QuantizationScheme(targets=["Linear"], weights=qa)
    initialize_observer(m, "weight")
    observe(m, "weight")
    return m


 def make_hessian(n, seed):
    g = torch.Generator(DEV).manual_seed(seed)
    A = torch.randn(n, n, generator=g, device=DEV)
    return A @ A.T + torch.eye(n, device=DEV)


 def bench(fn, warmup=1, iters=3):
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(iters):
        fn()
    torch.cuda.synchronize()
    return (time.perf_counter() - t0) / iters


 def report(scheme_name, shape_name, name, t):
    print(f"{scheme_name:10s} {shape_name:8s} {name:34s} {t * 1e3:10.1f} ms",
          flush=True)


 def report_speedup(scheme_name, shape_name, name, speedup):
    print(f"{scheme_name:10s} {shape_name:8s} {name:34s} {speedup:10.2f}x",
          flush=True)


 def quantize_modules(modules, quant_args, hessians):
    """Call the tensor-level GPTQ API for a stack of same-shaped modules."""
    qparams = [module.weight_observer.get_qparams() for module in modules]
    global_scales = None
    if qparams[0]["global_scale"] is not None:
        global_scales = torch.stack(
            [qparam["global_scale"].reshape(-1)[0] for qparam in qparams]
        )
    return quantize_weight(
        weights=torch.stack([module.weight for module in modules]),
        hessians=torch.stack(hessians),
        scale=torch.stack([qparam["scale"] for qparam in qparams]),
        zero_point=torch.stack([qparam["zero_point"] for qparam in qparams]),
        global_scale=global_scales,
        quant_args=quant_args,
    )


 def run(shape_names, scheme_names):
    torch.manual_seed(0)
    for scheme_name in scheme_names:
        qa_tpl = SCHEMES[scheme_name]
        for shape_name in shape_names:
            in_f, out_f, batch_sizes = SHAPES[shape_name]
            qa = copy.deepcopy(qa_tpl)
            t_build = time.perf_counter()
            Hs = [make_hessian(in_f, 100 + s) for s in range(max(batch_sizes))]
            torch.cuda.synchronize()
            report(scheme_name, shape_name, "hessian build (setup)",
                   time.perf_counter() - t_build)

            os.environ["LLMCOMPRESSOR_DISABLE_GPTQ_TRITON"] = "0"

            # 1. batched, B sweep (fast path, print early)
            t_b = None
            for B in batch_sizes:
                def fn(B=B):
                    ms = [make_mod(in_f, out_f, qa, s) for s in range(B)]
                    quantize_modules(ms, qa, [h.clone() for h in Hs[:B]])
                t = bench(fn, warmup=1, iters=3)
                report(scheme_name, shape_name, f"batched B={B}", t)
                if B == max(batch_sizes):
                    t_b = t

            # 2. single fused
            def fn_fused():
                m = make_mod(in_f, out_f, qa, 0)
                quantize_modules([m], qa, [Hs[0].clone()])
            t_fused = bench(fn_fused, warmup=2, iters=5)
            report(scheme_name, shape_name, "single fused kernel", t_fused)
            report(
                scheme_name,
                shape_name,
                f"loop{max(batch_sizes)} fused (extrap.)",
                t_fused * max(batch_sizes),
            )
            if t_b is not None:
                report_speedup(
                    scheme_name, shape_name,
                    (
                        f"speedup batched{max(batch_sizes)} vs "
                        f"loop{max(batch_sizes)}-fused"
                    ),
                    t_fused * max(batch_sizes) / t_b,
                )

            # 3. single eager, 1 rep (slow baseline)
            os.environ["LLMCOMPRESSOR_DISABLE_GPTQ_TRITON"] = "1"
            t_eager = bench(fn_fused, warmup=0, iters=1)
            os.environ["LLMCOMPRESSOR_DISABLE_GPTQ_TRITON"] = "0"
            report(
                scheme_name,
                shape_name,
                "single eager (kernel off, 1 rep)",
                t_eager,
            )
            report_speedup(
                scheme_name, shape_name, "speedup fused vs eager single",
                t_eager / t_fused,
            )
            if t_b is not None:
                report_speedup(
                    scheme_name, shape_name,
                    (
                        f"speedup batched{max(batch_sizes)} vs "
                        f"loop{max(batch_sizes)}-eager"
                    ),
                    t_eager * max(batch_sizes) / t_b,
                )


 if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--shapes", nargs="+", choices=tuple(SHAPES), default=tuple(SHAPES)
    )
    parser.add_argument(
        "--schemes", nargs="+", choices=tuple(SCHEMES), default=tuple(SCHEMES)
    )
    args = parser.parse_args()
    run(args.shapes, args.schemes)
diff --git a/kernel_microbenchmark_int4_results.txt b/kernel_microbenchmark_int4_results.txt
 INT4-g128  qwen_attn_qo hessian build (setup)                   236.8 ms
 INT4-g128  qwen_attn_qo batched B=1                              23.7 ms
 INT4-g128  qwen_attn_qo batched B=2                              42.4 ms
 INT4-g128  qwen_attn_qo single fused kernel                      22.4 ms
 INT4-g128  qwen_attn_qo loop2 fused (extrap.)                    44.8 ms
 INT4-g128  qwen_attn_qo speedup batched2 vs loop2-fused          1.06x
 INT4-g128  qwen_attn_qo single eager (kernel off, 1 rep)        402.6 ms
 INT4-g128  qwen_attn_qo speedup fused vs eager single           17.96x
 INT4-g128  qwen_attn_qo speedup batched2 vs loop2-eager         18.97x
 INT4-g128  qwen_attn_kv hessian build (setup)                     0.3 ms
 INT4-g128  qwen_attn_kv batched B=1                              11.0 ms
 INT4-g128  qwen_attn_kv batched B=2                              17.8 ms
 INT4-g128  qwen_attn_kv single fused kernel                      10.8 ms
 INT4-g128  qwen_attn_kv loop2 fused (extrap.)                    21.7 ms
 INT4-g128  qwen_attn_kv speedup batched2 vs loop2-fused          1.22x
 INT4-g128  qwen_attn_kv single eager (kernel off, 1 rep)        244.9 ms
 INT4-g128  qwen_attn_kv speedup fused vs eager single           22.60x
 INT4-g128  qwen_attn_kv speedup batched2 vs loop2-eager         27.50x
 INT4-g128  qwen_expert_gate_up hessian build (setup)                    19.1 ms
 INT4-g128  qwen_expert_gate_up batched B=1                              14.3 ms
 INT4-g128  qwen_expert_gate_up batched B=8                              92.3 ms
 INT4-g128  qwen_expert_gate_up batched B=32                            248.5 ms
 INT4-g128  qwen_expert_gate_up batched B=64                            518.4 ms
 INT4-g128  qwen_expert_gate_up batched B=128                          1062.3 ms
 INT4-g128  qwen_expert_gate_up single fused kernel                      14.1 ms
 INT4-g128  qwen_expert_gate_up loop128 fused (extrap.)                1802.7 ms
 INT4-g128  qwen_expert_gate_up speedup batched128 vs loop128-fused       1.70x
 INT4-g128  qwen_expert_gate_up single eager (kernel off, 1 rep)        254.1 ms
 INT4-g128  qwen_expert_gate_up speedup fused vs eager single           18.05x
 INT4-g128  qwen_expert_gate_up speedup batched128 vs loop128-eager      30.62x
 INT4-g128  qwen_expert_down hessian build (setup)                     7.6 ms
 INT4-g128  qwen_expert_down batched B=1                               8.8 ms
 INT4-g128  qwen_expert_down batched B=8                              58.2 ms
 INT4-g128  qwen_expert_down batched B=32                            194.8 ms
 INT4-g128  qwen_expert_down batched B=64                            430.0 ms
 INT4-g128  qwen_expert_down batched B=128                           893.9 ms
 INT4-g128  qwen_expert_down single fused kernel                       8.8 ms
 INT4-g128  qwen_expert_down loop128 fused (extrap.)                1129.7 ms
 INT4-g128  qwen_expert_down speedup batched128 vs loop128-fused       1.26x
 INT4-g128  qwen_expert_down single eager (kernel off, 1 rep)        104.6 ms
 INT4-g128  qwen_expert_down speedup fused vs eager single           11.85x
 INT4-g128  qwen_expert_down speedup batched128 vs loop128-eager      14.98x
 INT4-g128  llama_attn_qo hessian build (setup)                     0.4 ms
 INT4-g128  llama_attn_qo batched B=1                              77.3 ms
 INT4-g128  llama_attn_qo batched B=2                             144.5 ms
 INT4-g128  llama_attn_qo single fused kernel                      74.2 ms
 INT4-g128  llama_attn_qo loop2 fused (extrap.)                   148.3 ms
 INT4-g128  llama_attn_qo speedup batched2 vs loop2-fused          1.03x
 INT4-g128  llama_attn_qo single eager (kernel off, 1 rep)        860.2 ms
 INT4-g128  llama_attn_qo speedup fused vs eager single           11.60x
 INT4-g128  llama_attn_qo speedup batched2 vs loop2-eager         11.90x
 INT4-g128  llama_attn_kv hessian build (setup)                     0.2 ms
 INT4-g128  llama_attn_kv batched B=1                              34.5 ms
 INT4-g128  llama_attn_kv batched B=2                              62.1 ms
 INT4-g128  llama_attn_kv single fused kernel                      33.9 ms
 INT4-g128  llama_attn_kv loop2 fused (extrap.)                    67.8 ms
 INT4-g128  llama_attn_kv speedup batched2 vs loop2-fused          1.09x
 INT4-g128  llama_attn_kv single eager (kernel off, 1 rep)        520.6 ms
 INT4-g128  llama_attn_kv speedup fused vs eager single           15.36x
 INT4-g128  llama_attn_kv speedup batched2 vs loop2-eager         16.77x
 INT4-g128  llama_mlp_gate_up hessian build (setup)                     0.2 ms
 INT4-g128  llama_mlp_gate_up batched B=1                             235.1 ms
 INT4-g128  llama_mlp_gate_up batched B=2                             464.8 ms
 INT4-g128  llama_mlp_gate_up single fused kernel                     235.4 ms
 INT4-g128  llama_mlp_gate_up loop2 fused (extrap.)                   470.8 ms
 INT4-g128  llama_mlp_gate_up speedup batched2 vs loop2-fused          1.01x
 INT4-g128  llama_mlp_gate_up single eager (kernel off, 1 rep)        766.7 ms
 INT4-g128  llama_mlp_gate_up speedup fused vs eager single            3.26x
 INT4-g128  llama_mlp_gate_up speedup batched2 vs loop2-eager          3.30x
 INT4-g128  llama_mlp_down hessian build (setup)                     0.3 ms
 INT4-g128  llama_mlp_down batched B=1                             486.0 ms
 INT4-g128  llama_mlp_down single fused kernel                     482.1 ms
 INT4-g128  llama_mlp_down loop1 fused (extrap.)                   482.1 ms
 INT4-g128  llama_mlp_down speedup batched1 vs loop1-fused          0.99x
 INT4-g128  llama_mlp_down single eager (kernel off, 1 rep)       2249.3 ms
 INT4-g128  llama_mlp_down speedup fused vs eager single            4.67x
 INT4-g128  llama_mlp_down speedup batched1 vs loop1-eager          4.63x
diff --git a/kernel_microbenchmark_nvfp4_results.txt b/kernel_microbenchmark_nvfp4_results.txt
 NVFP4-g16  qwen_attn_qo hessian build (setup)                   588.9 ms
 NVFP4-g16  qwen_attn_qo batched B=1                              24.1 ms
 NVFP4-g16  qwen_attn_qo batched B=2                              41.3 ms
 NVFP4-g16  qwen_attn_qo single fused kernel                      22.1 ms
 NVFP4-g16  qwen_attn_qo loop2 fused (extrap.)                    44.2 ms
 NVFP4-g16  qwen_attn_qo speedup batched2 vs loop2-fused          1.07x
 NVFP4-g16  qwen_attn_qo single eager (kernel off, 1 rep)        337.3 ms
 NVFP4-g16  qwen_attn_qo speedup fused vs eager single           15.26x
 NVFP4-g16  qwen_attn_qo speedup batched2 vs loop2-eager         16.33x
 NVFP4-g16  qwen_attn_kv hessian build (setup)                     1.1 ms
 NVFP4-g16  qwen_attn_kv batched B=1                              11.5 ms
 NVFP4-g16  qwen_attn_kv batched B=2                              18.6 ms
 NVFP4-g16  qwen_attn_kv single fused kernel                      11.3 ms
 NVFP4-g16  qwen_attn_kv loop2 fused (extrap.)                    22.7 ms
 NVFP4-g16  qwen_attn_kv speedup batched2 vs loop2-fused          1.22x
 NVFP4-g16  qwen_attn_kv single eager (kernel off, 1 rep)        323.0 ms
 NVFP4-g16  qwen_attn_kv speedup fused vs eager single           28.52x
 NVFP4-g16  qwen_attn_kv speedup batched2 vs loop2-eager         34.80x
 NVFP4-g16  qwen_expert_gate_up hessian build (setup)                    48.6 ms
 NVFP4-g16  qwen_expert_gate_up batched B=1                              14.5 ms
 NVFP4-g16  qwen_expert_gate_up batched B=8                              94.2 ms
 NVFP4-g16  qwen_expert_gate_up batched B=32                            250.8 ms
 NVFP4-g16  qwen_expert_gate_up batched B=64                            530.0 ms
 NVFP4-g16  qwen_expert_gate_up batched B=128                          1084.8 ms
 NVFP4-g16  qwen_expert_gate_up single fused kernel                      14.7 ms
 NVFP4-g16  qwen_expert_gate_up loop128 fused (extrap.)                1883.5 ms
 NVFP4-g16  qwen_expert_gate_up speedup batched128 vs loop128-fused       1.74x
 NVFP4-g16  qwen_expert_gate_up single eager (kernel off, 1 rep)        337.4 ms
 NVFP4-g16  qwen_expert_gate_up speedup fused vs eager single           22.93x
 NVFP4-g16  qwen_expert_gate_up speedup batched128 vs loop128-eager      39.81x
 NVFP4-g16  qwen_expert_down hessian build (setup)                     9.2 ms
 NVFP4-g16  qwen_expert_down batched B=1                               9.3 ms
 NVFP4-g16  qwen_expert_down batched B=8                              59.5 ms
 NVFP4-g16  qwen_expert_down batched B=32                            199.8 ms
 NVFP4-g16  qwen_expert_down batched B=64                            444.2 ms
 NVFP4-g16  qwen_expert_down batched B=128                           916.9 ms
 NVFP4-g16  qwen_expert_down single fused kernel                       9.1 ms
 NVFP4-g16  qwen_expert_down loop128 fused (extrap.)                1168.6 ms
 NVFP4-g16  qwen_expert_down speedup batched128 vs loop128-fused       1.27x
 NVFP4-g16  qwen_expert_down single eager (kernel off, 1 rep)        131.6 ms
 NVFP4-g16  qwen_expert_down speedup fused vs eager single           14.41x
 NVFP4-g16  qwen_expert_down speedup batched128 vs loop128-eager      18.36x
 NVFP4-g16  llama_attn_qo hessian build (setup)                     5.8 ms
 NVFP4-g16  llama_attn_qo batched B=1                              77.4 ms
 NVFP4-g16  llama_attn_qo batched B=2                             146.4 ms
 NVFP4-g16  llama_attn_qo single fused kernel                      76.4 ms
 NVFP4-g16  llama_attn_qo loop2 fused (extrap.)                   152.8 ms
 NVFP4-g16  llama_attn_qo speedup batched2 vs loop2-fused          1.04x
 NVFP4-g16  llama_attn_qo single eager (kernel off, 1 rep)        730.3 ms
 NVFP4-g16  llama_attn_qo speedup fused vs eager single            9.56x
 NVFP4-g16  llama_attn_qo speedup batched2 vs loop2-eager          9.98x
 NVFP4-g16  llama_attn_kv hessian build (setup)                     5.7 ms
 NVFP4-g16  llama_attn_kv batched B=1                              34.9 ms
 NVFP4-g16  llama_attn_kv batched B=2                              61.9 ms
 NVFP4-g16  llama_attn_kv single fused kernel                      35.5 ms
 NVFP4-g16  llama_attn_kv loop2 fused (extrap.)                    71.0 ms
 NVFP4-g16  llama_attn_kv speedup batched2 vs loop2-fused          1.15x
 NVFP4-g16  llama_attn_kv single eager (kernel off, 1 rep)        680.3 ms
 NVFP4-g16  llama_attn_kv speedup fused vs eager single           19.15x
 NVFP4-g16  llama_attn_kv speedup batched2 vs loop2-eager         21.99x
 NVFP4-g16  llama_mlp_gate_up hessian build (setup)                     5.7 ms
 NVFP4-g16  llama_mlp_gate_up batched B=1                             236.5 ms
 NVFP4-g16  llama_mlp_gate_up batched B=2                             467.2 ms
 NVFP4-g16  llama_mlp_gate_up single fused kernel                     237.3 ms
 NVFP4-g16  llama_mlp_gate_up loop2 fused (extrap.)                   474.6 ms
 NVFP4-g16  llama_mlp_gate_up speedup batched2 vs loop2-fused          1.02x
 NVFP4-g16  llama_mlp_gate_up single eager (kernel off, 1 rep)        894.4 ms
 NVFP4-g16  llama_mlp_gate_up speedup fused vs eager single            3.77x
 NVFP4-g16  llama_mlp_gate_up speedup batched2 vs loop2-eager          3.83x
 NVFP4-g16  llama_mlp_down hessian build (setup)                   110.9 ms
 NVFP4-g16  llama_mlp_down batched B=1                             486.4 ms
 NVFP4-g16  llama_mlp_down single fused kernel                     484.5 ms
 NVFP4-g16  llama_mlp_down loop1 fused (extrap.)                   484.5 ms
 NVFP4-g16  llama_mlp_down speedup batched1 vs loop1-fused          1.00x
 NVFP4-g16  llama_mlp_down single eager (kernel off, 1 rep)       3010.9 ms
 NVFP4-g16  llama_mlp_down speedup fused vs eager single            6.21x
 NVFP4-g16  llama_mlp_down speedup batched1 vs loop1-eager          6.19x
diff --git a/phase_memory.py b/phase_memory.py
 import contextlib
 import functools
 import importlib

 import torch


 def start_allocation_history():
    baseline = torch.cuda.memory_allocated()
    torch.cuda.memory._record_memory_history(
        enabled="all", context="all", stacks="python", max_entries=1_000_000
    )
    return baseline


 def stop_allocation_history(baseline):
    snapshot = torch.cuda.memory._snapshot()
    torch.cuda.memory._record_memory_history(enabled=None)
    allocated = baseline
    peak = baseline
    peak_event = None
    for event in snapshot["device_traces"][torch.cuda.current_device()]:
        if event["action"] == "alloc":
            allocated += event["size"]
        elif event["action"] == "free_requested":
            allocated -= event["size"]
        if allocated > peak:
            peak = allocated
            peak_event = event

    if peak_event is None:
        return None
    return {
        "reconstructed_peak_mib": round(peak / 2**20, 1),
        "trigger_allocation_mib": round(peak_event["size"] / 2**20, 1),
        "frames": [
            {
                "file": frame["filename"],
                "line": frame["line"],
                "function": frame["name"],
            }
            for frame in peak_event.get("frames", [])[:12]
        ],
    }


 class PhaseMemoryTracker:
    def __init__(self):
        self.current_phase = "pipeline"
        self.phase_peaks = {}
        self.peak_allocated = 0
        self.peak_reserved = 0
        self.peak_allocated_phase = None
        self.peak_reserved_phase = None
        self._patches = []
        torch.cuda.reset_peak_memory_stats()

    def _record_segment(self):
        allocated = torch.cuda.max_memory_allocated()
        reserved = torch.cuda.max_memory_reserved()
        phase_peak = self.phase_peaks.setdefault(
            self.current_phase, {"allocated": 0, "reserved": 0}
        )
        phase_peak["allocated"] = max(phase_peak["allocated"], allocated)
        phase_peak["reserved"] = max(phase_peak["reserved"], reserved)
        if allocated > self.peak_allocated:
            self.peak_allocated = allocated
            self.peak_allocated_phase = self.current_phase
        if reserved > self.peak_reserved:
            self.peak_reserved = reserved
            self.peak_reserved_phase = self.current_phase

    def _switch(self, phase):
        self._record_segment()
        torch.cuda.reset_peak_memory_stats()
        previous = self.current_phase
        self.current_phase = phase
        return previous

    @contextlib.contextmanager
    def phase(self, phase):
        previous = self._switch(phase)
        try:
            yield
        finally:
            self._switch(previous)

    def patch(self, owner, name, phase):
        original = getattr(owner, name)

        @functools.wraps(original)
        def wrapped(*args, **kwargs):
            with self.phase(phase):
                return original(*args, **kwargs)

        setattr(owner, name, wrapped)
        self._patches.append((owner, name, original))

    def patch_sequential_batches(self, owner):
        original = owner._get_batches

        @functools.wraps(original)
        def wrapped(*args, **kwargs):
            desc = kwargs.get("desc", args[3] if len(args) > 3 else "")
            phase = (
                "propagation_forward"
                if "Propagating" in desc
                else "calibration_forward"
            )
            for batch in original(*args, **kwargs):
                # The generator remains suspended inside this context while its
                # caller executes subgraph.forward(model, **inputs).
                with self.phase(phase):
                    yield batch

        owner._get_batches = wrapped
        self._patches.append((owner, "_get_batches", original))

    def stop(self):
        self._record_segment()
        for owner, name, original in reversed(self._patches):
            setattr(owner, name, original)
        self._patches.clear()

    def result(self):
        return {
            "peak_allocated_mib": round(self.peak_allocated / 2**20, 1),
            "peak_allocated_phase": self.peak_allocated_phase,
            "peak_reserved_mib": round(self.peak_reserved / 2**20, 1),
            "peak_reserved_phase": self.peak_reserved_phase,
            "phase_peaks_mib": {
                phase: {
                    "allocated": round(peaks["allocated"] / 2**20, 1),
                    "reserved": round(peaks["reserved"] / 2**20, 1),
                }
                for phase, peaks in self.phase_peaks.items()
            },
        }


 def install_phase_tracker():
    oneshot_module = importlib.import_module("llmcompressor.entrypoints.oneshot")
    from llmcompressor.modifiers.gptq import base as base_module
    from llmcompressor.modifiers.gptq import gptq_quantize as quantize_module
    from llmcompressor.pipelines.sequential import pipeline as sequential_pipeline

    tracker = PhaseMemoryTracker()
    tracker.patch(oneshot_module, "linearize_moe", "moe_linearization")
    modifier = base_module.GPTQModifier
    tracker.patch(modifier, "on_initialize", "initialization")
    tracker.patch(modifier, "calibrate_module", "calibration")
    tracker.patch(modifier, "on_sequential_epoch_end", "qparam_setup")
    tracker.patch(modifier, "on_calibration_end", "calibration_end")
    tracker.patch(modifier, "on_finalize", "finalization")
    tracker.patch(modifier, "compress_module_list", "batch_setup")
    tracker.patch_sequential_batches(sequential_pipeline)
    tracker.patch(sequential_pipeline, "trace_subgraphs", "subgraph_tracing")

    if hasattr(modifier, "_compress_batch"):
        tracker.patch(modifier, "_compress_batch", "quantization")
        if hasattr(base_module, "_apply_activation_ordering"):
            tracker.patch(base_module, "_apply_activation_ordering", "permutation")
    else:
        tracker.patch(base_module, "quantize_weight", "quantization")
        tracker.patch(
            quantize_module, "_apply_activation_ordering", "permutation"
        )

    return tracker
diff --git a/run_comparison.sh b/run_comparison.sh
 #!/usr/bin/env bash
 set -uo pipefail

 script_dir=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
 pr_repo=${PR_REPO:-$(cd -- "$script_dir/.." && pwd)}
 main_repo=${MAIN_REPO:-/tmp/llm-compressor-main}
 benchmark=${1:-all}
 log_dir=${LOG_DIR:-/tmp/gptq-memory-benchmarks}
 mkdir -p "$log_dir"

 run_case() {
    local script=$1
    local root=$2
    local label=$3
    shift 3
    local case_log="$log_dir/${script%.py}_${label}.log"

    printf '===== %s / %s =====\n' "$script" "$label"
    (
        cd "$root"
        PYTHONPATH="$root/src" python "$script_dir/$script" \
            --label "$label" --debug-phases "$@"
    ) > "$case_log" 2>&1
    local status=$?
    tail -n 1 "$case_log"
    if [[ $status -ne 0 ]]; then
        printf 'status=%s; last errors:\n' "$status"
        tail -n 30 "$case_log"
    fi
 }

 run_suite() {
    local script=$1
    run_case "$script" "$pr_repo" triton_batched --backend triton --batching
    run_case "$script" "$pr_repo" triton_unbatched --backend triton
    run_case "$script" "$pr_repo" eager_batched --backend eager --batching
    run_case "$script" "$pr_repo" eager_unbatched --backend eager
    run_case "$script" "$pr_repo" triton_linear_targets \
        --backend triton --linear-sequential-targets
 }

 case "$benchmark" in
    llama3)
        run_suite bench_llama3.py
        ;;
    qwen3-moe)
        run_suite bench_qwen3_moe.py
        ;;
    qwen3-full)
        run_suite bench_qwen3_full.py
        ;;
    all)
        run_suite bench_llama3.py
        run_suite bench_qwen3_moe.py
        ;;
    *)
        printf 'usage: %s [all|llama3|qwen3-moe|qwen3-full]\n' "$0" >&2
        exit 2
        ;;
 esac
diff --git a/weight_hashes.py b/weight_hashes.py
 import hashlib
 import json

 import torch


 def write_quantized_weight_hashes(model, output_path):
    hashes = {}
    for name, module in model.named_modules():
        scheme = getattr(module, "quantization_scheme", None)
        if scheme is None or getattr(scheme, "weights", None) is None:
            continue
        weight = getattr(module, "weight", None)
        if weight is None:
            continue
        value = weight.detach().contiguous().cpu().view(torch.uint8).numpy()
        hashes[name] = {
            "shape": list(weight.shape),
            "dtype": str(weight.dtype),
            "sha256": hashlib.sha256(value.tobytes()).hexdigest(),
        }

    with open(output_path, "w") as output:
        json.dump(hashes, output, sort_keys=True)
Implementation	Hessian offload	Batching	Peak allocated	Peak reserved	Runtime
PR	No	No	19.20 GiB	20.38 GiB	17.83 s
PR	Yes	No	19.20 GiB	20.06 GiB	139.19 s
PR	No	Yes	19.20 GiB	20.58 GiB	19.81 s
PR	Yes	Yes	19.20 GiB	20.48 GiB	140.77 s
Main	No	N/A	19.31 GiB	20.58 GiB	264.04 s
Main	Yes	N/A	19.31 GiB	20.06 GiB	380.67 s
Implementation	Hessian offload	Batching	Compression setup	Permutation	Quantization	Peak phase
PR	No	No	17.89 GiB	18.43 GiB	19.20 GiB	Quantization
PR	Yes	No	17.89 GiB	18.43 GiB	19.20 GiB	Quantization
PR	No	Yes	18.01 GiB	18.44 GiB	19.20 GiB	Quantization
PR	Yes	Yes	18.01 GiB	18.44 GiB	19.20 GiB	Quantization
Main	No	N/A	16.79 GiB	17.99 GiB	19.31 GiB	Quantization
Main	Yes	N/A	16.79 GiB	17.99 GiB	19.31 GiB	Quantization
Implementation	Hessian offload	Batching	Peak GPU allocated	Peak GPU reserved	Peak CPU RSS	Runtime
PR	No	No	7.01 GiB	7.25 GiB	3.01 GiB	7.45 s
PR	Yes	No	3.45 GiB	3.61 GiB	10.72 GiB	17.05 s
PR	No	Yes	20.18 GiB	28.42 GiB	2.95 GiB	4.44 s
PR	Yes	Yes	15.91 GiB	24.13 GiB	13.59 GiB	16.11 s
Main	No	N/A	7.02 GiB	7.30 GiB	2.36 GiB	133.23 s
Main	Yes	N/A	3.45 GiB	3.61 GiB	10.62 GiB	146.29 s
Implementation	Hessian offload	Batching	Compression setup	Permutation	Quantization	Overall peak phase
PR	No	No	6.92 GiB	6.95 GiB	7.01 GiB	Quantization
PR	Yes	No	2.63 GiB	2.66 GiB	2.73 GiB	MoE linearization (3.45 GiB)
PR	No	Yes	17.67 GiB	20.17 GiB	20.18 GiB	Quantization
PR	Yes	Yes	13.40 GiB	15.90 GiB	15.91 GiB	Quantization
Main	No	N/A	6.80 GiB	6.92 GiB	7.02 GiB	Quantization
Main	Yes	N/A	2.51 GiB	2.64 GiB	2.73 GiB	MoE linearization (3.45 GiB)
Implementation	Hessian offload	Batching	Peak GPU allocated	Peak GPU reserved	Runtime
PR after in-place permutation	No	No	6.95 GiB	7.16 GiB	7.59 s
PR after in-place permutation	Yes	No	3.45 GiB	3.61 GiB	16.49 s
PR after in-place permutation	No	Yes	16.18 GiB	20.66 GiB	4.96 s
PR after in-place permutation	Yes	Yes	15.90 GiB	20.14 GiB	14.11 s
Main	No	N/A	7.02 GiB	7.30 GiB	133.23 s
Main	Yes	N/A	3.45 GiB	3.61 GiB	146.29 s
Implementation	Hessian offload	Batching	Batch setup	Quantization	Overall peak phase
PR after in-place permutation	No	No	6.82 GiB	6.95 GiB	Quantization
PR after in-place permutation	Yes	No	2.53 GiB	2.66 GiB	MoE linearization (3.45 GiB)
PR after in-place permutation	No	Yes	10.76 GiB	16.18 GiB	Quantization
PR after in-place permutation	Yes	Yes	10.48 GiB	15.90 GiB	Quantization
Main	No	N/A	6.80 GiB	7.02 GiB	Quantization
Main	Yes	N/A	2.51 GiB	2.73 GiB	MoE linearization (3.45 GiB)
Implementation	Hessian offload	Batching	Peak GPU allocated	Peak GPU reserved	Runtime
PR after in-place permutation	No	No	18.43 GiB	19.61 GiB	17.75 s
PR after in-place permutation	Yes	No	18.43 GiB	19.30 GiB	127.80 s
PR after in-place permutation	No	Yes	18.43 GiB	19.82 GiB	19.80 s
PR after in-place permutation	Yes	Yes	18.43 GiB	19.71 GiB	131.66 s
Main	No	N/A	19.31 GiB	20.58 GiB	264.04 s
Main	Yes	N/A	19.31 GiB	20.06 GiB	380.67 s
Implementation	Hessian offload	Batching	Batch setup	Quantization	Peak phase
PR after in-place permutation	No	No	16.97 GiB	18.43 GiB	Quantization
PR after in-place permutation	Yes	No	16.97 GiB	18.43 GiB	Quantization
PR after in-place permutation	No	Yes	17.25 GiB	18.43 GiB	Quantization
PR after in-place permutation	Yes	Yes	17.25 GiB	18.43 GiB	Quantization
Implementation	Hessian offload	Batching	Peak GPU allocated	Peak GPU reserved	Runtime
PR exact parity	No	No	19.20 GiB	20.38 GiB	18.20 s
PR exact parity	Yes	No	19.20 GiB	20.06 GiB	133.54 s
PR exact parity	No	Yes	19.20 GiB	20.94 GiB	17.55 s
PR exact parity	Yes	Yes	19.20 GiB	20.06 GiB	133.36 s
Main	No	N/A	19.31 GiB	20.58 GiB	266.39 s
Main	Yes	N/A	19.31 GiB	20.06 GiB	408.70 s
Implementation	Hessian offload	Batching	Peak GPU allocated	Peak GPU reserved	Peak CPU RSS	Runtime
PR exact parity	No	No	7.01 GiB	7.23 GiB	2.37 GiB	7.39 s
PR exact parity	Yes	No	3.45 GiB	3.61 GiB	10.50 GiB	16.31 s
PR exact parity	No	Yes	16.18 GiB	20.73 GiB	2.94 GiB	5.87 s
PR exact parity	Yes	Yes	15.89 GiB	20.14 GiB	10.44 GiB	15.96 s
Main	No	N/A	7.02 GiB	7.30 GiB	2.93 GiB	128.42 s
Main	Yes	N/A	3.45 GiB	3.61 GiB	10.14 GiB	146.28 s