probe_dit_throughput.py

#!/usr/bin/env python3
"""
probe_dit_throughput.py
Runs a 200-300M-param DiT-class architecture at 256px on a single CUDA GPU (target: RTX 4090),
binary-searches max batch size under 24GB VRAM, runs 100 training steps with CUDA timing
and nvidia-smi VRAM sampling, prints benchmark report.

Usage:
    pip install torch torchvision einops
    python probe_dit_throughput.py [--target-vram-gb 24] [--steps 100] [--resolution 256]
    python probe_dit_throughput.py --dry-run            # print config, no allocation, exit 0

Note: This does NOT require flash-attn. Uses torch.nn.functional.scaled_dot_product_attention
which dispatches to FlashAttention-2 on sm_89 automatically (PyTorch 2.2+).

On MPS / CPU (non-CUDA): prints a warning and runs a tiny smoke pass (batch=1, steps=2)
to verify the model code is functional. Full throughput numbers require CUDA.
"""
from __future__ import annotations

import argparse
import sys


# ---------------------------------------------------------------------------
# Argument parsing — intentionally BEFORE heavy imports so --help is instant
# ---------------------------------------------------------------------------

def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=(
            "DiT throughput probe for single-GPU 256px training. "
            "Targets CUDA (RTX 4090). On MPS/CPU a tiny smoke pass runs instead."
        ),
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=(
            "Examples:\n"
            "  python probe_dit_throughput.py --dry-run\n"
            "  python probe_dit_throughput.py --depth 18 --steps 50\n"
            "  python probe_dit_throughput.py --target-vram-gb 20 --resolution 512\n"
        ),
    )
    parser.add_argument(
        "--target-vram-gb", type=float, default=22.0,
        help="VRAM budget for batch search in GB (default 22.0 — leaves 2 GB headroom on a 24 GB 4090). Must be > 0.",
    )
    parser.add_argument(
        "--steps", type=int, default=100,
        help="Number of timed training steps (default 100). Must be > 0.",
    )
    parser.add_argument(
        "--resolution", type=int, default=256,
        help="Image resolution in pixels (default 256; try 512 for cliff measurement). Must be > 0.",
    )
    parser.add_argument(
        "--depth", type=int, default=18,
        help="Transformer depth (18 -> ~250M params, 24 -> ~330M params). Must be > 0.",
    )
    parser.add_argument(
        "--hidden", type=int, default=1024,
        help="Transformer hidden size (default 1024). Must be > 0.",
    )
    parser.add_argument(
        "--heads", type=int, default=16,
        help="Attention heads (default 16). Must be > 0.",
    )
    parser.add_argument(
        "--dry-run", action="store_true",
        help=(
            "Print the planned model config (depth/hidden/heads/resolution/target-vram) "
            "and validate all args without allocating tensors or running the benchmark loop. "
            "Exits 0 on success."
        ),
    )
    return parser


def _validate_args(args: argparse.Namespace) -> list[str]:
    """Return a list of validation error strings (empty = OK)."""
    errors: list[str] = []
    int_fields = [
        ("--steps", args.steps),
        ("--resolution", args.resolution),
        ("--depth", args.depth),
        ("--hidden", args.hidden),
        ("--heads", args.heads),
    ]
    for flag, val in int_fields:
        if val <= 0:
            errors.append(f"{flag} must be > 0 (got {val})")
    if args.target_vram_gb <= 0:
        errors.append(f"--target-vram-gb must be > 0 (got {args.target_vram_gb})")
    return errors


def _estimate_params_M(depth: int, hidden: int) -> float:
    """Rough param count in millions (matches empirical DiT formula)."""
    # Each block: norm1+norm2 (tiny), attn_qkv (3*hidden^2), attn_proj (hidden^2),
    # mlp 2x (hidden * 4*hidden), adaLN (hidden * 6*hidden)
    per_block = (3 * hidden * hidden + hidden * hidden
                 + 2 * hidden * 4 * hidden
                 + hidden * 6 * hidden)
    # Embeddings + final layer (rough)
    other = hidden * hidden * 8
    return (depth * per_block + other) / 1e6


