interference_benchmark.py

"""
Interference Benchmark
Measures how long-prompt prefill affects active short-request decode.

Scenario:
  t=0.0:  2 short requests arrive (prompt=32 tokens, output=64 tokens)
  t=0.1:  2 long requests arrive  (prompt=512 tokens, output=32 tokens)
  t=0.3:  2 more short requests arrive (prompt=32 tokens, output=64 tokens)

Question: how much do the long requests slow down the short requests?

Metrics:
  - short_tbt_mean: time between tokens for short requests (ms)
  - short_tbt_max:  worst-case TBT for shorts (decode stall indicator)
  - short_ttft:     time to first token for short requests
  - long_ttft:      time to first token for long requests
  - short_completion: total latency for short requests
"""

import csv
import gc
import statistics
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import List, Optional

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE  = torch.float16 if DEVICE == "cuda" else torch.float32
RESULTS = Path("results")
RESULTS.mkdir(exist_ok=True)


@dataclass
class Request:
    id: int
    label: str                          # "short_early", "long", "short_late"
    prompt_ids: torch.Tensor
    max_output_tokens: int
    arrival_time: float

    generated_ids: List[int] = field(default_factory=list)
    token_times: List[float] = field(default_factory=list)  # timestamp of each token
    prefill_done: bool = False
    prefill_pos: int = 0
    first_token_time: float = -1.0
    complete_time: float = -1.0
    past_key_values: Optional[object] = None
    last_tok: Optional[torch.Tensor] = None

    @property
    def prompt_len(self):
        return self.prompt_ids.shape[0]

    def done(self, max_out):
        return len(self.generated_ids) >= min(max_out, self.max_output_tokens)

    def ttft_ms(self):
        if self.first_token_time < 0: return -1.0
        return (self.first_token_time - self.arrival_time) * 1000.0

    def latency_ms(self):
        if self.complete_time < 0: return -1.0
        return (self.complete_time - self.arrival_time) * 1000.0

    def tbt_ms_list(self):
        """Time between successive tokens, in ms."""
        if len(self.token_times) < 2:
            return []
        return [(self.token_times[i] - self.token_times[i-1]) * 1000.0
                for i in range(1, len(self.token_times))]


def make_workload():
    """
    Fixed interference scenario:
    2 short at t=0.0, 2 long at t=0.1, 2 short at t=0.3
    """
    requests = []
    rid = 0

    # 2 short early
    for _ in range(2):
        requests.append(Request(
            id=rid, label="short_early",
            prompt_ids=torch.randint(0, 50000, (32,)),
            max_output_tokens=64,
            arrival_time=0.0,
        ))
        rid += 1

    # 2 long
    for _ in range(2):
        requests.append(Request(
            id=rid, label="long",
            prompt_ids=torch.randint(0, 50000, (512,)),
            max_output_tokens=32,
            arrival_time=0.1,
        ))
        rid += 1

    # 2 short late
    for _ in range(2):
        requests.append(Request(
            id=rid, label="short_late",
            prompt_ids=torch.randint(0, 50000, (32,)),
            max_output_tokens=64,
            arrival_time=0.3,
        ))
        rid += 1

    return requests


def cleanup():
    gc.collect()
    if DEVICE == "cuda":
        try: torch.cuda.empty_cache()
        except: pass


@torch.no_grad()
def do_prefill(model, req, start, end):
    assert end > start
    chunk = req.prompt_ids[start:end].unsqueeze(0).to(DEVICE)
    t0 = time.perf_counter()
    if req.past_key_values is None:
        out = model(chunk, use_cache=True)
    else:
        out = model(chunk, use_cache=True, past_key_values=req.past_key_values)
    torch.cuda.synchronize()
    elapsed = time.perf_counter() - t0
    req.past_key_values = out.past_key_values
    req.prefill_pos = end
    req.last_tok = out.logits[:, -1, :].argmax(dim=-1, keepdim=True)
    del chunk, out
    return elapsed


