File size: 3,003 Bytes
9a57d8a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | """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)
|