PyTorch MPS 경로, 그리고 부서지는 자리

#!/usr/bin/env python3
"""The same bf16 checkpoint decoded two ways on one Mac: PyTorch on MPS through
transformers, and mlx-lm. Same bytes per token, same pool, two paths."""
import sys, time
MODEL = "unsloth/Llama-3.2-1B-Instruct"     # bf16 safetensors, tied embeddings
PROMPT = "Explain, in about 300 words, why unified memory matters for language models."
N = 128


def torch_mps():
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer
    tok = AutoTokenizer.from_pretrained(MODEL)
    model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16).to("mps").eval()
    msgs = [{"role": "user", "content": PROMPT}]
    ids = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt", return_dict=True)["input_ids"].to("mps")
    nbytes = sum(p.numel() * p.element_size() for p in model.parameters())
    with torch.no_grad():
        model.generate(ids, max_new_tokens=8, do_sample=False)          # warm-up
        torch.mps.synchronize(); t0 = time.perf_counter()
        out = model.generate(ids, max_new_tokens=1, do_sample=False)
        torch.mps.synchronize(); ttft = time.perf_counter() - t0
        t0 = time.perf_counter()
        out = model.generate(ids, max_new_tokens=N, do_sample=False, min_new_tokens=N)
        torch.mps.synchronize(); dt = time.perf_counter() - t0
    gen = out.shape[1] - ids.shape[1]
    print(f"torch {torch.__version__} MPS via transformers: {nbytes/1e9:.2f} GB of bf16 weights; prompt {ids.shape[1]} tokens; "
          f"TTFT {ttft:.3f}s; {gen} tokens in {dt:.2f}s = {gen/(dt-ttft):.1f} tok/s (decode, TTFT removed)")


def mlx_path():
    import mlx.core as mx
    from mlx_lm import load, stream_generate
    from mlx_lm.sample_utils import make_sampler
    model, tok = load(MODEL)                                             # HF bf16 loaded as-is
    msgs = [{"role": "user", "content": PROMPT}]
    prompt = tok.apply_chat_template(msgs, add_generation_prompt=True)
    sampler = make_sampler(temp=0.0)
    for _ in stream_generate(model, tok, prompt, max_tokens=8, sampler=sampler): pass   # warm-up
    last = None
    for r in stream_generate(model, tok, prompt, max_tokens=N, sampler=sampler): last = r
    print(f"mlx {mx.__version__} mlx-lm: prompt {len(prompt)} tokens; prompt {last.prompt_tps:.0f} tok/s; "
          f"decode {last.generation_tps:.1f} tok/s; peak {last.peak_memory:.2f} GB")


if __name__ == "__main__":
    {"torch": torch_mps, "mlx": mlx_path}[sys.argv[1]]()

# office, M3 Ultra, 2026-09-15:
#   torch 2.11.0 MPS via transformers 5.14.1: 2.47 GB bf16; prompt 52; TTFT 0.024s; 42.1 tok/s, then 31.8 on a second run
#   mlx 0.32.2 / mlx-lm 0.31.3:                prompt 2350 tok/s; decode 190.0 tok/s; peak 2.56 GB
#   ceiling for 2.47 GB per token at 638 GB/s: 258 tok/s

#!/usr/bin/env python3
"""Run with PYTORCH_ENABLE_MPS_FALLBACK=1 so the missing op falls back instead of raising."""
import time, torch


def timeit(fn, n=3):
    fn(); torch.mps.synchronize()
    t = time.perf_counter()
    for _ in range(n):
        fn()
    torch.mps.synchronize()
    return (time.perf_counter() - t) / n


x = torch.randn(8192, 8192, device="mps")
for name, a in (("fp32", x), ("fp16", x.half()), ("bf16", x.bfloat16())):
    dt = timeit(lambda: a @ a, 5)
    print(f"matmul 8192^2 {name} on mps: {2*8192**3/dt/1e12:.2f} TFLOP/s")

W = torch.randn(32768, 32768, device="mps", dtype=torch.bfloat16); v = torch.randn(32768, device="mps", dtype=torch.bfloat16)
dt = timeit(lambda: W @ v, 10)
print(f"matvec over {W.numel()*2/1e9:.2f} GB bf16 on mps: {dt*1e3:.2f} ms = {W.numel()*2/dt/1e9:.0f} GB/s")

s = torch.randn(2048, 2048); s = s @ s.T
s_mps = s.to("mps")
print(f"eigh 2048^2 on cpu:                      {timeit(lambda: torch.linalg.eigh(s))*1e3:6.1f} ms")
print(f"eigh 2048^2 on mps (CPU fallback, env):  {timeit(lambda: torch.linalg.eigh(s_mps))*1e3:6.1f} ms  -> result on {torch.linalg.eigh(s_mps)[0].device}")

try:
    torch.randn(64, 64, device="mps", dtype=torch.float64)
except TypeError as e:
    print("float64:", e)

# office, torch 2.12.0, 2026-09-15:
# matmul 8192^2 fp32 18.05 / fp16 21.09 / bf16 21.88 TFLOP/s
# matvec over 2.15 GB bf16: 3.55 ms = 605 GB/s
# eigh 2048^2 on cpu 371.3 ms; on mps with fallback 350.2 ms -> result on mps:0
# float64: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead.
# without the env var, eigh raises: "The operator 'aten::_linalg_eigh.eigenvalues' is not currently implemented for the MPS device ..."