@torch.no_grad()
def do_decode_step(model, req, now):
    t0 = time.perf_counter()
    out = model(req.last_tok, use_cache=True, past_key_values=req.past_key_values)
    torch.cuda.synchronize()
    elapsed = time.perf_counter() - t0
    req.past_key_values = out.past_key_values
    req.last_tok = out.logits[:, -1, :].argmax(dim=-1, keepdim=True)
    req.generated_ids.append(req.last_tok.item())
    req.token_times.append(now + elapsed)
    del out
    return elapsed


def free_req(req):
    if req.past_key_values is not None:
        del req.past_key_values
        req.past_key_values = None
    if req.last_tok is not None:
        del req.last_tok
        req.last_tok = None


# ── FCFS ─────────────────────────────────────────────────────────────────────

class FCFSServer:
    def __init__(self, model):
        self.model = model

    def process(self, requests):
        queue = sorted(requests, key=lambda r: r.arrival_time)
        now = 0.0
        for req in queue:
            now = max(now, req.arrival_time)
            elapsed = do_prefill(self.model, req, 0, req.prompt_len)
            now += elapsed
            req.generated_ids.append(req.last_tok.item())
            req.token_times.append(now)
            req.first_token_time = now
            req.prefill_done = True
            for _ in range(req.max_output_tokens - 1):
                elapsed = do_decode_step(self.model, req, now)
                now += elapsed
            req.complete_time = now
            free_req(req)
        cleanup()
        return requests


# ── Continuous Batching with eager prefill ────────────────────────────────────

class EagerContBatchServer:
    """
    Admits requests immediately with full prefill.
    This blocks decode for active requests during prefill of new ones.
    """

    def __init__(self, model, max_batch_size=6):
        self.model = model
        self.max_bs = max_batch_size

    def process(self, requests):
        pending = sorted(requests, key=lambda r: r.arrival_time)
        active  = []
        completed = []
        now = 0.0
        ptr = 0

        def admit(current_now):
            nonlocal ptr
            total_e = 0.0
            while (ptr < len(pending) and
                   pending[ptr].arrival_time <= current_now + total_e and
                   len(active) < self.max_bs):
                req = pending[ptr]
                ptr += 1
                # FULL prefill — this blocks all active decode requests
                e = do_prefill(self.model, req, 0, req.prompt_len)
                total_e += e
                req.generated_ids.append(req.last_tok.item())
                req.first_token_time = current_now + total_e
                req.token_times.append(current_now + total_e)
                req.prefill_done = True
                active.append(req)
            return total_e

        now += admit(now)

        while ptr < len(pending) or active:
            if not active:
                if ptr < len(pending):
                    now = pending[ptr].arrival_time
                    now += admit(now)
                continue

            # Remove completed
            for req in list(active):
                if req.done(req.max_output_tokens):
                    req.complete_time = now
                    active.remove(req)
                    completed.append(req)
                    free_req(req)

            if not active:
                continue

            # Decode one step for each active request (sequential)
            max_e = 0.0
            for req in active:
                e = do_decode_step(self.model, req, now)
                max_e = max(max_e, e)
            now += max_e

            # Check completions
            for req in list(active):
                if req.done(req.max_output_tokens):
                    req.complete_time = now
                    active.remove(req)
                    completed.append(req)
                    free_req(req)

            # Admit new arrivals (blocks decode for existing!)
            now += admit(now)

        cleanup()
        return requests


# ── ChunkedPrefill with decode interleaving ───────────────────────────────────

