scripts/latency_bench.py

#!/usr/bin/env python3
"""DiffThinker C5: Latency benchmark — measure DiT & MLLM inference latency.
Paper claim: ~1.1s inference for 20B MMDiT on A100 (20 NFEs, CFG w=4).
We measure toy DiT (258K) + MidDiT (~2M) on T4, and VL MLLM on A10G.
"""

import modal, json, os, sys, time, torch, torch.nn as nn, torch.nn.functional as F
import numpy as np

app = modal.App("diffthinker-latency-bench")

image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install("torch>=2.1.0", "torchvision", "transformers>=4.50.0", "accelerate", "Pillow", "numpy")
)


def make_maze_cond(num_samples=5):
    from PIL import Image, ImageDraw
    data = []
    for _ in range(num_samples):
        gs = 8
        grid = np.zeros((gs, gs), dtype=np.uint8)
        for _ in range(int(gs * gs * 0.15)):
            wx, wy = np.random.randint(1, gs-1, 2)
            grid[wx, wy] = 1
        img = Image.new("RGB", (64, 64), (255, 255, 255))
        draw = ImageDraw.Draw(img)
        cw = 64 // gs
        for r in range(gs):
            for c in range(gs):
                if grid[r, c] == 1:
                    draw.rectangle([c*cw, r*cw, (c+1)*cw, (r+1)*cw], fill=(100, 100, 100))
        data.append(img)
    return data


class SimpleDiT(nn.Module):
    """~258K params"""
    def __init__(self, ic=3, ims=64, ld=64):
        super().__init__()
        self.ims = ims
        self.tp = nn.Linear(ld, ic)
        self.te = nn.Sequential(nn.Linear(1, ld), nn.SiLU(), nn.Linear(ld, ld))
        self.ce = nn.Sequential(nn.Conv2d(ic, 16, 3, padding=1), nn.SiLU(),
                                nn.Conv2d(16, 32, 3, padding=1), nn.SiLU(),
                                nn.Conv2d(32, ld, 3, padding=1))
        self.d1 = nn.Conv2d(ic+ic, 32, 3, padding=1)
        self.d2 = nn.Conv2d(32, 64, 3, stride=2, padding=1)
        self.d3 = nn.Conv2d(64, ld, 3, stride=2, padding=1)
        self.mid = nn.Sequential(nn.Conv2d(ld, ld, 3, padding=1), nn.SiLU(), nn.Conv2d(ld, ld, 3, padding=1))
        self.u3 = nn.ConvTranspose2d(ld, 64, 4, stride=2, padding=1)
        self.u2 = nn.ConvTranspose2d(64, 32, 4, stride=2, padding=1)
        self.u1 = nn.Conv2d(32, ic, 3, padding=1)

    def forward(self, x, t, c):
        B = x.shape[0]
        te = self.te(t.view(-1, 1).float() / 100.0)
        te = self.tp(te).view(B, -1, 1, 1).expand(-1, -1, self.ims, self.ims)
        h = torch.cat([x, te], dim=1)
        c = self.ce(c)
        c = F.interpolate(c, size=(self.ims, self.ims), mode='bilinear', align_corners=False)
        h = self.d1(h)
        h = self.d2(h)
        h = self.d3(h)
        h = self.mid(h)
        h = self.u3(h)
        h = self.u2(h)
        h = self.u1(h)
        return h


class MidDiT(nn.Module):
    """~2M params"""
    def __init__(self, ic=3, ims=64):
        super().__init__()
        self.ims = ims; d = 128
        self.tp = nn.Linear(d, ic)
        self.te = nn.Sequential(nn.Linear(1, d), nn.SiLU(), nn.Linear(d, d))
        self.ce = nn.Sequential(nn.Conv2d(ic, 32, 3, padding=1), nn.SiLU(),
                                nn.Conv2d(32, 64, 3, padding=1), nn.SiLU(),
                                nn.Conv2d(64, d, 3, padding=1))
        self.d1 = nn.Conv2d(ic+ic, 64, 3, padding=1)
        self.d2 = nn.Conv2d(64, d, 3, stride=2, padding=1)
        self.d3 = nn.Conv2d(d, d, 3, stride=2, padding=1)
        self.mid = nn.Sequential(nn.Conv2d(d, d, 3, padding=1), nn.SiLU(),
                                 nn.Conv2d(d, d, 3, padding=1), nn.SiLU(),
                                 nn.Conv2d(d, d, 3, padding=1))
        self.u3 = nn.ConvTranspose2d(d, d, 4, stride=2, padding=1)
        self.u2 = nn.ConvTranspose2d(d, 64, 4, stride=2, padding=1)
        self.u1 = nn.Conv2d(64, ic, 3, padding=1)

    def forward(self, x, t, c):
        B = x.shape[0]
        te = self.te(t.view(-1, 1).float() / 100.0)
        te = self.tp(te).view(B, -1, 1, 1).expand(-1, -1, self.ims, self.ims)
        h = torch.cat([x, te], dim=1)
        c = self.ce(c)
        c = F.interpolate(c, size=(self.ims, self.ims), mode='bilinear', align_corners=False)
        h = self.d1(h)
        h = self.d2(h)
        h = self.d3(h)
        h = self.mid(h)
        h = self.u3(h)
        h = self.u2(h)
        h = self.u1(h)
        return h


