"""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)