class ChunkedPrefillServer:
    """
    Prefill split into chunks. Decode runs between chunks.
    Active decode requests are NOT blocked during prefill of new ones.
    """

    def __init__(self, model, chunk_size=64, max_batch_size=6):
        self.model = model
        self.chunk = chunk_size
        self.max_bs = max_batch_size

    def process(self, requests):
        pending    = sorted(requests, key=lambda r: r.arrival_time)
        prefilling = []
        decoding   = []
        completed  = []
        now = 0.0
        ptr = 0

        def admit():
            nonlocal ptr
            while (ptr < len(pending) and
                   pending[ptr].arrival_time <= now and
                   len(prefilling) + len(decoding) < self.max_bs):
                prefilling.append(pending[ptr])
                ptr += 1

        admit()

        while ptr < len(pending) or prefilling or decoding:
            if not prefilling and not decoding:
                if ptr < len(pending):
                    now = pending[ptr].arrival_time
                    admit()
                continue

            # ONE chunk of prefill per prefilling request
            for req in list(prefilling):
                start = req.prefill_pos
                end   = min(start + self.chunk, req.prompt_len)
                if start >= req.prompt_len:
                    prefilling.remove(req)
                    continue
                elapsed = do_prefill(self.model, req, start, end)
                now += elapsed
                if end >= req.prompt_len:
                    req.generated_ids.append(req.last_tok.item())
                    req.first_token_time = now
                    req.token_times.append(now)
                    req.prefill_done = True
                    prefilling.remove(req)
                    decoding.append(req)

            # Decode ONE step for ALL decoding requests
            for req in list(decoding):
                if req.done(req.max_output_tokens):
                    req.complete_time = now
                    decoding.remove(req)
                    completed.append(req)
                    free_req(req)
                    continue

            max_e = 0.0
            for req in decoding:
                e = do_decode_step(self.model, req, now)
                max_e = max(max_e, e)
            if max_e > 0:
                now += max_e

            for req in list(decoding):
                if req.done(req.max_output_tokens):
                    req.complete_time = now
                    decoding.remove(req)
                    completed.append(req)
                    free_req(req)

            admit()

        for req in decoding:
            req.complete_time = now
            completed.append(req)
            free_req(req)

        cleanup()
        return requests


# ── Analysis ──────────────────────────────────────────────────────────────────