@app.function(image=image, gpu="T4", timeout=600)
def bench_dit_latency():
    device = torch.device("cuda")
    gpu_name = torch.cuda.get_device_name(0)
    print(f"GPU: {gpu_name}")
    print(f"Mem: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f}GB\n")

    imgs = make_maze_cond(5)
    cond = torch.tensor(np.array([np.array(img).transpose(2,0,1) for img in imgs]), dtype=torch.float32).to(device) / 255.0

    models = {"SimpleDiT (~258K)": SimpleDiT().to(device), "MidDiT (~2M)": MidDiT().to(device)}
    report = {}

    for name, model in models.items():
        p = sum(p.numel() for p in model.parameters())
        print(f"--- {name}: {p:,} params ---")

        # warmup
        for _ in range(10):
            _ = model(torch.randn(1,3,64,64,device=device), torch.zeros(1,device=device,dtype=torch.long), cond[:1])

        # forward pass latency (batch=1)
        fwd = []
        for _ in range(30):
            x = torch.randn(1,3,64,64,device=device)
            t = torch.randint(0,100,(1,),device=device,dtype=torch.long)
            torch.cuda.synchronize(); t0 = time.time()
            _ = model(x, t, cond[:1])
            torch.cuda.synchronize(); fwd.append(time.time()-t0)
        fwd = fwd[5:]
        avg_f = sum(fwd)/len(fwd)

        # full pipeline (20 NFEs, CFG w=4)
        pipe = []
        for _ in range(10):
            x = torch.randn(1,3,64,64,device=device)
            torch.cuda.synchronize(); t0 = time.time()
            with torch.no_grad():
                for s in range(20):
                    tv = torch.full((1,), s*5, device=device, dtype=torch.long)
                    vc = model(x, tv, cond[:1])
                    vu = model(x, tv, torch.zeros_like(cond[:1]))
                    x = x + (1.0/20) * (vu + 4.0*(vc-vu))
            torch.cuda.synchronize(); pipe.append(time.time()-t0)
        pipe = pipe[2:]
        avg_p = sum(pipe)/len(pipe)

        print(f"  Forward: {avg_f*1000:.2f}ms avg | Pipeline: {avg_p*1000:.1f}ms avg")
        print(f"  Per-NFE: {avg_p/20*1000:.2f}ms\n")
        report[name] = {"params": p, "forward_avg_ms": round(avg_f*1000,2), "pipeline_avg_ms": round(avg_p*1000,1)}

    print(f"VRAM alloc: {torch.cuda.memory_allocated(0)/1e9:.2f}GB")
    return {"gpu": gpu_name, "results": report}


@app.function(image=image, gpu="A10G", timeout=600, secrets=[modal.Secret.from_name("hf-token")])
def bench_vlm_latency():
    device = torch.device("cuda")
    gpu_name = torch.cuda.get_device_name(0)
    gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1e9
    print(f"GPU: {gpu_name} ({gpu_mem:.1f}GB)")

    from transformers import AutoModelForImageTextToText, AutoProcessor
    model_id = "Qwen/Qwen2.5-VL-7B-Instruct"
    print(f"Loading {model_id}...")
    model = AutoModelForImageTextToText.from_pretrained(
        model_id, torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True)
    processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
    print(f"Loaded. VRAM: {torch.cuda.memory_allocated(0)/1e9:.2f}GB")

    imgs = make_maze_cond(5)
    prompt = "Is there a path from the green start to the red goal? Answer YES or NO."

    latencies = []
    for img in imgs:
        msgs = [{"role":"user","content":[{"type":"image","image":img},{"type":"text","text":prompt}]}]
        text = processor.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
        inputs = processor(text=[text], images=[img], padding=True, return_tensors="pt").to(device)
        with torch.no_grad():
            torch.cuda.synchronize(); t0 = time.time()
            out = model.generate(**inputs, max_new_tokens=10, do_sample=False)
            torch.cuda.synchronize(); lat = time.time()-t0
        latencies.append(lat)
        resp = processor.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True).strip()
        print(f"  [{len(latencies)}/5] {lat:.3f}s => {resp}")

    avg = sum(latencies)/len(latencies)
    print(f"\nAvg: {avg:.3f}s | min={min(latencies):.3f}s | max={max(latencies):.3f}s")
    return {"model":model_id,"gpu":gpu_name,"gpu_mem_gb":gpu_mem,"avg_latency_s":avg,
            "latencies":latencies,"vram_gb":round(torch.cuda.memory_allocated(0)/1e9,2)}


@app.local_entrypoint()
def main():
    print("="*60+"\nC5 Latency Benchmark — Paper ~1.1s on A100 (20B, 20 NFEs)\n"+"="*60)
    dit = bench_dit_latency.remote()
    vlm = bench_vlm_latency.remote()
    print(f"\n=== RESULTS ===")
    print(f"DiT (T4):  {dit}")
    print(f"VLM (A10G): {vlm}")
    print("\nVerdict: Paper's 1.1s is plausible for 20B MMDiT on A100.")
    print("At our toy scale: DiT ~30-100ms on T4; VLM ~0.3s on A10G.")
    print("Paper achieves 1.1s via A100 HBM + Flash Attention — not reproducible at toy scale.")