| """Generation speed on one GPU: DiffuRefill-1B vs autoregressive baselines. |
| |
| Same job for every model: a 32-token prompt, 256 new tokens, greedy, batch 1 and |
| batch 32, bf16; tokens/s counts only new tokens, after a warm-up run. |
| |
| ours K tokens committed per forward pass (the K most confident masked |
| positions, left to right within a 32-token block). K=1 is the |
| one-token-per-pass order that keeps text coherent; larger K trades |
| quality for passes. No KV cache: every pass re-reads the full row. |
| AR Hugging Face generate() with its KV cache. |
| """ |
| import os, sys, time, json |
| import torch |
| import torch.nn.functional as F |
|
|
| MODE = os.environ.get("SP_MODE", "ours") |
| NEW, PROMPT = 256, 32 |
| OUT = os.environ.get("SP_OUT", "/tmp/speed.json") |
|
|
|
|
| def timed(fn): |
| torch.cuda.synchronize(); t = time.perf_counter(); fn(); torch.cuda.synchronize() |
| return time.perf_counter() - t |
|
|
|
|
| res = {} |
| if MODE == "ours": |
| os.environ.setdefault("BG_DOCMASK", "0") |
| sys.path.insert(0, "/tmp") |
| from big_common import Net, MASK_ID, VOCAB |
| from huggingface_hub import hf_hub_download |
| g = torch.load(hf_hub_download("Asilarkness/DiffuRefill-1B", "final/DiffuRefill-1B-200k-swa.pt", local_dir="/tmp/dlm"), |
| map_location="cpu", weights_only=False) |
| net = Net(); net.load_state_dict({k: v.float() for k, v in g["model"].items()}) |
| net = net.cuda().to(torch.bfloat16).eval() |
|
|
| @torch.no_grad() |
| def gen(B, K, block=32): |
| x = torch.randint(0, 50000, (B, PROMPT + NEW), device="cuda") |
| x[:, PROMPT:] = MASK_ID |
| for b0 in range(PROMPT, PROMPT + NEW, block): |
| for _ in range(block // K): |
| lg = net(x)[:, b0:b0 + block, :VOCAB - 2].float() |
| conf, tok = F.softmax(lg, -1).max(-1) |
| masked = x[:, b0:b0 + block] == MASK_ID |
| conf = torch.where(masked, conf, torch.full_like(conf, -1)) |
| idx = conf.topk(K, -1).indices |
| blk = x[:, b0:b0 + block] |
| blk.scatter_(1, idx, tok.gather(1, idx)) |
| return x |
|
|
| for B in (1, 32): |
| for K in (1, 2, 4, 8): |
| gen(B, K) |
| t = timed(lambda: gen(B, K)) |
| res["ours B%d K%d" % (B, K)] = B * NEW / t |
| print("SPEED ours B=%d K=%d: %.0f tok/s (%.2fs)" % (B, K, B * NEW / t, t), flush=True) |
| else: |
| from transformers import AutoModelForCausalLM |
| m = AutoModelForCausalLM.from_pretrained(MODE, torch_dtype=torch.bfloat16).cuda().eval() |
| for B in (1, 32): |
| ids = torch.randint(100, 20000, (B, PROMPT), device="cuda") |
| kw = dict(max_new_tokens=NEW, min_new_tokens=NEW, do_sample=False, pad_token_id=0) |
| with torch.no_grad(): |
| m.generate(ids, **kw) |
| t = timed(lambda: m.generate(ids, **kw)) |
| res["%s B%d" % (MODE, B)] = B * NEW / t |
| print("SPEED %s B=%d: %.0f tok/s (%.2fs)" % (MODE, B, B * NEW / t, t), flush=True) |
| json.dump(res, open(OUT, "w")) |
| print("SPEED DONE", flush=True) |
|
|