def analyze(requests, server_name):
    """Compute per-label metrics including TBT."""

    by_label = {}
    for req in requests:
        lbl = req.label
        if lbl not in by_label:
            by_label[lbl] = []
        by_label[lbl].append(req)

    row = {"server": server_name}

    for lbl, reqs in by_label.items():
        ttfts = [r.ttft_ms() for r in reqs if r.ttft_ms() >= 0]
        lats  = [r.latency_ms() for r in reqs if r.latency_ms() >= 0]

        all_tbt = []
        for r in reqs:
            all_tbt.extend(r.tbt_ms_list())

        row[f"{lbl}_ttft_mean_ms"]  = round(statistics.mean(ttfts), 1) if ttfts else 0
        row[f"{lbl}_ttft_max_ms"]   = round(max(ttfts), 1) if ttfts else 0
        row[f"{lbl}_latency_mean_ms"] = round(statistics.mean(lats), 1) if lats else 0

        if all_tbt:
            row[f"{lbl}_tbt_mean_ms"]   = round(statistics.mean(all_tbt), 2)
            row[f"{lbl}_tbt_p50_ms"]    = round(sorted(all_tbt)[len(all_tbt)//2], 2)
            row[f"{lbl}_tbt_p95_ms"]    = round(sorted(all_tbt)[int(0.95*(len(all_tbt)-1))], 2)
            row[f"{lbl}_tbt_max_ms"]    = round(max(all_tbt), 2)
            row[f"{lbl}_tbt_stall_count"] = sum(1 for t in all_tbt if t > 50)
        else:
            row[f"{lbl}_tbt_mean_ms"]   = 0
            row[f"{lbl}_tbt_p50_ms"]    = 0
            row[f"{lbl}_tbt_p95_ms"]    = 0
            row[f"{lbl}_tbt_max_ms"]    = 0
            row[f"{lbl}_tbt_stall_count"] = 0

    # Total throughput
    all_completed = [r for r in requests if r.complete_time > 0]
    if all_completed:
        dur = max(r.complete_time for r in all_completed) - \
              min(r.arrival_time for r in requests)
        total_tok = sum(len(r.generated_ids) for r in all_completed)
        row["throughput_rps"] = round(len(all_completed) / dur, 3) if dur > 0 else 0
        row["throughput_tps"] = round(total_tok / dur, 2) if dur > 0 else 0
        row["wall_s"] = round(dur, 3)

    return row


def print_result(row):
    server = row["server"]
    print(f"\n  [{server}]")

    for lbl in ["short_early", "long", "short_late"]:
        ttft = row.get(f"{lbl}_ttft_mean_ms", 0)
        tbt  = row.get(f"{lbl}_tbt_mean_ms", 0)
        tbt_max = row.get(f"{lbl}_tbt_max_ms", 0)
        stalls = row.get(f"{lbl}_tbt_stall_count", 0)
        lat  = row.get(f"{lbl}_latency_mean_ms", 0)

        print(f"    {lbl:14s}: TTFT={ttft:>7.1f}ms  "
              f"TBT mean={tbt:>7.2f}ms  "
              f"TBT max={tbt_max:>7.1f}ms  "
              f"stalls={stalls:>3d}  "
              f"latency={lat:>7.0f}ms")

    print(f"    throughput: {row.get('throughput_rps',0):.2f} rps  "
          f"{row.get('throughput_tps',0):.1f} tps")


def save_csv(path, rows):
    if not rows: return
    all_keys = []
    seen = set()
    for row in rows:
        for k in row:
            if k not in seen:
                all_keys.append(k)
                seen.add(k)
    with open(path, "w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=all_keys)
        writer.writeheader()
        for row in rows:
            writer.writerow({k: row.get(k, "") for k in all_keys})
    print(f"\nSaved -> {path}")


def main():
    print("=" * 72)
    print("INTERFERENCE BENCHMARK")
    print("How does long-prompt prefill affect active short-request decode?")
    print(f"Device: {DEVICE}  Dtype: {DTYPE}")
    if DEVICE == "cuda":
        print(f"GPU: {torch.cuda.get_device_name(0)}")
    print("=" * 72)

    print("\nScenario:")
    print("  t=0.0: 2 short requests (32 tok prompt, 64 tok output)")
    print("  t=0.1: 2 long  requests (512 tok prompt, 32 tok output)")
    print("  t=0.3: 2 short requests (32 tok prompt, 64 tok output)")
    print("\nQuestion: who protects short-request decode from long-prompt interference?")

    print("\nLoading gpt2...")
    tok = AutoTokenizer.from_pretrained("gpt2", local_files_only=True)
    if tok.pad_token is None: tok.pad_token = tok.eos_token
    model = AutoModelForCausalLM.from_pretrained(
        "gpt2", torch_dtype=DTYPE, local_files_only=True).to(DEVICE).eval()
    n = sum(p.numel() for p in model.parameters()) / 1e6
    print(f"  {n:.0f}M params on {DEVICE}")

    all_rows = []

    # Run each strategy with the SAME workload
    for server_name, ServerClass, kwargs in [
        ("FCFS",
         FCFSServer, {}),
        ("EagerContBatch",
         EagerContBatchServer, {"max_batch_size": 6}),
        ("ChunkedPrefill_c64",
         ChunkedPrefillServer, {"chunk_size": 64, "max_batch_size": 6}),
        ("ChunkedPrefill_c32",
         ChunkedPrefillServer, {"chunk_size": 32, "max_batch_size": 6}),
        ("ChunkedPrefill_c16",
         ChunkedPrefillServer, {"chunk_size": 16, "max_batch_size": 6}),
    ]:
        # Generate fresh workload each time (same seed = same requests)
        torch.manual_seed(42)
        requests = make_workload()

        server = ServerClass(model, **kwargs)

        wall_t0 = time.perf_counter()
        server.process(requests)
        wall_s = time.perf_counter() - wall_t0

        row = analyze(requests, server_name)
        row["wall_s"] = round(wall_s, 3)
        all_rows.append(row)

        print_result(row)

    save_csv(RESULTS / "interference_benchmark.csv", all_rows)

    # ── Summary comparison ────────────────────────────────────────────────────
    print("\n" + "=" * 72)
    print("INTERFERENCE RESILIENCE COMPARISON")
    print("=" * 72)

    print(f"\n-- Short-early TBT (requests already in decode when longs arrive) --")
    print(f"{'Server':<24s} {'TBT mean':>10s} {'TBT max':>10s} {'Stalls':>7s}")
    print("-" * 55)
    for r in all_rows:
        print(f"{r['server']:<24s} "
              f"{r.get('short_early_tbt_mean_ms',0):>9.2f}ms "
              f"{r.get('short_early_tbt_max_ms',0):>9.1f}ms "
              f"{r.get('short_early_tbt_stall_count',0):>7d}")

    print(f"\n-- Long request TTFT --")
    print(f"{'Server':<24s} {'TTFT mean':>10s}")
    print("-" * 36)
    for r in all_rows:
        print(f"{r['server']:<24s} "
              f"{r.get('long_ttft_mean_ms',0):>9.1f}ms")

    print(f"\n-- Short-late TTFT (arrive after longs) --")
    print(f"{'Server':<24s} {'TTFT mean':>10s}")
    print("-" * 36)
    for r in all_rows:
        print(f"{r['server']:<24s} "
              f"{r.get('short_late_ttft_mean_ms',0):>9.1f}ms")

    # Key findings
    print(f"\n-- Key findings --")

    fcfs = next(r for r in all_rows if r["server"] == "FCFS")
    eager = next(r for r in all_rows if r["server"] == "EagerContBatch")
    chunked64 = next(r for r in all_rows if r["server"] == "ChunkedPrefill_c64")

    se_tbt_fcfs  = fcfs.get("short_early_tbt_mean_ms", 0)
    se_tbt_eager = eager.get("short_early_tbt_mean_ms", 0)
    se_tbt_chunk = chunked64.get("short_early_tbt_mean_ms", 0)

    se_max_fcfs  = fcfs.get("short_early_tbt_max_ms", 0)
    se_max_eager = eager.get("short_early_tbt_max_ms", 0)
    se_max_chunk = chunked64.get("short_early_tbt_max_ms", 0)

    print(f"\n  1. Short-early TBT mean:")
    print(f"     FCFS:           {se_tbt_fcfs:.2f}ms")
    print(f"     EagerContBatch: {se_tbt_eager:.2f}ms")
    print(f"     ChunkedPrefill: {se_tbt_chunk:.2f}ms")

    if se_max_eager > 0 and se_max_chunk > 0:
        ratio = se_max_eager / se_max_chunk
        print(f"\n  2. Short-early TBT max (decode stall):")
        print(f"     EagerContBatch: {se_max_eager:.1f}ms")
        print(f"     ChunkedPrefill: {se_max_chunk:.1f}ms")
        print(f"     ChunkedPrefill reduces worst-case stall by {ratio:.1f}x")

    long_ttft_eager = eager.get("long_ttft_mean_ms", 0)
    long_ttft_chunk = chunked64.get("long_ttft_mean_ms", 0)
    print(f"\n  3. Long request TTFT:")
    print(f"     EagerContBatch: {long_ttft_eager:.1f}ms (full prefill immediately)")
    print(f"     ChunkedPrefill: {long_ttft_chunk:.1f}ms (split into chunks)")
    if long_ttft_chunk > 0:
        print(f"     ChunkedPrefill trades {long_ttft_chunk - long_ttft_eager:.1f}ms long TTFT"
              f" for better short decode protection")


if __name__ == "__main__":
    main()