[torchao] W8A8-INT GPU Profiling Result

log.txt

repro:
```
import gc, time, torch
from collections import defaultdict
from contextlib import contextmanager
from transformers import AutoModelForCausalLM, AutoTokenizer
from torchao.quantization import Int8DynamicActivationInt8WeightConfig, quantize_
import warnings, logging

warnings.filterwarnings("ignore", category=UserWarning, module="torchao")
logging.getLogger("torch._inductor.select_algorithm").setLevel(logging.CRITICAL)

MODEL_ID = "Qwen/Qwen3-8B"
PROMPT = "Quantization reduces memory footprint by" * 128
WARMUP = 3
mb = lambda n: n / 1024**2


@contextmanager
def timed_hooks(model):
    """Attach per-layer sync timers; yields a dict {name: [ms, …]}."""
    pending, times, handles = {}, defaultdict(list), []

    def attach(mod, name):
        def pre(m, _):
            torch.cuda.synchronize()
            pending[id(m)] = time.perf_counter()
        def post(m, _, __):
            torch.cuda.synchronize()
            times[name].append((time.perf_counter() - pending.pop(id(m))) * 1e3)
        handles.extend([mod.register_forward_pre_hook(pre), mod.register_forward_hook(post)])

    attach(model.model.embed_tokens, "embedding")
    for layer in model.model.layers:
        attach(layer.self_attn, "attention")
        attach(layer.mlp, "ffn")
    attach(model.lm_head, "lm_head")
    try:
        yield times
    finally:
        for h in handles:
            h.remove()


def kv_mb(past):
    if past is None:
        return 0.0
    pairs = zip(past.key_cache, past.value_cache) if hasattr(past, "key_cache") else past
    return mb(sum(t.numel() * t.element_size() for kv in pairs for t in kv if t is not None))


def run(label, tokenizer, quantize_fn=None):
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID, dtype=torch.bfloat16, device_map="cuda").eval()
    if quantize_fn:
        quantize_fn(model)

    # Hooks must be attached BEFORE compile so tracer sees graph-break points
    with timed_hooks(model) as times:
        model = torch.compile(model)
        torch.cuda.synchronize(); torch.cuda.empty_cache()
        w_mb = mb(torch.cuda.memory_allocated())
        inputs = tokenizer(PROMPT, return_tensors="pt").to("cuda")

        with torch.no_grad():
            for _ in range(WARMUP):
                model(**inputs)
            times.clear()  # discard warmup
            torch.cuda.synchronize(); torch.cuda.reset_peak_memory_stats()
            out = model(**inputs)

    timing = {k: sum(v) for k, v in times.items()}
    k = kv_mb(out.past_key_values)
    p = mb(torch.cuda.max_memory_allocated())
    result = dict(label=label, timing=timing,
                  weight_mb=w_mb, kv_mb=k, act_mb=max(0, p - w_mb - k), peak_mb=p)
    del model; gc.collect(); torch.cuda.empty_cache()
    return result


def print_results(results):
    C, labels = 16, [r["label"] for r in results]
    hdr = lambda t: f"\n{t:<20}" + "".join(f"{l:>{C}}" for l in labels)
    row = lambda l, v: print(f"{l:<20}" + "".join(f"{x:>{C}}" for x in v))
    sep = lambda: print("-" * (20 + C * len(results)))

    print(hdr("Compute Time (ms)")); sep()
    for k in ("embedding", "attention", "ffn", "lm_head"):
        row(k, [f"{r['timing'].get(k, 0):.1f}" for r in results])
    sep()
    totals = [sum(r["timing"].values()) for r in results]
    row("total",   [f"{t:.1f}" for t in totals])
    row("speedup", ["1.00×"] + [f"{totals[0]/t:.2f}×" for t in totals[1:]])

    print(hdr("Memory (MiB)")); sep()
    for key, lbl in [("weight_mb","weights"), ("kv_mb","kv cache"),
                     ("act_mb","activations"), ("peak_mb","peak alloc")]:
        row(lbl, [f"{r[key]:.1f}" for r in results])


if __name__ == "__main__":
    print(f"GPU  : {torch.cuda.get_device_name(0)}")
    print(f"Model: {MODEL_ID} | warmup: {WARMUP}")
    tok = AutoTokenizer.from_pretrained(MODEL_ID)
    print_results([
        run("BF16", tok),
        run("INT8 W8A8", tok, lambda m: quantize_(m, Int8DynamicActivationInt8WeightConfig())),
    ])
```


result:
```
Compute Time (ms)               BF16       INT8 W8A8
----------------------------------------------------
embedding                        0.1             0.1
attention                       38.0           230.3
ffn                             56.2            33.7
lm_head                          5.8             2.8
----------------------------------------------------
total                          100.1           266.9
speedup                        1.00×           0.38×

Memory (MiB)                    BF16       INT8 W8A8
----------------------------------------------------
weights                      15622.6          8416.7
kv cache                       108.0           108.0
activations                    237.2           684.1
peak alloc                   15967.8          9208.8
```