def _dry_run(args: argparse.Namespace) -> None:
    """Print planned config and exit 0. No tensor allocation."""
    latent_res = args.resolution // 8
    patch_size = 2
    n_patches = (latent_res // patch_size) ** 2
    est_params = _estimate_params_M(args.depth, args.hidden)

    print("=" * 60)
    print("DiT Throughput Probe — DRY RUN (no allocation)")
    print("=" * 60)
    print(f"  depth            : {args.depth}")
    print(f"  hidden           : {args.hidden}")
    print(f"  heads            : {args.heads}")
    print(f"  resolution       : {args.resolution}px")
    print(f"  latent_res       : {latent_res}x{latent_res}")
    print(f"  n_patches        : {n_patches}")
    print(f"  est_params_M     : ~{est_params:.0f}M (rough formula)")
    print(f"  target_vram_gb   : {args.target_vram_gb:.1f} GB")
    print(f"  steps            : {args.steps}")
    print(f"  patch_size       : {patch_size}")
    print("=" * 60)
    print("[DRY RUN] Config validated. No tensors allocated. Exiting.")
    sys.exit(0)


# ---------------------------------------------------------------------------
# Heavy imports — deferred so --help / --dry-run never pay the import cost
# ---------------------------------------------------------------------------

def _import_heavy():
    """Import torch and einops; raise ImportError with a helpful message if missing."""
    missing = []
    try:
        import torch  # noqa: F401
    except ImportError:
        missing.append("torch")
    try:
        import einops  # noqa: F401
    except ImportError:
        missing.append("einops")
    if missing:
        raise ImportError(
            f"Missing required packages: {', '.join(missing)}. "
            "Install with: pip install torch einops"
        )
    import torch
    import einops
    return torch, einops


# ---------------------------------------------------------------------------
# Device selection with MPS/CPU graceful fallback
# ---------------------------------------------------------------------------

def _select_device(torch):
    """Return (device, is_cuda, backend_name)."""
    if torch.cuda.is_available():
        return torch.device("cuda"), True, "CUDA"
    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return torch.device("mps"), False, "MPS"
    else:
        return torch.device("cpu"), False, "CPU"


# ---------------------------------------------------------------------------
# Minimal DiT-class model (~250M params)
# Follows Peebles & Xie 2022 DiT structure: patchify -> N x (attn + MLP) -> unpatchify
# Numbers chosen to land at ~250M params:
#   depth=18, hidden=1024, heads=16, mlp_ratio=4  => ~250M
#   depth=24, hidden=1024, heads=16, mlp_ratio=4  => ~330M
# ---------------------------------------------------------------------------

def _build_model_classes(torch_mod, rearrange):
    """
    Build and return MinimalDiT class using the provided torch module and rearrange fn.
    This is called at runtime after heavy imports succeed.  The returned class is
    also exposed at module-level (see the ``try`` block below) so that test-suite
    imports via ``importlib`` can access ``MinimalDiT`` and ``SinusoidalEmbedding``
    as plain attributes without triggering the CUDA benchmark.
    """
    import math
    import torch.nn as nn
    import torch.nn.functional as F
    _torch = torch_mod

    class SinusoidalEmbedding(nn.Module):
        """Timestep sinusoidal embedding."""
        def __init__(self, dim: int, max_period: int = 10000):
            super().__init__()
            self.dim = dim
            half = dim // 2
            freqs = _torch.exp(
                -math.log(max_period) * _torch.arange(half, dtype=_torch.float32) / half
            )
            self.register_buffer("freqs", freqs)

        def forward(self, t: "torch.Tensor") -> "torch.Tensor":
            args = t[:, None].float() * self.freqs[None]
            return _torch.cat([_torch.cos(args), _torch.sin(args)], dim=-1)

    class DiTBlock(nn.Module):
        """Single DiT block: adaLN-Zero conditioning + self-attn + MLP."""
        def __init__(self, hidden: int, heads: int, mlp_ratio: float = 4.0):
            super().__init__()
            self.norm1 = nn.LayerNorm(hidden, elementwise_affine=False, eps=1e-6)
            self.norm2 = nn.LayerNorm(hidden, elementwise_affine=False, eps=1e-6)
            self.attn_qkv = nn.Linear(hidden, 3 * hidden, bias=True)
            self.attn_proj = nn.Linear(hidden, hidden, bias=True)
            mlp_hidden = int(hidden * mlp_ratio)
            self.mlp = nn.Sequential(
                nn.Linear(hidden, mlp_hidden, bias=True),
                nn.GELU(approximate='tanh'),
                nn.Linear(mlp_hidden, hidden, bias=True),
            )
            # adaLN modulation: 6 scalings per block (shift/scale for norm1/norm2 + gate_msa/gate_mlp)
            self.adaLN_modulation = nn.Sequential(
                nn.SiLU(),
                nn.Linear(hidden, 6 * hidden, bias=True),
            )
            self.heads = heads
            self.head_dim = hidden // heads

        def forward(self, x: "torch.Tensor", c: "torch.Tensor") -> "torch.Tensor":
            B, N, C = x.shape
            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = \
                self.adaLN_modulation(c).chunk(6, dim=-1)

            # Attention
            x_norm = self.norm1(x) * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1)
            qkv = self.attn_qkv(x_norm)
            q, k, v = rearrange(qkv, 'b n (three h d) -> three b h n d', three=3, h=self.heads).unbind(0)
            # Uses FlashAttention-2 on sm_89 via SDPA dispatch (PyTorch >= 2.2)
            attn_out = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0)
            attn_out = rearrange(attn_out, 'b h n d -> b n (h d)')
            x = x + gate_msa.unsqueeze(1) * self.attn_proj(attn_out)

            # MLP
            x_norm2 = self.norm2(x) * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1)
            x = x + gate_mlp.unsqueeze(1) * self.mlp(x_norm2)
            return x

    class MinimalDiT(nn.Module):
        """
        Minimal DiT: image latents -> patches -> transformer -> patches -> image latents.
        No text conditioning to keep the probe self-contained.
        Uses AdaLN-Zero with just timestep conditioning.
        """
        def __init__(
            self,
            resolution: int = 256,
            patch_size: int = 2,
            in_channels: int = 4,       # standard 4-ch latent (VAE f8)
            hidden: int = 1024,
            depth: int = 18,
            heads: int = 16,
            mlp_ratio: float = 4.0,
        ):
            super().__init__()
            self.patch_size = patch_size
            self.in_channels = in_channels
            self.hidden = hidden

            # Latent resolution after VAE (8x downsample)
            latent_res = resolution // 8
            # Number of patches
            n_patches = (latent_res // patch_size) ** 2
            self.n_patches = n_patches

            # Patch embedding
            self.patch_embed = nn.Conv2d(
                in_channels, hidden,
                kernel_size=patch_size, stride=patch_size, bias=True
            )
            self.pos_embed = nn.Parameter(_torch.zeros(1, n_patches, hidden))

            # Timestep embedding
            self.t_embed = nn.Sequential(
                SinusoidalEmbedding(hidden),
                nn.Linear(hidden, hidden, bias=True),
                nn.SiLU(),
                nn.Linear(hidden, hidden, bias=True),
            )

            # Transformer blocks
            self.blocks = nn.ModuleList([
                DiTBlock(hidden, heads, mlp_ratio) for _ in range(depth)
            ])

            # Final layer
            self.final_norm = nn.LayerNorm(hidden, elementwise_affine=False, eps=1e-6)
            self.final_linear = nn.Linear(hidden, patch_size * patch_size * in_channels * 2, bias=True)
            self.final_adaLN = nn.Sequential(nn.SiLU(), nn.Linear(hidden, 2 * hidden, bias=True))

            self._init_weights()

        def _init_weights(self):
            # Zero-init the adaLN output layer (DiT paper)
            for block in self.blocks:
                nn.init.zeros_(block.adaLN_modulation[-1].weight)
                nn.init.zeros_(block.adaLN_modulation[-1].bias)
            nn.init.zeros_(self.final_adaLN[-1].weight)
            nn.init.zeros_(self.final_adaLN[-1].bias)
            nn.init.zeros_(self.final_linear.weight)
            nn.init.zeros_(self.final_linear.bias)
            nn.init.normal_(self.pos_embed, std=0.02)

        def forward(self, x: "torch.Tensor", t: "torch.Tensor") -> "torch.Tensor":
            """
            x: (B, 4, H, W) latent at H=W=resolution//8
            t: (B,) timestep indices in [0, 1000)
            returns: (B, 4, H, W) noise prediction
            """
            B, C, H, W = x.shape

            # Patchify
            h = self.patch_embed(x)         # (B, hidden, H//p, W//p)
            h = rearrange(h, 'b c h w -> b (h w) c')  # (B, N, hidden)
            h = h + self.pos_embed

            # Timestep embedding
            c = self.t_embed(t)              # (B, hidden)

            # Transformer blocks
            for block in self.blocks:
                h = block(h, c)

            # Final layer
            shift, scale = self.final_adaLN(c).chunk(2, dim=-1)
            h = self.final_norm(h) * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
            h = self.final_linear(h)         # (B, N, p*p*C*2) — predict mean+logvar

            # Unpatchify (predict noise only, take first half)
            p = self.patch_size
            h_out = h[..., :p*p*C]           # (B, N, p*p*C)
            h_out = rearrange(
                h_out, 'b (gh gw) (p1 p2 c) -> b c (gh p1) (gw p2)',
                gh=H//p, gw=W//p, p1=p, p2=p, c=C
            )
            return h_out

    # Bundle SinusoidalEmbedding as an attribute of MinimalDiT so callers that
    # only grabbed MinimalDiT can still reach it via _build_model_classes().
    MinimalDiT._SinusoidalEmbedding = SinusoidalEmbedding
    return SinusoidalEmbedding, MinimalDiT


# ---------------------------------------------------------------------------
# Module-level class exposure for test-suite imports
# (torch + einops are available in the test venv; the try-block is a no-op
#  on CUDA boxes and on this MPS dev box; it only becomes a silent skip if
#  running py_compile or --help/--dry-run in an environment without torch.)
# ---------------------------------------------------------------------------

try:
    import torch as _torch_global
    from einops import rearrange as _rearrange_global
    SinusoidalEmbedding, MinimalDiT = _build_model_classes(_torch_global, _rearrange_global)
except Exception:  # pragma: no cover  (torch/einops absent — fine for --help/--dry-run)
    SinusoidalEmbedding = None  # type: ignore[assignment,misc]
    MinimalDiT = None  # type: ignore[assignment,misc]


# ---------------------------------------------------------------------------
# VRAM probe via nvidia-smi
# ---------------------------------------------------------------------------

def get_peak_vram_gb(torch) -> float:
    """Query current GPU VRAM usage via nvidia-smi."""
    import subprocess
    try:
        out = subprocess.check_output(
            ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"],
            timeout=5
        ).decode().strip().split("\n")[0]
        return float(out) / 1024.0
    except Exception:
        # Fallback to torch
        return torch.cuda.max_memory_allocated() / 1e9


# ---------------------------------------------------------------------------
# Kernel warning check
# ---------------------------------------------------------------------------

def check_kernel_warnings(torch, device):
    """Check for known sm_89 kernel issues and report."""
    import torch.nn as nn
    import torch.nn.functional as F

    warnings = []
    cc_major, cc_minor = torch.cuda.get_device_capability(device)
    sm = cc_major * 10 + cc_minor
    print(f"  GPU compute capability: sm_{sm} (major={cc_major}, minor={cc_minor})")

    # 1. FP8 via Transformer Engine on sm_89
    try:
        import transformer_engine  # noqa
        try:
            import transformer_engine.pytorch as te
            # Try a minimal FP8 Linear
            layer = te.Linear(16, 16)
            x = torch.randn(1, 16, device=device, dtype=torch.bfloat16)
            with te.fp8_autocast(enabled=True):
                _ = layer(x)
            print("  [OK]    TransformerEngine FP8: runs on this device")
        except Exception as e:
            warnings.append(f"TransformerEngine FP8: {e}")
            print(f"  [WARN]  TransformerEngine FP8 fallback triggered: {e}")
    except ImportError:
        print("  [INFO]  TransformerEngine not installed (expected) — FP8 training path not available")
        warnings.append("TransformerEngine not installed — FP8 training path unavailable on sm_89 for diffusion")

    # 2. SDPA FlashAttention dispatch
    try:
        with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
            q = torch.randn(1, 8, 64, 64, device=device, dtype=torch.bfloat16)
            k = torch.randn(1, 8, 64, 64, device=device, dtype=torch.bfloat16)
            v = torch.randn(1, 8, 64, 64, device=device, dtype=torch.bfloat16)
            _ = F.scaled_dot_product_attention(q, k, v)
        print("  [OK]    SDPA FlashAttention (BF16): runs cleanly on sm_89")
    except Exception as e:
        warnings.append(f"SDPA FlashAttention BF16: {e}")
        print(f"  [WARN]  SDPA FlashAttention BF16 failed: {e}")

    # 3. Check torch.compile availability (used by Nitro-T)
    try:
        tiny = nn.Linear(8, 8).to(device)
        compiled = torch.compile(tiny, mode='reduce-overhead')
        x_t = torch.randn(1, 8, device=device)
        _ = compiled(x_t)
        print("  [OK]    torch.compile (reduce-overhead): works on sm_89")
    except Exception as e:
        warnings.append(f"torch.compile: {e}")
        print(f"  [WARN]  torch.compile failed: {e}")

    # 4. FP8 matmul via cuBLASLt (lower-level, Ada does have hardware FP8)
    if sm >= 89:
        print("  [INFO]  sm_89 (Ada): hardware FP8 tensor cores present but TE high-level FP8 training")
        print("          requires sm_90+ (Hopper) for stable FP8 matmul scaling. Use BF16 for training.")

    return warnings


# ---------------------------------------------------------------------------
# Binary search for max batch size
# ---------------------------------------------------------------------------

def find_max_batch(
    torch,
    model,
    resolution: int,
    target_vram_gb: float,
    device,
    dtype,
    min_batch: int = 4,
    max_batch: int = 256,
) -> int:
    """Binary search for maximum batch that fits under target_vram_gb."""
    import torch.nn.functional as F
    lo, hi = min_batch, max_batch
    best = min_batch
    lat_res = resolution // 8
    model.train()

    while lo <= hi:
        mid = (lo + hi) // 2
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats(device)
        try:
            x = torch.randn(mid, 4, lat_res, lat_res, device=device, dtype=dtype)
            t = torch.randint(0, 1000, (mid,), device=device)
            noise = torch.randn_like(x)
            with torch.autocast(device_type='cuda', dtype=dtype):
                pred = model(x, t)
            loss = F.mse_loss(pred, noise)
            loss.backward()
            model.zero_grad(set_to_none=True)
            peak_gb = torch.cuda.max_memory_allocated(device) / 1e9
            if peak_gb <= target_vram_gb:
                best = mid
                lo = mid + 1
            else:
                hi = mid - 1
        except torch.cuda.OutOfMemoryError:
            hi = mid - 1
        finally:
            torch.cuda.empty_cache()

    return best


# ---------------------------------------------------------------------------
# Tiny smoke pass for MPS / CPU
# ---------------------------------------------------------------------------

def run_smoke(torch, MinimalDiT, device, backend_name: str, args: argparse.Namespace) -> None:
    """Run a minimal 2-step forward+backward to verify model code is functional."""
    import torch.nn.functional as F

    print(f"\n{'='*70}")
    print(f"DiT Throughput Probe — SMOKE PASS ({backend_name})")
    print(f"{'='*70}")
    print(f"[WARN] Full throughput benchmark requires CUDA. Running tiny smoke (batch=1, steps=2).")
    print(f"       Backend: {backend_name} | Resolution: {args.resolution}px")
    print()

    dtype = torch.float32  # MPS/CPU: use float32 (bfloat16 may not be fully supported on all MPS)
    lat_res = args.resolution // 8

    model = MinimalDiT(
        resolution=args.resolution,
        hidden=args.hidden,
        depth=args.depth,
        heads=args.heads,
    ).to(device=device, dtype=dtype)
    n_params = sum(p.numel() for p in model.parameters()) / 1e6
    print(f"[Model] Parameters: {n_params:.1f}M  Depth: {args.depth}  Hidden: {args.hidden}")

    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
    model.train()

    smoke_steps = 2
    print(f"[Smoke] Running {smoke_steps} forward+backward steps on {backend_name}...")
    for step_idx in range(smoke_steps):
        x = torch.randn(1, 4, lat_res, lat_res, device=device, dtype=dtype)
        t = torch.randint(0, 1000, (1,), device=device)
        noise = torch.randn_like(x)
        pred = model(x, t)
        loss = F.mse_loss(pred, noise)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)
        print(f"  step {step_idx+1}/{smoke_steps}  loss={loss.item():.4f}")

    print()
    print(f"[SMOKE PASS] Model code is functional on {backend_name}. Full benchmark needs CUDA.")
    print(f"{'='*70}")


# ---------------------------------------------------------------------------
# Main benchmark (CUDA path)
# ---------------------------------------------------------------------------

def run_benchmark(
    torch,
    MinimalDiT,
    resolution: int = 256,
    target_vram_gb: float = 22.0,  # leave 2GB headroom on 24GB 4090
    n_steps: int = 100,
    warmup_steps: int = 10,
    depth: int = 18,
    hidden: int = 1024,
    heads: int = 16,
):
    import torch.nn.functional as F

    device = torch.device("cuda")
    dtype = torch.bfloat16

    gpu_name = torch.cuda.get_device_name(device)
    total_vram = torch.cuda.get_device_properties(device).total_memory / 1e9
    print(f"\n{'='*70}")
    print(f"DiT Throughput Probe — {gpu_name}")
    print(f"{'='*70}")
    print(f"Target VRAM budget: {target_vram_gb:.1f} GB / {total_vram:.1f} GB total")
    print(f"Resolution: {resolution}px | Latent: {resolution//8}x{resolution//8}")
    print(f"Precision: {dtype}")
    print()

    # --- Kernel warnings ---
    print("[Kernel Check]")
    kernel_warnings = check_kernel_warnings(torch, device)
    print()

    # --- Build model ---
    print("[Model]")
    model = MinimalDiT(
        resolution=resolution,
        hidden=hidden,
        depth=depth,
        heads=heads,
    ).to(device, dtype=dtype)
    n_params = sum(p.numel() for p in model.parameters()) / 1e6
    print(f"  Parameters: {n_params:.1f}M")
    print(f"  Depth: {depth}, Hidden: {hidden}, Heads: {heads}")
    print(f"  Patches: {model.n_patches} ({resolution//8 // 2}x{resolution//8 // 2})")
    print()

    # --- Find max batch ---
    print("[Batch Search]")
    print(f"  Binary-searching max batch under {target_vram_gb:.0f} GB VRAM...")
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
    max_batch = find_max_batch(torch, model, resolution, target_vram_gb, device, dtype)
    print(f"  Max batch size: {max_batch}")
    print()

    # --- Warmup ---
    lat_res = resolution // 8
    model.train()
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats(device)
    print(f"[Warmup — {warmup_steps} steps at batch {max_batch}]")
    for _ in range(warmup_steps):
        x = torch.randn(max_batch, 4, lat_res, lat_res, device=device, dtype=dtype)
        t = torch.randint(0, 1000, (max_batch,), device=device)
        noise = torch.randn_like(x)
        with torch.autocast(device_type='cuda', dtype=dtype):
            pred = model(x, t)
        loss = F.mse_loss(pred, noise)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)
    torch.cuda.synchronize(device)
    print("  Warmup complete.")
    print()

    # --- Timed benchmark ---
    print(f"[Benchmark — {n_steps} steps at batch {max_batch}]")
    torch.cuda.reset_peak_memory_stats(device)
    step_times = []
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)

    for step_idx in range(n_steps):
        x = torch.randn(max_batch, 4, lat_res, lat_res, device=device, dtype=dtype)
        t = torch.randint(0, 1000, (max_batch,), device=device)
        noise = torch.randn_like(x)

        start_event.record()
        with torch.autocast(device_type='cuda', dtype=dtype):
            pred = model(x, t)
        loss = F.mse_loss(pred, noise)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)
        end_event.record()
        torch.cuda.synchronize(device)

        step_ms = start_event.elapsed_time(end_event)
        step_times.append(step_ms)

    # --- Metrics ---
    peak_vram_allocated = torch.cuda.max_memory_allocated(device) / 1e9
    peak_vram_smi = get_peak_vram_gb(torch)

    step_times_ms = sorted(step_times)
    p50_ms = step_times_ms[len(step_times_ms) // 2]
    p95_ms = step_times_ms[int(len(step_times_ms) * 0.95)]
    mean_ms = sum(step_times) / len(step_times)
    steps_per_sec = 1000.0 / mean_ms
    images_per_sec = steps_per_sec * max_batch

    print()
    print(f"{'='*70}")
    print(f"RESULTS")
    print(f"{'='*70}")
    print(f"  arch             : MinimalDiT (DiT-class, AdaLN-Zero, RectFlow-compatible)")
    print(f"  params_M         : {n_params:.1f}")
    print(f"  depth/hidden     : {depth} / {hidden}")
    print(f"  resolution       : {resolution}px (latent {lat_res}x{lat_res}, {model.n_patches} patches)")
    print(f"  batch_size       : {max_batch}")
    print(f"  step_ms (mean)   : {mean_ms:.1f}")
    print(f"  step_ms (p50)    : {p50_ms:.1f}")
    print(f"  step_ms (p95)    : {p95_ms:.1f}")
    print(f"  steps/sec        : {steps_per_sec:.2f}")
    print(f"  images/sec       : {images_per_sec:.1f}")
    print(f"  peak_vram_GB     : {peak_vram_allocated:.2f} (torch.cuda)")
    print(f"  peak_vram_smi_GB : {peak_vram_smi:.2f} (nvidia-smi)")
    print()
    if kernel_warnings:
        print("  kernel_warnings:")
        for w in kernel_warnings:
            print(f"    - {w}")
    else:
        print("  kernel_warnings  : none")
    print(f"{'='*70}")

    # Machine-readable summary line for easy parsing
    print(
        f"\n[SUMMARY] arch=MinimalDiT params_M={n_params:.1f} batch={max_batch} "
        f"img_per_s={images_per_sec:.1f} step_ms={mean_ms:.1f} "
        f"peak_vram_GB={peak_vram_allocated:.2f} "
        f"kernel_warnings={len(kernel_warnings)}"
    )

    return {
        "arch": "MinimalDiT",
        "params_M": n_params,
        "batch": max_batch,
        "img_per_s": images_per_sec,
        "step_ms": mean_ms,
        "peak_vram_GB": peak_vram_allocated,
        "kernel_warnings": kernel_warnings,
    }


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    parser = _build_parser()
    args = parser.parse_args()

    # --- Validate numeric args (all paths, including --dry-run) ---
    errors = _validate_args(args)
    if errors:
        for err in errors:
            print(f"[ERROR] {err}", file=sys.stderr)
        sys.exit(1)

    # --- Dry run: no heavy imports, no allocation ---
    if args.dry_run:
        _dry_run(args)
        # _dry_run calls sys.exit(0), but be explicit:
        sys.exit(0)

    # --- Heavy imports ---
    try:
        torch, einops_mod = _import_heavy()
    except ImportError as exc:
        print(f"[ERROR] {exc}", file=sys.stderr)
        sys.exit(1)

    from einops import rearrange

    # --- Device selection ---
    device, is_cuda, backend_name = _select_device(torch)

    # --- Resolve model classes ---
    # The module-level try-block already built MinimalDiT when torch was importable.
    # If that block failed (edge case: torch appeared later), rebuild now.
    _MinimalDiT = globals().get("MinimalDiT")
    if _MinimalDiT is None:
        _, _MinimalDiT = _build_model_classes(torch, rearrange)

    if not is_cuda:
        print(
            f"[WARN] CUDA not available (backend: {backend_name}). "
            "Full throughput benchmark targets a CUDA GPU (RTX 4090). "
            "Running a tiny smoke pass instead to verify model correctness.",
            file=sys.stderr,
        )
        run_smoke(torch, _MinimalDiT, device, backend_name, args)
    else:
        run_benchmark(
            torch,
            _MinimalDiT,
            resolution=args.resolution,
            target_vram_gb=args.target_vram_gb,
            n_steps=args.steps,
            depth=args.depth,
            hidden=args.hidden,
            heads=args.heads,
        )